静态搜索树比二分查找快40倍

背景

在计算机科学中，在一个有序数组中进行查找是最基础的问题之一。经典解法是二分查找，其时间复杂度为 O(log n)，但受限于 CPU 缓存和分支预测的实际性能，二分查找在现代硬件上往往无法充分利用内存带宽。本文（2024 年发布）旨在通过构建一种名为“静态搜索树”（S+ tree）的数据结构，将有序数组的查找吞吐量提升至二分查找的 40 倍以上。该工作基于 Algorithmica 上关于静态 B 树的介绍，并对其进行了极限优化。

该问题的输入是一个已排序的 32 位无符号整数数组 vals: Vec<u32>，需要支持查询：返回数组中第一个不小于查询值 q 的元素，若不存在则返回 u32::MAX。输出指标是吞吐量，即每秒能处理的独立查询数量，通常以纳秒/查询（ns/query）表示。

该研究的动机来自生物信息学：对 DNA 序列（如人类基因组 30 亿碱基）建立后缀数组，然后在其中进行快速查找。静态搜索树是加速后缀数组搜索的第一步。此外，论文假设数据是静态的（如固定参考基因组），因此可以预先构建最优的树结构和内存布局。

核心内容

1. 基准与 Eytzinger 布局

作为基线，使用 Rust 标准库的二分查找。Khuong 和 Morin 2017 年的论文指出，Eytzinger 布局在实践中表现最佳。该布局将数据按完全二叉树方式重排：根节点位于中间位置，子节点位于 1/4 和 3/4 位置，依次类推。在内存中，根节点放在索引 1，左子节点在 2i，右子节点在 2i+1。这种布局使得前几次查找所需的节点彼此靠近，可高效缓存。通过预取未来 4 次迭代所需的缓存行，可以隐藏大部分内存访问延迟。

实验表明，当数组大小在 L2 缓存（256KB）内时，Eytzinger 布局与二分查找性能相当；当数组远大于 L3 缓存（12MB）时，Eytzinger 布局比二分查找快约 4 倍，这得益于一次预取 4 个缓存行。

2. 优化 find 操作（单个查找）

作者对单个查找操作进行了多轮优化，针对树上的遍历过程（从根到叶，共 log₂n 层），逐步消除指令：

• 线性搜索：当节点大小（B）较小时，可以在节点内进行线性搜索，而无需二分比较。
• 自动向量化：利用编译器自动生成 SIMD 指令。
• Trailing zeros：利用 tzcnt 指令快速定位分支。
• Popcount：使用位计数指令。
• 手动 SIMD：使用 AVX2 指令集（128/256 位）进行向量化比较，一次处理多个元素。

3. 优化搜索过程（批量查询）

为了最大化吞吐量，作者将多个独立查询组织成批次，并利用 CPU 的乱序执行和流水线：

• 批处理（Batching）：一次处理多个查询，让 CPU 同时执行多个遍历，最大化预取效果。
• 预取（Prefetching）：提前发起对下一层节点的内存预取，隐藏延迟。
• 指针算术优化：通过改变指针计算方式，减少指令数。例如，使用“up-front splat”将偏移量计算提前，或使用基于字节的指针避免乘法。
• 跳过预取（Skip Prefetch）：在树的上层，由于缓存命中率高，可以跳过预取指令，节省带宽。
• 交错（Interleave）：将多个查询的遍历步骤交错执行，更充分地利用 CPU 执行单元。

4. 优化树布局（内存布局）

• Left-tree 布局：采用左倾树（所有节点左对齐），减少内存碎片。
• 节点大小 B=15：选择 15 个元素作为一个节点，因为 15 是 2¹⁵-1，可以配合 SIMD 一次性比较，且节点大小正好为 64 字节（一个缓存行）。
• 数据结构大小：通过精心计算，使得树的总大小尽可能对齐到 2MB 大页，减少 TLB 压力。
• 总结：通过调整节点大小和布局，进一步减少缓存未命中次数。

5. 前缀分区（Prefix Partitioning）

当数据集非常大时，树的根节点附近节点无法全部放入缓存。作者提出将树的前几层（前缀）进行分区，使得根节点附近的节点可以紧凑存储在连续的缓存行中，从而减少缓存缺失。

• 完整布局：整体树按 Eytzinger 布局。
• 紧凑子树（Compact Subtrees）：将树的前几层单独存储为一棵紧凑的子树，剩余部分保持原样。
• 最佳结合：紧凑第一层：只对第一层进行紧凑存储，其余部分使用标准布局。
• 重叠树：利用前缀重叠，使得多个查询可以共享部分遍历路径。
• 人类数据：测试实际人类基因组数据，验证性能提升。
• 前缀映射：通过额外映射表加速前缀的查找。
• 总结：前缀分区在数据集超过 L3 缓存时效果显著，可额外提升 2~3 倍速度。

6. 多线程对比

作者将单线程最优实现与多线程（如 4 线程）的二分查找进行对比，静态搜索树在多核场景下仍然保持明显优势，且能够更好地利用内存带宽。

7. 未来工作

文中列出了多项未来优化方向，包括分支搜索、插值搜索、数据压缩、返回原始索引、范围查询、查询排序、后缀数组搜索等。

关键要点

• 问题定义：对有序静态 32 位整数数组进行下界查询（lower_bound），目标是最大化吞吐量（ns/query）。
• 性能基准：Eytzinger 布局的二分查找比标准二分查找快约 4 倍，但仍有提升空间。
• 核心优化手段：
  • 使用 AVX2 手动 SIMD 进行节点内线性搜索，节点大小选择 15（恰好一个缓存行）。
  • 批处理查询，利用 CPU 乱序执行和预取隐藏延迟。
  • 优化指针算术，消除冗余指令。
  • 采用前缀分区，将树的前几层紧凑存储，减少缓存未命中。
• 最终结果：在大型数据集（远超 L3 缓存）上，静态搜索树比标准二分查找快约 40 倍，比 Eytzinger 布局快约 10 倍。
• 工具与技巧：使用 2MB 大页减少 TLB 压力；所有分配向上取整到 32MB 以确保大页分配；深入分析汇编代码，逐条指令优化。
• 适用场景：静态数据、高吞吐量查询、数据规模较大（如后缀数组搜索、数据库索引、基因组学）。

意义与影响

• 理论意义：该工作展示了如何将经典的数据结构（B+ 树）与现代 CPU 特性（SIMD、预取、乱序执行）深度