【CUDA编程入门③】经典算子实战：Reduce归约——从朴素实现到高性能优化的演进之路

【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个人要往同一个账户里存钱，大家排队一个一个存------能快得了吗？

这个版本只有教学意义，实际项目中基本不会用。

---

三、版本二：共享内存 + 树形归约------团队内部先算好

既然大家都去抢一个"全局账户"太慢，那能不能先在小组内部算好，再派一个代表去全局汇总？

这就是共享内存树形归约的思路。

思路

1. 每个Block内的线程，先把自己的局部和写到共享内存
2. 在共享内存中做树形归约------两两相加，逐步合并
3. 最后每个Block得到一个"部分和"
4. 只有每个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步（s=4）：[a0+a4, a1+a5, a2+a6, a3+a7, _, _, _, _]
第2步（s=2）：[a0+a4+a2+a6, a1+a5+a3+a7, _, _, _, _, _, _]
第3步（s=1）：[全部加起来, _, _, _, _, _, _, _]

像一棵树一样，从叶子节点逐步"归"到根节点。

性能提升有多大？

假设有1024个Block：

• 版本一：1024个线程同时atomicAdd → 1024次原子操作
• 版本二：每个Block算好，只有1个线程去atomicAdd → 1024次还是1024次？不对，是每个Block 1次，总共也是1024次......

等等，那性能提升在哪？

提升在两方面：

1. 共享内存比全局内存快多了------大部分加法是在共享内存里做的，不是全局内存
2. 全局原子加的竞争少了------虽然次数一样，但每个Block内部的归约已经在共享内存里高效完成了，全局原子加的次数从N次降到了gridDim.x次（N是元素数，通常远大于gridDim.x）

举个具体的例子：N=100万，block size=256

• 版本一：100万次原子操作？不，是(100万/256)=3906次原子操作
• 等等，两个版本都是3906次？那为什么版本二快？

答案是：版本一每个线程做一次atomicAdd，3906个线程同时抢，冲突严重；版本二每个Block做一次，也是3906次，但每次加的数是局部归约后的结果，全局内存的写入量少了很多，而且共享内存中的归约速度远快于全局内存累加。

更重要的是，版本一的local_sum += in[i]是每个线程各自做的，累加的时候每次都读全局内存；而版本二的树形归约是在共享内存里做的------快的根本不是原子操作的数量，而是中间计算的效率。

---

四、版本三：Warp Shuffle优化------最后一步的极致优化

树形归约的最后几步，其实是在同一个Warp内完成的。

比如block size=256：

• 128→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一个结果
    // 注意：这里只归约到 s>32，也就是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算子：

1. 版本一：全局原子加 ------ 简单但慢，全局序列化是瓶颈
2. 版本二：共享内存树形归约 ------ 每个Block内部先归约，大幅减少全局原子的竞争
3. 版本三：Warp Shuffle优化 ------ 最后一级归约用Shuffle，省共享内存和同步
4. 版本四：多级归约 ------ 超大规模数据的解决方案

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
• 更多内容持续更新中......