这是理解__shfl_down_sync()最好的方法。假设 Warp 中 32 个线程Lane0 lane(0) Lane1 lane(1) Lane2 lane(2) ... Lane31 lane(31)这里的lane(i)表示线程 i 最初持有的数据。例如vallane(threadIdx.x);那么val: 0 1 2 3 4 5 6 7 ... 31第1轮offset16执行val__shfl_down_sync(mask,val,16);含义Lane0 Lane16 Lane1 Lane17 Lane2 Lane18 ... Lane15 Lane31结果Lane0 lane(0)lane(16) Lane1 lane(1)lane(17) Lane2 lane(2)lane(18) ... Lane15 lane(15)lane(31) Lane16~31 不重要此时前16个线程已经各自保存了2个元素的和。第2轮offset8执行val__shfl_down_sync(mask,val,8);例如Lane0 Lane8而 Lane8 当前已经不是原始值了Lane8 lane(8)lane(24)因此Lane0 (lane(0)lane(16)) (lane(8)lane(24))即Lane0 lane(0)lane(8)lane(16)lane(24)继续展开Lane1 lane(1)lane(9)lane(17)lane(25) Lane2 lane(2)lane(10)lane(18)lane(26) ... Lane7 lane(7)lane(15)lane(23)lane(31)此时Lane0~7 每个线程拥有4个元素之和第3轮offset4执行val__shfl_down_sync(mask,val,4);Lane0Lane0 Lane4而Lane4 lane(4)lane(12)lane(20)lane(28)因此Lane0 lane(0)lane(8)lane(16)lane(24) lane(4)lane(12)lane(20)lane(28)整理Lane0 lane(0)lane(4)lane(8)lane(12) lane(16)lane(20)lane(24)lane(28)已经包含8个元素同理Lane1 1,5,9,13,17,21,25,29 Lane2 2,6,10,14,18,22,26,30 Lane3 3,7,11,15,19,23,27,31第4轮offset2执行val__shfl_down_sync(mask,val,2);Lane0Lane0 Lane2Lane2 当前拥有2,6,10,14,18,22,26,30因此Lane0 0,4,8,12,16,20,24,28 2,6,10,14,18,22,26,30得到Lane0 0,2,4,6,8,10,12,14 16,18,20,22,24,26,28,30共16个元素第5轮offset1执行val__shfl_down_sync(mask,val,1);Lane0Lane0 Lane1Lane1 当前拥有1,3,5,7,9,11,13,15 17,19,21,23,25,27,29,31于是Lane0 0,2,4,6,8,10,12,14 16,18,20,22,24,26,28,30 1,3,5,7,9,11,13,15 17,19,21,23,25,27,29,31最终Lane0 lane(0)lane(1)lane(2)...lane(31)即 Warp 所有线程数据之和。用树来理解第一轮0←16 1←17 2←18 ... 15←31第二轮0←8 1←9 ... 7←15第三轮0←4 1←5 2←6 3←7第四轮0←2 1←3第五轮0←1画成树Lane0 / \ Lane0 Lane16 / \ Lane0 Lane8 / \ Lane0 Lane4 / \ Lane0 Lane2 / \ 0 1实际上是32 ↓ 16组 ↓ 8组 ↓ 4组 ↓ 2组 ↓ 1组每轮把参与求和的数据量扩大一倍offset16 → 每线程2个数 offset8 → 每线程4个数 offset4 → 每线程8个数 offset2 → 每线程16个数 offset1 → 每线程32个数所以你可以记住一个规律第 k 轮结束后 Lane0 保存了 2^k 个 lane 的部分和最后一轮结束Lane0 保存 32 个 lane 的总和这就是__shfl_down_sync()Warp Reduction 的本质。为什么不用个函数直接把所有的lane值相加这是个非常好的问题。你可以理解为GPU 硬件只提供了“从另一个 lane 读取寄存器值”的能力Shuffle并没有提供“把整个 Warp 的所有寄存器自动求和”的通用指令。CPU 的情况例如 CPU 有cab;对应一条加法指令ADD甚至还有POPCNT直接统计二进制中 1 的个数。因为这些操作非常常见CPU 专门做了硬件支持。GPU Warp 内部Warp 有Lane0 Lane1 Lane2 ... Lane31每个 Lane 都有自己的寄存器R0 R1 R2 ... R31硬件提供的是一个 Crossbar 网络----------- Lane0 ---| | Lane1 ---| Crossbar | Lane2 ---| | ... -----------因此支持__shfl_down_sync(...)本质把 Lane17 的寄存器值送给 Lane1或者把 Lane5 的寄存器值送给 Lane2这种任意线程间的数据交换。为什么不直接做 Warp Sum假设增加一个硬件指令warp_sum(val);那么硬件需要支持32个输入 ↓ 求和树 ↓ 输出结果类似 / \ / \ / \ ...对于32输入 64输入AMD Wavefront 128输入未来GPU都要专门设计电路。而 Shuffle 更通用__shfl_down_sync()__shfl_up_sync()__shfl_sync()__shfl_xor_sync()既可以实现SumMaxMinScanPrefix SumFFTMatrix 运算硬件只需要提供Warp Crossbar即可。实际上新 GPU 已经有了近几年 NVIDIA 已经发现Warp Reduction太常用了。因此增加了专门指令__reduce_add_sync()例如unsignedmask0xffffffff;intsum__reduce_add_sync(mask,val);效果Warp内所有线程的 val 求和等价于for(offset16;offset0;offset1){val__shfl_down_sync(mask,val,offset);}但这是较新的 GPU 架构才支持。CUDA Cooperative Groups还有更高级的写法#includecooperative_groups.hnamespacecgcooperative_groups;autowarpcg::tiled_partition32(cg::this_thread_block());intsumcg::reduce(warp,val,cg::plusint());看起来就是warp.reduce()实际上底层还是shuffle ↓ reduction treeAMD HIP 为什么仍然常见 __shfl_down因为NVIDIA Warp 32 AMD Wavefront 64不同架构支持情况不同。为了兼容性 性能 可移植性很多 HPC 代码仍然写for(intoffsetwarpSize/2;offset0;offset1){val__shfl_down(val,offset);}这样CUDA 能跑HIP 能跑不依赖新硬件更深层原因GPU 的设计哲学是提供最基础的数据交换能力 让程序员组合出各种算法而不是提供大量固定功能指令所以硬件提供Shuffle你自己实现Sum Max Min Scan Sort FFT这样硬件更简单、用途更广。因此历史上没有一个通用的warp_sum()函数。不过在新 CUDA 中实际上已经有__reduce_add_sync()只是它出现得比较晚而__shfl_down_sync()已经成为 GPU 编程中最经典、最通用的 Warp Reduction 实现方式。其他Lane的值是什么这是很多人第一次看 Warp Reduction 时容易疑惑的地方。关键点只有 Lane0 的结果是我们真正关心的。Lane1、Lane2、Lane3… 在归约过程中也会不断计算但它们保存的是各种“部分和”最后并不是完整结果。假设初始Lane0 lane(0) Lane1 lane(1) Lane2 lane(2) ... Lane31 lane(31)执行for(intoffset16;offset0;offset1){val__shfl_down_sync(0xffffffff,val,offset);}offset 16 后Lane0 016 Lane1 117 Lane2 218 ... Lane15 1531offset 8 后Lane0 081624 Lane1 191725 Lane2 2101826 ... Lane7 7152331offset 4 后Lane0 0481216202428 Lane1 1591317212529 Lane2 26101418222630 Lane3 37111519232731每个 Lane 包含 8 个元素。offset 2 后Lane0Lane0 Lane2得到Lane0 02468101214 161820222426283016个元素。Lane1Lane1 Lane3得到Lane1 13579111315 1719212325272931也是16个元素。最后一轮 offset 1Lane0Lane0 Lane1得到Lane0 0123... 31完整32个元素。Lane1Lane1 Lane2注意Lane2 在上一轮已经变成Lane2 2468101214 1618202224262830实际上包含从2开始的一组16个元素因此Lane1 (135...31) (246...30)结果Lane1 1234... 31少了 lane(0)Lane2同理Lane2 Lane3最后Lane2 234... 31少了0 1Lane3最后Lane3 345... 31少了0 1 2规律最终循环结束后Lane0 sum(0~31) Lane1 sum(1~31) Lane2 sum(2~31) Lane3 sum(3~31) ... Lane30 lane(30)lane(31) Lane31 lane(31)可以画成Lane0 0123...31 Lane1 123...31 Lane2 23...31 Lane3 3...31 ... Lane30 3031 Lane31 31像一个三角形0123...31 123...31 23...31 ... 31所以 Warp Reduction 的经典写法for(...)val__shfl_down_sync(...);执行完后if(laneId0){// 使用结果}因为只有 Lane0 整个 Warp 的总和其他 Lane 的值通常直接丢弃不会再使用。
__shfl_down_sync()用法理解
这是理解__shfl_down_sync()最好的方法。假设 Warp 中 32 个线程Lane0 lane(0) Lane1 lane(1) Lane2 lane(2) ... Lane31 lane(31)这里的lane(i)表示线程 i 最初持有的数据。例如vallane(threadIdx.x);那么val: 0 1 2 3 4 5 6 7 ... 31第1轮offset16执行val__shfl_down_sync(mask,val,16);含义Lane0 Lane16 Lane1 Lane17 Lane2 Lane18 ... Lane15 Lane31结果Lane0 lane(0)lane(16) Lane1 lane(1)lane(17) Lane2 lane(2)lane(18) ... Lane15 lane(15)lane(31) Lane16~31 不重要此时前16个线程已经各自保存了2个元素的和。第2轮offset8执行val__shfl_down_sync(mask,val,8);例如Lane0 Lane8而 Lane8 当前已经不是原始值了Lane8 lane(8)lane(24)因此Lane0 (lane(0)lane(16)) (lane(8)lane(24))即Lane0 lane(0)lane(8)lane(16)lane(24)继续展开Lane1 lane(1)lane(9)lane(17)lane(25) Lane2 lane(2)lane(10)lane(18)lane(26) ... Lane7 lane(7)lane(15)lane(23)lane(31)此时Lane0~7 每个线程拥有4个元素之和第3轮offset4执行val__shfl_down_sync(mask,val,4);Lane0Lane0 Lane4而Lane4 lane(4)lane(12)lane(20)lane(28)因此Lane0 lane(0)lane(8)lane(16)lane(24) lane(4)lane(12)lane(20)lane(28)整理Lane0 lane(0)lane(4)lane(8)lane(12) lane(16)lane(20)lane(24)lane(28)已经包含8个元素同理Lane1 1,5,9,13,17,21,25,29 Lane2 2,6,10,14,18,22,26,30 Lane3 3,7,11,15,19,23,27,31第4轮offset2执行val__shfl_down_sync(mask,val,2);Lane0Lane0 Lane2Lane2 当前拥有2,6,10,14,18,22,26,30因此Lane0 0,4,8,12,16,20,24,28 2,6,10,14,18,22,26,30得到Lane0 0,2,4,6,8,10,12,14 16,18,20,22,24,26,28,30共16个元素第5轮offset1执行val__shfl_down_sync(mask,val,1);Lane0Lane0 Lane1Lane1 当前拥有1,3,5,7,9,11,13,15 17,19,21,23,25,27,29,31于是Lane0 0,2,4,6,8,10,12,14 16,18,20,22,24,26,28,30 1,3,5,7,9,11,13,15 17,19,21,23,25,27,29,31最终Lane0 lane(0)lane(1)lane(2)...lane(31)即 Warp 所有线程数据之和。用树来理解第一轮0←16 1←17 2←18 ... 15←31第二轮0←8 1←9 ... 7←15第三轮0←4 1←5 2←6 3←7第四轮0←2 1←3第五轮0←1画成树Lane0 / \ Lane0 Lane16 / \ Lane0 Lane8 / \ Lane0 Lane4 / \ Lane0 Lane2 / \ 0 1实际上是32 ↓ 16组 ↓ 8组 ↓ 4组 ↓ 2组 ↓ 1组每轮把参与求和的数据量扩大一倍offset16 → 每线程2个数 offset8 → 每线程4个数 offset4 → 每线程8个数 offset2 → 每线程16个数 offset1 → 每线程32个数所以你可以记住一个规律第 k 轮结束后 Lane0 保存了 2^k 个 lane 的部分和最后一轮结束Lane0 保存 32 个 lane 的总和这就是__shfl_down_sync()Warp Reduction 的本质。为什么不用个函数直接把所有的lane值相加这是个非常好的问题。你可以理解为GPU 硬件只提供了“从另一个 lane 读取寄存器值”的能力Shuffle并没有提供“把整个 Warp 的所有寄存器自动求和”的通用指令。CPU 的情况例如 CPU 有cab;对应一条加法指令ADD甚至还有POPCNT直接统计二进制中 1 的个数。因为这些操作非常常见CPU 专门做了硬件支持。GPU Warp 内部Warp 有Lane0 Lane1 Lane2 ... Lane31每个 Lane 都有自己的寄存器R0 R1 R2 ... R31硬件提供的是一个 Crossbar 网络----------- Lane0 ---| | Lane1 ---| Crossbar | Lane2 ---| | ... -----------因此支持__shfl_down_sync(...)本质把 Lane17 的寄存器值送给 Lane1或者把 Lane5 的寄存器值送给 Lane2这种任意线程间的数据交换。为什么不直接做 Warp Sum假设增加一个硬件指令warp_sum(val);那么硬件需要支持32个输入 ↓ 求和树 ↓ 输出结果类似 / \ / \ / \ ...对于32输入 64输入AMD Wavefront 128输入未来GPU都要专门设计电路。而 Shuffle 更通用__shfl_down_sync()__shfl_up_sync()__shfl_sync()__shfl_xor_sync()既可以实现SumMaxMinScanPrefix SumFFTMatrix 运算硬件只需要提供Warp Crossbar即可。实际上新 GPU 已经有了近几年 NVIDIA 已经发现Warp Reduction太常用了。因此增加了专门指令__reduce_add_sync()例如unsignedmask0xffffffff;intsum__reduce_add_sync(mask,val);效果Warp内所有线程的 val 求和等价于for(offset16;offset0;offset1){val__shfl_down_sync(mask,val,offset);}但这是较新的 GPU 架构才支持。CUDA Cooperative Groups还有更高级的写法#includecooperative_groups.hnamespacecgcooperative_groups;autowarpcg::tiled_partition32(cg::this_thread_block());intsumcg::reduce(warp,val,cg::plusint());看起来就是warp.reduce()实际上底层还是shuffle ↓ reduction treeAMD HIP 为什么仍然常见 __shfl_down因为NVIDIA Warp 32 AMD Wavefront 64不同架构支持情况不同。为了兼容性 性能 可移植性很多 HPC 代码仍然写for(intoffsetwarpSize/2;offset0;offset1){val__shfl_down(val,offset);}这样CUDA 能跑HIP 能跑不依赖新硬件更深层原因GPU 的设计哲学是提供最基础的数据交换能力 让程序员组合出各种算法而不是提供大量固定功能指令所以硬件提供Shuffle你自己实现Sum Max Min Scan Sort FFT这样硬件更简单、用途更广。因此历史上没有一个通用的warp_sum()函数。不过在新 CUDA 中实际上已经有__reduce_add_sync()只是它出现得比较晚而__shfl_down_sync()已经成为 GPU 编程中最经典、最通用的 Warp Reduction 实现方式。其他Lane的值是什么这是很多人第一次看 Warp Reduction 时容易疑惑的地方。关键点只有 Lane0 的结果是我们真正关心的。Lane1、Lane2、Lane3… 在归约过程中也会不断计算但它们保存的是各种“部分和”最后并不是完整结果。假设初始Lane0 lane(0) Lane1 lane(1) Lane2 lane(2) ... Lane31 lane(31)执行for(intoffset16;offset0;offset1){val__shfl_down_sync(0xffffffff,val,offset);}offset 16 后Lane0 016 Lane1 117 Lane2 218 ... Lane15 1531offset 8 后Lane0 081624 Lane1 191725 Lane2 2101826 ... Lane7 7152331offset 4 后Lane0 0481216202428 Lane1 1591317212529 Lane2 26101418222630 Lane3 37111519232731每个 Lane 包含 8 个元素。offset 2 后Lane0Lane0 Lane2得到Lane0 02468101214 161820222426283016个元素。Lane1Lane1 Lane3得到Lane1 13579111315 1719212325272931也是16个元素。最后一轮 offset 1Lane0Lane0 Lane1得到Lane0 0123... 31完整32个元素。Lane1Lane1 Lane2注意Lane2 在上一轮已经变成Lane2 2468101214 1618202224262830实际上包含从2开始的一组16个元素因此Lane1 (135...31) (246...30)结果Lane1 1234... 31少了 lane(0)Lane2同理Lane2 Lane3最后Lane2 234... 31少了0 1Lane3最后Lane3 345... 31少了0 1 2规律最终循环结束后Lane0 sum(0~31) Lane1 sum(1~31) Lane2 sum(2~31) Lane3 sum(3~31) ... Lane30 lane(30)lane(31) Lane31 lane(31)可以画成Lane0 0123...31 Lane1 123...31 Lane2 23...31 Lane3 3...31 ... Lane30 3031 Lane31 31像一个三角形0123...31 123...31 23...31 ... 31所以 Warp Reduction 的经典写法for(...)val__shfl_down_sync(...);执行完后if(laneId0){// 使用结果}因为只有 Lane0 整个 Warp 的总和其他 Lane 的值通常直接丢弃不会再使用。
