在纯Rust中实现SIMD加速算法：以ChaCha20和ChaCha12为例的经验分享

欢迎关注 猩猩程序员 公众号

一切始于周末开始时，我想看看如何让我的新加密算法ChaCha12-BLAKE3在现代CPU上运行得更快。但有一个前提：代码必须既快速又安全且可审计，不能像大多数加密库那样拥有几千行的汇编代码。我喜欢我的加密代码是"安全、快速、可审计。选择三个。"

尽管过程中有些小插曲，但我还是惊讶于在纯Rust中实现ChaCha20 / ChaCha12的SIMD加速竟然如此简单。仅用了两天时间，我就得到了接近手工编写的汇编代码的效果，但这段代码是可读的、内存安全的、可测试的、可审计的、可更新的------这些都是合并任何代码到代码库之前应该具备的先决条件，但汇编代码往往做不到这一点。

以下是我从中学到的内容。

对于感兴趣的人，代码可以在GitHub上找到：GitHub - ChaCha12-BLAKE3（在chacha12文件夹中）。

SIMD是什么？

SIMD代表单指令多数据（Single Instruction, Multiple Data）：是可以在更大的数据向量上操作的CPU指令。

CPU通常处理64位的数据，我们称之为"标量指令"。而SIMD指令允许CPU处理更大的数据，最多可达512位，适用于amd64的AVX-512指令集。我们称这些为"向量指令"。

以下是伪代码示例，展示如何将4个uint64相加：

ini 复制代码

// 而不是这样做：
let mut a = [1, 2, 3, 4];
for n in &a {
    *n += 10;
}

// 应该这样做：
let mut vector = u64x4::from_array([1, 2, 3, 4]); // 一个包含4个uint64的256位向量
let x = u64x4::splat(10); // 创建一个包含4个uint64的256位向量：(10, 10, 10, 10)
let vector = vector + x;
// vector = u64x4(11, 12, 13, 14);

与生成可能需要高开销的循环相比，向量化代码通常会编译为大约3条指令。

需要注意的是，SIMD指令可能比标量指令消耗更多的电力（从而产生更多的热量），并且可能导致一些旧款Intel CPU的降频（至少是这样），这会对性能产生负面影响。

思考SIMD

使用SIMD指令可以概括为三个步骤：

加载 -> 计算 -> 存储

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

ini 复制代码

// 将8次值为1的int64加载到一个512位向量中
let v1 = _mm512_set1_epi64(1);

// 将（非对齐的）int64数组加载到512位向量中
let v2 = _mm512_loadu_epi64([1, 2, 3, 4, 5, 6, 7, 8]);

然后，你执行计算（加法、异或、减法等）：

ini 复制代码

// 并行地加上8个64位的数据
let v_result = _mm512_add_epi64(v1, v2);
// v_result = __m512i(2, 3, 4, 5, 6, 7, 8, 9)

最后，将结果存储回内存：

ini 复制代码

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在ChaCha20中的应用

以ChaCha20为例，它在32位（4字节）字上操作，组成512位（64字节）的数据块（16 * 32位 = 512位 = 64字节）。

因此，如果我们有256位向量可用，我们将同时操作8个数据块（256 / 32 = 8），因此我们的输入数据将被拆分成8块，达到单核的满速处理，适用于输入数据为8 * 64 = 512字节或更多的情况。

BLAKE3也是如此，它同样在32位字上操作。BLAKE3在拥有AVX-512指令集的机器上，对于输入数据为16KiB或更多时，能够达到单核的满速：它将输入数据"拆分"成16个块（称为块），每个块为1024字节，并使用AVX-512指令并行处理这些16个块，每次操作计算16个32位字的状态。在具有AVX2（256位向量）指令集的机器上，它对于输入数据为8KiB或更多时，能够达到单核的满速，因为256位向量仅提供8个32位通道。

确定目标

实现SIMD加速代码需要时间，并且增加了维护负担，因此你应该了解你的代码将在哪些环境中运行，以便集中精力。

如果你的代码仅在高端Intel / AMD服务器上运行，那么集中精力优化AVX-512可能已经足够。

另一方面，如果你的代码大多数时间运行在消费者机器上，那么集中精力优化AVX2和NEON可能是最好的选择。

此外，如今实现SSE2 SIMD并没有太大意义，因为自2015年以来，大多数处理器都支持AVX2。

在我的情况下，加密学几乎无处不在。例如，ChaCha被Go和Rust的随机数生成器所使用。因此，我决定实现AVX2、AVX-512、NEON和WASM SIMD128加速。

我能够这样做是因为ChaCha被特别设计成能够与不同的SIMD指令集一起扩展，因此将ChaCha从AVX2移植到AVX-512等，只需要更改一个变量的值和几个内在函数名。

CPU特性检测

SIMD加速代码取决于CPU上是否支持相关的指令集。

在Rust中有几种不同的方式来进行CPU特性检测。

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

scss 复制代码

fn foo() {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    {
        if is_x86_feature_detected!("avx2") {
            return unsafe { foo_avx2() };
        }
    }
    // 不使用AVX2的回退实现
}

这种方法需要标准库，但在处理低级代码时可能不总是可用。

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

css 复制代码

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

这些方法有一些更深奥的方式，但我不建议使用它们，因为它们可能会混淆你的包的使用者，尤其是当你的包作为另一个包的依赖时。

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

例如，Cargo.toml：

ini 复制代码

[features]
default = ["std"]

# 启用使用标准库来检测支持的平台的CPU特性
std = []

选择你的纯Rust SIMD实现

有几种方法可以在纯Rust中使用SIMD指令。

Rust标准库中的实验性simd模块：目前仅在Rust nightly版本中可用。
wide crate：这是一个第三方库，模拟了simd模块，但目前仅限于256位向量。
pulp crate：这是SIMD的高级抽象，类似于SIMD的rayon库。
Rust

标准库中的arch模块：最底层的实现，适用于稳定版Rust，无需任何依赖。

自动向量化

我想特别讨论一下LLVM的自动向量化。尽管进行了几次尝试，但很难实现比基本的XOR操作更快的方式：

scss 复制代码

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

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

测试

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

可移植的SIMD即将到来

Rust的可移植SIMD可能是目前Rust夜间版本中最令人兴奋的特性之一。

这将大大减轻开发人员的维护负担，使他们能够编写高效、易于维护的代码。

欢迎关注 猩猩程序员 公众号