sample by xbyak_aarch64

1. Intro to SVE
富岳のA64FXを
触ってみた
2020/8/8
ARM入門勉強会
光成滋生

2. • 自己紹介
• A64FXとSVE
• qemuの設定
• xbyak_aarch64で簡単なサンプル
• 述語レジスタ
• ループアンロール
• レジスタ割り当て
• 逆数近似
• レジスタリネーム
目次
2 / 24

3. • https://github.com/herumi/
• mcl/bls ; 暗号ライブラリの開発
• Ethereum 2などで利用されている
• xbyak ; Intel用のJITアセンブラ
• Intelの深層学習ライブラリoneDNNなどで利用されている
• TensorFlow, PyTorchなどのCPU向けバックエンド
• Intel AMX(2020/6/27)の仕様公開と同時にpull reqが来た
• https://github.com/herumi/xbyak/pull/95
• xbyak_aarch64 ; aarch64用のJITアセンブラ by 富士通
• https://github.com/fujitsu/xbyak_aarch64/
• 移植の設計アドバイス、バグとり・機能追加など
• 注 : 勉強中の身 / 中のプロの人ではありません
@herumi
3 / 24

4. • 富士通が開発したスパコン富岳用CPU
• Arm v8-A命令セット+SVEを採用した最初のCPU
• SVE ; SIMD命令セット
• https://static.docs.arm.com/ddi0584/a/DDI0584A_a_SVE_supp_armv8A.pdf
• A64FXでは32個の512-bit SIMDレジスタ ; z0, ..., z31
• int8 x 64, int32 x 16, float x 16, double x 8など
様々なデータ型の並列処理が可能
• 16個の述語(predicate)レジスタ ; p0, ..., p15
• 後述
• AVX-512を知っている人向けの説明 ; maskレジスタ相当
A64FX
4 / 24

5. • 富士通が開発中の深層学習用ライブラリ
• https://github.com/fujitsu/dnnl_aarch64
• IntelのoneDNNの富岳版
• A64FX向けxbyak_aarch64とxbyak_Translatorを利用
• ResNet50 + TensorFlow v2.1.0で
DNNL_aarch64を使うとTraining(FP32)が9.2倍速
• 環境 A64FX 2.2GHz 48core HBM 32GB
• 2020/8/4 MIRU2020松岡聡氏の招待講演 資料p.30より
5 / 24
DNNL_aarch64

6. • 3命令タイプ
• 2命令+述語タイプ
• 3命令の積和(dstをsrcとして利用する)
• movprfx (dstをsrcとして利用する命令の補助)
• 4命令の積和
• movprfxはμOPレベルでは
pack処理されて一つのアーキテクチャ命令になる
SVEの命令概略
op(dst, src1, src2); // dst = op(src1, src2);
op(dst, pred, src); // dst = op(dst, src) with pred
fmad(dst, pred, src1, src2); // dst = dst * src1 + src2
movprfx(dst, pred, src3);
fmadd(dst, pred, src1, src2); // dst = src3 * src1 + src2
6 / 24

7. • Procedure Call Standard for the ARM 64-bit
Architecture (AArch64) with SVE support
• https://developer.arm.com/documentation/100986/latest
• caller save ; 関数を呼び出した側がレジスタを保存
• callee save ; 呼び出された関数の中でレジスタを保存
• SVEレジスタ
• z0, ..., z7 ; free
• z8, ..., z23 ; callee save
• z24, ..., z31 ; caller save
• 述語レジスタ
• p0, ..., p3 ; free
• p4, ..., p15 ; callee save
呼び出し規約
7 / 24

8. • qemu-aarch64 + aarch64-linux-gnu-g++でテスト可能
• インストールの詳細はたとえば
• https://github.com/fujitsu/xbyak_aarch64/tree/master#execution-environment
• テスト
• qemuの注意点
• ライブラリのパスをQEMU_LD_PREFIXで指定
• A64FXのSVEは512-bitなのでqemuオプションで明記
SVEのエミュレータ
>cat t.cpp
#include <stdio.h>
int main() { puts("hello"); }
>aarch64-linux-gnu-g++ t.cpp
>env QEMU_LD_PREFIX=/usr/aarch64-linux-gnu qemu-aarch64 ¥
-cpu max,sve512=on ./a.out
8 / 24

9. • 配列の計算
• sqrAdd
• 2個の配列x, yの各要素に対してz[i]=x[i]^2+y[i]をするCの関数
• これをxbyak_aarch64を使って実装する
• コンパイラに-I <xbyak_aarch64>オプションを指定
• 注 : 私はintrinsicや.sでの書き方をよく知らない
簡単なループ
void sqrAdd(float *z, const float *x, const float *y, size_t n)
{
for (size_t i = 0; i < n; i++) {
z[i] = x[i] * x[i] + y[i];
}
}
9 / 24
git clone -b master git@github.com:fujitsu/xbyak_aarch64

10. • メイン部分
SVEによるsqrAdd(z, x, y, n);
const auto& out = x0;// 読みやすいようレジスタ名のaliasをつける
const auto& src1 = x1;
const auto& src2 = x2;
const auto& n = x3;
Label cond;
mov(x4, 0); // ループ変数を0に初期化
b(cond); // condラベルに無条件ジャンプ
Label lp = L();
ld1w(z0.s, p0/T_z, ptr(src1, x4, LSL, 2));// z0 = src1[x4 << 2]
ld1w(z1.s, p0/T_z, ptr(src2, x4, LSL, 2));// z1 = src2[x4 << 2]
fmla(z1.s, p0/T_m, z0.s, z0.s); // z1 += z0 * z0
st1w(z1.s, p0, ptr(out, x4, LSL, 2)); // out[x4 << 2] = z1
incd(x4); // x4 += 16
L(cond);
whilelt(p0.s, x4, n); // while (x4 < n)なら
b_first(lp); // lpラベルにジャンプ
ret(); // 関数終了
10 / 24

11. • SVEレジスタの各要素を処理する(1)か否(0)かを指定
• 例 ld1w(z.s, p/T_z, ptr(src));
• z = *src; // float x 16個読み込み 「.s」はfloat型
• i番目の各要素(i = 0, ..., 15)について
• 述語レジスタp[i] = 1ならz[i] = src1[i]
• p[i] = 0ならT_z(zero)を指定しているのでz[i] = 0
• T_zを指定しなければp[i]の値を変更しない
述語レジスタ
src x0 x1 x2 x3...
z.s x0 0 x2 x3...
p 1 0 1 1
11 / 24

12. • 「x4 + i < n」が成り立つ添え字までp[i] = 1にする
• 例 x4 + 16 <= nならi = 0, ..., 15についてp[i] = 1
• 全てのデータが有効
• x4 + 3 = nならi ≦ 2についてp[i] = 1, その他p[i] = 0
• p[i] = 0の部分はデータを読まない・書かない
• 読み書き属性が無い領域でも大丈夫
whilelt(p.s, x4, n);
i 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
p 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
i 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
p 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0
12 / 24

13. • メイン部分(再掲)
• x4 + 16 ≦ nである限りp0[i] = 1 for i = 0, ..., 15
• ループの最終ではp0[i] = 1 for i <(n % 16), p0[i] = 0(otherwise)
• メリット
• SVEが256bitや1024bitでも同じコードで動く
• SVEはScalable Vector Extensionの略
ループの終わり部分
Label lp = L();
ld1w(z0.s, p0/T_z, ptr(src1, x4, LSL, 2));// z0 = src1[x4 << 2]
ld1w(z1.s, p0/T_z, ptr(src2, x4, LSL, 2));// z1 = src2[x4 << 2]
fmla(z1.s, p0/T_m, z0.s, z0.s);
st1w(z1.s, p0, ptr(out, x4, LSL, 2)); // out[x4 << 2] = z1
incd(x4); // x4 += 16
L(cond);
whilelt(p0.s, x4, n); // while (x4 < n)なら
b_first(lp); // lpラベルにジャンプ
13 / 24

14. • 1024個のfloat(ループは1024/16=64回)
• 1ループあたり4.3nsec @ FX700
• 2GHzなら8.6clk(clock cycle)
ベンチマーク
C SVE
360nsec 277nsec
14 / 24

15. • 先程のコードはループごとにp0レジスタを更新する
• ループ最終以外は定数(全て1)
• 先程のループの前に次のコードを追加する
述語レジスタへの依存除去
ptrue(p0.s); // p0を全て1にする
Label skip;
b(skip);
Label lp = L();
ld1w(z0.s, p0/T_z, ptr(src1));
add(src1, src1, 64);
ld1w(z1.s, p0/T_z, ptr(src2));
add(src2, src2, 64);
fmla(z1.s, p0/T_m, z0.s, z0.s);
st1w(z1.s, p0, ptr(out));
add(out, out, 64); // 512-bit(64byte)ずつ増やす
sub(n, n, 16); // カウンタを16ずつ減らす
L(skip);
cmp(n, 16);
bge(lp); // n >= 16であるかぎりループ
15 / 24

16. • 約2倍の高速化
• 述語処理は結構重かった?
• ただしSVEが512bit固定になった(対応は可能)
ベンチマーク2
C SVE p0固定
360nsec 277nsec 134nsec
16 / 24

17. • 単純ループで他の要素に依存関係がない
• N=2,3,4 ; レジスタを(z0, z1), (z2, z3), ...として使う
ループアンロール
Label lp = L();
for (int i = 0; i < N; i++) {
ld1w(ZReg(i * 2).s, p0/T_z, ptr(src1, i));
ld1w(ZReg(i * 2 + 1).s, p0/T_z, ptr(src2, i));
fmla(ZReg(i * 2 + 1).s, p0/T_m, ZReg(i * 2).s, ZReg(i * 2).s);
st1w(ZReg(i * 2 + 1).s, p0, ptr(out, i));
}
add(src1, src1, 64 * N);
add(src2, src2, 64 * N);
add(out, out, 64 * N);
sub(n, n, 16 * N);
L(skip);
cmp(n, 16 * N);
bge(lp);
17 / 24

18. • ループアンロール
• Cに比べて約4.5倍の高速化
ベンチマーク3
C SVE p0固定 N=2 N=3 N=4
360 277 134 107 90 79
18 / 24

19. • ちょっと面白い現象の紹介
• floor命令 ; frintm(dst, p, src); // dst = floor(src);
• 除算命令 ; fdivr(dst, p, src); // dst = src / dst;
• 98clk latency!
• fadd, fmul, fmadなどは9clk, 論理演算は3~4clk
レジスタ割り当てと速度の変化
void func(float *z, const float *x, const float *y, size_t n) {
for (size_t i = 0; i < n; i++) {
z[i] = 1 / (floor(x[i]) + y[i]);
}
}
19 / 24

20. • frecpeは1/2^9程度の近似演算(4clk)
• frecpsはNewton-Raphson法の補正計算(9clk)
• frecps + fmulをもう一度すると精度がfloatに近い
• これを使うと大分速くなるのでは?
逆数近似命令
frecpe(t1, x); // t1 = xの逆数近似
frecps(t2, x, t1); // t2 = 2 - x t1
fmul(x, t1, t2); // x = (2 - x t1)t1
frecpe(t1, x); // t1 = xの逆数近似
frecps(t2, x, t1); // t2 = 2 - x t1
fmul(t1, t1, t2); // t1 = (2 - x t1)t1
frecps(t2, x, t1); // t1の再補正
fmul(x, t1, t2); // better 1/x
20 / 24

21. • https://github.com/herumi/misc/blob/master/sve/inv.cpp
• ? あんまり速くならない
• frecps x 2はfrecps + fmul足しただけなのにclkかかり過ぎ?
• 51 - 33 = 18 = 9 x 2?
ベンチマーク
fdiv frecps x 1 frecps x 2
clk 100 33 51
21 / 24

22. • かなり速くなった
レジスタ割り当てを変えてみた
fdiv A : frecps x 1 A': frecps x 1 B : frecps x 2 B': frecps x 2
clk 100 33 10.9 51 11.2
ld1w(z0, p0/T_z, ptr(src1));
ld1w(z1, p0/T_z, ptr(src2));
frintm(z2, p0, z0); // floor(src1[i])
fadd(z0, z1, z2); // floor(src1[i]) + src2[i]
fdivr(z0, p0, one);
frecpe(z1, z0);
frecps(z2, z0, z1);
fmul(z0, z1, z2);
frecpe(z1, z0);
frecps(z3, z0, z1);
fmul(z0, z1, z3);
frecpe(z1, z0);
frecps(z2, z0, z1);
fmul(z1, z1, z2);
frecps(z2, z0, z1);
fmul(z0, z1, z2);
frecpe(z1, z0);
frecps(z3, z0, z1);
fmul(z1, z1, z3);
frecps(z3, z0, z1);
fmul(z0, z1, z3);
A →A'
B →B'
22 / 24

23. • レジスタ名を変更しなくても大丈夫(この違いはなんだろう)
• A, A', B, B'は同じ
frintm→faddにしたら
fdiv A : frecps x 1 A': frecps x 1 B : frecps x 2 B': frecps x 2
frintm 100 33 10.9 51 11.2
add 100 7.5 7.5 11.1 11.1
ld1w(z0, p0/T_z, ptr(src1));
ld1w(z1, p0/T_z, ptr(src2));
fadd(z2, z0, z0); // src1[i] + src[i]
fadd(z0, z1, z2); // (src1[i] + src[i]) + src2[i]
23 / 24
Z0
Z1
Z2 Z0 Z1 Z2 Z1 Z2 Z0
Z3 Z3
rename
こうしないとfrintmでは遅くなる
faddなら大丈夫
register dependency

24. • 調べたことをまとめ中
• https://github.com/herumi/blog
• いろいろ実験コード
• https://github.com/herumi/misc/tree/master/sve
• exp, tanh for AVX-512をSVEに移植したものなど
• 富岳用DNNLのeltwiseまわりの実装にcommit
• https://github.com/fujitsu/dnnl_aarch64/
その他
24 / 24