嵌套数据并行的完全扁平化

# 嵌套数据并行的完全扁平化 来源：https://futhark-lang.org/blog/2026-07-31-full-flattening.html 发布于 2026 年 7 月 31 日 先说结论：这篇长博客文章描述了一项历时多年努力的成果——我们现在可以，除了编译器 bug 和少数无关紧要的边界情况外，将*任意*Futhark 程序编译为并行 GPU 代码。 ### 引言 (https://futhark-lang.org/blog/2026-07-31-full-flattening.html#introduction) Futhark 中最有趣的实现挑战之一是如何将*嵌套并行*（parallel operations can contain other parallel operations）映射到*扁平并行*。这个挑战之所以存在，是因为许多有趣的硬件（如 GPU）并不能高效地支持嵌套并行，即使在 CPU 上，任意嵌套也可能带来显著的运行时开销。长期以来，Futhark 只支持扁平化一个有限（尽管重要）的程序子集，称为*规则嵌套并行*（以区别于*不规则*嵌套并行），这意味着有些程序我们无法编译为 GPU 代码。2022 年，我开始了一个“完全扁平化”的 (https://github.com/diku-dk/futhark/pull/1740) 实验性实现，截至撰写本文时，它仍是 Futhark 仓库中开放时间最长的拉取请求。我当时已经足够构建出一个概念验证，但所需的工作量巨大，导致这项工作一度停滞——直到现在。主要得益于学生们的努力，我们目前已接近完成 Futhark 中扁平化的重新实现。主要贡献来自 Cornelius Sevald-Krause (https://di.ku.dk/english/staff/?pure=en/persons/700865)，他在其学士论文 (https://futhark-lang.org/student-projects/cornelius-bsc-thesis.pdf) 中实现了函数和 `match` 的提升；尤其要感谢 Amirreza Hashemi，他的硕士工作像旋风一样完成了*其他所有内容*。其中包括各种关键优化，以便在功能和性能上达到与旧扁平化器 (https://github.com/diku-dk/futhark/blob/a3f103a43256b847cfd0a13dca70e818b2f34204/src/Futhark/Pass/ExtractKernels.hs) 的同等水平。我无法用语言表达，能够为一个充满优化启发式、已有十年历史的编译器 pass 提出一个系统化且可维护的替代方案，是多么令人印象深刻。这篇文章将介绍扁平化的基本问题，梳理其中的一些挑战，并解释我们如何在即将合并的 Futhark 编译器工作中解决这些问题。扁平化有着悠久的学术传承，我会尽量链接相关的工作，但本文并非对先前工作的系统性综述。我会忽略编译器为了实现这一切而需要做的大部分非平凡的簿记工作，而是专注于程序变换背后的直觉。 我应该事先说明这项工作的目标：新的扁平化变换不应降低那些旧实现已经能很好处理的程序的性能，但我们不太关心那些之前无法处理的程序的性能。正如我们将看到的，要在任意不规则的嵌套并行上实现*真正*良好的性能，本身就是一个重大的研究问题，而我们打算在未来的时间里逐步解决它。另外，这是一篇很长的文章。我尽量让内容保持在相对高的层次，但如果你觉得内容过于冗长，可以直接跳到文末的反思部分。 ### 扁平化的基础 (https://futhark-lang.org/blog/2026-07-31-full-flattening.html#basics-of-flattening) 为了引出扁平化，我们假设一个有点漫画式的 GPU 视图。扩展到完整的 GPU 功能并不会从根本上改变故事；只是需要更复杂的例子而已。此外，尽管这篇文章聚焦于 GPU，你可以在脑海中将其替换为“高效但受限的处理器”——这些问题在任何向量处理中都会出现。用函数式编程的语言来说，我们首先假设 GPU 可以高效地执行一个 `map`，其中函数包含顺序代码，并且所有中间结果的内存都可以预先分配。为了避免一直写 `zip`/`unzip`，我们还假设我们的 `map` 可以接受任意数量的数组。例如，下面这个就是可以在 GPU 上高效执行的： ``` map (\x y -> x + y) xs ys ``` （顺便说一句，在描述接下来的变换时，如果 `map` 像 Futhark IR 中那样先接受数组*再*接受 lambda，事情会简单得多——但我不想在一篇博客文章中走那么远。） 很容易想象，这可以变成一个 GPU 程序，`map` 的每次迭代对应一个线程。除了带有顺序代码的单个 `map`，我们还允许任意嵌套的 `map`。所以下面这个也是高效的： ``` map (\xs ys -> map (\x y -> x + y) xs ys) xss yss ``` 同样很容易想象，这些 `map` 可以被折叠，最内层 `map` 的每次迭代对应一个 GPU 线程。除了任意顺序代码外，我们还允许最内层操作是*单个*并行操作，比如 scan、reduction、transpose——具体是什么并不重要。那么下面这个也是高效的： ``` map (\xs -> reduce (+) 0 xs) xs ``` 本质上，“高效 GPU 代码”意味着在某个单一并行原语操作（或顺序代码）之上叠放一层层 `map`。这些模式之所以高效，是因为我们知道如何从这些模式生成高效的 GPU*kernel*。这与骨架 (https://en.wikipedia.org/wiki/Skeleton_(computer_programming)) 有些关系，但我们通常不使用这个术语，因为我们打算将这些模式作为编译器的推理机制，而不是用户可见的编程模型。在实际编译器中，我们当然有专门的构造来可靠地表达“一个内部含有特定内容的 `map`”，但为了尽量减少本文的语法，我主要会使用 `map` 叠放的形式。 当然，Futhark 程序员可能不会写出看起来那样的程序。这些大胆而思想自由的人可能反而会这样写： ``` map (\xs -> let y = reduce (+) 0 xs let zs = map (+y) xs in zs) xss ``` 这确实符合我们讨论过的模式之一：如果我们丢弃内部 `map` 和 `reduce` 的并行性，那么这就是一个包含顺序代码的 `map`。在编译并行语言时，这是一个重要的选择：你不*必须*利用所有的并行性。从语义上讲，大多数并行语言并不*要求*并行执行才能正确；这只是实现可以选择的选项。当然，我们也不应过度——一个不利用*任何*并行性的并行语言很可能会被指责为违反了契约 (https://futhark-lang.org/blog/2022-01-27-cost-models-are-contracts.html)。 在这种情况下，我们可以把程序重写为语义相同、但用高效模式表达的形式。为了清楚地说明发生了什么，首先让我们添加一些类型注解，并假设 `xs` 的形状是 `[n][m]`： ``` map (\xs -> let y : i32 = reduce (+) 0 xs let zs : [m]i32 = map (+y) xs in zs) xss ``` 然后我们将外层 `map` 拆分成两个独立的 `map`，第一个返回一个中间结果数组，传递给第二个 `map`： ``` let ys = map (\xs -> reduce (+) 0 xs) xss let zss = map (\xs y -> map (+y) xs) xss ys ``` 这两个 `map` 现在符合我们知道如何转换为高效 GPU 代码的模式。我们在这里做的变换通常称为*循环分布*或*裂变*，它是循环融合的对应物。对于纯函数式语言，以下等式成立： ``` map f ∘ map g = map (f∘g) ``` （一旦允许多个输入数组，这条规则会更加复杂，但我们先忽略这一点。）通常我们从左到右应用这条规则以进行融合，避免将中间数组存储在内存中；但我们也可以从右到左应用它，这时它为我们提供了一种方法，将复杂的 `map` 分解为一系列只执行单一操作（或复合操作——我们可以按所需的程度进行裂变，后面会再谈）的简单 `map`。这是我们用来将任意复杂的数据并行程序切割成更小、更易管理的块的主要技术。据我所知，这是 Guy Blelloch 在 90 年代初作为其 NESL (https://www.cs.cmu.edu/~scandal/nesl.html) 工作的一部分首创的，尽管它当时（现在也是）有时被称为“向量化”，因为其输出在完全应用时由内部带有平凡操作的 `map` 和 `scan` 组成，可以被视为“向量操作”。 然而，在拥有一个可工作的扁平化算法的所有组成部分之前，我们还需要解决两个主要问题。 ### 不规则的中间数组 (https://futhark-lang.org/blog/2026-07-31-full-flattening.html#irregular-intermediate-arrays) 考虑这个程序： ``` map (\m -> let xs : [m]i64 = iota m let y : i64 = reduce (+) 0 xs in y) ns ``` 回忆一下 `iota m` 产生一个大小为 `m` 的数组 `[0, ..., m-1]`。这个程序与上面有同样的问题，即并行操作并非完美嵌套。所以让我们尝试裂变这个 `map`： ``` let xss : [n][???]i64 = map (\n -> iota n) ns let ys = map (\xs -> reduce (+) 0 xs) xss ``` 注意 `xss` 绑定类型中的问号——我们应该在那里填什么大小？直觉上，操作 `map (\n -> iota n) ns` 会产生一个不规则数组，其中各行的长度不同，取决于输入数组 `ns` 的值。有人可能期望 `map (\n -> iota n) [1,2,3]` 产生数组 `[[0],[0,1],[0,1,2]]`，但不幸的是，Futhark 不允许不规则数组；无论是在源语言中还是在我们的 IR 中。其中的原因是出于性能考虑，因为一旦有了不规则数组，你要么需要将数组表示为指针数组，要么会失去随机访问索引的能力。当目标是向量处理器时，这两者都不可接受。这意味着我们不想在源语言或 IR 中允许不规则数组。 这个问题正是我们最初扁平化变换局限性的根源：它无法处理*不规则的嵌套并行*，即并行度在外层并行维度的各次迭代之间变化的情况。在上面的例子中，`m` 是最外层 `map` 的局部变量，因此 `iota` 和 `reduce`（两者的并行宽度都是 `m`）是不规则的。我们唯一的选择就是将这些代码顺序化。现在，在实际编译器中，`iota` 会被融合进 `reduce`（使用一种我不会展开的表示），但如果我们忽略这一点，似乎我们违反了一个约束，即所有中间结果的内存必须预先分配——而这里我们需要一个大小为 `m` 的数组，其中 `m` 对每个线程可能不同。 在这种情况下，问题可以通过一个 inspector 来解决，它提取程序的一个*切片*，用于计算所有中间大小（但不计算其他值），然后用来计算内存需求并在内存中组织中间结果。这有时是可以做到的（粗略地说，当没有顺序循环的棘手问题时），但这是一种尽力而为的事情，而且它并不会带回并行性。 无论如何，处理不规则数组的能力显然对于处理不规则嵌套并行这一总体目标至关重要。解决方案是找出一种将不规则数组表示为扁平（一维）数组的方法，同时仍然允许我们需要的操作并行执行。核心构建块是将二维不规则数组编码为一个*数据向量*和一个*标志向量*，两者大小相同。数据向量只是数组各行拼接在一起的结果，而标志向量则指示其对应元素是否是新一行的开始。例如： - **嵌套的不规则数组：** `[[1], [2,3], [4,5,6]]` - **值向量：** `[1,2,3,4,5,6]` - **标志向量：** `[true, true, false, true, false, false]` （限制为二维数组并不失一般性：我们总是可以折叠额外的内部维度，并用辅助数组告诉我们如何重建它们。为简单起见，我们不再深入细节。） 我们称其为*分段数组*，而“不规则的行”被称为*段*。上面的数组有三个段，大小分别为 1、2 和 3。在实践中，我们经常使用一些额外的辅助数组，用于跟踪每个段的大小（*形状向量*，在此例中为 `[1,2,3]`）、每个段在数据向量中的起始偏移（*偏移向量*，此处为 `[0,1,3]`），以及将数据数组的每个元素关联到其段索引（`[0,1,1,2,2,2]`）和段内索引（`[0,0,1,0,1,2]`）的数组。不过，这些都可以从标志向量高效地计算出来。 这种表示上的基本操作是*分段扫描*，它可以借助标准 `scan` 定义如下： ``` def segscan 't [n] (op: t -> t -> t) (ne: t) (flags: [n]bool) (vals: [n]t) : [n]t = let pairs = scan (\(v1, f1) (v2, f2) -> let f = f1 || f2 let v = if f2 then v2 else op v1 v2 in (v, f)) (ne, false) (zip vals flags) let (res, _) = unzip pairs in res ``` 我们可以将其应用于上面不规则数组的表示，以获得对每个段进行（包含式）扫描的结果，并拼接成单个数组： ``` > segscan (+) 0 [true, true, false, true, false, false] [1,2,3,4,5,6] [1, 2, 5, 4, 9, 15] ``` 这是表示分段数组 `[[1,2,5]++[4,9]++[15]]` 的值数组，使用原始标志数组不变。由于扫描在并行算法 (https://www.cs.cmu.edu/~guyb/papers/Ble93.pdf) 中是如此重要的构建块，一旦我们能扫描这些分段数组，我们也能做各种其他事情也就不足为奇了。例如，我们可以先执行分段扫描，然后提取每个被扫描段的最后一个元素，从而进行分段*归约*： ``` def segreduce 't [n] (op: t -> t -> t) (ne: t) (flags: [n]bool) (vals: [n]t) : ?[m].[m]t = segscan op ne flags vals |> zip (rotate 1 flags) |> filter (.0) |> map (.1) ``` ``` > segreduce (+) 0 [true, true, false, true, false, false] [1,2,3,4,5,6] [1, 5, 15] ``` 类似地，我们可以构造一个分段 `iota`，它接受一个类型为 `[]i64` 的*形状向量*作为输入，并在道德上对这个数组映射 `iota`。该构造基于一个观察：普通 `iota` 可以写成对包含 `n` 个 1 的数组进行排他前缀和 (https://futhark-lang.org/examples/exclusive-prefix-sum.html)： ``` def iota n = expresum (replicate n 1) ``` 分段 `iota` 则只是对某个分段数组进行分段扫描，其中数据向量全为 `1`，标志向量在新 `iota` 开始时为 `true`，这可以通过对形状数组进行排他前缀和来确定。由于我们的分段扫描是包含式的而非排他式的，我们必须在最后从每个元素中减去 1： ``` def segiota [k] (ns: [k]i64) : ?[m].[m]i64 = let m = i64.sum ns let offsets = expresum ns let flags = spread m false offsets (replicate k true) in segscan (+) 0 flags (replicate m 1) |> map (\x -> x - 1) ``` 顺便说一句，形状向量的分段 `iota` 给出了分段表示的段内索引。 本质上，每当我们谈论分段 `foo` 时，我们指的是这样一种操作：它直观地执行将操作 `foo` 进行 `map` 的工作，要么消费、要么产生一个分段数组。只要这些分段操作符合我们关于什么可以转换为高效（或仅仅可执行）GPU 代码的规则，那就没问题——而且你可能会注意到，所有分段操作