使用RUST进行SIMD编程

使用RUST进行SIMD编程

david david ●
●
251 次阅读
Rust 杂谈

本文介绍如何在纯 Rust 中编写 SIMD 加速代码。AMD Zen 5 是首款拥有完整 512 位数据通路的 CPU，开发者终于可以放心使用 AVX-512 指令。SIMD工作流程分为三步：加载、计算、存储，关键是减少内存访问。文章推荐使用 std::arch 模块实现无依赖的 SIMD 代码，同时提醒编译器会自动向量化常见操作，无需过度优化。作者期待可移植 SIMD特性进入稳定版，届时只需实现一次算法即可跨平台运行。使用好512位数据通路和SIMD指令确实是能将CPU的真正能力发挥出来，因此我们应该学习下Rust，并将它用在合适的领域中。我们可以看到很多AI相关的Python库的底层都是C/C++和Rust开发的，并且使用Rust开发的库数量还在上升。

我最近体验了 AMD Zen 5 CPU（AWS 的 m8a 实例），哇……太棒了。即使不谈 GPU 和 NPU，未来 5 年的CPU发展也会非常令人期待！

在 m8a.2xlarge 虚拟实例上，纯 Rust 实现的 ChaCha20 运行速度达到 5.1 GB/s，ChaCha12 达到 6.7GB/s，BLAKE3 达到 10.8 GB/s，相当不错！

对于不了解的人来说，Zen 5 是第一代拥有完整 512位数据通路的（AMD）CPU。对于像我这样的普通开发者而言，这意味着我们终于可以放心使用 AVX-512 SIMD指令，而不用担心降频和其他令人头疼的问题。在拥有 256 位数据通路的 Zen 4 上，512 位 SIMD指令是"双周期执行"的。而早期的 Intel CPU 在使用 AVX-512 指令时会因为能耗过高（从而导致发热）而降低频率，这反而导致性能比不使用 AVX-512 加速时更差。

你可以在这篇精彩的《Zen 5 AVX-512 深度剖析》和《Zen 5 的 AVX-512 频率特性》中了解更多相关内容。大多数开发者可能并不太关心这些细节，无非就是"哦，我的电脑现在更快了"，但对于 Rust 开发者来说，这可能是我们做梦都不敢想的最佳年终礼物。

为什么？

因为 Rust 让你可以非常轻松地为热点代码路径添加 SIMD 加速，而无需处理汇编语言。你把数据加载到 SIMD 寄存器中，然后就像操作普通变量一样编写代码！AVX-512代码只需不到一天的工作量，就能带来超过 10 倍的性能提升。所以这里是一篇关于如何用纯 Rust（无需 nightly 版本）编写 SIMD 加速代码的入门指南，毕竟当软件运行得更快时，我们所有人都会受益。

你可以在 GitHub 上查看一个针对 x86、ARM64 和 WebAssembly 平台进行 SIMD 加速的生产代码示例：https://github.com/skerkour/chacha20-blake3

SIMD是什么

SIMD 代表单指令多数据（Single Instruction, Multiple Data）：即能够对更大数据向量进行操作的 CPU 指令。

CPU 通常处理最大 64 位的值，我们称这些为"标量指令"。而 SIMD 指令则允许 CPU 处理更大的值，对于 amd64 的 AVX-512 指令集来说最大可达 512 位。我们称这些为"向量指令"。

下面是一个伪代码示例，展示如何将 4 个 uint64 值各加 10：

// instead of doing this:
let mut a = [1, 2, 3, 4];
for n in &a {
    *n += 10;
}

// do this
let mut vector = u64x4::from_array([1, 2, 3, 4]); // a 256-bit vector of 4 uint64
let x = u64x4::splat(10); // create a 256-bit vector of 4 uint64: (10, 10, 10, 10)
let vector = vector + x;
// vector = u64x4(11, 12, 13, 4);

向量化代码不会生成执行成本较高的循环，而是大约只编译成 3 条指令。有一点值得注意：SIMD 指令可能比标量指令消耗更多的电量。

使用SIMD视角来思考

使用 SIMD 指令的工作流程可以概括为三个步骤：

加载 -> 计算 -> 存储

首先，你将数据从内存加载到向量寄存器中。

// loads 8 times the int64 with value 1 in a 512-bit vector
let v1 = _mm512_set1_epi64(1);

// loads the (unaligned) int64 array with 8 elements into a 512-bit vector
let v2 = _mm512_loadu_epi64([1, 2, 3, 4, 5, 6, 7, 8]);

然后进行任何你想要进行的计算例如add, xor, subtract

let result = [0i64, 8];
_mm512_storeu_epi64(result.as_mut_ptr(), v_result);
// result = [2, 3, 4, 5, 6, 7, 8, 9]

最后你要保存会内存中。

重要的是要理解，从内存加载数据和将数据存储到内存都有（相对）巨大的延迟成本，因此应该尽可能减少这类操作。让数据保持在 SIMD 寄存器中是更好的选择。

因此，了解目标指令集有多少可用的 SIMD 寄存器非常重要。例如，NEON 在 arm64 上提供 32 个 128 位寄存器：v0 到 v31。这意味着你可以同时保存多达 32 个 128位向量来执行运算，而无需访问"慢速"内存。

通常有两种方式可以使用 SIMD 指令加速算法。

第一种方式是找出算法中可以并行执行的操作，但这是针对特定算法的，实现起来通常更复杂。

第二种方式更通用也更容易实现，就是将输入"分割"成多个块，每个块包含 X 个数据块，其中 X 是可用的通道数量，这样你就可以并行计算这 X 个数据块。

simd_parallel_blocks.png

例如，ChaCha20 操作的是由 16 个 32 位字组成的 512 位块（16 × 32 位 = 512 位 = 64 字节）。

因此，如果我们有 256 位向量可用，我们将并行操作 8 个块（通道）（256 / 32 = 8），因此我们的输入数据块将是 8 个块长，对于 8 × 64 = 512 字节或更大的输入，可以达到单核全速运行。

另一个例子是 BLAKE3，它也操作 32 位字。在具有 AVX-512 指令的机器上，BLAKE3 对于 16KiB 或更大的输入可以达到单核全速运行：它将输入"分割"成 16 个每个 1024 字节的块（称为chunk），然后使用 AVX-512 指令并行处理这 16 个块，每次操作计算 16 个 32 位的状态字。在具有 AVX2（256 位向量）的机器上，对于 8KiB 或更大的输入可以达到单核全速运行，因为 256位向量中只有 8 个 32 位通道可用。

了解自己的目标主机

实现 SIMD 加速代码需要时间，也会增加维护负担，因此你应该了解你的代码将在哪里运行，以便集中精力。如果你的代码专门在高端 Intel/AMD 处理器（例如服务器）上运行，那么专注于 AVX-512 可能就足够了。

另一方面，如果你的代码主要在消费级机器上运行，那么专注于 AVX2 和 NEON 可能是最佳选择。另外，现在实现 SSE2 SIMD 已经没有意义了，因为 2015 年以后生产的大多数处理器都支持 AVX2。

CPU特性发现

SIMD 加速代码依赖于运行它的 CPU 上可用的指令集。在 Rust 中有几种不同的方式来提供 CPU 特性检测。

第一种是运行时检测，使用 std::arch 模块提供的宏：

fn foo() {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    {
        if is_x86_feature_detected!("avx2") {
            return unsafe { foo_avx2() };
        }
    }
    // fallback implementation without using AVX2
}

这种方法需要标准库，而在编写底层代码时标准库可能并不总是可用的。

第二种是使用编译时特性检测：

#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]

还有一些更冷门的方法，比如使用 cargo features，但我不推荐这样做，因为这会让你的包的使用者感到困惑，尤其是当你的包是他们正在使用的包的依赖的依赖的依赖时。

由于运行时检测依赖于标准库，而某些项目（例如嵌入式软件）可能无法使用标准库，我建议默认提供运行时检测，同时提供一个 Cargo feature让你的包的使用者可以选择仅使用编译时特性检测，这样他们可以精确地指定代码将在哪些 CPU 上运行。

类似这样：

Cargo.toml

[features]
default = ["std"]

# enables the use of the standard library for CPU features detection on supported platforms
std = []
fn my_function() {
    // use runtime detection
    #[cfg(feature = "std")]
    {
        #[cfg(target_arch = "x86_64")]
        if is_x86_feature_detected!("avx512f") {
            return my_function_avx512();
        }

        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        if is_x86_feature_detected!("avx2") {
            return my_function_avx2();
        }
    }

    // use compile-time detection
    #[cfg(not(feature = "std"))]
    {
        #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
        return my_function_avx512();

        #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "avx2"))]
        return my_function_avx2()
    }

    // scalar fallback for when SIMD acceleration is not available
    return my_function_generic();
}

选择适合自己的纯Rust SIMD实现

在纯 Rust 中使用 SIMD 指令有几种不同的方式。

标准库中的实验性 simd 模块。遗憾的是，它目前只能在 Rust nightly 版本中使用。我们将在本文后面介绍这个模块。

wide crate，这是一个为稳定版 Rust 复制 simd 模块功能的第三方 crate，但目前仅限于 256 位向量。我无法使用它，因为它引入了太多依赖。

use wide::*;

fn main() {
    let a = u32x4::splat(1);
    let b = u32x4::from([1, 2, 3, 4]);
    let result = a + b;
    assert_eq!(result.to_array(), [2, 3, 4, 5]);
}

如果你不介意额外的依赖，这是我推荐的方法。

pulp crate 是对 SIMD 的高级抽象，如果你喜欢的话，可以把它看作 SIMD 界的 rayon。和 wide 一样，我无法使用它，因为它引入了太多依赖。我也不太喜欢运行时检测 SIMD指令的方式，因为它目前无法用于 no_std 目标。

use pulp::Arch;

fn main() {
    let mut v = (0..1000).map(|i| i as f64).collect::<Vec<_>>();
    let arch = Arch::new();

    arch.dispatch(|| {
        for x in &mut v {
            *x *= 2.0;
        }
    });

    for (i, x) in v.into_iter().enumerate() {
        assert_eq!(x, 2.0 * i as f64);
    }
}

最后，还有 Rust 标准库中的 arch 模块。

arch 的子模块：x86、x86_64、aarch64等等暴露了每个平台可用的原始内联函数（例如 _mm512_add_epi32）和向量寄存器（例如 __m512i）。

这是最底层的方式，会导致更多重复代码，但它也是目前唯一一个在稳定版 Rust 上无需任何依赖就能工作的方式。因此这是我为我的实现所选择的方式。

自动向量化

有一个重要的问题我想讨论：LLVM 执行的自动向量化。

例如，尽管我尝试了很多次，但很难实现一个比基本方式更快的方法来对两个缓冲区进行异或操作：

input_block
    .iter_mut()
    .zip(keystream)
    .for_each(|(plaintext, keystream)| *plaintext ^= *keystream);

确实，编译器能够识别这种模式，并根据可用的指令集自动生成向量化实现。

编译器获得的信息越多（例如块/分块的大小...），编译器就越能执行自动向量化等优化。一如既往，智能的Rust编译器和 LLVM 在这里为我们保驾护航，让我们的生活更轻松。

我的建议是，除非你有确凿的证据表明这是一个瓶颈，否则不要费心为常见操作（如异或/相加两个缓冲区）手动实现 SIMD 优化。编译器很可能会为你自动向量化，或者至少输出高效的代码。

测试

不要忘记在启用和不启用不同 SIMD 指令集的情况下测试你的实现。

你可以使用 RUSTFLAGS 环境变量来选择性地禁用 CPU 特性：

# run tests for generic (no SIMD acceleration) code
RUSTFLAGS="-C target-cpu=native -C target-feature=-avx2,-avx512f" make test
# run tests for AVX2 code
RUSTFLAGS="-C target-cpu=native -C target-feature=-avx512f" make test
# run tests for AVX-512 code
make test

请注意，GitHub Actions 目前不支持 AVX-512 指令，因此你需要在自己的机器上运行 AVX-512 测试。

可移植 SIMD即将到来（希望如此吧）

可移植 SIMD（Rust 的 simd 模块）可能是目前 nightly 版本中最令人兴奋的 Rust 特性之一。它将大大减轻那些希望提供快速高效且易于维护代码的开发者的维护负担。

它将允许我们使用高级代码为每种向量大小只实现一次算法，例如使用 u32x8 来操作具有 8 个 32 位通道的 256 位向量，然后 Rust 编译器会在编译时选择针对不同 CPU架构使用什么具体指令，并自动回退到标量运算。代码与 wide 类似，但没有任何第三方依赖，并且支持最高 512 位的向量（而 wide 只支持 256 位）。

fn main() {
    // a 128-bit vector for all the platforms that support 128-bit registers
    let a = u32x4::splat(1);
    let b = u32x4::from([1, 2, 3, 4]);
    let result = a + b;
    assert_eq!(result.to_array(), [2, 3, 4, 5]);
}

这真是太棒了：

首先因为我们将不再需要费心学习每个不同平台/向量大小的特定内联函数名称。

其次因为它将大大简化我们的代码。例如，我不得不用 128 位向量实现 ChaCha20 两次。一次是为 NEON（arm64），一次是为 wasm32 的simd128。这并不是很困难，因为代码几乎相同，只是类型和内联函数的名称不同，但这意味着需要测试、维护和文档化更多的代码。

有了可移植 SIMD，我只需要基于 u32x4 类型实现一次，这是一个具有 4 个 32 位通道的 128 位向量，然后 Rust 会将其编译为任何具有 128 位向量指令的平台的优化代码（arm64 上的 NEON、x86上的 SSE2、wasm32 上的 simd128等等）。它还将大大简化 SIMD 代码的测试，因为使用 u32x4 的平台无关实现可以在任何支持 128 位向量的平台上测试，而 std::arch 模块则需要特定的硬件才能运行测试。

我真的迫不及待地想要这个特性登陆 Rust 稳定版！

译者感受

Rust平台如果能实现一个抽象的SIMD确实是个非常好的事情，因为在2025年的12月，自己使用了Rust重写了一个网络代理以便团队进行Web的A/B测试，其中使用到了RS算法。RS算法自己使用的是reed-solomon-erasure ，它里面就使用了SIMD的技术进行了算法加速，团队这使用的便携接入设备都是树莓派4B的小盒子，上面的软件时间比较久远（第一个版本是2009年左右开发的），在迁移到新的版本上，明显感到了网络速度的提升。后经过测试对比发现，使用了SIMD，在数据校验过程中提升了3.5倍左右。

在眼下算例紧缺的情境下，能充分利用好CPU的SIMD来完成一些事情，确实能帮助我们减少一些成本。

原文：SIMD programming in pure Rust