性能优化及其错误实践：以 SIMD 卷积为例

Performance optimization, and how to do it wrong

Source | HN Comments

文章分享了作者使用 SIMD 指令优化卷积运算的经验，并重点介绍了优化过程中遇到的问题和解决方案。作者首先尝试了多种性能分析工具，但效果不佳。通过简化代码，发现分支和寄存器溢出是性能瓶颈。作者通过循环展开、编译时单态化等技术，最终解决了问题，并实现了显著的性能提升。文章强调了现代 CPU 性能优化的复杂性，以及内联、分支等因素对性能的影响，并建议通过简化问题、学习汇编等方式进行调试。

[正文]

/blog 2025-03-02

| 性能优化及其错误实践 |

最近我尝试使用 SIMD 指令优化卷积，但原本以为简单的任务最终花费了我数天时间，各种问题接踵而至。事后看来，其中一些问题可以理解，但另一些则完全令人困惑。虽然具体示例是直接卷积，但这些考虑因素几乎适用于任何具有热循环的代码。

注意： 这篇博文主要根据记忆编写，因为我没有保留所讨论代码的每个版本。基准测试中的值是对实际值的粗略重现。

背景

我在 burn 上工作，最近想优化 burn-ndarray CPU 后端上的直接卷积。对于卷积，你需要将一个二维内核在输入特征图上移动，并对所有输入通道中的所有值求和。这对于每个输出通道都会重复。输入可以有 padding 像素的零填充围绕实际数据，并且内核可以以步进方式移动（即一次移动两个像素）。有很多算法具有不同的权衡，但我决定使用直接卷积，因为它们没有内存开销，并且在正确实现时仍然非常高效。基本轮廓是你有很多嵌套循环，一些边界检查和一个非常频繁执行的 fused-multiply-add (FMA) instruction。我的初始实现看起来像这样 1（简化）：

iter_range_par!(0, batch_size * oc_blocks).for_each(|k| {
 for oh in 0..out_height {
  for ow_block in 0..ow_blocks {
   seq!(N in 0..8 {
    let mut acc~N = bias;
});
   for ic in 0..in_channels {
    for kh in 0..k_height {
     if !in_bounds_h {
      continue;
}
     for kw in 0..k_width {
      let f0 = vload(&weights[[ic, kh, kw, oc]]);
      seq!(N in 0..8 {
       if in_bounds_w {
        let i~N = splat(x[[ic, ih, iw + N]]);
        acc~N = E::vmuladd(simd, i~N, f0, acc~N);
}
});
}
}
}
   seq!(N in 0..8 {
    if ow + N >= out_width {
     continue;
}
    vstore(&mut out[[ow + N, oc]], acc~N);
});
}
}
});

在这个实现中，我使用了一些技术。除了 SIMD 加载和 fmadds 之外，我还使用了来自 this paper 的优化循环顺序和寄存器阻塞（使用 seq 宏）技术。我完成了实现，执行了一个基准测试，结果...它更慢了。实际上，比幼稚的非向量化 2 实现慢了两倍以上（~670ms vs ~300ms）。

开始调查

为了做到这一点，我尝试使用各种性能分析器，cargo-flamegraph，samply，并且在经过大量的绝望之后，使用了 AMD μProf。在花费了几天时间试图从这些分析器中获得有用的信息（并使 μProf 能够工作）之后，我意识到它并没有给我带来任何进展。火焰图和热点似乎根本没有任何意义。

那么下一步是什么？

简化问题

好吧，我所有的分析尝试都没有成功。因此，让我们尝试将代码简化为基准测试中实际需要的代码。基准测试使用未填充、非步进、非膨胀和未分组的卷积，因此我剥离了所有填充检查和所有步幅/膨胀计算 - 它更快了，但仍然很慢。

还剩下一个分支需要消除：寄存器阻塞循环中的边界像素检查。

为了检查，我缩短了循环，只考虑高达 8 的最后一个倍数的像素。这会产生不正确的结果，但应该有助于调试性能。

let ow_blocks = out_width / ow_b // changed from `div_ceil`
//...
for ow_block in 0..ow_blocks {
 seq!(N in 0..8 {
  let mut acc~N = bias;
});
 for ic in 0..in_channels {
  for kh in 0..k_height {
   for kw in 0..k_width {
    let f0 = vload(&weights[[ic, kh, kw, oc]]);
    seq!(N in 0..8 {
     let i~N = splat(x[[ic, ih, iw + N]]);
     acc~N = E::vmuladd(simd, i~N, f0, acc~N);
});
}
}
}
 seq!(N in 0..8 {
  vstore(&mut out[[ow + N, oc]], acc~N);
});
}

执行基准测试后，现在的代码在单线程上比旧代码在多线程上_快得多！_

Benchmarking - conv2d-input_16x512x512_weight_16x3x3_stride_1
―――――――― Result ―――――――――
 Timing   full
 Samples   40
 Mean    205.12ms
 Variance  69.420µs
 Median   203.24ms
 Min     201.12ms
 Max     207.23ms
―――――――――――――――――――――――――

问题似乎是寄存器溢出（作为以前专注于 GPU 的人，我很震惊地发现现代 CPU 只有 16 个寄存器）和分支过多之间的混合。这就是为什么性能分析器没有给我带来任何有用的结果的原因。现代 CPU 中的分支实在太复杂了，无法用性能分析器热点有意义地表示。这可能是本文最大的收获：分支比你想象的要糟糕得多，因为 CPU 每个周期只能预测一个 3 分支 4。循环中的单个 if 语句足以阻止该周期内解码任何进一步的指令。由于最佳性能需要每个周期 2 个 FMA 指令（它们需要 1 个周期，延迟为 5 个周期，并且 Zen 4 有 2 个 FMA 单元），因此在每个指令上都有一个分支会大大降低性能。这在 Zen 5 上可能会有所不同，但请记住，除了需要预测的 ow 边界检查之外，我们仍然有其他分支。因此，实际上比 50% 的性能更糟糕。

好的，现在我们有了一个好的起点。让我们开始添加东西，看看性能何时开始变差。

首先，我们需要处理寄存器阻塞代码之后的剩余像素。为此，我们将使用一种技术，我们将在接下来的几次中更多地使用：

为什么有一个循环，而你可以有两个？

我之前提到过，我缩短了循环，只处理寄存器阻塞因子的干净倍数。因此，处理这些剩余像素的方法是... - 添加另一个循环。

我们添加第二个未阻塞循环，该循环从第一个循环的末尾开始，一直运行到特征图的边缘。由于它没有展开，因此我们不需要添加任何边界检查。

for ow_block in 0..ow_blocks {
 //...
}
for ow in ow_blocks * 8..out_width {
 let mut acc = bias;
 for ic in 0..in_channels {
  for kh in 0..k_height {
   for kw in 0..k_width {
    let f0 = vload(&weights[[ic, kh, kw, oc]]);
    let i0 = splat(x[[ic, ih, iw + N]]);
    acc = E::vmuladd(simd, i0, f0, acc);
}
}
}
 vstore(&mut out[[ow + N, oc]], acc~N);
}

运行基准测试，它仍然很快 - 耶！运行两个循环比在每次迭代中检查我们是否在边界内要高效得多。

添加回其他变量

为了添加回填充、步幅和膨胀，而不会再次降低性能，我决定使用 compile-time monomorphization 来消除常见的零填充和/或单位步幅/膨胀情况。因此，我使用在原始卷积实现中看到的技术，由 Justin Moore 添加，以启用单位步幅卷积的自动向量化。通过添加一个 if 语句来检查步幅和膨胀是否都是 1，我们允许编译器将此值 constant propagate 到该分支中。内部循环被提取到单独的内联函数中。此技巧允许在原始的非 SIMD 实现中自动向量化非步进卷积。

for ow_block in 0..ow_blocks {
 let ow = ow_block * ow_b + ow_start;
 // Trick the compiler into constant propagating stride/dilation
 #[allow(clippy::if_same_then_else)]
 if (1, 1, 1, 1) == (stride_h, stride_w, dilate_h, dilate_w) {
  conv2d_inner(
   simd, &x, &weights, &mut out, bias, oh, ow, oc, ic_off, stride_h, stride_w,
   dilate_h, dilate_w, k_height, k_width, pad_h, pad_w,
);
} else {
  conv2d_inner(
   simd, &x, &weights, &mut out, bias, oh, ow, oc, ic_off, stride_h, stride_w,
   dilate_h, dilate_w, k_height, k_width, pad_h, pad_w,
);
}
}

通过 const generic bool 添加回填充支持，该 const generic bool 将填充设置为 0。这允许编译器再次常量传播它。

fn run_conv2d</*...*/, const PAD: bool>(/*...*/){
 ///...
 if !PAD {
  pad_h = 0;
  pad_w = 0;
}
}

简单又容易。

让我们再次运行基准测试！

Benchmarking - conv2d-input_16x512x512_weight_16x3x3_stride_1
―――――――― Result ―――――――――
 Timing   full
 Samples   40
 Mean    8.136s (+3868%)
 Variance  75.115µs
 Median   8.042s (+3861%)
 Min     8.020s (+3890%)
 Max     8.341s (+3929%)
―――――――――――――――――――――――――

哦。哦，天啊。发生了什么？

当编译器出错时

为了解释刚刚发生的事情，我需要添加另一个之前没有提到的小背景细节。为了使用现代 SIMD 功能，代码使用带有 pulp 的运行时功能选择。它的工作方式是，pulp 使用类似 #[target_feature(enable = "avx2")] 的东西来注释一个函数，基于可用的功能。这告诉编译器允许使用 avx2 功能，即使目标通常不包括 avx2。但是，只有内联函数才会启用这些功能，而非内联函数调用则不会（这是伏笔）。

这就是 samply 开始变得真正有用的地方。运行它允许我查看每个函数的汇编并找到热点。这一次它们实际上是有意义的！ samply 告诉我，我将所有时间都花在了调用 FMA 的行和 FMA 本身上。所以我看了一下汇编 - 哦，不！

movapsxmm6,xmmword[rsp+0x2e0]
movapsxmm7,xmmword[rsp+0x2f0]
movapsxmmword[rsp+0x170],xmm15
movapsxmmword[rsp+0x160],xmm14
movapsxmm0,xmmword[rsp+0x180]
movapsxmmword[rsp+0xb0],xmm0
movapsxmm0,xmmword[rsp+0x190]
movapsxmmword[rsp+0xa0],xmm0
movrcx,rbx
learbx,qword[rsp+0x320]
movrdx,rbx
movr8,rdi
movr9,r14
call0x4edb0

为什么它使用 SSE 寄存器而不是 AVX 寄存器？
为什么它把它们放在堆栈上？
它使用 call 来执行 _mm256_fmadd_ps intrinsic 是什么原因？

事实证明：这些事情是相关的。接下来是我认为对这里发生的事情的一个非常准确的猜测。

你看，编译器对内联函数的大小有限制。#[inline(always)] 告诉编译器忽略大小限制，并且_几乎_我所有的函数都被标记为 #[inline(always)]。但是，_最外层_的函数没有。

这些是我认为接下来发生的步骤：

添加这些内联分支导致函数的大小超过了 Rust 的内联限制
Rust 大纲（这是一个词吗？）从 pulp 的包装函数中划分出顶层函数，该函数标有 #[target_feature]
编译器现在将我的函数视为常规函数，导致它回退到默认功能集 (x86-64-v1)
这意味着 256 位寄存器在我的函数中不再可用。
由于 _mm256_fmadd_ps 是一条需要 256 位寄存器的 AVX2 指令，因此编译器现在必须动态调用它并通过堆栈传输数据。这很慢。非常慢。

我有点不确定最后一步，也许一些更了解编译器内部原理的人可以启发我内在函数不再内联的实际原因。

幸运的是，解决方案比这一系列事件简单得多：使用 #[inline(always)] 注释顶层函数。

vbroadcastssymm9,dword[r12+r11*1]
learcx,qword[r11+r12*1]
vfmadd231psymm6,ymm8,ymm10

好多了。基准测试呢？

Benchmarking - conv2d-input_16x512x512_weight_16x3x3_stride_1
―――――――― Result ―――――――――
 Timing   full
 Samples   40
 Mean    230.12ms
 Variance  69.420µs
 Median   232.24ms
 Min     224.12ms
 Max     236.23ms
―――――――――――――――――――――――――

太棒了！

完成优化

对于未填充的卷积，性能很好，但对于填充的卷积来说仍然会很差，因为我们需要_在每次循环迭代中_检查我们是否在填充中。为了解决这个问题，我们可以使用我们之前用于 out_width 的相同技术：所有距离边缘超过 padding 的像素保证始终在边界内，因此我们可以从 padding_h 到 out_height - padding_h 和 padding_w 到 width - padding_w 运行一个没有边界检查的循环，然后为_确实_进行边界检查的边界像素运行第二个循环。这比检查每个像素要快得多，因为大多数像素始终在边界内。

if (pad_h, pad_w) != (0, 0) {
 let v_borders = (0..pad_h)
  .chain(out_height.saturating_sub(pad_h)..out_height)
  .cartesian_product(0..out_width);
 let h_borders = (0..out_height)
  .cartesian_product((0..pad_w).chain(out_width.saturating_sub(pad_w)..out_width));
 for (oh, ow) in v_borders.chain(h_borders) {
  //...
}
}

最后的想法

现代 CPU 很奇怪，性能并不总是显而易见的。内联很脆弱，添加单行代码可能会完全改变程序的编译方式，甚至没有注意到。性能分析器并不总是有帮助，特别是当您的问题比使用慢速函数或过于频繁地分配内存等问题复杂时。我的建议是尝试找出事情何时开始变坏，并学习一些基本的汇编，这样你就可以发现诸如堆栈加载/存储过多之类的事情（这表明寄存器溢出）。

我希望这有助于其他人比我更快地处理他们的性能问题。

再次感谢 samply。即使性能数据没有太多意义，能够轻松查看任何给定函数的汇编也非常有用。

此实现的最终代码可以在 here 找到。为了好玩，这是所有优化后的最终基准测试，使用多线程：

Benchmarking - conv2d-input_16x512x512_weight_16x3x3_stride_1
―――――――― Result ―――――――――
 Timing   full
 Samples   40
 Mean    42.731ms
 Variance  5.115µs
 Median   42.906ms
 Min     38.162ms
 Max     47.554ms
―――――――――――――――――――――――――

我最初使用了一个累加器数组，但很快发现 Rust 无法将常量数组索引优化为一组寄存器访问，而是从堆栈访问该值。为了简洁起见，我省略了这个版本。 ↩
实际上它是使用我们稍后会提到的技巧自动向量化的 ↩
在 Zen 5 上是两个 ↩
感谢 /u/caelunshun 的纠正 ↩