開発部 R&D ユニットの村田です。OSSの開発をしております。本記事では、Ruby で書かれたマンデルブロ集合を計算するメソッドを実行時に Julia に変換して実行するとめっちゃ速くなる (場合がある)、という話をします。

はじめに

Ruby 3.1 では YJIT がマージされ、Rails アプリケーションが速くなりました。今後のバージョンアップがとても楽しみですね。ただし、Ruby のデータ処理対応を進めている身としては、データ処理や数値計算がより高速になって欲しいと思っています。

データ処理や数値計算を高速化する試みとして、Python では NUMBA というライブラリが開発されています。NUMBA は、メソッド単位でバイトコードを LLVM を用いてネイティブコードにコンパイルすることでメソッド実行を高速化します。ただメソッドをネイティブコードに変換するのではなく、実行時にメソッドに渡された引数の型情報を利用した最適化も行います。この仕組みはとても面白く、Ruby でも似たようなことをやりたいと常々思っていました。

ところで、データ処理や数値計算に向いていて速いプログラミング言語といえば Julia ですよね。Julia の裏側には LLVM があります。Julia では、実行時に関数単位で引数の型に対して最適化されたネイティブコードを LLVM よって生成してから実行しています。つまり、NUMBA とだいたい同じ機構が働いています。

Ruby のメソッドを Julia に書き換えて実行できれば、NUMBA のような機構を Ruby でも実現できそうです。

実験

Ruby のメソッドを Julia に書き換えて実行する仕組みを試作しました。実装は以下のリンク先のスクリプトです。

github.com

実験はマンデルブロ集合を計算する次のメソッド群を対象としました。

def abs2(z)
  z.real * z.real + z.imag * z.imag
end

def mandel(z)
  c = z
  maxiter = 80
  (1 ... maxiter).each do |n|
    return n - 1 if abs2(z) > 4
    z = z**2 + c
  end
  maxiter
end

def mandelbrot
  ary = []
  (-1 .. 1.0).step(0.1) do |i|
    (-2 .. 0.5).step(0.1) do |r|
      ary << mandel(Complex(r, i))
    end
  end
  ary
end

スクリプトを実行すると、この mandelbrot メソッドの Ruby 版と Julia に書き換えたものの実行時間を計測して表示します。

Julia への書き換えは次のようにメソッドオブジェクトを変換用メソッド jl_trans に渡すだけです。

include JLTrans
jl_trans method(:mandelbrot)

実験用スクリプトの --show-jl オプションを渡すと、書き換え後の Julia コードが表示されます。上記の Ruby コードは次のような Julia コードに変換されます*1。

$ ruby -Ilib examples/jltrans.rb --show-jl
function mandelbrot()
    ary = []
    for i in RbCall.RubyRange(-1, 0.1, 1.0)
        for r in RbCall.RubyRange(-2, 0.1, 0.5)
            append!(ary, mandel(Complex(r, i)))
        end
    end
    ary
end

function mandel(z::ComplexF64)
    c = z
    maxiter = 80
    for n in RbCall.RubyRange(1, 1, maxiter)
        if abs2(z) > 4
            return n - 1
        end

        z = z ^ 2 + c
    end
    maxiter
end

function abs2(z::Any)
    real(z) * real(z) + imag(z) * imag(z)
end

実験結果

さて、スクリプトの実行結果を見てみましょう。

$ ruby -v
ruby 3.1.2p20 (2022-04-12 revision 4491bb740a) [x86_64-linux]

$ julia -v
julia version 1.7.2

$ ruby -Ilib examples/jltrans.rb
         name,   min,   max,  mean,   std
mandelbrot_rb, 5.803, 6.347, 6.021, 0.096
mandelbrot_jl, 0.265, 4.612, 0.276, 0.051

$ ruby --enable=mjit --mjit-min-calls=1 -Ilib examples/jltrans.rb
         name,   min,   max,  mean,   std
mandelbrot_rb, 4.875, 6.283, 5.113, 0.197
mandelbrot_jl, 0.259, 4.256, 0.275, 0.047 

$ ruby --enable=yjit --yjit-call-threshold=1 -Ilib examples/jltrans.rb
         name,   min,   max,  mean,   std
mandelbrot_rb, 4.214, 4.836, 4.372, 0.078
mandelbrot_jl, 0.259, 4.256, 0.268, 0.047

上から順に、デフォルト状態、MJIT 有効状態、そして YJIT 有効状態です。出力されている数値はそれぞれ実行時間の最小値、最大値、平均値、標準偏差です。値の単位はミリ秒です。評価対象のコードを合計実行時間が2秒を超えるか呼び出し回数が5回に達するまで呼び出し続け、毎回の実行時間を計測し、その計測した実行時間の統計量を表示しています。 mandelbrot_rb は Ruby で実装されたメソッド、mandelbrot_jl はそれを Julia に書き換えた関数群による計測結果です。

JIT なしの場合に平均6.0ミリ秒かかっていたメソッドの実行時間が MJIT によって平均5.1ミリ秒に、YJIT によって平均4.3ミリ秒に短縮されています。JIT の効果はすごいですね。一方、Julia に書き換えた関数群の実行時間は平均0.27秒で、JIT なしの場合に対して約22倍の高速化が達成されています。Julia すごい！最初に選んだ題材を Julia が有利になりそうなコードにしておいてよかったw

Ruby to Julia 実行時変換の仕組み概説

Ruby のメソッドを Julia に書き換えて実行する仕組みは、次のような手順で Ruby のメソッドを Julia に書き換えます。

対象となる Ruby メソッドの AST を作る
AST に型チェッカーを適用し、ノードに対して可能な限り型づけを行う
型つき AST を捜査して Julia のコード生成を行う
Ruby to Julia ブリッジの eval メソッドを使って Julia 側で生成されたコードを評価する
Julia 側に定義された関数を呼び出す Ruby 側のメソッドを定義する

これらの手順は、Ruby to Julia トランスパイラ (1〜3)、Ruby to Julia ブリッジ (4)、インターフェイス (5) の3パートに大きく分けられます。

Ruby to Julia トランスパイラを構成するメソッドの AST を作る解析器、AST に型づけを行う型チェッカー、そして Julia コード生成器は Yadriggy というライブラリを利用して実装しました。Yadriggy は、Ruby 文法のサブセットになるような文法を持つ DSL を実装するための仕組みを提供してくれるライブラリです。

Ruby to Julia ブリッジは、私が開発している pycall.rb の Julia 版に相当するものです。たとえば、上記の変換後の Julia コードで参照されている RbCall.RubyRange はブリッジによって提供されているものです *2。

この仕組みの将来性

現時点では、トランスパイラとブリッジの双方について、今回の実験対象である mandelbrot メソッドを処理するために必要な最小限の機能のみが実装されています。どの程度実用的なのかはまだ分かりませんが、機能を増やしていくと面白いことはできそうな気がします。

例えば、Julia の多次元配列を Ruby の配列に変換せずに Ruby から直接操作できるようなラッパーをブリッジに実装することで、numo-narray や numpy.rb の代わりのライブラリとして実用的に使える可能性があります。

また、numo-narray と Julia の多次元配列が MemoryView を介してデータ共有できるようにして、さらに、トランスパイラが numo-narray のデータを適切に扱えば、numo-narray を利用して Ruby で実装されている計算処理をデータ変換のコストを発生させずに Julia で高速実行できるようになるかもしれません。

Julia 自体が数値計算機能を持っていることと、pycall.rb では見込めなかった処理の高速化が可能になる点がとても面白いと思います。

まとめ

Ruby のメソッドを Julia に書き換えて実行する仕組みを試作しました。マンデルブロ集合を計算する処理を対象に実行時間の変化を計測したところ約22倍の高速化が実現されました。Ruby to Julia トランスパイラと Ruby to Julia ブリッジの双方にメリットがありそうなので、開発を続けて適用可能範囲を広げていきたいです *3。

おしらせ

Speeeでは一緒にサービス開発を推進してくれる仲間を大募集しています！もしSpeeeに興味を持っていただいた方は以下で社内メンバーのカジュアル面談を公開しているので、お気軽にご連絡ください💁

tech.speee.jp

Speeeでは様々なポジションで募集中なので、「どんなポジションがあるの？」と気になってくれてた方は、こちらチェックしてみてください！もちろんオープンポジション的に上記に限らず積極採用中です！！！

*1:引数に不要な型アノテーションが残っていたりしますが、試作品なので見逃してください。

*2:これは、Ruby の Range#step と同じ値を生成する Julia の StepRangeLen オブジェクトを作る関数です。

*3:RubyKaigi 2022 に応募して採択してもらえたら、この実験の詳細や続報を話したいなと思っています

Speee DEVELOPER BLOG

Speee開発陣による技術情報発信ブログです。メディア開発・運用、スマートフォンアプリ開発、Webマーケティング、アドテクなどで培った技術ノウハウを発信していきます！

Ruby を Julia に変換して実行すると速くなる (場合がある)

はじめに

実験

実験結果

Ruby to Julia 実行時変換の仕組み概説

この仕組みの将来性

関連研究

まとめ

おしらせ