CUDA之Warp Shuffle详解

之前我们有介绍shared Memory对于提高性能的好处，在CC3.0以上，支持了shuffle指令，允许thread直接读其他thread的寄存器值，只要两个thread在 同一个warp中，这种比通过shared Memory进行thread间的通讯效果更好，latency更低，同时也不消耗额外的内存资源来执行数据交换。主要包含如下API：

[cpp]  view plain  copy

1. int __shfl(int var, int srcLane, int width=warpSize);   

[cpp]  view plain  copy

1. int __shfl_up(int var, unsigned int delta, int width=warpSize);   

[cpp]  view plain  copy

1. int __shfl_down(int var, unsigned int delta, int width=warpSize);   

[cpp]  view plain  copy

1. int __shfl_xor(int var, int laneMask, int width=warpSize);   

[cpp]  view plain  copy

1. float __shfl(float var, int srcLane, int width=warpSize);   

[cpp]  view plain  copy

1. float __shfl_up(float var, unsigned int delta, int width=warpSize);   

[cpp]  view plain  copy

1. float __shfl_down(float var, unsigned int delta, int width=warpSize);   

[cpp]  view plain  copy

1. float __shfl_xor(float var, int laneMask, int width=warpSize);   

[cpp]  view plain  copy

1. half __shfl(half var, int srcLane, int width=warpSize);   

[cpp]  view plain  copy

1. half __shfl_up(half var, unsigned int delta, int width=warpSize);   

[cpp]  view plain  copy

1. half __shfl_down(half var, unsigned int delta, int width=warpSize);   

[cpp]  view plain  copy

1. half __shfl_xor(half var, int laneMask, int width=warpSize);  

这里介绍warp中的一个概念lane，一个lane就是一个warp中的一个thread，每个lane在同一个warp中由lane索引唯一确定，因此其范围为[0,31]。在一个一维的block中，可以通过下面两个公式计算索引：

laneID = threadIdx.x % 32

warpID = threadIdx.x / 32

例如，在同一个block中的thread1和33拥有相同的lane索引1。

Variants of the Warp Shuffle Instruction

有两种设置shuffle的指令：一种针对整型变量，另一种针对浮点型变量。每种设置都包含四种shuffle指令变量。为了交换整型变量，使用过如下函数：

int __shfl(int var, int srcLane, int width=warpSize);

该函数的作用是将var的值返回给同一个warp中lane索引为srcLane的thread。可选参数width可以设置为2的n次幂，n属于[1,5]。

eg：如果shuffle指令如下：

int y = shfl(x, 3, 16);

则，thread0到thread15会获取thread3的数据x，thread16到thread31会从thread19获取数据x。

当传送到shfl的lane索引相同时，该指令会执行一次广播操作，如下所示：

 

另一种使用shuffle的形式如下：

int __shfl_up(int var, unsigned int delta, int width=warpSize)

该函数通过使用调用方的thread的lane索引减去delta来计算源thread的lane索引。这样源thread的相应数据就会返回给调用方，这样，warp中最开始delta个的thread不会改变，如下所示：

 

第三种shuffle指令形式如下：

int __shfl_down(int var, unsigned int delta, int width=warpSize)

该格式是相对__shfl_down来说的，具体形式如下图所示：

 

最后一种shuffle指令格式如下：

int __shfl_xor(int var, int laneMask, int width=warpSize)

这次不是加减操作，而是同laneMask做抑或操作，具体形式如下图所示：

 

所有这些提及的shuffle函数也都支持单精度浮点值，只需要将int换成float就行，除此外，和整型的使用方法完全一样。

我们这里以reduction为例，看一下相比于使用shared memory进行通信的性能差异。

算法背景：为了简单起见，我们计算每32个int型变量元素的元素和。假设一个数组包含n个元素（e.g. n = 1 << 20），每32个元素计算一个和，则输出结果为n/32个int型变量。在编程中，block的大小就是32（刚好是一个warp），grid的大小是n / 32。

第一种，利用shared memory进行reduction：

[cpp]  view plain  copy

1. __global__ void reduce0(int *dst, int *src, const int n) {  
2.     __shared__ int sdata[WARP_SIZE*2];  
3.     int tidGlobal = threadIdx.x + blockDim.x * blockIdx.x;  
4.     int tidLocal = threadIdx.x;  
5.   
6.     sdata[tidLocal] = src[tidGlobal];  
7.   
8.     //actually the sync is no need here because only a warp exists in a block.  
9.     __syncthreads();  
10.   
11.     //only 32 threads in a block  
12.     //reduce in a warp  
13.     if (tidLocal < 32)  
14.         sdata[tidLocal] += sdata[tidLocal+16];  
15.     __syncthreads();  
16.     if (tidLocal < 32)  
17.         sdata[tidLocal] += sdata[tidLocal+8];  
18.     __syncthreads();  
19.     if (tidLocal < 32)  
20.         sdata[tidLocal] += sdata[tidLocal+4];  
21.     __syncthreads();  
22.     if (tidLocal < 32)  
23.         sdata[tidLocal] += sdata[tidLocal+2];  
24.     __syncthreads();  
25.     if (tidLocal < 32)  
26.         sdata[tidLocal] += sdata[tidLocal+1];  
27.     __syncthreads();  
28.   
29.     if (tidLocal == 0)  
30.         dst[blockIdx.x] = sdata[0];  
31.   
32. }  

几点说明：

1. 为了使warp内没有分支，32个线程都做加法操作（多分配点shared memory空间即可）。
2. 一个warp内的32个线程执行是同步的，因此不用担心写后读的错误。
3. 其实，一个block内只有一个warp，因此，所有的同步函数在这里都可以省略，条件语句if(tidLocal < 32)也可以省略。

第二种，利用shuffle进行通信：

[cpp]  view plain  copy

1. __global__ void reduce1(int *dst, int *src, const int n) {  
2.     int tidGlobal = threadIdx.x + blockDim.x * blockIdx.x;  
3.     int tidLocal = threadIdx.x;  
4.   
5.     int sum = src[tidGlobal];  
6.     //actually the sync is no need here because only a warp exists in a block.  
7.     __syncthreads();  
8.   
9.     for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) {  
10.         sum += __shfl_down(sum, offset);  
11.     }  
12.   
13.     if (tidLocal == 0)  
14.         dst[blockIdx.x] = sum;  
15.   
16. }  

几点说明：

1. 我们利用shuffle来做warp内的通信，因此没有用到shared memory。
2. 关于shuffle的操作含义，可以参考"cuda programming guide".

性能测试：

利用nvvp，我们来分析一下两个kernel的执行时间：

[cpp]  view plain  copy

1. ==31758== NVPROF is profiling process 31758, command: ./a.out  
2. Device 0: "Tesla K20c"  
3. check right!  
4. check right!  

[cpp]  view plain  copy

1. ==31758== Profiling application: ./a.out  
2. ==31758== Profiling result:  
3. Time(%)      Time     Calls       Avg       Min       Max  Name  
4.  80.87%  2.5935ms         1  2.5935ms  2.5935ms  2.5935ms  [CUDA memcpy HtoD]  
5.   8.07%  258.76us         1  258.76us  258.76us  258.76us  reduce0(int*, int*, int)  
6.   6.02%  192.90us         1  192.90us  192.90us  192.90us  reduce1(int*, int*, int)  
7.   5.04%  161.73us         2  80.866us  80.866us  80.866us  [CUDA memcpy DtoH]  

我们可以看到，在这个小例子中，使用shuffle可以提升性能25%左右。除了可以利用shuffle来做warp内的reduction操作，还可以进行scan，broadcast等操作。