跳到正文
Hacker News· ravenical·· 8 天前AI 评分22

Rust 如何实现内部也安全的 SIMD 抽象

Safe SIMD in Rust, even on the inside

AI 导读

Rust 通过 CPU 特性令牌（零大小类型）与宏封装，让 SIMD intrinsics 的 unsafe 代码只需审计一处，即可安全调用 AVX2 等指令。

正文 · AI 翻译

Rust 的 SIMD 抽象 并不像我期望的那样安全。直到现在。

众所周知，直接使用原始 SIMD 内联函数并不愉快。

你想写的是 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] 这样的数据加载成每个内联函数所接受的形式
  • 为运行所在的 CPU 选择最佳实现

幸运的是，Rust 提供了许多 SIMD 抽象，替你处理了所有这些，让你只需写 a + b。

只是有一个小问题。在内部，它们仍然充满了 unsafe。它并没有消失，只是被隐藏了。大量此类代码潜伏在表面之下，偶尔会出些岔子。

或者说，它们曾经如此。直到现在。

我们为什么还需要 ‘unsafe’？

在很长一段时间里，你无法避免把对每个内联函数（例如 _mm256_add_ps）的调用包进 unsafe 里，因为在你运行的 CPU 上不支持某个内联函数时调用它是非法的。

所以你不得不有某种机制来追踪每个内联函数需要哪些指令、你能访问哪些指令，并交叉比对它们来决定调用某个函数是否安全。

手动做很繁琐，用代码生成器做又很复杂，总是容易出错，而且每个内联函数周围都需要 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] 的函数内部调用内联函数，但链条中某处仍然必须有 unsafe。

另一个问题更根本。你不能把 #[target_feature] 放在你的类型的 + 实现上，因为 + 必须始终可用。所以使用这种机制，我们无法拥有 a + b。

引理：CPU 特性令牌

要理解最终解决方案如何工作，你首先需要理解 CPU 特性检测如何工作。

通常，检查像 AVX2 这样的 CPU 特性是在运行时使用 is_x86_feature_detected!("avx2") 完成的。但我们绝对不想每次把两个数相加时都运行这个检查——那会彻底拖垮性能。我们想只检查一次，然后向编译器证明从那时起使用 AVX2 指令是安全的。

相反，我们可以使用一个不可伪造的令牌把这个证明编码进类型系统：一个带有私有内部字段的零大小类型。获得这个令牌的唯一方式是调用一个执行 CPU 特性检查的函数。如果检查通过，该函数就把令牌交给你：

pub struct Avx2(());

fn detect_avx2() -> Option<Avx2> {
    if is_x86_feature_detected!("avx2") {
        Some(Avx2(()))
    } else {
        None
    }
}

而且因为它是零大小类型，传递这个令牌没有运行时开销。它纯粹作为编译期证明而存在。

结果是，只要你拥有一个 Avx2 结构体的实例，你就可以确信系统上可以使用 AVX2 指令。

关键洞见

编译器并不知道，但这个函数调用是安全的：

#[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. 将 token 映射到错误的 #[target_feature]
  2. 允许从安全上下文中调用 unsafe fn

而这两种失败模式都很容易检查。

所以现在我们可以不用 unsafe 来调用 add_avx2(token, a, b)，但这仍然无法让我们达到 a + b。我们如何解决那个问题呢？

泛型来救场

我们不能用 #[target_feature] 来标注 a + b 的实现，因为它必须可以从任何地方安全调用。我们也不能将 token 传入函数，因为它接受 a 和 b，但不接受 token。

但即使我们能做到，那也会导致一个相当丑陋的 API。我们希望 a + b 始终能工作，并自动使用最佳的 SIMD 指令，而用户无需操心 token。

我们可以使用泛型同时解决这两个问题：通过定义一个对可用指令集泛型的 f32x8 类型，我们可以在其上实现加法，既在其中夹带一个 token，又为每个 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 有话要说

不幸的是，编写 a + b 并让它降低为 SIMD 指令存在一个根本问题：函数调用开销。

调用函数并非免费，但相当便宜——只需少量 CPU 指令。但少量指令比我们刚刚用于实现加法的一条指令要多得多！

因此，如果中间有函数调用，加法性能将会骤降。而性能正是使用 SIMD 的全部意义所在！

编译器通常很擅长通过 内联 来消除这种开销。它基本上是将你调用的函数的实现复制粘贴到调用它的函数中，这样就没有函数了，也没有开销了。

但 #[target_feature] 注解会打乱这一切。编译器无法将带有 #[target_feature] 注解的函数内联到没有该注解的函数中，因为所需特性在其中不可用！

猜猜什么不能有 #[target_feature] 注解？没错。

那么，我们如何让 a + b 与 SIMD 协同工作呢？

内联它，用火内联它！

我们不能将 #[target_feature] 放在实现 a + b 的函数上，但我们可以将它放在调用 a + b 的函数上！

然后我们可以使用 内联 将 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 级别进行抽象，你甚至不需要宏。

问题在于，每当你在 SIMD 类型上调用 a + b 时，都必须从一个带有 #[inline(always)] 或 #[target_feature] 的函数中调用它，否则代码仍能编译，但性能会急剧下降。

想亲自看看吗？打开这个示例，从中移除 #[target_feature]，然后看着生成的汇编变成一场绝对的恐怖秀。

我不确定对此能做些什么。对于任何用 SIMD 实现 a + b 的方法来说，这个限制似乎都相当根本。

Struct Target Features RFC 为 add_avx2(token, a, b) 和 add::<S: Simd>(token, a, b) 解决了这个问题，但我目前还看不到通往 a + b 的路径。

那这让我们处于什么境地呢？

尽管所有 SIMD 代码都固有内联这一难题，我们还是设法提供了一个极其令人愉悦的 SIMD 抽象，且不安全代码的数量低得惊人，前所未见。

你可以在 fearless_simd v0.5 中找到这些想法的生产版本，现已推出，就在你附近的包注册表中！还有一个小示例，可以看看它在生产中是如何组合在一起的。

用于实现它的宏也公开暴露了出来，因此你可以轻松地混合搭配像 a + b 这样的高级操作和平台特定的 intrinsics，以充分利用硬件。

fearless_simd 中有不止一个 unsafe 块，因为它还提供了 safe_unaligned_simd crate 的功能，但这也是以比原版低得多的不安全代码量实现的。

对我来说，使用高级 SIMD 抽象的障碍一直是它们带来的大量 unsafe。这很吓人，也很难证明其合理性。

但现在 Rust 中的 SIMD 可以真正无所畏惧了。

致谢

我很震惊自己竟然是第一个将其投入生产的人（据我所知），因为我肯定不是第一个想到这个的人。

CPU 特性令牌是一个古老而常见的想法。pulp crate 已经使用它们多年，但它们依赖于围绕 intrinsics 手写的不安全包装器，并且偶尔会出错。

使用泛型生成多个实现也是一个古老的想法。它是 8 年前最初的 fearless_simd 概念的一部分。比它更早的 simdeez crate 也似乎使用了类似的东西。

将令牌与一个委托给 rustc 的单一安全包装器相结合的关键洞见也并非我所独有。仅在 fearless_simd crate 的背景下，Raph Levien 就对此进行过实验，而 Daniel McNab 在几个月前创建了一个比我的更精细的实现。

Daniel 的方法允许对每个单独的 CPU 特性进行细粒度跟踪，而不是 fearless_simd 使用的那几个固定的 CPU 特性级别。它更具表现力，但代价是复杂性，而且他的方法从未被合并，因为没有其他维护者站出来审查它。我仍然希望它有一天能作为独立 crate 发布。

感谢 Daniel 和 Laurenz Stampfl 审查了我对 fearless_simd 的所有 PR，它们很大，快速的审查真的非常感激！

本文最初发表在我的 Medium 上。

来源：Hacker News · shnatsel.github.io

开源生态部署工程教程实践#教程/实践#开源/仓库#部署/工程