c-bata web

日本語の技術ブログ by @c-bata . 英語の記事はMediumに書いています。

Intel AVX (Advanced Vector Extensions) によるSIMD演算

SIMD(Single Instruction Multiple Data)は、名前の通り複数のデータを1命令で処理する方式を指します。各プロセッサコアに異なる命令を供給する機構が必要なMIMDと比べ、SIMDは必要なトランジスタ数が少なく小さな面積でプロセッサを設計することができるため、GPUにおける並列処理のパラダイムとしても主流になっています (SIMTは今は扱わない)。

Intel CPUのMMX命令やSSE命令、AVX命令はSIMD方式の演算命令になります。動画配信チームからML系の研究チームに異動したのですが、研究成果を実用レベルのアプリケーション実装に落とし込めるよう今回はこのAVX命令について整理しました。この記事の実行環境はOSがmacOS, cpuは intel x86_64 です。AVXの解説書としては日本語だと下記の書籍がかなり整理されていますが、基本的に x86 前提のコードなので参考にする際は適宜置き換えてください。

AVX命令入門―Intel CPUのSIMD命令を使い倒せ

AVX命令入門―Intel CPUのSIMD命令を使い倒せ

SSEを用いた行列の加算演算

x64ではインラインアセンブリが利用できないため、アセンブリ命令をマクロで記述できるようにしたイントリンシックを利用します。マクロを覚えるのが面倒でx64アセンブリに慣れているなら直接関数シンボルを記述してgccでオブジェクトファイルを生成し、ldでリンクする方法もとれなくはないはずです。ただしポータビリティの観点からも基本的にはイントリンシックを素直に覚えるのがいいでしょう。

SIMD命令の1つであるSSEをまずは使ってみます。ヘッダファイル xmmintrin.h にSSE用のアセンブリに展開するためのマクロが入っているので、そちらincludeして次のように行列加算を実行します。

#include <stdio.h>
#include <xmmintrin.h>

int main(void) {
    __m128 a = {1.0f, 2.0f, 3.0f, 4.0f};
    __m128 b = {1.1f, 2.1f, 3.1f, 4.1f};
    float c[4];
    __m128 ps = _mm_add_ps(a, b); // add
    _mm_storeu_ps(c, ps); // store

    printf("  source: (%5.1f, %5.1f, %5.1f, %5.1f)\n",
            a[0], a[1], a[2], a[3]);
    printf("  dest. : (%5.1f, %5.1f, %5.1f, %5.1f)\n",
           b[0], b[1], b[2], b[3]);
    printf("  result: (%5.1f, %5.1f, %5.1f, %5.1f)\n",
           c[0], c[1], c[2], c[3]);
    return 0;
}

SSEの命令では128ビットのレジスタが利用できます。浮動小数点数 float型は4 bytes(=32 bits)であることから4つの要素を一度に演算できます。これを実行してみます。

$ gcc -o avx -Wall -O0 main.c
$ ./avx-
  source: (  1.0,   2.0,   3.0,   4.0)
  dest. : (  1.1,   2.1,   3.1,   4.1)
  result: (  2.1,   4.1,   6.1,   8.1)

AVX (AVX2)を用いた行列の加算演算

SSEではSIMD用のレジスタが128ビットであったため、floatであれば4要素までしか一度に演算できませんでした。AVX命令はビット幅が256ビットへ拡張され、3オペランド命令書式が採用され演算性能を大幅に向上することができました。またメモリオペランドを用いた演算時にアドレスのアライメントを要求しなくなったためメモリ配置の自由度があがりました(その場合性能が少し低下します)。さらにその後追加されたAVX2では浮動小数点だけでなく、整数の演算がサポートされました。

基本的にAVX2がサポートされている環境ではこちらを利用することが望ましいでしょう。AVXやAVX2のアセンブリに展開するためのマクロを使うときは immintrin.h をincludeします。

#include <stdio.h>
#include <immintrin.h>

int main(void) {
    __m256 a = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
    __m256 b = {1.1f, 2.1f, 3.1f, 4.1f, 5.1f, 6.1f, 7.1f, 8.1f};
    __m256 c;

    c = _mm256_add_ps(a, b);

    printf("  source: (%5.1f, %5.1f, %5.1f, %5.1f, %5.1f, %5.1f, %5.1f, %5.1f)\n",
            a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7]);
    printf("  dest. : (%5.1f, %5.1f, %5.1f, %5.1f, %5.1f, %5.1f, %5.1f, %5.1f)\n",
            b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]);
    printf("  result: (%5.1f, %5.1f, %5.1f, %5.1f, %5.1f, %5.1f, %5.1f, %5.1f)\n",
           c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]);
    return 0;
}

AVXではSIMD用のレジスタが256ビットになっているため、floatの演算であれば 256 bits / 32 bits(floatは4bytes) で8要素を一度に演算できます。浮動小数点 float ではなく倍精度浮動小数点 double を使いたいなら __m128d 、整数を使いたいなら __m128i を使用します。倍精度浮動小数点は8bytesを使うため4要素しか扱えないことに注意してください。

参照: https://www.codingame.com/playgrounds/283/sse-avx-vectorization/what-is-sse-and-avx

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

$ gcc -o avx2 -Wall main.c-
main.c:9:9: error: always_inline function '_mm256_add_ps' requires target feature 'xsave', but would be inlined into function 'main' that is compiled without support for 'xsave'
    c = _mm256_add_ps(a, b);
        ^
1 error generated.

コンパイルに失敗しました。 こちらの記事 によるとCPUがサポートしている機能は次のコマンドにより確認できます。

$ sysctl -a | grep machdep.cpu.features
machdep.cpu.features: FPU VME DE PSE TSC MSR PAE MCE CX8 APIC SEP MTRR PGE MCA CMOV PAT PSE36 CLFSH DS ACPI MMX FXSR SSE SSE2 SS HTT TM PBE SSE3 PCLMULQDQ DTES64 MON DSCPL VMX EST TM2 SSSE3 FMA CX16 TPR PDCM SSE4.1 SSE4.2 x2APIC MOVBE POPCNT AES PCID XSAVE OSXSAVE SEGLIM64 TSCTMR AVX1.0 RDRAND F16C

たしかにSSEやSSE2、SSE4.1、SSE4.2、AVX1.0 は存在しますが、AVX2がここにはありません。ただしCPUとしてはAVX2をサポートされている場合があり、次のコマンドで出てきた命令は特別なコンパイラーオプションを与えることで使用できるようです。

$ sysctl -a | grep machdep.cpu.leaf7_features
machdep.cpu.leaf7_features: SMEP ERMS RDWRFSGS TSC_THREAD_OFFSET BMI1 AVX2 BMI2 INVPCID SMAP RDSEED ADX IPT SGX FPU_CSDS MPX CLFSOPT

AVX2 が何らかのコンパイラーオプションを与えることで使用できることがわかりました。こちらの記事 によるとgccでは -mavx2 オプションが利用できます。

$ gcc -o avx2 -mavx2 -Wall main.c 
$ ./avx2
  source: (  1.0,   2.0,   3.0,   4.0,   5.0,   6.0,   7.0,   8.0)
  dest. : (  1.1,   2.1,   3.1,   4.1,   5.1,   6.1,   7.1,   8.1)
  result: (  2.1,   4.1,   6.1,   8.1,  10.1,  12.1,  14.1,  16.1)

AVX2の命令と利用できるレジスタ解説

先程のAVX2を利用したプログラムのアセンブリを出力しておく。これをもとにAVX2のレジスタ構成と命令セットを解説

$ gcc -S -mavx2 -Wall -O0 avx2.c
$ less avx2.s

まず注目するべきなのは次のアセンブリ群。

LCPI0_0:
        .long   1065353216              ## float 1
        .long   1073741824              ## float 2
        .long   1077936128              ## float 3
        .long   1082130432              ## float 4
        .long   1084227584              ## float 5
        .long   1086324736              ## float 6
        .long   1088421888              ## float 7
        .long   1090519040              ## float 8

... 略

## %bb.0:
        ... 略
        vmovaps LCPI0_0(%rip), %ymm0    ## ymm0 = [1.000000e+00,2.000000e+00,3.000000e+00,4.000000e+00,5.000000e+00,6.000000e+00,7.000000e+00,8.000000e+00]
        vmovaps %ymm0, 160(%rsp)

オブジェクトファイル内の LCPI0_0 セクションに書き込まれたfloat配列の情報を ymm0 と呼ばれるCPU上のレジスタ領域にコピーしていることがCのコードから想像できる。レジスタ構成は次のようになる。

参照: https://de.wikipedia.org/wiki/Advanced_Vector_Extensions

vmovaps命令によるレジスタ配置が続いたあと、次の命令が存在する。

        vaddps  %ymm1, %ymm0, %ymm0

オペランドが3つあることから AVX の命令であること、 add という文字から加算演算を意味することが予測できる。Intelのページ によるとsingle-precisionな浮動小数点(つまりfloat)の加算演算を行いその結果を第3オペランドで指定したレジスタに格納する。つまりymm1とymm0とよばれる256bitsのレジスタ領域に格納されている8つのfloat値を加算し、その結果がymm0に格納されることを意味している。

加算演算はこれで完了なので、あとは適切にその結果を変数 c に格納すればいい。少しアセンブリが冗長なように見えなくもないがなにかそこはもう少し読むと理由があるかも。パフォーマンスが思ったほど上がらなければ確認したほうがいいかも。

発展: AVX-512

colfaxresearch.com

たまたま見つけたこちらのページによるとどうやらAVX2はSkylakeにおいてもはやLegacyな命令セットで、AVX-512は512 bitのレジスタ幅を持っているようです。AVX2に比べてかなりの性能向上が期待できます。ただ $ sysctl -a | grep machdep.cpu.leaf7_features の結果にAVX-512が現れていないので対応状況から調べる必要がありそうです。 simdjson などもまだAVX2を使っているようなのでなにか対応状況などの背景から利用が難しいのかも。

また図を見る限りAVX-512の中にもかなりいろいろな派生があるようで、まとめると長くなりそうなのでそちらの調査結果はまた別記事にて整理しようと思います。