Fixstars Tech Blog /proc/cpuinfo
🔍

近似最近傍探索ライブラリFaissの4bit PQアルゴリズムについて、ARM CPU上での動作を60倍程度高速化しました

Techblog管理者
に公開
C++
Arm
SIMD
Neon
Faiss
tech

TL;DR: Issue, PR

こんにちは。ソリューション第二事業部の今泉です。 先日東京大学の松井先生と共同でFacebook AI Research(以下FAIR)が公開している近似最近傍探索ライブラリFaissの4bit PQアルゴリズムのARM CPU(aarch64)上での動作を60倍程度高速化しました。 本稿ではまず近似最近傍探索やFaissについて軽く紹介した後、その高速化内容について解説を行います。

近似最近傍探索について

まず最近傍探索とは、「複数のベクトルからなる集合 \mathit{Vs} が存在し、あるベクトル \boldsymbol{x} に対して最も近い要素 \boldsymbol{v} \in \mathit{Vs} を求める」という処理です。 近似最近傍探索はその名の通り「最も近い かもしれない ベクトルを求める」処理となります。 近似の手法によってある程度精度は落ちますが、その分非近似の探索に比べて高速に処理することが期待できます。

近似最近傍探索の応用例としては、類似画像検索や類似文書検索などに用いられます。 画像や文書などの特徴量をベクトルにしておくことで、類似するベクトルの探索問題に帰着させているわけです。

近似最近傍探索についての解説は松井先生のこちらの資料が詳しいです。

Faissについて

Faiss(faceと同じ発音)はFAIRが公開している近似最近傍探索を高速に行うことができるMITライセンスのライブラリで、C++およびPythonで使うことができます。 Faissは大量のベクトルの近似最近傍ベクトルを一度に求める(つまり、上述の \boldsymbol{x} が多数ある)のが得意なライブラリとなります。 FaissにはHNSWやPQ、IVFなど様々な近似最近傍アルゴリズムやそれらを組み合わせたものが実装されており、またGPU実装もあるためVRAMに乗り切る分にはGPUで高速に処理することもできます。 Faissは特にx86向けに強く高速化されており、今回はそこにARM向けの高速化を行った形となります。

4bit PQアルゴリズムについて

Faissで実装されているアルゴリズムの1つに今回高速化した4bit PQアルゴリズムがあります。 松井先生のブログ記事が詳しいのでそちらを参照していただきたいのですが、ざっくりまとめると「量子化のサイズを小さくしてデータをCPUのSIMDレジスタに乗せることでテーブル引きをシャッフル命令で高速に処理する」というものです。

高速化の概要

この節ではまずFaissがどのように4bit PQを実装しているかを説明し、そしてそれをどのように高速化していったか述べます。 平たく言えばスカラー処理のコードをSIMD化しました、というお話になります。

Faissにおける4bit PQの実装: simdlib

Faissにおける4bit PQの実装ではAVX2のintrinsicを直接使うのではなく、独自のラッパーである simdlib を用いて記述されています。 この simdlib が導入された経緯は以下の2点のようです。

  1. コード中で変数の型を明示したかった

    • IntelのSIMD命令のintrinsicは整数型が全部同じ型を使うので、例えばAVX2の整数型が16bit非負整数16個なのか8bit非負整数32個なのか、というのはどのようにintrinsicを使っているかで判断するしかなく、これを嫌ったようです。

  2. AVX2のintrinsicが読みにくいのを嫌った

    • 他の箇所でSSEやAVXのintrinsicを直接使ったコードがたくさんあるのでこれはちょっと意外でした。

この simdlib にはデータを配列で持ちスカラー処理を行う実装(simdlib_emulated)もあり、AVX2が無効の場合にはこちらが使われます。

simdlib_emulated の高速化

上述の通り4bit PQは simdlib で実装されているので、この simdlib にARM SIMD版を実装することで4bit PQをARM上で高速化することができそうです。 ですが、実は今回simdlib_emulated.h も高速化しているので、先にそちらから説明します。

元々の実装をプロファイラにかけたとき、比較的上位に std::invoke が現れたので不思議に思っていたのですが、 simdlib_emulated はstd::function を用いて実装されていました。 std::function は関数のinline化を阻害するので、最内のループで何回も呼び出す部分に使うと当然遅くなります。 また、他にも各関数の引数が値渡しになっていたため、関数呼び出しのたびにコピーコストが発生していました。 これをconst参照で受けるように変更し、コピーコストを削減しました。

これらの変更により、非AVX2環境で4倍程度実行時間が高速化されました。 尤も、この変更はどちらかというと「遅い実装だった simdlib_emulated を妥当な速度になるようにした」という色が強いです。

simdlib_neon の実装

というわけで simdlib をARM SIMD化していきます。 といっても、基本的にはAVX2版があるので同じインターフェースになるよう移植すればよいのですが、特筆すべき点を以下に挙げます:

というわけで、無事移植でき、上述の高速化後の simdlib_emulated 比で15倍程度高速化できました。 16要素のSIMDで15倍程度出てると考えると4bit PQのSIMDとの親和性の高さが伺えますね。 元々のコードではaarch64上で遅い simdlib_emulated が動いていたので、そこからだと合わせて60倍程度の高速化が達成できました。

まとめ

というわけで、Faissの4bit PQをaarch64上で60倍程度高速化した話でした。 内訳としては、まず simdlib_emulated を4倍速くして、そこから更にSIMD化することで15倍程度高速化した、という形でした。 今回 simdlib_emulated を妥当な速度まで高速化したので、今後他のアーキテクチャ向けに simdlib を移植してもおそらく60倍みたいなことにはならないと思います。 本稿では4bit PQの高速化について紹介しましたが、今回の松井先生からのFaiss高速化のご依頼では他にもいくつか高速化やパッケージングなどに取り組みました。一覧はこちらとなりますのでご興味のある方はご覧ください。

フィックスターズはソフトウェア高速化を専門としている会社です。 OSSの高速化や、最新のC++やRustで記述されたソフトウェアの高速化も大歓迎です。 また、ソフトウェア高速化に関する講演等も高速化と合わせて承っております。 お気軽にお問い合わせください。

Fixstars Tech Blog /proc/cpuinfo
Fixstars Tech Blog /proc/cpuinfo

このブログは、高速化のエキスパートであるフィックスターズのエンジニアたちが、技術的な知見や検証結果を自由に発信しているテックブログです。「ソフトウェアの高速化・最適化」と「最新ハードウェア/AI技術の活用」に焦点を当てた、専門性の高い記事が中心です。

Discussion