みなさーん、プログラマの石川某くんはいつもデスマにはまっていていつJITの勉強しているんだろと思ったことはありませんか？その秘密は、mruby-meta-circular。Cも機械語も使いません。だから、忙しい某君にもできるのです。サンプルCoDeを無料で送ります。続きはWebで
というわけで(どういうわけだ(定番の突っ込み))、mrubyとmrubyのVMであるRITE VMの命令だけでTracing JITの仕組みを解説しようと思います。
mrubyのJITと基本的な構造は大体同じです。
インタープリタ
初めにベースになるVMを見てみましょう。おそらく世の中で2番目にたくさん作られているプログラム、フィボナッチ級数それだけが動くよう注意深く命令セットが選ばれています。これから先、このVMをFibVMと呼びます。また、mrubyのVMはRITE VMです。
この2つを区別することが非常に重要です。
class FibVM
include RiteOpcodeUtil
OPTABLE_CODE = Irep::OPTABLE_CODE
OPTABLE_SYM = Irep::OPTABLE_SYM
def initialize
# For Interpriter
@stack = [] # スタック(@spより上位をレジスタとして扱う)
@callinfo = [] # メソッド呼び出しで呼び出し元の情報を格納
@pc = 0 # 実行する命令の位置
@sp = 0 # スタックポインタ
@cp = 0 # callinfoのポインタ
@irep = nil # 現在実行中の命令列オブジェクト
@irepid =nil # 命令列オブジェクトのid(JIT用)
end
def eval(irep)
@irep = irep
@irepid = @irep.id
while true
# 命令コードの取り出し
cop = @irep.iseq[@pc]
case OPTABLE_SYM[get_opcode(cop)]
# 何もしない
when :NOP
# MOVE Ra, RbでレジスタRaにレジスタRbの内容をセットする
when :MOVE
@stack[@sp + getarg_a(cop)] = @stack[@sp + getarg_b(cop)]
# LOADL Ra, pb でレジスタRaに定数テーブル(pool)のpb番目の値をセットする
when :LOADL
@stack[@sp + getarg_a(cop)] = @irep.pool[getarg_bx(cop)]
# LOADI Ra, n でレジスタRaにFixnumの値 nをセットする
when :LOADI
@stack[@sp + getarg_a(cop)] = getarg_sbx(cop)
# LOADSELF Ra でレジスタRaに現在のselfをセットする
when :LOADSELF
@stack[@sp + getarg_a(cop)] = @stack[@sp]
# ADD Ra, Rb でレジスタRaにRa+Rbをセットする
when :ADD
@stack[@sp + getarg_a(cop)] += @stack[@sp + getarg_a(cop) + 1]
# SUB Ra, n でレジスタRaにRa-nをセットする
when :SUBI
@stack[@sp + getarg_a(cop)] -= getarg_c(cop)
# EQ Ra でRaとR(a+1)を比べて同じならtrue, 違うならfalseをRaにセットする
when :EQ
val = (@stack[@sp + getarg_a(cop)] == @stack[@sp + getarg_a(cop) + 1])
@stack[@sp + getarg_a(cop)] = val
# JMP nでpcをnだけ増やす。ただし、nは符号付き
when :JMP
@pc = @pc + getarg_sbx(cop)
next
# JMPNOT Ra, nでもしRaがnilかfalseならpcをnだけ増やす。ただし、nは符号付き
when :JMPNOT
if !@stack[@sp + getarg_a(cop)] then
@pc = @pc + getarg_sbx(cop)
next
end
# メソッドの先頭で引数のセットアップする命令。面倒なので詳細は省略
when :ENTER
# SEND Ra, mid, anumでRaをレシーバにしてシンボルmidの名前のメソッドを
# 呼び出す。ただし、引数はanum個あり、R(a+1), R(a+2)... R(a+anum)が引数
when :SEND
a = getarg_a(cop)
mid = @irep.syms[getarg_b(cop)]
n = getarg_c(cop)
newirep = Irep::get_irep(@stack[@sp + a], mid)
if newirep then
@callinfo[@cp] = @sp
@cp += 1
@callinfo[@cp] = @pc
@cp += 1
@callinfo[@cp] = @irep
@cp += 1
@sp += a
@pc = 0
@irep = newirep
@irepid = @irep.id
next
else
args = []
n.times do |i|
args.push @stack[@sp + a + i + 1]
end
@stack[@sp + a] = @stack[@sp + a].send(mid, *args)
end
# RETURN Raで呼び出し元のメソッドに戻る。Raが戻り値になる
when :RETURN
if @cp == 0 then
return @stack[@sp + getarg_a(cop)]
else
@stack[@sp] = @stack[@sp + getarg_a(cop)]
@cp -= 1
@irep = @callinfo[@cp]
@irepid = @irep.id
@cp -= 1
@pc = @callinfo[@cp]
@cp -= 1
@sp = @callinfo[@cp]
end
else
printf("Unkown code %s \n", OPTABLE_SYM[get_opcode(cop)])
end
@pc = @pc + 1
end
end
end
VMのレジスタは@stackに、メソッドの呼び出し履歴は@callinfoに、プログラムカウンタは@pcに格納されています。mrubyのソースコード中のsrc/vm.cでmrb_context_runで実装されているRITE VMのと比べるとよく似ていると感じられると思います。
ではフィボナッチを動かしてみましょう
def fib(n)
if n == 1 then
1
elsif n == 0 then
1
else
fib(n - 1) + fib(n - 2)
end
end
def fibt
fib(20)
end
a = Irep::get_irep(self, :fibt)
vm = FibVM.new
p vm.eval(a)
p fibt
$ time ../../mruby/bin/mruby jit-1.rb
10946
10946
real 0m0.630s
user 0m0.592s
sys 0m0.016s
Tracing JITは命令の実行と同時にそれに対するネイティブコードを生成します。ただし、今回はネイティブコードといってもx86等の機械語ではなく、RITE VMになります。
実行回数を数えるコードを付け加える
世の中にはパレートの法則というものがあって、プログラムの2割が8割の実行時間を費やしているのだそうです。まあ、そんな感じもするなーと思います。
コード生成はそれなりに重い処理なので、全部の命令に対してネイティブコードを生成するのはあほらしいです。そういうことで、命令の実行回数を測ってたくさん実行されるようならネイティブコードを生成するようにします。
命令の実行回数によってはネイティブコードを生成しないことから勘のいい方は気づかれたかもしれませんが、Tracing JITではコンパイル後もすべてがネイティブコードで実行するのではなく、必要に応じてVM(インタープリタ)に戻ってくることができます。このことから、コード生成をサポートしない命令があっても良かったり、大胆な最適化をして条件に当てはまらない場合は、VMで実行するなんてことができます。
ただし、VMに戻る場合はVM側の状態を正しく設定する必要があります。速いコードを生成しようとすると大変ですが、とりあえずVMと全く同じ働きをするコードを生成するようにすれば勝手にVM側の状態が正しく設定されていますので簡単です。具体例はあとからお見せします。
VMにJITコンパイラを付け加える前段階として、実行回数を数える処理とJITコンパイラで使う変数を付け加えます。こんなコードになります。
class FibVM
include RiteOpcodeUtil
OPTABLE_CODE = Irep::OPTABLE_CODE
OPTABLE_SYM = Irep::OPTABLE_SYM
def initialize
# For Interpriter
@stack = []
@callinfo = []
@pc = 0
@sp = 0
@bp = 0
@cp = 0
@irep = nil
@irepid =nil
# For JIT
@prof_info = {} # プロファイラ用
@proc_tab = {} # 生成したProcを入れておく
@entry = nil # 現在コンパイル中のコードの先頭(RITEのコード)
# 現在コンパイル中かのフラグも兼ねる
@max_using_reg = 2 # 現在使用しているレジスタの最大番号
@code = [] # コンパイルして生成されたコード
@pool = [] # コンパイルして生成された定数テーブル
end
def eval(irep)
@irep = irep
# 命令列の生アドレスを得る
@irepid = @irep.id
# irepごとにエントリーを作成する
@prof_info[@irepid] ||= []
@proc_tab [@irepid] ||= []
while true
# ネイティブコードの実行(当然今はない)
cop = @irep.iseq[@pc]
# @pcに対応するRITE VMコードがない場合
# (今はまだコード生成しないから常に無い)
if !@proc_tab[@irepid][@pc] then
# 今のpcの命令の実行回数１回増やす
@prof_info[@irepid][@pc] ||= 0
@prof_info[@irepid][@pc] += 1
times = @prof_info[@irepid][@pc]
if times > 20 then
# コンパイラの宝石箱や
end
end
case OPTABLE_SYM[get_opcode(cop)]
when :NOP
以下同文
mrubyでは実行中の命令を指し示すには次の２つのデータを指定する必要があります
- irep 命令列、メソッド・ブロックごとに1つづずある
- pc irep中のどの命令を実行しているか示す。
そのため、実行回数を記憶するprof_infoはirepとpcの2つの要素をとる2次元配列です。(正確には配列を要素とするハッシュテーブルです。)どういったirepやpcが来るかあらかじめわからないから、その都度必要に応じて要素を増やしています。@prof_info[@irepid] ||= []はRubyの慣用句で、@prof_info[@irepid]が空(nil)の時、配列で初期化します。
@irepと@irepidの違いを説明します。@irepはmrubyのIREPをオブジェクト化したものですがオブジェクトとして扱うためヘッダが外側にあり、同じIREPでもアドレスが異なる場合があります。そこで、IREPの生アドレスをIrep#idというメソッドで得られるようにし、IREPが同じかどうかを厳密に判断したい場合はそちらを使います。@irepidは@irep.idが格納されています。
コンパイラの宝石箱
先ほどのプログラム中のコメント「でコンパイラの宝石箱や」と書いたところの詳細を見てみます。ここは、コード生成の心臓部です。
if @entry == nil then
@entry = @pc
@code.push mkop_AB(OPTABLE_CODE[:MOVE], 0, 1)
end
case OPTABLE_SYM[get_opcode(cop)]
when :NOP
when :MOVE
tmp = @max_using_reg
@max_using_reg += 1
@code += gen_get_reg(tmp, getarg_b(cop))
@code += gen_set_reg(getarg_a(cop), tmp)
@max_using_reg -= 1
when :LOADL
sidx = add_pool(@irep.pool[getarg_bx(cop)])
tmp = @max_using_reg
@max_using_reg += 1
@code.push mkop_ABx(OPTABLE_CODE[:LOADL], tmp, sidx)
@code += gen_set_reg(getarg_a(cop), tmp)
@max_using_reg -= 1
when :LOADI
tmp = @max_using_reg
@max_using_reg += 1
@code.push mkop_AsBx(OPTABLE_CODE[:LOADI], tmp, getarg_sbx(cop))
@code += gen_set_reg(getarg_a(cop), tmp)
@max_using_reg -= 1
when :LOADSELF
tmp = @max_using_reg
@max_using_reg += 1
@code += gen_get_reg(tmp, 0)
@code += gen_set_reg(getarg_a(cop), tmp)
@max_using_reg -= 1
when :ADD
tmp0 = @max_using_reg
tmp1 = tmp0 + 1
@max_using_reg += 2
@code += gen_get_reg(tmp0, getarg_a(cop))
@code += gen_get_reg(tmp1, getarg_a(cop) + 1)
@code.push mkop_ABC(OPTABLE_CODE[:ADD], tmp0, ADD_SYM, 1)
@code += gen_set_reg(getarg_a(cop), tmp0)
@max_using_reg -= 2
when :SUBI
tmp0 = @max_using_reg
@max_using_reg += 1
@code += gen_get_reg(tmp0, getarg_a(cop))
@code.push mkop_ABC(OPTABLE_CODE[:SUBI], tmp0, SUB_SYM, getarg_c(cop))
@code += gen_set_reg(getarg_a(cop), tmp0)
@max_using_reg -= 1
when :EQ
tmp0 = @max_using_reg
tmp1 = tmp0 + 1
@max_using_reg += 2
@code += gen_get_reg(tmp0, getarg_a(cop))
@code += gen_get_reg(tmp1, getarg_a(cop) + 1)
@code.push mkop_ABC(OPTABLE_CODE[:EQ], tmp0, EQ_SYM, 1)
@code += gen_set_reg(getarg_a(cop), tmp0)
@max_using_reg -= 2
when :ENTER
# Do nothing
when :JMP
# Do nothing
when :JMPNOT
exit_code = gen_exit(0)
tmp0 = @max_using_reg
@max_using_reg += 1
@code += gen_get_reg(tmp0, getarg_a(cop))
off = exit_code.size
if @stack[@sp + getarg_a(cop)] then
@code.push mkop_AsBx(OPTABLE_CODE[:JMPNOT], tmp0, off)
@code += exit_code
else
@code.push mkop_AsBx(OPTABLE_CODE[:JMPIF], tmp0, off)
@code += exit_code
end
@max_using_reg -= 1
else
# サポートされていない命令
stop_compile
end
区切って詳細を説明します。
コードの先頭
if @entry == nil then
@entry = @pc
# You are not expected to understand this
@code.push mkop_AB(OPTABLE_CODE[:MOVE], 0, 1)
end
@entryは現在コンパイル中のコードの先頭の@pcの値が入ります。これがnilということは、コンパイルしていないという意味です。つまり、@entryはコンパイル中かどうかのフラグも兼ねています。
次のMOVE命令はちょっと難しいです。@pcや@stackをインスタンス変数としてアクセスするために、self(R0)を引数として渡されたcallerのself(これはR1に格納される)で書き換えます。Rubyレベルではselfを書きかえることはできませんが、RITEレベルでは平気で出来るのです。
MOVE命令
when :MOVE
tmp = @max_using_reg
@max_using_reg += 1
@code += gen_get_reg(tmp, getarg_b(cop))
@code += gen_set_reg(getarg_a(cop), tmp)
@max_using_reg -= 1
MOVE命令のコンパイルです。ちなみに、NOP命令はやはり何もコードを出しません。
gen_get_reg, gen_set_regはCodeGenモジュールのメソッドで、FibVMのレジスタの読み書きを行います。CodeGenモジュールはコード生成時によく使う生成パターンのメソッドをまとめたモジュールです。詳しくは後から説明します。
この辺から今書いているコードはRITE VMのレジスタにアクセスしたいのか、FibVMのレジスタにアクセスしたいのか意識していないと混乱します。
@max_using_regは現在値が入っているもっとも大きなレジスタの番号でそれより大きな番号のレジスタをgen_get_regなどで作業用レジスタとして使おうという魂胆です。つまり、RITE VMはレジスタマシンですが、スタックマシンとして使っているわけです。RITE VMがスタックマシンだとすると@max_using_regはスタックポインタです。
LOADL命令(定数テーブルの生成)
when :LOADL
sidx = add_pool(@irep.pool[getarg_bx(cop)])
tmp = @max_using_reg
@max_using_reg += 1
@code.push mkop_ABx(OPTABLE_CODE[:LOADL], tmp, sidx)
@code += gen_set_reg(getarg_a(cop), tmp)
@max_using_reg -= 1
add_poolは定数テーブル(pool)に値を格納するためのCodeGenモジュールのメソッドです。プールのインデックスを返します。もし、すでに同じ値のエントリーがあれば再利用します。
add_poolはこんな感じの定義です。
# 定数テーブルに定数を追加する。すでにある場合は再利用する
def add_pool(val)
if idx = @pool.index(val) then
return idx
else
idx = @pool.size
@pool.push val
return idx
end
end
JMP命令、JMPNOT命令 と ガード
when :JMP
# Do nothing
when :JMPNOT
exit_code = gen_exit(0)
tmp0 = @max_using_reg
@max_using_reg += 1
@code += gen_get_reg(tmp0, getarg_a(cop))
off = exit_code.size
if @stack[@sp + getarg_a(cop)] then
@code.push mkop_AsBx(OPTABLE_CODE[:JMPNOT], tmp0, off)
@code += exit_code
else
@code.push mkop_AsBx(OPTABLE_CODE[:JMPIF], tmp0, off)
@code += exit_code
end
@max_using_reg -= 1
JMP命令は何も命令を生成しないのに注意してください。Tracing JITは順番に実行した命令を並べるだけなのでジャンプ命令は意味がないのです。
無条件分岐は無視してもいいですが、条件分岐は無視できません。しかし、通常の条件分岐のようにコンパイルするわけでもないです。JMPNOTではFibVMの真偽を決めるレジスタ@stack[@sp + getarg_a(cop)]のコンパイル時の値とコンパイルされたRITE VMのコードを実行するときの真偽を決めるレジスタの値を比較します。この値が同じなら生成したコードを引き続き実行できます。もし、異なればこの先のコードは実行する条件を満たしていないのでFibVMに戻るコードを生成します。
gen_exitはCodeGenモジュールのメソッドでFibVMに戻る命令列を生成します。
このように、レジスタなどの値が期待された値か調べて、期待外れならFibVM(より上のレベルのVM)に戻るようにする仕組みのことをガードと言います。ガードは条件分岐だけではなく、足し算で型が整数かどうかチェックするとかいろいろなところで使います。先ほど示したADD命令のコード生成は本当は型チェックのガードを入れないといけないのですが、サボっています。
コードの終わり
else
# サポートされていない命令
stop_compile
end
命令を実行するたびに対応するコードを生成すると説明しましたが、いつまで生成するか説明していませんでした。いつまでか？、こんな場合です
- コード生成をサポートしていない命令を実行した。SEND/RETURN命令はサポートしていません。理由はサポートするとめんどくさいからです。
- 使用頻度が低い命令を実行した
- その命令のコードがすでに生成されていた
こういう場合、ちょうどガードが失敗したときのようにFibVMに戻る処理のコードを生成します。それだけではなく、コンパイルで使った各種変数のリセットも行います。具体的なコードは以下の通りです。
def stop_compile
if @code.size > 1 then
@code += gen_exit(0)
@proc_tab[@irepid][@entry] = Irep.new_irep(@code, @pool, CodeGen::SYMS, 10, 2).to_proc
end
# Reset working
@entry = nil
@max_using_reg = 2
@code = []
@pool = []
end
Irep.new_irep(@code, @pool, CodeGen::SYMS, 10, 2).to_procで、生成したコード、定数テーブル等から、Procオブジェクトを生成してcallメソッドで呼び出せるようになります。
生成したコードの実行
# もし現在のPCの命令にコンパイルされたコードが存在したら
a = @proc_tab[@irepid][@pc]
if a then
if @entry then
# もし、コンパイル中ならコンパイルを中止する
stop_compile
end
# コード実行
@proc_tab[@irepid][@pc].call(self)
end
# コードを取り込む。コンパイルされたコードを実行すると@pcが変化する
# 可能性のあることに注意
cop = @irep.iseq[@pc]
コード生成、コード解釈に先だって、現在の命令がすでにコードが生成されているか調べます。さらに、現在コンパイル中かも調べ、もしコンパイル中ならコンパイラを即座に終了します。そして、生成されたコード(正確にはそれをProcオブジェクト化したもの)を実行します。実行のためのcall命令にはself(FibVMのインスタンス)の引数を1つ渡します。
コードを実行して戻ってくると@pcや@irepも次実行すべき命令にちゃんと設定されているのでそのまま命令をフェッチしてコード生成・コード解釈をおこなうことができます。
生成コードライブラリ
module CodeGen
include RiteOpcodeUtil
# Reg Mapping
# R0 SELF
# R1 SP
# R2 WORKING
# :
# @max_using_reg
SYMS = [:@stack, :@sp, :@pc, :+, :-, :[], :[]=, :p, :==]
STACK_SYM = 0
SP_SYM = 1
PC_SYM = 2
ADD_SYM = 3
SUB_SYM = 4
AREF_SYM = 5
ASET_SYM = 6
P_SYM = 7
EQ_SYM = 8
OPTABLE_CODE = Irep::OPTABLE_CODE
OPTABLE_SYM = Irep::OPTABLE_SYM
# FibVMのレジスタの内容をRITE VMのレジスタに格納する
def gen_get_reg(dst, src)
tmp0 = @max_using_reg
tmp1 = tmp0 + 1
tmp2 = tmp1 + 1
[
mkop_ABx(OPTABLE_CODE[:GETIV], tmp0, STACK_SYM),
mkop_AsBx(OPTABLE_CODE[:LOADI], tmp1, src),
mkop_ABx(OPTABLE_CODE[:GETIV], tmp2, SP_SYM),
mkop_ABC(OPTABLE_CODE[:ADD], tmp1, ADD_SYM, 1),
mkop_ABC(OPTABLE_CODE[:SEND], tmp0, AREF_SYM, 1),
mkop_AB(OPTABLE_CODE[:MOVE], dst, tmp0),
]
end
# RITE VMのレジスタの内容をFibVMのレジスタに格納する
def gen_set_reg(dst, val)
tmp0 = @max_using_reg
tmp1 = tmp0 + 1
tmp2 = tmp1 + 1
[
mkop_ABx(OPTABLE_CODE[:GETIV], tmp0, STACK_SYM),
mkop_AsBx(OPTABLE_CODE[:LOADI], tmp1, dst),
mkop_ABx(OPTABLE_CODE[:GETIV], tmp2, SP_SYM),
mkop_ABC(OPTABLE_CODE[:ADD], tmp1, ADD_SYM, 1),
mkop_AB(OPTABLE_CODE[:MOVE], tmp2, val),
mkop_ABC(OPTABLE_CODE[:SEND], tmp0, ASET_SYM, 2),
]
end
# RITE VMに戻るコードを生成する
def gen_exit(rreg)
code = []
tmp0 = @max_using_reg
@max_using_reg += 1
# @pcを設定するコード
code += [
mkop_AsBx(OPTABLE_CODE[:LOADI], tmp0, @pc),
mkop_ABx(OPTABLE_CODE[:SETIV], tmp0, PC_SYM),
]
# 戻り値の設定(戻り値はFibVMのレジスタに入っていることに注意)
code += gen_get_reg(tmp0, rreg)
# 戻る
code.push mkop_A(OPTABLE_CODE[:RETURN], tmp0)
@max_using_reg -= 1
code
end
end
コード生成ライブラリです。疲れたので興味のある方は読んでください。ブラックボックスでもJITコンパイラの理解には影響ないと思います。
結果
$ time ../../mruby/bin/mruby jit-1.rb
10946
10946
real 0m0.630s
user 0m0.592s
sys 0m0.016s
$ time ../../mruby/bin/mruby jit-3.rb
10946
10946
real 0m0.550s
user 0m0.483s
sys 0m0.046s
0.1秒ほど速くなってる。めでたしめでたし。
と思ったら、これはmrubyのJITの結果で、オリジナルのmrubyではこの通り
$ time bin/mruby ../../work/mruby-meta-circular/sample/jit-1.rb
10946
10946
real 0m0.890s
user 0m0.842s
sys 0m0.030s
$ time bin/mruby ../../work/mruby-meta-circular/sample/jit-3.rb
10946
10946
real 0m0.940s
user 0m0.919s
sys 0m0.015s
JIT版の方が遅い…
動かしたい、という酔狂な方はこちらをご覧ください
JITなし https://github.com/miura1729/mruby-meta-circular/blob/master/sample/jit-1.rb
JITあり https://github.com/miura1729/mruby-meta-circular/blob/master/sample/jit-1.rb
ただし、mruby-meta-circularが組み込まれていないと動きません
おまけ(練習問題)
実はここで説明したJITコンパイラは少なくとも１つバグがあります。それはなんでしょう？もちろん、フィボナッチはそのバグを踏みません。
正解者はもれなくいち早く答えを知ることができます(つまり賞品の類はないです)。答えはクリスマス明けにでも書きます。
もっと余裕のある方は、SEND/RETURN命令もコードを生成できるようにしてみると面白いと思います。これの正解は書く予定はありません。