Lesson 08 · CUDA 性能调优 · 第一个访存瓶颈算子的极限压榨
上一节把 roofline 钉到了真卡上。这一节换一种瓶颈:tiled GEMM 是计算瓶颈, 而 reduction(归约)——把一个大数组加成一个标量——是教科书级的访存瓶颈算子。 它在 AI 里无处不在:softmax 的分母、LayerNorm 的均值/方差、loss、梯度范数,本质都是 reduction。 本节只针对 A100(sm_80),把它从朴素版一路优化到逼近 HBM 带宽极限。
ncu 的 DRAM Throughput 验证你吃到了 A100 ~1.5–2.0 TB/s 的几成。
把 N 个 float 加成 1 个标量。读 N×4 字节,算 N−1 次加法。算术强度:
算术强度 = FLOPs / Bytes ≈ (N-1) / (4N) ≈ 0.25 FLOP/ByteA100 的脊点约 12.6 FLOP/Byte。0.25 远远在脊点左侧 → 铁定的访存瓶颈。这立刻给了你一条铁律:
N×4 / 带宽。比如 N=2²⁸(1 GB),A100 ~1.8 TB/s → 下限约 0.55 ms。你的 kernel 离它多远,就是优化空间。
经典写法: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];
__syncthreads();
}
if (tid == 0) g_out[blockIdx.x] = sdata[0];两个病:(1) tid % (2*s) 让一个 warp 里只有零星线程干活,其余空转却仍占用调度——
这正是第 6 节说的 warp 分支发散。
(2) sdata[tid+s] 这种跨步访问,在 shared 里会撞 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]
注意上面循环走到 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 步全在寄存器里完成。
BLOCK/32 个槽,几乎无 bank conflict。__reduce_add_sync(sm_80 原生指令)还能把整 warp 求和压成一条指令。
前面每个线程只读 1 个元素。但启动 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();
// …接第 3、4 节的块内归约…为什么是 grid-stride 而不是"每个线程连读一段连续区间"?因为同一拍里 warp 的 32 个线程地址必须相邻才合并。
跳 stride = gridDim.x*BLOCK 保证每一轮 32 个线程读的是连续的 32 个 float。[3]
每个 block 产出一个部分和,grid 有 G 个 block → 得到 G 个部分和。怎么合成最终一个标量?三种路子:
| 做法 | 说明 | 何时用 |
|---|---|---|
| 两趟 kernel | 第一趟 N→G,第二趟 G→1(同一 kernel 再跑一次) | 最稳,最常用,易调试 |
| atomicAdd | 每个 block 算完 atomicAdd(g_out, partial) | G 不大时;一趟搞定 |
| cooperative groups 网格同步 | grid.sync() 一个 kernel 内做完两级 | 进阶,需 launch 支持 |
A100 上 grid-stride 已经把 G 压得很小(几百~上千),所以 atomicAdd 收尾或两趟 kernel 都很快,差别可忽略——
瓶颈始终是第一趟那次"完整读一遍数组",收尾只是零头。
reduction 不看 Compute%,要盯内存吞吐。编译跑 profile(A100 用 sm_80):
nvcc -O3 -arch=sm_80 reduce.cu -o reduce
ncu --set full --section MemoryWorkloadAnalysis ./reduce三个关键读数,按这个顺序判断:
| 指标 | 看什么 | 健康值(A100) |
|---|---|---|
DRAM Throughput | 占峰值带宽的百分比 = 你离极限多远 | ✅ 优化后应 ≥ 80% |
Memory Throughput / SOL | Speed of Light 里 Memory 那条 | ✅ 高且 Compute 低 = 对了 |
Bank Conflicts | shared 内冲突数 | ✅ 应为 0(第 3 节已治) |
N×4/带宽 → 跑 kernel 测实际耗时 →
实测带宽 = 读字节/耗时 → 对比峰值。当 DRAM Throughput 上到 80%+ 且 Duration 逼近理论下限,这个访存瓶颈算子就算调到头了。
再快不会来自 reduction 本身,而是融合(把 reduction 和上游算子合成一个 kernel,省掉中间结果的读写)。
| 阶段 | 动作 | 治的病 / 用的前节知识 |
|---|---|---|
| 定位 | 算术强度 0.25 → 访存瓶颈 | Roofline(L1) |
| 第一刀 | sequential addressing | 分支发散(L6)+ bank conflict(L5) |
| 第二刀 | warp shuffle 收尾 | 省 shared 往返与 sync |
| 第三刀 | grid-stride 多元素 | 合并访问(L3)+ 摊薄开销 |
| 验证 | ncu 看 DRAM Throughput ≥80% | 读 SOL 报告(L2) |
点选项看反馈。这些题交错了 roofline、合并、bank conflict、occupancy——这正是检验你是否真懂的方式。
__shfl_down_sync 的 mask 为什么是 0xffffffff?block 大小到底选 256 还是 512?
softmax/LayerNorm 里的 reduction 怎么和指数/归一化融合?把问题抛来。我是你的老师。
📘 Mark Harris, "Optimizing Parallel Reduction in CUDA"(NVIDIA 经典 slide) ——逐版本演示 7 步优化(交错→顺序→首次加载即相加→展开→多元素),本节的骨架就来自它,只是把收尾换成现代的 warp shuffle 并落到 A100。