你的Rust代码加速的九个SIMD规则（第二部分）

从将数据摄入速度提高7倍的`range-set-blaze`箱中获得的一般教训

一个蟹将计算委托给小蟹 — 来源：https://openai.com/dall-e-2/. 其他所有图片作者提供。

感谢西雅图 Rust 社区的 Ben Lichtman (B3NNY) 在 SIMD 方面给我指了明确的方向。

这是有关在 Rust 中创建 SIMD 代码的文章的第2部分。 (参见第1部分.) 我们将讨论规则 7 至 9:

7. 使用 Criterion 基准测试选择算法，并发现 LANES (通道) 几乎总是 32 或 64。
8. 使用 as_simd 将最佳 SIMD 算法集成到项目中，用于 i128/u128 的特殊代码，以及额外的上下文基准测试。
9. 使用一个可选的 cargo 功能，将最佳 SIMD 算法从项目中移除（暂时）。

回顾规则 1 到 6:

使用 Rust 的夜版和 core::simd，Rust 的实验性标准 SIMD 模块。
CCC: 检查、控制和选择计算机的 SIMD 功能。
学习 core::simd，但有选择性地。
构思候选算法。
使用 Godbolt 和 AI 来理解您的代码的汇编，即使您不懂汇编语言。
使用内联泛型推广到所有类型和 LANES，(当不起作用时) 使用宏，以及 (当不起作用时) 使用特性。

这些规则是基于我尝试加快 Rust 包range-set-blaze的速度而总结出来的，该包用于操作“密集”整数集合。

请回想一下来自第1部分的规则 6，该规则演示了如何使 Rust 的 SIMD 算法完全泛型化以适用于各种类型和通道。接下来，我们需要选择算法和设置通道。

规则 7：使用 Criterion 基准测试选择算法，并发现 LANES (通道) 几乎总是 32 或 64。

在此规则中，我们将看到如何使用流行的criterion crate 对我们的算法和选项进行基准测试和评估。在 range-set-blaze 的上下文中，我们将评估以下内容:

5 个算法 — Regular、Splat0、Splat1、Splat2、Rotate
3 个 SIMD 扩展级别 — sse2 (128 位)，avx2 (256 位)，avx512f (512 位)
10 个元素类型 — i8、u8、i16、u16、i32、u32、i64、u64、isize、usize
5 个通道数 — 4、8、16、32、64
4 个输入长度 — 1024、10,240、102,400、1,024,000
2 个 CPU — 带有 avx512f 的 AMD 7950X，带有 avx2 的 Intel i5–8250U

基准测试衡量了运行每个组合所需的平均时间。然后我们计算吞吐量，单位为兆字节/秒。

请查看新的相关文章以了解如何开始使用 Criterion。那篇文章还展示了如何使用 Criterion 测量编译器设置（例如 SIMD 扩展级别）的影响。

运行基准测试会生成一个 5000 行的 *.csv 文件，开头如下：

Group,Id,Parameter,Mean(ns),StdErr(ns)vector,regular,avx2,256,i16,16,16,1024,291.47,0.080141vector,regular,avx2,256,i16,16,16,10240,2821.6,3.3949vector,regular,avx2,256,i16,16,16,102400,28224,7.8341vector,regular,avx2,256,i16,16,16,1024000,287220,67.067vector,regular,avx2,256,i16,16,32,1024,285.89,0.59509...

这个文件适合使用电子表格数据透视表或类似Polars的数据框架工具进行分析。

算法和通道

以下是一个 Excel 数据透视表，显示每种算法的吞吐量（兆字节/秒）与 SIMD 通道数的对比。该表对 SIMD 扩展级别、元素类型和输入长度进行了平均。

在我的 AMD 台式机上：

你的Rust代码加速的九个SIMD规则（第二部分）四海第2张

在一台英特尔笔记本电脑上：

你的Rust代码加速的九个SIMD规则（第二部分）四海第3张

这些图表显示 Splat1 和 Splat2 效果最好。它们还显示通道数始终在 32 或 64 的情况下更好。

例如，sse2（128 位宽）如何处理 64 个 i64（4096 位宽）通道？Rust 的 core::simd 模块通过自动有效地将 4096 位分成 32 个每个 128 位的块，使这种魔术变得可能。同时处理这 32 个 128 位块（显然）能够实现超越独立处理 128 位块的优化。

SIMD 扩展级别

让我们将通道数设置为 64，并比较 AMD 机器上不同 SIMD 扩展级别的性能。该表对元素类型和输入长度进行了平均。

你的Rust代码加速的九个SIMD规则（第二部分）四海第4张

在我的 AMD 机器上，使用 64 个通道时，sse2 最慢。将 avx2 与 avx512f 进行对比，结果并不一致。同样，算法 Splat1 和 Splat2 效果最好。

元素类型

接下来，我们将 SIMD 扩展级别设置为 avx512f 并比较不同的元素类型。我们将 LANES 设置为 64，并对输入长度进行了平均。

你的Rust代码加速的九个SIMD规则（第二部分）四海第5张

我们可以看到按位、32 位和 64 位元素的处理速度最快（不过，每个元素来看，较小的类型更快）。Splat1 和 Splat2 是最快的算法，其中 Splat1 稍微更好。

输入长度

最后，让我们将元素类型设置为 i32，比较输入长度和吞吐量。

你的Rust代码加速的九个SIMD规则（第二部分）四海第6张

我们可以看到，所有的SIMD算法在100万个输入上表现差不多。对于短输入来说，Splat1似乎比其他算法更好。

此外，短输入看起来比长输入更快。这可能是缓存的结果，也可能是基准测试中丢弃了不对齐的数据导致的。

基准测试结论

基于这些基准测试，我们将使用Splat1算法。暂时，我们将LANES设置为32或64，但请参考下一个规则中的一些复杂情况。最后，我们建议用户将它们的SIMD扩展级别设置为至少avx2。

规则8：使用`as_simd`将最佳的SIMD算法集成到项目中，使用`i128` / `u128`的特殊代码，并进行上下文基准测试。

`as_simd`

在添加SIMD支持之前，RangeSetBlaze的主要构造函数是from_iter：

let a = RangeSetBlaze::from_iter([1, 2, 3]);

然而，SIMD操作最好用于数组，而不是迭代器。此外，从数组构建RangeSetBlaze通常是一种自然的做法，所以我添加了一个新的from_slice构造函数：

#[inline]    pub fn from_slice(slice: impl AsRef<[T]>) -> Self {        T::from_slice(slice)    }

新的构造函数会对每个整数调用其自己的from_slice方法执行内联调用。对于除i128 / u128之外的所有整数类型，这个方法会执行以下操作：

let (prefix, middle, suffix) = slice.as_simd();

Rust的夜间版as_simd方法可以安全而快速地将切片转换为：

不对齐的prefix —— 我们像以前一样用from_iter处理它。
middle，一个对齐的Simd结构的数组块
不对齐的suffix —— 我们像以前一样用from_iter处理它。

将middle看作是将输入整数分成大小为16（或者其他LANES设置）的块。然后，我们通过is_consecutive函数迭代这些块，寻找连续的true。每个连续的区间形成一个单独的范围。例如，一个由从1000到1159（包括）的160个连续整数组成的区间将被识别并替换为一个单独的Rust RangeInclusive 1000..=1159。然后，这个范围通过from_iter比处理160个单独的整数要快得多。当is_consecutive返回false时，我们将使用from_iter处理块的单个整数。

`i128` / `u128`

对于core::simd不处理的类型（即i128 / u128），我们该如何处理数组呢？暂时，我只是使用较慢的from_iter来处理它们。

上下文基准测试

最后一步是在主要代码的上下文中对SIMD代码进行基准测试，最好使用代表性数据。

该range-set-blaze库已经包含了基准测试。其中一个基准测试衡量了在不同程度的分块性下摄入1,000,000个整数的性能。平均分块大小从1（无分块）到100,000个分块不等。让我们运行这个基准测试，通过将LANES设置为4、8、16、32和64。我们将使用Splat1算法和SIMD扩展级别avx512f。

对于每个分块大小，柱形图显示了摄入1,000,000个整数的相对速度。对于每个分块大小，最快的LANES设置为100%。

你的Rust代码加速的九个SIMD规则（第二部分）四海第7张

我们发现，当分块大小为10和100时，LANES=4最好。然而，当分块大小为100,000时，LANES=4比最佳结果差了4倍。在另一种极端情况下，当分块大小为100,000时，LANES=64看起来效果不错，但与100和1000相比，它分别差了1.8倍和1.5倍。

我决定将LANES设置为16。这是最适合分块大小为1000的情况。此外，它从来没有比最佳结果差过1.25倍。

有了这个设置，我们可以运行其他基准测试。下面的图表显示了各种范围集库（包括range-set-blaze）在相同的任务上运行——摄入1,000,000个不同分块性的整数。纵坐标是毫秒，数值越低越好。

你的Rust代码加速的九个SIMD规则（第二部分）四海第8张

在分块大小为1000的情况下，现有的RangeSetBlaze::into_iter方法（红色）已经比HashSet（橙色）快30倍。注意，刻度是对数刻度。使用avx512f，新的SIMD优化的RangeSetBlaze::into_slice算法（浅蓝色）比HashSet快230倍。使用sse2（深蓝色），它比HashSet快220倍。使用avx2（黄色），它比HashSet快180倍。在这个基准测试中，与RangeSetBlaze::into_iter相比，avx512f的RangeSetBlaze::into_slice快7倍。

我们还应该考虑到最坏情况，即摄入没有分块的数据。我运行了那个基准测试。结果显示现有的RangeSetBlaze::into_iter比HashSet慢了2.2倍。新的RangeSetBlaze::into_slice比HashSet慢了2.4倍。

因此，总体而言，新的SIMD代码为假定为分块的数据提供了巨大的优势。如果这个假设是错误的，它会变慢，但并非灾难性的。

有了SIMD代码集成到我们的项目中，我们准备好发布了，对吗？很遗憾，不是这样的。因为我们的代码依赖于Rust的nightly版本，所以我们应该将其作为可选的项。我们将在下一个规则中看到如何做到这一点。

规则9：将最佳SIMD算法（暂时）从项目中分离出来，采用可选的cargo功能。

我们美观的新SIMD代码依赖于Rust的nightly版本，而nightly版本是会经常更改的。要求用户依赖Rust的nightly版本是残酷的。（此外，当出现问题时收到投诉会很烦人。）解决方案是将SIMD代码隐藏在一个cargo功能的背后。

特性，特性，特性 — 在与SIMD和Rust一起工作的上下文中，单词“特性”有三种不同的使用方式。首先，“CPU/目标特性” — 这些描述了CPU的功能，包括它支持的SIMD扩展。参见target-feature和is_x86_feature_detected!。其次，“夜间特性门限” — Rust使用特性门限来控制Rust夜间版中新语言特性的可见性。例如：#![feature(portable_simd)]。第三，“货运特性” — 这些允许任何Rust crate或库提供/限制对其部分功能的访问。当你在Cargo.toml中添加一个依赖项时，你会看到这些。

range-set-blaze crate为使夜间依赖的SIMD代码可选采取的步骤如下：

在Cargo.toml中，定义与SIMD代码相关的货运特性：

[features]from_slice = []

在lib.rs文件的顶部，使用夜间特性portable_simd，以依赖于from_slice货运特性：

#![cfg_attr(feature = "from_slice", feature(portable_simd))]

使用条件编译属性，例如#[cfg(feature = “from_slice”)]，有选择地包含SIMD代码。这包括测试。

/// Creates a [`RangeSetBlaze`] from a collection of integers. It is typically many/// times faster than [`from_iter`][1]/[`collect`][1]./// On a representative benchmark, the speed up was 6×.////// **Warning: Requires the nightly compiler. Also, you must enable the `from_slice`/// feature in your `Cargo.toml`. For example, with the command:**/// ```bash///  cargo add range-set-blaze --features "from_slice"/// ```////// **Caution**: Compiling with `-C target-cpu=native` optimizes the binary for your current CPU architecture,/// which may lead to compatibility issues on other machines with different architectures./// This is particularly important for distributing the binary or running it in varied environments./// [1]: struct.RangeSetBlaze.html#impl-FromIterator<T>-for-RangeSetBlaze<T>#[cfg(feature = "from_slice")]#[inline]pub fn from_slice(slice: impl AsRef<[T]>) -> Self {    T::from_slice(slice)}

如上文档所示，向文档中添加警告和注意事项。
使用--features from_slice来检查或测试你的SIMD代码。

cargo check --features from_slicecargo test --features from_slice

使用--all-features来运行所有测试，生成所有文档，并发布所有货运特性：

cargo test --all-features --doccargo doc --no-deps --all-features --opencargo publish --all-features --dry-run

结论

所以，这就是向你的Rust代码添加SIMD操作的九个规则。这个过程的简单性反映了core::simd库的优秀设计。在适用的地方，你应该始终使用SIMD吗？最终，是的，当这个库从Rust夜间版转移到稳定版时。目前，只在性能关键的情况下使用SIMD，或者使其使用可选。

有关改进Rust中SIMD体验的想法吗？core::simd的质量已经很高；主要需求是稳定它。

感谢您加入我进行SIMD编程之旅。我希望如果您有一个适合SIMD的问题，这些步骤能帮助您加速它。

请关注VoAGI的Carl。我写有关Rust和Python的科学编程、机器学习和统计方面的文章。我倾向于每月写一篇文章。