说实话,并行归约算法是我这两年觉得最被低估的算法之一。前阵子帮一个物理引擎项目做性能剖析,粒子受力更新已经用上了很好的分段处理,反而是统计系统总能量、平均速度这种“看起来没技术含量”的小函数拖了后腿。单线程 for 循环求一亿个 float 的和,耗时高得离谱;我顺手开了几个线程,结果并发版比单线程还慢,忙活了一下午才意识到,问题全出在最后那一步合并上。
并行归约算法的核心并不复杂:把一组元素通过某个二元操作(求和、取最大值、按位与等)缩成一个值。但真要在 GPU 或多核 CPU 上把它做快,牵扯到数据依赖、内存带宽、线程调度、浮点精度,各种问题会一起冒出来。这篇东西我不会只讲理论,而是从一次实际性能优化出发,把并行归约里的原理、工程实现、容易踩的坑都梳理一遍。适合写图像统计、物理仿真、数值计算、机器学习预处理逻辑的人参考,尤其是正在和 GPU 上规约性能较劲的朋友。
1. 为什么“多开几个线程”算不对并行归约
1.1 看似简单的拆分为什么没用
先还原一下我当时写的“并发求和”。数组有 N 个 float,我想用两个线程:线程 A 从 0 加到 N/2,线程 B 从 N/2 加到 N,最后再让主线程把两个部分和相加。
cpp复制float sumA = 0.f, sumB = 0.f;
thread tA([&] { for (int i = 0; i < N / 2; ++i) sumA += data[i]; });
thread tB([&] { for (int i = N / 2; i < N; ++i) sumB += data[i]; });
tA.join();
tB.join();
float total = sumA + sumB;
跑出来的结果确实比单线程快了一点,但幅度远低于预期。当时我犯的最典型错误是只把“前半段 + 后半段”当成并行的全部,根本没意识到:真正限制性能的不是前 N/2 个数怎么累加,而是两个线程结束后产生的两个部分结果怎么合并。
这个例子极有代表性。如果只有最后一次合并,两个线程最终汇合时是一次加法,好像微不足道。但如果你保持这种“两两分裂再合并”的思路,在几百上千个线程的 GPU 上做同样的事情,每个线程都往同一个全局变量里加,那才是灾难。线程越多,合并竞争越严重,最后性能反而比单线程还差。
1.2 串行求和里藏着一条延迟链
为什么朴素单线程 for 循环会慢?先看它的运算形态:
code复制tmp = 0
tmp = tmp + data[0]
tmp = tmp + data[1]
tmp = tmp + data[2]
...
第 i 次累加用到了第 i-1 次累加的结果,这是严格的数据依赖链。CPU 或者 GPU 的浮点加法本身有延迟,哪怕现代处理器能够每个时钟周期发射好几个加法指令,这条依赖链上的每次加法都必须等上一次算出结果,所以延迟被原原本本地串起来。
假设加法延迟是 4 个周期,N 个数的串行求和就需要约 4N 个周期。吞吐量再高,也都被这条依赖链卡死。我当年试着在循环里写多个局部累加变量,让编译器做指令级并行,速度确实提升明显,但本质还是没有打破“一个人从头到尾算完”的模式。线程并行只是把这个延迟链拆成了 P 段,每段仍然是一条较短的链,最终还是需要一个合并步骤。
并行归约算法要解决的核心矛盾就是这个:N 个数,靠串行规则只能一条链走到底;但如果加法满足结合律和交换律,我们就可以改变求和顺序,让很多人各自算一块,再把结果用树形或分层方式合并,从而把整个计算过程里“必须顺序等待”的长度压到很低。
1.3 归约为什么能重排
这里有一个常常被一笔带过但其实很关键的点:并不是所有输入和操作都能随便重排。归约算法真正依赖的是二元操作满足结合律,比如:
- 加法:
(a + b) + c == a + (b + c) - 最大值:
max(max(a, b), c) == max(a, max(b, c)) - 位与、位或、乘法和最小值同理
只要具备这个性质,a0 + a1 + a2 + a3 既可以按 ((a0+a1)+a2)+a3 顺序算,也可以先算 a0+a1 和 a2+a3,再做一次总合并。并行归约正是利用了这种可重排特性,把原来严格的串行结构打散成“各算各的,最后集中合并”。
至于交换律,在大多数数据并行场景里也是成立的,所以可以让线程随机接任务。严格地说,交换律只在你想彻底打破顺序时才有必要;但为了让各线程负载均衡,我们通常会假设数据顺序不影响最终语义。对整数加法和 max/min 这自然是成立的,对浮点加法则要注意精度问题,这个我在后面单独说。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 从 n 步到 log2(n) 步:树形归约的核心模型
2.1 两两配对的直观模型
并行归约的教科书模型是树形归约。假设有 8 个数,你先把它分成 4 组,每组两个元素同时做加法:得到 4 个中间结果;再把这 4 个中间结果分成 2 组相加;最后剩下两个结果加一次。总共只需要 3 轮,而不是 7 轮。
code复制第 1 轮: a0+a1, a2+a3, a4+a5, a6+a7
第 2 轮: (a0+a1)+(a2+a3), (a4+a5)+(a6+a7)
第 3 轮: (a0+a1+a2+a3) + (a4+a5+a6+a7)
如果数据量是 N,足够多的线程同时工作,树的高度就是 ceil(log2(N))。8 是 3 轮,1024 是 10 轮,一百万个数也只需要大约 20 轮。这就把“必须顺序等待”的步数从 O(N) 降到了 O(log N)。
很多初学者第一次看到这个模型时会兴奋:那是不是意味着并行归约能无限快?不是。关键限制在于,每一层都要等待上一层所有运算完成,同时机器上的计算单元数量是有限的。更现实的问题是,你必须先把 N 个数从内存里读出来,这个操作本身就至少需要 O(N) 时间。树形归约压缩的不是“读数据的时间”,而是“读完数据之后,把一堆局部值汇总成最终值所花费的同步轮数”。
2.2 一次遍历是绕不过去的下限
这个认知非常重要。在绝大多数并行求和、求最大值场景里,数据量远大于处理单元数量,比如 1 亿个 float 让 1024 个线程来处理,每个线程要扫好几万个数。这样每个线程内部仍然是串行累加,只不过线程间互不干扰;只有当每个线程手里的局部段处理完,树形结构才开始真正发挥作用。
所以我在评估一个并行归约实现时,第一件事永远是算内存带宽能不能打满。比如 1 亿个 float 占 400MB,假设 GPU 读主存的带宽是 1.5TB/s,那么光是把数据完整读一遍的理论时间就是:
code复制400MB / 1500GB/s ≈ 0.27ms
如果你的内核跑到了 0.35ms,我会说它已经相当优秀,因为还要算上读写部分结果、线程同步和启动开销。反过来,如果你测出来是 2ms,那问题几乎一定不是“树形归约不够聪明”,而是数据读取方式没有让带宽打满。
这也是我在生产环境里减少许多无谓调优的经验:先盯带宽,再盯归约树。很多人一上来就优化共享内存里的合并逻辑,结果 profile 一看,全局内存读取根本没合并,那才是浪费时间。
2.3 最大值的树形归约与求和完全同构
求和需要结合律和交换律,求最大值同样需要。max(1, 5, 3, 2) 和 (1+5+3+2) 的差别只在初始值怎么设、两两合并时用哪个运算符。树形归约的每一层都可以复用同一个框架,只要把“加法”替换成“max”。
所以很多现成库把归约做成一个模板函数,接受一个自定义的二元操作,比如 CUDA CUB 的 BlockReduce、C++ 标准库的 std::reduce、Thrust 的 thrust::reduce。模板化的好处是求和、最大值、点积、自定义结构体合并都能走同一套并行逻辑。
但这带来一个容易被忽略的问题:自定义操作符必须真的满足结合律。如果你在归约时做的是求“二维向量相加”,那没问题;如果你写了一个“把当前值追加到链表尾部”的操作,它虽然看上去像归约,但无法并行重排。遇到这种情况应该在建链阶段就另想办法,而不是硬套并行归约。
3. 线程束与共享内存里的工程实现:一份可跑的 CUDA 归约
3.1 先把数据切成很多段
纸上谈兵到这里差不多了,直接上工程实现。下面这段是我自己项目里经常使用的一种 CUDA 分段归约模式,主要分两阶段:
- 把输入数组按
gridDim.x * blockDim.x分成若干大段,每个线程按跨度累加自己的那一段。 - block 内把各自线程的部分汇总到共享内存,用树形归约得到该 block 的部分结果,写回一个
partial数组。 - 之后可以由一个 block 对上一步的
partial数组再做一次归约,或者直接在核函数最后用原子操作合并。
第一次写并行归约的人,最容易直接上第三种“全局原子操作”。但 GPU 上每个线程都去原子加一个全局变量,会让全局内存原子单元成为巨大瓶颈。正确做法是每个 block 先归约出一个数,再让 block 之间做最终合并。这样需要原子操作的次数大约是 block 数量,而不是线程数量。
cpp复制__global__ void blockReduceKernel(const float* input,
float* partial,
size_t n) {
// 动态共享内存,大小在启动核函数时设置
extern __shared__ float sdata[];
size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
size_t stride = gridDim.x * blockDim.x;
// 第一阶段:每个线程按大步长遍历自己负责的元素。
// 这种访问方式让相邻线程访问相邻地址,全局内存读取是可以合并的。
float sum = 0.0f;
for (size_t i = tid; i < n; i += stride) {
sum += input[i];
}
sdata[threadIdx.x] = sum;
__syncthreads();
// 第二阶段:共享内存里的树形归约。
for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
if (threadIdx.x < s) {
sdata[threadIdx.x] += sdata[threadIdx.x + s];
}
__syncthreads();
}
if (threadIdx.x == 0) {
partial[blockIdx.x] = sdata[0];
}
}
启动方式类似:
cpp复制int threads_per_block = 256;
int blocks_per_grid = 1024; // 根据数据量和 GPU 占用情况调整
size_t smem_bytes = threads_per_block * sizeof(float);
blockReduceKernel<<<blocks_per_grid, threads_per_block, smem_bytes>>>(
device_data, device_partial, n);
__syncthreads() 是这段代码里面最重要的东西。共享内存是 block 内所有线程共享的,如果没有同步,线程 A 在树形归约第一轮读到的 sdata 可能还没被线程 B 写进去。每一轮结束之后都必须加一道栅栏,保证所有线程都完成当前这轮的写操作,才能进入下一轮。漏掉任何一个 __syncthreads(),在小规模数据上可能侥幸跑对,数据量稍大、调度稍有变化就会出随机错误。
3.2 让每个线程串行累加尽量多的数
上面的代码里第一阶段用一个 for 循环 for (i = tid; i < n; i += stride) 来遍历。这个写法看起来每个线程跳着访问,但实际上相邻线程在同一时刻访问相邻地址,GPU 可以把这些访问合并成少量大型内存事务。这个模式对带宽利用至关重要。
如果你反过来写一个朴素的“每个线程负责连续的一段”,比如线程 0 读 0 到 9999,线程 1 读 10000 到 19999,理论上没问题,但在 GPU 上相邻线程同一时刻读取的数据地址相隔 10000 个 float,无法合并,内存系统会浪费大量带宽。在多核 CPU 上倒是不一定差,因为 CPU 有独立的缓存行机制,但在 GPU 归约内核中,我强烈建议使用大步长连续遍历。
第一次实现跑通后,我通常会加一个简单优化:每个线程内部维护多个累加变量,手动做 4 路或 8 路展开。这样能隐藏内存访问延迟,也能让加法依赖链有更多独立并行链。
cpp复制float sum0 = 0.f, sum1 = 0.f, sum2 = 0.f, sum3 = 0.f;
size_t n_aligned = n / 4 * 4;
for (size_t i = tid; i < n_aligned; i += stride * 4) {
sum0 += input[i];
sum1 += input[i + stride];
sum2 += input[i + 2 * stride];
sum3 += input[i + 3 * stride];
}
for (size_t i = n_aligned + tid; i < n; i += stride) {
sum0 += input[i];
}
float sum = (sum0 + sum1) + (sum2 + sum3);
注意这里最后的 (sum0 + sum1) + (sum2 + sum3) 也是一个小型树形归约,而不是写成 sum0 + sum1 + sum2 + sum3 这种连续依赖链。在浮点语义上两者结果不完全一样,但后者会引入 3 次串行加法延迟,前者两次加法可以并行执行,效率更高。
3.3 线程束 Shuffle:把树形归约再压短
共享内存版本适合教学,也适合绝大多数场景。不过如果你用的是现代 NVIDIA GPU,还可以利用 warp shuffle 指令,让一个线程束内的 32 个线程直接交换寄存器数据,而不经过共享内存。
大致思路是这样的:第一步已经让每个线程有了自己的部分和,接下来不再写共享内存,而是调用 __shfl_down_sync 之类的指令,把相邻线程的值“往下传”,然后在寄存器里累加。
cpp复制for (int offset = 16; offset > 0; offset >>= 1) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
}
这段循环结束以后,线程 0 里就保存了该 warp 内 32 个线程之和。它比共享内存树形归约少了访问共享内存的开销,也不需要那么多 __syncthreads,因为 warp 内的线程天然是锁步执行的。再往上合并不同 warp 的结果时,仍然需要共享内存。
对于初学者,我不建议一上来就追这个优化。先用共享内存版本跑通、确认结果正确,再用 profiler 去看瓶颈。如果共享内存版本已经能达到 80% 以上带宽利用率,再追求 shuffle 通常只有锦上添花的效果。
4. 从求和切到最大值:算子替换背后的边界条件
4.1 初始值决定了正确性
求和换成求最大值时,很多人觉得只要把 + 换成 max 就行了,但漏掉了初始值的处理。求和的初始值是 0,因为 a + 0 = a;而求最大值的恒等元是负无穷,因为 max(a, -INFINITY) = a。
如果初始化成 0,在一组全为负数的数据上求最大值,结果会错误地变成 0,而不是真正的最大值。我在测试天气数据时遇到过这个问题,温度数组里全是零下,结果返回 0,排查了半天。
解决方案一般是:确保数组非空,然后把归约初始值设为 -FLT_MAX 或 -INFINITY。如果数组可能为空,最后要单独判断结果仍然等于负无穷的情况。最小值则相反,初始值为 +FLT_MAX 或 +INFINITY。
在 CUDA 里还有一点要注意:max(float, float) 在语义上可能会被编译成比较指令,而 fmaxf 是 CUDA 提供的数学函数,它对 NaN 的处理规则不同。二者并不总是可以完全互相替代。如果你希望“忽略 NaN,只要正常数值里的最大值”,就要明确使用对应的函数;如果希望“NaN 参与传播,结果也是 NaN”,就要使用普通比较表达式。这个决策应该在项目里统一,否则会有隐蔽错误。
4.2 结合律足够,但浮点数加法不精确
整数加法、整数最大值是完全精确的,改变计算顺序不改变结果。但浮点加法不满足真正的结合律,因为浮点数的表示精度有限:
code复制(16777216.0f + (-16777216.0f)) + 1.0f → 1.0f
16777216.0f + ((-16777216.0f) + 1.0f) → 0.0f
这导致一个现象:同一个数组,串行求和的结果和并行归约的结果可能在最后几位有差异。这本身是浮点计算的正常现象,但如果你的项目需要可重复结果,或者数值很敏感,就要针对性地处理。
我常用的折中方案是分层固定归约顺序:每个线程先累加固定长度的数据段,然后固定线程数、固定共享内存归约树形层级,这样多次运行的结果能保持稳定。它不改变浮点误差的大小,但能从机制上保证结果可复现。如果精度要求更高,可以在局部累加时使用 Kahan 补偿求和;但 Kahan 求和里面存在较强的数据依赖,会降低并行度,一般只用在最后一层小规模归约上,而不是所有线程的原始遍历阶段。
4.3 点积等组合操作也是归约
归约操作不限于一个输入数组。机器学习里经常要算两个向量的点积:
code复制dot = sum_i(a[i] * b[i])
把乘法和加法放一起看,它是一个先做元素级乘法、再做 sum 归约的复合流程。实现上可以先让每个线程读一对 a[i]、b[i],乘起来累加进局部和,然后走和普通求和完全一样的树形归约。基本框架不用改,只需把“读一个数并累加”改成“读两个数并累加乘积”。
同理,计算均值和方差、计算 L2 范数、统计直方图里的最大值,都可以用同一套并行归约策略。本质上你应该抽象出“局部算子”和“合并算子”:局部线程读到的多个元素先结合成一个中间结果,中间结果再按树形合并。局部阶段可以做得重一点,合并阶段专注于少量数据即可。
5. 我用一块 GPU 跑出来的四个性能瓶颈
5.1 别把内存带宽浪费在错误的数据排布上
我最早做基准测试时,用的是一个 1 亿 float 数组,理论读取量 400MB。我第一版内核跑了大约 1.8ms,看起来很“快”,但我算了下,假设 GPU 带宽 1.5TB/s,理论下限只有 0.27ms,差了接近 7 倍。
Profile 结果出来,问题不是共享内存归约,而是全局读取没有完全合并。原因是我为了“让每个线程访问连续数据”,把数组按 blockIdx.x * blockDim.x + i 这种顺序切成很多小块,每个线程一次性读很长一段连续数据,结果相邻线程在同一时刻访问的地址相隔很远。
修正方式就是前面代码里写的:相邻线程访问相邻地址,再用大步长跳过循环。只改了这一个读取模式,1 亿 float 的归约就降到了接近 0.4ms。所以如果你发现一个并行归约实现数据量增加时性能变得特别差,先去查数据访问模式。
5.2 共享内存 stage 之间的同步真的不能省
还有一次,同事为了“优化”,把共享内存树形归约里的几个 __syncthreads() 删掉了一部分。小数据规模下测试结果碰巧正确,因为线程块数量少,调度模式刚好没踩到问题。等换到大数据量,结果开始随机出现偏差,有时候就差几个数,特别像是精度问题,其实是典型的共享内存数据竞争。
排查这类 bug 的经验是:用一个很小的数组,比如 1024 个元素,反复跑一万次,对比单线程 CPU 归约结果。如果一万次里出现哪怕一次不一致,那基本可以断定是数据竞争;如果每次都一致,那更可能是浮点精度或边界条件。
后来我用 compute-sanitizer 检查才发现未同步的共享内存读写。这个工具会在检测到 race 时直接报告,比靠肉眼观察结果靠谱得多。写 CUDA 归约内核的调试阶段,把同步检查打开是一个非常值得养成的习惯。
5.3 分支发散最后几轮影响有限,但别在全局遍历阶段发散
树形归约最后几轮里,随着步长 s 不断减半,实际参与运算的线程数量越来越少:
code复制s=128 时前 128 个线程在干活
s=64 时前 64 个线程在干活
...
在 NVIDIA GPU 上,一个 warp 是 32 个线程,所以总会存在一个 warp 只有一半或四分之一线程在工作的情况。这种“分支发散”会降低该 warp 的执行效率,但因为最后几轮操作次数呈对数减少,所以对整个核函数的影响通常很小。很多优化文章里把它说得很严重,实际测试里它的影响往往排在带宽之后。
真正不该发散的是第一阶段的大循环。如果你的判断条件依赖数据内容,比如 if (input[i] > 0) sum += input[i],在线程束内产生不同分支,性能可以掉得很明显。条件归约在语义上虽然可行,但最好先用元素级过滤改写,让它成为一个“先掩码,再归约”的流程,避免 warp 里一部分线程执行加法、另一部分等待。CPU 并行归约同样会受到分支预测惩罚影响。
5.4 block 的最终结果要避免每线程一次全局原子操作
我见过一些实现里没有中间的 partial 数组,每个线程得到局部结果后就 atomicAdd(&result, local_sum)。数据量小时没感觉,数据量一大,全局原子操作的重试和等待会让内核卡在内存系统上。
我建议采用两级甚至三级合并流程:第一级每个 block 归约出一个 partial,第二级只启动一个 block 把若干个 partial 再统一归约。这样全局原子操作次数大致从“线程数量”降到“block 数量”。如果 block 数量还是很大,可以再做第二次多块归约,最后再让一个 block 收尾。
在 CPU 版本里也是同理。如果几百个线程各自算出局部和,最后直接 sum += local_sum 可能因为竞争导致结果错误;正确的做法是先把每个线程的局部和写到一个数组,再在主线程里循环合并,或者使用 omp critical / 原子变量。全局临界区不是不能有,但只应该出现在“最后一个阶段”,而不是每个元素都进临界区。
6. 多核 CPU 场景:伪共享、临界区与库的抉择
6.1 OpenMP 帮你隐藏了最终合并逻辑
在 CPU 上用 OpenMP 做并行归约,是最简洁的路径之一:
cpp复制double total = 0.0;
#pragma omp parallel for reduction(+:total)
for (int i = 0; i < n; ++i) {
total += data[i];
}
编译器会为每个线程生成一个私有副本,先并行累加,再在结束处用一种相对高效的方式合并。它会自动处理局部变量可见性,所以比你自己写 thread_local 再锁要安全得多。但要注意,reduction(+:total) 的合并策略在编译器和运行库实现中不完全相同,多次运行通常不会改变浮点结果,但不同编译器之间可能与串行结果略有差异。
如果数据量非常大,我更倾向于让每个线程先按固定步长遍历数组,而不是让 OpenMP 自动分配连续区间。这里还是那个规律:线程之间访问模式会影响缓存命中率和内存带宽。CPU 受缓存行影响,连续区间分配在某些架构下反而更好;没有统一标准,必须实测。一般我会先用块状分配,让线程尽量访问连续内存,因为 CPU 缓存更友好。
6.2 手动实现时,伪共享比锁更隐蔽
如果出于某些原因不能使用 OpenMP 或者标准库,比如你在做嵌入式、游戏引擎线程池,那就要手写归约。最安全的做法是每个线程只维护一个局部和,最后把所有局部和放入数组,由主线程做最终合并。
一个非常隐蔽的问题是伪共享。假设你定义:
cpp复制float partial_sum[8];
8 个 float 只占 32 字节,而现代 CPU 缓存行通常是 64 字节。这意味着两个相邻线程各自写 partial_sum[0] 和 partial_sum[1] 时,因为两个变量可能落在同一个缓存行里,其中一个线程的写入会导致整个缓存行在多个核之间反复失效。表面上两个线程没有共享变量,但缓存行层面发生了隐式共享,性能会急剧下降。
解决办法是对齐每个线程的局部槽位到 64 或 128 字节:
cpp复制struct alignas(64) ThreadPartial {
float sum;
};
ThreadPartial partials[256];
如果只是做一次归约,最终合并由主线程串行处理 P 个部分和,其实没有太大性能风险。需要担心的是反复调用归约时,每一次都触发伪共享抖动,那就必须把这个 cache line 对齐做对。
6.3 最后的合并用什么:锁、原子操作还是单线程
P 个线程都把局部和算完以后,合并 P 个数有很多选择:
- 让主线程循环加 P 次,简单,但主线程可能成为瓶颈。
- 用
std::atomic<float>在 8 个线程里各fetch_add一次,8 次原子竞争,通常可以接受。 - 把 P 个局部和写到一个数组,再启动一小段递归归约,模拟 GPU 上的树形合并。
我的经验是:当 P 小于 32 时,直接单线程合并局部数组最简单,不会成为性能瓶颈;P 达到几百上千时,再考虑分两层树形合并。很多多线程框架里“最后一步用栅栏加单线程”其实就够了,因为前面每个线程都处理了一大片数据,合并阶段的操作数只有线程数那么少。
如果我们在更上一层做分布式归约,问题就变成节点之间的 all-reduce。经典做法是 ring-all-reduce:每个节点先做本地归约,然后让所有节点形成一个环,把局部块传给下一个节点并累加,经过 P-1 轮后每个节点都拿到完整结果。这套思路和 GPU block 归约是同构的,只是网络传输成了新的瓶颈。平时只写单机代码的话,不需要掌握到那么深,但理解了分层归约思想后,分布式场景也要容易很多。
6.4 标准库已经替你造好了轮子
C++17 之后,标准库提供了 std::reduce 和并行执行策略:
cpp复制#include <execution>
float sum = std::reduce(std::execution::par_unseq,
data.begin(), data.end(), 0.0f);
它不一定比手写 GPU kernel 快,但和手写 CPU 线程池相比,正确性和可移植性都更好。标准库的底层实现会根据平台自动选择合适的并行方式,也可能在内部做向量化。如果你只是希望从串行改成并行,这是最省心的方式。
但现成库也有局限。CUDA 生态里的 Thrust、CUB 都很优秀,尤其 CUB 的 BlockReduce 提供了很多我自己很难达到的调优细节。然而,如果你需要把归约结果继续留在某个自定义 kernel 里使用,而不想把中间数组读回 CPU,那么写一个和自己的算子深度融合的归约 kernel 会更好。库给你的是通用性,你换来的是性能定制空间。这个决策没有绝对答案,跟我写过的大多数优化问题一样,需要回到实际场景做基准测试。
7. 什么时候别自己写归约,以及我的最终选择
经常有人问我:学并行归约是不是就要手写 CUDA kernel?其实不是。如果你只是需要一个可靠、快速、跨平台的归约,我更倾向先用现成库解决。Thrust、CUB、oneTBB、OpenMP、标准库并行策略,任意一个都经过大量测试,比草率手写更稳。
但有以下几种情况,我仍然会手动实现归约:
- 数据规模固定,且归约被高频调用,比如每个物理帧都要对几十万粒子的速度做归约。此时自定义 kernel 可以省去泛型抽象带来的额外开销,也可以结合粒子更新逻辑做融合。
- 需要把归约结果用在 GPU kernel 内部,不希望中间结果频繁访存。
- 归约算子不是简单的加法或 max,而是项目自己的复杂结构体合并,现成库的二进制组合子会降低可读性,甚至无法使用。
- 需要精确控制浮点处理顺序,保证结果可复现。
如果决定手动写,我建议按这样的顺序推进:先写出最简单、逐块遍历的版本,用单线程结果校验;再检查全局内存访问是否合并;第三步看共享内存树形归约是否存在 race;最后再考虑 unroll、shuffle、向量化加载这类进阶优化。不要一上来就把所有技巧叠上去,否则出了问题根本不知道是哪一步改坏的。
关于并行归约算法,我个人的体会是:它的核心价值不在“求和还能快多少”,而在于它提供了一套优雅的思维模型——把一个看起来强顺序依赖的流程,通过结合律重排成可并行的树形流程。这套思维在 CPU 多线程、GPU kernel、分布式 all-reduce、设计数据库聚合算子时都能复用。下次再有人拿“数太多求和慢”来找你,先别急着写数行并发代码,想清楚数据怎么读、分几层合并,再动手不迟。
