静态搜索树比二分搜索快 40 倍 – 深入优化细节

静态搜索树比二分搜索快 40 倍 – 深入优化细节

TL;DR – 静态搜索树的吞吐量比二分搜索快至 40 倍

静态搜索树（S+ 树）在有序 32 位整数数组上进行下界查询时，速度可比 Rust 标准库的二分搜索 快 40 倍，在 4 GB 数据集上达到 每查询 27 ns。这种加速来源于缓存感知的节点布局、基于 SIMD 的 find、激进的批处理、预取、指针算术技巧以及查询层的交错执行。

1. 问题与基准

目标。 给定一个有序的 Vec<u32> 和一批查询，返回每个查询的最小元素 ≥ 查询值（若不存在则返回 u32::MAX）。衡量指标为 吞吐量——每秒查询数，或以每查询的纳秒数计。

基准。 使用 Rust 的 binary_search 在普通数组上进行搜索，以及使用 Eytzinger 布局（按层序存储的二叉树）作为参考。在 1 GB 输入上，Eytzinger 布局比二分搜索快约 6 倍（≈200 ns/查询 vs 1150 ns/查询），因为它可以提前预取四个缓存行。

2. SIMD 加速的节点搜索 (find)

线性扫描很慢

在节点的 16 个有序值上进行朴素线性扫描，在 1 GB 数据集上大约需要 240 ns/查询，比 Eytzinger 基准还慢。

自动向量化消除分支

改为计数 < q 的元素数量，而不是提前返回，使编译器能够生成 AVX2 SIMD 比较、popcount 和 popcnt 指令。该版本比线性扫描 快 2 倍，并且超过了 Eytzinger 布局。

“自动向量化的版本比线性 find 快两倍，现在明显优于 Eytzinger 布局！” – 作者评论。

手写 SIMD 与 popcnt

手写的 intrinsics 用单条 vpmovmskb 加 popcnt 替代编译器生成的 popcnt+tzcnt 序列。最终的 find_popcnt 函数将每次节点搜索的指令数降低到 5 条指令，使 S‑tree 达到 115 ns/查询，比自动向量化版本 快 2 倍。

3. 批处理查询

同时处理多条查询可以让多个内存请求并行进行。批大小为 16 时，单查询的延迟从 115 ns 降至 45 ns（≈2.5 倍加速）。CPU 的 12 条行填充缓冲区限制了进一步提升；更大的批次收益递减。

4. 预取

在每一层结束后对子节点执行显式 prefetch_index，可以降低 L2‑到‑L3 的失效惩罚。当数据放在 L1 时预取影响不大，对 L2 有轻微帮助，而在数据超过 L2（≈1 GB）时 将 45 ns/查询 降至 30 ns/查询。

5. 指针算术优化

  • 提前 SIMD splat – 将 Simd::splat 移出内层循环可省掉几周期，但对吞吐量影响不大。
  • 基于字节的指针 – 将节点偏移以字节为单位存储，并使用 byte_add，去掉每次访问的昂贵 <<6 移位，削减几条指令，使 1 GB 数据的吞吐量从 31 ns 提升到 28 ns/查询。
  • 合并 /2 与 *64 – 最终的 find_splat64 直接把 popcount 结果乘以 32，省去不必要的 and $‑64 掩码。

6. 树布局变体

左最大 vs 右最小节点

存储 左子树的最大值（左最大）而不是右子树的最小值，可为多数查询省去一次额外的缓存行读取，使交错吞吐量从 24 ns 提升到 22 ns/查询。

节点大小（15 vs 16 个值）

每节点使用 15 个值（分支因子 16）可以简化索引算术（x << 4），对适合 L3 的数据有轻微加速，但结构开销会从 6.25 % 翻倍至 13.3 %。

内存布局实验

  • Packed（默认） – 层层紧密相邻存放。
  • Reversed – 层的顺序反转（如原 Algorithmica 论文所示）。
  • Full array – 使用隐式偏移的平面数组。 这些布局均未能超越默认的 packed 布局；额外的算术开销抵消了任何缓存行对齐的收益。

7. 前缀划分（降低树高）

可以按最高的 b 位将输入划分，针对每个前缀构建独立的 S‑tree。这样每个查询可以省去前 ⌈b/⌈log₂(B+1)⌉⌉ 层的搜索。

  • 简单全布局 – 内存开销大（最高 2 倍），对小输入只有有限的加速。
  • 紧凑子树 – 将每棵子树填充到最大子树的大小，降低总体开销，但需要在查询时跟踪子树编号，略微降低查询速度。
  • 紧凑首层 – 只压缩根层；空间开销与紧凑布局相当，查询速度与全布局相同。
  • 重叠树 – 相邻前缀共享存储；内存使用显著下降，性能与紧凑首层变体相当。
  • 前缀映射 – 使用查找表将前缀映射到子树起始偏移，几乎达到最小开销（≈1/16），且查询时间接近未划分的树。

总体来看，在随机数据上划分并未超越简单的交错查询；在偏斜的人类基因组 k‑mer 数据上，额外的索引查找会抵消收益。

8. 多线程扩展

在 6 核上运行优化后的 S‑tree，单查询时间从 27 ns（单线程）降至 7 ns，约 4 倍加速，受限于整体内存带宽。较慢的单线程方法最终会收敛到相同的吞吐量，因为内存子系统成为瓶颈。

9. 结论

  • 吞吐量提升： 1150 ns → 27 ns 每查询（≈40×），在 4 GB 有序数组上实现。
  • 关键技术： SIMD 密集的 find、大批量处理、预取、指针算术技巧、交错层执行、左最大节点布局。
  • 空间‑时间权衡： 15 值节点将开销翻倍但对缓存适配数据有轻微加速；划分降低树高但增加内存开销和查找成本。
  • 未来方向： 对偏斜划分使用分支搜索、插值搜索以减少 RAM 访问、16 位压缩以提升分支因子、范围查询以及与后缀数组搜索的结合。

10. 社区精彩反馈

“Eytzinger 布局的主要好处在于二分搜索前几步所需的所有值都彼此靠近，因而可以高效缓存……” – kazinator（HN 评论）

“我的第一直觉是 https://en.wikipedia.org/wiki/Van_Emde_Boas_tree” – stevefan1999（HN 评论）

“为什么图 1 对两个不同的曲线使用了相同的颜色和样式？” – petters（HN 评论）

11. 参考文献

  • Khuong, P‑V., & Morin, P. (2017). Array Layouts for Comparison‑Based Searching. ACM JEA, 22.
  • Abrar, M. H., & Medvedev, P. (2024). PLA‑Index: A K‑Mer Index Exploiting Rank Curve Linearity. WABI 2024.
  • Algorithmica 博客关于 S‑trees 与静态 B‑trees 的文章 (https://en.algorithmica.org/hpc/data-structures/s-tree/).

Sources