こんにちは、 simomu です。今日はRuby JIT Challenge をやってみた話をします。
Ruby JIT Challenge について
Ruby 3.3 からは RJIT と呼ばれる実験的な JIT コンパイラが登場しました。既に導入されている Ruby の JIT コンパイラである YJIT は Rust で実装されているのに対し、RJIT は Ruby 自身で書かれています。JIT コンパイラが Ruby 自身で記述されているため、コンパルを行うメソッドを上書きすることによって既に用意されている JIT コンパイラとは異なるコンパイラに差し替えることが可能です。
Ruby JIT Challenge は、この RJIT の基盤を使い、チュートリアルを通して Ruby の JIT コンパイラを自分の手で作成していくものです。
Ruby JIT Challenge では以下のようなフィボナッチ数を求めるプログラムを JIT コンパイルできるようにすることをゴールとしています。
def fib(n)
if n < 2
return n
end
return fib(n - 1) + fib(n - 2)
end
Ruby JIT Challenge を少しだけ触ってみる
チュートリアルの最初は nil を返すだけのメソッドをコンパイルするのが目標です。
def none
nil
end
none
none
p none
RJIT はメソッドが一定回数呼び出された時にマシンコードにコンパイルされ、Ruby JIT Challenge 上での設定では3回目の呼び出しで JIT コンパイルされます。
Ruby JIT Challenge ではあらかじめアセンブラが用意されていて、以下のような Ruby コードでアセンブリを書くことができます。
asm = Assembler.new
asm.mov(:r8, 1) # mov r8, 1
asm.mov(:r9, 2) # mov r9, 1
asm.add(:r8, ,:r9) # add r8, r9
これらを利用し、nil を返すだけのメソッドをコンパイルするコードは以下のようになります。
STACK = [:r8, :r9, :r10, :r11]
def compile(iseq)
stack_size = 0
# 中略
in :putnil
asm.mov(STACK[stack_size], 0x04)
stack_size += 1
in :leave
asm.add(:rsi, C.rb_control_frame_t.size)
asm.mov([:rdi, C.rb_execution_context_t.offsetof(:cfp)], :rsi)
asm.mov(:rax, STACK[stack_size - 1])
asm.ret
# 中略
end
Ruby のオブジェクトの内部表現は VALUE という形式で表現されていて、アセンブリでも VALUE を扱う必要があります。MRI の実装の↓を見ると、Ruby の nil を VALUE で表現すると 0x04 になることがわかるため、レジスタにはこの値を入れます。
in :putnil
asm.mov(STACK[stack_size], 0x04)
これを --rjit-dump-disasm オプションをつけて実行すると、実際に生成されたマシンコードをダンプした上で実行されます。
$ bin/ruby --rjit-dump-disasm test/none.rb
0x55555a222000: mov r8, 4
0x55555a222007: add rsi, 0x40
0x55555a22200b: mov qword ptr [rdi + 0x10], rsi
0x55555a22200f: mov rax, r8
0x55555a222012: ret
nil
ここで、本当に自作した JIT コンパイラによって実行された結果なのかどうかを確かめてみます。先程のコンパイラの実装で nil を生成するバイトコード命令の putnil のコンパイル時に、nil ではなく true 生成するように変更してみます。nil のときと同様に true も VALUE として表現する必要があり、MRI の実装の↓を見ると true は 0x14 であることがわかります。
これを参考に、nil の代わりに true を生成するように修正します。
in :putnil
asm.mov(STACK[stack_size], 0x14)
次にテストコードを修正し、2回目の none の実行結果も標準出力するように変更します。2回目の none の時点ではまだ JIT コンパイルが行われておらず nil が返ってきますが、3回目の実行時点では JIT コンパイルされたマシンコードを実行するため、true が返ってくるようになってしまうはずです。
def none
nil
end
none
p none # まだ JIT コンパイルされていないため、インタプリタ実行され `nil` が返る
p none # JIT コンパイルされて実行されるため、`true` が返るようになってしまう
実際に実行していみると以下のようになります。3回目の実行以降は true が返るようになり、自分で実装した JIT コンパイラによって実行されていることが確認できます。
$ bin/ruby --rjit-dump-disasm test/none.rb
nil
0x55555a222000: mov r8, 0x14
0x55555a222007: add rsi, 0x40
0x55555a22200b: mov qword ptr [rdi + 0x10], rsi
0x55555a22200f: mov rax, r8
0x55555a222012: ret
true
Ruby JIT Challenge を進めるうえで必要なこと
上記のような形で Ruby で JIT コンパイラを作成していき、最終的にはフィボナッチ数を求めるプログラムをコンパイルできるようにしていくわけですが、
- x86_64 のアセンブリそのもの
- Ruby のバイトコード命令列の中身
- Ruby のスタックフレームに関することとメソッド呼び出しの手順
などの知識が必要になります。特に、条件分岐を行う命令である branchunless のコンパイルが難しく、後ろの命令をコンパイルしないとジャンプ先のアドレスがわからない等の問題を解決するために、いわゆるバックパッチのような仕組みを実装する必要があります。
社内勉強会資料より引用
ヒントも詳細に書かれているので、丁寧に読み進めれば問題なさそうではありますが、チャレンジする場合はメソッド呼び出し・条件分岐あたりが大きな山場になると思います。
まとめ
今回は Ruby JIT Challenge を触ってみました。Ruby の JIT コンパイラだけではなく、Ruby のバイトコード命令や RubyVM のスタックフレームの管理などを学ぶ入口にもなったのかもしれないという感想です。
もともと社内勉強会のネタとして触り始めたものでしたが、 実際の社内勉強会では時間が全く足りず後半のメソッド呼び出しや条件分岐文のコンパイルの説明は省略せざるを得ませんでした。機会があれば社内勉強会資料もどこかで公開しようと思います。
最後に私の実装を置いておきますが、メソッド呼び出しや条件分岐文のコンパイル等が自力で行う難易度が高かったため、ほぼヒント通りの実装になっています。
余談
Ruby JIT Challenge をやっている最中に、 README.md にかかれているヒントのコードにミスがあったので修正の PR を出しました。
メソッド呼び出し命令の一つである opt_send_without_block において、 Ruby のスタックフレームである cfp ( Control Frame Pointer ) をプッシュした後に、呼び出すメソッドのレシーバを cfp にセットする命令を修正したものです。
Ruby でメソッドを呼び出すときには、レシーバ→引数の順番でスタックに積むバイトコード命令が来るようです。ヒント上でスタックとして利用しているレジスタからレシーバを取得する際には STACK[stack_size - C.vm_ci_argc(cd.ci) - 1] となっていて、スタックのトップから呼び出すメソッドの引数の数だけ下にあるものを取得してきています。レシーバを cfp にセットする場合は この値を sub するのではなく mov するのが正しそうというものでした。
おそらくオリジナルの RJIT だと↓あたりに相当する部分だと思われますが、あまり自信はありません。
Discussion