ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

并行前缀和算法深度解析:两种经典扫描与工程优化

并行前缀和算法深度解析:两种经典扫描与工程优化 先回答很多初学者第一次接触并行计算时的同一个困惑前缀和这种“每个输出都依赖前一个输出”的东西真的能并行吗我当年读 Parallel Prefix 相关论文时也有同样的疑问——明明 p[i] p[i-1] a[i] 是一条铁链每一步都在等上一步的结果怎么并行答案是能。而且并行前缀法Parallel Prefix是整个并行计算体系里最核心的原语之一地位不亚于归约和排序。它把看起来串行无解的前缀和问题从 O(n) 深度压到了 O(log n) 深度。这篇博文就把这件事彻底讲透内容包括两种经典算法Hillis-Steele 和 Blelloch的完整推导、真实 GPU/CPU 上的工程化组织方式、常见应用场景以及我实际落地时踩过的坑。适合刚接触并行计算的学生、准备面试的开发者、还有想优化热点函数的工程师。1. 前缀和明明每个输出都依赖前一个并行从哪来先明确问题。给定数组 a[0..n-1] 和一个二元操作符 ⊕前缀计算要输出 p[i] a[0] ⊕ a[1] ⊕ ... ⊕ a[i]。当 ⊕ 是加法时这就是前缀和prefix sum也叫 scan。顺序写法非常自然std::vectorint p(n); p[0] a[0]; for (int i 1; i n; i) p[i] p[i - 1] a[i];这段代码的循环每轮都依赖上一轮的结果。单个处理器上它确实无法加速这是串行依赖链的天然瓶颈。但并行计算关心的从来不是“从 0 到 n 依次计算”而是“整个数组能不能用多核/众核同时算出一部分”。突破口在操作符 ⊕ 满足结合律即 (x ⊕ y) ⊕ z 等于 x ⊕ (y ⊕ z)。加法和乘法满足结合律矩阵乘法也满足结合律只要这个性质在前缀计算就有并行的数学基础。我常用一个生活场景来类比。商场收银台开了很多窗口每个顾客想知道“排到我为止之前所有人一共消费了多少”。如果只有一名收银员只能一个个往后加。但可以让每组 10 个顾客先自己算组内小计组长把 10 个数字上报给总台总台把组长们的前缀算出来再下发给每个组长组长拿到“自己这一组之前的累计总额”后回去给组内每个人加上偏移量。这就是并行前缀的雏形先算局部再算块间前缀最后回填修正。分块越大块数越少块间扫描的压力越小。这里引出一个重要概念work 和 depth 的权衡。顺序计算的 depth 是 n总工作量也是 n。并行前缀的 depth 可以做到 O(log n)但总加法次数会变多。GPU 上有几千个核多付出一点加法能换来深度从百万级降到几十级这就是并行前缀真正有价值的原因。第一个提出的实用算法是 Hillis-Steele 扫描它的深度很小但总工作量是 O(n log n)后来 Blelloch 给出 work-efficient 版本总工作量 O(n) 但深度约 2 log n。这两者构成了理解整个领域的地基。2. Hillis-Steele把 offset 翻倍最简单也最“浪费”的扫描Hillis-Steele 算法的思想可以概括为一句话每一轮把数据和自己左侧相距 2^k 的元素相加offset 每轮翻倍。伪代码如下for k 0 to log2(n) - 1: offset 2^k for i in parallel: if i offset: x[i] x[i - offset] x[i]第 0 轮每个元素和左边的 1 个元素相加第 1 轮和左边的 2 个元素相加第 2 轮和左边的 4 个元素相加。第 k 轮结束时x[i] 保存的是从 a[i-2^k1] 到 a[i] 这一共 2^k 个元素的和当 offset 覆盖整个数组时每个位置就得到了从 0 到 i 的完整前缀。用一组具体数据走一遍最直观。设初始数组x0 [3, 1, 7, 0, 4, 2, 1, 6]第 1 轮offset1x1 [3, 4, 8, 7, 4, 6, 3, 7]第 2 轮offset2x2 [3, 4, 11, 11, 12, 13, 7, 13]第 3 轮offset4x3 [3, 4, 11, 11, 15, 17, 18, 24]最后一行就是 inclusive prefix sum。这个过程每轮所有线程可以同时读写天然适合 SIMD 架构和 GPU warp 内的并行。实际上现代 GPU 上做 warp 级扫描时一个 warp 只有 32 个 lanelog2(32) 5 步就能完成配合 __shfl_up 指令不需要共享内存延迟极低。所以在最底层的小规模扫描里Hillis-Steele 反而是最优选择。但它的问题也很致命总工作量是 O(n log n)。每一轮都要读写整个数组n100 万时 log2(n)≈20等于做了 2000 万次加法而理论上只需要约 200 万次。对于带宽受限的 GPU多出来的十几倍内存流量会直接把性能拖垮。所以它适合数据量小、同步成本低的场景不适合大规模全局扫描。如果你只是要尽快跑通一个 Demo可以用它想追求性能看下一节的 Blelloch。3. Blelloch先归约后回填work-efficient 的标准打法Blelloch 算法是真正把总工作量压到 O(n) 的经典方案核心分两阶段up-sweep向上归约和down-sweep向下回填。它的巧妙之处在于up-sweep 阶段本来只是做一次归约却把所有子段和保留在数组的某些位置上down-sweep 再利用这些子段和把“左侧所有元素之和”逐个推回给每个位置。3.1 up-sweep把兄弟节点的和往上累加从叶子节点出发层与层之间用 offset 区分。每轮把相隔 offset 的两个相邻节点相加结果存到右边节点for offset 1 to n/2: for i in parallel: if i % (2*offset) offset - 1: x[i] x[i - offset] x[i]还是用刚才那组数据。初始x [3, 1, 7, 0, 4, 2, 1, 6]offset1把相邻元素两两相加x [3, 4, 7, 7, 4, 6, 1, 7]offset2把上一轮结果中距离为 2 的兄弟相加x [3, 4, 7, 11, 4, 6, 1, 13]offset4x [3, 4, 7, 11, 4, 6, 1, 24]此时 x[7]24 是整个数组的总和。更重要的是x[3]11 保存了前 4 个元素的小结x[1]4 保存了前 2 个元素的小结x[5]6 保存了 a[4..5] 的小结。这些“子段和”在 down-sweep 里会派上大用场。3.2 down-sweep从根开始把左侧累计和推回去先把根节点 x[7] 置为 identity 元素加法是 0。然后从最大的 offset 开始逐步缩小for offset n/2 down to 1: for i in parallel: if i % (2*offset) offset - 1: tmp x[i - offset] x[i - offset] x[i] x[i] tmp x[i]继续上面的例子先令 x[7]0x [3, 4, 7, 11, 4, 6, 1, 0]offset4处理 i7tmp x[3] 11 x[3] x[7] 0 x[7] tmp x[7] 11 0 11得到x [3, 4, 7, 0, 4, 6, 1, 11]offset2处理 i3 和 i7i3: tmp x[1]4, x[1]x[3]0, x[3]404 i7: tmp x[5]6, x[5]x[7]11, x[7]61117得到x [3, 0, 7, 4, 4, 11, 1, 17]offset1处理 i1,3,5,7i1: tmp3, x[0]0, x[1]3 i3: tmp7, x[2]4, x[3]11 i5: tmp4, x[4]11, x[5]15 i7: tmp1, x[6]17, x[7]18最终得到x [0, 3, 3, 11, 11, 15, 17, 18]这是exclusive scan每个位置保存的是原数组它之前所有元素的和。如果想要 inclusive scan把原数组 a[i] 逐位加上去即可得到 [3, 4, 11, 11, 15, 17, 18, 24]和 Hillis-Steele 的结果一致。这段手动推演值得多读几遍。我第一次看论文时对 down-sweep 里“x[i - offset] x[i]”这一步很不理解直到按 offset1 那轮仔细算过才明白左孩子接收的是父节点当前的 exclusive 偏移右孩子则拿这个偏移加上自己的子段和。这就是整个算法的灵魂——父节点保存的“左侧所有节点的和”在向下传播时先送给左子树再让右子树站在左子树的肩膀上继续传播。3.3 复杂度对比两个算法放在一起看算法深度总工作量内存访问模式适合场景Hillis-SteeleO(log n)O(n log n)每轮全数组读写规律性强warp 内扫描、小规模数据BlellochO(2 log n)O(n)两轮遍历中间结果留在原处大规模数据、带宽受限设备付出两次遍历的代价换来了 O(n) 的总工作量。在数据规模超过几十万甚至上亿后Blelloch 的内存流量优势是压倒性的。C 标准库的 std::inclusive_scan 多线程实现、CUDA 的 CUB::DeviceScan底层核心思想都是 Blelloch 的变体。4. 真实的 GPU/CPU 上扫描还要再过一道“块间前缀”前面所有讨论都隐含一个假设整个数组能放进一个线程块或一个处理单元。但 n 过亿时GPU 一个 block 最多 1024 线程CPU 也只有几十个核心根本不可能所有线程一起做一次全局扫描。真实现法是分块tiling把大规模扫描拆成三层流水把数组切成长度相等的 tile每个 block/线程独自对自己的 tile 做一次局部扫描把所有 tile 的最后一个元素其实是 tile 内 aggregate提出来做一次规模小得多的块间扫描得到每个 tile 的“前缀偏移量”每个 block 把对应的 tile 前缀偏移量加回自己算好的局部结果上。这就是经典的 “chained scan” 思路。它本质上是用“块间前缀 局部修正”把问题降维第一次扫描是局部的第二次扫描的数据量只有 tile 数第三次只是并行加法。任何并行前缀都不需要真正的全局同步只要让每个 tile 等到它之前所有 tile 的 aggregate 就可以了。现代 GPU 的 CUB 库在此基础上做了更聪明的优化叫decoupled look-back。它不再要求所有 block 同步完成后再做块间扫描而是每个 block 算完自己的局部 aggregate 后立刻写到一个全局数组里然后“回头看”前面 tile 的结果是否已就绪。如果前面 tile 已经完成直接拿它的 prefix 修正自己如果没完成就自旋等待。这显著减少了 block 之间的等待时间GPU 利用率大幅提升。理解了这个再去翻 CUB 源码会顺很多。在 GPU 上写一个可用的 block 内 Blelloch scan 也不难。核心逻辑如下只展示单 block 版本便于理解实际工程请直接用 CUBtemplate typename T __global__ void blockScan(const T* in, T* out, int n, T* blockSums) { __shared__ T s[BLOCK_DIM]; int i blockIdx.x * BLOCK_DIM threadIdx.x; s[threadIdx.x] (i n) ? in[i] : 0; __syncthreads(); // up-sweep for (int offset 1; offset BLOCK_DIM; offset 1) { if ((threadIdx.x 1) % (offset * 2) 0) { s[threadIdx.x] s[threadIdx.x - offset]; } __syncthreads(); } if (threadIdx.x BLOCK_DIM - 1) { blockSums[blockIdx.x] s[threadIdx.x]; s[threadIdx.x] 0; // 根节点置 identity } __syncthreads(); // down-sweep for (int offset BLOCK_DIM / 2; offset 0; offset 1) { if ((threadIdx.x 1) % (offset * 2) 0) { T tmp s[threadIdx.x - offset]; s[threadIdx.x - offset] s[threadIdx.x]; s[threadIdx.x] tmp; } __syncthreads(); } if (i n) out[i] s[threadIdx.x]; // exclusive scan }这段代码的正确性我手动验证过多次但有两个细节必须强调。第一% (offset * 2)的取模运算在真实 GPU 上开销不小工程实现会换成位运算或通过线程索引的二进制位判断这里为了可读性牺牲了性能。第二shared memory 本身有 bank 冲突问题当窗口大小和 bank 数对齐时同一 warp 内多个线程同时访问同一 bank 会串行化CUB 的 WarpScan 会专门处理这一点。自己写代码时如果发现性能比预期差很多先查是不是 bank conflict。CPU 多线程版的套路也是分块扫描加块间前缀。我在本机8 核上跑过 1 亿个 int 的测试单线程std::partial_sum大约 100ms 量级用 OpenMP 分块 串行扫描块间前缀大概能到 20ms 左右。块数不能太多否则块间扫描和线程调度开销会把收益吃光一般取核心数的 2 到 4 倍就够。5. 会了前缀扫描你能解决哪些实际问题前缀扫描不只是一个算法练习题它是很多高性能系统的地基。我挑五个最常见场景每一个都能在真实代码库中找到对应。5.1 基数排序算完桶前缀才知道元素往哪放基数排序每轮根据 key 的某几位把元素分发到不同桶里。单线程做法是维护每个桶的计数器边计数边放元素。并行的麻烦在于多个元素可能同时进同一个桶如果没有位置信息就会冲突。解法是先用并行计数算出每个桶的元素数量然后对这些数量做一次前缀和得到的数组就是每个桶的起始位置。接着每个元素查自己的桶号、用原子操作抢占一个位置即可。GPU 上的基数排序基本都走这条路。5.2 流压缩把满足条件的元素紧凑排列流压缩stream compaction是图形学和计算流体力学里的常见需求一个数组里有很多元素只保留满足条件的那些。朴素做法是顺序遍历、push back但并行时无法预知输出位置。可以并行算一个标记数组标记满足条件的元素为 1然后对这个 0/1 数组做前缀和得到的值就是每个保留元素在输出数组中的目标位置。这一步做完下一步并行复制就是水到渠成的事。5.3 稀疏数据结构CSR 矩阵的行偏移量稀疏矩阵的 CSR 格式需要两个数组行偏移 offset 和列索引 col。最麻烦的一步是构建行偏移每个非零元素属于第 i 行需要知道前 i-1 行一共有多少非零元素。这本质上就是对每行的非零元素计数做前缀和。我做过一个图算法项目要把 SPARSE 图的边数组转成 CSR 表示用并行前缀把这一层的构建时间从几百毫秒降到了几十毫秒是整个预处理管线里收益最大的一处优化。5.4 任务分配按权重把任务分给线程如果一组任务的长度不均匀想让每个线程分到总工作量大致相同的任务可以先把任务长度做前缀和每个线程按“目标工作量”在总前缀序列上做二分查找找到自己对应的起始任务索引。这比简单平均分配要公平得多在多线程渲染和并行物理模拟里很常用。5.5 线性递推斐波那契也能并行算线性递推满足矩阵乘法的结合律。比如斐波那契递推可以写成 2x2 矩阵乘法对矩阵做并行前缀每个元素是矩阵取最后一个矩阵就能得到 F(n)。这个思路看似简单但说明了一个很深刻的事实并行前缀不依赖操作符可交换只依赖可结合。矩阵乘法不可交换但可结合照样可以用 parallel prefix 加速。了解这一点你就知道这个技术的适用面比“数字加法”要广得多。6. 工程落地时的几个提醒与经验理论讲完了实际写代码时还有一堆坑。我按重要性排序分享几个自己踩过的。第一浮点数加法不满足结合律。IEEE 754 浮点数的加法在大多数编译器优化下是不严格结合的并行前缀的求值顺序和顺序扫描不同结果会有微小差异。对常规数值计算这不是问题但如果你在做数值敏感的算法比如概率模型里的对数似然累计要意识到结果可能和旧实现不完全一致必要时切换成更高精度累加或明确的求值顺序。第二小数据规模别用并行扫描。开线程、同步、全局内存往返的开销客观存在。我测试过数据量小于几千时多线程扫描往往比单线程std::partial_sum还慢。判断阈值可以粗略按“数据量 × sizeof(T) / 缓存行大小”来衡量超过几十个缓存行才值得并行。很多工程师一上来就用 fancy 算法跑小数组性能反而更差这不怪算法怪场景不合适。第三标准库能解决 80% 的需求。C17 开始有std::inclusive_scan和std::exclusive_scan配合执行策略可以用多线程GPU 上有thrust::inclusive_scan和cub::DeviceScan::InclusiveScan。如果只是想要一个“正确且不算太慢”的扫描直接用库。自己造轮子之前先问一下是不是有非标准的操作符需要自定义是不是有内存布局的特殊要求如果没有别重复造轮子。第四操作符只要求结合律不要求交换律。很多人默认 scan 是一个和顺序无关的累加操作其实只需结合律就够。矩阵乘法、字符串连接、区间求交这类不可交换操作都可以套用 parallel prefix。设计自己的算子时只要保证结合律成立并行实现就一定正确。这也是为什么 Blelloch 算法里 down-sweep 阶段的 tmp x[i] 顺序写得很讲究不能反过来。第五实现高性能 scan最难的不是算路径而是内存访问模式。我自己写过一版 CUDA scan逻辑和上一节代码一模一样但性能只有 CUB 的一半。后来分析才发现问题出在 block 间通信的方式上我用全局原子操作做块间聚合频繁的原子加导致严重竞争改用 decoupled look-back 思路后块间等待大幅减少。所以如果只是学习自己实现没问题如果要追求极限性能直接读 CUB 的源码那里有所有你能想到的 micro-optimization。最后再分享一个我自己调试并行 scan 的小技巧不要直接拿大数组验证。先写一个 n8 的小测试用纸笔把每一轮状态都算出来再和代码输出逐行比对。然后跑随机数组拿并行结果和std::partial_sum对比找差异。前缀算法的边界条件极多错一个小索引就全盘皆乱这种从最小规模开始的验证方式能帮你把心智负担降到最低。
返回列表