束内规约与块内规约问题

写在前面:规约问题在 cuda 编程中应用非常广泛,笔者最近在研究 faster transformer 源码,趁此机会结合 nivida 官方的代码对规约手段进行总结。
1 应用背景 关于规约的定义,相信能读到这篇文章的读者都不陌生,笔者在早期的文章中也介绍过一些规约方法,基本思想都是折半规约,主要应用于较大元素规模的向量规约,有兴趣的读者可以移步【cuda编程】cuda编程中的并行规约问题。
本文要介绍的规约场景与之前有所不同,主要应用于矩阵规约,也就是说本文假设的输入变量的维度是 2 维的,形状为 [batch_size, hidden_units],规约之后的输出变量形状为 [batch_size, ]。
接下来,本文将以规约求和为例介绍两种规约方式:束内规约、块内规约。
2 束内规约 束内规约,也就是在一个线程束内对某个变量进行规约。我们知道 cuda 架构下指令是以线程束(相邻的 32 个线程)为基本单元执行的,线程束内也可以通过束内洗牌指令进行通信,所以这提供了一个很好的束内规约思路。下面是 nvidia 提供的基础的一个规约设备函数。
template __inline__ __device__t warpreducesum(t val){ for(int mask = 16; mask > 0; mask >>= 1) val += __shfl_xor_sync(final_mask, val, mask, 32); return val;} 这个设备函数可以求出当前线程所在线程束的指定变量的规约和,原理涉及洗牌指令的计算逻辑,不再赘述。
当矩阵宽度 hidden_units 较小时,通常可以使用一个 warp 处理一行数据,一个 block 内可以处理多行数据,笔者给出具体的核函数如下:
// 一个 warp 处理一行数据template__global__ void matrix2dwarpreducesum(const t* inp, t*out, const uint32_t hidden_units) { uint32_t tid = threadidx.x; uint32_t lane_id = tid % 32; uint32_t warp_id = tid / 32; uint32_t warp_num = blockdim.x / 32; uint32_t offset = blockidx.x * warp_num * hidden_units + warp_id * hidden_units; t val = 0.0f; for (uint32_t i=lane_id; i(d_x, d_y, hidden_units);} 先确定 block_size,这里笔者直接取 128,由于是一个 warp 处理一行数据,所以一个 block 可以处理 warp_num 行数据,总共需要 grid_size 个 block。
核函数内部首先计算当前线程所在的 warp 编号 warp_id 用来定位当前处理元素在哪一行,然后确定线程在 warp 内的编号 lane_id 用来定位该线程具体处理那些元素。由于矩阵宽度 hidden_units 实际肯定还是比 32 大的,所以不可能说一个线程只处理一个元素,因此每个线程会处理多个元素,步长为 32,例如当 hidden_units 为 128 时,lane_id = 0 的线程将处理位置为 0、32、64、96 的四个元素,lane_id = 1 的线程将处理位置为 1、33、65、97 的四个元素,以此类推,这个计算过程是没有并行的。循环计算一轮后,对线程束内每个线程的 val 进行束内规约就可以得到一行元素的规约和。
3 块内规约 块内规约,就是在一个线程块内求规约值,通常块内规约会通过束内规约来实现,以下是 nvidia 提供的一个块内规约设备函数。
template __inline__ __device__t blockreducesum(t val){ static __shared__ t shared[32]; int lane = threadidx.x & 0x1f; int wid = threadidx.x >> 5; val = warpreducesum(val); if(lane == 0) shared[wid] = val; __syncthreads(); val = (threadidx.x > 5 )) ? shared[lane] : (t)0.0f; val = warpreducesum(val); return val;} 规约思路分为两步,首先通过束内规约求出当前线程所在 warp 的规约值,存入 shared 中,然后把 warpsum 赋值给 threadidx.x 小于 32 的线程内的变量 val,这 32 个线程正好也在一个线程束内,然后再执行一次束内规约就得到块内规约值,计算思路非常巧妙。
另外针对块内规约的问题,官方 cub 库其实提供了 api,开发者可以导入头文件 cub/cub.cuh 后直接使用,注意低版本的 cuda 不支持此 api。我们来看下 api 的调用方式。
#include templatestruct sumop { __device__ __forceinline__ t operator()(const t& a, const t& b) const { return a + b; }};template