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

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版があるので同じインターフェースになるよう移植すればよいのですが、特筆すべき点を以下に挙げます:

• SIMDレジスタ型2要素をまとめて扱う
  • 当たり前ですが、AVX2のラッパーとして生まれた simdlib は256bit SIMDを前提に設計されています。
  • 一方、ARM SIMDは最大で128bitなので、 simdlib_neon 内では2本のSIMDレジスタ型の変数をまとめてクラスにしています。
    • より具体的には、uint16x8x2_t などの型をメンバに用いています。
• 基底クラス simd256bit の排除
  • AVX2版ではunion{ __m256i i; __m256 f; } をメンバに持つクラスsimd256bitの派生クラスとして simd16uint16 などのクラスが作ってあります(simdlib_emulated も同様の設計思想に基づくので配列の共用体で同様の実装になっています)。
    • これにより他のレーン幅のSIMDクラスとのreinterpretな変換が simd256bit のコンストラクタをつくることで実現できます。
  • 一方、ARM SIMDのintrinsicでは uint16x8_t など、要素の型とレーンの数に応じて最初から異なるSIMDレジスタ型が提供されているので、このように単純にはいきません(type punningに目を瞑ればたぶんいけるのですが…)。
    • こうなると simd256bit を継承する設計はあまりうれしくないので、各クラス間の基底クラスは削除しました。
    • 従来 simd256bit のメンバ関数として実装されていたものは適宜 detail 名前空間上に移動し、各クラスは処理を委譲するだけにしてあります。
  • 今回はそれぞれの型(uint8x16x2_t , uint16x8x2_t , uint32x4x2_t , float32x4x2_t)からそれぞれの型へのreinterpretな変換を同名の関数として事前に定義しておき、従来 simd256bit の派生クラスだった各クラスがコンストラクタに来た場合は当該関数を呼び出すことでreinterpretな変換を行えるようにしています。
    • simd16uint16 などの従来 simd256bit の派生クラスだった型を判定するためにis_simd256bit というメタ関数を追加しています。
• ARM SIMDに存在しないintrinsicを仮定したコード
  • AVX2版で _mm256_movemask_epi8 を用いたコードはマスク情報を32bit非負整数型で返すようになっていますが、ARM SIMDにはそのようなintrinsicは存在しないため、マスク情報をうまく32bit非負整数型に詰める必要があります。
    • これは頑張って実装しました。無いものは無いので
      • ただし、この実装によってbig endianのサポートができなくなっています。まぁbig endianのaarch64 CPU見たことありませんが…

というわけで、無事移植でき、上述の高速化後の 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で記述されたソフトウェアの高速化も大歓迎です。 また、ソフトウェア高速化に関する講演等も高速化と合わせて承っております。 お気軽にお問い合わせください。

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください。