この記事は、「ElixirからRustの関数をつかう → はやい」の続編だ。今回は、前回の記事で最速だった Rust によるマルチスレッドプログラムを、少ない労力で SIMD 化して、さらなる高速化を図る。
実はコード自体は少し前に書いてあったのだが、なかなか本文を書き進める時間がとれず、今回は、駆け足で説明することになってしまった。説明不足な点があると思うので、疑問とか、試してみたけどうまくいかないとかあれば、コメント欄で質問してほしい。
前回やったこと
前回 の記事では、Elixir から Rust の関数を呼ぶことで、円周率の近似値を求める計算を高速化した。計算方法は「数値積分法の長方形近似」という、とても素朴な方法だった。
これを選んだのは次のような理由だった。
- 実装が非常に簡単
- 計算が適度に重い
- マルチコア CPU を活かしたプログラムを書きやすい
以下の段取りで、計算にかかった時間を計りながら開発を進め、１のプログラムよりも４のプログラムの方が 約65倍速い という結果が得られた。
- Elixir：単一の軽量プロセスで計算
- Elixir：複数の軽量プロセスでマルチコア計算
- Elixir：HiPE でネイティブコード化
- Elixir から Rust のマルチスレッドプログラムを呼び出す
当時、使用した環境は以下の通り。
- Core i7 プロセッサ搭載の2012年製 Mac mini
- Elixir 1.2.0 ＋ Erlang/OTP 18.2.1
- Rust 1.5.0 安定版
今回も同じ Mac mini を使うが、Rust のバージョンは上がっている（後述）。今回は Elixir は使用しない。具体的な計算方法や、前回使用した Elixir と Rust のプログラムについては、前回 の記事を参照していただきたい。
今回やること
今回は、前回最速だった Rust によるマルチスレッドプログラムを少し書き換えて、CPU が持つ SIMD(Single Instruction Multiple Data) 機構に対応させる。私の Mac に搭載されている CPU は Intel の AVX 命令に対応しているのだが、これなら、機械語の１命令で64ビット浮動小数点演算を４つ並列で実行できるので、計算がさらに高速化するはずだ。
今回は以下のバージョンの Rust を使用する。
- Rust 1.12.0 安定版
- Rust 1.14.0 nightly 版（2016-10-15）
ちなみに OS だが、この Mac にトリプルブートでインストールされている FreeBSD 10.3-RELEASE、Arch Linux、Mac OS X Yosemite El Capitan を使用した。どの OS でも同じ結果だった。
SIMD 化の２つの方法
Rust で SIMD 命令を使う方法は２つある。
- コンパイラの最適化に任せる
- コンパイラの SIMD intrinsic を使う
１は rustc コンパイラがバックエンドに使っている LLVM の最適化に任せる方法だ。cargo build --release または rustc -O とするだけで、配列に対する単純なループ処理くらいなら、自動的に SIMD を使った機械語命令へと最適化してくれる。この方法は、安定版の Rust を含めた全てのリリースチャネルで使用できるが、1.9.0 だけは、リリースビルドの作成時に問題があったようで、この機能がオフになってしまっている。なので、そのバージョンは避けたほうがいい。（現時点の最新版は 1.12.0）
また、最適化の内容について、詳しくは LLVM ドキュメントの「Auto-Vectorization in LLVM」を参照してほしい。
２は rustc コンパイラに用意されている SIMD 関連の intrinsic を使う方法だ。Intrinsic は、Rust で書かれた外部関数のような見た目をしているが、実際はコンパイラの組み込み関数だ。それを呼び出すと、対応する LLVM の中間コードが生成されるしくみになっている。
2016年10月現在では、SIMD 関連の intrinsic は非安定扱いなので、利用者側のコードで feature gate を通してこの機能を利用可能にする必要がある。そのため、Rust の安定版リリースとベータ版リリースでは、この機能は利用できない。つまり、SIMD intrinsic を利用するには、nightly 版の Rust を使用するか、適切な config パラメータを指定して Rust をソースコードからビルドする必要がある。今回は rustup で簡単にインストールできる nightly 版を使用する。
なお、Rust の SIMD 関連機能の開発状況は、Tracking issue for SIMD support #27731 で確認できる。SIMD 関連の intrinsic は、現状ではドキュメントがないようだが、こことか、こことか のコードを読むと雰囲気はつかめるだろう。
Intrinsic はプリミティブな関数なので、そのままでは使いづらい。これをラップした SIMD クレート を使うと、演算子オーバーロードなどの Rust ならではの抽象化機構を活かしたプログラミングが行える。今回はこれを使用する。
寄り道：コンパイラの最適化による SIMD 化を体験
その前に、コンパイラの最適化による SIMD 化を体験してみよう。
例として、巨大な配列 x と y の saxpy、つまり、a * x[i] + y[i] を計算するプログラムを使用する。以下の記事で紹介されていたプログラムを参考にした。
ただ、今回はそれを、コマンドライン引数の入力と結果の出力をするように修正してある。これは、コンパイラによる過度な最適化を防ぐためのものだ。もし計算結果を出力しないと、最悪の場合、最適化によって計算自体が省かれてしまう。
// -*- coding:utf-8-unix -*-
extern crate rand;
extern crate time;
use rand::{Rand, Rng, SeedableRng, StdRng};
use std::error::Error;
use std::ops::{Add, Mul};
fn print_help() {
println!("
Usage:
saxpy <size of array> [<item numbers in array z> ...]
");
}
// 時間計測のコードを挿入するマクロ
macro_rules! timeit {
($label: expr, $code: expr) => ({
let start = time::get_time();
let ret = $code;
let end = time::get_time();
let microsecs = (end.sec - start.sec) as f64 * 1_000_000.0 +
(end.nsec - start.nsec) as f64 / 1_000.0;
println!("{}: {} micro-seconds", $label, microsecs);
ret
})
}
fn main() {
// コマンドライン引数を数値に変換
match parse_args() {
Err(e) => {
println!("\nError: {}", e.description());
print_help();
}
Ok((size, item_numbers)) => {
// 乱数生成器を初期化（毎回同じシードを使用）
let seed = [1, 2, 3, 4];
let mut rng: StdRng = SeedableRng::from_seed(&seed[..]);
// a, x, y を乱数で初期化
let a: f64 = rng.gen();
let (x, y) = timeit!("init x and y", {
(init_array(&mut rng, size),
init_array(&mut rng, size))
});
// z を 0 で初期化
let mut z = timeit!("init z", vec![0.0; size]);
// saxpy、つまり、a * x[i] + y[i] を計算し、z[i] に格納する。
timeit!("saxpy", saxpy(a, &x, &y, &mut z));
// 結果を表示
println!("");
for i in item_numbers {
println!("{}: x:{:.4}, y:{:.4}, z:{:.4}", i, x[i], y[i], z[i]);
}
}
}
}
fn parse_args() -> Result<(usize, Vec<usize>), Box<Error>> {
let mut args = std::env::args().skip(1);
// 引数が1つもなければエラー
let size_str = try!(args.next()
.ok_or::<Box<Error>>(From::from("Please specify the size of the array.")));
// 配列の要素数を数値に変換できなかったらエラー
let size = try!(size_str.parse::<usize>());
let mut ns = Vec::with_capacity(args.len());
for n_str in args {
// 配列の要素番号を数値に変換できなかったらエラー
let n = try!(n_str.parse::<usize>());
// 要素番号が配列の範囲外だったらエラー
if n >= size {
return Err(From::from(format!("The item with the item number does not exist. \
(size of the array: {}, item number: {})\n\
NOTE: The array is 0-based.",
n,
size)));
}
ns.push(n);
}
Ok((size, ns))
}
fn init_array<R: Rng, T: Rand>(rng: &mut R, size: usize) -> Vec<T> {
// rng.gen_iter() は乱数を無限に生成するイテレータ。
// take(size) で size 個分の乱数を収集し Vec<T> に入れる。
rng.gen_iter().take(size).collect()
}
fn saxpy<T>(a: T, x: &[T], y: &[T], z: &mut [T])
where T: Copy + Add<T, Output = T> + Mul<T, Output = T>
{
for ((&xi, &yi), zi) in x.iter().zip(y.iter()).zip(z.iter_mut()) {
*zi = a * xi + yi;
}
}
SIMD 化の対象となるコード
saxpy() 関数のこのループ計算が、LLVM の最適化により SIMD 化される。
for ((&xi, &yi), zi) in x.iter().zip(y.iter()).zip(z.iter_mut()) {
*zi = a * xi + yi;
}
慣れないとわかりにくい書き方になっているが、こう書くのとほぼ同じだ。
for i in 0..x.len() {
z[i] = a * x[i] + y[i];
}
こういう単純な繰り返し処理は、機械的な書き換えで SIMD 化しやすい。
プロジェクトの作成
Rust の安定版を使用するので、もし Rust がインストールされてなかったら、こちらの記事を参考にインストールしてほしい。
インストールできたら、適当なディレクトリに移動して、プロジェクトを作成する。
$ cd 適当なディレクトリ
$ cargo new saxpy --bin
$ cd saxpy
そして、Cargo.toml に以下の依存クレートを追加する。
[dependencies]
rand = "0.3.14"
time = "0.1.35"
最後に、src/main.rs の内容を、上記のプログラムで置き換えれば OK だ。
最適化あり、target-feature 指定なしで実行する
最適化あり（--release フラグのみ指定）で実行しよう。
$ cargo run --release -- 50000000 0 100 49999999
-- 以降は、プログラムに与えるコマンドライン引数で、50000000（５千万）は配列の大きさ、それに続く数字は、配列 X、Y、Z について、要素番号 0、100、49999999 の値を表示しろという指示だ。
実行に成功すると、以下のように表示される。計算結果だけでなく、配列の初期化や計算にかかった時間もマイクロ秒で表示される。
init x and y: 738325.583 micro-seconds
init z: 133700.386 micro-seconds
saxpy: 99319.297 micro-seconds
0: x:0.1363, y:0.8950, z:0.9653
100: x:0.3087, y:0.4003, z:0.5596
49999999: x:0.3535, y:0.5883, z:0.7707
- x と y の初期化（init）処理の実行時間は、全項目を乱数で初期化するという計算の重い処理なので、全体の実行時間の76%を占めている。
- z の 0.0 による初期化もそれなりに時間がかかり、全体の約14%。
- saxpy の計算に要した時間は全体の約10%だった。
以下のようにすると、ビルドと同時に、アセンブリコードをファイルに保存できる。
$ RUSTFLAGS='--emit asm' cargo build --release
全ての x86_64 プロセッサは SSE2 命令に対応しているので、SIMD 命令が含まれているはずだ。
$ grep xmm target/release/saxpy.s
movdqu 1352(%rsp), %xmm0
movdqu 1368(%rsp), %xmm1
movdqa %xmm1, 160(%rsp)
movdqa %xmm0, 144(%rsp)
movdqa 144(%rsp), %xmm0
movdqa 160(%rsp), %xmm1
movapd 176(%rsp), %xmm2
movapd %xmm2, 96(%rsp)
movdqa %xmm1, 80(%rsp)
movdqa %xmm0, 64(%rsp)
movaps .LCPI8_0(%rip), %xmm0
movaps %xmm0, 1392(%rsp)
movdqa .LCPI8_1(%rip), %xmm0
movdqa %xmm0, 1408(%rsp)
xorps %xmm0, %xmm0
cvtsi2sdq %rax, %xmm0
mulsd .LCPI8_3(%rip), %xmm0
cvtsi2sdl %ecx, %xmm1
divsd .LCPI8_4(%rip), %xmm1
addsd %xmm0, %xmm1
movsd %xmm1, 352(%rsp)
...
$ grep -c xmm target/release/saxpy.s
99
やはり、SIMD 命令が含まれている。命令数は99個あった。
最適化あり、target-feature 指定ありで実行する
この Mac は、AVX 命令まで対応しているので、さらに最適化できる。RUSTFLAGS に -C target-feature=+avx を追加しよう。なお、この Mac の場合、-C target-cpu=native と指定しても同じ結果になる。
$ RUSTFLAGS='--emit asm -C target-feature=+avx' cargo run --release -- 50000000 0 100 49999999
init x and y: 731802.102 micro-seconds
init z: 132771.348 micro-seconds
saxpy: 95226.379 micro-seconds
0: x:0.1363, y:0.8950, z:0.9653
100: x:0.3087, y:0.4003, z:0.5596
49999999: x:0.3535, y:0.5883, z:0.7707
saxpy の計算のところが、約4%速くなった。
アセンブリコードを調べよう。
$ grep xmm target/release/saxpy.s
vxorps %xmm0, %xmm0, %xmm0
vcvtsi2sdq %rax, %xmm0, %xmm0
vmulsd .LCPI8_2(%rip), %xmm0, %xmm0
vcvtsi2sdl %ecx, %xmm0, %xmm1
vdivsd .LCPI8_3(%rip), %xmm1, %xmm1
vaddsd %xmm1, %xmm0, %xmm0
vmovsd %xmm0, 384(%rsp)
vmovq %rax, %xmm1
vmovq %rax, %xmm0
vmovdqa %xmm1, 464(%rsp)
vpunpcklqdq %xmm1, %xmm0, %xmm0
vmovdqa %xmm0, 960(%rsp)
vxorps %xmm0, %xmm0, %xmm0
vcvtsi2sdq %rax, %xmm0, %xmm0
vmulsd .LCPI8_2(%rip), %xmm0, %xmm0
vcvtsi2sdl %ecx, %xmm0, %xmm1
vdivsd .LCPI8_3(%rip), %xmm1, %xmm1
...
$ grep -c xmm target/release/saxpy.s
53
命令が変わったのにお気づきだろうか（例：addsd → vaddsd）。SSE2 から AVX になったことで、より効率のよい命令が使えるようになり、命令数が53に減っている。
SIMD クレートで円周率の近似計算を高速化する
では、いよいよ本題の、円周率の近似計算を SIMD 化しよう。
SMID クレート経由で、SIMD intrinsics を使うので、nightly 版の Rust が必要だ。私が使用したのは、1.14.0 の 2016-10-15 nightly ビルド版だ。rustup で全く同じバージョンをダウンロードするには、以下のようにする。
$ rustup update nightly-2016-10-17
日付が後ろにずれていることに注意。どうやら nightly 版の Rust は日曜日以外の毎晩ビルドされ、そのバイナリパッケージは翌日に作られるようだ。
適当なディレクトリに移動して、プロジェクトを作成する。
$ cd 適当なディレクトリ
$ cargo new calcpi --bin
$ cd calcpi
$ mkdir src/pi_fallback
$ mkdir src/pi_simd
このプロジェクトは nightly 版でビルドするように設定する。
$ rustup override add nightly-2016-10-17
$ rustup override list
/usr/home/tatsuya/rust-projects/calcpi nightly-2016-10-17-x86_64-unknown-freebsd
Cargo.toml の編集
Cargo.toml を以下のように編集する。
[dependencies]
simd = { version = "0.1.1", optional = true }
[features]
# Enable to use simd acceleration.
simd-accel = ["simd"]
依存しているクレートとして simd をオプションで指定し、さらに、simd-accel というフィーチャーを定義した。この後、src/lib.rs の中でコンパイル時に条件をチェックし、このフィーチャーが指定され、かつ、CPU が AVX 命令に対応している時だけ、SIMD 化が行われるようにする。
元の関数
SIMD 化前の calc_pi_range() 関数は以下の通り。計算内容について、詳しくは 前回 の記事を参照のこと。
pub fn display_name() -> String {
"calc_pi() with no simd support.".to_string()
}
pub fn calc_pi_range(n: u64, offset: u64, count: u64) -> f64 {
let w = 1.0 / (n as f64);
let mut s = 0.0;
for i in offset..(offset + count) {
let x = (i as f64) * w;
s += (1.0 - x * x).sqrt();
}
4.0 * w * s
}
このロジックだと、LLVM の自動 SIMD 化は働かなかった。
SIMD 化した関数
SIMD クレートに用意されている構造体やメソッドを使って SIMD 化するとこうなる。
use simd::x86::avx::{f64x4, AvxF64x4};
pub fn display_name() -> String {
"calc_pi() with avx simd support.".to_string()
}
pub fn calc_pi_range(n: u64, offset: u64, count: u64) -> f64 {
let w = 1.0 / (n as f64);
let mut sv = f64x4::splat(0.0);
let o64 = offset as f64;
let mut iv = f64x4::new(o64, o64 + 1.0, o64 + 2.0, o64 + 3.0);
for _ in 0..(count / 4) {
let xv = iv * f64x4::splat(w);
sv = sv + (f64x4::splat(1.0) - xv * xv).sqrt();
iv = iv + f64x4::splat(4.0);
}
let s = (0..4).fold(0.0, |acc, i| acc + sv.extract(i));
4.0 * w * s
}
f64x4 が f64（64ビット浮動小数点数型）を４並列で演算するための モジュール 構造体で、avx モジュール配下にある。calc_pi_range() では、f64x4 型の変数を見分けやすいよう、sv や iv のように、名前に v を付けている。
-
f64x4::splat(0.0)なら、４つとも0.0にセットされる。 -
f64x4::new(o64, o64 + 1.0, o64 + 2.0, o64 + 3.0)なら、４つがそれぞれ指定した値にセットされる。 - SIMD クレートは演算子オーバーロードをしてくれるので、
sv = sv + (f64x4::splat(1.0) - xv * xv).sqrt()のように書ける。便利。 -
(0..4).fold(0.0, |acc, i| acc + sv.extract(i))として、変数sの４つの値を順に取り出し、その合計を求めている。
呼び出し元の関数
これらを呼び出す側は以下の通り。
#![cfg_attr(feature = "simd-accel", feature(cfg_target_feature))]
#[cfg(feature = "simd-accel")]
extern crate simd;
use std::thread;
#[cfg(feature = "simd-accel")]
#[cfg(target_feature = "avx")]
#[path = "./pi_simd/calcpi.rs"]
mod calcpi;
#[cfg_attr(feature = "simd-accel", cfg(not(target_feature = "avx")))]
#[path = "./pi_fallback/calcpi.rs"]
mod calcpi;
const MAX_THREADS: u64 = 64;
// Require Rust 1.11.0 or newer because of sum()
pub fn calc_pi_parallel(n: u64, num_threads: u64) -> Result<f64, String> {
try!(check_pi_params(n, num_threads));
println!("Using {}", calcpi::display_name());
let len = n / num_threads;
let handles: Vec<_> = (0..num_threads)
.map(|i| thread::spawn(move || calcpi::calc_pi_range(n, len * i, len)))
.collect();
let results = handles.into_iter().map(|h| h.join().unwrap());
let pi = results.into_iter().sum();
Ok(pi)
}
fn check_pi_params(n: u64, num_threads: u64) -> Result<(), String> {
if num_threads <= 0 || num_threads > MAX_THREADS {
Err(format!("Invalid num_threads {}. It must be > 0 and <= {}",
num_threads,
MAX_THREADS))
} else if n % num_threads != 0 {
Err(format!("n must be a multiple of num_threads. (n = {}, num_threads = {})",
n,
num_threads))
} else if (n / num_threads) % 4 != 0 {
Err(format!("n / num_threads must be a multiple of 4. (n / num_threads = {})",
n / num_threads))
} else {
Ok(())
}
}
-
#[cfg(feature = "simd-accel")]、#[cfg(target_feature = "avx")]を使って、simd-accel フィーチャーが指定され、かつ、CPU が AVX 命令に対応している時だけ、SIMD 化されたcalc_pi_range()関数をコンパイルするようにしている。そうでない時は、元の SIMD 化されていない関数をコンパイルする。 -
calc_pi_parallel()は、プログラムの実行時にnum_threadsで指定された数のスレッドを立ち上げ、calc_pi_range()を呼び出す。詳しくは 前回 の記事を参照してほしい。
最後に main() 関数。
extern crate calcpi as pi;
fn main() {
let num_threads = 8;
let n = 500_000_000 * num_threads * 4;
match pi::calc_pi_parallel(n, num_threads) {
Ok(pi) => println!("pi = {}", pi),
Err(e) => println!("ERROR: {}", e),
}
}
- 四分円を
50億分割160億分割して近似する。（前回は10億分割） - ８スレッドに並列化。８つの論理コア（４コア x Hyper Threading）を使い切る
- もし simd-accel フィーチャーが指定されていたら、CPU コア毎に SIMD による４並列演算を行う。（その切り替えは、lib.rs 側でコンパイル時に実施する）
というわけで、最大32並列で、浮動小数点演算を回していこう。
実行する → はやい
では実行しよう。まずは、SIMD なし。８スレッドによる８並列演算。
$ rustc --version
rustc 1.14.0-nightly (98a3502da 2016-10-15)
$ cargo build --release
$ time ./target/release/calcpi
Using calc_pi() with no simd support.
pi = 3.141592653714508
./target/release/calcpi 132.06s user 0.02s system 788% cpu 16.747 total
続いて、SIMD あり。８スレッド x SIMD による32並列演算。
$ RUSTFLAGS='-C target-feature=+avx' cargo build --release --features=simd-accel
$ time ./target/release/calcpi
Using calc_pi() with avx simd support.
pi = 3.1415926537147785
./target/release/calcpi 65.80s user 0.00s system 792% cpu 8.303 total
小数点以下13桁目以降が微妙に違うのは、計算の順番が異なることによる誤差だろうか。なにはともあれ、約２倍の高速化を達成できた！
まとめ
- Rust で SIMD 化するには２つの方法がある
- コンパイラの最適化に任せる → 安定版 Rust を含む全てのリリースチャネルで利用できる。ごく簡単なパターンなら SIMD 化される。
- コンパイラの SIMD intrinsic を使う → 現時点では非安定。SIMD クレートを使うと、Rust らしいプログラミングスタイルのまま SIMD 化できる。
- 今回の計算内容（円周率の近似計算）と私の実行環境では、SIMD 化によって、計算速度が約２倍に高速化された。
- 今回のように重い計算処理では、SIMD 化の恩恵を受けやすい。