SIMT数据依赖与并行前缀和：从原理到流压缩实战

做 GPU 编程的人几乎都会在某个版本里撞上同一堵墙明明我的 kernel 逻辑很顺可跑出来结果就是不对或者偶尔对、偶尔错。把代码翻来覆去找半天最后发现是数据依赖在作祟。这个问题的本质跟 SIMT单指令多线程的执行模型咬得非常紧。搞清楚“指令流里的数据依赖”到底怎么回事很多奇奇怪怪的 bug 和性能瓶颈其实都能迎刃而解。这篇文章我会从 SIMT 的执行模型讲起把数据依赖的类型、同步手段、以及一个非常实用的“前缀和解决数据依赖”的思路全部串起来。不管你是刚上手 CUDA 的新手还是已经写过几个 kernel 但老是觉得心里没底的进阶用户这篇文章都能给你一点实测下来的经验参照。最后我会用一个流压缩的实战例子具体演示怎么用前缀和把一串串行依赖改造成并行友好的代码。1. 先搞清楚SIMT 到底怎么执行你的代码1.1 锁步执行是依赖问题绕不开的起点SIMT 全称是 Single Instruction Multiple Threads。很多人把它类比成 SIMD单指令多数据但实际上两者差别很大。SIMD 是让一组数据在同一个处理单元里做同一件事而 SIMT 是把线程组织成 warp在硬件上以 warp 为单位执行指令。以 NVIDIA GPU 为例一个 warp 固定是 32 个线程。这 32 个线程执行的是同一条指令但它们各自操作自己的数据而且每个线程都有自己的寄存器、自己的程序计数器状态。关键点在于“锁步执行”。所谓锁步是指在一个 warp 内部32 个线程在同一个时钟周期内执行同一条指令谁也不能领先别人跳过当前这条指令。如果一个 warp 里发生了分支硬件会把执行路径分成多个“活跃掩码”状态分别执行这时候某些线程会处于不活跃状态但它们不会跑远。这个机制带来的直接后果是如果你在 warp 内线程之间做数据交换且指令之间没有插入额外的同步操作那这条交换指令执行时数据可能还没准备好。因为后续指令的执行不保证前一条指令对另一个线程已经完全可见。听起来有点绕但举一个简单的例子就明白了。假设线程 0 写入变量 a线程 1 紧接着读取 a。从逻辑上看线程 1 读到的应该是线程 0 写入的新值但在锁步模型里线程 1 的读取指令可能在线程 0 的写入指令尚未提交到寄存器或内存时就执行了于是读到旧值。1.2 三种数据依赖都要防但最狠的是 RAW数据依赖在计算机体系结构里被分为三类RAWRead After Write写后读、WARWrite After Read读后写、WAWWrite After Write写后写。在 SIMT 环境里这三类问题都会出现但表现和代价完全不同。RAW 是最常见的也是最容易造成结果错误的。一个线程要读某个位置的值而另一个线程要先写这个位置。如果不加同步读方拿到的可能是旧值。在 SIMT 场景里典型的 RAW 就是 warp 内的相邻线程依赖比如第 i 个线程需要第 i-1 个线程的结果。很多并行算法比如扫描、流压缩、某些 BLAS 的 kernel都会有这种依赖一旦处理不当数据错位、结果漂移这些头疼的问题全来了。WAR 在 SIMT 里更多体现在共享内存或全局内存的重复使用上。某个线程先读取一个值之后另一个线程覆盖写入这个值。如果覆盖过早读方会读到新值逻辑就乱了。这种问题经常出现在共享内存缓冲复用的时候尤其是手工做 double buffer 或者池化内存时。WAW 则发生在多个线程往同一位置写数据时最后谁赢取决于执行顺序这种问题基本只能靠原子操作或者重新设计布局来规避。在硬件层面GPU 本身也有一定的调度能力。一个 warp 内部的指令发射是有序的所以对于同一个线程内的指令只要符合程序顺序硬件会保证依赖正确。但跨线程的依赖genuinely是程序员必须显式处理的。记住这句话在 SIMT 里面线程间依赖默认是不成立的你必须在代码里明确地告诉硬件“等一下我要用别人的结果”。2. 处理数据依赖的常规手段同步、延迟与原子2.1 用同步指令做“显式刹车”最简单直接的方式就是同步。在 CUDA 里面块内同步用__syncthreads()warp 内同步用__syncwarp()。__syncthreads()会让 block 内所有线程都到达这个点之后才允许继续往下走。这本质上是一个执行屏障它保证了在屏障之前的内存访问在屏障之后的线程都能看到前提是你用对了内存比如共享内存或者加了 volatile 的全局内存。__syncwarp()则只同步当前 warp 内的线程开销比__syncthreads()小得多。要注意的是如果代码里出现了分支而__syncthreads()放在了分支内部那问题就大了。比如if (threadIdx.x % 2 0) { __syncthreads(); }这个写法在 CUDA 里是非法的因为奇数线程可能不会到达这个屏障而偶数线程卡在那里等一个永远不会来的伙伴整个 block 就死锁了。这个坑我踩过不止一次排查起来真的让人头秃。从性能角度看同步指令也不是免费午餐。__syncthreads()会强制 block 内所有 warp 在屏障点汇合这会破坏指令级并行和 warp 调度的自由度用得过多会让 kernel 明显变慢。所以经验法则是能减少同步次数就尽量减少不要在每个循环里都塞一个。2.2 Warp Shuffle让数据在寄存器之间直接流动同步是最直白的手段但不是效率最高的。在 warp 内部CUDA 提供了一组 shuffle 指令可以让线程直接在寄存器层面交换数据而不经过共享内存也不需要显式的同步。这一套指令包括__shfl_sync、__shfl_up_sync、__shfl_down_sync、__shfl_xor_sync。举个例子假设每个线程都有一个数thread i 想拿 thread i-1 的数你可以这样写unsigned mask 0xffffffff; int pred __shfl_up_sync(mask, value, 1);这里的mask是参与操作的线程掩码必须把真正在跑的线程都标为 1否则会出错。__shfl_up_sync让 lane 编号为 i 的线程拿到 lane i-1 的value越界的 lane 0 会拿到自己原来的值。这个操作完全在寄存器网络里完成不需要走共享内存也不涉及__syncthreads()代价比同步小得多。但要注意shuffle 的同步是 warp 级的而且是隐式的它要求在使用时所有在 mask 中标记的线程都必须执行到这一行。如果你在分支里用了 shuffle那必须保证整个 warp 都进入这个分支或者通过__activemask()之类的函数处理好掩码。这又是一个容易被忽视的坑分支活跃度不一致的时候shuffle 的结果是不可预知的。另外shuffle 只能解决 warp 内的数据交换。如果依赖跨越了 warp比如 thread 33 依赖 thread 32 的结果那还得回到共享内存把数据过渡一下。2.3 原子操作简单但别滥用当多个线程要对同一个地址做读改写时原子操作是最省心的选择。atomicAdd、atomicMin、atomicCAS这些函数在硬件层面保证了操作的不可分割性省去手动加锁的麻烦。但原子操作不是银弹。首先全局原子操作的吞吐量非常有限当大量线程同时访问同一地址时会在硬件层面串行化性能断崖式下跌。其次原子操作只保证了单个操作的原子性并不能保证一个复合逻辑的原子性。比如你要实现“读取计数器的值如果小于 N 就加 1”这种逻辑两个线程可能同时读到同一个旧值然后同时加 1计数就错了。这种场景需要 CASCompare And Swap配合循环自己构造自定义原子逻辑。在 SIMT 的数据依赖场景中原子操作适合处理“多个线程把结果汇总到一个全局位置”这种需求。比如统计一个数组中满足条件的元素个数让每个线程对计数器做atomicAdd就很合适。但如果是一个依赖链比如每个线程的结果要作为下一个线程的输入那原子操作就无能为力了——因为这不是“并发更新同一个地址”而是“串行依赖”的问题。这种时候就需要换一种思路用并行前缀和来把依赖链解开。3. 前缀和把串行依赖变成并行前缀3.1 为什么前缀和能解决数据依赖前缀和Prefix Sum也叫扫描Scan它做的事情是给一个输入数组输出一个数组其中每个元素是输入数组当前位置之前所有元素的和inclusive 版本或者不包括当前位置exclusive 版本。数学表达就是output[i] input[0] input[1] ... input[i]。你可能会问这跟数据依赖有什么关系我举一个实际场景假设有 N 个任务每个任务的复杂程度不同需要分配不同的线程数。你要为每个任务计算它在最终线程数组中的起始偏移。任务 i 的起始偏移是前 i-1 个任务占用的线程数之和这天然就是一个前缀和。如果每个任务占用的线程数还依赖于上一个任务的结果那就形成了一个串行依赖链直接写循环会让所有线程空等性能极差。用前缀和解决问题的核心思想是把这种“一步步推”的串行过程转化为“可并行化计算的累积和”。你可以一次性把所有任务的占用数先算好再用并行前缀和快速得到偏移量。依赖从“结果依赖”变成了“数据依赖”而数据依赖的前缀和本身可以被并行算法加速。3.2 并行前缀和的两个主流算法要让前缀和真正并行起来有两个经典算法Hillis-Steele 和 Blelloch。Hillis-Steele 算法更容易理解。它采用逐步加倍步长的策略每一轮里每个线程把自己前面 offset 距离的元素加到自己头上。用一个循环控制步长从 1 递增到 N/2。每一轮的时间复杂度是 O(log N)但总工作量是 O(N log N)因为每个元素在每一轮都参与了计算。这个算法在数据量小、线程数少的时候挺好用代码简单也不容易出错。但数据量大时它的额外计算量比较浪费。Blelloch 算法则分两个阶段先做一趟自底向上的规约up-sweep也叫 reduce phase再做一趟自顶向下的下推down-sweep也叫 scan phase。整体时间复杂度是 O(N)总工作量是 O(N log N) 的常规版本但实际操作中Blelloch 更适合 GPU因为它每个元素的计算次数更少而且更适合线程块内部的并行规约。代价是实现复杂一点需要处理好每一步的 offset 和活跃线程范围。我自己在工程里通常这样选如果只是在一个 block 内做几百个元素的扫描Hillis-Steele 就够用了代码短debug 容易。如果是跨 block 的大规模扫描那就得走 Blelloch或者用现成的库比如 CUB 的BlockScan、DeviceScan不要自己重复造轮子。3.3 前缀和解决依赖的典型套路两步走在实践中用前缀和解决 SIMT 指令流里的数据依赖有个通用的两步套路。第一步把依赖关系抽象成“每个线程/每个元素贡献一个数值”。比如流压缩stream compaction里每个元素要么输出要么丢弃那“贡献数值”就是 0 或 1表示这个位置是否需要输出一个结果。更复杂的场景可以是一个动态计算出来的任务权重。第二步对这些贡献值做一次 exclusive 前缀和得到每个元素在输出数组里的起始位置。这一步做完之后每个输出位置的写入地址就变成独立的了。之前那种“我必须知道前面有几个元素输出才知道自己写到哪里”的串行依赖现在变成了“先并行算前缀和再并行写入”。依赖链被彻底砍断。这个思路在 GPU 上几乎无处不在。BFS 图的 frontier 扩展、稀疏矩阵压缩、排序里的 rank 计算、甚至一些机器学习预处理里的采样底层都能看到这样一个模式。理解了前缀和这个工具很多看着像串行的算法都能被改造成“前缀和并行写入”的形态。4. 实战用前缀和改进流压缩4.1 场景描述与基线版本流压缩是我平时最常用到前缀和的场景之一。假设有个数组input每个元素有一个有效性标记valid。我们要把所有有效元素按顺序拷贝到输出数组output的前面最后输出元素个数count。最简单的基线版本是这样__global__ void stream_compact_simple(const int* input, const bool* valid, int* output, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) return; int count 0; for (int i 0; i idx; i) { if (valid[i]) count; } // 这里需要把idx前面所有有效元素数算出来串行方式巨慢 if (valid[idx]) { output[count - 1] input[idx]; } }这个版本从逻辑上没错但每个线程都要遍历自己之前的所有元素总复杂度是 O(n^2)。当 n 到了几十万跑一次 kernel 比喝水还慢。更麻烦的是这种写法并没有真正体现 SIMT 的并行性只是把 CPU 上串行的逻辑搬到了 GPU 上。4.2 用前缀和改造的详细步骤改造分三步走第一步每个线程计算自己的 valid 值转成 0 或 1。第二步对整个数组做 exclusive 前缀和得到每个有效元素在输出数组里的偏移。这里的意义是如果valid[idx] true那prefix[idx]就是这个元素在输出数组里的目标位置。第三步所有有效线程把自己的input[idx]写入output[prefix[idx]]同时让最后一个线程把总有效数写到count地址。关键代码可以这样组织__global__ void stream_compact_scan(const int* input, const bool* valid, int* output, int* count, int n) { extern __shared__ int shared[]; int idx blockIdx.x * blockDim.x threadIdx.x; int tid threadIdx.x; shared[tid] (idx n valid[idx]) ? 1 : 0; __syncthreads(); // Hillis-Steele 前缀和就在block内做 for (int offset 1; offset blockDim.x; offset 1) { int v 0; if (tid offset) v shared[tid - offset]; __syncthreads(); shared[tid] v; __syncthreads(); } // 现在shared[tid] 是 inclusive 前缀和我们自己转换成 exclusive int exclusivePrefix (tid 0) ? 0 : shared[tid - 1]; // 注意这里还缺了 block 之间的前缀偏移完整实现需要跨block scan if (idx n valid[idx]) { output[exclusivePrefix] input[idx]; } }上面这段代码为了演示省略了跨 block 的部分。实际应用中如果你只有一个 block或者元素数量在一个 block 范围内这已经能工作。但真实大数据集通常需要多个 block这时就需要先算出每个 block 的有效元素总数再对 block 总数做一次前缀和给每个 block 一个统一的起始偏移再在每个 block 内部做前缀和并加上这个偏移。多 block 的实现我一般会分两个 kernel第一个 kernel 统计每个 block 的有效数同时把 block 内的局部前缀和算好写进全局内存第二个 kernel 对 block 计数做前缀和然后每个 block 在自己的局部偏移上累加 block 起始偏移。这样每一步都是并行的总复杂度 O(n/p p)p 是 block 数性能比 O(n^2) 强太多了。4.3 性能观察与参数调优心得我用这个流压缩 kernel 实测过一组数据n 100 万GTX 3060 上运行。基线版本跑了大约 12 毫秒前缀和版本只有不到 0.2 毫秒快了两个数量级。差距更大的是稳定性基线版本的有效元素分布越靠后单个线程的循环次数越长整个 kernel 的时间波动很大前缀和版本几乎不受数据分布影响时间非常平稳。参数上block 大小我习惯开到 256这个值在大多数 GPU 上都能保证较好的占用率。共享内存大小需要blockDim.x * sizeof(int)注意动态共享内存要加上合适的 size。另外在求exclusivePrefix的时候很多新手会忘了处理 tid 0 的情况直接把 shared[tid-1] 拿来用结果第一个元素越界。这个细节非常容易栽跟头。还有一个优化点如果你确定一个 warp 内的数据规模很小可以用__shfl_up_sync做 warp 内的前缀和再配合少量共享内存归约 warp 之间的部分。这种优化在 block 内元素数量很大时收益明显但代码复杂度上了一个台阶。我建议先跑通简单版本再用 profiler 看瓶颈确认是 scan 阶段之后再考虑进一步微调不要一上来就写最高级的版本。5. 常见问题与排查实录5.1 竞争条件看起来结果是对的但偶尔错这是数据依赖问题最让人头疼的表现。程序跑 1000 次999 次结果正确就有一次不正确或者换了更高端的 GPU 就出错。这种间歇性 bug 绝大多数都是因为线程间依赖没有同步。排查思路有三条第一检查共享内存的读写顺序。每个__syncthreads()前后要确认哪些线程在读写共享数组它们是否所有人都确实执行到了屏障点。最好画一个时间线把每个线程在每个阶段的读取和写入列出来。第二检查 shuffle 的 mask。__shfl_sync这一族指令的第一个参数 mask 一定要是实际参与线程的掩码。如果你在分支里用了动态掩码比如__activemask()在某些情况下返回的掩码可能包含还没到达该指令的线程导致未定义行为。稳妥的办法是用0xffffffff但前提是你能保证整个 warp 都到达。第三检查 volatile 和内存栅栏。如果你在全局内存上做跨 block 的数据传递或者一个 block 等另一个 block 的数据那就不能用__syncthreads()解决了得用__threadfence()配合原子标志位或者依赖不同的 kernel 启动边界来天然保证一致性。5.2 __syncthreads() 里的分支陷阱我在前面已经提过一次但这里还是要专门拿出来说因为它是新手最常见也是最隐蔽的错误。__syncthreads()必须保证 block 内所有线程都能到达这意味着它不能放在任何可能被部分线程跳过的分支里。但这个问题不只是“不要在 if 里放 syncthreads”这么简单。有时候你看代码觉得所有线程都会走同一个分支但实际情况并不一定。比如循环的次数对每个线程不同那在循环里放 syncthreads 就是定时炸弹。最后一行迭代结束后有的线程还在循环里有的已经出来了直接死锁。解决方式有两个一是把同步移到分支外面二是用__syncthreads_and、__syncthreads_or这类的带谓词同步指令。这些指令可以在条件不满足时自动处理线程的阻塞但使用场景有限别指望它能替代所有分支里的同步。5.3 性能优化小技巧减少同步次数比优化同步实现更重要很多人在调 kernel 性能时第一反应是优化同步函数的实现比如用__syncwarp代替__syncthreads。这确实有点效果但更根本的优化思路是减少同步点本身。一个典型做法是把同步从循环内部提到循环外部。有些人在循环里做局部归约时每个迭代都放一个__syncthreads()其实很多归约算法并不需要每轮都同步。比如树状归约中上一层的结果只在下一层被使用那就可以把同步放在每轮迭代的最后而不是每轮迭代的读写之间各放一次。关键是要分析清楚当前这轮迭代里线程 A 会不会在下一轮迭代之前使用线程 B 在这一轮迭代中写入的值。如果不会那中间那个同步就可以省掉。另一个技巧是充分利用 warp 级并行性。如果一个 block 内可以分成多个独立的 warp 组各自处理不同的数据片段那么 warp 之间的依赖就不存在了你甚至不需要跨 warp 同步只要在最后合并时用一次 block 级同步。5.4 前缀和实现中几个容易翻车的细节前缀和看着简单但实现起来有太多细节。首先是 inclusive 和 exclusive 的转换。Hillis-Steele 原生算出的是 inclusive 前缀和输出output[i] input[0] ... input[i]。但流压缩需要的是 exclusive即第 i 个元素之前有多少有效元素。这两者相差一个当前元素的值转换时要注意边界条件。通常我会让每个线程保存自己的 inclusive 值然后计算 exclusive 时再看一眼自己的输入值或者 neighbor 值。其次是对齐问题。共享内存数组的对齐对性能影响很大。CUDA 里共享内存的 bank 是 32 个每个 bank 4 字节。如果你访问共享内存的 stride 不是 1而是 2 或 4就会造成 bank conflict性能瞬间掉一半。前缀和算法在 scan 阶段经常出现 stride 不连续的情况这时候要注意 padding 技巧比如把共享内存数组多申请几个 int然后在奇怪的 offset 上插入占位数据。还有一个很隐蔽的问题block 之间结果的合并。多 block 做 scan 时block 1 的起始偏移完全依赖于 block 0 的总和。如果这个偏移没有算对整个输出就错位。我经常看到有人把 block 内的 scan 算对了但漏了把 block 起始值加进去结果输出数组前半段正常后半段全部偏了一个固定量。这种 bug 特别难排查因为你不一定每次都触发。我的做法是写一个小规模测试用例比如 3 个 block每个 block 4 个线程手工验证一遍 offset再跑真实数据。技巧调试 scan 类 kernel 的时候别一上来就跑大数据。先用一个非常小的 n让每个线程的职责都在你的脑子里过一遍把中间结果 print 或写回全局内存比对确认无误后再放大。省下的时间绝对比你想象得多。写在最后一点自己的体会做 SIMT 编程这几年我最大的感受是数据依赖不是一个“遇到了再解决”的问题而是一个在设计算法时就应该提前规划的问题。每当你看到一个串行的 for 循环第一反应不应该是“GPU 能不能直接跑”而应该是“这个依赖能不能被重新表达成可并行的形态”。前缀和只是其中一把钥匙掌握了它你会发现自己面对很多看似并行的难题都会变得豁然开朗。另外还想多说一句不要迷信网上的各种“花哨”写法先跑通最简单的正确版本再用 profiler 逐段分析。数据依赖相关的 bug 往往神出鬼没最朴素的同步和最小范围的数据交换才是保证正确性的根基。踩过几次坑之后你会发现GPU 编程调的最多的不是性能而是耐心。