CRuby向けのLLVMベースのJITコンパイラを書いている話

LLRBというRuby向けのメソッドJITコンパイラを書いている

github.com

RubyKaigi 2015の最後のキーノート@evanphxが「LLVMでCRubyのコードをインライン化するメソッドJITを実装したら速いんじゃね」みたいな発表をしていたのを覚えているだろうか。

LLRBというのはまさにそれを実装しているプロジェクトであり、少なくとも現時点で「LLVMでCRubyのコードをインライン化するメソッドJIT」と言える状態まで実装でき、ものによっては効果が出る状態になったので公開した。

なんで書いてるの

言語を自分で実装するとその言語に関する理解が大分深まる、というのをHamlの実装とかCコンパイラとかで体験していて、僕が一番好きな言語はRubyなのでRubyでもそれをやっておきたい、というのがあった。また、Rubyは遅いと言われがちだが、どこに改善可能な点が眠っているのかにも興味がある。

また、ささださんがRuby3のゴールの一つにJust-in-Time compilationをあげていて、そのスライドで1つの選択肢として言われていた"Use LLVM"というパートを検証するものでもある。ただし、MatzがLLVMはDependency的に微妙と言っており、本体にこの方向性で入ることは多分なく、あくまで実験的なプロジェクトになる。

どのように動くか

READMEを書いてたら大分疲れたので、この記事はその日本語訳にしておく。

LLRBのビルドプロセスでは、CRubyの命令やそれに使われるいくつかの関数をLLVM bitcodeという形式にプリコンパイルしている。

 ________     _________     ______________
|        |   |         |   |              |
| CRuby  |   | CRuby   |   | CRuby        |
| C code |-->| LLVM IR |-->| LLVM bitcode |
|________|   |_________|   |______________|

LLRBのJITのためのプロファイラが開始されると、Rubyの実行中にJITコンパイルがスケジューリングされる。 JITコンパイルの際、LLRBはYARV ISeqをLLVM IRにコンパイルするが、その際にプリコンパイルしておいたLLVM bitcodeをリンクするので、LLVM Passによる関数のインライン化やその後の最適化が行なえる。

最適化を行なった後、LLVM IRはLLVMのMCJITエンジンによって機械語に変換され、それを呼び出すためのCの関数ポインタとしてLLRBに返される。コンパイル対象のISeqに対し、後に説明する方法によって変化が加えられ、この関数ポインタが実行されるようになり最適化が完了する。

 ______     ______________      __________      ___________      _________
|      |   |              |    |          |    |           |    |         |
| YARV |   | ISeq Control |    | LLVM IR  |    | Optimized |    | Machine |
| ISeq |-->| Flow Graph   |-*->| for ISeq |-*->| LLVM IR   |-*->| code    |
|______|   |______________| |  |__________| |  |___________| |  |_________|
                            |               |                |
                            | Link          | Optimize       | JIT compile
                      ______|_______     ___|____          __|____
                     |              |   |        |        |       |
                     | CRuby        |   | LLVM   |        | LLVM  |
                     | LLVM Bitcode |   | Passes |        | MCJIT |
                     |______________|   |________|        |_______|

実際にパフォーマンスが向上するか?

基本的なインライン化は実装済なので、その効果を見ていく。

以下のようなruby/benchmark/bm_loop_whileloop.rbのコードについて考える。

def while_loop
  i = 0
  while i<30_000_000
    i += 1
  end
end

このメソッドのYARV ISeq (LLRBのコンパイル対象) はこうなっている。

> puts RubyVM::InstructionSequence.of(method(:while_loop)).disasm
== disasm: #<ISeq:while_loop@(pry)>=====================================
== catch table
| catch type: break  st: 0015 ed: 0035 sp: 0000 cont: 0035
| catch type: next   st: 0015 ed: 0035 sp: 0000 cont: 0012
| catch type: redo   st: 0015 ed: 0035 sp: 0000 cont: 0015
|------------------------------------------------------------------------
local table (size: 1, argc: 0 [opts: 0, rest: -1, post: 0, block: -1, kw: -1@-1, kwrest: -1])
[ 1] i
0000 trace            8                                               (   1)
0002 trace            1                                               (   2)
0004 putobject_OP_INT2FIX_O_0_C_
0005 setlocal_OP__WC__0 3
0007 trace            1                                               (   3)
0009 jump             25
0011 putnil
0012 pop
0013 jump             25
0015 trace            1                                               (   4)
0017 getlocal_OP__WC__0 3
0019 putobject_OP_INT2FIX_O_1_C_
0020 opt_plus         <callinfo!mid:+, argc:1, ARGS_SIMPLE>, <callcache>
0023 setlocal_OP__WC__0 3
0025 getlocal_OP__WC__0 3                                             (   3)
0027 putobject        30000000
0029 opt_lt           <callinfo!mid:<, argc:1, ARGS_SIMPLE>, <callcache>
0032 branchif         15
0034 putnil
0035 trace            16                                              (   6)
0037 leave                                                            (   4)
=> nil

LLRBのコンパイラは、これを以下のようなLLVM IRにコンパイルする。

define i64 @llrb_exec(i64, i64) {
label_0:
  call void @llrb_insn_trace(i64 %0, i64 %1, i32 8, i64 52)
  call void @llrb_insn_trace(i64 %0, i64 %1, i32 1, i64 52)
  call void @llrb_insn_setlocal_level0(i64 %1, i64 3, i64 1)
  call void @llrb_insn_trace(i64 %0, i64 %1, i32 1, i64 52)
  br label %label_25

label_15:                                         ; preds = %label_25
  call void @llrb_insn_trace(i64 %0, i64 %1, i32 1, i64 52)
  %2 = call i64 @llrb_insn_getlocal_level0(i64 %1, i64 3)
  call void @llrb_set_pc(i64 %1, i64 94225474387824)
  %opt_plus = call i64 @llrb_insn_opt_plus(i64 %2, i64 3)
  call void @llrb_insn_setlocal_level0(i64 %1, i64 3, i64 %opt_plus)
  br label %label_25

label_25:                                         ; preds = %label_15, %label_0
  %3 = call i64 @llrb_insn_getlocal_level0(i64 %1, i64 3)
  call void @llrb_set_pc(i64 %1, i64 94225474387896)
  %opt_lt = call i64 @llrb_insn_opt_lt(i64 %3, i64 60000001)
  %RTEST_mask = and i64 %opt_lt, -9
  %RTEST = icmp ne i64 %RTEST_mask, 0
  br i1 %RTEST, label %label_15, label %label_34

label_34:                                         ; preds = %label_25
  call void @llrb_insn_trace(i64 %0, i64 %1, i32 16, i64 8)
  call void @llrb_set_pc(i64 %1, i64 94225474387960)
  call void @llrb_push_result(i64 %1, i64 8)
  ret i64 %1
}

上記のLLVM IRでcallされている関数は全てLLVM bitcodeにプリコンパイルされており、LLRBはJITコンパイル時にそれをリンクする。そのため以下のようなインライン化と最適化がLLVMによって行なわれる。

define i64 @llrb_exec(i64, i64) #0 {
  ...

land.lhs.true.i:                                  ; preds = %label_25
  %49 = load %struct.rb_vm_struct*, %struct.rb_vm_struct** @ruby_current_vm, align 8, !dbg !3471, !tbaa !3472
  %arrayidx.i = getelementptr inbounds %struct.rb_vm_struct, %struct.rb_vm_struct* %49, i64 0, i32 39, i64 7, !dbg !3471
  %50 = load i16, i16* %arrayidx.i, align 2, !dbg !3471, !tbaa !3473
  %and2.i = and i16 %50, 1, !dbg !3471
  %tobool6.i = icmp eq i16 %and2.i, 0, !dbg !3471
  br i1 %tobool6.i, label %if.then.i, label %if.else11.i, !dbg !3475, !prof !3380

if.then.i:                                        ; preds = %land.lhs.true.i
  call void @llvm.dbg.value(metadata i64 %48, i64 0, metadata !2680, metadata !3361) #7, !dbg !3477
  call void @llvm.dbg.value(metadata i64 60000001, i64 0, metadata !2683, metadata !3361) #7, !dbg !3478
  %cmp7.i = icmp slt i64 %48, 60000001, !dbg !3479
  %..i = select i1 %cmp7.i, i64 20, i64 0, !dbg !3481
  br label %llrb_insn_opt_lt.exit

if.else11.i:                                      ; preds = %land.lhs.true.i, %label_25
  %call35.i = call i64 (i64, i64, i32, ...) @rb_funcall(i64 %48, i64 60, i32 1, i64 60000001) #7, !dbg !3483
  br label %llrb_insn_opt_lt.exit, !dbg !3486

llrb_insn_opt_lt.exit:                            ; preds = %if.then.i, %if.else11.i
  %retval.1.i = phi i64 [ %..i, %if.then.i ], [ %call35.i, %if.else11.i ]
  %RTEST_mask = and i64 %retval.1.i, -9
  %RTEST = icmp eq i64 %RTEST_mask, 0

  ...
}

いろいろインライン化されている上に僕の書いたCSSがクソなので読みづらいが、簡単に説明するとRubyVMの状態を取得し、<が再定義されているかどうかをチェックし、再定義されていなければicmp sltという命令が実行されるようになっている。インライン化されているので、llrb_insn_opt_ltという本来の関数の呼び出しのオーバーヘッドもなくなっている。 これはYARVのinsns.defの定義をほとんどそのまま持ってきてインライン化しているだけなので、実装コストが低いというメリットがある。

この最適化で、以下のようなベンチマークで、

ruby = Class.new
def ruby.script
  i = 0
  while i< 30_000_000
    i += 1
  end
end

llrb = Class.new
def llrb.script
  i = 0
  while i< 30_000_000
    i += 1
  end
end

LLRB::JIT.compile(llrb, :script)

Benchmark.ips do |x|
  x.report('Ruby') { ruby.script }
  x.report('LLRB') { llrb.script }
  x.compare!
end

以下のようにパフォーマンスが改善することがわかる。

# x86_64 GNU/Linux, Intel(R) Core(TM) i7-4790K CPU @ 4.00GHz
Calculating -------------------------------------
                Ruby      2.449  (± 0.0%) i/s -     13.000  in   5.308125s
                LLRB      8.533  (± 0.0%) i/s -     43.000  in   5.040016s

Comparison:
                LLRB:        8.5 i/s
                Ruby:        2.4 i/s - 3.48x  slower

LLRBの設計はどうなっているか

C拡張として作成

どうやらYARVは最初C拡張として作られていたらしい。JITコンパイラをC拡張として作るのは大分無理があると思ってたけど、YARVがC拡張でできるならJITコンパイラもC拡張でできるのではないかと思ったのでそうしてみた。実際、CRubyのforkとして開発をするのに比べると、コアと疎結合になるだけでなく、bundlerとかbenchmark-ipsとかその辺のgemをカジュアルに使えるので開発はやりやすい。

しかし実際のところ、基本的な演算子のメソッドが再定義されてるかどうかを持つ変数のシンボルが当然exportされてないことに気付き、そこでCRubyに変更を加えず開発することは諦めた。ので、k0kubun/rubyのllrbブランチにしかインストールできない。それでも、シンボルのexport以外には何もしないという縛りを設けてあるので、今後CRubyのバージョンが上がっても追従はしやすい気がする。

YARVに手を加えない保守的なアプローチ

YARVは大分長い間運用されている信頼性のあるVMなので、仮にJITコンパイラを導入したとしても移行期間中はベースになるYARVがそのまま使える状態の方が安心できると思っている。なので、LLRBはYARV ISeqからコンパイルする方式を取り、またYARVのコア自体には一切変更を加えないようになっている。

じゃあそれでどうやってJITを達成しているかというと、YARVにはopt_call_c_functionという命令があり、「ネイティブコンパイルしたメソッドを起動。」という説明が書いてある。これをJITに使わない理由はない。

というわけで、LLRBは前述した方法でCの関数ポインタを取得してそれをopt_call_c_functionオペランドとし、全てのメソッドを以下のようなISeqに変換する。

0000 opt_call_c_function
0002 leave

opt_call_c_functionYARVの内部例外を扱える *1ので、この命令を使えば大抵のことは実装できるし、少なくとも僕がテストを書いた範囲ではthrowとかbreakとかreturnが動いている。

しかし、ISeqを書き変える方法だと考慮しないといけないことがある。それは、YARVが内部例外を補足するcatch tableがプログラムカウンターの位置に依存して内部例外をハンドルすることである。 プログラムカウンターが0か2にいればすむようにcatch table自体を書き変えると分岐できなくなるケースがありそうだったので、LLRBはcatch tableには手を加えず、逆にネイティブコードの中でプログラムカウンターを従来通り動かす方を選択した。

なので、プログラムカウンターがそのような位置にあっても適切にleaveできるよう、余った位置にleave命令埋めをしている。これはshyouheiさんのDeoptimization Engineのnop埋めのパクリである。

正直それで問題が起きないのかはよくわかってないので他の人のツッコミ待ちだが、僕がテストを書いた範囲ではrescueとかensureが動いている。

サンプリングベースの軽量なプロファイラ

LLRBはプロファイラとしてサンプリングプロファイラを採用している。CPU時間で一定のインターバルおきにバックトレースのサンプリングを行ない、バックトレースの一番上に出現する頻度が高いものをコンパイル対象として選択する。

stackprofとかperfもサンプリングプロファイラであり、この手法自体は既に広く使われているため信頼性がある。

また、stackprofと同じくrb_postponed_job_register_one APIを使う + GC実行中かどうかを判定していて、おそらくコンパイルしても安全と思われる瞬間にコンパイルが行なわれる。

低いコンパイルコスト

上述したように、CRubyのC関数はビルドプロセスでLLVM bitcodeにプリコンパイルされているので、LLRBがJITコンパイルを実行する際ランタイムでCの関数のパースやコンパイルを行なうオーバーヘッドは発生しない。 また、RubyのASTからコンパイルするわけではなく、ある程度コンパイルされたYARV ISeqからコンパイルを開始しているので、おそらくこれが一番コンパイルコストが低い手法になると思われる。

TODOs

  • まだ全ての命令に対して最適化を実施できていないので、実際のアプリケーションで使うとパフォーマンスが劣化する。なので、まずはその対応から…
  • ISeqのインライン化に対応したい
  • まだコンパイルに対応してないYARV命令がいくつかある
    • expandarray, reverse, reput, defineclass, once, opt_call_c_function
  • コンパイル中に作るRubyのオブジェクトのうちGCを考慮できてない奴がある気がする

ビルド・使用方法

このへんを参照

気持ち

まだ大分未完成なんだけど、3か月くらい経ちモチベーションの維持が難しくなってきたので公開した。この辺で得られた知見をRubyKaigiとかで話せたらいいなと思っていたけど、まだ僕のCFPはacceptされていないので、まあ無理だったらRubyConfで会いましょう。

*1:YARV内部の例外についてはRuby under a microscopeとかに詳しく書いてあるのでここでは説明しない