靜態搜尋樹比二分搜尋快 40 倍 – 深入探討最佳化

靜態搜尋樹比二分搜尋快 40 倍 – 深入探討最佳化

TL;DR – 靜態搜尋樹在吞吐量上比二分搜尋快最高可達 40 倍

靜態搜尋樹（S+ 樹）在已排序的 32 位元整數陣列上執行 lower‑bound 查詢的速度，可比 Rust 標準函式庫的 binary search 快 40 倍，在 4 GB 資料集上達到 每查詢 27 ns。此加速來自快取感知的節點佈局、基於 SIMD 的 find、積極的批次處理、預取、指標算術技巧，以及查詢層的交錯執行。

1. 問題與基準

目標。 給定一個已排序的 Vec<u32> 與一批查詢，回傳每個查詢的最小元素 ≥ 該查詢（或 u32::MAX）。衡量指標為 吞吐量——每秒查詢數，亦以每查詢的奈秒數表示。

基準。 以 Rust 的 binary_search 在普通陣列上執行，以及以層序（breadth‑first）方式儲存的 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

自行撰寫的內建函式將編譯器產生的 popcnt‑加‑tzcnt 序列取代為單一 vpmovmskb 後接 popcnt。最終的 find_popcnt 函式將每次節點搜尋的指令數減至 5 條指令，使 S‑tree 的查詢時間降至 115 ns/查詢——相較自動向量化版本提升 2 倍。

3. 批次查詢

同時處理多筆查詢可讓多個記憶體請求同時在飛行中。批次大小為 16 時，單筆查詢的延遲從 115 ns 降至 45 ns（≈2.5 倍加速）。CPU 的 12 條 line‑fill 緩衝限制了進一步的提升；更大的批次會出現報酬遞減。

4. 預取

在每層之後對子節點執行 prefetch_index，可降低 L2‑to‑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 右最小節點

儲存 左子樹的最大值（left‑max）而非右子樹的最小值，可為許多查詢減少一次額外的快取線抓取，將交錯吞吐量從 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 倍），對小型輸入僅有 modest 的速度提升。
  • 緊湊子樹 – 將每棵子樹填充至最大子樹大小，降低開銷但需在每筆查詢追蹤子樹，稍微減慢查詢。
  • 緊湊第一層 – 只壓縮根層；空間開銷與緊湊佈局相同，查詢速度與完整佈局相當。
  • 重疊樹 – 在相鄰前綴間共享儲存空間；記憶體使用大幅下降，效能與緊湊第一層變體相當。
  • 前綴映射 – 使用查找表將前綴映射到其子樹的起始偏移，實現近乎最小的開銷（≈1/16），同時保持接近未分割樹的查詢時間。

總體而言，分割在隨機資料上不會超過簡單的交錯查詢，而在偏斜的人類基因組 k‑mer 資料上，額外的索引查找會侵蝕效益。

8. 多執行緒擴展

在 6 核心上執行最佳化的 S‑tree，單筆查詢時間從 27 ns（單執行緒）降至 7 ns，約 4 倍加速，受限於總體記憶體頻寬。較慢的單執行緒方法最終會收斂到相同的吞吐量，因為記憶體子系統成為瓶頸。

9. 結論

  • 吞吐量提升： 1150 ns → 27 ns 每查詢（≈40×），在 4 GB 已排序陣列上。
  • 關鍵技術： SIMD 重度 find、大批次處理、預取、指標算術技巧、交錯層執行、以及 left‑max 節點佈局。
  • 空間‑時間權衡： 15 值節點將開銷加倍，但對快取適配資料有 modest 的速度提升；分割降低樹高卻增加記憶體開銷與查找成本。
  • 未來方向： 針對偏斜分割的分支搜尋、插值搜尋以減少 RAM 存取、16 位元打包以提升分支因子、範圍查詢，以及與 suffix‑array 搜尋的整合。

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 與 static B‑trees (https://en.algorithmica.org/hpc/data-structures/s-tree/).

Sources