Reduction 优化(单卡 A100)

Lesson 08 · CUDA 性能调优 · 第一个访存瓶颈算子的极限压榨

上一节把 roofline 钉到了真卡上。这一节换一种瓶颈:tiled GEMM 是计算瓶颈, 而 reduction(归约)——把一个大数组加成一个标量——是教科书级的访存瓶颈算子。 它在 AI 里无处不在:softmax 的分母、LayerNorm 的均值/方差、loss、梯度范数,本质都是 reduction。 本节只针对 A100(sm_80),把它从朴素版一路优化到逼近 HBM 带宽极限。

1. 先定瓶颈:reduction 注定是访存瓶颈

把 N 个 float 加成 1 个标量。读 N×4 字节,算 N−1 次加法。算术强度:

算术强度 = FLOPs / Bytes ≈ (N-1) / (4N) ≈ 0.25 FLOP/Byte

A100 的脊点约 12.6 FLOP/Byte。0.25 远远在脊点左侧 → 铁定的访存瓶颈。这立刻给了你一条铁律:

2. v0 朴素版:interleaved addressing(慢在哪)

经典写法:block 内每个线程搬一个元素进 shared,然后"树形"两两相加。第一版用交错寻址:

__shared__ float sdata[BLOCK];
int tid = threadIdx.x;
sdata[tid] = g_in[blockIdx.x*BLOCK + tid];   // 合并加载 ✅
__syncthreads();

for (int s = 1; s < BLOCK; s *= 2) {
    if (tid % (2*s) == 0)              // ❌ 取模 → warp 内严重分支发散
        sdata[tid] += sdata[tid + s];        // ❌ 跨步访问 → bank conflict
    __syncthreads();
}
if (tid == 0) g_out[blockIdx.x] = sdata[0];

两个病:(1) tid % (2*s) 让一个 warp 里只有零星线程干活,其余空转却仍占用调度—— 这正是第 6 节说的 warp 分支发散。 (2) sdata[tid+s] 这种跨步访问,在 shared 里会撞 bank conflict。

v1:先只治发散(留着 bank conflict 当对照)

不动跨步访问,只把"取模筛线程"换成"把活跃线程压到 warp 前部":index = 2*s*tid。 小步长时活跃 lane 在 warp 内连续 → 不再发散;但 sdata[index+s] 仍是跨步 → bank conflict 没治。 这一版单独存在,是为了让你在 benchmark 里看清"去发散"本身值多少。

for (int s = 1; s < BLOCK; s *= 2) {
    int index = 2*s*tid;                  // ✅ 活跃线程连续 → 无发散
    if (index < BLOCK)
        sdata[index] += sdata[index + s];   // ❌ 仍跨步 → bank conflict
    __syncthreads();
}

3. v2 第一刀:sequential addressing(干掉发散与 bank conflict)

把循环反过来——从大步长往小走,让活跃线程永远是前一半连续的 tid:

for (int s = BLOCK/2; s > 0; s >>= 1) {
    if (tid < s)                          // ✅ 活跃线程是连续的 0..s-1
        sdata[tid] += sdata[tid + s];
    __syncthreads();
}

同一个 warp 的 32 个线程要么全活跃、要么全休眠 → 无分支发散。 而且 sdata[tid] 与 sdata[tid+s] 都是连续地址 → 无 bank conflict。[1]

4. v3 加载即相加:first add during load

v2 的循环第一轮就有一半线程闲着(s=BLOCK/2 时只有前一半 tid 干活)。 与其让它们空转,不如让每个线程在写进 shared 之前先读两个元素加起来—— 于是每个 block 覆盖 2*BLOCK 个元素,grid 直接减半,固定开销摊薄一倍:

int base = blockIdx.x*(BLOCK*2) + tid;
float a = g_in[base];
float b = g_in[base + BLOCK];             // 两段都连续 → 仍合并访问
sdata[tid] = a + b;                       // ✅ 加载即相加,省掉第一轮的一半空转
__syncthreads();
// …接 v2 的 sequential 循环…

这是从"块内归约"转向"每线程多干活"的第一步,也是下一刀 grid-stride 的引子。

5. v4 第二刀:warp shuffle —— 最后 32 个线程不走 shared

注意上面循环走到 s ≤ 32 时,只剩一个 warp 在干活。这时还在用 __syncthreads() 和 shared memory,纯属浪费——同一 warp 内的线程本就是锁步执行的,可以直接交换寄存器。 A100 上用 __shfl_down_sync:

// 当 s 降到 32:不再回写 shared,改用 warp 内寄存器洗牌
float v = sdata[tid];
for (int off = 16; off > 0; off >>= 1)
    v += __shfl_down_sync(0xffffffff, v, off);   // 寄存器→寄存器,零 shared、零 sync
if (tid == 0) g_out[blockIdx.x] = v;

__shfl_down_sync(mask, v, off):让 lane tid 直接读到 lane tid+off 寄存器里的 v, 全程不碰 shared memory,也不需要 __syncthreads()。[2] 最后一个 warp 的 5 步全在寄存器里完成。

6. v5 第三刀:grid-stride —— 一个线程先吞多个元素

前面每个线程只读 1~2 个元素。但启动 block、加载、做 log₂ 轮归约这些开销是固定的, 摊在 1 个元素上太亏。更进一步是让每个线程在进 shared 之前,先用一个 grid-stride 循环把 N/总线程数 个元素加起来:

int i = blockIdx.x*BLOCK + tid;
int stride = gridDim.x * BLOCK;          // 整个 grid 的线程数
float sum = 0;
while (i < N) {                          // 每次跳一整个 grid → 始终合并访问
    sum += g_in[i];
    i += stride;
}
sdata[tid] = sum;                         // 进 shared 时已经是局部和
__syncthreads();
// …接 v2/v4 的块内归约 + warp shuffle 收尾…

为什么是 grid-stride 而不是"每个线程连读一段连续区间"?因为同一拍里 warp 的 32 个线程地址必须相邻才合并。 跳 stride = gridDim.x*BLOCK 保证每一轮 32 个线程读的是连续的 32 个 float。[3]

7. 收尾:G 个部分和怎么变一个标量

每个 block 产出一个部分和,grid 有 G 个 block → 得到 G 个部分和。怎么合成最终一个标量?三种路子:

A100 上 grid-stride 已经把 G 压得很小(几百~上千),所以 atomicAdd 收尾或两趟 kernel 都很快,差别可忽略—— 瓶颈始终是第一趟那次"完整读一遍数组",收尾只是零头。

8. 真机验证:实测带宽 + ncu 看 DRAM Throughput

reduction 不看 Compute%,要盯内存吞吐。配套代码 code/0008-reduce.cu 把六版串成一个程序, 用 CUDA event 计时逐版打印带宽与加速比。先看本机 3060 Laptop(sm_86,~330 GB/s 峰值)的实测:

nvcc -O3 -arch=sm_86 0008-reduce.cu -o reduce   # 3060;A100 用 -arch=sm_80
./reduce

读这张表的三个要点:(1) v3/v4 各贡献最大的两跳——加载即相加摊薄开销、warp shuffle 省掉最后 5 轮同步; (2) v4 已到 ~310 GB/s ≈ 峰值 92%,撞上 3060 的带宽墙; (3) v5 grid-stride 反而略慢(0.92x)——不是写错了,而是带宽已封顶,grid 大小不再是瓶颈。 同一条代码阶梯换到 A100,v5 才会反超。

原生 Linux 上还可以用 ncu 看更细的读数(WSL 下 ncu 内核级 profiling 不可用,只能用上面的 event 计时):

ncu --set full --section MemoryWorkloadAnalysis ./reduce

三个关键读数,按这个顺序判断:

9. 整条优化链回顾

10. 练习:reduction 决策(混合前几节)

点选项看反馈。这些题交错了 roofline、合并、bank conflict、occupancy——这正是检验你是否真懂的方式。

主源推荐(本节精读)

📘 Mark Harris, "Optimizing Parallel Reduction in CUDA"(NVIDIA 经典 slide) ——逐版本演示 7 步优化(交错→顺序→首次加载即相加→展开→多元素),本节的骨架就来自它,只是把收尾换成现代的 warp shuffle 并落到 A100。