您能相信编译器会优化您的代码吗？ - 详解

大家好！我是大聪明-PLUS！

对于 SIMD 的工作原理，有三个层次的理解（至少我目前处于第 3 级）：

1. 编译器很聪明！它们会自动矢量化所有代码！
2. 编译器很笨，自动矢量化很脆弱，很容易因为无关的代码修改而破坏它。手写特定的 SIMD 指令总是更好的选择。
3. 手写 SIMD 确实很困难——你必须为每种处理器架构编写不同的代码。此外，你可能意识到编译器用汇编语言编写标量代码比你写得更好。你凭什么认为自己能在 SIMD 中打败编译器呢？毕竟 SIMD 指令和限制更加奇怪。编译器只是工具而已。如果代码是以可矢量化的形式编写的，它们就能可靠地进行矢量化。

我最近从 2 级升到了 3 级，注意到编译器使用的模型在我的脑海中闪过一个念头。在这篇文章中，我想解释一下适用于优化 Rust 或 C++ 等静态语言的编译器的通用结构。之后，我会将此结构应用于自动矢量化。

我还没有从事过生产优化编译器后端的工作，所以以下内容不会绝对完美，但这些模型绝对有用，至少对我来说！

❯ 通过编译器的眼睛

我们首先要解决的难题是理解编译器如何看待代码。


优化的单位是函数。我们来看一个简单的函数，如下所示：

fn sum(xs: &[i32]) -> i32 {
  let mut total = 0;
  for i in 0..xs.len() {
    total = total.wrapping_add(xs[i]);
  }
  total
}

在某些伪编译器中它看起来像这样：

fn sum return i32 {
  param xs_ptr: ptr
  param xs_len: size
  local total: i32 = 0
  local i: size = 0
  local x: i32
loop:
  branch_if i >= xs_len :ret
  load x base=xs_ptr offset=i
  add total x
  add i 1
  goto :loop
ret:
  return total
}

这里最重要的特征是存在两种“实体”：

第一种是程序内存，大致可以理解为一个字节数组。编译器通常不太理解内存的内容，因为它在所有函数之间共享，而不同的函数对内存内容的解释可能不同。第二种是局部变量。局部变量不是字节——它们是整数，它们遵循编译器已经理解的数学运算。 例如，如果编译器遇到如下循环：

param n: u32
local i: u32 = 0
local total: u32
local tmp
loop:
  branch_if i >= n :ret
  set tmp i
  mul tmp 4
  add t tmp
  goto :loop
ret:
  return total

他意识到在每个循环中 tmp 都会变成 i*4，并将代码优化为：

param n: u32
local i: u32 = 0
local total: u32
local tmp = 0
loop:
  branch_if i >= n :ret
  add t tmp
  add tmp 4  # replace multiplication with addition
  goto :loop
ret:
  return total

之所以有效，是因为这里的一切都只是数字。如果我们执行相同的计算，但所有数字都位于内存中，编译器将很难确定转换是否正确。如果 n 和 total 的存储空间实际上被覆盖了怎么办？如果 tmp 与当前函数中根本不存在的内容相交怎么办？

然而，数学局部变量的世界和内存字节的世界之间存在一座桥梁——加载和存储指令。加载指令获取内存中的一系列字节，将这些字节解释为整数，并将该整数存储在局部变量中。存储指令执行相同的操作，但顺序相反。通过将内存中的某些内容加载到本地文件中，编译器可以深入了解其中的内容。这样，编译器无需跟踪内存的全部内容。它只需要检查在特定时刻是否需要加载某些内容。因此，编译器对事物的理解程度并不那么高，它一次只能推理一个函数，并且只能推理该函数中的局部变量。

❯ 将编译器的面孔塞进代码中

编译器是短视的。这可以通过为编译器提供更多上下文来解决，这是两个主要优化的目标。

第一个主要优化是“内联”。它替换了函数调用的“主体”。这样做的好处不是消除函数调用的开销；这相对来说很小。重要的是，调用者和被调用者的本地数据现在位于同一作用域内，编译器可以同时优化它们。

让我们再看一下这段 Rust 代码：

fn sum(xs: &[i32]) -> i32 {
  let mut total = 0;
  for i in 0..xs.len() {
    total = total.wrapping_add(xs[i]);
  }
  total
}

表达式 xs[i] 实际上是一个函数调用。索引操作会在访问数组元素之前执行边界检查。将其插入 sum 函数后，编译器会发现这是死代码并将其删除。

如果你查看常见的优化函数，就会发现它们通常看起来像是在删除一些愚蠢到没人会写的东西，因此目前尚不清楚这类优化是否值得实现。但事实上，在函数内联之后，会出现很多愚蠢的事情，因为函数通常处理一般情况，而在特定位置，已经存在足够多的条件来“拒绝”某些极端情况。

第二个主要的优化是 SRA（聚合的标量替换）。这是对我们已经见过的“让我们使用 load 来避免推理内存，而是推理本地数字”这一想法的概括。

如果你有一个像这样的函数：

fn permute(xs: &mut Vec) {
  ...
}

编译器很难推理出这一点。它接收一个指向包含复杂结构的内存的指针，因此推理该结构的演变非常困难。编译器可以做的是从内存中加载该结构，用一组标量局部变量替换聚合：

fn permute(xs: &mut Vec) {
  local ptr: ptr
  local len: usize
  local cap: usize
  load ptr xs.ptr
  load len xs.len
  load cap xs.cap
  ...
  store xs.ptr ptr
  store xs.len len
  store xs.cap cap
}

这使得编译器能够再次推理。SROA 类似于内联，但针对的是内存而不是代码。

❯ 可能性与不可能性

使用编译器的这个思维模型：

1. 优化各项功能；
2. 可以使用嵌入；
3. 擅长注意局部变量之间的关系并基于它们重建代码；
4. 能够对记忆进行有限的推理（即决定何时可以安全地加载或保存）；

我们可以通过讨论零成本抽象来描述哪些代码可以轻松优化，哪些代码无法优化。

为了启用内联，编译器需要知道实际调用的是哪个函数。如果直接调用某个函数，编译器几乎肯定会尝试将其内联。如果调用是间接的（通过函数指针或虚函数表），编译器在大多数情况下将无法内联该函数。即使是间接调用，编译器有时也可以推断指针值并取消虚拟化调用，但这取决于其他地方的成功优化。

这就是为什么在 Rust 中，每个函数都有一个唯一的零大小类型，不需要向环境进行“表示”。这确保了编译器始终能够内联代码，并将这种抽象的“成本”降低到零，因为任何优秀的优化编译器都会否定它。

高级语言可能更喜欢始终使用指针表示函数。实际上，最终代码在许多情况下都具有同等的可优化性。但是源代码不会指示这是否是可优化的情况（指针仅在编译时清除）或真正的动态调用。在 Rust 中，保证可优化和潜在可优化之间的区别反映在源语言中：

// Compiler is guaranteed to be able to inline call to `f`.
fn call1(f: F) {
  f()
}
// Compiler _might_ be able to inline call to `f`.
fn call2(f: fn()) {
  f()
}

因此，第一条规则是使大多数调用静态解析以允许内联。函数指针和动态调度会阻止内联。内存

间接寻址可能会给编译器带来问题。 例如：

struct Foo {
  bar: Bar,
  baz: Baz,
}

Foo 的结构对编译器完全透明。

但是，这里：

struct Foo {
  bar: Box,
  baz: Baz,
}

事情不太清楚。分配给 Foo 内存的内容通常不会迁移到 Bar 内存。同样，在很多情况下，编译器可以基于唯一性推断出块，但这并不能保证。
现在一个好的家庭作业是研究 Rust 的迭代器，并理解它们为什么看起来是这个样子。为什么map的签名和值是这个样子的？

#[inline]
fn map(self, f: F) -> Map
where
  Self: Sized,
  F: FnMut(Self::Item) -> B,
{
  Map::new(self, f)
}

关于内存的另一个重要点是，编译器通常无法更改整体布局。SROA 可以将一些数据加载到一组局部变量中，然后这些变量可以（例如）用“一对指针”替换“指针和索引”的表示形式。但最终，SROA 必须将“指针和索引”重新打包，并将此表示形式存储回内存中。这是因为内存布局在所有函数之间共享，因此函数无法单方面规定自己的规则。

综上所述，这些观察结果提供了对代码工作原理的基本理解。

• 思考一下数据在内存中的布局。编译器在这里几乎不做任何工作，主要只是将字节放置在你指定的位置。为了提高缓存效率，请尽量压缩数据，减少间接访问，并使用常见的访问模式。
• 当编译器能够看到代码时，它们能够更好地理解代码。确保大多数调用在编译时已知并且可以内联；剩下的交给编译器。

❯ SIMD

让我们将对可优化代码的一般理解应用到自动矢量化中。我们将优化一个计算两个字节片段之间最长公共前缀的函数。

一个简单的实现如下：

use std::iter::zip;
// 650 milliseconds
fn common_prefix(xs: &[u8], ys: &[u8]) -> usize {
  let mut result = 0;
  for (x, y) in zip(xs, ys) {
    if x != y { break; }
    result += 1
  }
  result
}

如果您已经对自动矢量化有了一定的了解，或者您查看了汇编输出，您可能会发现，该函数每次只对一个字节进行操作，这比它应该的速度慢得多。让我们解决这个问题！

SIMD 会同时对多个值进行操作。直观地说，我们希望编译器一次比较一堆字节，但我们当前的代码并没有表达这一点。让我们通过一次处理 16 个字节，然后分别处理剩余部分来明确结构：

// 450 milliseconds
fn common_prefix(xs: &[u8], ys: &[u8]) -> usize {
  let chunk_size = 16;
  let mut result = 0;
  'outer: for (xs_chunk, ys_chunk) in
    zip(xs.chunks_exact(chunk_size), ys.chunks_exact(chunk_size))
  {
    for (x, y) in zip(xs_chunk, ys_chunk) {
      if x != y { break 'outer; }
      result += 1
    }
  }
  for (x, y) in zip(&xs[result..], &ys[result..]) {
    if x != y { break; }
    result += 1
  }
  result
}

有趣的是，这已经快了一些，但还不够。具体来说，SIMD 必须并行且相同地处理一个块中的所有值。在上面的代码中，我们有一个 break，这意味着处理第 n 对字节依赖于第 n-1 对。让我们通过禁用快捷方式来解决这个问题。我们将检查整个字节块是否匹配，但我们不关心哪个特定的字节不匹配：

// 80 milliseconds
fn common_prefix3(xs: &[u8], ys: &[u8]) -> usize {
  let chunk_size = 16;
  let mut result = 0;
  for (xs_chunk, ys_chunk) in
    zip(xs.chunks_exact(chunk_size), ys.chunks_exact(chunk_size))
  {
    let mut chunk_equal: bool = true;
    for (x, y) in zip(xs_chunk, ys_chunk) {
      // NB: &, unlike &&, doesn't short-circuit.
      chunk_equal = chunk_equal & (x == y);
    }
    if !chunk_equal { break; }
    result += chunk_size;
  }
  for (x, y) in zip(&xs[result..], &ys[result..]) {
    if x != y { break; }
    result += 1
  }
  result
}

此版本最终支持矢量化，将执行时间缩短了近一个数量级。现在我们可以使用迭代器来压缩此版本。

// 80 milliseconds
fn common_prefix5(xs: &[u8], ys: &[u8]) -> usize {
  let chunk_size = 16;
  let off =
    zip(xs.chunks_exact(chunk_size), ys.chunks_exact(chunk_size))
      .take_while(|(xs_chunk, ys_chunk)| xs_chunk == ys_chunk)
      .count() * chunk_size;
  off + zip(&xs[off..], &ys[off..])
    .take_while(|(x, y)| x == y)
    .count()
}

请注意，这段代码与我们最初的代码有显著的不同。我们不会盲目地依赖编译器优化。相反，我们清楚在这种情况下需要哪些具体的优化，并编写代码来利用它们。

具体来说，对于 SIMD：

• 我们正在改变处理元素块的算法。
• 在每个块内，我们确保没有任何分支并且所有元素都得到平等对待。

❯ 结论

编译器只是工具。虽然有时会发生相当一部分“乐观”的转换，但优化编译器的大部分影响来自于在某些前提条件下保证的优化。编译器是短视的——它们难以推理当前函数之外的代码以及未存储在局部变量中的值。内联和标量聚合替换是两种解决这个问题的优化方法。零成本抽象通过表达保证优化的机会来发挥作用。