最近、筆者はRaBitQという新しい近似最近傍探索アルゴリズムを使ったさまざまな試みを行っています。このアルゴリズムを提案した論文の著者はすでにC++実装を提供しており、実行速度はかなり速いです。筆者はこれをRustに書き換えようとしました（いわゆるRiiR（Rewrite it in Rust）です）が、実装結果は元のアルゴリズムよりも大幅に遅いことが判明しました。この記事では、アルゴリズムのパフォーマンスを段階的に改善する方法をご紹介します。

環境の準備

データセット

最も重要なのは、適当なデータセットをいくつか用意することです。前述の論文では、sift_dim128_1m_l2とgist_dim960_1m_l2という2つのデータセットについての結果が示されています。このデータセットでは、128項目と960項目というのがよく使われるサイズであり、1,000,000個のデータポイント（ベクトル）があれば性能を測るのに十分だということです。つまり、これらのデータセットを使うことで、アルゴリズムの性能をしっかりと評価できるということです。データセットはこちらからダウンロードできます。（こちらのサイトはTLSに対応しておらず、FTPダウンロードしか提供していません。）

これらのデータセットは、一般的なベクトル形式であるfvecs/ivecsを使用しています。

| dim (4 bytes) | vector (4 * dim bytes) |
| dim (4 bytes) | vector (4 * dim bytes) |
...
| dim (4 bytes) | vector (4 * dim bytes) |

読み取り／書き込み用のスクリプトは筆者のGistから取得できます。

プロファイリングツール

Rustコードのプロファイリングにはsamplyを使用しています。このツールはFirefox Profilerに追加して使用できます。プロファイリングの結果は、クラウドにアップロードすることで他者と共有することもできます。こちらはGistで公開されているC++版のプロファイリング例です。ビューはFlameGraphとCallTreeが最も一般的です。パフォーマンスイベント権限を付与し、mlockの上限を増やす必要があります。

echo '1' | sudo tee /proc/sys/kernel/perf_event_paranoid
sudo sysctl kernel.perf_event_mlock_kb=2048

GodBoltのCompiler Explorerも、C++とRustの関数のアセンブリコードを比較するのに役立ちます。

Cargoプロファイル

リリースビルドにデバッグ情報を含めるには、Cargo.tomlに新たなプロファイルを追加します。

[profile.perf]
inherits = "release"
debug = true
codegen-units = 16

コンパイルコストとランタイム速度は、プロファイリングにおけるユーザーエクスペリエンスを大きく左右します。

• cargo buildはコンパイル速度は速いものの、Pythonよりもコードの実行速度は遅い場合がある
• cargo build --releaseは、コードの実行速度は速いものの、コンパイルには時間がかかる場合がある ベンチマーキングにはopt-level = 3を使用することを推奨します。 以下の設定の使用が推奨されているのを目にしましたが、筆者の場合はコンパイル速度が遅くなり、パフォーマンスの改善は全く見られませんでした。

codegen-units = 1
lto = "fat"
panic = "abort"

ベンチマークツール

Criterionは、統計に基づく優れたベンチマークツールです。関連するベンチマークコードをすべて格納するために、新たなリポジトリを作成しました。ベンチマークコードは同じリポジトリに入れる必要があるようです。

注目すべき点として、ベンチマークの結果はあまり安定していません。筆者の場合はコードに手を加えていない状態で±10%の差異が見られました。ノートパソコンを使用している場合、高温によりCPUがアンダークロックし、さらに差異が広がる可能性があります。

関数のベンチマークには、いくつかの異なるパラメータを使用することをおすすめします。この場合、筆者は異なる項目のベクトルを使用します。すべての項目の結果がプラスであれば、通常は改善が効果的であることを意味します。

指標

最初からいくつかの指標を追加しましょう。指標をチェックすることで、多くのバグやパフォーマンスの問題を見つけることができます。筆者は、現行の要件が単純であるため直接AtomicU64を使用していますが、後でPrometheusを使った指標に移行するかもしれません。 指標／ロギング／トレースが多すぎるとパフォーマンスに影響が出る可能性があるため、これらを追加する際には注意してください。

リソース

ベンチマークの際に、筆者はエンドツーエンドのQPSが極めて不安定であることに気がつきました。コードを再コンパイルしていなくても、翌朝にはベンチマークの結果に15%の改善または悪化が見られる場合もあります。また、筆者はVSCodeの拡張機能「rust-analyzer」を使用しているため、CPUが完全にアイドル状態ではなく、CPUの使用率は低いもののベンチマークの結果に大きな影響を及ぼすことが分かりました。ちなみに、筆者は8つの高性能コアと8つの高効率コアで構成されたIntel Core i7-13700Kを使用しており、プログラムはシングルスレッドです。

筆者はtasksetを使用してプロセスを実行するCPUを指定しています。そうすることで、種類の異なるコアのスケジューリングの影響を受けることはありません。

第13〜14世代のIntel Core CPUは電圧が非常に高いため不安定です。筆者はこの問題をBIOSで修正しました。

クラウドの仮想マシンはCPU温度の影響を受けないかもしれませんが、クラウドプロバイダーにはCPUのスロットリングやオーバーコミットに関する独自のポリシーがある場合があります。

段階的な改善

単純な実装から始める

筆者の最初のリリースは、nalgebraという代数ライブラリをベースにRaBitQアルゴリズムを実装しました。主な理由は、RaBitQアルゴリズムにおいて重要なステップである直交行列を取得するためにQR分解を使う必要があるからです。また、成熟した線形代数ライブラリは行列やベクトルを操作するのに役立つ関数を多く提供するため、アルゴリズムを実装しやすいという利点もあります。Pythonでnumpyを使わずに行列の乗算、射影、分解に関するアルゴリズムを実装することを考えてみてください。まさに悪夢です。

nalgebraはこのようなシナリオ向けに最適化されているため、パフォーマンスは良いはずだと考えました。しかし、ベンチマークは筆者の予想よりもかなり遅い結果となりました。numpyで再実装した方がはるかに速いでしょう:(

プロファイリングによると、f32::clone()の呼び出しが多数あります。全体の時間の約33%を占めており、query_one関数に注目すると44%となります。この結果から、一部のベクトルについてメモリを事前に割り当て、同じイテレーション内で再利用するという非常に一般的な裏技が使えることが分かります。したがって、(x - y).norm_squared()を使う代わりに、(x - y)の結果を格納する別のベクトルをあらかじめ宣言する必要があり、結果的にx.sub_to(y, &mut z); z.norm_squared()となります。コミット23f9affをご覧ください。

ほとんどの代数ライブラリと同様に、このベクトルは列優先で行列を格納するため、行よりも列に対して処理を順番に繰り返し実行した方が速いことを意味します。イテレーションの前に行列を転置しなくてはならないのは少し面倒であり、ベクトルと行列の乗算がすべてコンパイル時に項目のミスマッチエラー（1 x dynまたはdyn x 1）を検出できるとは限りません。

CPUターゲット

RaBitQは、バイナリの内積距離を用いておおよその距離を推定します。距離は以下により計算します。

fn binary_dot_product(x: &[u64], y: &[u64]) -> u32 {
    assert_eq!(x.len(), y.len());
    let mut res = 0;
    for i in 0..x.len() {
        res += (x[i] & y[i]).count_ones();
    }
    res
}

ここでは、u64::count_ones()は組み込み関数であるため、通常は特別な設定なしで使用できると考えましたが、実際にはコンパイル時にpopcnt機能を有効化する必要があります。RUSTFLAGS="-C target-feature=+popcnt"を使うことで有効化できますが、筆者はRUSTFLAGS="-C target-cpu=native"を好みます。後者は現在のCPUが対応するすべてのCPU機能を有効化しますが、バイナリ互換性が失われます。現時点では、特定の環境でのみ使用するため問題ではないと考えています。以下のセクションでも、AVX2機能を有効化するためにこの環境変数が必要になります。

CPUの機能は以下のコマンドを使用することで確認できます。

rustc --print=cfg -C target-cpu=native | rg target_feature

SIMD

最近傍探索において重要な関数は距離関数であり、この場合はユークリッド距離です。通常は、平方根の計算を避けるためにL2二乗距離を使用します。単純な実装は以下のようになります。

{
    y.sub_to(x, &mut residual);
    residual.norm_squared()
}

プロファイリングの後、f32::clone()がまだ残っていることに気づきました。nalgebraのソースコードを確認してみると、何らかの理由により多数のcloneがありました。そこで、SIMDを手書きすることにしました。幸い、hnswlib（人気のHNSW実装）がすでにこれを実装しています。

これにより、距離の計算においてf32::clone()がなくなり、SIFTでQPSが28%改善します。コミット5f82fccをご覧ください。

筆者のCPUはAVX512に対応していないため、AVX2バージョンを使用しています。Steamでハードウェアの統計データを確認できますが、「Other Settings」にSIMDのサポートが記載されています。100%のユーザーがSSE3に対応し、94.61%のユーザーがAVX2に対応していますが、AVX512Fに対応しているユーザーは13.06%に過ぎません。当然この統計データは偏っており、ほとんどのクラウドIntel CPUはAVX512に対応しています。ゲームプレイヤーがすべてのユーザーを代表することはできません。

SIMDを使用する上で非常に役立つガイドとしてIntel Intrinsics Guideが挙げられます。オンラインだと使い勝手が良くないため、サイトをダウンロードすることをおすすめします。組み込み関数の「レイテンシ」と「スループット」は必ず確認してください。そうしないと、通常のバージョンよりもコードが遅くなる場合があります。

x86 Intrinsics Cheat Sheetというリソースもあります。こちらは筆者のような初心者が使いやすい内容になっています。

@ashvardanianが、テール要素問題を解決する「マスクロード」（AVX512が必要）に関する記事を投稿しています。

コードを他のプラットフォームでも実行できるようにするには、以下を行います。 rust #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] { if is_x86_feature_detected!("avx2") { // AVX2 version } else { // normal version } }

SIMD向けにより良いcfgを書くのに役立つクレートがいくつかありますが、今のところは複雑にせず、シンプルにしておきましょう。

SIMDに関する補足

SIMDは金槌のようなものなので、今度はコードの中の「釘」をもっと見つける必要があります。

コミットf114fc1でAVX2を使用してbinarize_vector関数を書き換えたところ、GistのQPSが32%改善した

元のC++版に比べて、この実装も分岐がありません。opt-level=3を有効化すると、コンパイラによってこれを最適化することができます。アセンブリをご覧ください。

@andrewaylett pointed out that opt-level=3 can optimize this

opt-level=3でこれを最適化できると@andrewaylettが指摘

- let shift = if (i / 32) % 2 == 0 { 32 } else { 0 };
+ let shift = ((i >> 5) & 1) << 5;  // opposite version because I store the binary in u64 with different endian from the C++ version

-let shift = if (i / 32) % 2 == 0 { 32 } else { 0 };

+let shift = ((i >> 5) & 1) << 5; // C++版とは異なるエンディアンを使用してバイナリをu64に格納するため、逆バージョン

@NovaX first pointed out that it's equivalent to i & 32, which is more readable.

より可読性に優れたi & 32と同等であると、@NovaXが最初に指摘

スカラー量子化

コード内のf32::clone()を減らすため、筆者はもっと多くのnalgebra関数を手動で実装することにしました。minおよびmax関数が最も一般的なものです。nalgebraバージョンは以下のようになります。

let lower_bound = residual.min();
let upper_bound = residual.max();

これは以下により行います。

fn min_max(vec: &[f32]) -> (f32, f32) {
    let mut min = f32::MAX;
    let mut max = f32::MIN;
    for v in vec.iter() {
        if *v < min {
            min = *v;
        }
        if *v > max {
            max = *v;
        }
    }
    (min, max)
}

以前は便利なのでf32::min()とf32::max()を使用していましたが、昇順／降順以外のベクトルについてはifの方がパフォーマンスが優れています。

次のようにメソッドチェーンでベクトルに対して処理を数回繰り返し実行し、複数のイテレーションでの集計によりスカラー量子化を計算する代わりに、

let y_scaled = residual.add_scalar(-lower_bound) * one_over_delta + &self.rand_bias;
let y_quantized = y_scaled.map(|v| v.to_u8().expect("convert to u8 error"));
let scalar_sum = y_quantized.iter().fold(0u32, |acc, &v| acc + v as u32);

1回のループでこれを実行できます。

{
    let mut sum = 0u32;
    for i in 0..vec.len() {
        let q = ((vec[i] - lower_bound) * multiplier + bias[i]) as u8;
        quantized[i] = q;
        sum += q as u32;
    }
    sum
}

スカラー量子化については、f32をu8に変換できることは分かっているため、to_u8().unwrap()の代わりにas u8を使用できます。 コミットaf39c1cおよびコミットd2d51b0は、GistのQPSを31%改善しました。 以下の部分もSIMDで書き換えることができ、これによりGistのQPSが12%改善します。

• 最小／最大：コミットc97be68およびコミットe5a4af0
• スカラー量子化：コミット28efe09

tr_mulをベクトル射影であるSIMDに置き換えることも試してみました。ここでは、nalgebraがBLASを使用することが判明したため、パフォーマンスは変わりません。

新たな代数クレート：faer

f32::clone()の問題について調べている際に、faerという新たなRust代数クレートを見つけました。多数のSIMDにより最適化されており、行と列のイテレーションパフォーマンスを向上させます。QR分解もnalgebraよりかなり高速です。こちらのコミット0411821はトレーニング部分を高速化します。

また、コミット0d969bd以降はColRefやRowRefラッパーなしでこれらのベクトルを通常のスライスとして使用できます。

最初からfaerを使っていれば、多くのトラブルを回避できていたでしょう。いずれにせよ、この経験から多くを学びました。

バイナリ内積

バイナリ内積についてはpopcntがすでに解決済みと思っていましたが、FlameGraphによると、count_ones()の呼び出しがbinary_dot_productの実行時間に占める割合は7%にすぎないようです。AVX512にはvpopcntq命令がありますが、AVX2シミュレーションの方が一般的なので、筆者はこちらを使用しています。

こちらはAVX2を用いたpopcnt実装の良い参考になります。コミットedabd4aはこれをRustで再実装したもので、GistのQPSを11%改善します。この裏技は、ベクトルに256以上の項目がある場合しかうまくいきません。つまり、バイナリ表現に256ビット必要ということです。

インライン

#[inline]属性は慎重に使う必要があります。この属性をすべてのSIMD関数に追加することで、GistのQPSが5%改善します。

IO

ここで少し背景情報を追加する必要があります。

現行の実装は、k平均法を用いてベクトルのクラスタリングを行い、セントロイドをメモリに格納するIVFアルゴリズムがベースになっています。クエリベクトルは、l2_squared_distance(query, centroid)がより小さいクラスタとしか比較されません。

探索する最近傍クラスタの数を制御するn_probeというパラメータがあります。n_probeの値が大きいと再現率は上がりますが、QPSは下がります。

RaBitQは、バイナリ内積を使用しておおよその距離を推定します。しきい値よりも小さい場合は元のL2二乗距離と同じランキングにリランキングし、それとともにしきい値を更新します。

筆者は以前、n最近傍を選択するだけで並べ替えを行わないslice::select_nth_unstableを使用していました。クエリから遠いクラスタを探索するとリランキング率が上がり、L2二乗距離の計算が増えます。選択したn番目のクラスタを並べ替えることで、GistのQPSが4%改善しました。

別の裏技として、各クラスタのベクトルをセントロイドまでの距離順に並び替える方法もあります。こちらのコミットea13ebcも、GistのQPSを4%改善しました。

各ベクトルのおおよその距離の推定に使用するメタデータがいくつかあります。

• factor_ip: f32
• factor_ppc: f32
• error: f32
• x_c_distance_square: f32

以前、筆者は4つのVec<f32>を使用してこれらのメタデータを格納していましたが、計算を行う際にはそれぞれにvector[i]が必要なため、IOには適していませんでした。コミットbb440e3でこれらを一つのstructに組み合わせることで、GistのQPSが2.5%改善しました。これは4xf32であるためうまく行き、そのためC表現を直接使用することができます。

#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
#[repr(C)]
struct Factor {
    factor_ip: f32,
    factor_ppc: f32,
    error_bound: f32,
    center_distance_square: f32,
}

残念ながら、faerはu64ベクトルに対応していません。したがって、ベクトルのバイナリ表現をVec<Vec<u64>>に格納する必要があります。コミット48236b2でこれをVec<u64>に変更することで、GistのQPSが2%改善しました。

const generics

C++版は、テンプレートを使用して異なる項目のコードを生成します。この機能はRustにもありますが、異なる項目用のコードの再コンパイルは、少数の固定された項目しかない会社の中など、特定のユースケースでしか行えない可能性があるため、筆者は試してみませんでした。公開ライブラリでは、ユーザーが自ら再コンパイルする必要がないよう、一般解を提供する方が良いでしょう。

その他のツール

bounds-check-cookbookでは、safe Rustで境界チェックを無くす方法の例をいくつか紹介しています。

筆者はPGOとBOLTを試してみましたが、改善は得られませんでした。

jemallocやmimallocに変えてもパフォーマンスは改善しません。

まとめ

• SIMDは適切に使用すれば素晴らしい
• 特に大きなデータセットではIOも重要

データセットGistを使用したパフォーマンスは、現時点ではC++版と同じです。筆者はSIMDをより頻繁に使用しますが、C++版はconst genericsを使用しています。

参考文献

• Algorithmica / HPC

この著作物のライセンスは、Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Licenseに基づいて付与されています。