cuda Reduction

有一段时间没更了,之前一直在想记录博客的意义何在,也没想出个所以然,但是不能三分钟热度,所以还是把之前想发的存货找一找

这段时间在准备工作的事情,复习了下cuda,cuda的reduction是非常常见的,而众多的方法中尤以warp shfl的方法最为简便和高效,这里整理了下代码,简述思路,供自己以后查阅,也欢迎大家交流

  • 顺便一提,其实最近已经有了采用Tensor Core来进行reduction和scan操作的工作出现了,可以参考论文《Accelerating Reduction and Scan Using Tensor Core Units》,取得了比cub库更好的效果。期待随着WMMA API的进化,能让Tensor Core为除了MatMul以外的算法提供更好的计算效果。 *

言归正传,基于warp shuffle 的reduction分为三部分:warp粒度的reduction,block粒度的reduction和整个device的reduction

warp Reduction

其中warp粒度的reduction比较简洁,但也是整个reduction部分的核心,代码如下。
warpSize就是warp的大小,一般为32,这里是一个device函数,虽然代码简单,但是可以计算得到这32个线程的和

1
2
3
4
5
6
7
8
9
#define FULL_MASK 0xffffffff
#define WARPSIZE 32

__device__ int warp_shfl(int val) {
for(int offset=WARPSIZE;offset>0;offset>>=1){
val = __shfl_down_sync(FULL_MASK,val,offset,WARPSIZE);
}
return val
}

block reduction

block的reduction 则会把每个warp计算的结果再进行一次reduction,由于当前卡的blockSize最大均不超过1024=32*32,所以我们可以将计算的结果在进行一次warpReduction得到计算的结果。

当然在进行block reduction之前我们需要先把结果存储到shared memory中做一个暂存

若每个线程负责多个元素(元素数目较多时),则需要先将每个线程负责的元素做一个加法,代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
__global__ blockReduction(int *in,int *out,int n){
int tx = threadIdx.x;
int idx = blockIdx.x*blockDim.x+tx;
int nThread = blockDim.x*gridDim.x;
int sum = 0;

__shared__ int sh[32];
int laneid = tx &0x1f;
int wid = tx >> 5;

//每个线程负责的元素先加起来
for(int i=0;i<n;i+=nThread){
sum+=in[i];
}
//warp shfl
sum = warp_shfl(sum);

if(laneid==0) sh[wid]=sum;
__syncthreads();

sum = tx<blockDim.x/WARPSIZE?sh[laneid]:0;

if(wid==0) sum=warp_shfl(sum);
//结果存到全局内存
if(tx==0) out[blockIdx.x]=sum;
}

device reduction

block内部由于有共享内存,做reduction比较方便,block之间一般采用将结果存到全局内存(上面代码的最后一行),在用单一block的进行进一步规约的形式,这里不再贴出代码,跟上面的形式大同小异

以上是关于reduction的一些实现,更快速地利用张量核心的计算见本文最开始提到的论文,欢迎大家交流讨论。