近似最近傍探索ライブラリFaissをArm SVEに対応させました

ソリューション第五事業部の今泉です。 この度Meta Researchが公開している近似最近傍探索ライブラリFaissをArm SVEに対応させるPR(#2886)を作成し、無事マージされました。 まず本稿ではFaissやArm SVEについて、現状の高速化具合、FaissをArm SVEに対応させるに至った経緯について述べます。 FaissのSVE化にまつわる技術的詳細はもう1つ別の記事を後日執筆予定です。

近似最近傍探索、FaissとArm SVEについて

この節では近似最近傍探索とFaiss、そしてArm SVEについて簡単に説明します。

近似最近傍探索について

あるベクトル空間の中からそのベクトル空間に含まれないクエリベクトルに近いものを高速かつ・または省メモリに探索するアルゴリズムです(厳密な最近傍ベクトルを求めるわけではありません)。 ここで、文書や画像など世の中には多数のテータがありますが、こうしたデータからうまく特徴量(元データの特徴を表すベクトル)を抽出することで近似最近傍探索の対象とすることができます。 これにより特徴量に基づいて類似した文書や画像などを探索する、といった応用ができます。

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

Faissについて

FaissはMeta Researchが公開している有名な近似最近傍探索ライブラリのOSSです。 C++とPythonで使うことができ、多数の近似最近傍探索アルゴリズムが実装されています。 特にx86_64とCUDAにおけるバッチ探索(複数のクエリベクトルをまとめて検索する処理)に向けて高速化されています。 3年前にARMv8向けに特定アルゴリズムを高速化しました。

Arm SVEについて

Arm SVE(Scalable Vector Extension)はArmv8の拡張命令セットであるベクトル命令です。 また、Armv9ではSVEを拡張したSVE2が基本命令セットに含まれるため、Armv9のCPUでは常にSVEが使えます。 立ち位置としてはAdvanced SIMD(ARM NEON)の後継の命令セットと見做して差し支えないかと思います。 富岳にも載っている富士通A64fxや、2022年5月から使用可能になったAWS EC2 c7gインスタンスに載っているAWS Graviton3などの他、Armv9が最近のフラッグシップスマートフォンのCPUを中心に採用されてきており、SVEが使えます。 データ並列を高速に処理するための命令セットとしては現時点ではまだNEONほど多くのCPUで使えるわけではありませんが、今後使える場所が確実に増えていくでしょう。

経緯

この節ではFaissをArm SVEに対応させるに至った経緯について簡単に述べます。

弊社ではエンジニア有志で社内勉強会を定期・不定期を問わず開催しているのですが、2022年の定期勉強会としてArm SVE勉強会を開催しました。 2年前の5月末にSVEが使えるAWS EC2インスタンスであるC7gインスタンスが公開されたこともあり、参加者各々が色々なコードをSVEで記述して知見を共有する会となりました。 この際私は3年前にArm NEONで高速化したFaissの存在を思い出し、FaissのSVE化を試行しその実装知見を共有しました。

その後AWS Japan様からお声がけいただき、AWS 秋の Amazon EC2 Deep Dive 祭り 2022というセミナーイベントに登壇させていただき、FaissのSVE化をケーススタディとしてGraviton3の性能を引き出すためにArm SVEが重要であること、intrinsicsを用いたArm SVEプログラミングの概要について発表させていただきました。 この際の発表内容については上記リンク先に動画・スライド共に公開されておりますので、ぜひご覧ください。

上述のセミナーにおいてSVE版のFaissについて公式リポジトリへのPRを予定しているとしていたのですが、その後諸般の都合で中々公開にこぎつけられず、半年以上経過してようやくFaissの公式リポジトリに対してPRの作成に至りました。 結局実装してから公開するまでに1年近くかかってしまったわけですが、PR作成後も4ヶ月ほどfat binary packageに関する議論などが続き、最終的にはCI環境が整うまでpendingとなってしまいました。 CI環境が整ったのはPR作成からほぼ1年後、今年の5月末。8ヶ月近くぶりの連絡を受け取った頃には私は弊社GitHub Organizationへのアクセス権や開発環境などを既に失っていたため、それらの再確保から作業を再開し、今回晴れて実装から2年越しに公式リポジトリにマージされた次第です。

まとめ

簡単でしたが、Meta Researchの近似最近傍探索ライブラリFaissをArm SVEに対応させた話でした。

今後の展望としては、実装のSVE化を順次進めていきたいです。 現時点ではコンパイラフラグでSVEを有効化できるようにしただけで、auto vectorizationによるSIMD化が行われている一部の関数のみがSVE化した状態であり、SVEのintrinsicを用いた実装を行ったわけではありません。 いくつかの関数については過去に私がSVE化したりインターン生にSVE化してもらったものがあるので、最新のFaissに対応させて別途PRを作成予定です。 参考程度ですが、昨年時点ではGraviton3上でSIFT1Mデータセットをnprobeが16のIndexIVFPQ + IndexHNSWFlatで探索するタスクが2.25倍程度高速化していた実績があるので、乞うご期待ください。 また例えば3年前に扱った4bitPQについてはまだSVE化の目処が立っていないため、継続して作業していく予定です。 今後もより多くの実装をSVEに対応させ、FaissをSVE環境下でさらに高速にしていきたいですね。

また、弊社はArm SVEを用いたソフトウェア高速化も承っております。 Arm SVEを用いたソフトウェア高速化にお悩みの方はぜひご連絡ください。

最後に繰り返しとなりますが、Arm SVEそれ自体やFaissのArm SVE化に関する技術的詳細や、上述の別MRでどの程度高速化したかについては後日別の記事で扱う予定ですので、しばらくお待ち下さい。

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