このブログは、株式会社フィックスターズのエンジニアが、あらゆるテーマについて自由に書いているブログです。
ソリューション第五事業部の今泉です。 2か月前に公開した記事の通りMeta Researchが公開している近似最近傍探索ライブラリFaissをArm SVEに対応させました。これは無事v1.9.0のリリースに入ったようです。
先日intrinsicsを用いていくつかの関数を実装し高速化したPR(facebookresearch/faiss#3933)を作成したので、本記事では高速化具合やArm SVE周りの知見をまとめようと思います。
このブログを読みに来る皆様が一番知りたいであろう情報を先にお伝えします。 今回もSIFT1Mデータセットに対する探索の実行時間計測を行いました。 比較対象はfaissの現在の `main` ブランチのSVEビルド(これには自動ベクトル化によるSVE命令とNEONのintrinsicsによるSIMD命令が含まれます)と今回私が作ったPRのSVEビルドです。 計測環境はAWS EC2インスタンスのGraviton 3(c7g.large)とGraviton 4(r8g.large)ですが、Graviton 3では最大2.4倍程度の高速化を達成したのに対して、Graviton 4では最大でも1.3倍程度の改善に留まる結果となりました。
この節ではArm SVEのintrinsicsを用いたプログラミングについて、実装から得られた知見を中心に述べます。
先述の通り、当該のPRは実行環境によって高速化度合いに大きな差があります。 これは一体何が原因なのでしょうか? 実際にはきちんと調べきれたわけではないのですが、考えられる大きな要因の一つとして、Graviton 3のSVEベクトル長は256bitなのに対してGraviton 4のSVEベクトル長は128bitである、という点が挙げられると思います。
Graviton 3、およびそのベースとなったNeoverse V1ではSVEのベクトル長は256bitでした。 つまりCPUには256bitレジスタが積んであり、演算器も256bit演算器が積んであるということです。 そしてArm SVEに対応したCPUでは、Arm NEON(Advanced SIMD)を使う際SVEレジスタの下位128bitをNEONのレジスタとして扱います。 つまり、Graviton 3ではNEONを使っているとCPUダイの結構な面積(少なくとも256bitレジスタの半分)が使えなくなるということです。 またNeoverse V1の命令レイテンシ・スループットの表を眺めると、多くの命令においてSVEに対してNEONが2倍のスループット値を出すようになっているので、一見理論性能的にはNEONでも問題なく性能が出せるように見えますが、メモリの読み書きについてはこの差が小さくなっているため、時間あたりのデータ量の観点ではSVEに分があることがわかります(同じスループットの値でもSVEの方が一度に2倍のデータを扱うため、データ量の観点ではSVEの値がNEON比で2倍になることに注意)。 これはベクトル長が長いことの恩恵と言えるでしょう。 こうした点が効いているのか、私の経験上でもGraviton 3ではしっかりSVEを使ったほうがパフォーマンスが向上する印象があります。 (これは裏を返せばSVEを使わないとGraviton 3は性能が出し切れないということでもあります)
一方、Graviton 4およびそのベースとなったNeoverse V2ではSVEのベクトル長は128bitです。 これは近年のArmの方針のようで、多くの既存のprebuiltバイナリはSVEを用いていないため、「SVEを使うと速いCPU」を作っても多くのバイナリは(残念ながら)現時点ではその恩恵に預かれません。 NEONも十分に速く動作するCPUを作りたい、となると”余剰なダイの余白”はもったいないので、そのような部分が存在しないようにSVEのベクトル長を最低長の128bitに設定して、代わりに同時発行可能命令数を倍に増やすことで「NEONでもSVEでも速いCPU」を目指す、という考えなのだと思います。 ところでベクトル長が短いということはメモリアクセス周りを改善するのは難しくなります。 実際Neoverse V2のレイテンシ・スループットの表を眺めると、基本的にNEONとSVEに差はありません。むしろpredicate loadに至っては遅いです。
ということで、今後のArm CPUがどう転ぶかはわからないのですが、SVEを使うと速くなる…と言い切れるCPUというのは、しばらく出ないかもしれません。 個人的には面白みに欠ける状況なのですが…スカラー処理や既存のpacked SIMDな処理が高速に動作したほうが当然恩恵は大きいので、現実的なことを考えると致し方ないのかもしれません。 とはいえモバイル向けならいざ知らず、せめてサーバー向けCPUなら256bitレジスタくらいは積んでほしいですが…。 またCPUに関わらずSVEが幾分か効きそうなアプリケーションとしては、gather/scatterやcompactのようなSVEの新規命令を使う類のものが(従来はスカラー処理でなんとかしていた部分なので)ありそうです。
Arm SVEでsaxpyを書くと以下のようになります:
#include <arm_sve.h>
void saxpy(float* z, float a, const float* x, const float* y, std::size_t n){
const std::size_t lanes = svcntw();
for(std::size_t i = 0; i < n; i += lanes) {
const auto mask = svwhilelt_b32_u64(i, n);
const auto xi = svld1_f32(mask, x + i);
const auto yi = svld1_f32(mask, y + i);
const auto zi = svmla_n_f32(mask, yi, xi, a);
svst1_f32(mask, z + i, zi);
}
}
しかし、実は
WHILELT
命令はあまり速くない
といった点から、実際にはArm SVEでも以下のように端数部の処理を別に書いた方が速度が出ることが多いです:
#include <arm_sve.h>
void saxpy(float* z, float a, const float* x, const float* y, std::size_t n){
const std::size_t lanes = svcntw();
std::size_t i = 0;
for(; i + lanes < n; i += lanes) {
const auto mask = svptrue_b32(); // ループ内はマスクしない
const auto xi = svld1_f32(mask, x + i);
const auto yi = svld1_f32(mask, y + i);
const auto zi = svmla_n_f32(mask, yi, xi, a);
svst1_f32(mask, z + i, zi);
}
const auto mask = svwhilelt_b32_u64(i, n); // 最後の1回だけマスクする
const auto xi = svld1_f32(mask, x + i);
const auto yi = svld1_f32(mask, y + i);
const auto zi = svmla_n_f32(mask, yi, xi, a);
svst1_f32(mask, z + i, zi);
}
ループの回数が少なく、ループのbodyが命令キャッシュにピッタリ収まったときなどはもしかすると遅くなることもあるかもしれません。
最後は実際のデータを用いて実測するしかないですが、一般にSIMDな処理を書くときはデータサイズ n
は十分長いはずなので、このようにした方が速くなることが多いでしょう。
SVEのベクトル長はCPU毎に異なるので、基本的にコンパイル時点ではArm SVEのレジスタ型のバイト数はわかりません。
つまり、 sizeof(svfloat32_t)
を得ることができません(コンパイルエラーになります)。これが原因で以下の操作ができません:
svfloat32_t* ptr;
に対して ptr + 1
すると、 ptr
の指すアドレスの値は(単純化すれば) reinterpret_cast<svfloat32_t*>(reinterpret_cast<std::uintptr_t>(ptr) + sizeof(svfloat32_t))
相当の値になるはずです。ところで sizeof(svfloat32_t)
はコンパイル時に決まらないので、このような操作も同様に許可されません(コンパイラがオフセットを計算できない)。svfloat32_t arr[2];
に対して arr[1]
するのは *(arr + 1)
と同義です。上述の通りポインタに対する算術演算は許可されていません。以下のような構造体を考えるとわかります:
struct S { svfloat32_t vec; // Sの先頭0バイト目から???バイト目まで int another_member; // Sの先頭何バイト目から始まるのかコンパイル時に決定不可 };
このような仕様から、レジスタブロッキングのためのコードやラッパーライブラリの実装が素直にできません。
一応-mse-vector-bits=
コンパイラオプションやarm_sve_vector_bits
attributeなどを使うことでコンパイル時にレジスタ長を固定すること自体はでき、これによって上述の制限は全て消えます。
ですが、異なるベクトル長を持つCPU向けにベクトル長に依存したコードを書く場合当然それぞれ実装を別に記述しなければなりません。
オプションがサポートしている範囲でも5通り(128
, 256
, 512
, 1024
, 2048
bit毎)の実装を書かねばならず、加えてこれら以外のベクトル長を持つCPU向けにはそもそもオプション指定できないため、個人的にはこのオプションはあまり積極的に採用しようとは思えない存在です。
そもそもSVEのSはScalable(可変長ベクトル)のSですし…
この節ではFaissをSVE化するにあたっての難しさや設計上の選択に関して述べます。
弊社のソフトウェア高速化業務は多くの場合、実行環境が決まっていて、そのプロセッサのハードウェアスペックの限界に向けて処理を詰めていきます。 多くの場合、入力されるデータのサイズについてもある程度決まっていることが多く、それ故込み入った高速化を行うことができます。
一方、Faissは実行環境についてあまり仮定を置けません。わかり易い例で言えばSVEのベクトル長は実行時まで確定しませんし、他にもレジスタリネーミングに使われるレジスタの数がいくつあるのか、データキャッシュはどの程度あるのか、各命令のレイテンシとスループットはいくつなのか、といったことはわからないのです。 また、実際に入力されるデータセットも小さいものから大きいものまで様々です。 高速化にあたって命令レベル並列性を稼ぐためにループアンロールしてレイテンシを隠蔽する、といったことをよく行いますが、何段のアンロールが適切なのかはCPUによってまちまちです。 自分の中にある程度の指針ができるまではしばらく悩みましたし、Graviton 3では速くなるコードがGraviton 4では遅くなるなんてこともあり、どの環境でも(それなりに)速く動作するコードを書くことの難しさを改めて感じました。
また速度以前の問題として、一見問題無い実装でもGraviton 3では問題なく動作するがGraviton 4では正常に動作しない、といった事象に複数回遭遇しました。今回の問題について実際のところ本質的な原因がなんなのかはしっかりと調べたわけではないのですが、いずれにせよレジスタサイズの異なるCIマシンを複数台揃えてきちんとCIを回し続けるなどの施策が肝要だと思います。
C++のライブラリであるFaissがAVX2やSVEのような拡張命令に対してFaissがどのように有効・無効を切り替えているかというと、各拡張命令の有効・無効毎に異なるバイナリとしてビルドします。
現在のFaissだとArm上では -DFAISS_OPT_LEVEL=sve
を付けてビルドすることで libfaiss.a
と libfaiss_sve.a
が生成されるため、好みの環境向けにリンクするライブラリを選べます。
さて、FaissはC++のライブラリですが、SWIGとラッパーモジュールによるPython bindingsを内包しており、Pythonのパッケージとして使うこともできます。
PythonパッケージとしてのFaissは _swigfaiss.so
と _swigfaiss_sve.so
をパッケージ内に内包し、実行時にどの動的ライブラリからimportするかを自動で選択します。
ここで先述の話に戻るのですが、もし -mse-vector-bits
コンパイラオプションでベクトル長を固定してしまうと、 _swigfaiss_sve128.so
や _swigfaiss_sve256.so
などが大量に生成されてしまい、SVE対応のためにArm版FaissのPythonパッケージのバイナリサイズが急激に増えてしまいます。
既存のPRでのやり取りなどを踏まえてこの選択は取らないほうが良いだろうと判断したため、現在のFaissではベクトル長は固定していません。
上述の通りFaissではSVEのベクトル長を固定していないので、配列やクラスメンバとしてSVEのレジスタ型を扱うことができません。
例えば3年前にNEONで高速化したsimdlibは、SIMDレジスタをメンバに持つクラスにoperator overloadや各種メンバ関数を実装することでSIMDを比較的扱いやすくしたライブラリでした。 つまり直接simdlibをSVEに対応させることはできません。 この問題はベクトル長に依存しない新規なラッパーライブラリを開発・利用することで(AVX-512対応すらも含めた)解決が図れるのですが、比較的規模も大きい上に実験的な内容のため、業務の片手間で進めるには困難を伴いました。 そこで学生アルバイトの方にお願いして実装検討を進めてもらっていたのですが、残念ながら形になる前に他のタスクをお願いすることが決まってしまったため、現時点では僅かながら実装知見が得られた程度の状況です。 それ故4bitPQはまだSVEに対応させられていない状況なのですが、この問題はMeta Researchの人たちも把握し始めているので、もしかすると画期的な解決がなされるかもしれません。
また、レジスタブロッキングのためのコードでは配列を用いることが多いため、こちらもSVE対応はやや難易度が高いです。 しかしこちらについては昨年インターンシップに取り組んだ片山悠哉さんの活躍で 配列を用いることなくコンパイル時定数に応じた数のSVEレジスタを取り回す 方法が確立済みで、この手法を用いたPRについても近日中にFaiss公式リポジトリに作成予定です。
Meta Researchの近似最近傍探索ライブラリFaissをArm SVE向けに高速化してきた話でした。 こうして振り返ってみると、(社外に公開していない頃から含めて)2年くらいはSVEでFaissの高速化を続けてきて、結構な実装知見が溜まっていたのだなと感じます。 そして、こうした作業の成果がなんとか世に出せてホッとしています。
弊社はArm SVEを用いたソフトウェア高速化も承っております。 Arm SVEを用いたソフトウェア高速化にお悩みの方はぜひご連絡ください。
keisuke.kimura in Livox Mid-360をROS1/ROS2で動かしてみた
Sorry for the delay in replying. I have done SLAM (FAST_LIO) with Livox MID360, but for various reasons I have not be...
Miya in ウエハースケールエンジン向けSimulated Annealingを複数タイルによる並列化で実装しました
作成されたプロファイラがとても良さそうです :) ぜひ詳細を書いていただきたいです!...
Deivaprakash in Livox Mid-360をROS1/ROS2で動かしてみた
Hey guys myself deiva from India currently i am working in this Livox MID360 and eager to knwo whether you have done the...
岩崎システム設計 岩崎 満 in Alveo U50で10G Ethernetを試してみる
仕事の都合で、検索を行い、御社サイトにたどりつきました。 内容は大変参考になりま...
Prabuddhi Wariyapperuma in Livox Mid-360をROS1/ROS2で動かしてみた
This issue was sorted....