静的検索ツリー：バイナリサーチの40倍高速（2024年）

目次 1 はじめに 1.1 問題提起 1.2 動機 1.3 推奨される読書 1.4 バイナリサーチとEytzingerレイアウト 1.5 Hugepages 1.6 ベンチマークに関する注意 1.7 キャッシュライン 1.8 SツリーとBツリー 2 findの最適化 2.1 線形 2.2 自動ベクトル化 2.3 末尾のゼロ 2.4 Popcount 2.5 手動SIMD 3 検索の最適化 3.1 バッチ処理 3.2 プリフェッチ 3.3 ポインタ演算 3.3.1 事前スプラット 3.3.2 バイトベースのポインタ 3.3.3 最終バージョン 3.4 スキッププリフェッチ 3.5 インターリーブ 4 ツリーレイアウトの最適化 4.1 左ツリー 4.2 メモリレイアウト 4.3 ノードサイズ（B=15） 4.3.1 データ構造サイズ 4.4 まとめ 5 プレフィックスパーティショニング 5.1 完全なレイアウト 5.2 コンパクトなサブツリー 5.3 両方の長所：コンパクトな第1レベル 5.4 オーバーラップするツリー 5.5 人間のデータ 5.6 プレフィックスマップ 5.7 まとめ 6 マルチスレッド比較 7 結論 7.1 今後の作業 7.1.1 分岐検索 7.1.2 補間検索 7.1.3 データをより小さくパッキングする 7.1.4 元のデータでインデックスを返す 7.1.5 範囲クエリ 7.1.6 ソートクエリ 7.1.7 サフィックス配列検索 この記事では、Algorithmicaで紹介された、ソート済みデータの高スループット検索のための静的検索ツリー（S+ツリー）を実装します。ここでは、主にAlgorithmicaで提示されたコードを出発点として、それを限界まで最適化します。大部分は、その投稿の「今後の作業」のアイデアを取り入れて実装します。そして、可能な限り命令を削減するために、アセンブリコードをたくさん確認します。最後に、スループットを最適化するための大きな追加機能として、バッチ処理があります。 すべてのソースコード（ベンチマークおよびプロットコードを含む）は、githubにあります：RagnarGrootKoerkamp/static-search-tree。 r/programming、hacker news、twitter、bsky、またはyoutubeで議論してください。 1 はじめに 見出しへのリンク 1.1 問題提起 見出しへのリンク 入力。n個の32ビット符号なし整数valsのソート済みリスト：Vec<u32>。 出力。クエリqをサポートするデータ構造で、q以上のvalsの最小要素を返します。そのような要素が存在しない場合はu32::MAXを返します。オプションで、この要素のインデックスも返すことができます。 メトリック。スループットを最適化します。つまり、1秒あたりに回答できる（独立した）クエリの数です。典型的なケースは、十分な長さのクエリ：&[u32]を入力として受け取り、対応するanswers：Vec<u32>を返す場合です。 1注：通常、クエリ/秒ではなく、ns/query（または単にns）として逆スループットを報告します。これは、クエリあたりの償却時間（平均時間ではない）と考えることができます。 ベンチマーク設定。今のところ、入力とクエリの両方が単純な一様ランダムサンプリングされた31ビット整数であると仮定します2。 コード。コードでは、これは次のようにモデル化できます。 1 2 3 4 5 6 7 8 9 trait SearchIndex { /// お互いの関数で実装されるデフォルト実装を持つ2つの関数。 fn query_one(&self, query: u32) -> u32 { Self::query(&[query])[0] } fn query(&self, queries: &[u32]) -> Vec<u32> { queries.iter().map(|&q| Self::query_one(q)).collect() } } コードスニペット 1：ソリューションが実装する必要があるトレイト。 1.2 動機 見出しへのリンク このプロジェクトを単に楽しみのために行うことに加えて、いくつかの高い目標があります。バイオインフォマティクスの大きな目標の1つは、DNA（たとえば、単一のヒトゲノム（30億塩基対/文字）またはそれらのいくつかを）インデックス化するための効率的なデータ構造を作成することです。そのようなデータ構造の1つがサフィックス配列（wikipedia）であり、入力文字列のサフィックスをソートします。古典的には、サフィックス配列をバイナリサーチすることで、文字列が出現する場所を見つけることができます。 このプロジェクトは、サフィックス配列検索を高速化するための最初のステップです。 また、入力データは静的であると仮定することに注意してください。通常、固定参照ゲノムを使用するためです。 1.3 推奨される読書 見出しへのリンク この問題に対する古典的な解決策はバイナリサーチであり、次のセクションで簡単に説明します。このおよびその他の検索レイアウトに関する優れた論文は、KhuongとMorin（2017）による「Array Layouts for Comparison-Based Searching」です。Algorithmicaにも、その論文に基づいたケーススタディがあります。 この記事では、Algorithmicaのフォローアップ記事「static B-trees」で紹介されたS+ツリーに焦点を当てます。私の時間の都合上、この記事に精通しているとほとんど仮定します。 また、CPUを限界までプッシュするベンチマークが含まれている、CPUパフォーマンスの入門書（作業中）を読むこともお勧めします。そこで得られたメトリックを、最適化の試みを理解するためのベースラインとして使用します。 SIMD命令を調べる際には、Intel Intrinsics Guideも役立ちます。ここではAVX2命令のみを使用し、Intelを想定していることに注意してください。また、AVX512命令（特に私のラップトップにはないため）の利用可能性が低いとは想定していません。 1.4 バイナリサーチとEytzingerレイアウト 見出しへのリンク ベースラインとして、Rust標準ライブラリのバイナリサーチ実装を使用します。 1 2 3 4 5 6 7 8 9 10 pub struct SortedVec { vals: Vec<u32>, } implement SortedVec { pub fn binary_search_std(&self, q: u32) -> u32 { let idx = self.vals.binary_search(&q).unwrap_or_else(|i| i); self.vals[idx] } } コードスニペット 2：Rust標準ライブラリのバイナリサーチ。 配列レイアウト論文（KhuongとMorin 2017）の主な結論は、Eytzingerレイアウトが実際には最良のものの1つであるということです。このレイアウトはメモリ内の値を並べ替えます。バイナリサーチは実質的にデータ上のバイナリツリーであり、ルートは中央ノード、次に位置1/4nと3/4nのノード、次に1/8n、3/8n、5/8n、7/8n、といった具合です。Eytzingerレイアウトの主な利点は、バイナリサーチの最初のステップに必要なすべての値が近接しているため、効率的にキャッシュできることです。ルートをインデックス1に配置し、インデックスiのノードの2つの子を2iと2i+1に配置します。これは、インデックス2iまたは2i+1のどちらが必要かを知る前に、次のキャッシュラインを効果的にプリフェッチできることを意味します。これはさらに進んで、4回のイテレーションで必要となる正確な値であるインデックス16iから16i+15を含むキャッシュラインをプリフェッチすることができます。これは、ツリーをたどる際のレイテンシを効果的に隠すことができます。 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 pub struct Eytzinger { /// ツリーのルートはインデックス1にあります。 vals: Vec<u32>, } implement Eytzinger { /// L：プリフェッチするレベル数。 pub fn search_prefetch<const L: usize>(&self, q: u32) -> u32 { let mut idx = 1; while (1 << L) * idx < self.vals.len() { idx = 2 * idx + (q > self.get(idx)) as usize; prefetch_index(&self.vals, (1 << L) * idx); } // 最後の数回のイテレーションでは、もはやプリフェッチは必要ありません。 while idx < self.vals.len() { idx = 2 * idx + (q > self.get(idx)) as usize; } let zeros = idx.trailing_ones() + 1; let idx = idx >> zeros; self.get(idx) } } コードスニペット 3：Eytzingerレイアウトの検索の実装、L=4レベルのプリフェッチ付き。 これら2つをプロットすると、Eytzingerレイアウトは配列がL2キャッシュ（私の場合256KB、中央の赤い線）に収まる場合はバイナリサーチと同等のパフォーマンスを示しますが、配列がL3キャッシュ（12MB）よりもはるかに大きくなると、バイナリサーチよりもはるかに優れたパフォーマンスを発揮し始めます。最終的に、Eytzinger検索は約4倍高速になり、一度に4回のキャッシュラインのプリフェッチをメモリから実行できることにうまく対応しています。 図1：入力サイズの増加に伴うバイナリサーチとEytzingerレイアウトのクエリ スループット。1GBの入力では、バイナリサーチは約1150ns/クエリを必要としますが、Eytzingerは200ns/クエリで6倍高速です。（この2つの青い線については申し訳ありません。上の線がバイナリサーチ、下の線がEytzingerです。すべてのプロットを機能させるのに十分な時間がかかり、すべての色をカスタマイズするのは面倒なので、デフォルトの色を循環させて使用しています。残念ながら、ここでは2つの同じ色になりました。少なくとも、色はプロット間で一貫しています。） 1.5 Hugepages 見出しへのリンク すべての実験で、ツリーを通常の4KBページではなく、デフォルトで2MBのhugepagesを使用して割り当てます。これにより、翻訳ルックア