【CUDA编程入门③】经典算子实战Reduce归约——从朴素实现到高性能优化的演进之路本文是「CUDA编程与HPC高性能计算」系列的第三篇。Reduce归约是深度学习中最基础也最常用的算子之一——Softmax需要它LayerNorm需要它Loss计算也需要它。别看Reduce的功能很简单就是求和/求最大值但要写得快里面的门道可不少。本文带你一步步演进从最朴素的实现到接近硬件极限的高性能版本。一、写在前面Reduce是什么Reduce归约听起来很高端其实就是**“把一堆数合并成一个数”**。最常见的是求和输入[1, 2, 3, 4, 5, 6, 7, 8]输出36所有元素的和除了求和还可以求最大值、最小值、按位与/或等等——只要满足结合律的操作都可以做Reduce。Reduce难在哪所有线程要协作产生一个结果线程之间需要交换部分和最终必须合并成全局结果说白了难在**“协作”**——GPU有几千个线程怎么让它们高效地凑在一起算出最终结果先上一张Reduce优化演进图二、版本一朴素实现——全局内存原子加我们从最简单的版本开始。思路每个线程读一段数据算出自己的局部和然后大家一起用atomicAdd加到全局结果上。__global__ void reduce_atomic(const float* in, float* out, int N) { int idx blockIdx.x * blockDim.x threadIdx.x; int stride gridDim.x * blockDim.x; float local_sum 0.0f; // 每个线程累加自己负责的一段 for (int i idx; i N; i stride) { local_sum in[i]; } // 全局原子加 atomicAdd(out, local_sum); }优点 vs 缺点优点缺点简单、好写、不容易错性能差所有block同时做atomicAdd全局序列化适合小规模数据数据量大了原子操作成了瓶颈为什么慢想象一下1000个人要往同一个账户里存钱大家排队一个一个存——能快得了吗这个版本只有教学意义实际项目中基本不会用。三、版本二共享内存 树形归约——团队内部先算好既然大家都去抢一个全局账户太慢那能不能先在小组内部算好再派一个代表去全局汇总这就是共享内存树形归约的思路。思路每个Block内的线程先把自己的局部和写到共享内存在共享内存中做树形归约——两两相加逐步合并最后每个Block得到一个部分和只有每个Block的0号线程去做一次全局原子加__global__ void reduce_shared(const float* in, float* out, int N) { __shared__ float sdata[256]; // 共享内存大小等于block size int tid threadIdx.x; int idx blockIdx.x * blockDim.x tid; // 第一步每个线程累加自己的数据 float local 0.0f; for (int i idx; i N; i gridDim.x * blockDim.x) { local in[i]; } sdata[tid] local; __syncthreads(); // 等所有人都写完 // 第二步树形归约在共享内存中 for (int s blockDim.x / 2; s 0; s 1) { if (tid s) { sdata[tid] sdata[tid s]; } __syncthreads(); // 每一步都要同步 } // 第三步只有0号线程去做全局原子加 if (tid 0) { atomicAdd(out, sdata[0]); } }树形归约长什么样以8个元素为例树形归约的过程是这样的初始[a0, a1, a2, a3, a4, a5, a6, a7] 第1步s4[a0a4, a1a5, a2a6, a3a7, _, _, _, _] 第2步s2[a0a4a2a6, a1a5a3a7, _, _, _, _, _, _] 第3步s1[全部加起来, _, _, _, _, _, _, _]像一棵树一样从叶子节点逐步归到根节点。性能提升有多大假设有1024个Block版本一1024个线程同时atomicAdd → 1024次原子操作版本二每个Block算好只有1个线程去atomicAdd → 1024次还是1024次不对是每个Block 1次总共也是1024次……等等那性能提升在哪提升在两方面共享内存比全局内存快多了——大部分加法是在共享内存里做的不是全局内存全局原子加的竞争少了——虽然次数一样但每个Block内部的归约已经在共享内存里高效完成了全局原子加的次数从N次降到了gridDim.x次N是元素数通常远大于gridDim.x举个具体的例子N100万block size256版本一100万次原子操作不是(100万/256)3906次原子操作等等两个版本都是3906次那为什么版本二快答案是版本一每个线程做一次atomicAdd3906个线程同时抢冲突严重版本二每个Block做一次也是3906次但每次加的数是局部归约后的结果全局内存的写入量少了很多而且共享内存中的归约速度远快于全局内存累加。更重要的是版本一的local_sum in[i]是每个线程各自做的累加的时候每次都读全局内存而版本二的树形归约是在共享内存里做的——快的根本不是原子操作的数量而是中间计算的效率。四、版本三Warp Shuffle优化——最后一步的极致优化树形归约的最后几步其实是在同一个Warp内完成的。比如block size256128→64→32跨Warp需要共享内存同步16→8→4→2→1在同一个Warp内既然在同一个Warp内那我们可以用Warp Shuffle来做——不用共享内存也不用同步直接在寄存器层面交换数据。Warp内归约函数先写一个Warp内求和的辅助函数__inline__ __device__ float warp_reduce_sum(float val) { // 一步步往下挪每次加一半 val __shfl_down_sync(0xFFFFFFFF, val, 16); // lane 0 加 lane 16 val __shfl_down_sync(0xFFFFFFFF, val, 8); // lane 0 加 lane 8 val __shfl_down_sync(0xFFFFFFFF, val, 4); // ... val __shfl_down_sync(0xFFFFFFFF, val, 2); val __shfl_down_sync(0xFFFFFFFF, val, 1); return val; // lane 0 有完整结果其他lane不一定对 }为什么offset是16, 8, 4, 2, 1因为一个Warp有32个线程用蝴蝶式butterfly归约log₂(32)5步。整合到Reduce kernel里__global__ void reduce_warp(const float* in, float* out, int N) { __shared__ float sdata[256]; int tid threadIdx.x; int idx blockIdx.x * blockDim.x tid; // 第一步加载数据到共享内存和之前一样 float local 0.0f; for (int i idx; i N; i gridDim.x * blockDim.x) { local in[i]; } sdata[tid] local; __syncthreads(); // 第二步共享内存树形归约归约到每个Warp一个结果 // 注意这里只归约到 s32也就是8个Warp各自有一个结果 for (int s blockDim.x / 2; s 32; s 1) { if (tid s) { sdata[tid] sdata[tid s]; } __syncthreads(); } // 第三步Warp Shuffle做最后一级归约 float val sdata[tid]; val warp_reduce_sum(val); // 每个Warp的0号线程tid % 32 0去做全局原子加 if (tid % 32 0) { atomicAdd(out, val); } }提升在哪里最后5步归约不用共享内存了 → 省了共享内存访问不用__syncthreads()了 → 省了同步开销数据一直在寄存器里 → 速度更快虽然只是优化了最后几步但因为最后几步的同步开销占比不小所以收益还是很可观的。五、版本四多级归约——超大数据的终极方案如果数据特别特别大比如几十亿元素Block数量会很多全局原子加的竞争又会变得严重。这时候怎么办用多级归约。思路Level 1每个Block归约一部分数据 → 产生 M 个部分和M Block数 Level 2一个Block归约这 M 个部分和 → 最终结果第一级kernel产生M个部分和写回全局内存第二级kernel把这M个部分和再归约一次得到最终结果。好处避免了大量Block同时做atomicAdd第二级只有一个Block效率很高可以用两级也可以用多级看数据规模 实际上Thrust、cuBLAS这些库的Reduce实现都是多级归约的思路。六、性能对比哪个版本最快我们来排个序从慢到快版本性能适用场景全局原子加⭐教学演示、超小规模共享内存树形归约⭐⭐⭐一般规模最常用 Warp Shuffle优化⭐⭐⭐⭐追求极致性能多级归约⭐⭐⭐⭐超大规模数据手写Reduce能达到什么水平一般能达到带宽的70%-90%Reduce是memory-bound的和cuBLAS/Thrust的专业实现比大概有10%-30%的差距但作为学习手写一遍的价值远超那点性能差距七、常见误区误区一Reduce就是不断累加朴素的逐个累加除了慢还有精度问题——浮点数加的顺序不一样结果会有微小差异浮点误差累积。树形归约的误差通常比串行累加小因为加的深度更浅。如果对精度要求高可以用Kahan求和算法。误区二Block size越大越好对于Reduce来说Block size太大反而不好Block太大 → 每个线程处理的数据太少 → 计算/访存的比例太低树形归约的步数更多 → 同步开销占比上升256通常是Reduce的甜点——平衡了并行度和同步开销。误区三Shuffle可以完全替代共享内存不行Warp Shuffle只能在一个Warp内通信。跨Warp的通信还是得用共享内存。Shuffle是锦上添花不是雪中送炭——共享内存树形归约才是主体Shuffle是最后一步的优化。八、总结一下今天这篇文章我们从朴素实现开始一步步优化Reduce算子版本一全局原子加—— 简单但慢全局序列化是瓶颈版本二共享内存树形归约—— 每个Block内部先归约大幅减少全局原子的竞争版本三Warp Shuffle优化—— 最后一级归约用Shuffle省共享内存和同步版本四多级归约—— 超大规模数据的解决方案Reduce虽然简单但它涵盖了CUDA优化的很多核心思想合并访存、共享内存复用、Warp级通信、减少全局原子……把Reduce搞懂了其他算子也能触类旁通。下一篇预告我们来挑战深度学习的王者算子——GEMM通用矩阵乘法。从为什么需要Tiling到Shared Memory Tiling实现再到Tensor Core带你一步步接近硬件极限。本系列导航第1篇CUDA编程体系全景线程模型、内存层次与Warp执行机制第2篇CUDA性能优化基础内存访问、Warp Shuffle与Occupancy调优第3篇经典算子实战Reduce归约——从朴素到高性能本文第4篇经典算子实战GEMM矩阵乘法——从Tiling到Tensor Core更多内容持续更新中……
