改进 std::simd::swizzle_dyn

# 改进 std::simd::swizzle_dyn | Sergey "Shnatsel" Davidoff 来源：https://shnatsel.github.io/improving-std-simd-swizzle-dyn/ Fearless SIMD (https://github.com/linebender/fearless_simd)最近收到了一个拉取请求 (https://github.com/linebender/fearless_simd/pull/265)，开头是这样的： > Swizzles 是个大麻烦，我现在不想搞清楚如何泛化地实现它。😄 拉取请求的作者只想转换 RGBA 图像布局，例如 RGBA <-> BGRA，所以我们通过实现一个更简单的版本 (https://github.com/linebender/fearless_simd/pull/266) 来满足当前需求，该版本在 128 位块内进行字节洗牌。这样，即使在 128 位硬件上处理 512 位向量，我们也能轻松支持此操作，而 512 位硬件仍能通过原生指令发挥全部潜力。但这让我好奇：`std::simd` (https://doc.rust-lang.org/std/simd/index.html) 如何实现不限于 128 位块的任意洗牌？答案竟然是：**很差。** 在本文中，我们将着手改进一些方面。 ## 什么是洗牌？https://shnatsel.github.io/improving-std-simd-swizzle-dyn/#what-is-a-swizzle 洗牌（或称重排）重新排列数组中的元素。例如，输入数组 `[A,D,F,I,M,R,S,T]`，应用洗牌掩码 `[2,0,6,7,6,3,4,1]`，得到 `[F,A,S,T,S,I,M,D]`。其中的 `dyn` 表示洗牌掩码是动态的（非硬编码），因此你可以在运行时提供掩码 `[6,4,0,5,7,0,6,6]` 来得到不同的结果。这个例子只用了 8 个字节，即 64 位。实际硬件在 128 到 512 位的向量上实现了洗牌，`std::simd` 也是如此。 ## 理解实现 https://shnatsel.github.io/improving-std-simd-swizzle-dyn/#understanding-the-implementation 从高层次看，`std::simd::swizzle_dyn` 的当前实现 (https://github.com/rust-lang/rust/blob/d527bc9bfa297ca7fd7f5ae93781eeec42073170/library/portable-simd/crates/core_simd/src/swizzle_dyn.rs) 非常简单：如果可用则使用原生硬件操作，否则放弃并逐字节移动元素。即使给定大小存在硬件洗牌指令，`swizzle_dyn` 也不一定能干净地映射到硬件。它承诺越界索引的值会被设为 `0`，而 x86 硬件洗牌只是让它们回绕，因此 `swizzle_dyn` 需要额外工作来维持这一保证。等等……它的文档中这个警告是什么？> 注意，当前实现是在标准库构建时选择的，因此可能需要 `cargo build -Zbuild-std` 来解锁更好的性能，特别是对于更大的向量。计划中的编译器改进将允许使用 `#[target_feature]` 代替。哦，糟了。 ## std::simd 中的害群之马 https://shnatsel.github.io/improving-std-simd-swizzle-dyn/#the-black-sheep-of-std-simd Rust 使用 LLVM (https://en.wikipedia.org/wiki/LLVM) 来优化代码并将其转换为 CPU 可执行的指令。但 LLVM 不喜欢处理平台特定的操作。像“两个 SIMD 向量相加”这样的操作更容易泛化实现，对所有优化都在这种泛化形式上进行，然后在后续为特定 CPU 选择合适的指令。大多数 `std::simd` 操作只生成这种泛化形式（例如“两个 SIMD 向量相加”），这使得实现非常直接。但 `swizzle_dyn` (https://doc.rust-lang.org/std/simd/struct.Simd.html#method.swizzle_dyn) 并非如此。而它的实现 (https://github.com/rust-lang/rust/blob/d527bc9bfa297ca7fd7f5ae93781eeec42073170/library/portable-simd/crates/core_simd/src/swizzle_dyn.rs) 看起来像这样： ```rust #[cfg(target_feature = "neon"))] unsafe fn armv7_neon_swizzle_u8x16(bytes: Simd, idxs: Simd) -> Simd { use core::arch::arm::{uint8x8x2_t, vcombine_u8, vget_high_u8, vget_low_u8, vtbl2_u8}; unsafe { let bytes = uint8x8x2_t(vget_low_u8(bytes.into()), vget_high_u8(bytes.into())); let lo = vtbl2_u8(bytes, vget_low_u8(idxs.into())); let hi = vtbl2_u8(bytes, vget_high_u8(idxs.into())); vcombine_u8(lo, hi).into() } } ``` 我见过这个！这根本不是 LLVM 的魔法，尽管它完全难以理解。这些只是老式的 SIMD 内联函数！问题在于这些平台特定的内联函数被放在 `#[cfg(target_feature = ...)]` 下。正是这个 `cfg` 扼杀了性能，因为它会在构建过程的早期被解析，远早于考虑调用函数的上下文。`fearless_simd` 和所有其他不叫 `wide` (https://crates.io/crates/wide) 的 SIMD crate 都有一种解决方案 (https://shnatsel.github.io/safe-simd-in-rust-even-on-the-inside/)，但 `std::simd` 没有，并且它无法在不进行大规模 API 改动的情况下直接采用生态系统的解决方案。文档中的警告并没有充分说明问题的严重性。不仅需要 `-Z build-std`，你还必须传递 `RUSTFLAGS=-C target-cpu=` 才能获得不错的性能。但如果你这样做了，并且在比你选择的 CPU 更老的 CPU 上运行程序，程序将会崩溃。典型的解决方案是通过 `multiversion` (https://crates.io/crates/multiversion) 等 crate 在运行时选择最佳实现，但这对于 `swizzle_dyn` 不起作用，即使使用 `-Z build-std` 也不行。需要说明的是：我不是编译器工程师。我不会去构建任何计划中的编译器改进来修复这个问题。但我是 `fearless_simd` 的维护者，我对 SIMD 内联函数的了解已经超出了我原本的期望。所以这里有很多我可以改进的地方。 ## 没有 NEON？https://shnatsel.github.io/improving-std-simd-swizzle-dyn/#no-neon 阅读 `swizzle_dyn` 代码 (https://github.com/rust-lang/rust/blob/d527bc9bfa297ca7fd7f5ae93781eeec42073170/library/portable-simd/crates/core_simd/src/swizzle_dyn.rs) 时，第一件让我注意到的事情是它只在 ARM NEON 上实现了 128 位洗牌，尽管 NEON 确实有用于更大洗牌的硬件指令。结果发现 `std::simd` 实际上并不是作为 Rust 标准库的一部分开发的——它有自己的仓库 (https://github.com/rust-lang/portable-simd/)，开发在那里进行，结果会定期同步到标准库中。而缺少的 NEON 实现已经在那个仓库中实现了 (https://github.com/rust-lang/portable-simd/blob/9acfa97d5ce89df77c172b3d6942cafe688c18da/crates/core_simd/src/swizzle_dyn.rs#L118-L183)。这在浏览代码时完全看不出来。我差点浪费大量时间向错误的仓库中过时的代码提交拉取请求。在同步进来的文件中添加标题，说明如何为它们贡献代码，可以为我这种处境的人节省大量时间。 ## 加倍向量宽度 https://shnatsel.github.io/improving-std-simd-swizzle-dyn/#doubling-the-vector-width 假设我们想要执行 `[A,D,F,I,M,R,S,T].swizzle_dyn([2,0,6,7,6,3,4,1])`，但我们只有半宽度的向量。`swizzle_dyn` 将越界元素设置为零的特性在这里会帮助我们。为了处理前半部分 `[2,0,6,7]`，我们运行 `[A,D,F,I].swizzle_dyn([2,0,6,7])` 和 `[M,R,S,T].swizzle_dyn([-2,-4,2,3])`，得到 `[F,A,0,0]` 和 `[0,0,S,T]`。然后我们用廉价的按位或合并这两个中间结果，瞧！我们得到了 `[F,A,S,T]`！现在对后半部分重复相同操作，就完成了！可移植性保证很少能真正帮助优化某事而不是增加额外工作，但这次我们幸运了！这需要四次洗牌而不是一次，但仍然比逐个字节移动快得多。我测量了在 `fearless_simd` 内部对 AVX2 的 512 位向量实现此方法，性能提升了 4 倍。这现在已经合并 (https://github.com/rust-lang/portable-simd/pull/540) 到 `std::simd` 中，也将在下一个 `fearless_simd` 版本中可用。将这种方法用于 4 倍向量大小需要 16 次洗牌，所以我认为不值得，但可能仍然值得进行基准测试。 ## 优化 AVX2 https://shnatsel.github.io/improving-std-simd-swizzle-dyn/#optimizing-avx2 AVX2 很奇怪。它有 256 位向量，但大多数操作独立处理 128 位半部分。这对于大多数操作来说没问题，但对于可能跨越这些半部分移动值的洗牌来说就有问题了。所以 AVX2 实现需要两个步骤：在 128 位块内洗牌，然后使用与我们刚才所做的类似的技巧。另一个复杂之处是 `swizzle_dyn` 承诺越界索引的值会被设为 `0`，而 x86 硬件洗牌不这样做，所以我们必须在软件中实现这一点。当前的实现 (https://github.com/rust-lang/portable-simd/blob/9acfa97d5ce89df77c172b3d6942cafe688c18da/crates/core_simd/src/swizzle_dyn.rs#L206-L223) 看起来像这样： ```rust use x86::_mm256_permute2x128_si256 as avx2_cross_shuffle; use x86::_mm256_shuffle_epi8 as avx2_half_pshufb; let mid = Simd::splat(16u8); let high = mid + mid; // 这是顺序敏感的，LLVM 会按照你放置的顺序来安排。 // 大多数 AVX2 实现使用约 5 个“端口”，其中只有 1 或 2 个能够进行置换。 // 但“组合”步骤会降级为也能使用至少 1 个其他端口的操作。 // 所以这试图分散置换，使得组合流程通过“空闲”端口进行。 // 在重新排序之前，应在多个 AVX2 CPU 上进行对比基准测试 let hihi = avx2_cross_shuffle::<0x11>(bytes.into(), bytes.into()); let hi_shuf = Simd::from(avx2_half_pshufb( hihi, // 复制向量的上半部分 idxs.into(), // 这样只使用索引的 4 位仍然能选中字节 16-31 )); // 组合步骤中的零填充给出了“类似 NEON”的越界为零语义 let compose = idxs.simd_lt(high).select(hi_shuf, Simd::splat(0)); let lolo = avx2_cross_shuffle::<0x00>(bytes.into(), bytes.into()); let lo_shuf = Simd::from(avx2_half_pshufb(lolo, idxs.into())); // 重复，然后选择索引 < 16，覆盖之前组合步骤的 0-15 索引 let compose = idxs.simd_lt(mid).select(lo_shuf, compose); compose ``` 这相当聪明。核心洞察是：CPU 拥有许多不同的硬件电路，处理洗牌的电路与处理算术的电路完全分开，因此你可以同时运行它们。这被称为“指令级并行性 (https://en.wikipedia.org/wiki/Instruction-level_parallelism)”。上述代码经过精心调整，将工作分散到不同的电路上，这样就不会在单个电路（即“端口”）上堆积过多工作，从而利用这种并行性来加速。你可以通过使用 `cargo-show-asm` (https://crates.io/crates/cargo-show-asm) 查看生成的汇编代码，然后将结果输入 `llvm-mca` (https://llvm.org/docs/CommandGuide/llvm-mca.html)（机器代码分析器）来衡量其有效性，该工具会考虑所有这些效果来模拟在特定 CPU 模型上的执行。根据 llvm-mca 的分析，在早期的 AVX2 Intel CPU（即 Haswell 和 Broadwell）上，将操作分散到不同类型的端口非常重要。而在另一端，最近的 Intel Tiger Lake 和所有 AMD Zen 系列都塞满了各种电路，它们更倾向于总体工作量更少，你不再需要担心压垮某个电路。Intel Skylake 则无所谓，对所有形式的处理方式都同样平庸。很容易陷入争论应该优先考虑和优化哪个平台。一方面，最近的 CPU 现在更常见，因此为它们优化能使更多人受益。另一方面，较旧的 CPU 按今天的标准性能不足，需要尽可能多的帮助。我们可以通过使用一个巧妙的算法来完全避开这个问题，该算法消除了比较-选择步骤，从而在最近的 CPU 上提升性能，同时不会在老旧 CPU 上退化： ```rust use x86::_mm256_permute2x128_si256 as avx2_cross_shuffle; use x86::_mm256_shuffle_epi8 as avx2_half_pshufb; let lolo = avx2_cross_shuffle::<0x00>(bytes.into(), bytes.into()); let hihi = avx2_cross_shuffle::<0x11>(bytes.into(), bytes.into()); // 加上 0x60 会保留有效索引 0..=31 的低半字节和第 4 位。 // 更大的索引会设置高位，因此 VPSHUFB 提供了所需的越界归零。 let control = x86::_mm256_adds_epu8(idxs.into(), x86::_mm256_set1_epi8(0x60)); // 将索引的第 4 位移动到每个字节的符号位，用于 VPBLENDVB。 let select_high = x86::_mm256_slli_epi16::<3>(control); let from_low = avx2_half_pshufb(lolo, control); let from_high = avx2_half_pshufb(hihi, control); x86::_mm256_blendv_epi8(from_low, from_high, select_high).into() ``` llvm-mca 显示在 Haswell 和 Broadwell 上没有变化，但 Zen 1、Zen 2 和 Tiger Lake 提升了 23%，Zen 3 提升了 7-20%。Skylake 在 256 位洗牌上没有变化，但将 512 位洗牌分解为四个 256 位洗牌后提升了 20%。*我在这里忽略了吞吐量与延迟 (https://www.geeksforgeeks.org/computer-organization-architecture/computer-organization-and-architecture-pipelining-set-1-execution-stages-and-throughput/) 的区别，因为这种形式同时提升了两者。* 这个改动目前正在 `std::simd` 中审核 (https://github.com/rust-lang/portable-simd/pull/542)，`fearless_simd` 从一开始就会采用这种形式。 ## 这值得吗？https://shnatsel.github.io/improving-std-simd-swizzle-dyn/#was-it-worth-it 单独一个 SIMD 操作的性能提升并不一定能转化为使用它的算法的改进。例如，AVX2 只有 16 个 SIMD 寄存器来保存 CPU 可以直接操作的数据。如果你的操作运行更快但使用了更多寄存器，整个算法可能会耗尽寄存器，不得不将数据溢出到堆栈，这增加了代价高昂的加载/存储指令，从而抵消了收益甚至拖慢整个算法！我们需要测量一些使用 512 位洗牌的可行代码的性能。512 位就是 64 字节，因此 base64 (https://en.wikipedia.org/wiki/Base64) 编码/解码很自然地适合。我不知道有没有现成的基于洗牌的 Rust base64 实现可以让我直接使用。但让 LLM 生成一个原型 (https://github.com/Shnatsel/intrepid-base64) 很容易，仅供教育用途。以下是不同洗牌实现下的性能： ``` 标量回退： 1.6 GiB/s 编码，0.9 GiB/s 解码 AVX2 双宽度： 8.2 GiB/s 编码，5.5 GiB/s 解码 AVX2 dw 优化后： 9.4 GiB/s 编码，5.9 GiB/s 解码 ``` 仅本文所述的工作就带来了约 6 倍的改进！让我们看看它与 Sunny Young 的使用 std::simd 的 base64 (https://github.com/mcy/vb64/) 相比如何，后者使用了许多巧妙的技巧，你一定要去读一读相关文章 (https://mcyoung.xyz/2023/11/27/simd-base64/)： ``` AVX2 vb64: 5.9 GiB/s 编码，5.7 GiB/s 解码 ``` 通过模拟的洗牌暴力解决问题，并与一个通过巧妙的完美哈希 (https://mcyoung.xyz/2023/11/27/simd-base64/#simd-hash-table) 精心避免洗牌的实现达到同一量级，这已经相当不错了！而且我们的实现可以随意更换 base64 字母表，而 `vb64` 的完美哈希与标准 base64 字母表绑定。一个未使用 SIMD 的生产级实现 (https://crates.io/crates/base64) 运行速度为 2.7 GiB/s 编码和 2.3 GiB/s 解码，因此使用 SIMD 非常值得。我们的实现在 AVX-512 CPU 上运行速度约为 13 GiB/s，而 vb64 从 AVX-512 中获益远没有那么明显。一个最先进的 (https://arxiv.org/abs/1910.05109) AVX-512 实现被翻译成安全代码后……