Futhark 应如何向程序员暴露不规则数组？

# 如何向程序员展现 Futhark 的不规则数组？ 来源：https://futhark-lang.org/blog/2026-08-12-flatmap.html 发布于 2026 年 8 月 12 日 最近 Futhark 增加了一些强大的新特性，例如 扁平化非均匀数据并行（https://futhark-lang.org/blog/2026-07-31-full-flattening.html） 和 递归函数（https://futhark-lang.org/blog/2026-08-05-recursion.html）。后者明显是用户可见的，我们将在本文后面看到一些不错的用法，而扁平化主要是一种程序变换。简而言之，它让编译器能够处理诸如 `` map (\n -> i64.sum (iota n)) ns `` 这样的表达式，其中内层并行度（`n`）在外层 `map` 的不同迭代之间是不同的，并且能充分利用所有并行性。添加这一特性并没有从根本上改变可以编写的 Futhark 程序的类型或语义，它只是影响了我们能生成多好的 GPU 代码（在极端情况下，影响我们能否生成 GPU 代码）。在本文中，我们将看到如何通过一个相对较小的扩展，使语言从根本上更具表现力，以及我们现在可以用 Futhark 编写哪些类型的程序。 ### 不规则数组 (https://futhark-lang.org/blog/2026-08-12-flatmap.html#irregular-arrays) Futhark 的主要语言限制是不允许不规则数组。不规则数组是指多维数组中子数组的大小不同。例如，``[[1], [3,3]]`` 是不规则的，而 `[[1,2],[3,4]]` 是规则的。这个限制的根本原因在于编译效率的考虑，但从用户体验来看，这些数组在 Futhark 的 大小类型系统（https://futhark-lang.org/blog/2019-08-03-towards-size-types.html）中根本无法良好地类型化。然而，实现扁平化的主要挑战是处理在执行 循环分布（https://en.wikipedia.org/wiki/Loop_fission_and_fusion）时作为中间结果出现的不规则数组： `` map (\n -> i64.sum (iota n)) ns ↓ let tmps = map (\n -> iota n) ns let res = map (\tmp -> i64.sum tmp) tmps `` 这意味着扁平化变换实际上 完全能够处理不规则数组（https://futhark-lang.org/blog/2026-07-31-full-flattening.html#irregular-intermediate-arrays），通过将它们编码为规则数组，而编译器的其余部分无需知道这一点。 不规则数组在某些情况下确实很有用，所以问题在于我们如何在不*添加*对不规则数组的全面支持的情况下，在源语言中展现部分这种能力。作为我们想编写的那种代码的一个例子，这是一个用 NESL（https://www.cs.cmu.edu/~scandal/nesl.html） 编写的递归数据并行快速排序： `` function quicksort(a) = if (#a < 2) then a else let pivot = a[#a/2]; lesser = {e in a| e < pivot}; equal = {e in a| e == pivot}; greater = {e in a| e > pivot}; result = {quicksort(v): v in [lesser,greater]}; in result[0] ++ equal ++ result[1]; `` 注意主要技巧：对一个包含两个元素的数组进行映射，从而执行递归的 `quicksort` 调用，尽管在 NESL 中这是通过数组推导式而不是 `map` 函数完成的。这个数组在 NESL 中是不规则的，这当然在 Futhark 中行不通。这将是我们的起步示例，不过我们很快就会看到，快速排序实际上是其中比较简单的案例。 ### 初始的 `flatmap` SOAC (https://futhark-lang.org/blog/2026-08-12-flatmap.html#an-initial-flatmap-soac) Futhark 允许通过*二阶数组组合子*（SOAC）进行并行编程：这些函数包括 `map`、`reduce` 和 `scan`。我们通过一个新的 SOAC `flatmap` 来扩展它，它在语义上类似于 `map`，只是映射函数返回的结果大小可以不同，并且 `flatmap` 返回这些数组的拼接结果： `` val flatmap [n] 'a 'b : (f: (k: i64) -> (x: a) -> [k]b) -> [n]i64 -> [n]a -> ?[m].[m]b `` 注意大小类型的巧妙运用：`flatmap` 接受两个大小均为 `n` 的数组，其中第一个数组（类型为 `[n]i64`（https://futhark-lang.org/blog/2023-03-20-why-are-sizes-signed.html））包含 `f` 函数在对应的 `a` 数组上应返回的预期大小。大小 `k` 也会传递给映射函数，这样我们可以在其返回类型中使用它。`flatmap` 的结果具有存在性大小 `m`，因为我们无法预先知道它。（实际上我们可以：`m` 就是 `[n]i64` 数组的和。） 这与 Haskell 的 `concatMap` 类似，尽管意图是 `flatmap` 能充分利用 `n` 个输入元素之间的并行性*以及* `f` 中可能存在的任何并行性。然而，*在语义上* `flatmap` 并无特殊之处，很容易用顺序循环实现。即使是 `flatmap` 的并行实现（这个版本和下面的版本）也基本是直接的，因为它只是对所提供函数“提升”形式的一次调用，所以在本文中我们专注于 `flatmap` 是否提供了我们表达所需算法所需的灵活性。 这个 `flatmap` 足以在 Futhark 中实现递归快速排序： `` def quicksort [n] (xs: [n]i32) : [n]i32 = if n <= 1 then xs else let pivot = xs[0] let [m][k] (lesser: [m]i32, greater: [k]i32) = partition (<= pivot) (drop 1 xs) let sorted = flatmap (\k p -> quicksort (sized k (if p then lesser else greater))) [m, k] [true, false] in sized n (take m sorted ++ [pivot] ++ drop m sorted) `` 这段代码看起来比普通的递归快速排序稍微奇怪一些。奇怪之处在于构造一个包含两个元素的数组只是为了立即对它进行 `flatmap`，而不是执行两次递归调用，但这是为了并行运行这两个递归调用所必需的。在 NESL 中这个数组是不规则的，而在 Futhark 中我们构造一个布尔数组（`[true, false]`），并在 `flatmap` 函数内部使用条件表达式来选择 `lesser` 或 `greater` 数组。我预计这可能会成为在 Futhark 中编写分治程序的常用技巧。 ### 泛化 `flatmap` (https://futhark-lang.org/blog/2026-08-12-flatmap.html#generalising-flatmap) 尽管足以表达快速排序，但上面的 `flatmap` 有一个限制：我们必须能够预先确定每个结果的大小。这对于排序来说是足够的，因为对数组排序不会改变其大小，但对于像 Quickhull（https://en.wikipedia.org/wiki/Quickhull） 这样的算法来说并不实用，因为其结果的大小通常小于输入的大小。明显的扩展是修改 `flatmap`，使映射函数返回的数组大小是存在量化的： `` val flatmap [n] 'a 'b : (f: (x: a) -> ?[k].[k]b) -> [n]a -> ?[m].[m]b `` 现在 `k` 不再是 `f` 函数的参数，而是绑定在其返回类型中。这意味着我们不再知道结果大小的任何信息，这很快就被证明是不切实际的。因此，我们将 `flatmap` 进一步扩展，让它也返回 `f` 返回的每个数组的大小，以形状向量的形式： `` val flatmap [n] 'a 'b : (f: (x: a) -> ?[k].[k]b) -> [n]a -> ?[m].([n]i64, [m]b) `` 事实上，Futhark 预lude（https://futhark-lang.org/docs/prelude/doc/prelude/soacs.html#term:flatmap）中*真正的* `flatmap` 不仅返回形状向量，还返回扁平化无论如何都会计算的各种其他相关元数据，但为了本文的简洁性，我们只保留形状向量。 使用这个定义，我们可以实现 Quickhull。我不会给出 完整的程序（https://github.com/diku-dk/futhark/blob/master/tests/quickhull.fut），因为它涉及许多用于几何计算的函数，但递归部分可以这样写： `` def hull [n] (a: point) (b: point) (pts: [n]point) : []point = if n <= 1 then pts else let p = farthest a b pts let f i = let (x, y, pts') = if i == 0 then (a, p, filter (\q -> side a p q > 0) pts) else (p, b, filter (\q -> side p b q > 0) pts) in hull x y pts' let (shape, pts') = flatmap f [0, 1] in take shape[0] pts' ++ [p] ++ take shape[1] pts' `` 实际上，我们可以用一种稍微巧妙的方式来编写，完全不使用 `shape` —— 如果你感兴趣的话，可以看看链接。 ### 进一步泛化 `flatmap` (https://futhark-lang.org/blog/2026-08-12-flatmap.html#generalising-flatmap-further) 作为最后的泛化，尽管我们还没有具体的用例，我们也允许 `flatmap` 为每个输入元素返回一个*统一*的结果，就像普通的 `map` 一样。这样我们就得到了最终的类型： `` val flatmap [n] 'a 'b 'c : (f: (x: a) -> ?[k].([k]b, c)) -> [n]a -> ?[m].([n]i64, [m]b, [n]c) `` 每当我们对统一结果不感兴趣时，我们可以简单地将 `c` 实例化为 `()`，并且 prelude 中包含一个名为 `flatmap'` 的包装器来完成这一操作。 尽管大小类型目前已经是 Futhark 中相当成熟的特性，但当它们让我们能够如此清晰而精确地表达像 `flatmap` 这样的函数类型时，我仍然会感到惊喜。通过这个泛化，`flatmap` 现在已经达到了它的最终形态。我喜欢它是一个如此小的扩展，甚至都算不上“语言扩展”（尽管编译器 IR 必须修改），并且它的语义完全平凡。问题是，尽管语言本身不支持不规则数组，它在多大程度上足以清晰而优雅地表达，比如说， NESL 的并行算法库（https://www.cs.cmu.edu/~scandal/nesl/algorithms.html）？ ### 性能 (https://futhark-lang.org/blog/2026-08-12-flatmap.html#performance) 故事还有一个最后的波折：`flatmap` 有多快？在 扁平化文章（https://futhark-lang.org/blog/2026-07-31-full-flattening.html）中 我提到非均匀情况的性能尚未得到优先考虑，而在 递归文章（https://futhark-lang.org/blog/2026-08-05-recursion.md）中 我提到在数据并行语言中，递归似乎很容易造成非常高的成本。我们现在把两件低效的事情结合起来，所以刺猬精神（https://futhark-lang.org/blog/2026-02-09-hedgehogs.html）可能会有所欠缺。但它到底有多糟糕？我们很幸运地处于这样一个境地：`flatmap` 并没有从根本上增加语言的表达能力，因为手工扁平化总是可能的（尽管单调且容易出错），所以对于我们现在可以用递归分治实现的算法，我们确实有手工扁平化的版本。例如，让我们将 sorts 包中的快速排序与上面的递归快速排序进行比较： `` -- == -- entry: bench_flatmap bench_manual -- random input { [100000]i32 } entry bench_flatmap = quicksort -- from above module M = import "lib/github.com/diku-dk/sorts/quick_sort" entry bench_manual = M.qsort (i32.<=) `` 请注意，这个包中的 `qsort` 已知是相当低效的（我们通常建议在 GPU 上使用归并排序或基数排序），所以这里的排序在绝对意义上都远非快速。但让我们看看它有多糟糕： `` $ futhark bench --backend=cuda sortbench.fut Compiling sortbench.fut... Reporting arithmetic mean runtime of at least 10 runs for each dataset (min 0.5s). More runs automatically performed for up to 300s to ensure accurate measurement. sortbench.fut:bench_flatmap (no tuning file): [100000]i32: 231986μs (95% CI: [ 230655.0, 233282.1]) sortbench.fut:bench_manual (no tuning file): [100000]i32: 12079μs (95% CI: [ 12059.0, 12098.0]) `` 相当糟糕。我还没有调查为什么 `bench_flatmap` 这么慢，但我怀疑是因为扁平化产生了大量的小内核，以及递归调用的固有开销，每个激活记录维护着许多相当大的扁平化元数据和值的数组。手工扁平化的快速排序（`bench_manual`）作为排序而言非常慢，因为有过多的簿记工作，但它是用尾递归循环编写的，并且并行操作的数量最少。 这提出了一个有趣的研究问题：我们能否自动将这种递归分治函数编译成等价于手工扁平化的尾递归代码？是的，我记得我写过为什么我不喜欢尾调用优化（https://futhark-lang.org/blog/2026-01-20-why-not-tail-recursion.html），但在这种情况下，它似乎是获得良好性能所必需的。然而，这是将来的事；如果程序分析变得过于棘手或脆弱，我们可能需要一个专门的 SOAC，以编译器可以利用的方式表达各种分治模式。