rust で SIMD -- x86intrinsic を移植した話

これは rust advent calendar 2016 (2) の 1日目の記事です。

SIMD とは

SIMD (シムディー) とは、Single Instruction, Multiple Data の略で、１つの命令で(同じ種類の)演算を複数のデータに対して同時に行なうことを言います。

例えば、SSE2 には PADDW という命令があります。これはワード単位、すなわち、2byte 単位で加算を行なうという命令です。SSE では、16byte の幅を持つ xmm レジスタを用いて命令を発行することができます。この 16byte 幅のレジスタは、16x1byte, 8x2byte, 4x4byte, 2x8byte のように複数のデータをまとめたものとみなすことができます。PADDW 命令を xmm レジスタに対して発行すれば、2byte 単位の加算を一命令で 8 つ同時に行なうことができます。普通の CPU 命令は 1 つの命令に対して 1 つのデータしか演算を行わないので、単純には 8 倍速で加算が行えることになります。したがって、同じ演算を大量のデータに対して行なうときには、SIMD 命令を用いると高速に計算を行なうことができます。

rust で simd 演算

rust で SIMD 演算を行なうには、simd crate を使うのが今の所メジャーです。

一方、2016年11月現在、simd crate には AVX2 命令があまり実装されておらず、また SSE 命令にも実装されていないものがそこそこあります。筆者はちょうどこれらの命令が必要であったこと、および C/C++ において SIMD 命令を利用できる x86intrinsic 命令たちに慣れていたことから、rust で x86intrinsic と同様の関数群を持つ x86intrin crate を実装しました。現状 AVX2 まで、ごく一部の命令とMMX命令を除いて実装しています。一応私が必要な機能は全部揃ったので現状は満足しています。

x86intrin crate の使い方

x86intrin crate は、基本的に C/C++ での x86intrinsic から関数名、型名から先頭の _ を取り除いただけなので、C/C++ に慣れている人はすぐに使うことができると思います。ただし、nightly でしか動かないので、nightly を使ってください。

以下、C/C++ で x86intrinsic を使ったことがない人向けに説明します。

m128, m128d, m128i が xmm レジスタと同サイズ (16byte) のデータを表す型になります。それぞれ単精度浮動小数点数、倍精度浮動小数点数、整数を表すときに用います。単精度浮動小数点数は 4byte 、倍精度浮動小数点数は 8byte なので、xmm レジスタにはそれぞれ 4 つ、2 つのデータを同時に載せられます。整数は、1byte, 2byte, 4byte, 8byte のいずれかで、それぞれ 16, 8, 4, 2 個のデータを同時に載せることができます。

簡単な足し算

実際に簡単な足し算を計算してみます。今回、型はわざと明示しています。

extern crate x86intrin;

use x86intrin::*;

fn main() {
    // mm_setr_epi32(r0, r1, r2, r3) は、i32 を 4 つとって、m128i を作ります。
    // m128i の下位から r0, r1, r2, r3 が並びます。
    let a: m128i = mm_setr_epi32(1, 2, 3, 4);
    let b: m128i = mm_setr_epi32(4, 4, 4, 4);

    // 4x4byte を同時に足し算します。
    let c: m128i = mm_add_epi32(a, b);

    // ここは x86intrin の拡張で、as_i32x4().as_array() で、[i32; 4] な
    // 値を生成できます。println! とかには便利です。
    let vs: [i32; 4] = c.as_i32x4().as_array();
    println!("{} {} {} {}", vs[0], vs[1], vs[2], vs[3]);
}

これを実行すると、1 + 4, 2 + 4, 3 + 4, 4 + 4 が実行された結果である 5 6 7 8 と表示されます。

アセンブラを見る

実際に SIMD 命令が使われているかアセンブラで確認してみましょう。

リリースビルドだと完全に定数畳み込みが行なわれて SIMD 命令が消えるので、次のように#[inline(never)] をつけてビルドしました。

extern crate x86intrin;

use x86intrin::*;

#[inline(never)]
fn kotori(a: m128i, b: m128i) -> m128i {
    mm_add_epi32(a, b)
}

fn main() {
    let a: m128i = mm_setr_epi32(1, 2, 3, 4);
    let b: m128i = mm_setr_epi32(4, 4, 4, 4);
    let c: m128i = kotori(a, b);

    let vs: [i32; 4] = c.as_i32x4().as_array();
    println!("{} {} {} {}", vs[0], vs[1], vs[2], vs[3]);
}

適当に関数を作って、あとは適当に objdump -d とかした後でその関数を見てあげると、mm_add_epi32 に対応する paddd 命令が使われていることがわかります。

0000000100001260 <__ZN6intrin6kotori17hdacbdd83fa921175E>:
   100001260:   55                      push   %rbp
   100001261:   48 89 e5                mov    %rsp,%rbp
   100001264:   66 0f fe c1             paddd  %xmm1,%xmm0
   100001268:   5d                      pop    %rbp
   100001269:   c3                      retq
   10000126a:   66 0f 1f 44 00 00       nopw   0x0(%rax,%rax,1)

なお、cargo でビルドするとき、デフォルトでは x64 の上では SSE2 命令までしか使えないので、次のように RUSTFLAGS を付加してください。target-cpu=native で自分の CPU が持っている命令が使えるようになります。

$ RUSTFLAGS="-C target-cpu=native" cargo build --release

筆者は現在 Haswell という CPU を使っているので、AVX2 命令まで使えます。AVX2 まで使える場合、mm_add_epi32 は、SSE2 の paddd ではなく AVX2 命令の vpaddd が代わりに使用されます。

0000000100001260 <__ZN6intrin6kotori17hdacbdd83fa921175E>:
   100001260:   55                      push   %rbp
   100001261:   48 89 e5                mov    %rsp,%rbp
   100001264:   c5 f1 fe c0             vpaddd %xmm0,%xmm1,%xmm0
   100001268:   5d                      pop    %rbp
   100001269:   c3                      retq
   10000126a:   66 0f 1f 44 00 00       nopw   0x0(%rax,%rax,1)

マンデルブローの実装

実際に x86intrin crate を用いてマンデルブロー集合を計算し、実際に高速になるか試してみます。

SIMD 演算を使わない naive 版、xmm レジスタを使った simd4 版、ymm レジスタ (32byte 幅) を使った simd8 版を用意しました。

長いのでソースは[こちらの gist] (https://gist.github.com/mayah/fac2f92bb963c09166c19b2ad8b3dcbf) に。

手持ちの MacBook Pro Late 2013 (2.4 GHz Intel Core i5) での RUSTFLAGS="-C target-cpu=native" cargo bench の実行結果 (抜粋) がこちらです。なおこのマシンの CPU は Haswell です。

test mandel_naive  ... bench:   3,552,445 ns/iter (+/- 226,933)
test mandel_simd4  ... bench:   1,005,547 ns/iter (+/- 175,346)
test mandel_simd8  ... bench:     601,596 ns/iter (+/- 137,141)

simd4 版は 4 倍とまではいかないものの、3.5 倍ほど高速になっています。simd8 版も 6 倍ほど高速になっています。

4 倍や 8 倍にならなかった理由は、マンデルブロー集合は場所によって必要な計算回数が異なるので、同時に 4 点や 8 点計算する SIMD 版は、最も遅い点の部分に合わせて計算する必要があるため、というのが主な原因と考えていますが、ほかにも、通常命令と SSE, AVX 命令とでは IPC (Instructions Per Clock) が異なるなどの原因もあるでしょう。

x86intrin 実装の余談

以下は x86intrin crate を使いたいだけの人には余談です。

rust は LLVM の上に実装されています。LLVM のビットコードから x86 の命令へ変換するとき、LLVM 上の特定のビットコードの組み合わせを見ると対応する x86 命令を吐き出すという仕様になっています。そこで、一部の SIMD 命令を狙って吐き出すには、LLVM の気持ちになってあげる必要があります。

例えば、PANDN x, y という命令は、~x & y を計算するという命令ですが、これにそのまま対応するLLVMビットコード上の命令はありません。しかし、LLVM は、ones を全てのビットが 1 であるような値として (ones ^ x) & y に相当するLLVMビットコードを見つけると、PANDN 命令を出力するという実装になっています。これは LLVM の実装を知らないとわかりません。

また、SSE ファミリーの命令には、レジスタ内で値をコピーしたり順番を入れ替えたりする命令がありますが、ほとんどが rust が用意している simd_shuffle という platform-intrinsic で実装ができます。これも PANDN と同じで、LLVM が知っている simd_shuffle の引数を与えると対応する命令を吐き出すという形になっているため、clang に付属する emmintrin.h などのヘッダーを読んで simd_shuffle に対応する値を入れていきます。

そのほか、LLVM が用意している命令はバージョンアップとともにしばしば消え去り、より汎用的な命令で置き換えられている場合があります。例えば PALIGNR という命令を吐き出させるためには、一時は LLVM が用意している intrinsic (__builtin_ia32_palignr128 相当のもの) を呼び出せばよかったようで、手元の clang の header ではそうなっているようなのですが、気がついたら C/C++ 版は shuffle_vector (rust では simd_shuffle) を使う実装になっていました。このように、LLVM 側の変更を追わないと必要な実装ができないということが何度かありました。

安定化に向けて

先々週あたりから SIMD を安定版 rust で提供するためにどういう方向で実装していくかについて、長いスレッドで議論が交わされています。

https://internals.rust-lang.org/t/getting-explicit-simd-on-stable-rust/4380

まあ大方の予想通り、simd crate のような type safe な SIMD を提供するような意見が主流を占めていると（ざっと読んだ限りでは）思います。早く安定化してほしいですね。

2016-12-01