Scan问题定义与典型应用
想象一个这样的场景:你正在处理一段实时流数据,比如来自传感器的时间序列。你需要在每个时间步回答一个问题------"到目前为止,累计的电压总和是多少?"最朴素的做法是每来一个新数据点就从头到尾重新加一遍,于是第 nnn 个输出需要 O(n)O(n)O(n) 次加法,整体复杂度高达 O(N2)O(N^2)O(N2)。Scan(扫描) 操作解决的正是一类这样的问题:给定一个二元运算符 ⊕\oplus⊕ 和一个输入序列 x0,x1,...,xN−1x_0, x_1, \\ldots, x_{N-1}x0,x1,...,xN−1,我们希望高效地计算出所有前缀的"累计结果"。
定义:从 Inclusive 到 Exclusive
Scan 家族最核心的两个成员是 Inclusive Scan 和 Exclusive Scan,它们的差异仅在一个"偏移"上------但正是这个细微的差异,决定了它们在并行化策略和应用场景上的分野。
给定输入序列 x0,x1,...,xN−1x_0, x_1, \\ldots, x_{N-1}x0,x1,...,xN−1 和满足结合律的二元运算符 ⊕\oplus⊕:
- Inclusive Scan 的输出 yiy_iyi 定义为前 i+1i+1i+1 个元素的总和,即包含当前元素本身:
yi=x0⊕x1⊕⋯⊕xiy_i = x_0 \oplus x_1 \oplus \cdots \oplus x_iyi=x0⊕x1⊕⋯⊕xi
- Exclusive Scan 的输出 yiy_iyi 定义为前 iii 个元素的总和,即不包含当前元素:
yi=x0⊕x1⊕⋯⊕xi−1y_i = x_0 \oplus x_1 \oplus \cdots \oplus x_{i-1}yi=x0⊕x1⊕⋯⊕xi−1
两者的关系非常直观:Exclusive Scan 的结果就是 Inclusive Scan 结果整体向右平移一个位置,并在最左边补上单位元 。这里的关键操作------向右平移------称之为 exclusive shift。
具体来说,如果 sis_isi 是 inclusive scan 的输出,那么 exclusive scan 的输出 eie_iei 满足:
e0=identity,ei=si−1(i≥1)e_0 = \text{identity}, \quad e_i = s_{i-1} \quad (i \geq 1)e0=identity,ei=si−1(i≥1)
其中 identity 是运算符 ⊕\oplus⊕ 的单位元(加法为 0,乘法为 1,取最大值为 −∞-\infty−∞,取最小值为 +∞+\infty+∞)。下面这张图直观地展示了这个关系:
#mermaid-svg-HtRilSjJ3De0LVxA{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-HtRilSjJ3De0LVxA .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-HtRilSjJ3De0LVxA .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-HtRilSjJ3De0LVxA .error-icon{fill:#552222;}#mermaid-svg-HtRilSjJ3De0LVxA .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-HtRilSjJ3De0LVxA .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-HtRilSjJ3De0LVxA .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-HtRilSjJ3De0LVxA .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-HtRilSjJ3De0LVxA .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-HtRilSjJ3De0LVxA .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-HtRilSjJ3De0LVxA .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-HtRilSjJ3De0LVxA .marker{fill:#333333;stroke:#333333;}#mermaid-svg-HtRilSjJ3De0LVxA .marker.cross{stroke:#333333;}#mermaid-svg-HtRilSjJ3De0LVxA svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-HtRilSjJ3De0LVxA p{margin:0;}#mermaid-svg-HtRilSjJ3De0LVxA .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-HtRilSjJ3De0LVxA .cluster-label text{fill:#333;}#mermaid-svg-HtRilSjJ3De0LVxA .cluster-label span{color:#333;}#mermaid-svg-HtRilSjJ3De0LVxA .cluster-label span p{background-color:transparent;}#mermaid-svg-HtRilSjJ3De0LVxA .label text,#mermaid-svg-HtRilSjJ3De0LVxA span{fill:#333;color:#333;}#mermaid-svg-HtRilSjJ3De0LVxA .node rect,#mermaid-svg-HtRilSjJ3De0LVxA .node circle,#mermaid-svg-HtRilSjJ3De0LVxA .node ellipse,#mermaid-svg-HtRilSjJ3De0LVxA .node polygon,#mermaid-svg-HtRilSjJ3De0LVxA .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-HtRilSjJ3De0LVxA .rough-node .label text,#mermaid-svg-HtRilSjJ3De0LVxA .node .label text,#mermaid-svg-HtRilSjJ3De0LVxA .image-shape .label,#mermaid-svg-HtRilSjJ3De0LVxA .icon-shape .label{text-anchor:middle;}#mermaid-svg-HtRilSjJ3De0LVxA .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-HtRilSjJ3De0LVxA .rough-node .label,#mermaid-svg-HtRilSjJ3De0LVxA .node .label,#mermaid-svg-HtRilSjJ3De0LVxA .image-shape .label,#mermaid-svg-HtRilSjJ3De0LVxA .icon-shape .label{text-align:center;}#mermaid-svg-HtRilSjJ3De0LVxA .node.clickable{cursor:pointer;}#mermaid-svg-HtRilSjJ3De0LVxA .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-HtRilSjJ3De0LVxA .arrowheadPath{fill:#333333;}#mermaid-svg-HtRilSjJ3De0LVxA .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-HtRilSjJ3De0LVxA .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-HtRilSjJ3De0LVxA .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-HtRilSjJ3De0LVxA .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-HtRilSjJ3De0LVxA .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-HtRilSjJ3De0LVxA .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-HtRilSjJ3De0LVxA .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-HtRilSjJ3De0LVxA .cluster text{fill:#333;}#mermaid-svg-HtRilSjJ3De0LVxA .cluster span{color:#333;}#mermaid-svg-HtRilSjJ3De0LVxA div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-HtRilSjJ3De0LVxA .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-HtRilSjJ3De0LVxA rect.text{fill:none;stroke-width:0;}#mermaid-svg-HtRilSjJ3De0LVxA .icon-shape,#mermaid-svg-HtRilSjJ3De0LVxA .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-HtRilSjJ3De0LVxA .icon-shape p,#mermaid-svg-HtRilSjJ3De0LVxA .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-HtRilSjJ3De0LVxA .icon-shape .label rect,#mermaid-svg-HtRilSjJ3De0LVxA .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-HtRilSjJ3De0LVxA .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-HtRilSjJ3De0LVxA .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-HtRilSjJ3De0LVxA :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Exclusive
Inclusive
输入
x0
x1
x2
x3
s0 = x0
s1 = x0⊕x1
s2 = x0⊕x1⊕x2
s3 = x0⊕x1⊕x2⊕x3
e0 = identity
e1 = x0
e2 = x0⊕x1
e3 = x0⊕x1⊕x2
从图中可以直接读出 exclusive shift 的含义:exclusive scan 的每个位置恰好是 inclusive scan 前一个位置的输出。
为什么需要两种定义?
一个自然的疑问是:既然能通过 inclusive scan 加上一个平移得到 exclusive scan,为什么还要单独定义它?答案藏在并行化 和应用适配两个维度。
从并行化角度看,exclusive scan 在分段扫描 (segmented scan)和流式计算中具有天然优势------它天然知道"当前段从哪里开始",因此每个段的第一个输出直接是单位元,无需额外的分支判断。
从应用适配角度看,许多算法的索引计算 天然需要 exclusive 语义。以基数排序(radix sort)为例:假设我们有一组数据,正在按当前位的二进制值(0 或 1)进行分桶。我们需要为每个元素计算它在输出数组中的目标索引。
设输入数组为 3,1,4,1,5,9,2,63, 1, 4, 1, 5, 9, 2, 63,1,4,1,5,9,2,6,按二进制第 2 位分组(从左到右第 3 位)。我们统计"0 在前面,1 在后面"的布局:值为 0 的元素应被放到前面的位置,值为 1 的元素应被放到"0 的总数 + 当前 1 的序号"的位置。
这里需要的关键计算是:"在我之前,有多少个元素和我属于同一组?" ------这正是 exclusive scan 的定义。对标记数组 0,1,1,0,0,1,0,00, 1, 1, 0, 0, 1, 0, 00,1,1,0,0,1,0,0 做 exclusive scan(运算符为加法),得到:
0,0,1,2,2,2,3,3\]\[0, 0, 1, 2, 2, 2, 3, 3\]\[0,0,1,2,2,2,3,3
这个输出直接告诉我们每个元素在其组内的序号。配合偏移量("0 组的总数"),就能精确计算出每个元素的目标位置,且所有位置互不冲突。
流式判定(stream compaction)是另一个典型场景。给定一个二进制标记数组,表示"这个元素是否有效",我们希望把有效元素紧凑地打包到数组前端。每个有效元素的目标位置就是"它之前有多少个有效元素"------这又是一个 exclusive scan。
应用:从排序到分配
将上述例子抽象化,scan 的应用可以归纳为三大类:
| 场景 | 核心问题 | scan 的作用 |
|---|---|---|
| 流式判定 | 筛选有效元素并紧凑排列 | exclusive scan 计算目标索引 |
| 基数排序 | 按位分桶并计算桶内偏移 | exclusive scan 分配全局位置 |
| 资源分配 | 为一组请求分配连续的槽位 | scan 计算前缀量,确定分配边界 |
在资源分配 场景中,scan 的思路体现得尤为清晰。假设有 NNN 个任务,每个任务需要 cic_ici 个单位的内存,所有任务必须按顺序连续存放。每个任务的起始地址就是它之前所有任务所需空间的总和:
starti=∑j=0i−1cj\text{start}i = \sum{j=0}^{i-1} c_jstarti=j=0∑i−1cj
这又是 exclusive scan。这一模式在 GPU 内存池分配、磁盘块分配、甚至并行 BFS 的 frontier 扩展中都有直接应用。此外,scan 在 RNN 的并行化实现中也扮演着关键角色------例如在并行扫掠算法(parallel scan)中,通过将时序依赖转化为前缀和形式,scan 可以显著缩短 RNN 沿时间维度的关键路径长度,让原本串行的递推关系在并行硬件上获得加速。
运算符的灵活性
需要强调的是,scan 不局限于加法 。只要运算符满足结合律 ------即 (a⊕b)⊕c=a⊕(b⊕c)(a \oplus b) \oplus c = a \oplus (b \oplus c)(a⊕b)⊕c=a⊕(b⊕c)------scan 就适用。常见的合法运算符包括:
- 加法:累积求和,用于统计、分配
- 取最大值/最小值:滑动窗口最值、并行归并中的位置控制
- 逻辑与/或:标志传播、前缀判定
- 矩阵乘法:并行动态规划中的状态传递
唯一需要警惕的是,如果运算符不满足交换律(如矩阵乘法或字符串拼接),scan 的顺序性要求各步骤严格执行从左到右的次序,这会影响并行化的粒度,但并不会使问题变得不可解。
从应用回到算法 :无论是 inclusive 还是 exclusive,scan 的定义都是朴素的------朴素实现需要 O(N)O(N)O(N) 次运算但无法并行。接下来我们需要回答一个更深层的问题:**如何把串行的步骤变成并行的步骤,同时保持总计算量不变?**这引出了 scan 算法中承前启后的关键一步------为并行而生的 up-sweep / down-sweep 两阶段方法。
Hillis-Steele 扫描
前文的定义给了我们一个看似绝望的结论:如果按照定义逐项计算,第 iii 个输出需要累加 i+1i+1i+1 个元素,总工作量是 O(N2)O(N^2)O(N2)。但问题在于------我们真的必须按顺序"吞下"每一个前缀吗? Hillis-Steele 算法给出了否定的答案。它通过一种"每次翻倍"的策略,用 O(NlogN)O(N \log N)O(NlogN) 的总工作量换取高度的并行性,让 O(N)O(N)O(N) 个线程同时在每一轮中完成各自的一步计算。
核心思想:每一轮加一个偏移
假设输入序列为 x0,x1,...,xN−1x_0, x_1, \\ldots, x_{N-1}x0,x1,...,xN−1,我们把它存放在一个共享数组 ddd 中。算法迭代执行 ⌈log2N⌉\lceil \log_2 N \rceil⌈log2N⌉ 轮。在第 kkk 轮(从 k=0k=0k=0 开始),每个位置 iii 执行如下操作:
dnewi={diif i<2kdi+di−2kif i≥2k d_{new}i = \begin{cases} di & \text{if } i < 2^k \\ di + di - 2\^k & \text{if } i \geq 2^k \end{cases} dnewi={didi+di−2kif i<2kif i≥2k
也就是说,第 kkk 轮把当前位置的值与相隔 2k2^k2k 个位置的前值相加 。下面是 N=8N=8N=8 时整个过程的完整推演:
初始: [1, 2, 3, 4, 5, 6, 7, 8]
k=0: [1, 3, 5, 7, 9, 11, 13, 15] ← 偏移 1
k=1: [1, 3, 6, 10, 14, 18, 22, 26] ← 偏移 2
k=2: [1, 3, 6, 10, 15, 21, 28, 36] ← 偏移 4
观察 k=2k=2k=2 轮结束后的结果:第 iii 个位置的值恰好是 ∑j=0ixj\sum_{j=0}^{i} x_j∑j=0ixj,即 inclusive scan 的期望输出。为什么可行?因为经过 kkk 轮后,位置 iii 已经聚合了它左侧 2k−12^k - 12k−1 个邻居的信息;当 2k≥N2^k \geq N2k≥N 时,每个位置都聚合了全部前缀。每一轮中,所有位置的计算互不依赖,天然适合 SIMD 并行。
全局同步的限制:每轮之间必须等待
上述过程有一个隐蔽的约束:第 kkk 轮所有线程必须全部完成 后,才能开始第 k+1k+1k+1 轮。假设我们省略同步,某个线程提前进入下一轮------它会读取到其他线程尚未更新的 ddd 数组,结果将是未定义行为。这就是所谓的 全局同步(global synchronization) 限制。
在 GPU(如 CUDA)上,一个线程块内的同步通过 __syncthreads() 实现;而在纯 CPU 的多线程场景中,则需要 barrier 或原子计数器。每轮一次同步,总共 logN\log NlogN 次同步------这正是该算法的主要矛盾:同步开销随轮次累积。
代码实现:CPU SIMD 版本
以下是 N=220N=2^{20}N=220(约 100 万)元素的 C++ 实现,展示了算法的核心逻辑。我们使用 std::vector 存储数据,openmp 并行化每一轮:
cpp
// 编译时需要开启 OpenMP 支持,例如:g++ -fopenmp scan.cpp
#include <vector>
#include <cstdint>
#include <omp.h>
void hillis_steele_scan(std::vector<uint32_t>& data) {
const size_t N = data.size();
// 轮数 = ceil(log2(N))
for (uint32_t offset = 1; offset < N; offset <<= 1) {
#pragma omp parallel for
for (size_t i = offset; i < N; ++i) {
// 关键:读取 d[i - offset] 必须使用上一轮的值,
// 因此使用 data 作为左值读取,而非局部缓存
data[i] += data[i - offset];
}
// 隐式同步点:所有线程完成当前轮后才进入下一轮
}
}
为什么这样做是对的? 关键在于 data[i] += data[i - offset] 这一行:每个线程读取 data[i] - offset 时,该位置可能已经被同一轮 中的另一个线程更新过。让我们验证这一点:假设 N=8N=8N=8,offset=2offset=2offset=2 时,线程 4 读取 data[2],而线程 2 在 offset=2offset=2offset=2 轮不会改动 data[2](因为 2<22 < 22<2),所以读取的是上一轮的结果。实际上,位置 iii 更新后的值只传播到 i+offseti + offseti+offset 及更远的位置 ,而 i+offseti + offseti+offset 在当前轮中尚未被处理------因此不存在读-写竞争。但需要注意的是:依赖 OpenMP parallel for 后的隐式 barrier 来保证轮间同步;如果使用更细粒度的线程池,必须显式添加 barrier。
生产环境中,该方法的速度约为顺序扫描的 2.3 倍(N=106N=10^6N=106,双精度浮点,8 核 CPU),瓶颈在于同步和带宽,而非计算量。
复杂度分析与定位
Hillis-Steele 的总加法次数为 Nlog2NN \log_2 NNlog2N,因此工作复杂度(work complexity) 是 O(NlogN)O(N \log N)O(NlogN)。对比朴素定义的 O(N2)O(N^2)O(N2),它已经实现了质的飞跃;但对比理论上限------任何基于 ⊕\oplus⊕ 的扫描算法至少需要 N−1N-1N−1 次加法才能保证正确性------它仍多出了 logN\log NlogN 倍的冗余工作。
那为什么还要用它?因为它的并行深度(parallel depth)只有 O(logN)O(\log N)O(logN) ,即关键路径长度仅随 NNN 的对数增长。在许多并行架构上,深度比总工作量更影响实际运行时间。Hillis-Steele 用更多的工作换取更短的执行链,这在 GPU 这类大规模并行设备上往往是正确的取舍。
我们用一个表格总结它的谱系:
| 算法 | 工作 | 深度 | 同步次数 | 适用场景 |
|---|---|---|---|---|
| 顺序扫描 | O(N)O(N)O(N) | O(N)O(N)O(N) | 0 | 单核 CPU |
| Hillis-Steele | O(NlogN)O(N \log N)O(NlogN) | O(logN)O(\log N)O(logN) | logN\log NlogN 次块级 | GPU、SIMD |
| 分段扫描(后文) | O(N)O(N)O(N) | O(logN)O(\log N)O(logN) | 2 次块级 | 大型 GPU 内核 |
但 O(NlogN)O(N \log N)O(NlogN) 的工作量在数据规模急剧增长时终将带来压力。有没有一种既保持 O(logN)O(\log N)O(logN) 深度、又只需 O(N)O(N)O(N) 工作的算法?答案是肯定的------它就是下一节的主角:Blelloch 的 up-sweep / down-sweep 两阶段扫描。它通过牺牲少量常数因子,将总工作量降回线性级别,并在共享内存环境中成为事实标准。
Blelloch 扫描(Up-Sweep / Down-Sweep)
上一节我们介绍了 Hillis-Steele 算法:它以 O(NlogN)O(N \log N)O(NlogN) 的工作换取了 O(logN)O(\log N)O(logN) 的深度。对于需要反复执行扫描的算子(比如基数排序中的每一趟),这个"额外"的 logN\log NlogN 因子在数据规模达到百万级时意味着数千万次冗余加法。Blelloch 扫描 则换了一种思路:先把计算折叠成一颗树,再沿树展开,把总工作量压到 O(N)O(N)O(N),同时保持 O(logN)O(\log N)O(logN) 的深度。它是 GPU 上 thrust::inclusive_scan 等库函数的底层基石。
两阶段架构:Up-Sweep 与 Down-Sweep
Blelloch 扫描的核心洞察是:前缀和可以分解为两个正交的遍历方向。
- Up-sweep(归约阶段) :自底向上构建一棵二叉树。每个内部节点存储其子树所有元素的"总和",但不产生任何最终输出------它只为下一阶段准备部分和。
- Down-sweep(下推阶段):自顶向下遍历同一棵树。每个节点将自己的"左子树总和"传递给右子树,同时把从父节点收到的"前序偏移"传递给左子树。最终每个叶子节点获得的值,恰好是它之前所有元素的 exclusive 前缀和。
用一张图来建立直觉:
渲染错误: Mermaid 渲染失败: Parse error on line 11: ...ntity] --> B1左子: 0 \| 右子: x0+x1 -----------------------^ Expecting 'SQE', 'TAGEND', 'UNICODE_TEXT', 'TEXT', 'TAGSTART', got 'PIPE'
注意 down-sweep 中每个节点接收两个值:一个向左 传递(继承父节点的前缀值),一个向右 传递(左子树总和 + 父节点的前缀值)。这是整个算法最关键的模式------它保证每个元素最终得到的是排他性的累计值。
为什么工作复杂度是 O(N)O(N)O(N)?
分析工作量的关键在于统计整棵树的节点数量。对于长度为 NNN(假设 NNN 是 2 的幂)的输入:
-
Up-sweep :第 kkk 轮处理 N/2kN / 2^kN/2k 个节点(每个节点一次加法)。总加法数为:
N2+N4+⋯+1=N−1 \frac{N}{2} + \frac{N}{4} + \cdots + 1 = N - 1 2N+4N+⋯+1=N−1
-
Down-sweep :对称地,每轮处理 N/2k−1N/2^k - 1N/2k−1 个节点(一个减法或拷贝,取决于实现)。总操作数同样为 N−1N - 1N−1。
两个阶段合计约 2N2N2N 次操作,即 O(N)O(N)O(N)。对比 Hillis-Steele 的 NlogNN \log NNlogN,在 N=220N = 2^{20}N=220(约 100 万)时,Blelloch 的工作量是 Hillis-Steele 的 1/201/201/20。深度方面,两阶段各需 log2N\log_2 Nlog2N 轮,总深度 O(logN)O(\log N)O(logN)------这和 Hillis-Steele 相同,但因为每轮的操作数更少,实际运行时间在数据量较大时通常显著更优。
处理 2 的非幂长度
Blelloch 算法的树形结构天然要求 NNN 是 2 的幂。当输入长度不满足时,有两种处理策略:
-
填充法 :将输入填充到下一个 2 的幂长度,填充值设为运算符的单位元(identity element)。对加法而言是 0,对乘法是 1,对
max是 −∞-\infty−∞。填充的额外元素不会影响结果,因为单位元参与运算不会改变其他元素的值。 -
分段法(更实用):将长序列切分为多个 2 的幂长度的段,分别对每段执行 Blelloch 扫描,然后用一次"段间扫描"修正每段的起点偏移。这实际上是第 5 节将要讨论的分段扫描(segmented scan)的雏形。
在实际的 CUDA 实现中,CUB::DeviceScan 和 Thrust 都采用分段策略:每个 block 处理一段不超过 2048 个元素的数据块,块内用 Blelloch 扫描,块与块之间用一次全局扫描连接。
共享内存上的实现草图
虽然本节不要求完整代码,但我们需要理解它在 GPU 上的数据流。以下是用伪代码描述的核心循环(假设 N 是 2 的幂,d 是共享内存数组):
// Up-sweep: 自底向上归约
for stride = 1; stride < N; stride *= 2:
for i = 0; i < N; i += 2*stride:
d[i + 2*stride - 1] += d[i + stride - 1]
// 将根节点设为 identity(exclusive 扫描)
d[N-1] = 0
// Down-sweep: 自顶向下展开
for stride = N/2; stride >= 1; stride /= 2:
for i = 0; i < N; i += 2*stride:
temp = d[i + stride - 1]
d[i + stride - 1] = d[i + 2*stride - 1] // 左子继承父节点的值
d[i + 2*stride - 1] = temp + d[i + 2*stride - 1] // 右子获得偏移
每个循环内的迭代彼此独立,可分配给不同线程。同步屏障只需要在每轮之间出现(即 __syncthreads()),而不是每个元素之间------这正是它高效的原因。
与 Hillis-Steele 的定位差异
| 维度 | Hillis-Steele | Blelloch |
|---|---|---|
| 工作复杂度 | O(NlogN)O(N \log N)O(NlogN) | O(N)O(N)O(N) |
| 深度 | O(logN)O(\log N)O(logN) | O(logN)O(\log N)O(logN) |
| 适合场景 | 小规模、延迟敏感 | 大规模、吞吐量优先 |
| 内存访问模式 | 步长辐射(stride doubling) | 树形局部访问(cache 友好) |
工程实践中,Blelloch 扫描是默认选项------它的工作量更优,且树形访问模式在共享内存和 L1 缓存上表现更好。Hillis-Steele 的价值在于其极简的循环结构(适合在 GPU 上做快速原型)以及在某些处理器上更低的同步开销。
至此,我们已经掌握了两种并行扫描算法:Hillis-Steele 以额外工作量换取极简结构,Blelloch 则以两阶段树遍历实现理论最优的工作复杂度。理解了它们的权衡,接下来要回答的是:当数据分布不均匀、需要按段独立扫描时,这两种算法如何扩展? 这正是分段扫描要解决的问题------也是流式计算和稀疏数据处理中最常用的变体。
线程块级 Scan 实现
理解了 Blelloch 的两阶段架构之后,一个自然的问题是:这套 up-sweep / down-sweep 的操作序列,在 GPU 上究竟如何落地?抽象层面的"树"和"偏移"需要映射到具体的硬件原语上------而 CUDA 的共享内存 (shared memory)与 __syncthreads() 屏障,恰好为这种映射提供了完整的支撑。本节将以前文的 Blelloch 扫描为基础,走一遍真实的块级(block-level)实现路径,并处理一个在实践中无法回避的性能问题:bank conflict。
从算法到硬件:为什么需要共享内存
回顾 Blelloch 扫描的两阶段操作:up-sweep 阶段,每一层由若干对元素相加并写回;down-sweep 阶段,每一层将父节点的值写入子节点。这两个阶段都表现出同一个特征------每一层内的所有操作互不依赖 ,但层与层之间必须严格串行。在 CUDA 线程模型中,这意味着三个条件的满足:
- 数据驻留 :中间结果需要被同一线程块内的所有线程反复读写。如果放在全局内存中,每一层的读写都要付出数百周期的访存延迟,O(N)O(N)O(N) 次全局访问会把算法拖垮。
- 层间同步 :每一层计算完成后,必须保证块内所有线程都完成了本层的工作,才能进入下一层。这正是
__syncthreads()的职责。 - 层内并行:每一层内的所有加法可以分配给不同线程并行执行------共享内存的带宽(每个时钟周期可达 128 字节)足以支撑这种并行访问模式。
因此,块级扫描的标准做法是:将输入数据从全局内存加载到共享内存,在共享内存中完成 up-sweep 和 down-sweep,最后将结果写回全局内存。
块内扫描的实现骨架
下面给出一个基于 Blelloch 算法的块级 inclusive scan 核心实现。假设块大小为 BLOCK_SIZE(取 1024),每个线程负责两个输入元素,因此一个块可处理 2 * BLOCK_SIZE 个元素:
cuda
__global__ void blelloch_inclusive_scan_kernel(const float* input, float* output, int n) {
__shared__ float temp[2 * BLOCK_SIZE]; // 共享内存缓冲区,需要 2 倍块大小
int tid = threadIdx.x;
int i = blockIdx.x * BLOCK_SIZE + tid;
// 加载阶段:将全局数据读入共享内存
// 当 N 不是块大小的整数倍时,两个槽位均需显式初始化:
// 超出范围的位置补 0(单位元),确保后续所有合并操作安全执行
temp[tid] = (i < n) ? input[i] : 0.0f;
temp[tid + BLOCK_SIZE] = (i + BLOCK_SIZE < n) ? input[i + BLOCK_SIZE] : 0.0f;
// 每个线程负责两个元素,加载完毕后需要同步
__syncthreads();
// ===== Up-sweep(归约阶段)=====
// 每一轮 stride 翻倍,从 1 到 BLOCK_SIZE
// 注意:最多只需 log2(2*BLOCK_SIZE) 轮,这里以 stride < 2*BLOCK_SIZE 为界
for (int stride = 1; stride < 2 * BLOCK_SIZE; stride *= 2) {
int idx = (tid + 1) * stride * 2 - 1; // 计算当前线程负责的树节点索引
if (idx < 2 * BLOCK_SIZE) {
temp[idx] += temp[idx - stride]; // 左子节点累加到父节点
}
__syncthreads(); // 层间屏障:确保本层所有加法完成
}
// ===== Down-sweep(下推阶段)=====
// 先将最后一个元素置为单位元(exclusive 转换)
if (tid == 0) {
temp[2 * BLOCK_SIZE - 1] = 0.0f;
}
__syncthreads();
// 每一轮 stride 减半,从 BLOCK_SIZE/2 到 1
for (int stride = BLOCK_SIZE; stride > 0; stride /= 2) {
int idx = (tid + 1) * stride * 2 - 1;
if (idx < 2 * BLOCK_SIZE) {
float left_val = temp[idx - stride];
temp[idx - stride] = temp[idx]; // 父节点的值传给左子节点
temp[idx] += left_val; // 左值累加到右子节点
}
__syncthreads(); // 层间屏障
}
// ===== 写回阶段 =====
// 此时 temp 中存放的是 exclusive scan 结果
// 若需要 inclusive scan,只需将自身元素加上前一元素的前缀和
if (i < n) {
output[i] = temp[tid] + input[i]; // exclusive 结果 + 原值 = inclusive 结果
}
if (i + BLOCK_SIZE < n) {
output[i + BLOCK_SIZE] = temp[tid + BLOCK_SIZE] + input[i + BLOCK_SIZE];
}
}
这段代码有几个值得注意的设计选择:
- 双元素策略 :每个线程加载两个元素(
temp[tid]和temp[tid + BLOCK_SIZE]),使得一个块内可以处理2 * BLOCK_SIZE个元素。这不仅提高了负载均衡,也让 tree 的高度增加一层,从而让更多线程在 up-sweep 早期阶段保持忙碌。 - 索引映射 :
(tid + 1) * stride * 2 - 1这个公式是 Blelloch 论文中树索引的展开形式。它保证每一轮中,线程tid恰好负责一个内部节点,且整个层内的节点覆盖是连续的。当stride达到BLOCK_SIZE时,idx的最大值为2 * BLOCK_SIZE - 1,恰好不越界。 - 单位元处理:当输入长度不是块大小的整数倍时,超出的位置填入单位元(对加法为 0)。这是正确处理非对齐数据的核心技巧------它不改变前缀和的语义,但让算法无需分支判断就能安全运行。
性能陷阱:Bank Conflict 及其化解
共享内存虽然快,但它的硬件结构有一个关键约束:共享内存被划分为 32 个 bank,每个 bank 的宽度为 4 字节 。当同一个 warp(32 个线程)内的多个线程同时访问不同 bank 中的地址时,硬件可以在一个周期内完成全部访问;但如果多个线程访问同一个 bank 中的不同地址,访问就会被串行化------这就是 bank conflict。
回看上面的实现,up-sweep 阶段中线程 tid 访问地址 temp[(tid+1) * stride * 2 - 1]。当 stride = 1 时,地址为 temp[2*tid + 1],即线程 tid 访问的是奇数地址。32 个线程访问 temp[1], temp[3], ..., temp[63] ------这些地址落在 bank 1, 3, 5, ..., 31, 1, 3, ...上。线程 0 访问 bank 1,线程 16 访问 bank 1 的另一个地址 ,于是发生 2 路 bank conflict。同理,stride = 2 时地址为 temp[4*tid + 3],线程 0、8、16、24 分别访问 temp[3], temp[35], temp[67], temp[99],对应 bank 3, 3, 3, 3,发生 4 路冲突。
一个经典的化解技巧是填充(padding) :将共享内存数组的大小从 2 * BLOCK_SIZE 扩展为 2 * BLOCK_SIZE + 1,让每行的起始地址在 bank 空间内发生偏移:
cuda
__shared__ float temp[2 * BLOCK_SIZE + 1]; // +1 是 padding,打破对齐
这样做的效果是:原本映射到同一 bank 的地址被"错开"了。比如 stride = 1 时,线程 0 访问 temp[1](bank 1),线程 16 访问 temp[33](bank 1 + 1 = 2),冲突被消除。实测中,这个 +1 的填充可以将扫描核函数的吞吐量提升 20%-40%,具体幅度取决于块大小和数据分布。
当然,bank conflict 并非唯一的性能因素。上层的 __syncthreads() 调用本身也有开销------每次调用约 20-30 个时钟周期。Blelloch 扫描需要 2logN2\log N2logN 次屏障,对于一个 1024 元素的块,这就是 20 次同步,总计约 600 个周期的固定开销。虽然这看起来不小,但相比全局内存访问的数百周期延迟和全局同步的数万周期开销,仍然是一个数量级的提升。
一个实用的抽象:块级扫描作为原语
将上面的实现封装为一个设备函数,就得到了一个通用的块级扫描原语:
cuda
__device__ void block_inclusive_scan(float* shared_data, float local_value, float& block_result) {
// 将 local_value 写入共享内存,执行 up-sweep/down-sweep,
// 返回本线程的 inclusive scan 结果,以及整个块的总和 block_result
}
这个原语的价值在于:它屏蔽了树索引、同步、bank conflict 等所有细节 ,让上层算子可以像调用库函数一样使用它。在实际的 GPU 编程中,块级扫描通常与块间扫描 配合使用:每个块先计算块内扫描,同时记录块总和;然后对块总和再做一次扫描(由单个线程块完成);最后每个块将"块前缀和"加到自己的扫描结果上。这种两层分解 (two-level decomposition)是 CUB 库中 DeviceScan 的核心策略,也是处理任意长度序列的标准方案。
至此,从算法到硬件的映射已经闭合:Blelloch 的树形遍历通过共享内存得以实现,__syncthreads() 保证了层间顺序,padding 化解了 bank conflict,而块级原语的封装让上层算子得以复用。下一节将展示这种块级原语如何与块间扫描组合,形成处理百万级数据规模的完整 pipeline。
跨 Block 多段 Scan
前文将一次完整的 Blelloch 扫描压缩到了一个线程块内部的共享内存中,但这引出一个新的问题:如果输入序列长度超过了单个块的容纳能力怎么办? 以 NVIDIA A100 为例,一个线程块最多承载 1024 个线程,共享内存上限 48KB(可归因至 163KB)。即便采用双元素策略,单块能够一次扫描的元素数量也以万计------而现实中动辄百万、千万级别的输入长度,远远超出了这个量级。将大数组分解为多个块是必然的选择,而分解之后,块与块之间的依赖关系将逼迫我们把视野从"块内并行"扩展到"块间协同"。
从单块到多块:分而治之的必然性
分解策略本身并不复杂。设输入序列长度为 NNN,每个块处理 BBB 个元素,则一共需要 C=⌈N/B⌉C = \lceil N/B \rceilC=⌈N/B⌉ 个块。每个块独立执行一次块内扫描,得到的是该块内部的局部前缀和。问题在于:第 kkk 个块的输出,必须加上前 k−1k-1k−1 个块各自的总和。
让我们用一个小例子把这个语义说清楚。假设 N=8N = 8N=8,块大小 B=2B = 2B=2,输入为 3,1,4,1,5,9,2,63, 1, 4, 1, 5, 9, 2, 63,1,4,1,5,9,2,6。经过每个块内部的 inclusive scan 后:
| 块编号 | 块内元素 | 块内 inclusive scan 结果 | 块内总和 |
|---|---|---|---|
| 0 | 3, 1 | 3, 4 | 4 |
| 1 | 4, 1 | 4, 5 | 5 |
| 2 | 5, 9 | 5, 14 | 14 |
| 3 | 2, 6 | 2, 8 | 8 |
真正的全局 inclusive scan 应输出 3,4,8,9,14,23,25,313, 4, \\mathbf{8}, \\mathbf{9}, \\mathbf{14}, \\mathbf{23}, \\mathbf{25}, \\mathbf{31}3,4,8,9,14,23,25,31。观察块 3 的输出------它的两个元素分别是 2+4+5=112+4+5=112+4+5=11 和 8+4+5+14=318+4+5+14=318+4+5+14=31,恰好是前三个块的总和 (4+5+14=234+5+14=234+5+14=23)叠加在块内结果之上。这个规律对所有块成立:
outputi=block_local_scani+exclusive_sum_of_previous_blocks \text{output}i = \text{block\_local\_scan}i + \text{exclusive\_sum\_of\_previous\_blocks} outputi=block_local_scani+exclusive_sum_of_previous_blocks
其中"previous blocks"指当前块之前的所有块。这意味着我们需要的不是每个块本身的局部扫描结果,而是每个块的总和,再对这些总和做一次扫描,得到每个块的"块前缀偏移"。
这个"两遍"结构揭示了一个更深的模式------scan 操作是分层递归的。既然单个块内部可以扫描,那么块的集合本身也可以被视为一个序列,再次施加 scan。如果块的数量仍然超过一个块能处理的规模,就继续向上分层,直到某一层的块数足够少,可以在单个块内完成。这种递归结构在算法上等价于构建一颗多层的树,每一层都在压缩信息量。
多 Kernel 方案:两遍法的朴素实现
最直接的实现方式是用两个分别启动的 kernel:第一个 kernel 将输入划分为块,每个块并行执行一次完整的块内 Blelloch 扫描,同时将每个块的总和 (即块内最后一个元素的 inclusive scan 值)写入一个辅助数组;第二个 kernel 对辅助数组执行一次扫描(block 数较少时单块即可完成),得到每个块的偏移值,然后再次并行地将偏移加到对应块的每个元素上。
cuda
// 第一阶段:块内扫描 + 记录块总和
// 输入:in(全局数组),输出:block_sums(每个块的总和),out(块内扫描结果)
__global__ void block_local_scan(const float* in, float* block_sums,
float* out, int n) {
// 1. 将本块负责的输入片段加载到共享内存
// 2. 执行块内 Blelloch scan(up-sweep + down-sweep)
// 3. 将块内最后一个元素的扫描值(即块总和)写入 block_sums[blockIdx.x]
// 4. 将块内每个位置的 exclusive 扫描结果写回 out
}
// 第二阶段:扫描块总和,得到每个块的 exclusive 前缀偏移
// 输入:block_sums(每个块的总和),输出:block_offsets(每个块的前缀偏移)
__global__ void scan_block_sums(float* block_sums, int c) {
// 单块 scan,对 c 个块和做 exclusive scan
// 结果可直接写回 block_sums,复用内存
}
// 第三阶段:每个块加上自己的偏移,得到全局 inclusive scan 结果
// 输入:block_offsets(每个块的前缀偏移),输出:out(最终结果)
__global__ void add_block_offset(const float* block_offsets,
float* out, int n) {
// out[i] += block_offsets[blockIdx.x]
}
三个 kernel 依次启动,每次启动都意味着一次全局同步------即所有线程块必须全部完成当前 kernel 后,才能进入下一个 kernel。这种方案的优点是逻辑清晰、易于调试;代价是两次 kernel 启动的延迟(通常各为数微秒)以及中间数据在全局内存中的往返读写。
对于块数量很大的情况(例如百万级元素、块大小为 1024,则 C≈1000C \approx 1000C≈1000),第二阶段对 1000 个块和的扫描本身也需要跨块处理------此时可以递归地再次应用同样的策略,形成多级扫描。这就是 CUDA 库中 cub::DeviceScan 的基本结构:它在内部根据块数动态选择需要几级扫描。
单 Pass 方案:依赖管理的光谱
多 kernel 方案虽然直观,但它引入了一个重要的性能瓶颈:每个 kernel 的启动都要求整个 GPU 完成全局同步。在深度学习的训练循环中,一次 scan 的延迟可能直接叠加在关键路径上。如果能在一个 kernel 内完成所有工作,避免 kernel 间的全局屏障,延迟会显著降低。
单 Pass 方案的核心挑战在于依赖管理:块 3 需要知道块 0、1、2 的总和,而这些总和直到 block 0、1、2 各自完成扫描后才可用。我们面临三种依赖管理策略:
策略一:全局计数器 + 原子操作(锁自由同步) 。每个块在完成块内扫描后,将块总和写入全局数组,然后原子递增一个计数器。最后一个到达的块(计数器值达到 CCC)意识到所有块的总和均已就绪,它将负责对块和数组执行扫描,计算出每个块的偏移量,再将这些偏移量写入全局数组,最后递增第二个计数器通知其他块"偏移已就绪"。其余块则在第一个计数器上自旋等待,然后读取自己的偏移并完成最终修正。这个方案只需一个 kernel,但引入原子操作和自旋等待的开销,而且在块数很多时,最后一个块的负载会异常沉重------它必须串行地完成所有块和的扫描。
策略二:协作组(Cooperative Groups) 。CUDA 的协作组编程模型提供了 grid.sync() 原语------它在 kernel 内部提供全局同步,而无需退出 kernel。使用协作组时,kernel 必须通过 cudaLaunchCooperativeKernel 启动,并且所有块必须同时驻留在 GPU 上(否则死锁)。这要求块数不超过 GPU 的并发块容量。在这个约束下,块间扫描的写法与块内扫描几乎对称:
cuda
__global__ void cooperative_scan(const float* in, float* out,
float* block_sums, int n) {
namespace cg = cooperative_groups;
cg::grid_group grid = cg::this_grid();
// 1. 块内 scan,计算块总和
extern __shared__ float smem[];
// ... 块内 up-sweep/down-sweep ...
__syncthreads();
// 2. 块 0 收集所有块总和,做一次小规模 scan
if (threadIdx.x == 0) block_sums[blockIdx.x] = block_total;
grid.sync(); // 全部块完成第 1 步
// 3. 块 0 单线程扫描 block_sums,写回偏移
if (blockIdx.x == 0 && threadIdx.x == 0) {
// 串行扫描 C 个块和
float acc = 0;
for (int i = 0; i < gridDim.x; i++) {
float tmp = block_sums[i];
block_sums[i] = acc;
acc += tmp;
}
}
grid.sync(); // 块 0 完成偏移计算
// 4. 每个块加上自己的偏移
float offset = block_sums[blockIdx.x];
// ... 将 offset 加到本块所有输出上 ...
}
这里 grid.sync() 的使用让代码结构与逻辑结构严格对应:第一次同步确保所有块总和已写入,第二次同步确保偏移已计算完毕。协作组方案的优雅之处在于它在不引入原子操作的前提下实现了块间依赖管理,但代价是对硬件资源提出了占满 GPU 的硬性要求。
需要特别强调的是:使用协作组进行全局同步,必须保证所有线程块同时驻留在 GPU 上,否则会导致死锁 。因此,启动前必须确认 gridDim.x 不超过当前 GPU 的并发块容量(可通过 cudaOccupancyMaxActiveBlocksPerMultiprocessor 查询)。此外,上述实现中块 0 单线程扫描块和数组的时间复杂度为 O(C)O(C)O(C),当块数 CCC 很大(例如数千)时,这一步将成为整个 pipeline 的串行瓶颈。缓解方案是让块 0 内的多个线程并行扫描(利用 warp 级别的前缀和),将复杂度降至 O(logC)O(\log C)O(logC)。
策略三:decoupled look-back(解耦回看) 。这是 CUB 库实际采用的高级方案,它的思想是:块不等待所有前驱块完成,而是按需读取前驱块的部分结果。每个块在完成自身扫描后,写入两个值:块总和和该块的 exclusive 前缀偏移。后一个值在块的前驱尚未完成时无法确定,所以每个块维护一个"状态标志"。需要偏移的块从目标块读取状态------如果目标块已完成,则直接读取其偏移;如果尚未完成,则等待。这种细粒度的依赖管理以"可能阻塞"换取"尽早开始",使得后续块在前驱块完成前就开始自己的块内扫描,缩短了整体延迟。更多块可以在 GPU 上流水线式地同时推进,而非等待全局屏障。
三种方案形成了依赖管理的光谱:无条件全局同步 (多 kernel)、粗粒度全局同步 (协作组)、细粒度按需同步(decoupled look-back)。选择哪种取决于你的约束条件------如果延迟是首要关注点且块数受限于 GPU 并发容量,协作组是最佳平衡;如果数据规模远超单次驻留能力,则必须回到多 kernel 方案或实现 decoupled look-back;如果需要极致的性能并且愿意处理复杂的标志位逻辑,decoupled look-back 提供了理论上最优的调度灵活性。
分层递归:统一视角
无论选择哪种块间扫描策略,一个共同的抽象视角是层级递归。扫描操作在本质上是一个可以无限向上叠加的归约过程:
第 0 层:元素级 scan(块内)
第 1 层:块间 scan(块和序列)
第 2 层:超块间 scan(若块数仍超限)
...
每一层的输出是下一层的输入,每一层的块大小可以不同。这个视角与流式多处理器(SM)的硬件层次完美对应------第 0 层驻留在共享内存中,第 1 层驻留在全局内存中,如果 GPU 规模足够大,第 2 层可以驻留在 L2 缓存中。扫描的递归本质------一个输出依赖 ALL 输入的性质------决定了无论硬件如何演变,这个分层递归结构都是不可避免的。
在实践中,层数的选择取决于三个因素:输入规模、GPU 的 SM 数量、以及共享内存的容量。深度学习的算子融合框架(如 TensorRT、XLA)在编译期就会根据这些参数静态推导出最优的层级划分,从而避免运行时决策的开销。这种编译期规划能力,也是这些框架能够在保持灵活性的同时不牺牲性能的关键。
跨块的扫描设计最终归结为一个问题:如何通过依赖管理,让尽可能多的计算在同一时刻并行推进? 多 kernel 方案用最简单的方式确保正确性,协作组方案用统一原语简化调优,decoupled look-back 用精细的调度换取最低延迟。而分层递归的视角则提醒我们------无论块如何划分,scan 的递归本质不会改变。
案例:非 2 的幂长度处理
上一节的块级实现采用了双元素策略,即每个线程处理两个元素,从而让单块能够处理更多数据。但这种策略隐含了一个前提:输入长度恰好等于块容量的整数倍。而现实中的数组长度几乎总是非 2 的幂:基数排序中一趟扫描的数据量取决于待排序元素的分布,流式计算中每个时间窗口的大小受业务逻辑约束。当输入长度与线程块容量不匹配时,一个不加处理的朴素实现会在两个地方出错:读取越界------线程访问了共享内存中尚未写入数据的槽位;错误累加------垃圾值被当成了有效数据参与前缀和运算。
边界条件:长度不匹配的三种情形
先明确问题的具体形态。假设共享内存为每个线程分配了 2 个槽位(双元素策略),线程块内有 T 个线程,那么理论上可容纳的元素数为 Ncap=2TN_{cap} = 2TNcap=2T。当输入序列实际长度为 N<NcapN < N_{cap}N<Ncap 时,未使用的槽位共有 Ncap−NN_{cap} - NNcap−N 个。这些槽位的分布有两种可能:
- 尾部空缺 :前 NNN 个槽位存放有效数据,后 Ncap−NN_{cap} - NNcap−N 个槽位为空------这是将输入按顺序加载到共享内存后的自然结果。
- 交错空缺:如果加载时采用了某种带状映射(strided mapping),空缺位置会散布在整个数组中。
对于 up-sweep 阶段来说,问题出现在层间合并时。回顾 Blelloch 算法的 up-sweep:每一层将相距 2layer2^{\text{layer}}2layer 的两个节点相加,结果写入右节点。在标准实现中,一个线程负责一个加法操作,因此线程的活跃范围是 [0, N/2) 。当 NNN 恰好是 2 的幂时,每一层的线程数量都能整除,不会越界;但当 NNN 不是 2 的幂时,最后一层(或最后几层)的合并操作可能访问到超出 NNN 的索引------如果不对线程活跃范围加以控制,或者不对空槽位做特殊处理,这些线程就会读取到上一轮遗留的垃圾值。
更隐蔽的问题出现在 down-sweep 阶段的种子传播。down-sweep 要求每个内部节点持有"左子树的总和"作为右子树的种子。如果某个右子树对应的区间中包含空缺槽位,且空缺槽位中不是中性元(例如加法下的 0),那么从该节点传下去的种子就是错的,最终污染整个右半部分的前缀和结果。
Mask:让空槽位"隐形"的机制
解决上述问题的通用方法是用掩码标记有效元素 。具体来说,在将输入加载到共享内存时,同步维护一个等长的掩码数组 mask[]:
mask[i] = 1 表示 共享内存槽位 i 中存放的是有效输入
mask[i] = 0 表示 槽位 i 为空缺
掩码本身也需要参与扫描 。这里有一个关键的设计决策:掩码的扫描运算符是什么?如果原始运算符是加法(前缀和),掩码应该在同一个 up-sweep / down-sweep 过程中同步进行归约,但使用的规则不同------掩码的归约运算符是逻辑与(AND)。原因在于:一个节点对应的区间"全部有效"当且仅当该区间内所有槽位都有效。因此:
- up-sweep 阶段:
mask[idx] = mask[2*idx+1] & mask[2*idx+2] - down-sweep 阶段:掩码的种子传播同样遵循左子树掩码 AND 右子树掩码的规则
当 down-sweep 完成时,掩码数组的最后一个元素 (即全局的 AND 归约结果)告诉我们整个数组是否全部有效。每个输出位置 iii 的有效性由 mask[i] 直接给出。
有了掩码之后,处理空缺槽位就有两条路线可选:
路线一:填充中性元 。将空缺槽位写入运算符的中性元。对于加法,中性元是 0;对于乘法,中性元是 1;对于取最小值,中性元是 +∞+\infty+∞。这种方式的好处是------操作序列完全不需要改动。空缺槽位被中性元填充后,参与合并时不会影响任何有效元素的结果,前提是运算符满足结合律且中性元定义正确。
路线二:条件执行。在每个线程的合并操作前检查掩码:
cuda
// 仅当左右两个槽位都有效时才执行合并
if (mask[left] && mask[right]) {
value[idx] = value[left] + value[right];
mask[idx] = 1;
} else if (mask[left]) {
// 左有效右空缺:结果直接取左值
value[idx] = value[left];
mask[idx] = 1;
} else {
// 左右都空缺:保持空缺
mask[idx] = 0;
}
路线一的性能优势 是明显的,因为它避免了分支发散;而路线二的优势 在于不需要预先知道中性元------这在运算符为自定义结构体(如 <key, position> 的字典序比较)时至关重要。实践中,绝大多数 CUDA 库(包括 CUB 的 DeviceScan)采用填充中性元 的策略,因为加法、乘法、min、max 等常见运算符都有明确的中性元,而自定义运算符通常可以设计为支持中性元的结构。
线程越界与边界裁剪
掩码解决了"垃圾值参与运算"的问题,但还有一个独立的议题:线程本身的越界访问 。当一个线程块的线程数 TTT 大于实际需要的操作数时,多余的线程在循环加载阶段需要被显式地限制活跃范围:
cuda
// 加载阶段:仅当全局索引小于 N 时才从全局内存读取
// 越界槽位必须显式初始化,防止后续合并操作读取未定义值
int tid = threadIdx.x;
int global_idx = blockIdx.x * blockDim.x + tid;
if (global_idx < N) {
sdata[tid] = input[global_idx];
smask[tid] = 1;
} else {
sdata[tid] = 0; // 填充中性元
smask[tid] = 0; // 标记为无效
}
这段代码同时完成了两件事:中性元填充 和掩码标记 。从全局内存读取时,条件 global_idx < N 防止了越界读取;对超出范围的线程,显式写入中性元和掩码 0,保证后续 up-sweep / down-sweep 阶段的每一次合并操作都能安全执行------无论线程的活跃范围如何,它们读取的共享内存位置都已被初始化。
在 down-sweep 的输出阶段,同样的边界检查再次出现:
cuda
// 写出阶段:仅将有效的槽位回写到全局内存
if (global_idx < N) {
output[global_idx] = sdata[tid];
}
对齐舍入:另一种视角
填充中性元的一种特例是对齐舍入 (rounding up to power of two)。如果预先知道输入长度 NNN,可以计算出下一个 2 的幂 N′N'N′,然后将数组从 NNN 扩展到 N′N'N′,空缺部分全部用中性元填充。这种方式下,up-sweep / down-sweep 的每一层操作数量都是整数,不需要任何条件分支 ------所有线程在所有层中都保持活跃。代价是内存和计算量从 NNN 增加到 N′N'N′,当 NNN 接近 2k+12^{k} + 12k+1 时,浪费接近一倍。
掩码方案则在时间和空间上都更精细:它不扩展数组,仅仅是标记,因此额外的内存开销只有 O(N)O(N)O(N) 比特(或 O(N)O(N)O(N) 字节,取决于实现),且计算量的增加仅在空缺位置对应的节点上。对于长度接近 2 的幂的数组,掩码方案几乎没有浪费;对于长度恰好是 2 的幂的数组,掩码全为 1,等价于无掩码的原始实现。
综合来看,填充中性元 + 掩码追踪有效性的组合,既保证了运算的正确性,又保留了 O(N) 的工作复杂度------这是 Blelloch 扫描在非 2 幂输入下仍然成立的关键。而从本节暴露出的"输入长度与硬件结构不匹配"这一问题,也指引我们走向下一个层次:当数组规模超过单块容量时,我们需要在多块之间建立同样的边界协议------这正是前文所述跨块分段扫描中"段边界扫描"所要处理的精确问题。
至此,我们从 scan 的定义出发,一路走过了 Hillis-Steele 的并行化启蒙、Blelloch 的工作量优化、共享内存中的工程落地、跨块的协同扩展,再到非对齐长度的边界处理。这条路径勾勒出了并行 scan 从理论到实践的全貌------它既是一个优雅的算法家族,也是 GPU 编程中不可或缺的工程原语。无论你接下来要构建的是基数排序、流式判定还是资源分配器,scan 都将是你手中最锋利的工具之一。