这一切始于周末之初，我想要看看如何让我的新加密算法ChaCha12-BLAKE3在现代CPU上运行得更快。但有一个问题：代码需要既快又安全且可审计，不像大多数加密库中那些成千上万行汇编代码那样。我喜欢我的加密代码是“安全、快速、可审计。三选二”。
尽管过程中遇到了一些波折，但我对在纯Rust中实现ChaCha20/ChaCha12的SIMD加速感到极为惊讶。仅用两天时间就获得了与手工编写汇编代码相差无几的性能，但代码可读、内存安全、可测试、可审计且可更新——这本应是代码合并到代码库前的必要条件，但汇编代码并不具备这些特性。
以下是我所学到的内容。
对于好奇的读者，代码已发布在 GitHub 上：https://github.com/bloom42/chacha12-blake3(位于 chacha12 文件夹中)。
SIMD?
SIMD代表单指令多数据：CPU指令可以操作更大的数据向量。
CPU通常处理最多64位的值，我们称这些为“标量指令”。而SIMD指令允许CPU处理更大的值，对于amd64的AVX-512指令集，最多可达512位。我们称这些为“向量指令”。
以下是一个伪代码示例，我们希望将 4 个 uint64 值相加：
// 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指令可能比标量指令消耗更多电力(从而产生更多热量)，并在某些较旧的英特尔CPU上导致CPU降频(至少在某些情况下)，从而对性能产生负面影响。
采用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 the 8 64-bit lanes in parallel let v_result = _mm512_add_epi64(v1, v2); // v_result = __m512i(2, 3, 4, 5, 6, 7, 8, 9)
最后将结果存储回内存。
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 个数据块。
以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模块提供的宏：
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 特性让包的用户选择仅在构建时进行特性检测，以便他们能精确指定代码将运行的 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_avx512()
}
// pure-software fallback
return my_function_software();
}
选择您的纯 Rust SIMD 实现
在纯 Rust 中使用 SIMD 指令有几种不同的方式。
标准库中的实验性 simd 模块。不幸的是，它目前仅在 Rust 夜间版本中可用。我们将在本文后面介绍这个模块。
wide crate，这是一个第三方 crate，复制了稳定版 Rust 的 simd 模块，但目前仅限于 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 执行的自动向量化。
例如，尽管尝试过几次，但要实现比基本点积操作更快的两种缓冲区异或(XOR)方法非常困难：
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
可移植的SIMD(希望)即将到来
可移植的SIMD(Rust的simd模块)可能是Rust目前在夜间版本中提供的最令人兴奋的功能之一。
它将大大减轻开发人员维护快速、高效且易于维护的代码的负担。
它将使我们能够仅需一次实现算法，即可支持所有向量大小，使用高级代码如 u32x8 操作一个由 8 个 32 位通道组成的 256 位向量，然后 Rust 编译器将在编译时为不同 CPU 架构选择具体指令，并自动回退到纯软件实现。
该代码与 wide 类似，但不依赖任何第三方库，并支持长度达 512 位的向量(而 wide 仅支持 256 位)。
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]);
}
这真是令人惊叹，首先是因为我们无需再费心去学习每个不同平台/向量大小的内置函数的具体名称。
其次，这将极大简化我们的代码。例如，我曾两次实现ChaCha20算法，分别针对128位向量。一次是针对NEON(arm64)，另一次是针对wasm32的simd128。虽然代码几乎相同，只需更改类型和内置函数的名称，但这意味着需要测试、维护和文档化的代码量会增加。
借助可移植SIMD，我只需基于u32x4类型(即包含4个32位通道的128位向量)进行实现，Rust会将其编译为针对任何支持128位向量指令平台的优化代码(如arm64上的NEON、x86上的SSE2、wasm32上的simd128等)。
这也将极大简化SIMD代码的测试，因为使用u32x4的跨平台实现可在任何支持128位向量的平台上测试，而std::arch模块则需要特定硬件才能运行测试。
我真的迫不及待想看到这个功能在Rust稳定版中实现！
结语
你使用Rust的时间越长，就越能理解为什么它最终会席卷整个计算堆栈，从微控制器到大型服务器，涵盖WebAssembly、机器人、卫星以及其间的一切。
如前文文章所述，加密库中超过37%的漏洞是内存安全问题，因此很明显，汇编语言在加密代码中的应用正逐渐淘汰，而加密代码是数字时代最基础的部分之一，Rust是唯一合理的替代方案。
