编译器GPU性能

让 Futhark 远离 GPU

Futhark通常将数组程序编译为GPU内核，但某些工作负载在许多CPU核心上运行得更好。本文以稀疏雅可比计算为例，介绍当目标平台是CPU而非GPU时，编译器需要做出哪些改变。

FutharkTroels Henriksen6 分钟阅读

中文
中文日本語
复制

今年早些时候，Elias Smedegaard 完成了一篇BSc 论文，研究如何通过自动微分高效计算稀疏 Jacobian 矩阵。对这篇博文来说，不必确切理解它是什么、有什么用——它和自动微分有关，我还没写过这方面的博文。就本文而言，关键的一点是：他的方案中有一部分涉及对一张图进行着色，这张图对应矩阵的稀疏模式。

最优图着色是出了名的 NP-hard 问题，但我们不需要最优着色——只需要够好，并且能用高效算法实现。Elias 找到了一个相当快的贪心算法来做 distance-2 着色（如果你好奇，是这篇论文中的 Algorithm 3.1），但遗憾的是这个算法本质上是串行的。对 Futhark 来说这不是问题，因为 Futhark 支持相当高效的原地更新，问题在于 Elias 的整个程序有两个步骤：

  1. 给图着色。
  2. 用这个着色结果，以大规模并行方式计算 Jacobian。

我们非常希望第二步能利用并行硬件（比如 GPU），同时把第一步放在 CPU 上做。结果发现，Futhark 的一些实现选择让这件事变得很别扭。这篇博文讲的就是这个问题是怎么产生的、一个有原则的解决方案可能是什么样，以及为什么我们最后只是打了个补丁。

问题是怎么产生的

Futhark 的编译模型很简单。数组需要放在内存里，而如果你使用某条 GPU 流水线，那么所有数组都会被放进 GPU 内存。这并不是因为我们的内部表示无法表示 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 实现的 distance-2 着色算法，本质上就是几层嵌套的顺序循环，操作一些表示图和栈的数组。用 GPU 后端编译这段代码，性能相比顺序的 c 后端是灾难性的，轻松差出三个数量级甚至更多。改用 c 后端，性能就非常好（基本等同于用 C 写出来的水平），但这样一来，整个程序的第二部分也只能顺序执行。折中方案是用 Futhark 的 multicore 后端，它会生成并行的 CPU 代码——可我就是想在 GPU 上算那些 Jacobian！

一个有原则的解决方案可能是什么样

Futhark 不支持“分离编译”，即让程序的不同部分使用不同的后端，而且这看起来也是个相当笨拙的解决办法。真正最好的方案就是我上面暗示的：考察数组实际是如何（或者说在哪里）被使用的，据此决定把它放在 CPU 还是 GPU 内存里——如果程序需要，也许两边都放一份。

我们其实已经有一个优化在做类似的事情，由 Philip Børgesen 实现：如果某个简单顺序计算（不是数据！）的结果只在 GPU 上使用，就把它迁移到 GPU。目的是省下通信开销——如果计算本身微不足道，时间全花在来回复制几个数组元素上，那不如把这个顺序计算放进一个单线程的 GPU kernel 里跑，纯粹为了减少通信。遗憾的是，这个优化在这里用不上。首先，distance-2 着色算法太复杂（带循环），触发不了这个优化；就算我们关掉这些检查、强行让它跑在 GPU 上，结果也（毫不意外地）表明，用单个 GPU 线程跑这种循环密集的顺序代码慢得离谱。

理想情况下，我们可以设想一种与这种优化相对应的方案：它搬运的是数据，而不是计算，而且许多底层原理很可能是一样的。尤其要注意的是，必须避免过多的副本，同时也要避免让这些副本无谓地长期存在，因为它们会推高内存占用。我能看出这种优化大致会是什么样子，但也清楚它值得投入比我眼下所能腾出的更多精力和关注。

所以……

我们实际采用的取巧做法

我没有在编译器里搞什么聪明的机制，而是加了一个新的属性，叫 #[cpu_function]。在我最初的设计里，把它加在一个非内联函数上（通常用 #[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] 内部并非如此。修好这个问题倒不算太难，但类似的 bug 很可能还藏在别处。

最后一个问题是，对 #[cpu_function] 做 map 时应该发生什么。我随意决定此时忽略 #[cpu_function]，生成一个普通的并行提升函数，但也有人可以主张，这里应该以某种方式生成一个串行函数。

目前的进展

#[cpu_function] 虽然是个有点扎手的 hack，但确实解决了眼下的问题：现在我们既能用非常高效的串行代码做图着色，又能在同一个 Futhark 程序里做并行的事情。尤其是，图着色如今只占总运行时间的很小一部分，大头是随后的数值计算。#[cpu_function] 这个手法是粗粒度的——调用这些函数意味着昂贵的拷贝——但由于是批量拷贝而非逐元素拷贝，开销并不算大。

编译器需要的改动不算多，而且大多可以归为 bug 修复。即便到了那个理想中的美好未来，Futhark 编译器能够自动判断数组该存在哪里，我们多半也会保留这个属性，把它当作一个有用的提示。

由 DIKU 开发（版权、致谢、刺猬）

来源: Futhark← 返回首页