RubyKaigi 2022 セッション紹介「Method-based JIT compilation by transpiling to Julia」

はじめに

RubyKaigi2022に初参加してきましたので、そのセッションの中で面白かった内容を紹介しようと思います。

今回紹介するのは、KentaMurata氏のメソッドベースのJust-In-Timeコンパイルへの新しいアプローチとして、RubyのインフラストラクチャにJulia言語を使用した背景と仕組み、特徴についてになります。

スライドリンク

本題

現在Rubyでは数値計算が高速ではなく、Numo::NArrayやRed Arrowを使えば大きな数値計算が出来つつありますが、MJITやYJITが利用できてもあまり高速ではないという問題があるとのことでした。

理由として、これらのJITコンパイラがRubyの全てのセマンティクスを保持するためです。Rubyでは全てのメソッドが再定義可能で、再定義されたメソッドは直ちにコード実行に影響を与えます。例えば、以下のようなループの途中でも、injectメソッドや+演算子（メソッド）が再定義されていないかの確認が毎回行われます。

s = (1..10).inject { |a, x| a + x }

以上の特徴によって高速性が失われていましたが、数値計算の場合このRubyの動的性は数値計算アルゴリズムでは殆ど無意味なので、これを無視して計算を最適化できるほうが良いのでは？　というのがこのセッションの議題になります。

しかし、現在ではこのような最適化を行うためには、アルゴリズムをCの拡張ライブラリに書き換える必要があります。これをせずに、高速化を行う手段として挙げられたのがJuliaでした。Juliaはデータ処理や数値計算に向いていて高速な言語という特徴があります。

ここで、比較としてPythonの世界での解決策であるNumbaというライブラリの紹介がありました。NumbaはCPython用のJITコンパイラであり、PythonとNumPyのコードの一部を高速な機械語に変換します。もう少し具体的にコンパイルの流れを書きます。

• CPythonのバイトコードを解析
• NumbaIRを生成して書き換え
• 型を推論
• 型付IRに書き換え
• 自動並列化の実行
• LLVM IRの生成
• ネイティブコードにコンパイル

という流れで、CPythonを型付の中間表現に変換してネイティブコードを生成しています。

またNumbaには2つのモードがあります。オブジェクトモードというCPythonインタプリタのC APIを使用しCPythonの完全なセマンティクスを保持するモード。もう1つはnopythonモードというfloat64やnumpy配列などの特定のネイティブデータ型に特化した小さくて効率の良いネイティブコードを生成するモードです。ざっくり言うと処理をPython経由で行うか、CPUに直接命令するかの違いになります。

このNumbaのnopythonモードのようなものをRubyのJITコンパイラでも実現する手段として登場するのが、今回のセッションの本題であるJuliaでした。まず、RubyでNumbaライクなJITコンパイラを表現すると以下のようになります。

• Rubyメソッド
  • ASTの生成
  • 最適化
  • バイトコードの生成
• CRubyのバイトコード
  • IRの生成
  • タイプ推論
  • 最適化
• 型付けされたIR
• LLVM IRの生成
• ネイティブコードにコンパイル

これと同じようなことをやっているのが、Juliaになります（JuliaはNumbaともだいたい同じ機構が動いています）

• Juliaコード
  • ASTの生成
• Julia AST
  • タイプ推論
  • IRの生成
• Julia typed IR
• LLVM IRの生成
• ネイティブコードにコンパイル

そこで、RubyをJuliaへトランスパイルし、それ以降をJuliaのJITコンパイラに実行してもらって高速化します。

• Rubyメソッド
• Juliaコード
• Julia AST
• Julia typed IR
• LLVM IRの生成
• ネイティブコードにコンパイル

Juliaは最適化されたネイティブコードを生成してくれるため、高速です。この辺りの解説はセッションのスライド図と発表が大変分かりやすく、面白かったので是非資料や動画をご覧ください。

次に、RubyからJuliaへのトランスパイル方法についての紹介がありました。トランスパイルには、yadriggyというRubyメソッドのASTを構築し、構文と型をチェックするgemを使用しているとのことでした。セッションでは具体的なコードを紹介していましたが、大雑把にいうとRubyコードからASTを構築し、それを使ってJuliaコードを生成しているようでした。

Rubyコード → Ruby AST → 型チェッカーでノードに対して型付け → Juliaコード

これによって、Rubyの機械への命令を最適化させています。

またRubyとJuliaで実装の違うものがいくつかあり（例ではRangeを挙げていました）これらの対応をするには自分で変換コードを書く必要があるとのことでした。

これらの高速化の実験比較として、以下の計算を行っていました。

• マンデルブロ集合
• モンテカルロ法によるπの近似
• クイックソート
• 畳み込み
• ドット積

それぞれの結果は以下のようになっていました。RubyはおそらくYJITが有効。

• マンデルブロ集合
  • Ruby：平均3.326ms
  • Ruby to julia：平均171.667μs
• モンテカルロ法によるπの近似
  • Ruby：平均106.368ms
  • Ruby to julia：平均8.851ms
• クイックソート
  • Ruby：平均7.551ms
  • Ruby to julia：平均1.937ms
• 畳み込み（結果が複雑なので省略）
• ドット積
  • Ruby (N=10000)
    • 平均468.608μs
    • Ruby to julia (N=10000)
    • 平均3.759ms
  • Ruby to julia (N=10000, T=Float64)
    • 平均10.651μs

これらのように、一部の結果を除いて超高速に計算することが可能になるようでした。

紹介は以上になります。発表が分かりやすく、深掘りしたくなるような興味深い内容でした！Rubyの高速化についてこういう方法もあるんだなと、とても学びが多かったです。RubyKaigiではこのような発表がいくつもあり、すごくワクワクしました。