将 Futhark 留在 GPU 之外

# 让 Futhark 保持在 GPU 之外 来源：https://futhark-lang.org/blog/2026-10-02-cpu_function.html 发布于 2026 年 10 月 2 日 今年早些时候，Elias Smedegaard 完成了一篇学士论文（https://futhark-lang.org/student-projects/elias-bsc-thesis.pdf），内容是关于通过自动微分高效计算稀疏 Jacobian 矩阵（https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant#Jacobian_matrix）的。要理解这篇博文，不必弄清这究竟是什么、有什么用处——它与自动微分（https://autodiff.org/）有关，而我还没来得及就此写一篇博客。就本文而言，关键的细节在于：他解决方案的一部分涉及对一个对应于矩阵稀疏模式的图进行着色。注意，*最优*图着色是一个著名的 NP-hard 问题，但我们并不需要最优着色——我们只需要*不错*的着色，并且用高效的算法实现。Elias 找到了一个相当快的贪心距离-2 着色算法（如果你好奇的话，是这篇论文（https://epubs.siam.org/doi/10.1137/S0036144504444711）中的 Algorithm 3.1），但遗憾的是该算法本质上是顺序执行的。这对 Futhark 来说不是问题，因为 Futhark 支持相当高效的原地更新（https://futhark-lang.org/blog/2022-06-13-uniqueness-types.html），但问题在于 Elias 的整个程序有两步： 1. 对图进行着色。 2. 利用着色结果，借助大规模并行性计算 Jacobian。 我们*非常*希望第二步能够利用并行硬件（如 GPU），而第一步则在 CPU 上执行。 结果发现，Futhark 的一些实现选择让这件事变得很别扭。本篇博文讲述问题是如何产生的、一个有原则的解决方案可能是怎样的，以及为什么我们最终只是打了个补丁。 ### 问题是怎么产生的 Futhark 的编译模型相当简单。数组必须存放在内存中，而如果你使用 GPU 流水线之一，那么*所有*数组都会被放入 GPU 内存。这并不是因为我们的内部表示（https://futhark-lang.org/blog/2024-03-06-array-representation.html）无法表达 CPU 和 GPU 内存的混合，而纯粹是实现上的便利。这同样是一个安全的设计，因为 CPU 代码可以访问 GPU 内存（通过开销较大的 API），而 GPU 代码通常无法访问 CPU 内存。但可行并不意味着高效。考虑这个假想的 Futhark 循环： ``` loop sum = 0 for i < n do sum + xs[i] ``` 如果 `xs` 是一个存储在 GPU 内存中的数组，那么循环体中将包含代码，费力地从 GPU 内存中复制单个元素到 CPU 内存中的某处，然后才能将其加到 `sum`（大概率存放在寄存器中）。这*太慢了*！复制单个数组元素本身几乎不花时间，但设置通信以及阻塞等待复制完成的开销是毁灭性的——每个数组元素轻松就要好几微秒。 更好的办法是在循环开始之前把整个 `xs` 数组复制过去，以摊销通信成本——这确实是更好的解决方案，下文我们会讲到。但只要对 GPU 数组的索引出现在 CPU 代码中，我们就会生成一次非常慢的 GPU 内存读取，如果这发生在循环内部，那就是灾难性的。 Elias 实现的距离-2 着色算法本质上就是几个操作图和栈的数组的嵌套顺序循环。用 GPU 后端编译这段代码，与使用顺序的 `c` 后端相比，性能差得一塌糊涂——轻松差三个数量级甚至更多。如果我们改用 `c` 后端，性能则非常出色（基本就是你用 C 手写能得到的水平），但当然这样一来，整体程序的第二部分也只能顺序执行了。作为折中，我们可以使用 Futhark 的 `multicore` 后端来生成并行的 CPU 代码，但我真的很想在 GPU 上计算那些 Jacobian！ ### 一个有原则的解决方案可能是怎样的 Futhark 不支持"分开编译"，即让程序的不同部分使用不同的后端，而且这似乎也是一个相当笨拙的解决方案。相反，最好的方案就是我上面暗示过的：查看数组实际被*如何*（或*在哪里*）使用，并据此决定是把它们放进 CPU 内存还是 GPU 内存——或者在程序需要时同时冗余地存放在两处。 我们实际上已经有一个类似的优化（https://github.com/diku-dk/futhark/blob/c04d9c47185c20a7ce366a3b7b203f004412dc87/src/Futhark/Optimise/ReduceDeviceSyncs.hs），由 Philip Børgesen 实现，它会把简单的顺序*计算*（不是数据！）迁移到 GPU 上，只要计算结果本来就只在 GPU 上被使用。目标是节省通信——如果所做的工作微不足道，而所有时间都花在复制几个数组元素上，那么更好的做法是在一个单线程的平凡 GPU kernel 中完成顺序工作，只是为了减少通信。 可惜的是，这个优化在这里不适用。首先，距离-2 着色算法太复杂了（它有循环），无法触发这个优化；而且即使我们禁用这些检查，强制它在 GPU 上执行，结果也很明显——单个 GPU 线程运行这种重度循环的顺序代码慢得一塌糊涂。 在理想世界里，我们可以想象这个优化有一个对应版本，它迁移的是*数据*而不是*计算*，而且底层的许多原则很可能是相通的。尤其是，我们必须小心避免过多的复制，*同时*也要避免不必要地保留这些副本，因为它们会增加内存消耗。我已经能勾勒出这个优化大概的样子，但我也明白，它值得比我现在拥有的更多的时间和精力去认真对待。因此…… ### 我们实际采用的补丁 我没有试图在编译器里做些巧妙的东西，而是添加了一个新的 attribute（https://futhark-lang.org/blog/2020-06-28-attributes-in-futhark.html），名为 `#[cpu_function]`。在我最初的设计中，当放在一个不内联的函数（https://futhark-lang.org/blog/2024-10-28-inlining.html）上时（通常配合 `#[noinline]` 强制要求），`#[cpu_function]` 会导致该函数的函数体被编译成 CPU 代码，并且函数的所有输入、中间和输出数组都存放在 CPU 内存中，不论使用的是哪个编译器后端。调用该函数的代码负责确保数组实参位于正确的内存空间，编译器会插入代码来保证这一点。 以下是使用示例： ``` #[noinline] #[cpu_function] def seqsum [n] (xs: [n]i32) = loop sum = 0 for i < n do sum + xs[i] ``` 我这样做是因为我以为这会很容易。而它几乎确实是！我们用于内存的中间表示从来都是（而且是经过精心设计的）可以并发使用多种内存的（具体来说，我们的 `mem` 类型以一个内存空间为参数，我们也用它来表示 GPU 共享内存）。然而，我们此前从未大规模*使用*过这一功能，这导致了一些必须解决的问题。本着这种精神，我挑选的解决方案都是当时最方便的办法，直到将来我有时间再去做*那件有原则的事*。 我遇到的一个问题是，我们 GPU kernel 的代码生成假定函数体中引用的所有数组都已在 GPU 内存中，而现在已经不是这样了。例如，假设这是我们的程序： ``` #[noinline] #[cpu_function] def frob (n: i64) = iota n entry main n = let xs = frob n in map (\x -> x + 1) xs ``` `main` 中的 `map` 被编译成一个 GPU kernel，但数组 `xs` 由于 `#[cpu_function]` 而位于 CPU 内存中！这是不合法的。*有原则的*解决方案是实际检查 GPU kernel 内部访问了哪些数组，并在必要时把它们复制到 GPU 上，但把这件事*做对*其实并不容易： 1. 如果一个 GPU kernel 在循环中被调用，我们不应该在循环的每次迭代中都进行复制。 2. ……但我们也应该让数组副本的存活时间尽可能短。 我对这个问题的解决方案是稍微调整一下 `#[cpu_function]` 的含义：*输入*和所有中间结果都位于 CPU 内存中，但*结果*将位于 GPU 内存中，通过在函数结尾处的复制来实现。这基本上意味着编译器可以继续在很大程度上假定所有数组都位于 GPU 内存中，因为对于除标记了 `#[cpu_function]` 之外的所有内容，情况依然如此。 ……当然，那些函数接着就变成了问题所在。GPU 后端假定所有 `iota` 和 `replicate` 操作（这些是 IR 原语）都会把结果产生在 GPU 内存中，而在 `#[cpu_function]` 内部并不是这样。修复这个问题不算太难（https://github.com/diku-dk/futhark/blob/42808b162f6b85b76f89ca7eed687d27e952ef8f/src/Futhark/CodeGen/ImpGen/GPU.hs#L267-L270），但类似的问题很可能还潜伏在别处。 最后还有一个问题：当你 `map` 一个 `#[cpu_function]` 时应该发生什么。我武断地决定在这种情况下忽略 `#[cpu_function]`，生成一个普通的并行提升函数（https://futhark-lang.org/blog/2026-07-31-full-flattening.html），但也可以认为这种情况应该以某种方式产生一个顺序函数。 ### 我们目前的状态 虽然 `#[cpu_function]` 是一个相当粗糙的补丁，但它确实解决了眼前的问题：我们现在可以用非常高效的顺序代码进行图着色，*同时*还能在同一个 Futhark 程序中进行并行操作。特别是，图着色现在只占整体运行时间的很小一部分，运行时间主要被随后的数值计算所主导。`#[cpu_function]` 这个技术比较粗粒度，因为调用这些函数涉及开销较大的复制，但由于这些是*批量*复制而非逐元素复制，开销并没有那么大。对编译器所做的修改相对较小，而且大部分可以归类为 bug 修复。即使到了那个光辉的、有原则的未来——Futhark 编译器能够自动确定数组应该存放在哪里——我们很可能也会把这个 attribute 保留下来，作为一个有用的提示。