SVEの自動ベクタライズ

ARM SVE というベクタ拡張があるんだけど、それって公開されてるコンパイラでどのぐらい対応されるんや？という話題を見かけたので、ちょっと試してみたら良い感じだったのでついでにもう少し調べてみた。

http://d.hatena.ne.jp/w_o/20150423#1429775436

個人的には自動ベクタライザというのは信用してないんだけど、それはそれとして、ループ変換ツールとしての最適化パスの処理はまあ面白いと思ってるので。

GCCは、ビルドしてもいいけど、

https://developer.arm.com/tools-and-software/open-source-software/developer-tools/gnu-toolchain/gnu-a/downloads

ここからビルド済みバイナリが取得できるので、そちらのほうが簡単だと思う。

このGCCに、-march=armv8-a+sve -O3 を渡すと、SVEを使って自動ベクタライズしてくれるようになる。

自動ベクタライザを信用しない理由として、上のリンク先で上げてる通り、どうしても命令数がかなり膨らんでしまうというのがありますね。これは性能上のデメリットもあるんだけど、出るasmの見た目が汚なくなるから嫌という気持ちの問題も大きい。

ところが、SVE向けに自動ベクトル化で出したコードは、この命令数の膨張が無くて、見ためがスッキリしてたので好きになった。

int f(int *__restrict__ a,
      int *__restrict__ b,
      int *__restrict__ c,
      int n)
{
    for (int i=0; i<n; i++) {
        a[i] += b[i] * c[i];
    }
}

これが、

f:
.LFB0:
	.cfi_startproc
	cmp	w3, 0
	ble	.L2
	mov	x4, 0
	sxtw	x3, w3
	ptrue	p1.s, all
	whilelo	p0.s, xzr, x3
	.p2align 3
.L3:
	ld1w	z1.s, p0/z, [x0, x4, lsl 2]
	ld1w	z2.s, p0/z, [x1, x4, lsl 2]
	ld1w	z0.s, p0/z, [x2, x4, lsl 2]
	mad	z0.s, p1/m, z2.s, z1.s
	st1w	z0.s, p0, [x0, x4, lsl 2]
	incw	x4
	whilelo	p0.s, x4, x3
	bne	.L3
.L2:
	ret

vectorizeしない場合は、

f:
.LFB0:
	.cfi_startproc
	cmp	w3, 0
	ble	.L2
	mov	x4, 0
	.p2align 3
.L3:
	ldr	w6, [x0, x4, lsl 2]
	ldr	w5, [x1, x4, lsl 2]
	ldr	w7, [x2, x4, lsl 2]
	madd	w5, w5, w7, w6
	str	w5, [x0, x4, lsl 2]
	add	x4, x4, 1
	cmp	w3, w4
	bgt	.L3
.L2:
	ret

ループ前の処理が3命令増えているが、ループ中の命令数は全く同じで、関数レベルで見ても12->15 の 3命令しか増えていない。neonの場合は、これが47命令の+35命令になって、ここに貼るのも嫌なレベルだ。

一番大きいのは、ループ回数がSIMD幅で割り切れない時の端数の処理に必要なマスクをwhilelo 命令一発で作れるところだろう。

http://d.hatena.ne.jp/w_o/20160901#1472719309 に昔説明を書いたけど、実際これだけ自動ベクタライザがスッキリした命令出すのを見てかなり感動しましたね。

自動ベクタライザが出すasmは汚いという時代はもう終わったんだ。

あと

int f(int *__restrict__ a,
      int *__restrict__ b,
      int *__restrict__ c,
      int n)
{
    for (int i=0; i<n; i++) {
        if (a[i]) {
            a[i] += b[i] * c[i];
        }
    }
}

こういう条件付きストアがある場合でも

f:
.LFB0:
	.cfi_startproc
	cmp	w3, 0
	ble	.L2
	mov	x4, 0
	sxtw	x3, w3
	ptrue	p2.s, all
	whilelo	p1.s, xzr, x3
	.p2align 3
.L3:
	ld1w	z1.s, p1/z, [x0, x4, lsl 2]
	cmpne	p0.s, p2/z, z1.s, #0
	and	p0.b, p0/z, p0.b, p1.b
	ld1w	z0.s, p0/z, [x1, x4, lsl 2]
	ld1w	z2.s, p0/z, [x2, x4, lsl 2]
	mad	z0.s, p2/m, z2.s, z1.s
	st1w	z0.s, p0, [x0, x4, lsl 2]
	incw	x4
	whilelo	p1.s, x4, x3
	bne	.L3
.L2:
	ret