CUDA之Warp Shuffle详解

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

这里介绍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：

几点说明：

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

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

几点说明：

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

性能测试：

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

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