Reduction 与 Scan：warp shuffle、block 级与 grid 级归约

求和、求最大值、求范数——归约（reduction）把一个数组压成一个标量，是最基础的协作类算子，也是理解 warp 级原语和 shared memory 协作的最佳入口。后面的 softmax、LayerNorm、FlashAttention 里都嵌着归约。这一篇从朴素归约一步步优化，并用实测说明一个容易被忽略的事实：优化点对不对，取决于瓶颈在哪。

一、归约的并行结构：树形

串行求和是 \(O(n)\) 次相加。并行化的标准结构是二叉树归约：每一步把相邻元素两两相加，元素数减半，\(\log_2 n\) 步后剩一个。block 内 256 个元素 8 步归约完成。

难点不在算法，在于怎么让 warp 高效执行这棵树。同样的树形结构，不同的索引方式会带来分支发散和 bank conflict，性能差几倍。下面用一个隔离实验对比三种写法——为了让 block 内归约树本身成为瓶颈，实验让每个 block 反复执行归约树 20000 次，排除数据加载的干扰。

二、版本一：交错寻址 + 取模（最差）

最直觉的写法：第 \(s\) 步让 t % (2s) == 0 的线程把 sd[t] 和 sd[t+s] 相加。

for (int s = 1; s < blockDim.x; s *= 2) {
    if (t % (2*s) == 0) sd[t] += sd[t+s];
    __syncthreads();
}

两个问题（外加一个常被忽略的隐性成本）：

• 分支发散：t % (2*s) == 0 在 warp 内只有部分线程满足。第一步每 2 个线程才有 1 个活跃，warp 内严重发散（第 02 篇）。
• 取模与乘法本身的开销：t % (2*s) 每步都要做一次整数乘法和取模，整数取模在 GPU 上并不便宜。版本一相对版本二的差距不全来自发散，这部分指令开销也占一份。
• bank conflict：活跃线程访问的 sd[t] 间隔随 \(s\) 增大，容易落到同一 bank（第 05 篇）。

三、版本二：顺序寻址（消除发散与冲突）

改成从大步长往小步长走，让前 \(s\) 个连续线程活跃：

for (int s = blockDim.x/2; s > 0; s >>= 1) {
    if (t < s) sd[t] += sd[t+s];
    __syncthreads();
}

t < s 让活跃线程是连续的一段——前几步整个 warp 要么全活跃要么全空闲，没有 warp 内发散；访问 sd[t] 和 sd[t+s] 也是连续的，避免了 bank conflict。这是 Mark Harris 经典归约优化里的关键一步。

四、版本三：warp shuffle（去掉最后几步同步）

当活跃线程缩到 32 个以内（一个 warp），还在用 shared memory + __syncthreads() 就浪费了——warp 内线程可以用 shuffle 原语直接交换寄存器，不经过 shared，也不需要 block 级同步。

// 先用 shared 归约到 32 个
for (int s = blockDim.x/2; s >= 32; s >>= 1) {
    if (t < s) sd[t] += sd[t+s];
    __syncthreads();
}
// 最后一个 warp 用 shuffle 完成
if (t < 32) {
    float v = sd[t];
    for (int o = 16; o > 0; o >>= 1)
        v += __shfl_down_sync(0xffffffff, v, o);
    if (t == 0) sd[0] = v;
}

__shfl_down_sync(mask, v, o) 让每个 lane 读取 lane+o 的寄存器值 v，5 次就把 32 个值归约成 1 个。好处：省掉了最后 5 步的 __syncthreads()（block 级栅栏不便宜），warp 内通信走寄存器直连。

五、实测：三个版本的差距

RTX 3060 Ti，256 线程/block，隔离归约树（重复 20000 次）：

从最差到最好快 3.4 倍。顺序寻址消除发散和冲突拿到第一个 1.7 倍，warp shuffle 省掉同步再拿一个约 2 倍。这组数字量化了”索引方式”这个看似无关紧要的细节对协作类 kernel 的真实影响。

六、一个反直觉的事实：单遍归约里这些优化看不见

上面的差距是在隔离归约树时测的。如果直接做一次完整的大数组归约会怎样？再测一次：每个 block 用 grid-stride loop 把它负责的那部分元素加到寄存器，再做 block 内归约树，输出每 block 一个部分和。\(n=2^{24}\)：

三个版本完全一样。原因是单遍归约的时间几乎全花在从 global 读 1600 万个 float 上（受访存限制，约 330 GB/s），block 内那棵只处理 256 个元素的归约树相比之下微不足道。归约树的优化被访存彻底掩盖。

这正好印证 Roofline 篇 的结论：优化要对准瓶颈。单遍归约是 memory-bound 的，该优化的是访存（合并加载、float4 向量化加载、减少遍数），而不是归约树。归约树的优化只在树本身成为瓶颈时才有意义——比如 block 数很多、每个 block 数据很少，或归约被反复调用的场景。

七、grid 级归约：跨 block 怎么合并

block 之间不能直接同步（第 04 篇），所以全数组归约要把”每 block 一个部分和”再合并。两种常见做法：

• 两趟 kernel：第一趟每 block 输出一个部分和到一个小数组，第二趟（或在 host）归约这个小数组。简单、确定，适合需要可复现结果的场景。
• 原子加：每 block 算完直接 atomicAdd 到一个全局累加器。省一趟 kernel，但浮点原子加的累加顺序不确定，浮点不满足结合律，结果会有微小的 run-to-run 差异。对数值敏感的场景要权衡。

// block 内归约出 blockSum 后：
if (t == 0) atomicAdd(out, blockSum);   // 注意：浮点累加顺序不确定

八、Scan（前缀和）：归约的近亲

scan（前缀和）输出每个位置的累计值，比归约难一档，因为每个输出都依赖前面所有输入。并行 scan 的经典算法是 Blelloch 的 work-efficient 双阶段（up-sweep 归约 + down-sweep 分发），warp 内同样可以用 __shfl_up_sync 高效实现。scan 是 radix sort、流压缩（stream compaction）、稀疏结构构建的基础原语，CUB 库提供了高度优化的 BlockScan/DeviceScan 实现，工程中通常直接用 CUB 而非手写。

九、小结与下一步

• 归约是树形结构，难点在让 warp 高效执行：顺序寻址消除发散和 bank conflict，warp shuffle 省掉最后几步同步。
• 隔离归约树时三个版本差 3.4 倍；但单遍大数组归约受访存限制，三者实测一样（约 330 GB/s）——优化必须对准真实瓶颈。
• 跨 block 用两趟 kernel（确定）或原子加（省一趟但顺序不确定）合并。
• scan 比归约复杂，工程中优先用 CUB。

掌握了协作类算子的基本功，下一篇进入最核心、也最能体现 tiling 思想的算子——GEMM：从朴素实现到 shared memory tiling 与寄存器分块。