<<<CUDA C++ 学习路线grid · block · warp · lane
EN中文
阶段 3 · 优化实战08 / 15约 30 分钟cuda/06_reduction.cu

归约优化七步：一个 kernel 榨出 30 倍

从最朴素的版本一路优化到 warp shuffle，每一步都说清楚为什么变快

学完这一课你会
  • 掌握并行归约这个最重要的 GPU 原语
  • 逐步理解 divergence、bank conflict、idle thread、循环开销各自的代价
  • 学会用 warp shuffle 完全绕开共享内存
  • 建立「每次优化都要能说出消除了哪个具体瓶颈」的思维习惯

求和、求最大值、softmax 的分母、L2 范数、注意力的 rowsum……这些全都是归约（reduction）。它是 GPU 编程里最经典的教学案例，因为它把几乎所有优化技巧都串在了一条线上。下面这七个版本来自 Mark Harris 的经典分享，我们逐个拆解每一步到底消除了什么瓶颈。

V1 — 交错寻址：能跑，但每一步都在分叉

reduce_v1：baseline
1__global__ void reduce_v1(const float* in, float* out, int n) {2    extern __shared__ float sdata[];3    int tid = threadIdx.x;4    int i = blockIdx.x * blockDim.x + tid;5 6    sdata[tid] = (i < n) ? in[i] : 0.0f;7    __syncthreads();8 9    for (int s = 1; s < blockDim.x; s *= 2) {10        if (tid % (2 * s) == 0) {           // 问题就在这里11            sdata[tid] += sdata[tid + s];12        }13        __syncthreads();14    }15    if (tid == 0) out[blockIdx.x] = sdata[0];16}

tid % (2*s) == 0 让 warp 内只有一部分线程干活，且干活的线程是交错分布的——第一轮是 0、2、4、6…，warp 内 32 个线程一半工作一半空转，100% 分叉。更糟的是 % 是取模运算，在 GPU 上比较昂贵。

V2 — 换索引方式：消除 divergence

reduce_v2：让活跃线程连续
1for (int s = 1; s < blockDim.x; s *= 2) {2    int index = 2 * s * tid;                // 把「谁干活」改成「干哪一格」3    if (index < blockDim.x) {4        sdata[index] += sdata[index + s];5    }6    __syncthreads();7}

工作量完全没变，但活跃线程变成了 tid 从 0 开始的连续一段。这样前几个 warp 是全员工作、后几个 warp 是全员空闲——warp 内不再分叉。代价是引入了新问题：index = 2*s*tid 让共享内存访问的 stride 变成 2、4、8…，bank conflict 出现了。

V3 — 顺序寻址：消除 bank conflict

reduce_v3：把循环反过来
1for (int s = blockDim.x / 2; s > 0; s >>= 1) {   // 从大步长往小走2    if (tid < s) {3        sdata[tid] += sdata[tid + s];               // 活跃线程访问连续地址4    }5    __syncthreads();6}

V4 — 加载时先加一次：一半线程不再是摆设

看 V3 的第一轮：blockDim.x 个线程里只有一半在工作。也就是说，你启动的线程有一半从头到尾只干了「把数据搬进共享内存」这一件事。既然如此，不如让每个线程在加载阶段就读两个元素并相加——block 数量直接减半。

reduce_v4：first add during load
1int tid = threadIdx.x;2int i = blockIdx.x * (blockDim.x * 2) + tid;        // 每个 block 覆盖 2 倍数据3 4sdata[tid] = (i < n ? in[i] : 0.0f) + (i + blockDim.x < n ? in[i + blockDim.x] : 0.0f);5__syncthreads();6 7for (int s = blockDim.x / 2; s > 0; s >>= 1) {8    if (tid < s) sdata[tid] += sdata[tid + s];9    __syncthreads();10}

V5 — 展开最后一个 warp：省掉多余的屏障

当 s <= 32 时，只剩一个 warp 在工作。此时 __syncthreads() 是纯粹的浪费——它是 block 级屏障，却只有一个 warp 需要同步。把最后 6 轮单独拿出来展开处理即可。

reduce_v5：warp 级展开（Volta 安全版）
1for (int s = blockDim.x / 2; s > 32; s >>= 1) {     // 循环到 32 为止2    if (tid < s) sdata[tid] += sdata[tid + s];3    __syncthreads();4}5 6if (tid < 32) {7    // 注意：不能再依赖「warp 内隐式同步」，Volta 之后必须显式来8    float v = sdata[tid] + sdata[tid + 32];9    for (int offset = 16; offset > 0; offset >>= 1) {10        v += __shfl_down_sync(0xffffffff, v, offset);   // 寄存器直接换数据11    }12    if (tid == 0) out[blockIdx.x] = v;13}

V6 — 模板化完全展开

reduce_v6：blockSize 作为模板参数
1template <unsigned int blockSize>2__global__ void reduce_v6(const float* in, float* out, int n) {3    extern __shared__ float sdata[];4    unsigned tid = threadIdx.x;5    unsigned i = blockIdx.x * (blockSize * 2) + tid;6    unsigned gridSize = blockSize * 2 * gridDim.x;7 8    float sum = 0.0f;9    while (i < n) {                       // grid-stride：一个 block 吃掉任意多数据10        sum += in[i];11        if (i + blockSize < n) sum += in[i + blockSize];12        i += gridSize;13    }14    sdata[tid] = sum;15    __syncthreads();16 17    // blockSize 是编译期常量，这些 if 会被完全消解，循环彻底展开18    if (blockSize >= 512) { if (tid < 256) sdata[tid] += sdata[tid + 256]; __syncthreads(); }19    if (blockSize >= 256) { if (tid < 128) sdata[tid] += sdata[tid + 128]; __syncthreads(); }20    if (blockSize >= 128) { if (tid <  64) sdata[tid] += sdata[tid +  64]; __syncthreads(); }21 22    if (tid < 32) {23        float v = sdata[tid] + sdata[tid + 32];24        for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);25        if (tid == 0) out[blockIdx.x] = v;26    }27}28 29// 调用时要显式实例化30reduce_v6<256><<<blocks, 256, 256 * sizeof(float)>>>(d_in, d_out, n);

V7 — 全 warp shuffle：共享内存只用来跨 warp

reduce_v7：现代写法
1__inline__ __device__ float warpReduceSum(float v) {2    for (int off = warpSize / 2; off > 0; off >>= 1)3        v += __shfl_down_sync(0xffffffff, v, off);      // 全程在寄存器里4    return v;5}6 7__global__ void reduce_v7(const float* in, float* out, int n) {8    float sum = 0.0f;9    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n;10         i += blockDim.x * gridDim.x) sum += in[i];11 12    sum = warpReduceSum(sum);                 // 第一级：warp 内归约，零共享内存13 14    __shared__ float warpSums[32];            // 最多 1024/32 = 32 个 warp15    int lane = threadIdx.x % warpSize;16    int wid  = threadIdx.x / warpSize;17    if (lane == 0) warpSums[wid] = sum;       // 每个 warp 只写一个数18    __syncthreads();19 20    // 第二级：由第 0 个 warp 把各 warp 的结果再归约一次21    sum = (threadIdx.x < blockDim.x / warpSize) ? warpSums[lane] : 0.0f;22    if (wid == 0) sum = warpReduceSum(sum);23 24    if (threadIdx.x == 0) atomicAdd(out, sum);   // 也可写成 out[blockIdx.x]25}
版本关键改动消除的瓶颈相对 V1
V1tid % (2*s)—（baseline）1×
V2改成 index = 2*s*tidwarp divergence~2×
V3循环步长从大到小bank conflict~4×
V4加载时先加一次一半线程空转~8×
V5展开最后一个 warp多余的 block 屏障~12×
V6模板 + 完全展开 + grid-stride循环与地址计算开销~20×
V7warp shuffle 两级归约共享内存往返~30×，逼近带宽上限
生产环境的写法
1#include <cub/cub.cuh>2 3void* d_temp = nullptr;4size_t temp_bytes = 0;5 6// 第一次调用只是问「需要多少临时空间」7cub::DeviceReduce::Sum(d_temp, temp_bytes, d_in, d_out, n);8CUDA_CHECK(cudaMalloc(&d_temp, temp_bytes));9 10// 第二次才真正执行11cub::DeviceReduce::Sum(d_temp, temp_bytes, d_in, d_out, n);

自测

先自己在心里回答一遍，再展开对照。答不上来的说明这一段值得重读。