LLVMのauto-vectorizationとc2goasmによるGo Plan9 Assemblyの生成によるSIMD最適化

InfluxDBの開発チームはApache Arrowの技術に注目していて、ArrowのGo実装の開発にも積極的に参加しています。Stuart Carine (InfluxDBの開発チームメンバー)がApache ArrowのGo実装に取り入れたc2goasmとLLVMを使った最適化が以前話題になりました。

www.influxdata.com

c2goasm はClangで生成したアセンブリをGo Plan9 Assemblyに変換できる汎用的なコマンドラインツールです。しかしその活用事例はほとんどが、Goコンパイラーが行っていないSSEやAVX命令を用いたSIMD最適化です。この記事ではAVX2によるSIMD演算を実現するための c2goasm の使い方を整理します。最終的にどれくらい速くなるかでいうとfloat64のsum演算は10x以上高速化しました。

$ go test -bench .
goos: darwin
goarch: amd64
pkg: github.com/c-bata/sample-c2goasm
BenchmarkSumFloat64_256-4              5000000           282 ns/op
BenchmarkSumFloat64_1024-4             1000000          1234 ns/op
BenchmarkSumFloat64_8192-4              200000         10021 ns/op
BenchmarkSumFloat64_AVX2_256-4        50000000            23.5 ns/op
BenchmarkSumFloat64_AVX2_1024-4       20000000            95.9 ns/op
BenchmarkSumFloat64_AVX2_8192-4        2000000           904 ns/op
PASS
ok      github.com/c-bata/sample-c2goasm    10.911s

c2goasm

• Github: https://github.com/minio/c2goasm
• 作者minioによる解説: https://blog.minio.io/c2goasm-c-to-go-assembly-bb723d2f777f

c2goasm の使い方は次のとおりです。

1. C/C++で記述した関数からClangでアセンブリを生成
2. Goのバインディング用のシグネチャを定義
3. c2goasmのコマンドラインツールでGo Plan9 Assemblyを生成

ツールとしての作り込みが少し雑な印象はありハマりどころもありますが、この3つのステップを踏めばSIMDを使って最適化したC/C++のコードを少ないオーバーヘッドでGoから呼び出せます。

CgoによるCの関数呼び出し

Goにおいてデファクトスタンダードとなっている cgo は性能面において優れたソリューションではありません ¹ 。c2goasmで生成されたGo Plan9 Assemblyのサブルーチン呼び出しは、他のGoの関数呼び出しと同程度の効率で実行できるようです。

GithubリポジトリのREADMEにはcgoとの性能比較が紹介されています。 https://github.com/minio/c2goasm#benchmark-against-cgo

Cのコードの記述

float64のarrayに入った値の合計値を計算する関数を用意します。 SIMDによる演算で合計値の計算をするのは少しアルゴリズム的な工夫が必要です。

#include <immintrin.h>

void sum_float64_avx_intrinsics(double buf[], size_t len, double *res) {
    __m256d acc = _mm256_set1_pd(0);
    for (int i = 0; i < len; i += 4) {
        __m256d v = _mm256_load_pd(&buf[i]);
        acc = _mm256_add_pd(acc, v);
    }

    acc = _mm256_hadd_pd(acc, acc); // a[0] = a[0] + a[1], a[2] = a[2] + a[3]
    *res = _mm256_cvtsd_f64(acc) + _mm_cvtsd_f64(_mm256_extractf128_pd(acc, 1));
}

このようにintrinsicsを使って自分でかくこともできますが、実装が少し大変です。 LLVMのAuto-Vectorization にこの辺の最適化は任せられるなら楽ができます。

void sum_float64(double buf[], int len, double *res) {
    double acc = 0.0;
    for(int i = 0; i < len; i++) {
        acc += buf[i];
    }
    *res = acc;
}

これをClangでコンパイルしてみます。

$ clang -S -mavx2 -masm=intel -mno-red-zone -mstackrealign -mllvm -inline-threshold=1000 -fno-asynchronous-unwind-tables -fno-exceptions -fno-rtti -c sum_avx_intrinsics.c

オプションはこの辺を参考にしてください。

• https://github.com/minio/c2goasm#compatible-compilers
• https://clang.llvm.org/docs/genindex.html

実行すると ~.s とアセンブリファイルが生成されるので中を見るとこのような感じです。

 .section    __TEXT,__text,regular,pure_instructions
    .build_version macos, 10, 14
    .intel_syntax noprefix
    .globl  _sum_float64            ## -- Begin function sum_float64
    .p2align    4, 0x90
_sum_float64:                           ## @sum_float64
## %bb.0:
    ...
    vxorps  xmm0, xmm0, xmm0
    mov qword ptr [rsp + 40], rdi
    mov dword ptr [rsp + 36], esi
    mov qword ptr [rsp + 24], rdx
    vmovsd  qword ptr [rsp + 16], xmm0
    ...

vmovsdといった命令やレジスタ xmm0 を利用していることからLLVMがSIMD命令を活用していることがわかります。サブルーチン名 _sum_float64 はGoの関数定義で必要なので覚えておきます。

※ xmm0レジスタは128 bitsしかないレジスタのため、64 bitsを消費するdoubleの演算は2要素しか一度に演算できません。おそらくSSEの命令を利用しています。ここで ymm0 など128 bitsのレジスタを使い AVX2の命令を呼び出してほしいのですが、Clangのバージョンや最適化オプションによりうまくいったりイカなかったりするようです。詳細はGithubをご覧ください。

Goの関数定義

package c2goasm_sample

import "unsafe"

//go:noescape
func __sum_float64(buf, len, res unsafe.Pointer)


func SumFloat64Avx2(a []float64) float64 {
    var (
        p1 = unsafe.Pointer(&a[0])
        p2 = unsafe.Pointer(uintptr(len(a)))
        res float64
    )
    __sum_float64(p1, p2, unsafe.Pointer(&res))
    return res
}

シグネチャはアセンブリ内のサブルーチン名に _ を追加したものです。今回のようにもしサブルーチン名が _ から始まっていたら、 __ で始まることに注意してください。

https://github.com/minio/c2goasm/blob/0325a40cfd1fc6a5097e69eaf0292990eb6cee6a/arguments.go#L85

c2goasm実行

$ go get -u github.com/minio/asm2plan9s
$ go get -u github.com/minio/c2goasm
$ go get -u github.com/klauspost/asmfmt/cmd/asmfmt
$ c2goasm -a -f _lib/sum_avx_intrinsics.s sum_avx.s

これでGoアセンブリが生成されます。第2引数が出力ファイルですが、Goのファイル名を sum_avx.go としたときは拡張子を .s に変えただけの sum_avx.s を指定します。 https://github.com/minio/c2goasm/blob/0325a40cfd1fc6a5097e69eaf0292990eb6cee6a/c2goasm.go#L252

ここまでが一連の流れです。

ベンチマーク

Pure Goで実装した次の関数とベンチマークにより比較します。

package c2goasm_sample

func SumFloat64(a []float64) float64 {
    var sum float64
    for i := range a {
        sum += a[i]
    }
    return sum
}

ベンチマークのコードは次の通り。

package c2goasm_sample_test

import (
    "math/rand"
    "testing"
)

func init() {
    rand.Seed(0)
}

func initializeFloat64Array(n int) []float64 {
    var max float64 = 1024
    var min float64 = 0
    x := make([]float64, n)
    for i := 0; i < n; i++ {
        x[i] = rand.Float64() * (max - min) + min
    }
    return x
}

func benchmarkFloat64Sum(b *testing.B, n int) {
    x := initializeFloat64Array(n)
    b.ResetTimer()
    for i := 0; i < b.N; i++ {
        SumFloat64(x)
    }
}

func benchmarkFloat64SumAvx2(b *testing.B, n int) {
    x := initializeFloat64Array(n)
    b.ResetTimer()
    for i := 0; i < b.N; i++ {
        SumFloat64Avx2(x)
    }
}

func BenchmarkSumFloat64_256(b *testing.B) {
    benchmarkFloat64Sum(b, 256)
}

func BenchmarkSumFloat64_1024(b *testing.B) {
    benchmarkFloat64Sum(b, 1024)
}

func BenchmarkSumFloat64_8192(b *testing.B) {
    benchmarkFloat64Sum(b, 8192)
}

func BenchmarkSumFloat64_AVX2_256(b *testing.B) {
    benchmarkFloat64SumAvx2(b, 256)
}

func BenchmarkSumFloat64_AVX2_1024(b *testing.B) {
    benchmarkFloat64SumAvx2(b, 1024)
}

func BenchmarkSumFloat64_AVX2_8192(b *testing.B) {
    benchmarkFloat64SumAvx2(b, 8192)
}

実行結果

$ go test -bench .
goos: darwin
goarch: amd64
pkg: github.com/c-bata/sandbox-go/c2goasm
BenchmarkSumFloat64_256-4                5000000               277 ns/op
BenchmarkSumFloat64_1024-4               1000000              1205 ns/op
BenchmarkSumFloat64_8192-4                100000             10401 ns/op
BenchmarkSumFloat64_AVX2_256-4           2000000               768 ns/op
BenchmarkSumFloat64_AVX2_1024-4           500000              2872 ns/op
BenchmarkSumFloat64_AVX2_8192-4           100000             23946 ns/op
PASS
ok      github.com/c-bata/sandbox-go/c2goasm    10.474s

遅くなっちゃった...

追加調査

Apache Arrowでの利用事例がある以上、なにかコンパイラーオプションとかが理由で性能が十分に引き上げられていない。気になってる点は次のあたり

• xmm0とか使ってるし、vxorpsとかvmovsdとかもAVX2じゃなくてSSEとかの命令を使っていそう
  • LLVM Auto-Vectorizationのドキュメント読みつつ、Clangのコンパイルオプション見直したほうが良い。
  • https://github.com/apache/arrow/blob/master/go/arrow/math/float64_avx2_amd64.s とかはちゃんとAVX2使ってることが確認できる
• LLVMの最適化に頼らずintrinsicsを使って実装した処理で性能がちゃんとあがるのかも試してみたけど、実行時にセグフォで死んだ
  • https://github.com/c-bata/sample-c2goasm/blob/master/_lib/add_float32_avx2_intrinsics.s を見る限りアセンブリ生成まではおそらく問題なくて、Go側の呼び出し時のポインタの扱いに問題がある。
• Apache Arrowのパフォーマンスがちゃんとあがってるのかベンチマークとってみる。

追記: うまくいきました。

ひとまずClang 7.0.1+O2オプション+pragmaヒントでうまくいった。Clangのバージョンとか最適化オプションによって結構動かなくなってしまうみたい。細かい調査はまた今度。
試しにClangでアセンブリだしてみてだめなら自分でAVX2のイントリンシック使って書くとかになりそうhttps://t.co/zp22jGNEAo

— Masashi Shibata (@c_bata_) 2019年3月17日

Clang-7.0.1を利用する際には、次のコンパイルオプションを使用してください。

$ /usr/local/Cellar/llvm/7.0.1/bin/clang -S -O2 -mavx2 -masm=intel -mno-red-zone -mstackrealign -mllvm -inline-threshold=1000 -fno-asynchronous-unwind-tables -fno-exceptions -fno-rtti -c sum_float64.c

追記: SSEのコードはなぜ遅くなっていたのか。

SSEでも遅くはならないんじゃないかなというのが気になっていたけれど、なんとなく原因がイメージついた気がします。 InfluxDataの記事にあるとおりsumの計算を最適化してCで直接書くとAVX2のコードはこんな感じになります。

void sum_float64_avx_intrinsics(double buf[], size_t len, double *res) {
    __m256d acc = _mm256_set1_pd(0);
    for (int i = 0; i < len; i += 4) {
        __m256d v = _mm256_load_pd(&buf[i]);
        acc = _mm256_add_pd(acc, v);
    }

    acc = _mm256_hadd_pd(acc, acc); // a[0] = a[0] + a[1], a[2] = a[2] + a[3]
    *res = _mm256_cvtsd_f64(acc) + _mm_cvtsd_f64(_mm256_extractf128_pd(acc, 1));
}

重要なのはこの _mm256_hadd_pd の処理で こちらの説明にあるように詰め込まれている要素の加算 を担当します。128 bitsしかないSSE用のレジスタ xmm0 とかだと8 bytes 必要な double の値を詰め込んでも2つしかのらないので、それにたいして hadd_pd みたいな処理(これはAVX2の処理ですがそれに相当する処理がSSEにもあるのかと思います) をすると全く並列化されてないことになります。 細かいところはアセンブリが使ってる命令をみればはっきりしそうです。

プログラミング言語Go (ADDISON-WESLEY PROFESSIONAL COMPUTING SERIES)

---

1. GoのポインタをCに渡せないようにポインタのバリデーション処理が走るなど、GC側の都合による性能劣化が発生するようです。詳細は Why cgo is slow @ CapitalGo 2018 - Speaker Deck をご覧ください。↩