ZEAM開発ログ v.0.3.2 rayon によるSIMD(SSE2)マルチコアCPUによりOpenCL + GPUを上回るパフォーマンスが出た件

この記事は「ZEAM開発ログ v.0.3.1 非同期NIFにより1ms制約をクリアしてGPU駆動のパフォーマンスが改善された件」の続きです。

「ZEAM開発ログ 目次」はこちら

今回はベンチマーク高速化の第2弾です。Rustの並列実行ライブラリであるrayonを導入して，マルチコアCPU並列実行をやってみたいと思います。

背景その1〜Elixir/FlowとRustlerの相性が悪い！

Rustlerでプログラミングしていて気づいた事実です。

• Elixir単体の時はFlowでstages(並列実行で使用するコア数)を増やすときれいに速度が向上する
• それに対し，Rustlerではstages=1のときが最速で，stagesを増やしても速度が向上しない

そこでかねてより Rust 上で並列プログラミングしてみたいなと思っていました。

背景その2〜GPU駆動のボトルネック

また，OpenCL でGPUを駆動した場合にボトルネックになるのが，引数となる配列をGPUに転送し，結果で得られた配列をGPUから転送する処理であることがわかっています。だいたい実行時間の半分くらいを消費しています。OpenCL で CPU を並列実行した場合も同様です。

そこで，かねてより CPU を SIMD 命令を駆使して並列プログラミングしたときのパフォーマンスを試したいと思っていました。

調査〜Rustでの並列プログラミング/SIMD命令サポート事情

調査してみると，Rustでは並列実行ライブラリの開発がとても盛んであることがわかりました。私たちはその中でプログラミングが簡単な rayon を採用しました。 特徴としては，イテレーターと map 関数を使うようなプログラミングで，iter()の代わりに par_iter() と書くだけで並列実行してくれます。 @twinbee さんが調べてくださいました！ ありがとうございます！

また，Rust でイテレーターと map 関数を使うようなプログラミングの場合，x86_64 アーキテクチャでは自動的に SIMD 命令 (SSE2) を使用してコンパイルしてくれることもわかりました。実際に下記の記事を参考にして出力されるアセンブリコードを確認しました。

Rustの関数でSIMDをつかう → もっとはやい

コンパイルオプションを変えることで，AVX命令も使えます。今回，対応するマシンを持っていないので試せませんでしたが。。。

rayon を使った Rustler プログラミング

では，開発したコードを紹介します。

Elixir のプログラムはこんな感じです。

lib/logistic_map.ex

defmodule LogisticMap do
  ...
  def map_calc_t1(list, num, p, mu, _stages) do
  	list
  	|> Enum.to_list
  	|> LogisticMapNif.map_calc_t1(num, p, mu)
  	receive do
  		l -> l
  	end
  end
  ...
end

lib/logistic_map_Nif.ex

defmodule LogisticMapNif do
  use Rustler, otp_app: :logistic_map, crate: :logistic_map
  ...
  def map_calc_t1(_list, _num, _p, _mu), do:
  :erlang.nif_error(:nif_not_loaded)
  ...
end

native/logistic_map/Cargo.toml に rayon を下記のように追記します。

[dependencies]
...
rayon = "1.0"

そして native/logistic_map/src/lib.rs では次のようにします。

#[macro_use] extern crate rustler;
#[macro_use] extern crate lazy_static;

extern crate rayon;

use rustler::{Env, Term, NifResult, Encoder, Error};
use rustler::env::{OwnedEnv, SavedTerm};
use rustler::types::list::ListIterator;
use rustler::types::tuple::make_tuple;

use rayon::prelude::*;

mod atoms {
    rustler_atoms! {
        atom ok;
        //atom error;
        //atom __true__ = "true";
        //atom __false__ = "false";
    }
}

rustler_export_nifs! {
    "Elixir.LogisticMapNif",
    [("map_calc_t1", 4, map_calc_t1)],
    None
}

fn loop_calc(num: i64, init: i64, p: i64, mu: i64) -> i64 {
    let mut x: i64 = init;
    for _i in 0..num {
        x = mu * x * (x + 1) % p;
    }
    x
}

fn map_calc_t1<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let pid = env.pid();
    let mut my_env = OwnedEnv::new();

    let saved_list = my_env.run(|env| -> NifResult<SavedTerm> {
        let list_arg = args[0].in_env(env);
        let num      = args[1].in_env(env);
        let p        = args[2].in_env(env);
        let mu       = args[3].in_env(env);
        Ok(my_env.save(make_tuple(env, &[list_arg, num, p, mu])))
    })?;

    std::thread.spawn(move || {
        my_env.send_and_clear(&pid, |env| {
            let result: NifResult<Term> = (|| {
                let tuple = saved_list.load(env).decode::<(Term, i64, i64, i64)>()?;
                let iter: ListIterator = try!(tuple.0.decode());
                let num = tuple.1;
                let p = tuple.2;
                let mu = tuple.3;
                let res: Result<Vec<i64>, Error> = iter
                    .map(|x| x.decode::<i64>())
                    .collect();

                match res {
                    Ok(result) => Ok(result.par_iter().map(|&x| loop_calc(num, x, p, mu)).collect::<Vec<i64>>().encode(env)),
                    Err(err) => Err(err)
                }
            })();
            match result {
                Err(_err) => env.error_tuple("test failed".encode(env)),
                Ok(term) => term
            }
        });
    });
    Ok(atoms::ok().to_term(env))
}

今回の核心部分はここです。

Ok(result.par_iter().map(|&x| loop_calc(num, x, p, mu)).collect::<Vec<i64>>().encode(env)),

result.iter().map(...)となっていたところを，result.par_iter().map(...) とするだけです。超簡単ですね。

実行結果

実行結果は次のような感じです。

$ mix run -e "LogisticMapNif.init; LogisticMap.benchmark8(1); LogisticMap.benchmark8(8); LogisticMap.benchmarks_g2; LogisticMap.benchmarks_t1"
Compiling NIF crate :logistic_map (native/logistic_map)...
    Finished release [optimized] target(s) in 0.14s
LogisticMapNif_map_calc_list: 4100
LogisticMapNif_map_calc_binary: 3400
LogisticMapNif_map_calc_binary_to_binary: 2400
LogisticMapNif_call_ocl: 1

stages: 1
12.988499
stages: 8
18.014392
stages: 1
6.79009
stages: 1
6.191977

Rustler で単一スレッドで実行したときには12.99秒かかっているのに対し，Flow で Rustler を8並列で駆動したときには，18.01秒かかっています。かえって遅くなりましたね。

また非同期NIF呼出しでGPU駆動したときには6.79秒かかります。

これに対し，非同期NIF呼出しでrayonでCPU並列実行したときには6.192秒です。GPUを少し上回りました！

スレッドプールについて

今回のテスト環境は例によって次の環境です。

Mac Pro (Mid 1010)
2.8 GHz Quad-Core Intel Xeon
16GB 1066 MHz DDR3
ATI Radeon HD 5770 1024 MB

同じプログラムを Linux で実行すると遅くなることがわかっています。理由は，Linux の場合はスレッドの新規作成コストが高いからです。

これを回避するために，スレッドプールを使うというのが推奨されています。rayon では次のようにスレッドプールを作成します。

lazy_static! {
    static ref _THREAD_POOL: ThreadPool = rayon::ThreadPoolBuilder::new().num_threads(32).build().unwrap();
}

このコードは，Rust 単体では有効に機能して Linux 環境下で大幅に速度改善することを確認してます。

ところがなんと，Elixir / Rustler では有効に機能しないのですよね！

今後もこの問題の解決については継続的に研究したいと思っています。

まとめ

• Ruslterではrayonを用いて並列実行させることができます。
• その結果，MacではGPUを少し上回る実行速度が得られました！
• Linux ではスレッドの新規作成コストが高いので，実行速度が遅くなります。この問題の解決のため，スレッドプールが用いられます。
• Elixir / Rustler ではどういうわけか，スレッドプールを有効にできません。

次は「ZEAM開発ログ v.0.3.3 GPU駆動ベンチマークで時間を食っていた「ある処理」を最適化することで，驚きのパフォーマンス改善となった件」です。いよいよ「季節外れのアドベントカレンダー」のフィナーレです。お楽しみに！

p.s.「いいね」よろしくお願いします

よろしければ，ページ左上の や のクリックをお願いしますー
ここの数字が増えると，書き手としては「ウケている」という感覚が得られ，連載を更に進化させていくモチベーションになりますので，もっとElixirネタを見たいというあなた，私たちと一緒に盛り上げてください！