はじめに
・・だいぶ遅れてしまいましたが投稿します。
ArmアーキテクチャでのSIMD命令セットである、NEONとSVEについて、VLA設計の観点からまとめてみます。
自身の勉強も兼ねたテーマ設定ですが、よろしくお願いします。
ArmのSIMD命令(NEONとSVE)
Armアーキテクチャの代表的なSIMD命令セットにNEONとSVE、SVE2があります。これらは一般的には次のような特徴があると言われています。
-
NEON
- 導入世代:Armv7以降
- ベクトル幅:128bit固定
- 用途:画像処理・ディジタル信号処理などに利用される
-
SVE
- 導入世代:Armv8.2以降
- ベクトル幅:可変
- 用途:科学計算や大規模ベクトル計算などに利用される
-
SVE2
- 導入世代:Armv9以降
- ベクトル幅:可変
- 用途:NEONで得意だったマルチメディア処理が効率的に実行できる
ベクトル幅はCPU実装ごとに異なり、固定ではありません。
例えば、富士通のA64FXでは128bit・256bit・512bitといった複数のベクトル幅をサポートしています。
このような違いに対応するため、SVEではVLA(Vector Length Agnostic)というプログラミングモデルが採用されています。これにより、特定のベクトル幅に依存せずコードの記述が可能となります。結果として、以下の利点があります。
- 柔軟性:実装ごとのベクトル幅に合わせて自動的に最適化される
- 移植性:将来より広いベクトル幅を持つCPUが登場しても同じソースコードで対応可能
SVE2は、NEONで提供されていた整数演算、ビット操作、暗号処理などを統合し、マルチメディア処理を効率的に実行できるよう設計されています。
命令列の違い
簡単なコードでSIMD命令を見てみます。
void func(double* restrict a, double* restrict b, int size) {
for(int i = 0; i < size; ++i) {
b[i] += a[i];
}
}
NEON
$ clang -c -S -O3 -march=armv8-a+simd -o a_neon.s a.c
.LBB0_4: // =>This Inner Loop Header: Depth=1
ldp q0, q3, [x11, #-16]
subs x12, x12, #4
ldp q1, q2, [x10, #-16]
add x10, x10, #32
fadd v0.2d, v1.2d, v0.2d
fadd v1.2d, v2.2d, v3.2d
stp q0, q1, [x11, #-16]
add x11, x11, #32
b.ne .LBB0_4
NEONの場合、q0~q3の128bit固定幅のNEONレジスタを使います。
ldp/stpでpredicateレジスタ(q0, q1)によって常に決まった要素数(double x2)を処理します。
subs x12, x12, #4で、残りの要素数を固定幅単位で減算してループを回しています。
SVE
$ clang -c -S -O3 -march=armv8-a+sve -o a_sve.s a.c
.LBB0_4: // =>This Inner Loop Header: Depth=1
ld1d { z0.d }, p0/z, [x0, x10, lsl #3]
ld1d { z1.d }, p0/z, [x1, x10, lsl #3]
ld1d { z2.d }, p0/z, [x12, x10, lsl #3]
ld1d { z3.d }, p0/z, [x13, x10, lsl #3]
fadd z0.d, z0.d, z1.d
fadd z1.d, z2.d, z3.d
st1d { z0.d }, p0, [x1, x10, lsl #3]
st1d { z1.d }, p0, [x13, x10, lsl #3]
add x10, x10, x9
cmp x11, x10
b.ne .LBB0_4
SVEの場合、z0~z3の可変長のSVEレジスタを使います。
ld1d/st1dでpredicateレジスタ(p0)によって実装依存の要素数を処理します。
add x10, x10, x9で実際のベクトル幅に応じた要素数でループを回します。
p0によるマスク処理で端数要素も自然に処理でき、ループ末尾の境界処理が不要となります。
SVE2
$ clang -c -S -O3 -march=armv9-a+sve2 -o a_sve2 a.c
.LBB0_4: // =>This Inner Loop Header: Depth=1
ld1d { z0.d }, p0/z, [x0, x9, lsl #3]
ld1d { z1.d }, p0/z, [x1, x9, lsl #3]
ld1d { z2.d }, p0/z, [x11, x9, lsl #3]
ld1d { z3.d }, p0/z, [x12, x9, lsl #3]
fadd z0.d, z0.d, z1.d
fadd z1.d, z2.d, z3.d
st1d { z0.d }, p0, [x1, x9, lsl #3]
st1d { z1.d }, p0, [x12, x9, lsl #3]
incw x9
cmp x10, x9
.ne .LBB0_4
SVE2ではincwというベクトル幅に応じたインクリメントを自動的に行う便利な命令が導入されました。
レジスタはx9がアドレス計算、x11, x12がベースポインタとして割り当てられたことで、簡潔に書くことができています。
Intrinsicの違い
NEON
128bitの固定幅レジスタのため、float32x4型で常に4要素のfloatを使います。命令もすべてq(quadword)の128bit固定です。
void axpy_neon(float a, const float *x, float *y, int n) {
int i;
for (i = 0; i <= n - 4; i += 4) {
float32x4_t xv = vld1q_f32(&x[i]);
float32x4_t yv = vld1q_f32(&y[i]);
float32x4_t av = vmulq_n_f32(xv, a);
yv = vaddq_f32(yv, av);
vst1q_f32(&y[i], yv);
}
}
vld1q_f32 -> 4要素ロード
vmulq_n_f32 -> スカラーとの乗算
vaddq_f32 -> ベクトル加算
vst1q_f32 -> ストア
ループ展開は4要素に固定されています。そのため、4で割ったあまりの要素については、端数処理が別途必要になります。
SVE/SVE2
可変幅ベクトルでハードウェアに依存する128~2048bitまで可変の命令を扱えます。
ペイロードマスク(svbool_t pg)はどの要素を有効にするかを示し、端数処理も自動的にマスクで対応します。
void axpy_sve(float a, const float *x, float *y, int n) {
int i = 0;
svbool_t pg;
for(; i < n; i += svcntw()) {
pg = svwhilelt_b32(i, n);
svfloat32_t xv = svld1(pg, &x[i]);
svfloat32_t yv = svld1(pg, &y[i]);
svfloat32_t av = svmul_n_f32_x(pg, xv, a);
yv = svadd_f32_m(pg, yv, av);
svst1(pg, &y[i], yv);
}
}
svld1(pg, &x[i]) -> マスク付きロード
svmul_n_f32_x(pg, xv, a) -> スカラーとの乗算
svadd_f32_m(pg, yv, av) -> マスク付き加算
svst1(pg, &y[i], yv) -> マスク付きストア
svcntw()がこの環境で一度に処理できるfloat32の数を返すため、ループはその単位で進みます。そのため、端数処理を含めて一つのループで完結します。
Gather/Scatterの違い
さらにIntrinsicの違いを見るために、Gather/Scatter命令を見てみます。
Gather: 複数のメモリアドレスからデータを集めて、1本のベクトルレジスタにまとめてロードする命令
Scatter: ベクトルレジスタの値をインデックスで指定されたバラバラのメモリアドレスへ書き戻す命令
NEON
基本的に連続ロード・ストアのみサポート。
以下のように擬似的なGather関数を作ることは可能ですが、固定長なので要素数が固定されており、扱いにくいです。
float32x4_t gather_neon(const float *x, const int *idx) {
return (float32x4_t) {
x[idx[0]], x[idx[1]], x[idx[2]], x[idx[3]]
};
}
SVE
ネイティブにサポート。svld1_gather_index, svst1_scatter_indexでインデックスベクトルを使った非連続アクセスが可能です。
- gather
for (int i = 0; i < n; ) {
svbool_t pg = svwhilelt_b32(i, n);
svuint32_t idx = svld1(pg, &indices[i]);
svfloat32_t v = svld1_gather_index(pg, base, idx);
svst1(pg, &out[i], v);
i += svcntw();
}
- scatter
for (int i = 0; i < n;) {
svbool_t pg = svwhilelt_b32(i, n);
svuint32_t idx = svld1(pg, &indices[i]);
svfloat32_t v = svld(pg, &values[i]);
svst1_scatter_index(pg, base, idx, v);
i += svcntw();
}
Reductionの違い
もう一つ、Reduction演算について見てみます。Reductionとは、あるベクトルの要素の総和や乗算など、横方向に畳み込んで一つの値にする操作です。
NEON
固定幅なので手動でreduceする必要があります。あわせて端数のループをスカラ処理する必要があります。
size_t i = 0;
float32x4_t acc = vdupq_n_f32(0.0f);
for(; i + 4 <= n; i += 4) {
float32x4_t v = vld1q_f32(&a[i]);
acc = vaddq_f32(acc, v);
}
float sum = vaddvq_f32(acc);
for(; i < n; ++i) {
sum += a[i]; // 4で割り切れない端数のループを処理
}
SVE
svcntw()がこの環境で一度に処理できるfloat32の数を返すので、端数のループの扱いが不要となり、簡潔に書くことができます。
size_t i = 0;
svfloat32_t acc = svdup_f32(0.0f);
while(i < n) {
svbool_t pg = svwhilelt_b32(i, n); // predication 残りの要素に応じて有効レーンを決める
svfloat32_t v = svld1(pg, &a[i]);
acc = svadd_f32_m(pg, acc, v);
i += svcntw(); // SVEのレーン数だけ進める(可変長)
}
float sum = svaddv_f32(svptrue_b32(), acc); // reduction ベクトル->スカラー
終わりに
今回は、Arm SVEの最大の特徴であるVLA(Vector Length Agnostic)についてコード例から見てみました。
NEONが固定幅のSIMDを前提とするため、コンパイラは端数処理を別途生成することになりますが、SVEは可変幅を前提とした仕組みで、predicateやsvcntw命令によって端数処理を吸収することができます。
これによりコード(アセンブリ)が簡潔となることもわかりました。
読んでいただきありがとうございました。メリークリスマス&良いお年をお迎えください。