Intro to SVE 富岳のA64FXを触ってみた

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

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

• 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

• 富士通が開発したスパコン富岳用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

• 富士通が開発中の深層学習用ライブラリ
• 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

• 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

• 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

• 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

• 配列の計算
• 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

• メイン部分
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]
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

• 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

• 「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

• メイン部分(再掲)
• 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();
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

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

• 先程のコードはループごとに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

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

• 単純ループで他の要素に依存関係がない
• 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

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

• ちょっと面白い現象の紹介
• 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

• 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

• 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

• かなり速くなった
レジスタ割り当てを変えてみた
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));
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

• レジスタ名を変更しなくても大丈夫(この違いはなんだろう)
• 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
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

• 調べたことをまとめ中
• 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

Intro to SVE 富岳のA64FXを触ってみた

Recommended

Recommended

More Related Content

What's hot

What's hot (20)

Similar to Intro to SVE 富岳のA64FXを触ってみた

Similar to Intro to SVE 富岳のA64FXを触ってみた (20)

More from MITSUNARI Shigeo

More from MITSUNARI Shigeo (20)

Recently uploaded

Recently uploaded (14)

Intro to SVE 富岳のA64FXを触ってみた