在 Rust 中实现安全的 SIMD，连内部也不例外

# Rust 中安全的 SIMD——由内而外 | Sergey "Shnatsel" Davidoff 来源：https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/ **Rust 的 [SIMD 抽象](https://shnatsel.medium.com/the-state-of-simd-in-rust-in-2025-32c263e5f53d)一直不如我希望的那样安全。直到现在。** 原始的 SIMD intrinsics 不好用，这不是什么秘密。你想要写 `a + b`，而不是这个怪物： ``` 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]` 这样的数据加载成每个 intrinsic 都能接受的形式 - 为运行它的 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`？ https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/#why-do-we-even-need-unsafe 长期以来，你无法避免将对每个 intrinsic 函数（如 `_mm256_add_ps`）的调用包裹在 `unsafe` 中，因为在当前 CPU 不支持该功能时调用它是非法的。所以你*必须*有某种机制来追踪每个 intrinsic 需要哪些指令集，你又拥有哪些指令，并交叉比对它们来判断调用某个函数是否安全。手动做很繁琐，用代码生成器做又很复杂，而且总是容易出错，每个 intrinsic 周围都需要 `unsafe`。 这在 Rust 1.87 中发生了变化，编译器开始自己追踪所需的指令集，所以你可以这样写： ``` #[target_feature(enable = "avx2")] fn add_avx2(a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) // this is an avx2 intrinsic } ``` 看哪，没有 `unsafe`！...暂时还没有。你仍然无法用这种方式写 `a + b`。你最多只能写： ``` unsafe { add_avx2(a, b) } ``` 这只是把 `unsafe` 上移了一层。现在你可以在标有正确 `#[target_feature]` 的函数内部调用 intrinsics，但在调用链的某个地方仍然必须有 `unsafe`。 另一个问题更加根本。你无法在你的类型的 `+` 实现上标注 `#[target_feature]`，因为 `+` 必须在任何情况下都可用。所以用这种机制我们写不出 `a + b`。 ## 前置知识：CPU 特性令牌 https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/#lemma-cpu-feature-tokens 要理解最终的解决方案，你首先需要了解 CPU 特性检测的工作原理。通常，检查 AVX2 等 CPU 特性是在运行时使用 `is_x86_feature_detected!("avx2")` 完全成的。但我们绝对不想每次加两个数都运行这个检查——那会彻底拖垮性能。我们想要*检查一次*，然后向编译器证明从那以后使用 AVX2 指令是安全的。 我们可以用一个*不可伪造的令牌*将这个证明编码到类型系统中：一个带有私有内部字段的零大小类型（ZST）。获取这个令牌的*唯一*方式是调用一个执行 CPU 特性检查的函数。如果检查通过，函数就把令牌交给你： ``` pub struct Avx2(()); fn detect_avx2() -> Option<Avx2> { if is_x86_feature_detected!("avx2") { Some(Avx2(())) } else { None } } ``` 因为它是零大小类型，传递这个令牌没有运行时开销。它纯粹作为编译时的证明而存在。 关键在于，只要你拥有 `Avx2` 结构体的实例，你就可以确定系统上 AVX2 指令是可用的。 ## 核心洞见 https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/#the-key-insight 编译器并不知道，但这个函数调用是安全的： ``` #[target_feature(enable = "avx2")] fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } ``` 你*只有*在拥有 `Avx2` 令牌时才能调用这个函数，而这个令牌只有在系统支持 AVX2 指令时才能获得。如果我们能向编译器解释这是合法的（使用 `unsafe`），我们就可以只写一次 `unsafe`，然后到处复用。 我们需要的是一个可以安全调用的宏： ``` with_avx2!( fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } ) ``` 但它的展开结果是这样的： ``` fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { // SAFETY: Avx2 is available according to the token, // and we verified that the inner function is not an `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 中不包含的 intrinsic，**编译器会拒绝它！** 我们就这样在不需要任何自定义目标特性追踪的情况下，为 SIMD intrinsics 提供了一个安全的编程接口！ 虽然其中仍然有一个 `unsafe` 块，但它被封装在一个健全的 API 中，你无法滥用它来造成内存安全漏洞。在这个意义上，它就像 `println!` 一样，安全地抽象了不安全的代码。 这样你只需要审查和审计*这一个宏*，而不是成百上千个自定义的 `unsafe` 块。而在实现中我们*可能*出错的只有： 1. 将令牌映射到错误的 `#[target_feature]` 2. 允许从安全上下文调用 `unsafe fn` 这两种失败模式都相当容易检查。 现在我们可以调用 `add_avx2(token, a, b)` 而不需要 `unsafe` 了，但这仍然不能让我们实现 `a + b`。我们要怎么解决*那个*问题？ ## 泛型来救场 https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/#generics-to-the-rescue 我们无法在 `a + b` 的实现上标注 `#[target_feature]`，因为它必须在任何地方都可以安全调用。我们也不能将令牌传入函数，因为它只接受 `a` 和 `b`，不接受 `token`。 但即使我们能做到，这也会让 API 变得非常丑陋。我们希望 `a + b` 永远可用，并自动使用最好的 SIMD 指令，而用户完全不需要操心令牌的事。 我们可以用泛型同时解决这两个问题：通过定义一个对可用指令集泛型的 `f32x8` 类型，我们可以在其上实现加法，这个加法既*暗中携带令牌*，又为每种 SIMD 指令集创建独立的实现！ 代码如下： ``` pub trait Level {} #[derive(Clone, Copy)] pub struct Avx2(()); impl Level for Avx2 {} pub struct f32x8<L: Level> { // For simplicity we'll back this with an array in this example. // In production code we use native SIMD types for the level. data: [f32; 8], // The smuggled token! token: L, } /// implementation of `a + b` for Avx2 impl std::ops::Add for f32x8<Avx2> { type Output = Self; fn add(self, rhs: Self) -> Self::Output { // (type conversions abbreviated) // Use the Avx2 token to call our safe wrapper let result = add_avx2(self.token, self, rhs); Self { data: store_m256(result), token: self.token, } } } ``` 然后我们可以同样轻松地让它在任何其他指令集上工作，或者在完全没有 SIMD 的情况下工作： ``` #[derive(Clone, Copy)] pub struct NoSimd(()); impl Level for NoSimd {} /// implementation of `a + b` when no SIMD is available 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 想说几句 https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/#the-abi-would-like-a-word 不幸的是，要写 `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 工作呢？ ## 内联它，用火焰内联它 https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/#inline-it-inline-it-with-fire 我们无法在实现 `a + b` 的函数上放置 `#[target_feature]`，但我们可以把它放在*调用* `a + b` 的函数上！然后我们可以使用[内联](https://matklad.github.io/2021/07/09/inline-in-rust.html)将 `a + b` 的实现复制到调用它的函数中，最终它就落在了 `#[target_feature]` 上下文中。 所以调用链看起来是这样的： ``` #[target_feature(enable = "avx2")] fn do_stuff() { // TODO: some computation c = a + b; // TODO: some more computation } // which calls into... #[inline(always)] // Function body will be copied into the caller fn add(self, rhs: Self) -> Self::Output { // Use the Avx2 token to call our safe wrapper add_avx2(self.token, self, rhs); // return statement abbreviated } // which calls into... #[inline] // Function body will be copied into the caller if feasible #[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` 的方案来说似乎都是根本性的。[Struct Target Features RFC](https://shnatsel.medium.com/github.com/rust-lang/rfcs/pull/3525) 为 `add_avx2(token, a, b)` 和 `add::(token, a, b)` 解决了这个问题，但我暂时还看不到通往 `a + b` 的路径。 ## 现在我们处于什么位置？ https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/#where-does-that-leave-us 尽管所有 SIMD 代码都存在内联的固有问题，我们还是成功地提供了一个令人愉快的 SIMD 抽象，其 `unsafe` 代码量低到前所未有。 你可以在 `fearless_simd` [v0.5](https://crates.io/crates/fearless_simd) 中找到这些想法的生产版本，**现已在**你附近的包注册表**上架**！ [这里有一个小示例](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` 这样的高级操作和平台特定的 intrinsics，以充分利用硬件。 `fearless_simd` 中不止一个 `unsafe` 块，因为它还提供了 `safe_unaligned_simd`（https://crates.io/crates/safe_unaligned_simd）crate 的功能，但那也是在[显著更少](https://github.com/okaneco/safe_unaligned_simd/issues/51)的 `unsafe` 代码量下完成的。 对我来说，使用高级 SIMD 抽象的障碍一直是它们带来的 `unsafe` 代码量之巨大。这很可怕，而且难以合理化。但现在 Rust 中的 SIMD 可以真正做到无畏。 ## 致谢 https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/#acknowledgements 我很震惊我竟然是第一个将这个想法投入生产的人（据我所知），因为我肯定不是第一个想到这个的人。CPU 特性令牌是一个古老而常见的想法。`pulp`（https://crates.io/crates/pulp）crate 已经使用它们多年了，但它依赖于手写的 intrinsics 不安全包装器，并且偶尔会出错。使用泛型生成多个实现也是一个古老的想法。它是 8 年前[原始 `fearless_simd` 概念](https://linebender.org/blog/towards-fearless-simd/)的一部分。比它更早的 `simdeez`（https://crates.io/crates/simdeez）crate 似乎也[使用了类似的东西](https://docs.rs/simdeez/1.0.8/simdeez/)。 将令牌与单个安全包装器（委托给 `rustc`）结合的核心洞见也不是我独有的。仅在 `fearless_simd` crate 的范围内，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 特性级别。它更具表达力，但代价是复杂性，而且他的方案从未被合并，因为没有其他维护者站出来审查。我仍然希望它有一天能作为独立的 crate 发布。 感谢 Daniel 和 [Laurenz Stampfl](https://github.com/LaurenzV) 审查我所有对 `fearless_simd` 的 PR，它们都非常大，快速的审查真的很感激！ *本文最初发表在我的 [Medium](https://shnatsel.medium.com/safe-simd-in-rust-even-on-the-inside-c6f1ff381828) 上。*