高速な静的検索ツリー (S+ ツリー) の実装と最適化

1. イントロダクション

この記事では、Algorithmica に紹介された静的検索ツリー（S+ ツリー）の実装と、その限界まで行う最適化手法を解説します。主な目標は、スループット（1 秒あたりの問い合わせ数）の最大化です。

1.1 問題の定義

• 入力: ソートされた 32 ビット符号なし整数のリスト

  vals: Vec<u32>

  。
• 出力: クエリー $q$ に対して、

  vals

  で $q \geq q$ を満たす最小要素を返すデータ構造（存在しない場合は

  u32::MAX

  ）。インデックスも同時に返すことができる。
• メトリック: スループット（ns/query）の最適化。通常、アモルタイズされた時間を ns で報告します。
• ベンチマーク環境: 入力・クエリーは 31 ビット整数の一様ランダムサンプリング。

1.2 動機

• 本プロジェクトはバイオインフォマティクスにおける DNA データ（例：サフィックスアレイ）の高速検索への応用を目指しています。
• サフィックスアレイ検索の高速化に寄与します。
• 入力は静的（固定参照ゲノム）を前提としています。

1.3 推奨文献

• Khuong & Morin (2017): 「Comparison-Based Searching のためのアレイ配置」。実用的にはアイツギンガー配置が最適です。
• Algorithmica: S+ ツリーに関する論文（本文では既読を前提とします）。
• CPU パフォーマンス入門稿: ベンチマーク指標の基準となります。
• Intel Intrinsics Guide: SIMD 命令の確認に役立ちます（AVX2 想定、AVX512 は使用不可）。

1.4 バイナリ検索とアイツギンガー配置 (Eytzinger layout)

• 基準線として Rust 標準ライブラリのバイナリ検索を使用します。
• アイツギンガー配置の利点: メモリ内の値を再順序化し、二分木の構造に合わせて配置することでキャッシュ効率を最大化します（ルートはインデックス 1、子ノードは $2i, 2i+1$）。これによりキャッシュラインへのプレフェッチが容易になります。
• 性能向上: アレイサイズが増加するにつれてバイナリ検索よりも優れ、最終的に約 4 倍高速（約 200ns/query vs 1150ns/query）となります。

1.5 巨大ページ (Hugepages)

• ツリー割り当てに2MB の巨大ページを使用します（通常は 4kB）。
• TLB ヒット率を向上させ、キャッシュ効率を高めます。
• 実装上の注意: 自動的な巨大ページ化は正確な 2MB の倍数が必要ですが、32MB 未満の割り当てでは実際には 32MB の倍数に丸められるため、大きなデータセットで有効です。

1.6 ベンチマークに関する注記

• 測定単位: クエリーあたりの時間（ns/query）。
• 環境設定: i7-10750H (2.6GHz), L1:32KB/core, L2:256KB/core, L3:12MB, ハイパースレッド無効。
• 統計処理: 測定値の中央値を使用し、最小・最大を除外します。ノイズはキャッシュのアソシエティビティ効果による再現可能なものです。

1.7 キャッシュライン

• メモリ帯域幅の非効率な使用を解消するために、単一のキャッシュライン（64 バイト）に複数の値やレイアウト情報をパックする手法を検討しました。
• Figure 2 (参考): S-ツリーノードを用いたレイアウト変換により、1 つのキャッシュラインから複数回のイテレーションが可能になります。

1.8 S-ツリーと B-ツリー

• B-ツリー: データベースなどで一般的だが、ディスク IO を考慮した設計です。
• S-ツリー: バイナリ検索の高速化とアイツギンガー配置の利点を兼ね備えた構造。
• S+ ツリー (本文中では単に S-ツリー): 内部ノードで値を複製することで、葉ノードでのみ結果を探すように設計し、処理を簡素化しました。

2. 検索の最適化 (

find

の最適化)

2.1 線形探索

• ソートされたノード（16 個の値）内で $q \geq q$ を満たす最初のインデックスを線形スキャンで探す単純な実装です。
• 性能: 初期版はアイツギンガー配置よりも遅い場合がありました。

2.2 自動ベクトラライズーション (Auto-vectorization)

• 早期に

  return

  を避けることで、分岐予測の失敗を防ぎます。代わりに「小さい値の数」をカウントします。
• コンパイラーが自動的に SIMD 命令へ最適化されます。
• 性能: アイツギンガー配置を超え、線形探索より数倍高速になります。

2.3 末尾のゼロカウント (Trailing zeros)

• シェア専用ライブラリ (

  portable_simd

  ) を使用し、比較結果をビットマスクに変換して末尾のゼロ（最初の

  true

  の位置）を検出します。
• 問題点: Intel 固有命令との互換性や、データ型（u32 vs i32）の問題からアセンブリコードが冗長化しました。

2.4 ポップカウント (Popcount)

• レーンの順序を保持せず、「$q < q$」を満たすレーンの数そのものをカウントする方式に変更します。
• 末尾のゼロカウントよりも効率的ですが、依然として命令数が多くなりすぎます（xor, or などの処理が必要）。

2.5 マニュアル SIMD

• シェアライブラリのオーバーヘッドを取り払い、直接 Intel 固有命令 (

  _mm256_packs_epi32

  など) を使用します。
• 最適化結果: 比較命令の後処理を減らし、最終的に合計 5 つの命令だけで済むようになります。
• 性能: Figure 6 に示される通り、手動最適化により約 115ns/query から 140ns/query、そしてさらに改善されました。

3. 検索プロセスの最適化

3.1 バッチ処理 (Batching)

• 一度に複数のクエリー（例：128 個）を処理することで、メモリアクセスの待ち時間をアモルタイズします。
• 性能: バッチサイズが大きくなるにつれてスループットが向上し、約 12 クエリーの並列処理で飽和します。最終的に約 45ns/query (2.5 倍高速化) を達成しました。

3.2 プレフェッチング

• 次回のイテレーションに必要なデータを明示的に CPU に読み込ませます（

  prefetcht0

  ）。
• 効果: データが L3 キャッシュを超えた場合、スループットをさらに低下させずに約 30ns/query (45ns $\to$ 30ns) を達成します。

3.3 ポインタ算術

• アセンブリレベルでの最適化を行いました。
  • 事前スplat: ループ外でクエリー値を SIMD ベクトルに展開し、再利用します。
  • バイト単位のポインタ: 64 ビット分のシフト演算を避け、バイト単位でインデックス計算を行います。
  • 最終バージョン:

    /2

    と

    *64

    を結合して

    *32

    に簡素化し、命令数をさらに削減します。

3.4 スキッププレフェッチ (Skip prefetch)

• ツリーの最初の数層は L1 キャッシュに収まるため、その層のプレフェッチをスキップしても問題ありません。ただし、必ずしも効果的ではありませんでした。

3.5 インターリーブ (Interleave)

• CPU 計算（早期層）とメモリアクセス（後期層）のボトルネックを同時に処理するために、異なるバッチの異なる層を交互に処理する手法です。
• 実装: ツリーの全層をインタリーブし、各ステップで別のバッチを進めます。
• 性能: 約 29ns/query から24ns/queryへと改善され、現在の実装における主要な最適化手法の一つです。

4. ツリー構造の最適化

4.1 Left-tree（左最大ツリー）

• 従来の B+ ツリーや S-ツリーは「右の部分ツリーの最小値」を内部ノードに格納していましたが、これがキャッシュライン境界を越えてしまう問題がありました。
• 左最大ツリー: 内部ノードに「左の部分ツリーの最大値」を格納することで、答えを含むノードの直前ではなく、答えそのものを含むノードへ直接向かせるようにします。
• 性能: ランタイムがさらに低下し、約 22ns/query を達成しました。

4.2 メモリアレイ配置

• フルアレイ配置: オフセット管理を廃止し、ツリー全体を連続したメモリブロックとして格納します。インデックス計算が簡素化されますが、未使用領域の浪費が見られます。
• 逆順配置: 層の順序を逆転させましたが、これらも大きな性能向上は見られませんでした。

4.3 ノードサイズ ($B=15$)

• 分岐因子を 16 にし、ノードあたり 15 値を格納します。
• 利点: インデックス乗算 (

  x = (x<<4) + x

  $\to$

  x<<4

  ) が簡素化されます。
• 性能: データセットが L3 キャッシュ内に収まる場合のみわずかに高速ですが、L3 を超えると RAM スループットがボトルネックとなり効果は薄れます。
• オーバーヘッド: 複製によるスペース効率は悪化します（約 13.3% 増加）。

5. プレフィックスパーティション化 (Prefix partitioning)

5.1 フルレイアウト

• 入力値の上位 $b$ ビットで値をパーティション化し、複数の小さなツリーを作成します。
• 性能: 検索深さを減少させますが、部分ツリー間のパディングオーバーヘッドが発生するため、実用的ではなかった場合がありました。

5.2 コンパクトな部分ツリー

• 各部分ツリーに必要最低限のパディングのみを追加し、スペース効率を改善します。
• トレードオフ: スペース効率性は向上しますが、各クエリーの属する部分を追跡するためのインデックス計算コストが発生するため、速度は低下しました。

5.3 両者の良き点：コンパクトな第一層

• 最初の層（ルート）のみをコンパクト（変動分岐因子）にし、以降の層をフルツリー（固定分岐因子）として格納します。
• 効果: スペース効率性と速度のバランスが最も取れた実装となりました（Figure 24 参照）。

5.4 重なるツリー (Overlapping trees)

• 部分サイズが偏っている場合、隣接するツリー間でメモリアレイを共有することでオーバーヘッドを削減します。
• 実装:

  read_unaligned

  を使用してアンアラインされたアクセスに対応し、パディングの数を調整します。

5.5 ヒューマン由来データ (Human data)

• 人類ゲノムの k-mer データを用いたテストでは、一部の k-mer が圧倒的に多く存在しました（偏り）。
• この場合、重なるツリーやプレフィックスマップなどの高度なパーティション化が必要ですが、単純なツリーではスペースオーバーヘッドが大きくなり実用性が失われました。

5.6 プレフィックスマップ (Prefix map)

• 不均一な部分サイズを線形マッピングせず、最初にどの部分ツリーに属するかを指示する「プレフィックスマップ」を用意します。

6. マルチスレッドでの比較演算

• 本稿では単スレッドの最適化に焦点を当てていますが、バッチ処理は自然な並列化が可能です。
• ただし、メモリロックやキャッシュ排他的な挙動により、単純なマルチスレッド実装には注意が必要です（詳細は GitHub リポジトリ参照）。

7. 結論

まとめ

• 最大の改善要因:
  1. S+ ツリーの採用による構造効率化。
  2. 手動 SIMD (

     find_popcnt

     ) による検索ロジックの高速化。
  3. バッチ処理によるメモリアクセスの並列化。
  4. 左最大ツリーによるキャッシュライン利用効率の向上。
  5. インターリーブによる計算とメモリの同時進行。

今後の仕事 (Future work)

• ブランチー検索 (Branchy search): ツリー探索をさらに最適化。
• 挿入探索 (Interpolation search): データの分布に応じて検索範囲を調整。
• データのコンパクト化: さらに小さいデータ構造へのパッキング。
• インデックス返却: 元のデータの位置情報も返す実装。
• レンジクエリー & ソートクエリー: 範囲検索やソート処理との統合。

---

ソースコード: github:RagnarGrootKoerkamp/static-search-tree