在ARM64上加速gearhash（性能提升2倍）

# 在ARM64上加速gearhash（提速2倍） — sam reis 来源：https://sam.dev/blog/gearhash-on-arm64 **简而言之：**自版本`0.1.4`起，`gearhash`（https://crates.io/crates/gearhash）crate 新增了 NEON 后端，使其在典型分块大小下**在 ARM64 上大约快 2 倍**。在 aarch64 架构上该后端会被自动选择，并且向后兼容，因此 crate 的使用者只需更新即可。如果你对实现细节感兴趣，请继续阅读，或直接跳转到[最终结果](https://sam.dev/blog/gearhash-on-arm64#final-results)。 ## 一切如何开始 (https://sam.dev/blog/gearhash-on-arm64#how-it-all-started) 2019 年底，我在构建个人备份系统时，对一种称为**内容定义分块**的技术产生了兴趣。其核心思想是：不再在固定边界处分割文件，而是对文件运行滑动窗口哈希函数，当哈希值满足特定条件时触发分块。这样做的缺点是会产生长度可变的分块，但优点是分块策略对文件中间字节序列的插入或删除具有更强的适应性。 在此过程中，我发现了[FastCDC 论文](https://www.usenix.org/node/196197)。其基础构件是 GEAR 滚动哈希。由于我喜欢快速的事物，我花了大量时间研究如何将论文中发表的串行算法转换为 SIMD 算法。最终，我将其发布为`gearhash`（https://crates.io/crates/gearhash）——一个针对 SSE4.2 和 AVX2 进行了优化的小型 Rust crate。 编写该 crate 时，ARM64 并非一个值得优化的主要目标。AWS 一年前开始提供 ARM64 实例，但仅有第一代基于 Cortex-A72 核心的 Graviton A1 系列，主要面向横向扩展工作负载而非通用计算。Graviton2（首款具有竞争力的 ARM64 核心）在我首次提交代码的同月于 re:Invent 大会上宣布，但直到 2020 年 5 月才正式发布。Apple 在 2020 年 11 月发布了 M1 芯片。 时至今日，情况已大不相同。Apple 将 ARM64 推入了消费硬件主流；AWS 已推出多代 Graviton，并表示连续三年其新增的 CPU 容量中超过一半为 Graviton；GitHub Actions 在 2025 年为公开仓库提供了免费 ARM64 运行器。在所有这些平台上，`gearhash` crate 之前都只能回退到标量循环。 此外，`gearhash` 最初除了我自己几乎没有生产用户，但后来它已成为[Xet 客户端](https://github.com/huggingface/xet-core)的核心组成部分——这是 Hugging Face 用于 Hub 上大文件的存储协议，已取代 Git LFS 成为默认方案。对`gearhash`而言，这意味着现在每天有 1 万至 2 万次下载。 这种对 crate 的重新关注，促使我再次寻求优化的空间。 ## 解锁并行化的关键洞察 (https://sam.dev/blog/gearhash-on-arm64#the-insight-that-unlocked-parallelization) gear hash 内核定义为对 64 位无符号整数的串行函数： ``` hash = (hash << 1).wrapping_add(table[byte as usize]); ``` 两个特性使其难以向量化： 1. **它是一条串行依赖链**：每个字节的哈希值都依赖于前一个字节。从单个流中无法提取数据并行性。 2. **表查找是一种 gather 操作**：256 × 8 字节 = 2 KB，远超任何寄存器内排列操作的容量。每个字节都需要一次真实的内存加载。 在反复尝试后，我对第一个特性有了观察：哈希值宽度为 64 位，每个字节左移 1 位，因此 64 字节后初始值已被完全移出。这意味着你可以从缓冲区的任意偏移位置开始以`hash = 0`进行哈希计算，经过 64 字节的“预热”后，其哈希值将与从头开始的哈希值完全相同。 这使得我们可以将分块拆分为**条带**：将真实的传入哈希值作为 lane 0 的种子，通过计算每个条带前 64 字节的哈希值来初始化其他所有 lane，然后让所有条带同步运行。当某个 lane 报告匹配时，只需确定*哪个*匹配是最早的——这成为了实现中最复杂的部分。 ## 开始移植到 NEON (https://sam.dev/blog/gearhash-on-arm64#beginning-the-port-to-neon) 我从直接移植 SSE4.2 实现开始。`aarch64::uint64x2_t` 是两个 64 位通道，与 `x86_64::__m128i` 相同，因此 SSE4.2 结构几乎可以机械式映射过来。唯一无法映射的是掩码提取。NEON 没有与 `pmovmskb` 等价的指令，因此需要通过一次窄化移动和一次移动将通道比较结果移入标量寄存器，我将其包装在一个小的 `movemask` 辅助函数中。 结果令人失望：比标量代码慢 0.92 倍。 要理解原因，我们需要看看 ARM64 上的循环携带延迟。NEON 版本每次迭代执行如下操作： ``` add.2d v1, v1, v1 ; h << 1，LLVM 将其发射为自加 add.2d v1, v1, v_g ``` 在 Apple 核心上，每条指令耗时约 2 个周期（根据 [Dougall Johnson 的 M1 表](https://dougallj.github.io/applecpu/firestorm-simd.html)），因此每次迭代约 4 个周期，而每次迭代处理 2 个字节（每个 lane 一个），折合约每字节 2 个周期。 标量版本 `hash = (hash << 1) + table[b]` 编译为单条带移位的加法指令 `add x0, x1, x0, lsl #1`，延迟约 2 个周期。这也是约每字节 2 个周期。 这意味着向量版本在关键路径上每单位时间完成的工作量与标量版本相同，但额外还需要承担加载和掩码提取的开销。它无法取得优势。 要在 NEON 上取胜，必须缩短依赖链本身。 ## 缩短依赖链 (https://sam.dev/blog/gearhash-on-arm64#shortening-the-chain) 如果链是每 2 字节 2 个操作，为何不通过展开每字节更新的两个步骤并乘法处理，使其变为每 4 字节 2 个操作？ ``` h1 = (h << 1) + g0 h2 = (h << 2) + (g0 << 1) + g1 ``` 这样，`h2` 通过单次移位和单次加法依赖于 `h`，前提是你预先计算了 `G = (g0 << 1) + g1`。`G` 仅依赖于表查找，不依赖于 `h`，因此不在关键路径上。 结果：0.92× → 1.13×，有所改善，但距离预期的 2 倍还很远。 ### 隐藏在反汇编中的原因 (https://sam.dev/blog/gearhash-on-arm64#the-reason-hiding-in-the-disassembly) ``` add.2d v2, v1, v1 ; h << 1 add.2d v2, v3, v2 ; h1 = (h<<1) + g0 shl.2d v1, v1, #2 ; h << 2 add.2d v3, v3, v3 ; g0 << 1 add.2d v1, v1, v4 ; (h<<2) + g1 <-- 在 h 链上 add.2d v1, v3, v1 ; ... + (g0<<1) <-- 也在 h 链上 ``` 原来，LLVM 自己重新关联了表达式！我写的是 `(h << 2) + (G0 + G1)`，它却生成了 `((h << 2) + G1) + G0`。这当然是合法的变换，但它将第二个加法放回了依赖链上。 ### 有些取巧的修复方法 (https://sam.dev/blog/gearhash-on-arm64#the-somewhat-naughty-fix) 你无法阻止编译器重新关联加法，但你可以（尝试）阻止它看到加法操作。组合项由两个表查找构成，而表查找结果无论如何都会进入通用寄存器，因此组合操作可以在那里完成： ``` let (t00, t01) = (table[b00 as usize], table[b01 as usize]); let (t10, t11) = (table[b10 as usize], table[b11 as usize]); // 在标量寄存器中组合两个表项，使向量操作数 // 不透明，从而阻止编译器将下面的加法重新关联为 // 在循环携带的 h 链上的两个依赖向量加法。 let g1 = vcombine_u64( vcreate_u64((t00 << 1).wrapping_add(t01)), vcreate_u64((t10 << 1).wrapping_add(t11)), ); let h2 = vaddq_u64(vshlq_n_u64::<2>(h), g1); ``` 检查反汇编显示现在链上只有 `shl.2d` → `add.2d`。 结果：1.13× → 1.46×，终于开始有显著的加速，但仍然不够快！ ## 从延迟瓶颈到吞吐量瓶颈 (https://sam.dev/blog/gearhash-on-arm64#from-latency-bound-to-throughput-bound) 受一次展开两步结果的鼓舞，我尝试了四个中间状态，每个状态仍直接从 `h` 计算： ``` let h1 = vaddq_u64(vshlq_n_u64::<1>(h), g[0]); // hk == (h << k) + g[k-1] // ... let h4 = vaddq_u64(vshlq_n_u64::<4>(h), g[3]); ``` 这再次将每字节依赖链长度减半，因此我预期会有另一个大幅提升。然而测量后发现完全没有变化。 此时我怀疑瓶颈可能不再是迭代间的延迟，而仅仅是 CPU 吞吐量。遵循这一假设，我转向尝试减少指令数量。 循环现在每字节执行两次加载：一次加载字节本身，一次加载其表项。虽然后者无法避免，但我们现在可以用一次非对齐的 32 位加载替换那四次连续的字节加载，然后在每个步骤中通过移位取出每个字的一个字节。 结果：1.46× → 1.63×，向 2 倍目标又迈出了一大步。 ## 用一个分支测试四个状态 (https://sam.dev/blog/gearhash-on-arm64#testing-four-states-with-one-branch) 加载优化后，每次迭代中最大的指令块是边界测试：每一步，代码都通过掩码哈希值并与零比较来检查分块边界。虽然我已展开为每步四个中间状态，但仍是逐一探测它们。这意味着从向量单元进行四次独立的移动，每次移动都等待一个分支。 由于常见情况是无匹配，我们实际上可以做的，是在向量单元内组合四个测试并只进行一次移出。如果两个条带中的四个位置都不是边界，那么只需一次移入通用寄存器和一次分支，迭代就可以继续： ``` let t = vandq_u64( vandq_u64(vtstq_u64(h1, maskv), vtstq_u64(h2, maskv)), vandq_u64(vtstq_u64(h3, maskv), vtstq_u64(h4, maskv)), ); if movemask(t) == u64::MAX { i += UNROLL; continue; } ``` 只有在检查失败的罕见情况下，代码才会逐一查看四个状态。对于 Xet 客户端用于其 64 KiB 分块的 16 位掩码，这种情况大约每 8192 次迭代发生一次。 结果：1.63× → 1.81×。对此我相当满意，这也成了我暂时的终点。 ## 最终结果 (https://sam.dev/blog/gearhash-on-arm64#final-results) 综合所有优化，在该 crate 的 11 位基准测试掩码下，NEON 路径从直接移植的 0.92× 提升到了我发布为`0.1.4`（https://crates.io/crates/gearhash/0.1.4）版本中的 1.81×。 然而，在研究过程中我意识到（事后看来很明显），基准测试结果高度依赖于掩码密度。这是因为优化路径有固定的每次调用开销，而标量路径没有，更稀疏的掩码意味着更多的边界和更多的调用。 为了看看这有多大影响，我在 4 到 20 位设置位的掩码范围内运行了基准测试。 Apple M4 上标量和 NEON 路径的吞吐量与平均分块大小的关系。 NEON 在低于约 350 字节分块时比标量慢，在超过 64 KiB 后稳定在约 4,350 MB/s，而标量始终保持在约 2,000 MB/s 附近。 | 设置位 | 平均分块 | 标量 MB/s | NEON MB/s | 比率 | |--------|----------|-----------|-----------|------| | 4 | 16 B | 996 | 250 | 0.25× | | 8 | 256 B | 1864 | 1641 | 0.88× | | 11 | 2 KiB | 1981 | 3565 | 1.80× | | 16 | 64 KiB | 1999 | 4277 | 2.14× | | 20 | 1 MiB | 2000 | 4347 | 2.17× | 该 crate 自己的基准测试比率为 1.8×，位于曲线陡峭的部分，该曲线持续上升，直到在 64 KiB 到 1 MiB 之间平稳在约 2.17×。16 位那一行是 Xet 客户端使用的掩码，比率为 2.14×。 *附注：* 在平均分块低于约 350 字节时，每次调用开销超过了收益，NEON 路径逐渐变得比标量慢。我不期望有人会使用这种配置，但对于设置位较少的掩码，回退到标量路径似乎是一个弥合差距的廉价方法。 ## 下一步计划 (https://sam.dev/blog/gearhash-on-arm64#whats-next) NEON 现在比标量快约两倍，感觉是时候重新审视 x86 后端了。谁知道呢，我在 NEON 实现过程中学到的一些技巧或许可以迁移过去。敬请期待！