NumPyのint8行列積をRustで高速化し、20環境で検証しました
ベクトル検索の研究では、検索精度だけでなく、量子化したベクトルをCPU上でどのように計算するかも重要になります。
現在進めているベクトル検索 ANNプロジェクトの一環として、NumPyでは期待した速度が出なかったint8行列積をRustで実装しました。さらに、CPUごとに適した命令を選ぶ仕組みと環境別のwheelビルドを整え、20の実行環境で動作と性能を確認しました。
NumPyなら何でも速いわけではない
NumPyは、Python上のループをCで実装された処理へまとめて渡すことで、高速な計算を実現しています。特に浮動小数点の行列積では、OpenBLASやApple Accelerateなど、CPU向けに最適化された計算ライブラリが使われることがあります。
しかし、今回扱うのはint8同士の行列積です。NumPyの整数行列積はBLASへ委譲されず、NumPy自身の汎用ループで計算されます。浮動小数点の行列積で使われる高速なBLAS経路は利用できません。
今回必要だったのは、入力をint8でコンパクトに保持しながら、積和演算の結果をint32へ蓄積する処理です。NumPyそのものが遅いという話ではなく、データ型と演算内容に合わせた専用実装に効果があるのではないかと考えました。
必要な計算だけをRustで実装する
データの準備、実験、分析にはPythonとNumPyを使い、計算量が多い部分をRustで担当させています。
今回はPythonから呼び出せる拡張としてint8行列積を実装しました。Python側の実験方法を変えず、計算を行う部分だけをRustへ置き換えられます。
Pythonで試行錯誤し、処理内容が固まった部分をRustで高速化する。この役割分担を実際に動く形にできました。
int8のまま検索結果までつなげる
ベクトル検索で目指しているのは、int8で保持しているデータを検索時にも別の型へ変換せず、そのままクエリとの内積計算に使うことです。
検索対象をC連続配置のNumPy int8配列として用意すれば、Rust側はそのメモリを読み取り専用で借用できます。検索のたびにfloat32やint32の入力配列を作る必要はありません。クエリもint8のまま受け取り、積和演算だけをint32へ蓄積します。
一般的な正方行列積に加えて、検索に近い次の形の計算もRustで実装し、検証しています。
この形なら、検索対象のint8データをゼロコピーで参照しながら、各ベクトルとクエリの内積を計算できます。さらに、内積から上位k件の選択までをRust側で完結させれば、全件分のスコアをPythonへ戻す必要もありません。
データを小さく保持することと、検索時に余計な変換やコピーを発生させないことを、一つの処理としてつなげていきます。
CPUによって使える命令が違う
同じCPU向けのプログラムでも、実際に使える命令はCPUの種類や世代によって異なります。
- x86_64の基本実装ではSSE2を基準とし、対応CPUではAVX2やAVX-512を使う
- ARM64では標準のNEONを基本実装で使い、対応CPUではdot product命令を追加で使う
- それ以外のCPUアーキテクチャでは基本実装を使う
特定のCPUだけに合わせてビルドすれば高速化しやすい一方、そのバイナリを対応していないCPUで実行すると問題が起きます。配布先ごとに別のバイナリを用意する方法もありますが、運用は複雑になります。
そこで今回は、同じプラットフォーム向けのバイナリに複数の実装を持たせ、実行時にCPUの機能を調べて適切な経路を選ぶようにしました。対応していない命令は実行せず、利用できる範囲で最も適した経路へ切り替えます。
Python向けのwheelは環境ごとに用意する
CPU命令の違いは実行時に吸収できますが、Pythonから利用するwheelまで一つにまとめられるわけではありません。OS、CPUアーキテクチャ、Linuxで利用するlibcなどに合わせたwheelが必要です。
GitHub Actionsでは、OS、CPUアーキテクチャ、Linuxのlibcという軸でwheelをビルドしています。
- Linux: x86_64、x86、ARM64など複数のCPUアーキテクチャ
- Windows: x64、x86、ARM64
- macOS: IntelとApple Silicon
- Linuxのlibc: manylinuxとmusllinux
wheelは環境ごとに分かれますが、同じプラットフォームとCPUアーキテクチャであれば、AVX2やAVX-512などの違いに合わせてwheelを作り分ける必要はありません。一つのwheelに含まれるRustバイナリが、実行時に利用可能な命令を判定します。
また、Pythonの安定ABIであるabi3を使い、一つのwheelをPython 3.8以降から利用できるようにしています。abi3ではPythonのバージョンごとにwheelを作り分けずに済みます。今回の設計では、PythonとRustの境界を越えるのは行列積ごとに一度だけで、その呼び出しの中で大きな計算をまとめて行います。CPUごとの性能を引き出しながら、配布物の種類と境界呼び出しを増やさない構成です。
20の実行環境で確かめた
作ったバイナリを、Linux、macOS、Windowsにまたがる20の実行環境で検証しました。同じ物理マシン上の異なるOSやコンテナ、CIの仮想環境も含むため、20台のマシンという意味ではありません。
検証した範囲は、次のとおりです。
- Linux、macOS、Windowsの3 OS
- x86_64とARM64の2アーキテクチャ
- glibcとmuslの2種類の実行環境
- AVX2、AVX-512、dotprod、基本実装の各経路
すべての環境で、int32による参照結果とのビット一致を確認できました。整数演算では、int32の範囲に収まる限り、SIMD経路によって加算順序が変わっても同じ結果になります。浮動小数点演算では一般に成立しない強い確認を、今回の整数行列積では行えました。
代表的な結果は次のとおりです。1024×1024のint8行列同士を掛け合わせたときの計測時間で、複数回実行した中央値です。
| CPUと環境 | 選択された経路 | NumPy int8（参考） | Rust | 同一バイナリ内の比較 |
|---|---|---|---|---|
| Apple M5 / macOS ARM64 | dotprod | 581.41 ms | 8.00 ms | base 7.93 ms → dotprod 7.77 ms（同程度） |
| Ryzen 9 5950X / Linux x86_64 | AVX2 | 1325.91 ms | 4.43 ms | base 17.12 ms → AVX2 4.85 ms（3.5倍） |
| Ryzen 5 5500U / WSL2 | AVX2 | 1634.49 ms | 17.32 ms | base 58.19 ms → AVX2 17.78 ms（3.3倍） |
| Ryzen 5 5500U / Windows 11 | AVX2 | 2092.03 ms | 15.26 ms | base 57.63 ms → AVX2 15.73 ms（3.7倍） |
NumPy列で測っているのは、入力も出力もint8のまま行う a @ b です。この計算はint8の範囲を超えると折り返すため、処理時間の参考値として掲載しています。Rust実装はint8の入力をint32へ積算し、正しさについては別に a.astype(np.int32) @ b.astype(np.int32) の結果と比較しています。速度表のNumPy列とRust列は、同じ結果を返す実装同士の比較ではありません。
また、NumPyの整数行列積は単一スレッドですが、Rust列はRayonを使って出力行を並列に計算しています。そのためNumPyに対する倍率には、専用の整数カーネル、SIMD、マルチスレッド化の効果がすべて含まれます。右端のbaseとSIMDの比較は、どちらも同じRayon並列で実行し、同じRustバイナリ内で命令経路だけを切り替えています。SIMD自体の効果を見るには右端が適切です。
通常評価と経路を強制した評価は別々に測定しているため、Rustの時間には多少の差があります。
同じRyzen 5 5500UをWindows 11とWSL2から測定した結果では、基本実装が58.19 msと57.63 ms、AVX2経路が17.78 msと15.73 msでした。少なくとも今回の条件では、WSL2を介してもCPU命令の検出と高速化が機能し、Windowsネイティブと大きく離れない結果になりました。
同じ環境内で基本実装とSIMDを使う実装を比較すると、x86_64では約2.7倍から5.8倍の改善が見られました。AVX-512を利用できる環境の一つでは、AVX2経路と比べて約21%の上積みも確認できました。
一方、ARM64のdot product命令は、実機で基本実装と同程度から最大約1.4倍でした。ARM64の基本実装は、必須のNEONを使ってすでに自動ベクトル化されています。dot product命令は、そのNEON経路に対する追加の最適化です。
x86_64の基本実装がSSE2を基準としてAVX2やAVX-512と比較されるのに対し、ARM64ではNEONを使う基本実装とdot product命令を比較しています。同じbase対SIMDでも比較の出発点が異なるため、ARM64側の差が小さいことにも理由があります。
NumPyとの倍率は参考値として見る
一部の環境では、Rust実装がNumPyのint8行列積に対して非常に大きな速度差を示しました。ただし、NumPyとの倍率は、NumPyのバージョンやビルド時の最適化、CPU、OS、仮想化などの影響を受けます。
さらに、先ほど説明したとおり、NumPy列とRust列では出力型とスレッド条件も異なります。そのため、環境をまたいだ単純な倍率を競うのではなく、同じマシン上で基本実装からSIMD経路へ切り替えたときの変化を主な指標にしました。
今回の結果は「Rustなら常にNumPyより速い」というものではありません。計算内容が限定でき、データ型やCPU命令に合わせた実装ができる場合には、部分的なRust化が有効な選択肢になるという結果です。
ベクトル検索に使える計算基盤ができた
前回の記事で紹介したように、ベクトル検索の研究は1,000万件規模の検証、IVF系の探索、SQ8をはじめとする量子化へ広がっています。データが大きくなるほど、ベクトルを小さく保持することと、その形式のまま効率よく計算することの両方が重要になります。
今回、Pythonから使えるRust実装、CPU命令の実行時選択、環境ごとのwheelビルド、20の実行環境による検証までを一通り実現しました。さらに、検索対象ベクトルと一つのクエリとの内積を計算する形でも、int8データを変換せずに扱えることを確認しています。必要な計算だけをRust化し、命令セットの異なるCPUへ同じ実装を届けられる計算基盤になっています。
この計算を上位k件の選択まで融合し、ベクトル検索処理として仕上げていきます。また、CPU上で短時間の計算を行うAWS Lambdaでの活用も検証していきます。AWS Lambdaではx86_64とARM64のアーキテクチャを選択できます。アーキテクチャごとのwheelと、実行環境のCPU機能に合わせた実行時ディスパッチを、そのまま活用できるテーマです。