ZEAM開発ログ v.0.3.3 GPU駆動ベンチマークで時間を食っていた「ある処理」を最適化することで，驚きのパフォーマンス改善となった件

この記事は「ZEAM開発ログ v.0.3.2 rayon によるSIMD(SSE2)マルチコアCPUによりOpenCL + GPUを上回るパフォーマンスが出た件」の続きです。

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

今回はベンチマーク高速化の第3弾です。 @twinbee さんがGPU駆動ベンチマークの各処理の実行時間を計測していて「ある処理」がとても時間を食っていることがわかったので，最適化してみました。

背景〜 @twinbee さんの気づき

共同研究している @twinbee さんがロジスティック写像のベンチマークをいろいろいじって時間を計測していくうちに，あることに気づきました。

「実はベンチマーク本体よりも，Enum.to_list で大きく時間を食っているんじゃ...?!」

具体的には次のようなコードがあるのですが，

  def map_calc_g2(x, p, mu, _stages) do
    x
    |> Enum.to_list
    |> LogisticMapNif.call_ocl2(p, mu)
    receive do
        l -> l
    end
  end

この Enum.to_list がベンチマーク本体 LogisticMapNif...よりも時間を食っていたんですね。

解決のために〜範囲オブジェクトをRustlerで受け取るには？

そもそもなんでこんなところに Enum.to_list を入れているかというと，引数に与えられた範囲オブジェクト(1..100みたいなやつです)を Rustler でリストとして受け取れないからなのですね。実際，Enum.to_list を外すと，実行時の型エラーになります。

そこで焦点は「範囲オブジェクトをいかにして Rustler で受け取るか？」という点になります。

調べてみると範囲オブジェクトは Map の一種であることがわかりました。また，Rustlerで型判定をする方法もわかりました。そこで，範囲オブジェクトまたはリストを受け取ってベンチマークを実行するように改良できました。

Rustler のコード

Elixirのコードは次のようにします。同様に全てのベンチマークの冒頭の Enum.to_list を削除します。

  def map_calc_g2(x, p, mu, _stages) do
    x
    |> LogisticMapNif.call_ocl2(p, mu)
    receive do
        l -> l
    end
  end

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

Rustのコードは次のような感じです。

native/logistic_map/src/lib.rs の共通部分

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

extern crate ocl;
extern crate rayon;
extern crate scoped_pool;

use rustler::{Env, Term, NifResult, Encoder, Error};
use rustler::env::{OwnedEnv, SavedTerm};
use rustler::types::list::ListIterator;
use rustler::types::map::MapIterator;
use rustler::types::binary::Binary;

use rustler::types::tuple::make_tuple;
use std::ops::Range;

use rayon::prelude::*;
use rayon::ThreadPool;

use ocl::{ProQue, Buffer, MemFlags};

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

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

lazy_static! {
    static ref POOL:scoped_pool::Pool = scoped_pool::Pool::new(8);
}

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

整数型の範囲オブジェクトを読取るのはこんな感じです。力技のコードですね。取り急ぎ，型が合わないとパニックを発生させます。

fn to_range(arg: Term) -> Range<i64> {
    match arg.decode::<MapIterator>() {
        Ok(iter) => {
            let mut vec:Vec<(Term, Term)> = vec![];
            for (key, value) in iter {
                vec.push((key, value));
            }
            match vec[0].0.atom_to_string() {
                Ok(struct_k) => {
                    if struct_k == "__struct__" {
                        match vec[0].1.atom_to_string() {
                            Ok(struct_v) => {
                                if struct_v == "Elixir.Range" {
                                    match vec[1].1.decode::<i64>() {
                                        Ok(first) => {
                                            match vec[2].1.decode::<i64>() {
                                                Ok(last) => {
                                                    std::ops::Range {start: first, end: last + 1}
                                                },
                                                Err(_) => panic!("argument error"),
                                            }
                                        },
                                        Err(_) => panic!("argument error"),
                                    }
                                } else {
                                    panic!("argument error")
                                }
                            },
                            Err(_) => panic!("argument error"),
                        }
                    } else {
                        panic!("argument error")
                    }
                },
                Err(_) => panic!("argument error"),
            }
        },
        Err(_) => panic!("argument error"),
    }
}

追記: 範囲オブジェクトの変換，もっとシンプルになりました！ @tatsuya6502 さん， @statiolake さん， @lo48576 さん， @termoshtt さん，ありがとうございます！

fn to_range(arg: Term) -> Result<RangeInclusive<i64>, Error> {
    let vec:Vec<(Term, Term)> = arg.decode::<MapIterator>()?.collect();
    match (&*vec[0].0.atom_to_string()?, &*vec[0].1.atom_to_string()?) {
        ("__struct__", "Elixir.Range") => {
            let first = vec[1].1.decode::<i64>()?;
            let last = vec[2].1.decode::<i64>()?;
            Ok(first ..= last)
        },
        _ => Err(Error::BadArg),
    }
}

範囲オブジェクトもしくはリストを受け取ってベクトルに変換するのはこんな感じです。型が合わない場合は Bad Argumentエラーを発生させます。

fn to_list(arg: Term) -> Result<Vec<i64>, Error> {
    match arg.is_map() {
        true => Ok(to_range(arg).collect::<Vec<i64>>()),
        false => match arg.is_list() {
            true => {
                let iter: ListIterator = try!(arg.decode());
                let res: Result<Vec<i64>, Error> = iter
                    .map(|x| x.decode::<i64>())
                    .collect();

                match res {
                    Ok(result) => Ok(result),
                    Err(_) => Err(Error::BadArg)
                }
            },
            false => Err(Error::BadArg)
        },
    }
}

追記: @twinbee さんのその後の調査と @tatsuya6502 さん， @statiolake さんの指摘で，もっとスッキリすることがわかりました。ありがとうございます！

fn to_list(arg: Term) -> Result<Vec<i64>, Error> {
    match (arg.is_map(), arg.is_list() || arg.is_empty_list()) {
        (true, false) => Ok(to_range(arg)?.collect::<Vec<i64>>()),
        (false, true) => Ok(arg.decode::<Vec<i64>>()?),
        _ => Err(Error::BadArg),
    }
}

そのほかのコードは下記の通りです。

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 trivial(x: Vec<i64>, p: i64, mu: i64) -> ocl::Result<(Vec<i64>)> {
    let src = r#"
        __kernel void calc(__global long* input, __global long* output, long p, long mu) {
            size_t i = get_global_id(0);
            long x = input[i];
            x = mu * x * (x + 1) % p;
            x = mu * x * (x + 1) % p;
            x = mu * x * (x + 1) % p;
            x = mu * x * (x + 1) % p;
            x = mu * x * (x + 1) % p;
            x = mu * x * (x + 1) % p;
            x = mu * x * (x + 1) % p;
            x = mu * x * (x + 1) % p;
            x = mu * x * (x + 1) % p;
            x = mu * x * (x + 1) % p;
            output[i] = x;
        }
    "#;

    let pro_que = ProQue::builder()
        .src(src)
        .dims(x.len())
        .build().expect("Build ProQue");

    let source_buffer = Buffer::builder()
        .queue(pro_que.queue().clone())
        .flags(MemFlags::new().read_write())
        .len(x.len())
        .copy_host_slice(&x)
        .build()?;

    let result_buffer: Buffer<i64> = Buffer::builder()
        .queue(pro_que.queue().clone())
        .flags(MemFlags::new().read_write())
        .len(x.len())
        .build()?;

    let kernel = pro_que.kernel_builder("calc")
        .arg(&source_buffer)
        .arg(&result_buffer)
        .arg(p)
        .arg(mu)
        .build()?;

    unsafe { kernel.enq()?; }

    let mut vec_result = vec![0; result_buffer.len()];
    result_buffer.read(&mut vec_result).enq()?;
    Ok(vec_result)
}

fn call_ocl2<'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 p        = args[1].in_env(env);
        let mu       = args[2].in_env(env);
        Ok(my_env.save(make_tuple(env, &[list_arg, p, mu])))
    })?;

    POOL.spawn(move || {
        my_env.send_and_clear(&pid, |env| {
            let result: NifResult<Term> = (|| {
                let tuple = saved_list.load(env).decode::<(Term, i64, i64)>()?;
                let p = tuple.1;
                let mu = tuple.2;
                match to_list(tuple.0) {
                    Ok(result) => {
                        let r1: ocl::Result<(Vec<i64>)> = trivial(result, p, mu);
                        match r1 {
                            Ok(r2) => Ok(r2.encode(env)),
                            Err(_) => Err(Error::BadArg),
                        }
                    },
                    Err(err) => Err(err)
                }
            })();
            match result {
                Err(_err) => env.error_tuple("test failed".encode(env)),
                Ok(term) => term
            }
        });
    });
    Ok(atoms::ok().to_term(env))
}

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])))
    })?;

    POOL.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 num = tuple.1;
                let p = tuple.2;
                let mu = tuple.3;

                match to_list(tuple.0) {
                    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))
}

ポイントですが，次のように書換えることで，範囲オブジェクトも受取れるようになります。

                match to_list(tuple.0) {
                    Ok(result) => Ok(...),
                    Err(err) => Err(err)
                }

なお，lazy_static でスレッドプールを確保していますが，現状では有効にできていません。現在調査中です。

実験結果

ここまででベンチマークを取り直しました。

stages	benchmarks1	benchmarks3	benchmarks8
	pure Elixir	pure Elixir	Elixir/Rustler
	loop	inlining inside of Flow.map	loop, passing by list, with Window
1	47.472903→47.904713	37.091920→36.558484	9.487438→7.690833
2	22.943836→22.640058	19.870662→19.067411	24.325771→24.836797
4	13.509106→13.701835	13.612359→12.308965	18.880198→22.318805
8	12.146374→12.149914	14.655723→13.063872	17.437002→19.550625
16	12.180955→12.252772	10.867108→11.196587	19.406703→21.662009
32	12.366302→12.407663	11.144992→10.578510	22.267963→23.749637
64	12.297703→12.273788	13.347286→10.784512	28.936603→29.402247
128	12.196510→12.177492	11.158467→11.052482	38.189437→38.949820

benchmarks_g2	benchmarks_t1	benchmarks_empty	Python
Elixir/Rustler	Elixir/Rustler	Elixir/Rustler	Python
OpenCL(GPU), inlining	rayon	Ruslter empty	NumPy, CPU
6.879992→2.388171	1.748167	4.520966→1.859177	17.749182

C	Rust CPU	Rust OpenCL	Rust rayon
clang	Rust	Rust	Rust
CPU, loop	CPU(1), loop	OpenCL(GPU), inlining	CPU(multi), loop
2.727346	2.931215→2.926270	1.561099→1.545693	0.669060

Elixir_Rustler_* の速度向上が凄まじく，ネイティブコード(Rust_*)に肉迫する勢いです！
やはり Enum.to_list が相当時間を食っていたんですね。。。

別環境(GCE)で CuPy の環境を構築できたので比較したのですが，今までは Elixir/Rustler/OpenCL と Python/CuPy/Cuda がほぼ同等の速度だったのですが，今回は Elixir/Rustler/OpenCL の方が Python/CuPy/Cuda の3.67倍の速度向上になりました！ ついに Python 以上の数値計算能力を手に入れることができました！

まとめ

• GPU を使用した Elixir / Rustler コードは CPUのみで実行する Elixir に対して 4.43〜8.23倍，Python に対して 7.43〜9.64 倍速いです。
• CPU 並列実行した方が GPU で実行するよりも Elixir/Rustler で 1.37倍，Rust で4.37倍速いです。
• GPU を使用した Elixir / Rustler コードは，GPUを使用したネイティブコードと比べて 1.48〜1.54倍の速度低下に止まります。
• Elixir/Rustler/OpenCL は，Python/CuPy/Cuda より3.67倍速いです！

これで「季節外れのアドベントカレンダー」は千秋楽を迎えましたが，今後も継続的に Qiita には書いていきますので，応援よろしくお願いします！ また，fukuoka.ex での発表も続けますし，2018年のアドベントカレンダーにも参戦します！

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

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