RustElixirで線形回帰を高速化した話

(この記事は「fukuoka.ex Elixir／Phoenix Advent Calendar 2018」の9日目です)
昨日の記事は @mocamocaland さんの「PhoenixをGCPにデプロイする方法」でした．

---

どうもfukuoka.exキャストの@hisawayです．

ElixirからRustler(NIF)経由でRustのコードを呼び出したら，どれだけコードが速くなるか，線形回帰で検証しました． @zacky1972 さんの「ZEAM開発ログ」，特にHastegaシリーズの先行探検隊な位置付けです. HastegaというのはElixirからSIMD:マルチコアCPU/GPU 並列処理を行う機能の総称です．

より詳しいHastegaについてはZEAM開発ログ: Elixir マクロ + LLVM で超並列プログラミング処理系を研究開発中をご覧ください．

(タグがRustなのがきになってしょうがないです．タイトル以外はRustに統一してます)

概要

Elixirの速さ

ElixirはErlangVM(BEAM)上で動く仕様であるため，ネイティブコードではなく，バイトコードを出力します．
Python, Rubyと比べると速いんですが, 比較対象にCやC++を挙げてしまうと遅いです．

HiPEというオプションを有効にすることでネイティブコードでコンパイルすることができますが，速度を重視する際には，やはり元からスピードのある言語を使用する必要があります．
(HiPEはErlangのコンパイルオプションです．)
|> @sile 「HiPE(High Performance Erlang)について」

そこでRustの力を借ります．ElixirからRustを使おうって記事は先人の方が投稿しております．
|> @tatsuya6502 さん 「ElixirからRustの関数をつかう → はやい」
|> @twinbee さん 「Elixirから簡単にRustを呼び出せるRustler #1 準備編」

用語

NIF

Native Implemented Functionsの略で，Elixirの下で動いているErlangVM(BEAM)からネイティブコードを使う仕組みのことです．

Rustler

NIF経由でRustの関数を呼びだすElixirのモジュールです．Rustで記述した関数とElixirの関数を対応づけて記述することで楽にコーディングすることができます．

Rustler_export_nifs! {
  "Elixir.LinearRegressorNif", // Elixirのモジュール名を指定
  [
    //("Elixir's func-name, number of arguments, Rust's func)
    ("_fit", 5, nif_fit), 
  ],
  None
}

SIMD

Single Instruction Multiple Dataの略です．たくさんのデータに対して，１つだけ処理を与えるような命令をいいます．「このベクトルの全要素を2倍してよ」みたいな感じです．
RustはデフォルトでSIMD命令にコンパイルしてくれます．

題材について

高速化する対象が線形回帰なのは，Elixirで機械学習しようという話になったとき，古典的な機械学習で取り組みやすいとして， @piacere_ex さんがElixirで書かれた線形回帰のコードを提供してくださったことが理由です．

高速化の方針

以下の２つの方法で高速化を行いました．アルゴリズムの最適化はしません．

1. Rustler経由でネイティブコードの呼び出し及び，NIFの非同期呼び出し
2. Rust の Rayonを使用したCPUマルチコア並列実行

Rayonは連続して計算をするコードを簡単に並列計算仕様にできるデータ処理ライブラリです．
|> Crate Rayon

備考

今回は線形回帰を高速化させるのが目的ではなく，Elixirのネイティブコード化による速度向上を評価することが目的です．
もし劇的な高速化を図る際は，アルゴリズムから組み直してくださいね．

動作環境

PC

• MacBook Air(Early 2015)

  • Processor: 1.6 GHz Intel Core i5
  • Memory: 8 GB 1600 MHz DDR3
  • Graphics: Intel HD Graphics 6000 1536MB

• iMac Pro(2017)

  • Processor: 2.3 GHz Intel Xeon W (プロセッサ数 1，物理コア 18，論理コア 36)
  • Memory: 32 GB 2666 MHz DDR4
  • Graphics: Radeon Pro Vega 64 16368MB

• Mac Pro (Mid 2010)

  • Processor:
  • 2 x 3.46 GHz 6-Core Intel Xeon (プロセッサ数 2，物理コア数 6 x 2，論理コア数 12 x 2)
  • Memory: 16 GB 1066 MHz DDR3
  • Graphics: ATI Radeon HD 5770 1024MB
  • OS : Sierra 10.12.6

• GPGPUサーバー ユニットコム UCGPU-E2630V4- 32GB

  • Processor: 2.20GHz Intel Xeon E5-2630 v4 (プロセッサ数 2 物理コア数 10コア x 2，論理コア数 20コア x 2)
  • Memory: 32 GB 2400 MHz DDR4
  • Graphics: NVIDIA GeForce GTX 1080 Ti
  • OS : Linux Ubuntu 16.04

• 言語

  • Elixir 1.7.3(Mac Proは 1.7.4でした)
  • Erlang/OTP 21 [erts-10.1] [source] [64-bit] [smp:40:40] [ds:40:40:10] [async-threads:1] [hipe] [sharing-preserving]
  • rustc 1.28.0
  • Rustler 0.18.0

性能評価

1. inline: Elixirのみ，一部機能をインライン展開して最適化したコード
2. Rust: Rustlerを使って，ElixirからRustのコードを呼び出し
3. Rayon: Rustのコードに加えて一部を並列処理

inlineを基準として倍率を測定しています．

MacBook Air(Early 2015)

|| inline | Rust | Rayon|
|:-:|:-:|:-:|:-:|:-:|
|実行時間[s]|7.240421|4.860835|7.394763|
|倍率|(1.00)|1.48|0.97|

Mac Pro(Mid 2010)

|| inline | Rust | Rayon|
|:-:|:-:|:-:|:-:|:-:|
|実行時間[s]|7.066738|3.64614|8.535108 |
|倍率|(1.00)|1.93|0.82|

iMac Pro(2017)

|| inline | Rust | Rayon|
|:-:|:-:|:-:|:-:|:-:|
|実行時間[s]|4.159493|3.01175|7.962244|
|倍率|(1.00)|1.38|0.52|

GPGPUサーバー

|| inline | Rust | Rayon|
|:-:|:-:|:-:|:-:|:-:|
|実行時間[s]|6.643913|3.113606|6.690754|
|倍率|(1.00)|2.13|0.99|

評価総括

• 純粋にネイティブコードに落とすと，1.38 ~ 2.13 倍ほど速くなりました．
• 並列処理にすると，0.52 ~ 0.99倍と遅くなりました
• iMac Proだけinlineが異常に速い(原因は未調査)

実装

線形回帰のメインであるfittingを行う関数です．
大まかな流れは，

1. NIF経由で起動されたスレッドから別スレッドを立てる準備をする
2. 元のスレッドは制約があるのでさっさと終了する
3. 別スレッドの方でメインの計算をする

といった感じです．

すぐかはわかりませんが，後ほどリポジトリを公開する予定です！

native/linear_regressor_nif/src/lib.rs#L95-L145

fn nif_fit<'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 _x = args[0].in_env(env);
    let _y = args[1].in_env(env);
    let theta = args[2].in_env(env);
    let alpha = args[3].in_env(env);
    let iterations = args[4].in_env(env);
    Ok(my_env.save(make_tuple(env, &[_x, _y, theta, alpha, iterations])))
  })?;

  std::thread::spawn(move ||  {
    my_env.send_and_clear(&pid, |env| {
      let result: NifResult<Term> = (|| {
        let tuple = saved_list
        .load(env).decode::<(
          Term, 
          Term,
          Term,
          Num,
          i64)>()?; 
        
        let x: Vec<Vec<Num>> = try!(tuple.0.decode());
        let y: Vec<Vec<Num>> = try!(tuple.1.decode());
        let theta: Vec<Vec<Num>> = try!(tuple.2.decode());
        let alpha: Num = tuple.3;
        let iterations: i64 = tuple.4;

        let tx = transpose(&x);
        let m = y.len() as Num;
        let (row, col) = (theta.len(), theta[0].len());
        let tmp = alpha/m;
        let a :Vec<Vec<Num>> = vec![vec![tmp; col]; row]; 

        let ans = (0..iterations)
          .fold( theta, |theta, _iteration|{
           sub2d(&theta, &emult2d(&mult( &tx, &sub2d( &mult( &x, &theta ), &y ) ), &a))
          });

        Ok(ans.encode(env))
      })();
      match result {
          Err(_err) => env.error_tuple("test failed".encode(env)),
          Ok(term) => term
      }  
    });
  });
  Ok(atoms::ok().to_term(env))
}

解説

ElixirとRustの通信

ElixirからRustに引数を渡すのは1手間かかります．Elixirから渡されるデータはErlangVMが読むためのバイトコードなので，コンパイル時に型推論ができません．ちゃんと指定してあげましょう．

|>「Elixir + Rustlerでベクトル演算を高速化しよう〜Rustler初級者編 1 〜」

またElixirで引数を渡す際には，あらかじめリストの中身をすべて浮動小数点型にしましょう．

# Rust側の関数と対応づけ
def _fit(_x, _y, _theta, _alpha, _iteration), do: exit(:nif_not_loaded)

def to_float(num) when is_integer(num), do: num /1
def to_float(r) when is_list(r) do
    r
    |>Enum.map( &to_float(&1) )
end

def fit( x, y, theta, alpha, iterations ) do
   _fit(
      x |> to_float, 
      y |> to_float, 
      theta |> to_float, 
      alpha |> to_float,
      iterations)
    receive do
      l -> l
    end
  end

Rustの方から結果を受け取るにはreceiveを使います．

NIFの制約

NIFはElixirとの通信，内部でのデータ処理を1ms以下に抑えないと，パフォーマンスが大幅に低下します．
そのため，Rust側でスレッドを立てて，そこから計算結果を受信する構造になっています．

参照や所有権について

Rustではデータにつき，1つの変数が所有権をもっています．そのため，基本的に読み取り専用の参照で関数に渡すと速くなりますし，その後のコードが書きやすいです．

もし参照を渡さない場合は，関数に設定した引数に元の変数の所有権が移ってしまうためにその関数が処理を終えた後，再利用できなくなってしまいます．以前はこの対処に事前にclone()して関数に渡すという方法をとっていました．
今思えば，プログラマーに参照渡しを強制させられる仕様になっているように思えますね．

並列処理部

下記コードがメインの処理部分です．

native/linear_regressor_nif/src/lib.rs#L131-L134

let ans = (0..iterations)
  .fold( theta, |theta, _iteration|{
    sub2d(&theta, &emult2d(&mult( &tx, &sub2d( &mult( &x, &theta ), &y ) ), &a))
  });

マルチコアで並列実行する際にはiterをpar_iterに変更します．また，一番外側の処理を並列化すると効果が高いです．今回の場合だと，行列の各要素で減算を行うsub2d関数に変更を加えます．

変更前

native/linear_regressor_nif/src/lib.rs#L52-L56

pub fn sub2d(x: &Vec<Vec<Num>>, y: &Vec<Vec<Num>>) -> Vec<Vec<Num>>{
  x.iter().zip(y.iter())
  .map(|t| sub(&t.0.to_vec(), &t.1.to_vec()))
  .collect()
}

変更後

native/linear_regressor_nif/src/lib.rs#L52-L56

pub fn sub2d(x: &Vec<Vec<Num>>, y: &Vec<Vec<Num>>) -> Vec<Vec<Num>>{
  x.par_iter().zip(y.par_iter())
  .map(|t| sub(&t.0.to_vec(), &t.1.to_vec()))
  .collect()
}

Rayonはiter->par_iterと変えるだけで並列処理ができる便利なライブラリですが，zip()と併用するとめっちゃ遅いみたいなことをどこかで見た覚えがあるので，近いうちに調査しようと思ってます．

今後への繋ぎ

本稿では，Elixirのネイティブコード化，CPUマルチコアを生かした並列処理をコードに組み込んだ場合のパフォーマンスを測定しました．

上記の変更だけでは，パフォーマンスは劇的には伸びませんでしたね．その要因の1つに処理自体がそんなに重くないというのがありますが，手動で最適化・高速化しようとしてみると，とても手間がかかる上に結果がよろしくないというあまりおもんない結果でした．

そのため，Hastegaを有効に使うためには普通に動かす場合，CPUマルチコア，GPUマルチコアのどれが速くなるのか判断する機能が必要になりますね．並列処理させる方法としては，Elixirのプロセスもあげられるのでプロセス間の通信方法を改良する(Flowの改良?)とかもするかもしれません．

研究計画

今後の取り組みとして，

1. 大きめの線形的な擬似モデルを生成する機能を作る(現実的なモデルとして存在する規模かは一旦置いとく)

2. SIMD & CPUマルチコアにおける性能評価を再度行う．

3. RustのOpenCLラッパーであるCrate oclを使ったGPGPUのコードを評価を行います．

出来上がり次第記事にします．

最後に

研究の全体構想について知りたい方はこちらを合わせてお読みください．

「ZEAM開発ログ2018年総集編その1: Elixir 研究構想についてふりかえる(前編)」

2018年12月15日に公開予定:
「fukuoka.ex Elixir／Phoenix Advent Calendar 2018」15日目
「ZEAM開発ログ2018年総集編その2: Elixir 研究構想についてふりかえる(後編)」

追記

リポジトリを公開しました
https://github.com/zeam-vm/linear_regressor_cli

今週中(~2018/12/16まで)にコメントとご指摘いただいた内容と上記で示した研究計画を実施，それとドキュメントを整備する予定です．

---

明日の「fukuoka.ex Elixir／Phoenix Advent Calendar 2018」10日目の記事は， @curry_on_a_rice さんの「Elixirでニューラルネットを実装しようとした話」」です。こちらもお楽しみに！