Duane Merrill, Michael Garland (NVIDIA Corporation)
本文针对并行计算中的基础原语——前缀扫描(Prefix Scan),提出了一种高效、避免通信的单遍(single-pass)并行算法,名为"解耦回溯"(Decoupled Look-back)。
核心问题与研究目标:
在现代计算机系统中,前缀扫描的性能通常受限于数据移动成本(内存带宽),而非计算能力。传统的顺序扫描算法仅需 $2n$ 的数据移动($n$ 次读,$n$ 次写)。然而,现有的GPU并行扫描策略存在显著缺陷:
1. 多遍算法(如 Reduce-then-scan):通常需要约 $3n$ 的全局数据移动,且包含两次完整的内存遍历,无法用于处理器过载情况下的原位(in-situ)全局分配。
2. 单遍链式扫描(Chained-scan):虽然数据移动量为 $\sim 2n$,但相邻处理器间存在串行前缀依赖,导致信号传播延迟,阻止了内存I/O的完全饱和。
创新点:
1. 解耦回溯策略(Decoupled Look-back):通过执行有界的冗余计算,解除了本地计算与全局前缀传播延迟之间的耦合。这使得算法仅需约 $2n$ 的数据移动即可完成。
2. 高性能实现:该方法已在开源 CUB 库中实现。在现代 GPU 架构上,其吞吐量接近 memcpy(内存拷贝)操作的性能上限。
3. 广泛的适应性:该方法的单遍特性使其能够适应 (1) 原位(in-place)压缩行为,以及 (2) 在处理器过载计算中的原位全局分配。
扫描电路网络模型
并行扫描解决方案的研究历史悠久。扫描电路是快速加法器硬件的基础 [8, 24]。许多扫描并行化方法被表示为电路模型中的递归定义、无环数据流网络 [7, 26]。在此模型中,前缀扫描可被视为每个输出对应的归约树森林。
电路深度与大小:并行扫描网络的最小电路深度为 $\log_2 n$ 步,最小操作符数量为 $n-1$。然而,不存在深度为 $\log_2 n$ 且具备 $O(n)$ 工作效率的网络。Snir 的研究 [25] 指出,对于大小为 $s$、深度为 $d$ 的网络,存在下界 $d + s \ge 2n - 2$。随着电路深度受限,网络大小会迅速增加。
常见的扫描网络结构:
* 串行(链式)扫描 (Fig 1a):具有最大的 $n-1$ 深度,无并发计算,但其最小的 $n-1$ 大小使其适合作为处理器过载时的子组件。
* Kogge-Stone 结构 (Fig 1b):采用递归倍增方法,具有最小深度,但工作复杂度为 $O(n \log_2 n)$。适合 SIMD 架构(如 GPU Warp)。
* Sklansky 结构 (Fig 1c):采用递归的 scan-then-fan 方法,也实现了最小深度 $\log_2 n$,但工作复杂度同样为 $O(n \log_2 n)$。其高扇出特性有助于减少电路尺寸。
* Brent-Kung 结构 (Fig 1d):一种工作高效的策略,深度为 $2 \log_2 n$,大小为 $O(n)$。数据流呈现“沙漏”状,包含上行累积树(Upsweep)和下行传播树(Downsweep)。
* Reduce-then-scan (Fig 1e):为了消除 Brent-Kung 方法在输入超过片上内存时产生的 $\sim n$ 个活跃值的溢出/填充,可以在下行阶段重新计算中间值,代价是 $O(n)$ 的冗余计算,将下行行为从传播(propagation)转换为扫描(scan)。
当代 GPU 扫描策略
GPU 扫描实现通常由两层组织构成:(1) 全局粗粒度策略,(2) 局部细粒度策略(用于线程块内)。
3.1 块级(Block-wide)扫描策略
性能目标:由于目标是内存受限(memory-bound)的内核,块级策略的主要目标是足够高效,以便局部计算开销(线程间通信、同步、扫描操作)能被过载的内存子系统吸收。
混合策略:
* Brent-Kung 混合策略 (Fig 2a):嵌入 Kogge-Stone warp 扫描和传播扇出。CUDPP 中首次演示。
* Raking Reduce-then-scan 策略 (Fig 2b):计算委托给单个 warp,在寄存器中执行串行上行和下行。MatrixScan 中首次演示。
* Sklansky 混合策略 (Fig 2c):结合 Kogge-Stone warp 扫描和 Sklansky 扇出传播。
3.2 全局(Global)扫描策略
历史上 GPU 扫描主要采用三种策略:
Scan-then-propagate:如 CUDPP [10, 23] 和 Thrust [3]。采用高基数 Brent-Kung 数据流。产生 $\sim 4n$ 的全局数据移动。
Reduce-then-scan:如 MatrixScan [12]、B40C [1]、MGPU [2]。如图 3 所示,包含三个内核阶段:(1) 块级归约,(2) 块聚合扫描,(3) 块级下行扫描。全局数据移动为 $\sim 3n$。
Chained-scan:如图 4 所示,单遍方法 [11, 27]。线程块之间存在串行依赖链,需等待前驱的包含前缀(inclusive prefix)可用。全局数据移动为 $\sim 2n$。其性能受限于线程块间的信号传播延迟。例如在 Tesla C2050 上,信号延迟限制了吞吐量远低于内存带宽理论值。虽然增加分区大小可缓解,但受限于片上存储容量。
4.1 运作机制
基本原理:本方法是链式扫描(chained-scan)方法的推广,旨在显著减少前缀传播延迟。核心思想是解耦每个处理器对其直接前驱的单一依赖,代价是进行渐进式的冗余计算。与链式扫描固定的“回溯一个分区”不同,本方法允许处理器检查越来越远的前驱状态。
状态更新机制:当每个分区被处理时,其状态会用“分区聚合值”(aggregate)更新。聚合值仅是分区内元素的归约,可独立计算。由于聚合值易于获取,处理器可以自由地查看前驱分区记录的聚合值,并逐步累积,直到获知完整的“独占前缀”(exclusive prefix)。随后,分区状态会被更新为“包含前缀”(inclusive prefix)。此外,如果在回溯过程中偶然遇到前驱的“包含前缀”,则可以提前终止回溯阶段。
具体执行步骤:每个并行处理器按如下步骤操作:
1. 初始化分区描述符:每个分区的状态描述符包含以下字段:
* aggregate:记录由分区扫描的上行阶段(upsweep)计算的分区聚合值。
* inclusive_prefix:记录分区的包含前缀(由聚合值和从前驱累积的回溯值归约而成)。
* status_flag:描述分区当前状态的标志:
* X - 无效(Invalid):尚无可用信息。所有描述符初始化为 X。
* A - 聚合值可用(Aggregate available):aggregate 字段已记录。
* P - 前缀可用(Prefix available):inclusive_prefix 字段已记录。
aggregate 字段。执行内存栅栏(memory fence)并将 status_flag 更新为 A。拥有第一个分区的处理器将 aggregate 复制到 inclusive_prefix,更新状态为 P,并跳至步骤 6。使用解耦回溯确定分区的独占前缀:处理器初始化并维护一个运行中的 exclusive_prefix,从直接前驱开始,逐步检查更早的分区描述符。对于每个前驱,根据其 status_flag 有条件地执行:
X:阻塞(或继续轮询)直到状态不为 X。A:将该前驱的 aggregate 字段加到 exclusive_prefix 中,并继续检查更早的分区(look-back)。P:将该前驱的 inclusive_prefix 字段加到 exclusive_prefix 中,并终止回溯阶段。计算并记录分区包含前缀:处理器将 exclusive_prefix 加到 aggregate 上,结果记录到 inclusive_prefix 字段。执行内存栅栏并将 status_flag 更新为 P。
exclusive_prefix 整合到每个输出值中。并行性说明:处理器在步骤 1、3、5 和 6 中可以独立并行进行。仅在步骤 4(解耦回溯)中,处理器需要等待其前驱完成步骤 3(记录分区归约)。
4.2 属性
安全性(Safety):如果系统保证所有处理器的前向进展(forward-progress),算法将运行至完成。前向进展保证没有处理器会在步骤 4 无限期等待,因为每个前驱都可以自由地在步骤 3 记录其聚合值。
最小化等待(Minimal waiting):在提供公平或近乎公平调度的系统中,阻塞将是最小的。公平性确保所有处理器在大致相同的时间内记录其聚合值。因此,处理器通常可以自由累积前驱聚合值,极少需要阻塞等待。
常数界回溯(Constant-bound look-back):给定有限数量的处理元素,回溯量是常数。这一属性确保了最优的 $O(n)$ 总体工作复杂度。无论处理器被分配单个分区还是多个分区,每个分区至多检查 $p$ 个前驱即可到达同一处理器先前处理过的分区前缀。
加速信号传播(Accelerated signal propagation):在链式扫描中,共享 inclusive_prefix 只能解除直接后继的阻塞。而在本策略中,记录分区的 inclusive_prefix 能够解除所有后继处理器的阻塞(通过短路回溯)。
非交换算子支持:本方法仅将归约算子应用于连续的输入(或连续的部分归约)。
4.3 实例
图 6 展示了一个由 8 个处理器在 8 个分区上计算前缀和的执行快照(假设每个分区和为 2)。
* processor3 已完成。它使用了 3 个前驱的回溯窗口来确定其 exclusive_prefix,并将其状态设为 P。
* processor7 正在等待其直接前驱变为有效状态。
* 不同的处理器处于不同阶段,利用了短路机会(如遇到状态 P)。
4.4 适配与优化
虚拟处理器(Virtual processors):CUDA 等模型使用虚拟处理器抽象。若运行时调度顺序任意且无抢占,原方法可能死锁(步骤 4 中等待未被调度的前驱)。解决方案是为每个虚拟处理器提供一个标识符(如通过原子递增全局计数器获得),以保证每个被检查的前驱分区都已被活跃调度。
无栅栏描述符更新(Fence-free descriptor updates):步骤 3 和 5 的更新通常需要内存栅栏以防止编译器或硬件重排序。但这会带来性能惩罚。如果架构支持(如 64 位加载/存储),可以将 status_flag 和对应的值(如 32 位整数)合并到一个字中进行单次原子写入。例如 {A, 2} 的 64 位写入可保证一致性,且无需分别维护 aggregate 和 inclusive_prefix 字段(一个 value 字段即可)。
并行化回溯(Parallelized look-back):通过并行检查前驱,可进一步减少可变回溯的延迟。利用 SIMD 线程组(如 warp),一组 $t$ 个线程可同时检查 $t$ 个前驱的分区窗口:
* 线程组轮询直到各自的前驱不再是 X。
* 若所有前驱均为 A:各线程读取聚合值,计算局部归约加到运行中的 exclusive_prefix,窗口回滑 $p$ 个分区继续。
* 若至少一个前驱为 P:计算局部段归约(segmented-reduction),其中状态 P 的前驱作为段头。最后一个段的总和加到 exclusive_prefix,终止回溯。
原位压缩行为(In-place compaction behavior):该单遍策略适合实现具有压缩行为的并行算法(如 select-if, run-length-encode)。
* 原理:基于二进制标志数组的前缀和计算散射偏移量(scatter offset)。
* 原位优势:由于信号机制,每个处理器都能保证所有前驱处理器至少已经读取了它们的输入。这允许处理器将输出写入到压缩后的位置,而没有覆盖前驱尚未读取的输入的风险。这是首个在 GPU 模型上针对此类算法的原位解决方案。
数据集:
硬件配置:
软件配置:
memcpy 操作(作为性能上限)。设备级前缀和吞吐量:
对比结论:CUB 在所有架构和问题规模上的性能均达到或超过其他实现。
memcpy 的性能上限,表明已达到硬件带宽极限。StreamScan 的局限:尽管进行了自动调优,StreamScan 仍受限于串行前缀传播延迟。在 Fermi 和 Kepler 架构上,由于片上资源限制了分块大小,无法掩盖 L2 缓存往返延迟,导致无法匹配 memcpy 性能。
压缩类算法性能:
* 算法:包括 select-if, partition-if, reduce-by-key, run-length-encode。
* 结果 (Fig 12):CUB 利用单遍扫描策略将项选择(item-selection)与前缀扫描融合,无需额外的 I/O 或冗余计算。
* 对比 Thrust:CUB 的实现显著更快。例如,在 select_if, reduce_by_key, partition_if, run_length_encode 操作上,CUB 分别比 Thrust 快 4.1x, 7.1x, 3.5x, 和 3.8x。
本文提出的“解耦回溯”方法是对链式扫描方法的创新推广,显著降低了前缀传播延迟。通过允许并行处理器在必要时执行少量冗余工作来检查更远的前驱状态,该方法能够重叠前缀依赖的传播,从而完全饱和 DRAM 带宽。
主要结论:
1. 性能卓越:这是首个在 GPU 上能够匹配 memcpy 吞吐量的单遍前缀扫描算法。
2. 功能扩展:证明了该策略适用于实现具有原位(in-place)压缩行为的并行算法,无需为压缩输出分配单独存储。
3. 非确定性执行:与构建在静态数据流网络上的传统并行化不同,本方法的动态调度意味着对于伪结合算子(如浮点数加法),扫描结果可能因 CUB 应用的扫描算子数量和顺序的运行间差异而呈现非确定性。
该实现已包含在开源的 CUB 库中。