expの実装に見る
AVX-512とSVEの違い

2020/12/3

第9回HPC-Phys勉強会
光成滋生

目次
• expの近似計算• AVX-512による実装• SVEによる実装• レジスタリネーミング• SVEのfexpaを使った近似計算• 実装• 逆数近似命令

2 / 27

exp(float x);の近似計算（1/3）• 級数展開を使う•

𝑥2𝑥3𝑥
𝑒 =1+𝑥+ + +⋯26

• 要請
• 級数展開における𝑥は小さいのが望ましい• 2𝑛 （𝑛が整数）は高速に求められる• 式変形• 𝑒 𝑥 = 2𝑥 log2 𝑒• 𝑦 ≔ 𝑥 log 2 𝑒 を整数部分𝑛と小数部分𝑎に分ける（ 𝑎 ≤ 1/2）
• 𝑦 =𝑛+𝑎• 𝑒 𝑥 = 2𝑦 = 2𝑛+𝑎 = 2𝑛 2𝑎• 2𝑎 = 𝑒 𝑎 log(2)• 𝑏 ≔ 𝑎 log(2), 𝑏 ≤ log 2 /2 = 0.346

3 / 27

exp(float x);の近似計算（2/3）• floatなら1e-6程度まで求まれば十分• 5次の項まで計算• 𝑒𝑥𝑥2𝑥3𝑥4𝑥5=1+𝑥+ + + +2624120• 係数補正• 𝑓 𝑥 ≔ 1 + 𝑥 + 𝐵𝑥 + 𝐶𝑥 2 + 𝐷𝑥 3 + 𝐸𝑥 4 + 𝐹𝑥 5
• 𝑥 ∈ [−𝐿, 𝐿] where 𝐿 ≔ log(2)/2•𝐿𝐼 ≔ ‫׬‬−𝐿 𝑓 𝑥− 𝑒 𝑥 2 𝑑𝑥を最小化する(𝐵, 𝐶, 𝐷, 𝐸, 𝐹)を求める• 1次の項(𝐴)を1に固定しているのは0次と共用したいから• 誤差が概ね半分程度になる• Sollyaのremezよりは誤差が小さい?
• L2ノルムとL1ノルムどちらがよいかアプリ依存?
4 / 27

exp(float x);の近似計算（3/3）• まとめ
input : x
x = x * log_2(e)n = round(x) ; 四捨五入
a = x - nb = a * log(2)c = 1 + b(1 + b(B + b(C + b(D + E b))))y = c * 2^noutput : y• 配列に対する一括処理
void expv(float *dst, const float *src, size_t n) {for (size_t i = 0; i < n; i++) dst[i] = exp(src[i]);}• テーブルを使わないためSIMD化しやすい
5 / 27

AVX-512• Intelの512bit幅のSIMD命令セット• 32個の512bit AVXレジスタ• double x 8, float x 16, int32 x 16, int16 x 32, int8 x 64などの型• 概ね3オペランド• op(r1, r2, r3) ; r1 = r2 op r3
paddd zmm0, zmm1, zmm2 // zmm0 = zmm1 + zmm2 as 32-bit int
addps zmm0, zmm1, zmm2 // as float
addpd zmm0, zmm1, zmm2 // as double• d ; 32-bit int, ps ; float, pd ; double• 積和演算FMA• d = a * b + cが望ましいが4オペランドになるのでvfmadd132ps r1, r2, r3 // r1 = r1 * r3 + r2
vfmadd213ps r1, r2, r3 // r1 = r2 * r1 + r3
vfmadd231ps r1, r2, r3 // r1 = r2 * r3 + r1
6 / 27

Xbyak/Xbyak_aarch64• C++でCPUのニーモニックを記述• 実行時にそれに対応する機械語が生成されて実行する• 例 : 「整数nを足す関数」を生成する関数
struct Code : Xbyak::CodeGenerator {void genAdd(int n) {// rsiはLinuxでの関数の第一引数
lea(rax, ptr[rsi + n]);ret();} };• genAdd(5); // 「5を足す関数」が実行時に生成される• 生成された関数
lea rax, [rsi + 5] // + 5は即値
ret

• exp自体は静的な関数でよいがIntel oneDNNではパーツの一つとして実行時生成されている
7 / 27

AVX-512によるexp• round関数• vrndscaleps(dst, src, round_ctl)• round_ctl ; 2bitフラグ• 00 ; nearest even ; 一番近い偶数丸め• 01 ; equal or smaller ; 切り捨て
• 10 ; equal or larger ; 切り上げ
• 11 ; truncate ; 0方向に切り捨て• vrndscaleps(dst, src, 0) ; dst = round(src);• c = 1 + b(1 + b(B + b(C + b(D + E b))))• FMAを5回適用• 𝑦 = 𝑐 × 2𝑛 の部分• AVX2までは整数𝑛をビットシフトしてfloatに変換して...
• AVX-512ではvscalefps(dst, r1, r2) // dst = r1 * 2^r2
8 / 27

exp一つ分• 準備• 下記変数はzmmレジスタのエリアス名• log2_e, log2, one, B, C, D, Eは事前に定数設定• t1, t2は一時レジスタ, xが入出力（x = exp(x)）
genExpOne() {vmulps(x, log2_e);// x *= log2_e
vrndscaleps(t1, x, 0); // t1 = round(x)vsubps(x, t1);// x = x - t1 ; 小数部分avmulps(x, log2);// a * log2
vmovaps(t2, E);
vfmadd213ps(t2, x, D);
vfmadd213ps(t2, x, C);
vfmadd213ps(t2, x, B);
vfmadd213ps(t2, x, one);
vfmadd213ps(t2, x, one); // t2 = 1+b(1+b(B+b(C+b(D+E b))))vscalefps(x, t2, t1);// x = t2 * 2^t1}
9 / 27

端数処理(1/2)• ループnが16の倍数でないときの扱い• マスクレジスタk1, ..., k7• zmmレジスタのkに対応するbitが1なら処理, 0なら非処理• T_zを指定すると非処理の部分に0が入る• vmovups(zm0|k1|T_z, ptr[src]);
k1 = 0b00..01111111 (8bitの1)のとき
src | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9| a| b| ...
|← readしない→
zm0 | x0| x1| x2| x3| x4| x5| x6| x7| 0| | 0| 0|...

• readしない部分はメモリアクセスしない(アクセス違反しない)• マスクレジスタを指定しないときに比べてやや遅い• 16の倍数のときは指定しない方がよい
10 / 27

端数処理(2/2)• コア部分Label lp = L();// メインループ
vmovups(zm0, ptr[src]); // zm0 = *src
add(src, 64);// src += 64
genExpOne()// zm0 = exp(zm0)
vmovups(ptr[dst], zm0); // *dst = zm0
add(dst, 64);// dst += 64
sub(n, 16);// n -= 16jnz(lp);// if (n != 0) goto lp
L(mod16);// 端数処理
and_(ecx, 15);// n &= 15jz(exit);mov(eax, 1);shl(eax, cl);// eax = 1 << nsub(eax, 1);// eax = (1 << n) - 1
kmovd(k1, eax);// k1 = nビットの1vmovups(zm0|k1|T_z, ptr[src]); // 残り読み
genExpOne()// zm0 = exp(zm0)
vmovups(ptr[dst]|k1, zm0|k1); // 残り書き
L(exit);11 / 27

ベンチマーク• float x[16384];に対するstd::exp(float)との比較(clk)
std::exp fmath::exp
140.1
8.7
• Xeon Platinum 8280 2.7GHz, g++ 9.3.0 Ubuntu 20.04.1 LTS• genOneExpをループアンロール
genExpTwo() {vmulps(x0, log2_e);
vmulps(x3, log2_e);
vrndscaleps(t1, x0, 0);
vrndscaleps(t4, x1, 0);
vsubps(x0, t1);
vsubps(x3, t4);...

• fmath::exp 8.7clk→7.2clk• https://github.com/herumi/fmath/blob/master/fmath2.hpp
12 / 27

SVE• Armの可変長SIMD命令セット• 128～1024(?)bitレジスタ• double x 8, float x 16, int32 x 16, int16 x 32, int8 x 64などの型• 概ね3オペランド• A64FX（富岳のCPU）では32個の512bitレジスタz0, ..., z31
fadd z0.s, z1.s, z2.s // as float
add z0.b, z1.b, z2.b // as int8

• movprfx（dstをsrcとして利用）
movprfx(dst, pred, r1);
fmadd(dst, pred, r2, r3); // dst = r1 * r2 + r3

• predは述語レジスタ• マスクレジスタ相当• AVX-512と異なり常に指定
13 / 27

述語レジスタ• SVEレジスタの各要素を処理する(1)か否(0)かを指定• 例 ld1w(z.s, p/T_z, ptr(src));• z = *src; // float x 16個読み込み
src

x0

x1x2x3...

z.sx00x2x3...

p1011• i番目の各要素(i = 0, ..., 15)について• 述語レジスタp[i] = 1ならz[i] = src1[i]• p[i] = 0ならT_z(zero)を指定しているのでz[i] = 0• T_zを指定しなければp[i]の値を変更しない
14 / 27

whilelt(p.s, x4, n);• 「x4 + i < n」が成り立つ添え字までp[i] = 1にする
• 例 x4 + 16 <= nならi = 0, ..., 15についてp[i] = 1• 全てのデータが有効
i0 1 2 3 4 5 6 7 8 9 10

p 11111111111111213141511111• x4 + 3 = nならi ≦ 2についてp[i] = 1, その他p[i] = 0
i0 1 2 3 4 5 6 7 8 9 10

p 11100000000111213141500000• p[i] = 0の部分はデータを読まない・書かない• 読み書き属性が無い領域でも大丈夫
15 / 27

ループの終わり処理• メイン部分
Label lp = L();
ld1w(z0.s, p0/T_z, ptr(src1, x4, LSL, 2));// z0 = src1[x4 << 2]...
incd(x4);// x4 += 16L(cond);
whilelt(p0.s, x4, n); // while (x4 < n)ならb_first(lp);// lpラベルにジャンプ• 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の略• デメリット• あまり速くない ; 結局AVX-512のように分けた方がよい 16 / 27

sveによるgenExpOne• AVX-512とほぼ同じ
fmul(tz0, tz0, log2_e);
frintn(tz2, p, tz0); // round : float -> float
fcvtzs(tz1, p, tz2); // float -> int
fsub(tz2, tz0, tz2);
fmul(tz2, tz2, log2);
movprfx(tz0, p, E);fmad(tz0, p, tz2, D); // tz0 = tz2 * E + D
fmad(tz0, p, tz2, C);fmad(tz0, p, tz2, B);fmad(tz0, p, tz2, one);fmad(tz0, p, tz2, one);
fscale(tz0, p, tz1);// tz0 *= 2^tz1• 述語レジスタpはall 1に設定
• round関数はfrintn• fscale(x, p, n)のnがfloatではなくintなのでfcvtzsでintに変換
17 / 27

レジスタリネーミングが苦手?
• frintnなどの命令のあとの依存関係• レジスタ名を入れ換えると2倍速くなるケース(74→32nsec)• https://github.com/herumi/misc/commit/0362d5647f693be66
da841eea7ca333d0f5b5329

18 / 27

ループアンロール• 利用レジスタを増やしてループ展開// n = 1 or 2 or 3for (int i = 0;i<n;i++) fmul(t[0+i*3], t[0+i*3], log2_e);for (int i = 0;i<n;i++) frintn(t[2+i*3], p, t[0+i*3]);for (int i = 0;i<n;i++) fcvtzs(t[1+i*3], p, t[2+i*3]);for (int i = 0;i<n;i++) fsub(t[2+i*3], t[0+i*3], t[2+i*3]);for (int i = 0;i<n;i++) fmul(t[2+i*3], t[2+i*3], log2);
for (int i = 0;i<n;i++) {movprfx(t[0+i*3], p, E);fmad(t[0+i*3], p, t[2+i*3], D);}for (int i = 0;i<n;i++) fmad(t[0+i*3], p, t[2+i*3], C);
for (int i = 0;i<n;i++) fmad(t[0+i*3], p, t[2+i*3], B);
for (int i = 0;i<n;i++) fmad(t[0+i*3], p, t[2+i*3], one);for (int i = 0;i<n;i++) fmad(t[0+i*3], p, t[2+i*3], one);for (int i = 0;i<n;i++) fscale(t[0+i*3], p, t[1+i*3]);

19 / 27

ベンチマーク• float x[16384];に対するstd::exp(float)との比較(nsec)• FX700

std::exp
217.3

fmath::exp, N=119.0

N=211.8

N=38.4• A64FXはXeonに比べてレイテンシが大きい• https://github.com/fujitsu/A64FX• frintn, fmul, fadd; 9clk
• fdiv, 98clk• ループアンロールの効果が大きい

20 / 27

SVEのfexpa• pow(2, i/64) (i=0,...,63)のテーブル引き命令• floatのbit表現𝑓2• u=|s:1|e:8|f:23|, 𝑓 = −1 𝑠 2𝑒−127 (1 + 24 ), e:指数部, f:仮数部• tbl[i] := 「pow(2, i/64)」の下位23bit• 𝑥=2𝑖64

for 𝑖 = 0, … , 63は1 ≤ 𝑥 < 2なので常にx.e=127• fexp(x)の挙動
x.sx.ex.fy=8bit v=6bit
tbl[]

0y=8bittbl[v]

21 / 27

fexpaの使い方• 2ベキの近似なのでそれにもっていく• exp 𝑥 = 2𝑥 log2 𝑒 = 2𝑦 where 𝑦 ≔ 𝑥 log 2 𝑒• 𝑦 = 𝑛 + 𝑎 where 𝑛 ≔ floor(𝑦), 𝑎 ≔ 𝑛 − 𝑦, 0 ≤ 𝑎 < 1• fexpaの添え字vが2ベキになるには[1, 2)にある必要性• 𝑏 ≔ 1 + 𝑎とすると1 ≤ 𝑏 < 2, 𝑦 = 𝑛 − 1 + 𝑏
v=6bit𝑏.f
w=17bit• 𝑏.f上位6bit vを取り出して使う c = fexpa(b.f >> 17)• 𝑏 = 𝑏. f/224 =v /26 + w/224 , 𝑧 ≔ w/224• 2𝑏 =fexpa(v)⋅ 2 𝑧 , 2 𝑧 = 𝑒 log 2 𝑧 = 1 + log 2 𝑧 + log 2 𝑧 2 /2
• 差分が小さいので2次の項まででよい• https://github.com/herumi/misc/blob/master/sve/fmath-sve.hpp

22 / 27

fexpaによるexp• 主要コード概要
fmul(t0, t0, para.log2_e);
frintm(t1, p, t0);// floor : float -> float
fcvtzs(t2, p, t1);// n = float -> int
fsub(t1, t0, t1);// afadd(t1, t1, one);// b = 1 + alsr(t3, t1, 17);// v = b >> 17
fexpa(t3, t3);// c = fexpa(v)fscale(t3, p, t2);// t3 *= 2^nand_(t2.d, t1.d, not_mask17.d);fsub(t2, t1, t2);// z
movprfx(t0, p, coeff2);
fmad(t0, p, t2, coeff1);
fmad(t0, p, t2, one);// 1 + log2 z + (log2)^2/2 z^2fmul(t0, t3, t0);

23 / 27

ベンチマーク• float x[16384];に対するstd::exp(float)との比較(nsec)• ()内はfxpaを使わないバージョン
std::exp
217.3

N=1 N=2 N=3
18.2 11.3
8.4(19.0) (11.8) (8.4)• 気持ち速いかも?
• 定数レジスタの個数 7 vs 5
• 中間レジスタの個数 3 vs 4

• 雑感
• 最初の方法に比べてレジスタの依存関係が複雑になる• 定数レジスタが少ないのはよいがそこまでのメリットが• もう少し精度があれば• うまく速くなるレジスタ割り当てに試行錯誤（全数探索?）
24 / 27逆数近似計算• fdivの代わりに逆数近似命令(rcp)のあと補正する
t = rcp(x)1/x = 2 * t - x t^2• AVX-512// x : input/output// t : temporary// two : 2.0f
vrcp14ps(t, x);vfnmadd213ps(x, t, two); // x = -xt + 2vmulps(x, x, t);• SVE ; frecps(dst, src1, src2) ; // dst = 2 - src1 * src2
• 2回補正するとfloatに近い精度
25 / 27逆数近似計算• frintmあるいはaddの後のz0の逆数処理• https://github.com/herumi/misc/blob/master/sve/inv.cpp• ○で囲った部分が変?
• レジスタリネーミング?
// input : z0 (compute it with z0, z1, z2)
1. fdivr(z0.s, p0, one.s);
2. frecpe(z1.s, z0.s); frecps(z2.s, z0.s, z1.s);fmul(z0.s, z1.s, z2.s);
3. frecpe(z1.s, z0.s); frecps(z3.s, z0.s, z1.s);fmul(z0.s, z1.s, z3.s);
4. frecpe(z1.s, z0.s); frecps(z2.s, z0.s, z1.s);fmul(z1.s, z1.s, z2.s);
frecps(z2.s, z0.s, z1.s); fmul(z0.s, z1.s, z2.s);
5. frecpe(z1.s, z0.s); frecps(z3.s, z0.s, z1.s);fmul(z1.s, z1.s, z3.s);
frecps(z3.s, z0.s, z1.s); fmul(z0.s, z1.s, z3.s);clk
frintm add
100
100
33
7.9
10.9

7.9

51

11.0

11.2

11.0

26 / 27

まとめ• AVX-512とSVE(A64FX)• どちらも512-bitレジスタx32• SVEの方がレイテンシが大きい• ループアンロール重要• レジスタリネーミングに気をつける場合がある?

27 / 27