Rust 中的安全 SIMD，即使内部也安全

# 即使在内部也能实现安全的 Rust SIMD 编程 来源：https://shnatsel.medium.com/safe-simd-in-rust-even-on-the-inside-c6f1ff381828 Sergey "Shnatsel" Davidoff (https://shnatsel.medium.com/?source=post_page---byline--c6f1ff381828---------------------------------------) *Rust 的 SIMD 抽象 (https://shnatsel.medium.com/the-state-of-simd-in-rust-in-2025-32c263e5f53d）并不像我希望的那样安全——直到现在。* 众所周知，直接使用原始 SIMD 内联函数（intrinsics）非常令人不快。你希望写 `a + b`，而不是下面这种怪物： ```rust unsafe { #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "avx2"))] _mm256_add_ps(a, b) #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "sse", not(target_feature = "avx2")))] _mm_add_ps(a, b) #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] vaddq_f32(a, b) } ``` 看看它，多么丑陋！而且整个东西还包裹在 `unsafe` 中！这还算是一个简化的例子。它仍然没有处理： - 其他常见平台：AVX-512、32 位 ARM、WebAssembly - 没有 SIMD 的平台或像 RISC-V 这样的冷门平台 - 实际将数据（如 `&[f32]`）加载到每个内联函数接受的格式中 - 为正在运行的 CPU 选择最佳实现 幸运的是，Rust 提供了许多 SIMD 抽象 (https://shnatsel.medium.com/the-state-of-simd-in-rust-in-2025-32c263e5f53d），它们为你处理了所有这些，让你只需写 `a + b` 即可。但有一个小问题：它们内部仍然充满了 `unsafe`。这并没有消失，只是被隐藏了。大量 `unsafe` 代码潜伏在表面之下，偶尔会搞砸事情 (https://github.com/servo/pathfinder/issues/588)（比如这里 (https://github.com/wingertge/macerator/issues/31）和这里 (https://github.com/sarah-quinones/pulp/issues/28））。或者说，它们曾经是这样的——直到现在。 ## 为什么我们需要 `unsafe`？ 在很长一段时间里，你无法绕开将每个内联函数（如 `_mm256_add_ps`）的调用包裹在 `unsafe` 中的做法，因为当 CPU 不支持某个指令时调用它是非法的。所以你*必须*有某种机制来跟踪每个内联函数需要哪些指令，以及你能访问哪些指令，并交叉引用它们来决定调用给定函数是否安全。手工做的话很繁琐，用代码生成器做则很复杂，而且总是容易出错，需要在每个内联函数周围放上 `unsafe`。 这种情况在 Rust 1.87 版本中发生了变化，编译器开始自行跟踪所需的指令集，因此你可以这样写： ```rust #[target_feature(enable = "avx2")] fn add_avx2(a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) // 这是一个 avx2 内联函数 } ``` 看，没有 `unsafe` 了！……但还没完。你仍然不能直接用 `a + b` 来写。你能做到的最好程度是： ```rust unsafe { add_avx2(a, b) } ``` 这只是将 `unsafe` 向上移了一层。现在你可以在标注了正确 `#[target_feature]` 的函数内部调用内联函数，但链中仍然必须在某处有一个 `unsafe`。 另一个问题更根本。你不能在类型 `+` 操作的实现上放置 `#[target_feature]`，因为 `+` 必须始终可用。所以通过这种机制，我们仍然无法实现 `a + b`。 ## 引理：CPU 特性令牌 要理解最终的解决方案，你首先需要了解 CPU 特性检测是如何工作的。通常，运行时检查 AVX2 这样的 CPU 特性是通过 `is_x86_feature_detected!("avx2")` 完成的。但我们肯定不希望每次将两个数相加时都运行这个检查——那会完全拖垮性能。我们希望*只检查一次*，然后向编译器证明从那时起使用 AVX2 指令是安全的。 我们可以通过一个*不可伪造的令牌*将此证明编码到类型系统中：一个具有私有内部字段的零大小类型。获得这个令牌的*唯一*方式是调用一个执行 CPU 特性检查的函数。如果检查通过，该函数就会给你这个令牌： ```rust pub struct Avx2(()); fn detect_avx2() -> Option<Avx2> { if is_x86_feature_detected!("avx2") { Some(Avx2(())) } else { None } } ``` 由于它是零大小类型，传递这个令牌没有运行时开销。它纯粹作为编译时的证明存在。要点是：只要你拥有一个 `Avx2` 结构体的实例，你就可以确信系统上 AVX2 指令是可用的。 ## 关键洞见 编译器不知道，但这个函数可以安全调用： ```rust #[target_feature(enable = "avx2")] fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } ``` 你*只有*在拥有 Avx2 令牌时才能调用这个函数，而只有系统上 AVX2 指令可用时你才能获得该令牌。如果我们能向编译器解释这是有效的（使用 `unsafe`），那么我们可以只写一次 `unsafe` 并在任何地方重用它。 我们需要的是一个可以安全调用的宏： ```rust with_avx2!( fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } ) ``` 但在幕后展开成这样： ```rust fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { // 安全性：根据令牌，Avx2 可用， // 并且我们验证了内部函数不是 `unsafe fn` unsafe { inner(token, a, b) } #[target_feature(enable = "avx2")] fn inner(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } } ``` 现在如果你使用一个不在 AVX2 中的内联函数，**编译器会拒绝它！** 我们刚刚成功地为 SIMD 内联函数提供了一个安全的编程接口，而无需任何特制的目标特性跟踪！即使内部仍然有一个 `unsafe` 块，它也被封装在一个健全的 API 中，所以你无法滥用它导致内存安全错误。从这个意义上讲，它就像 `println!`，安全地抽象了不安全的代码。 这样，你只需要审查和审计*这一个宏*，而不是成百上千个特制的 `unsafe` 块。而且我们在实现中*可能*搞砸的地方只有两处： 1. 将令牌映射到错误的 `#[target_feature]` 2. 允许从安全上下文中调用一个 `unsafe fn` 这两种失败模式都很容易检查。 所以现在我们可以不用 `unsafe` 就调用 `add_avx2(token, a, b)`，但这仍然没有让我们达到 `a + b`。那我们要*如何*解决呢？ ## 泛型来援救 我们不能用 `#[target_feature]` 注释 `a + b` 的实现，因为它必须可以从任何地方安全调用。而且我们不能向函数传递令牌，因为它只接受 `a` 和 `b`，而不接受 `token`。但即使我们能做到，也会使 API 非常丑陋。我们希望 `a + b` 始终有效，并自动使用最佳的 SIMD 指令，用户永远不用与令牌打交道。 我们可以一次性用泛型解决这两个问题：通过定义一个在可用指令集上泛型的 `f32x8` 类型，我们可以在其上实现加法，既在内部走私一个令牌，*又*为每个 SIMD 指令集创建独立的实现！看起来像这样： ```rust pub trait Level {} #[derive(Clone, Copy)] pub struct Avx2(()); impl Level for Avx2 {} pub struct f32x8<L: Level> { // 为简单起见，本例中用数组作为后备。 // 生产代码中，我们使用针对该等级的原生 SIMD 类型。 data: [f32; 8], // 走私的令牌！ token: L, } /// `a + b` 对 Avx2 的实现 impl std::ops::Add for f32x8<Avx2> { type Output = Self; fn add(self, rhs: Self) -> Self::Output { // （类型转换省略） // 使用 Avx2 令牌调用我们的安全包装器 let result = add_avx2(self.token, self, rhs); Self { data: store_m256(result), token: self.token } } } ``` 然后我们可以轻松地让它适用于任何其他指令集，或者当根本没有 SIMD 时： ```rust #[derive(Clone, Copy)] pub struct NoSimd(()); impl Level for NoSimd {} /// `a + b` 在没有 SIMD 时的实现 impl std::ops::Add for f32x8<NoSimd> { type Output = Self; fn add(self, rhs: Self) -> Self::Output { let result = std::array::from_fn(|i| self.data[i] + rhs.data[i]); Self { data: result, token: self.token } } } ``` 我们一次性解决了安全*和*运行时指令选择的问题！再添加一个便利函数，它给出系统上最好的 `Level`，你就得到了几乎完美的 SIMD API！ ## ABI 有话要说 不幸的是，在编写 `a + b` 并将其降级为 SIMD 指令方面存在一个根本问题：**函数调用开销**。调用一个函数不是免费的，但相当便宜——只是少数几条 CPU 指令。但*少数*指令远比我们刚用于实现加法的*一条*指令要多！所以如果中间有一个函数调用，加法性能就会暴跌。而性能正是使用 SIMD 的全部意义所在！ 编译器通常通过内联 (https://matklad.github.io/2021/07/09/inline-in-rust.html) 很好地消除这种开销。它基本上将你调用的函数的实现复制粘贴到调用它的函数中，因此不再有函数，也不再有过开销。但是 `#[target_feature]` 注释会打乱这个机制。编译器不能内联一个具有 `#[target_feature]` 注释的函数到一个没有它的函数中，因为所需的特性在其中不可用！猜猜什么不能有 `#[target_feature]` 注释？是的。 那么*我们*要如何让 `a + b` 与 SIMD 一起工作呢？ ## 内联它，用火焰内联它！ 我们不能将 `#[target_feature]` 放在实现 `a + b` 的函数上，但*可以*将它放在调用 `a + b` 的函数上！然后我们可以利用内联将 `a + b` 的实现复制到调用它的函数中，这样它最终就处于一个 `#[target_feature]` 的上下文中。 所以调用链看起来像这样： ```rust #[target_feature(enable = "avx2")] fn do_stuff() { // TODO: 一些计算 c = a + b; // TODO: 更多计算 } // 它调用... #[inline(always)] // 函数体将被复制到调用者中 fn add(self, rhs: Self) -> Self::Output { // 使用 Avx2 令牌调用我们的安全包装器 add_avx2(self.token, self, rhs); // return 语句省略 } // 它又调用... #[inline] // 如果可行，函数体将被复制到调用者中 #[target_feature(enable = "avx2")] fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } ``` 这可行。你可以在正确的地方添加这些注释，并在 SIMD 层级上进行抽象，而且你甚至不需要宏 (https://docs.rs/fearless_simd/latest/fearless_simd/trait.Simd.html#tymethod.vectorize）。问题在于，每当你对 SIMD 类型调用 `a + b` 时，你都必须在一个带有 `#[inline(always)]` 或 `#[target_feature]` 的函数中进行，否则代码仍然可以编译，但性能会暴跌。想亲自试试？打开这个示例 (https://rust.godbolt.org/z/xo1fEfqnT)，移除 `#[target_feature]`，看看生成的汇编变成多么恐怖的东西。我不确定对此能做什么。对于任何用 SIMD 实现 `a + b` 的方法，这个限制似乎相当根本。结构体目标特性 RFC (https://shnatsel.medium.com/github.com/rust-lang/rfcs/pull/3525） 解决了 `add_avx2(token, a, b)` 和 `add(token, a, b)` 的问题，但我还没看到实现 `a + b` 的路径。 ## 那我们处于什么位置？ 尽管所有 SIMD 代码都固有的内联问题，我们仍然设法以一种惊人的低水平 `unsafe` 代码提供了非常令人愉快的 SIMD 抽象，这是前所未有的。你可以找到这些想法的生产版本，就在 `fearless_simd` (https://crates.io/crates/fearless_simd) v0.5 中，**现已可用**，就在你附近的包注册表中！下面是一个小例子 (https://github.com/linebender/fearless_simd/blob/d2411665f0726cd6c09bd6fd98af4af18e3c1778/fearless_simd/examples/sigmoid.rs)，看看它们在生产中如何组合在一起。用于实现它的宏也是公开暴露的，所以你可以轻松地混合和匹配 (https://github.com/linebender/fearless_simd/blob/e76b820154dc09f41bca8302670424d8ca079219/fearless_simd/examples/srgb.rs) 高级操作（如 `a + b`）和平台特定的内联函数，以充分利用硬件。 `fearless_simd` 中有不止一个 `unsafe` 块，因为它还提供了 `safe_unaligned_simd` (https://crates.io/crates/safe_unaligned_simd) 包的功能，但这也是以比原始代码**显著更低** (https://github.com/okaneco/safe_unaligned_simd/issues/51) 的 `unsafe` 代码量完成的。 对我来说，使用高级 SIMD 抽象的障碍一直是它们带来的大量 `unsafe`。这很可怕，也很难证明其合理性。但现在 Rust 中的 SIMD 可以真正地无所畏惧了。 ## 致谢 我很震惊我是第一个将其投入生产的人（据我所知），因为我肯定不是第一个想到这一点的人。CPU 特性令牌是一个古老而常见的想法。`pulp` (https://crates.io/crates/pulp) 包已经使用它们多年，但它们依赖于手写的围绕内联函数的 `unsafe` 包装器，并且偶尔会把它们弄错 (https://github.com/sarah-quinones/pulp/issues/28)。使用泛型生成多个实现也是一个古老的想法。它是 8 年前最初的 `fearless_simd` 概念 (https://linebender.org/blog/towards-fearless-simd/) 的一部分。更早的 `simdeez` (https://crates.io/crates/simdeez) 包似乎也使用了类似的东西 (https://docs.rs/simdeez/1.0.8/simdeez/)。 将令牌与单个委托给 rustc 的安全包装器结合的关键洞见也不是我独有的。就在 `fearless_simd` 包的上下文中，Raph Levien 已经尝试过它 (https://github.com/linebender/fearless_simd/commit/81b6ab4c8dd6f25064f539c49818c45c2f686815)，而 Daniel McNab (https://github.com/DJMcNab) 在几个月前创建了一个比我更精细的实现 (https://github.com/linebender/fearless_simd/pull/108)。Daniel 的方法允许对每个单独的 CPU 特性进行细粒度跟踪，而不是像 `fearless_simd` 那样只使用少数固定的 CPU 特性级别。它更具表现力，但代价是复杂性，而且他的方法从未被合并，因为没有其他维护者站出来审查它。我仍然希望有一天它能作为独立的包发布。 感谢 Daniel 和 Laurenz Stampfl (https://github.com/LaurenzV) 审查我对 `fearless_simd` 的所有 PR，它们很大，快速的审查真的非常感谢！