我第一次写cuda reduce kernel是在一个物理模拟项目里当时要先对一千万个float求和作为下一步统计的前置计算。刚把CUDA入门教程刷完我信心满满开一堆线程每人加一块最后再加到一起这不就完了结果那个朴素内核跑出来居然比CPU单线程还慢我盯着终端上的耗时愣了好一会儿——GPU怎么会输给CPU后来我才明白reduce看似简单但它是CUDA里最能体现并行思维和串行思维差异的一个例子。写完优化版之后再回头看性能差了几十倍不止。这篇文章就把我那次从慢到离谱到逼近带宽极限的完整过程拆开讲包括每个版本的代码、为什么快、以及跑偏之后怎么排查。适合刚学完CUDA基础、准备写第一个实用内核的人也适合写过reduce但没系统做过优化的开发者。1. 归约问题一个循环变成一棵树之后的事1.1 归约是什么从CPU上的一行循环说起归约reduction指的是通过一个满足结合律的二元运算符把一组数据折叠成一个结果的操作。求和、求最大值、求最小值、求均值、求点积、求方差这些都属于归约。在CPU上写起来就是一个循环float sum 0.f; for (int i 0; i n; i) { sum input[i]; }就这么简单。问题是GPU上不能照搬这个循环。你开的几千个线程如果都去执行同一个串行for循环那等于每个线程都在重复读整段数组结果不仅错误性能也完全没法看。所以GPU上的归约必须换一个思路让数据被并行地折叠而不是被一个线程逐个吃掉。这也解释了为什么很多从PyTorch自定义算子开始接触CUDA的人第一个要手写的内核往往就是reduce——softmax的分母、layer norm的均值方差、梯度累加底层全是归约。把它吃透了后面那些算子基本就是换换运算符和数据类型的事。1.2 并行归约的树状结构log N轮的意义并行归约本质上是一棵树。想象一场淘汰赛N个参赛者两两配对第一轮产生N/2个胜者第二轮产生N/4个胜者……直到最后只剩一个冠军。归约也是同样的过程把相邻的两个元素相加得到N/2个部分和再把这N/2个部分和两两相加得到N/4个……以此类推。每轮配对的数量减半总共需要约log2(N)轮。这个树状结构是所有reduce优化的出发点。它告诉我们两件事第一归约的并行深度是O(log N)不是O(N)所以理论上N很大时也能很快算完第二归约过程包含大量中间结果这些中间结果存放在哪、怎么被读取决定了性能的上限。还有一点容易被忽视浮点加法不满足精确的结合律(a b) c和a (b c)的结果可能有微小差异。并行归约改变的是累加顺序所以最终结果和CPU串行累加会有细微出入。工程上通常可以忽略但做科学计算、做对比验证的时候要心里有数免得排查半天以为是自己写错了。1.3 GPU线程组织给归约出的难题GPU的线程模型是三层线程thread、线程块block、网格grid。一个block内的线程可以通过共享内存shared memory直接互相通信也可以靠__syncthreads()做同步但不同block之间没有直接通信机制谁也看不见谁的数据。这就逼出了reduce kernel的两级架构第一级每个block先把分给它的那一段数据归约成一个部分和第二级再把所有block的部分和合并成最终的单一结果。第二级的合并方案有好几种后面会专门讲。先记住一个总体印象reduce的优化空间基本都集中在这两级归约的衔接处和每一级的树形折叠方式上。2. 第一版朴素内核为什么跑不过CPU分支发散与带宽浪费2.1 教科书版隔半归约框架正确但性能稀烂我先写的是最常见的教科书版本每个block处理一段数据先把数据读进共享内存然后循环里做隔半归约。__global__ void reduce_naive(const float* input, float* output, int n) { int tid threadIdx.x; int idx blockIdx.x * blockDim.x tid; // 声明动态共享内存 extern __shared__ float sdata[]; // 把全局数据读入共享内存 sdata[tid] input[idx]; __syncthreads(); // 树状归约步长从 blockDim.x/2 开始逐次减半 for (int stride blockDim.x / 2; stride 0; stride 1) { if (tid stride) { sdata[tid] sdata[tid stride]; } __syncthreads(); } if (tid 0) { output[blockIdx.x] sdata[0]; } }这个版本能跑出正确结果但慢得令人绝望。我当时的数组长度是2000万启动的block数和block大小组合也是一步步试的结果跑出来比CPU的O2循环还慢一个数量级。原因不复杂往下看就明白了。2.2 线程束发散有一半线程在摸鱼第一个大问题是线程束发散。GPU执行指令的最小单位是warp一个warp包含32个线程它们在同一时刻执行同一条指令。当代码里出现if (tid stride)这种分支时如果同一个warp里有的线程满足条件、有的不满足硬件会把它们拆成两条路径分别执行满足条件的线程先干活不满足的线程在旁边空等。在隔半归约的循环里stride从256减到128、64、32、16、8、4、2、1。前几轮还算正常但到stride16之后一个block里真正在干活的线程已经少于warp的大小了。到stride1那轮512个线程里只有第0个线程在加最后两个数剩下511个线程全在等__syncthreads()。这种绝大多数线程围观少数线程干活的场景就是分支发散最典型的形态。更麻烦的是__syncthreads()是要整个block的线程都到位才会放行的。这就意味着空转的线程不干活但同步等待的时间一点没少。每一轮大家都在等最慢的路径执行完性能自然被拖垮。2.3 数据复用为零全局访存延迟完全暴露第二个问题更致命每个线程只加载了一个float就进入了归约阶段。一个float是4字节2000万float就是80MB。这种每人拿一个数就跑的写法让每个线程的每次加法都依赖一次来自全局内存的读取完全没有把数据复用起来。GPU的全局内存吞吐是靠大量并发线程的访存请求叠加出来的。如果每个线程只读一次数据就开始做同步和归约那些读请求根本来不及攒成足够宽的传输队列内存带宽利用率会低得可怜。打个比方这就好比一条十车道的高速路每辆车只运一箱货就跑一趟车道再多也跑不出吞吐量。正确的做法是让每辆车装满货物再出发——也就是让每个线程一次处理多个元素。顺带说共享内存的作用也在这时候体现它是芯片上的高速内存带宽比全局内存高一个数量级延迟低得多。把数据先搬到共享内存再做树状归约至少能避免在归约过程中反复访问全局内存。但这个版本的问题在于搬一次、归约很多次指令开销和同步开销太大了框架对收益没吃到。3. 线程束级优化共享内存、bank冲突与最后的手动展开3.1 Grid-stride loop让每个线程先把一段数据算完意识到每个线程只算一个元素太浪费之后我的第一个改动是引入grid-stride loop。这个模式在CUDA高性能代码里非常常见核心思路是用固定数量的线程以blockDim.x * gridDim.x为步长去遍历整个数组每个线程在寄存器里累加自己负责的所有元素。__global__ void reduce_grid_stride(const float* input, float* output, int n) { int tid threadIdx.x; int idx blockIdx.x * blockDim.x tid; float sum 0.f; // 每个线程以网格总线程数为步长循环累加 for (int i idx; i n; i blockDim.x * gridDim.x) { sum input[i]; } // 累加结果写入共享内存再做块内归约 extern __shared__ float sdata[]; sdata[tid] sum; __syncthreads(); for (int stride blockDim.x / 2; stride 32; stride 1) { if (tid stride) { sdata[tid] sdata[tid stride]; } __syncthreads(); } // 剩下的归约在 warp 内完成后面细说 }这个改动带来了两个好处。一是数据复用率大幅提升每个从全局内存加载的数都被累加进同一个寄存器后续计算不用再访存。二是线程数量不再受数据量限制数组是1000也好是10亿也好都能用同一个内核处理只是每个线程循环的轮数不同。这一步做完耗时已经肉眼可见地降了下来。3.2 共享内存的bank冲突与padding解法共享内存虽然快但它有个奇怪的硬件特性它被分成32个bank存储体每个bank每个时钟周期可以服务一个地址。同一warp里的32个线程如果同时访问同一个bank的不同地址这些访问会被串行化速度直接降到1/32。这叫bank conflict。什么时候会产生bank冲突举个例子如果共享内存里相邻地址恰好间隔32的倍数那么线程0访问sdata[0]、线程1访问sdata[32]、线程2访问sdata[64]……它们全部命中同一个bank。在隔半归约里访问sdata[tid]和sdata[tid stride]的模式通常是连续的只要stride不是32的倍数一般没有冲突。但有些写法——比如按block内的线程号绑定到相距很远的位置——很容易踩中。最经典的规避办法是padding把共享内存数组的长度从blockDim.x改成blockDim.x 1。多出来的这一个float不参与逻辑计算只是把每一行数据的bank编号整体错开一位这样相邻线程访问的bank就不同了。代价几乎为零收益在某些size组合下能达到10%~20%。我见过不少人忽略这一步觉得多分配一个float有什么意义实测过之后才会服气。3.3 最后32个线程warp内不需要同步的手动展开当块内归约循环进行到stride32时一个block里只剩最后一个warp在参与。这时候有一个很重要的特性可以拿来用同一个warp内部是锁步执行的天然同步不需要__syncthreads()。而__syncthreads()的代价是让block里所有warp都停下来等即使它们已经不参与归约了。所以最后这32个线程的归约可以跳出循环改成手动展开if (tid 32) { float val sdata[tid] sdata[tid 32]; sdata[tid] val; } // 从这往下不需要 __syncthreads()因为都在同一个 warp 内 if (tid 16) { sdata[tid] sdata[tid 16]; sdata[tid] sdata[tid 8]; sdata[tid] sdata[tid 4]; sdata[tid] sdata[tid 2]; sdata[tid] sdata[tid 1]; } if (tid 0) { output[blockIdx.x] sdata[0]; }这段代码看起来像魔法其实逻辑很清楚stride从16、8、4、2到1每行都是一个独立的加法没有分支、没有循环、没有同步。编译器还能把这些加法重排成指令级流水线进一步隐藏延迟。这是reduce优化里收益高但需要理解原理的一步不建议在没跑通前面几版之前直接上。4. 实测对比与参数调优从8.2ms到0.33ms4.1 一个老GTX上的版本对比表为了让你对每一档优化到底能带来多少提升有个直观印象我放一张实测对比表。测试环境是RTX 3060 Laptop数据规模2000万个float80MB显存理论带宽约256GB/s只读一次的理论耗时下限大约是80MB / 256GB/s ≈ 0.31ms。每次跑10轮取中位数避免时钟波动干扰。版本耗时带宽利用率说明朴素隔半归约每人1个元素8.2ms约4%分支发散严重数据复用为零grid-stride loop 共享内存块内归约1.1ms约28%数据复用率提升但同步开销仍在展开 warp内手动归约0.45ms约68%省掉最后几轮同步和循环分支展开 float4向量化0.33ms约94%访存宽度翻四倍接近带宽上限注意绝对数字只在你的硬件和我的硬件接近时有参考价值更值得关注的是倍率关系从第一版到最后一版性能差了约25倍。而这四档之间每一档的改动其实都不大没有任何一个版本做了伤筋动骨的重写。我自己的体会是前两档解决的是算法框架到底对不对后两档解决的是硬件资源到底用没用好。很多人一开始就纠结float4和展开却连grid-stride loop都没写对等于还没学会走路就想跑。4.2 block size与grid size的选取思路block size怎么选是reduce新人问得最多的问题。通用的经验值是128到512其中256是我在大多数GPU上实测最稳的中间值。为什么block太小比如64会导致block数量过多末级归并和调度开销变大block太大比如1024共享内存占用和同步等待都会变重尤其到了最后几轮归约大量线程在空等同一个__syncthreads()。最靠谱的做法是写一个探针把block size从64、128、256、512、1024都跑一遍选耗时最小的那个。不要猜你的GPU架构可能和教程作者的不一样跑一次也就几秒。grid size的思路是让GPU尽量被填满。一个粗略的估算是grid总线程数约等于GPU的SM数量乘以每SM最大常驻线程数一般是2048或更多取决于架构和占用率。但配合grid-stride loop之后grid不需要卡得那么死多一点少一点影响不大只要不是小到只有几个block就行。4.3 多block结果合并atomicAdd与最后一个block方案块内归约完成之后每个block产出一个部分和接下来要把它们合并成一个数。常见方案有三种我分别说下适用场景第一种atomicAdd。每个block的0号线程把自己block的部分和原子累加到全局内存的一个变量上使用前记得cudaMemset把输出变量清零。这是最简单直接的做法现代GPU的atomicAdd对L2缓存做过优化block数量在几千个以内时开销很低。绝大多数项目我推荐先用它。if (tid 0) { atomicAdd(output[0], sdata[0]); }第二种二次内核。第一个内核写出一个包含所有block部分和的小数组第二个内核把这个小数组再做一次归约。好处是精度更好、逻辑纯粹坏处是多启动一次内核多一次全局内存读写。block数量很大、或者你在做需要精确归约顺序的场景时可以考虑。第三种让最后一个block收尾。思路是每个block把部分和写到全局数组然后用一个计数器记录到达的block数最后一个到达的block计数器值等于gridDim.x - 1负责把之前所有人的部分和读回来再做一次归约。这个方案避免了原子累加的性能损耗但实现复杂度明显更高收益在多数场景下不明显。除非你已经确认atomicAdd是瓶颈否则不建议一上来就用这个。顺带提一句精度atomicAdd的累加顺序是浮动的每次运行结果可能有百分位级别的微小差异这是浮点加法的正常现象不是bug。5. 排错手记结果不对与性能不及预期时的排查路径5.1 结果总是差一点先查同步与数据竞争如果你跑出来的结果偶尔对、偶尔不对或者总是和CPU参考值差一点九成是同步或边界问题。我踩过的坑主要有这几个输出变量没清零就直接atomicAdd结果是在随机地址上累加结果完全不可控。__syncthreads()放错位置比如在赋值共享内存之后忘了放导致归约时读到还没写好的数据。数组长度不是blockDim的整数倍时共享内存里有一部分是上一轮残留的数据没初始化就参与归约结果差一点是正常的。排查手法很简单先用小数组、单block、固定值去验证基本正确性再逐步放大。发现数据竞争类问题用compute-sanitizer老版叫cuda-memcheck直接跑一下compute-sanitizer --tool racecheck ./your_program它能把共享内存的竞争访问精确指出来比自己盯着代码猜高效多了。5.2 性能卡住不动Nsight与带宽利用率分析结果对了但速度上不去这时候别瞎猜直接开Nsight Computencu看两个核心指标Memory Throughput实际达到的内存带宽百分比和Achieved Occupancy。如果带宽利用率已经到90%以上说明瓶颈就是显存带宽再优化代码意义不大算法层面已经到头了如果只有50%上下大概率是访存模式有问题——比如没有使用合并访问、每个线程处理的数据太少、或者有未对齐访问。还要注意一个叫tail effect尾部效应的现象当数据量不能被grid总线程数整除时最后一批block里只有部分线程在干活其他人空转。数据规模凑巧不整齐的时候这个损失能有几个百分点。解决方法是可以让grid大小和数据量更好地匹配或者干脆让每个线程加载数据时做边界检查而不是依赖整除。5.3 float4向量化踩坑16字节对齐问题最后一个高频坑是float4向量化之后结果突然错乱。float4类型要求16字节对齐也就是说被转换的指针地址必须是16的倍数。cudaMalloc分配的内存一般满足对齐要求但如果你对数组做了偏移比如从input 1开始处理那这个地址极大概率不是16字节对齐的reinterpret_cast不会帮你检查结果就是读出来的数据是错位的程序还不报错。const float4* in4 reinterpret_castconst float4*(input); // 前提input本身16字节对齐 float4 v in4[i]; sum v.x v.y v.z v.w;规避方法如果输入数据确实需要偏移先把偏移量补齐到16字节的整数倍再开处理或者用cudaMalloc分配一个额外padding的缓冲区。另外结构体数组AoS直接转float4也很容易踩对齐坑因为结构体的大小不一定是16的倍数。确认对齐之后再向量化否则宁可先不用。提示优化reduce这类内存密集内核记住一个原则——先把算法框架做对再把访存模式做宽最后才轮到指令级微调。每一步保持一个变量在变方便量化这一改动到底带来多少收益。最终版本跑出来0.33ms已经接近RTX 3060 Laptop在这块卡上的带宽上限了。再往下抠就是换卡或者换精度的事对我来说写reduce kernel能到这个程度已经可以把精力放到下一个算子上去了。最后再分享一点个人体会reduce kernel的价值不在于求一个和本身而在于它是很多并行归约问题的原型。softmax的分母归一化、layer norm的均值和方差、矩阵乘的partial sum累加、梯度聚合底层全是这棵树。把这个内核从慢到快完整调一遍你会对共享内存、warp同步、访存合并这些概念有一个实打实的体感比看十篇理论文章都有用。当年那个比CPU还慢的第一版现在回想起来倒是个挺宝贵的起点——没有它我还真体会不到CUDA优化那些骚操作到底在解决什么问题。