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

はじめに

　この記事の発端は，ピュアElixirとElixir+Rustlerの実行速度がどれくらい違うのか，ベクトル演算を行うコードで比較してみようという感じで書かれています．
　
シリーズとしてはElixirありきでRust(Rustler)を使う人をメインに据えています．

今回は Rust, Rustler の仕様について触れていきます．

開発環境

• OTP-21.0.9
• Elixir(1.7.3)
• Rustler(0.18.0)
• Rust(1.28.0)

事前調査と方針

　Rust で行列演算を行うライブラリはすでにいくつかあります．

• Rusty-machine (rulinalg): Machine Learning library for Rust
• ndarray : Arrays inspired by numpy.
• nalgebra : Low dimensional linear algebra.
• servo/euclid : Basic linear algebra used by Servo.
• cgmath : Linear algebra for game development.

参照：「Rustで数値計算してみた話」

ですが，今回は1から演算を行う関数を自作しました．
自前の関数はジェネリクス(C++でいうとテンプレート)には対応していません．

環境構築

本稿では扱いません．
「Elixirから簡単にRustを呼び出せるRustler #1 準備編」
こちらを参照してください．

mix.exsファイルのapp_nameやモジュール名はリンク先のものに準拠します

注意点

本稿で扱うRustlerはv.0.18.0なので，mix.exsのdepsの設定には気をつけましょう．
主要な型に違いがあります．Elixir-NIF-Rustボイラープレート Ruster0.17.1の注意点

mix.exs

defp deps do
[
  { :Rustler,   "~> 0.18.0"},
]
end

基本構文

ruslterのプロジェクトを作成するとadd関数が作成されています．

lib.rs

fn add<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
        let num1: i64 = try!(args[0].decode());
        let num2: i64 = try!(args[1].decode());

        Ok((atoms::ok(), num1 + num2).encode(env))
}

解説

見慣れない子<'a>がいますね．ライフタイムといいます．
非同期NIF呼び出しの際には気を配る必要がありますが，今は無視して問題ありません．

EnvはBEAM(Erlang VM)との交信に必要です．
argsはElixirからNIF経由で送られてきたデータですね． インデックスでアクセスしているように配列になっています．
NifResult<T>はResult<T, Error>と同値で，処理内容とエラーをまとめたものです．

上記のaddを見てもらうとdecodeではじまり，encodeで終わっています．
そのままでは使えないようです．

decodeの型

Rustでは型推論ができますが， decodeは型指定しないといけません．
ためしにaddの型指定を消してみましょう．

error[E0277]: cannot add `()` to `()`
  --> src/lib.rs:47:31
   |
47 |         Ok((atoms::ok(), num1 + num2).encode(env))
   |                               ^ no implementation for `() + ()`
   |
   = help: the trait `std::ops::Add` is not implemented for `()`

おそらく空のタプルになっており，計算できなくなってしまいます．
今回は単純なものをつくるので型多相(Polymorphic)にはしていませんが，今後対応させていくつもりです．
（多分Rustではなくて，Elixirの方で）

エラーハンドリング

try!()はエラーのチェックをして，早期リターンをさせることができます．
似た構文として，？があります．
使い分けは使いたい型がプリミティブ型（初期で定義されている型）かどうかです．

プリミティブ型の例(try!)

• i64
• f64
• Rustler::ListIterator

プリミティブでない型(?)

• Vec
• tuple

ベクトル演算(内積)

設定

関数を足す際にはexportの設定をしましょう．

lib.rs

Rustler_export_nifs! {
  "Elixir.NifExample",
  [
    //("Elixir's func, number of arguments, Rust's func)
    ("add", 2, add),
    ("dot_product", 2, dot_product),    
  ],
  None
}

忘れないうちにElixirの方でも書きます．

defmodule NifExample do
  use Rustler, otp_app: :phx_Rust, crate: :example

  def add(_a, _b), do: exit(:nif_not_loaded)
  def def dot_product(_a, _b), do: exit(:nif_not_loaded)
end

コード

それではベクトル演算の関数を見ていきましょう．

1. decode()でVec<i64>に変換
2. zip()で各配列の同じインデックスを持つ要素をペアにして，1つのtupleにまとめます
3. forで内積計算
4. encode()してNifResultを返す

lib.rs

fn dot_product<'a>(env: Env<'a>, args: &[Term<'a>])-> NifResult<Term<'a>> {
  // 1
  let x: Vec<f64> = args[0].decode()?;
  let y: Vec<f64> = args[1].decode()?;
    
  // 2      
  let tuple = x.iter().zip(y.iter());
  
  // 3
  let mut ans = 0.0;
  for (i, j) in tuple{
    ans = ans + (i*j);
  }
  
  // 4
  Ok((atoms::ok(), ans).encode(env))
}

使い方

iex -S mixでコンパイル，　関数を呼び出しましょう．

iex(1)> NifExample.dot_product [1.0, 2.0, 3.0], [4.0, 5.0, 6.0]
{:ok, 32.0}

パフォーマンス測定

要素1千万or一億をもつリストの内積計算を行う速度を比較します．

コード

以下，パフォーマンス測定に使用したコードです．

example.ex

defmodule NifExample do
  use Rustler, otp_app: :phx_Rust, crate: :example

  # func with nif
  def dot_product(_a, _b), do: exit(:nif_not_loaded)
  def add(_a, _b), do: exit(:nif_not_loaded)
  
  # func for list
  def new(e1, e2) when e1 > e2, do: []
  def new(e1, e2) do
    [e1] ++ new(e1+1, e2)
  end

  # elixir func
  def dot_product_ex(r1, _r2) when r1 == [], do: 0
    def dot_product_ex(r1, r2) do
      [h1|t1] = r1
      [h2|t2] = r2
      (h1*h2) + dot_product_ex(t1, t2)
  end

  # Benchmark
  # none = elixir
  #   1 = call nif, Rustler
  def dp_benchmark do
    m = new(1, 10_000_000.0)
    :timer.tc( fn ->
      dot_product_ex(m, m) 
    end)
    |>elem(0)
    |>Kernel./(1_000_000)
  end

  def dp_benchmark1 do
    m = new(1, 10_000_000.0)
    :timer.tc( fn ->
      dot_product(m, m)
    end)
    |>elem(0)
    |>Kernel./(1_000_000)
  end
end

解説

new

自前でnewを定義しています．理由は，工夫しないとElixirにrange（例　1..10_000_000）を渡せないのでリストに変換する必要があるからです．しかし，Enum.to_listがとっても遅いので，自前の関数を使っています．
詳しくは以下のリンクをどうぞ．
@zacky1972
「ZEAM開発ログ v.0.3.3 GPU駆動ベンチマークで時間を食っていた「ある処理」を最適化することで，驚きのパフォーマンス改善となった件」

リンク先のto_rangeが私の環境(Rust 1.28.0)だとコンパイルエラーになるので，まだ使えていません．

dot_product_ex

Elixir側で書いた内積計算プログラムです．

dp_benchmark dp_benchmark1

無印側がElixirでの内積計算，1がついた方がRustでの内積計算です．

実行結果

※Enum.to_listが遅いとかいっておきながら，リストの作成時間は実行時間に入っていません．

要素数	dp_benchmark(Elixir)	dp_benchmark1(Rust)	倍率
一千万	1.094141	0.47988	2.28
	1.169583	0.518724	2.25
一億	37.05859	9.495681	3.90
	33.451313	11.876747	2.81

振れ幅大きいですね．．．なんとかできないしょうか

まとめと感想

まとめとしてはこれだけです．

• elixir上のベクトル演算はRustlerを使用することで高速化が見込めそう

感想

• C++erなら馴染みやすいと聞いたんですが，まだよくわかりません． （おそらく私のC++知識が古い？ため）
• 型推論がある割に静的型付けをしないとdecodeで詰んでしまうので，結局英語で書かれたAPIやらを見に行かないといけないあたり敷居が高い気がします．
• Rustに限った話ではないと思いますが，コンパイルエラーが読めるかどうかで生産性が大きく変わることを再認識しました．

参考

• @twinbee 「Elixirから簡単にRustを呼び出せるRustler #1 準備編」
• 「Rustlerドキュメント」
• 「プログラミング言語Rust」
• 「[Rust] zipでfor文 & 速度比較」