これは言語実装 Advent Calendar 2018の25日目の記事です。
今日2018年12月25日は、Rubyに初めてJITコンパイラが搭載されたRuby 2.6のリリース日です。めでたい!!
進捗はRuby 2.6 JIT - Progress and Futureに書いてある通りですが、ハイライトとしては、CPU計算負荷の高いOptcarrotというNESエミューレータのベンチマークで、以下のような性能向上を達成しました。
Optcarrot Lan_Master.nes
2.6.0+JIT: 86.6 fps
2.6.0: 54.6 fps - 1.59x slower
2.5.3: 48.5 fps - 1.78x slower
2.0.0: 34.6 fps - 2.50x slower
Ruby 2.6では、JITを有効にしなくても2.5に比較してそこそこ速くなっていますが、その上でJITを有効すると更に1.59倍速くなる状態になりました。僕が究極的に目的としているRailsアプリケーションの最適化には至ってないのは微妙なところですが、単なるマイクロベンチマークではないアプリケーションを約700行のコード最適化ロジックで1.59倍速くするまでには、様々な工夫が必要でした。現在は単にVMの命令を切り貼りしてコンパイルしているだけ ではありません 。
$ find . -name \*.erb -o -name \*.c |grep mjit_compile|xargs wc -l
254 ./mjit_compile.c
77 ./tool/ruby_vm/views/mjit_compile.inc.erb
54 ./tool/ruby_vm/views/_mjit_compile_ivar.erb
94 ./tool/ruby_vm/views/_mjit_compile_send.erb
36 ./tool/ruby_vm/views/_mjit_compile_pc_and_sp.erb
86 ./tool/ruby_vm/views/_mjit_compile_insn.erb
105 ./tool/ruby_vm/views/_mjit_compile_insn_body.erb
706 total
この記事では、Ruby 2.6で検討・検証・実装を行なった最適化のアイデアをまとめておきます。
前提として、Ruby 2.6のJITはバイトコードをメソッド単位でCのコードに変換し、その際Cのコードレベルでの最適化を行ない、Cコンパイラにコンパイル・コード生成をさせることでVM実行の時よりも良いコードを生成することを期待するというもので、今のところインラインアセンブラを使っていないので、その制約下でのみ行なえる最適化しか出てきません。
RubyElixirConf Taiwanで話した内容に近いので、その時の資料がある奴に関しては併記します。というか、実質その発表の日本語版 + preview3のアップデートを含めたもの、という感じになっています。
気付いたらこれ以降突然丁寧語じゃなくなってましたが、大目に見てください。
実装済のもの
メソッドディスパッチ処理のバイパス
以前Rubyのメソッドディスパッチの挙動をまとめた。メソッドキャッシュがヒットした時、以下のような挙動をする。
- レシーバのクラスがキャッシュキーのクラスと同じ
- キャッシュに入っている、あるメソッドのタイプ専用の関数ポインタを呼ぶ
関数ポインタの呼び出しをはさむと、関数ポインタを読んでくるオーバーヘッドと、その関数ポインタを越えてコンパイラが最適化ができないという問題がある。生成コードの中で、1の通りにキャッシュをチェックし、ヒットしたら2の関数を直接呼び出し、ヒットしない場合はVM実行にフォールバックするようにすると速くなる。キャンセルしないパスが2つあるとコンパイラの最適化が効きにくくなるため。
最適化されるメソッドの余計な型チェックの排除
Rubyは +, -, *, / や ==, [] といった基本的なオペレータの呼び出しの時、レシーバがInteger, String, Array といったビルトインのクラスである時に有利なVM命令を決め打ちでコンパイルする。そのため、例えば+を呼び出す時レシーバがIntegerやFloat、StringやArrayかといったチェックをした後、そうでなければ普通のメソッド呼び出し命令の処理をすることになっている。
バイトコード中のある命令でレシーバのクラスが何になっているかは一度VM実行するとキャッシュキーとして入るようになっているので、それがビルトインのクラスでなかった場合は、余計な型チェックをせず普通にメソッド呼び出しを実行するようにした。
VM命令の実装内でよく実行される関数をインライン化可能な形で呼び出す
基本的にCのコードは別オブジェクトファイルにある関数をプロトタイプ宣言して呼び出すよりは、同じCのファイルで定義がわかるようにしてコンパイルしてあげると速くなるという普通の話。命令を切り貼りしている都合JITの方がよく効くと思ったが、たとえばArrayの #[] 呼び出しでこれをやったところVMの性能も普通に上がったので、あまりJITに関係ない感じになった。
stack pointerの移動のスキップ
スタックベースのVMの場合、VMのスタックに値をpush/popする際はstack pointerをインクリメントしたりデクリメントしたりする必要があるけれど、ある命令がスタックのどの位置に値をつっこむ予定なのかは(少なくともRubyの現時点のVMでは)決定的なので、もし普通に配列のローカル変数を用意してその決まったインデックスに値をつっこむようなコードを生成すれば、実行時にspを動かす必要がなくなる。ただしメソッド呼び出しをする際は引数をVMのスタックに書き戻す必要がある。
しかしこれをやるのは言うほど簡単ではなくて、Rubyの例外はlongjmp / setjmpで実現されているのだが、通常VMが起動した時にsetjmpをした後、JITの関数が呼び出され、その中でlongjmpが行なわれた場合、VMを起動した時のスタックに戻るのでJITの関数のフレームにあったローカル変数はスタックから飛んでしまう。しかし、そのフレームで例外をrescueして再び処理が走ることがあるので、JIT関数で起動していたフレームは途中から突然VM実行に切り替えができるようにしておく必要がある。
これには3通りの対応が考えられる:
- longjmpする時にcallerを辿り、JIT関数のフレームについたらローカル変数だった値をVMのスタックに書き戻す
- 例外をcatchしうるフレームではこの最適化を諦め、VMのスタックを使う
- スタックが飛ばないことが保証できなくなるメソッド呼び出しの前に毎回setjmpをして、その後のlongjmpが、必要なスタックが飛ぶ直前で止まるようにする
1は実装が果てし無く面倒くさくて、簡単な実装はしてみたけど何かバグってたりして面倒になった。今は2と3が実装されているのだけど、片方でいいような気がする…。理由がすぐ思い出せないので気が向いたらまた後で考えることにする、申し訳ない。
VMからJIT関数を呼び出す時にsetjmpをしない
一つ前の最適化について、何も考えずに3のガードをやってるとsetjmpを差し込みまくることになるのだけど、これは遅いので、すでにsetjmpを呼び出してるはずのVMのフレームから直接JIT関数を呼び出す時(callerはVMなのでもともとVMのスタックを使っている)は、setjmpをしないという話。なおRubyの実装レベルではvm_execの呼び出しがその中でsetjmpの呼び出していることを意味している。
program counterの移動のスキップ
VMはprogram counterを移動して、それが指す命令を読んで各命令をディスパッチしている。が、JITで各命令の処理を切り貼りしていれば、普通に考えたらprogram counterの操作のためのメモリアクセスは余計ということになる。
しかしこれも簡単な話ではない。まずRubyがRubyレベルでバックトレースを作るようなメソッドを呼び出した時、これはprogram counterに依存して行番号を計算するので、callerでは常にprogram counterが正しい位置にないといけない。また、各メソッドにbegin rescueの区間に応じて例外処理をディスパッチするためのcatch tableというものがあり、これもprogram counterの位置に応じて処理が決まるので、calleeで例外が発生した時にprogram counterが正しい位置にいる必要がある。
修正方法のうち検討したのは2通り:
- callerのprogram counterが必要になった瞬間、マシンスタックやnativeのinstruction pointerからprogram counterを逆算する
- メソッド呼び出しや例外を上げる直前のみprogram counterを動かすようにする
1はまあすごい面倒な奥義みたいな感じで、GraalのJITを使うTruffleRubyはこれをやってると聞いたが、実装の労力がかかりすぎると思う。とりあえず2をやっているが、これも多少は面倒で、そもそもRubyのメソッド呼び出しというのはCの関数で行なえるので、任意の(メソッド呼び出しでない)VM命令が突然メソッド呼び出しをやっている可能性がある。なので、Ruby 2.6からは各命令に対しメソッド呼び出しが発生しうるかというメタデータを用意し、それが正しいかどうかassertするようなCIを走らせ、catch tableがないメソッド内かつ例外が発生しえない命令でのみprogram counterの移動をスキップするようにしている。
インライン化で最適化されそうなVM命令に __always_inline__ をつけまくる
https://speakerdeck.com/k0kubun/method-jit-compiler-for-mri?slide=87
https://speakerdeck.com/k0kubun/method-jit-compiler-for-mri?slide=88
コンパイラは関数のサイズが閾値を越えるとインライン化を諦めるようになっている(LLVMの実装しか読んだことないけど、まあこちらの実装をいじって計測してみた限りではGCCとかも多分そうだと思っている)。VM命令の実装は大きくなりがちだが、命令を切り貼りして合わせたことで、例えば + みたいな最適化される命令のオペランドをインライン化して渡せると、コンパイル時に計算が完了できるようになって速くなる。なので、そういう最適化がしやすそうな場所で強制的にインライン化するようにすると速くなることがある。
なお、VMのコードサイズが大きくなるので、今回はJITに渡すヘッダファイルでのみ __always_inline__ にするような実装にした。どの道VMではオペランドの値のインライン化ができないので意味もないと思っている。
JIT内ではstack pointerの検証を無効にする
stack pointerはVM命令のメタデータに応じて基本的にはVMのメインの関数が動かすことになっているのだが、一方で(メソッド呼び出しなどの複雑なスタック処理を含む命令向けに)各命令の実装内でも普通にstack pointerが動かせるようになっているので、VMがメソッド呼び出しを抜ける時にstack pointerが変な位置にないかをチェックするようになっているが、これが遅い。
JITではそもそもstack pointerに依存してなかったりするので、この重いスキップを無効にして速くしている。JITが使うスタックのインデックスはJITコンパイル時にassertされている。
生成コードのコンパクション
GC以外であんまりコンパクションって言わない気がするけど、内部で"JIT Compaction"と呼んでいるテクニックを実装している。RubyのJITはdlopenでコードをロードしているのだけど、この実装の仕様上1つのsoをロードする度に各2MBくらいのスキマができてしまうという問題があり、初期の実装ではメソッドとsoが一対一対応していたので各メソッドのJIT生成コード間の距離が遠いという問題があった。これは生成コードのロードにキャッシュが効きにくくなることを意味する。
そのため、全く速くなる余地も遅くなる余地もないようなメソッドを複数JITしまくるだけでどんどん遅くなるという問題が発覚した。これをRubyKaigiで発表したところ、shinhさんがELFオブジェクトを直接ロードするローダを作ってくださったが、ポータビリティや、その実装特有のトレードオフがあり、これの採用は見送った。
最終的には、メソッドが全てJITされたっぽいタイミングで全てのJITされたメソッドのオブジェクトファイルを一つのsoファイルにまとめてdlopenし直すという解決に帰着した。RubyのJITに生成コードのメモリ局所性対策を入れた話という記事に詳細を書いた。
そこまで丁寧に計測していないものの、shinhさんが実装したobjfcnにそこそこ近い性能向上が得られた。が、メソッドの数が増えた時のオーバーヘッドを無にできてはいないので、ローダ部分にまだ改善の余地はあるかもしれない。今のところ、後述するインライン化を駆使してコンパイル単位を変えてコードの読み出し回数自体を減らすことでどうにかするつもりでいる。
なおこれの現在の実装にはやや問題があり、全部のメソッドをつっこんだsoから生成したコードがコールスタックのどこにもない(生成コードを開放できる)瞬間はあまりないだろうという仮定のもと、JIT Compactionによってロードしてきたコードのdlcloseをインタプリタ終了まで完全にサボっている。なので、「メソッドが全てJITされたっぽいタイミング」が何度も来続けるようなアプリケーションを長時間走らせると、メモリリークのような挙動になってOOMに殺される状態になっている。辛い…。まあ、一応JITの現在のスケジューリングの構造上無限にリークし続けるわけではないはずなのでそれほどシビアではないと思っているけど、2.7では何らか対策をいれるつもり。
インスタンス変数アクセスのインライン化
インスタンス変数にアクセスする命令にもキャッシュキーとキャッシュ対象の値がある。具体的には、あるオブジェクトのインスタンス変数領域(ヒープかオブジェクト表現内埋め込み)のどのインデックスに値があるかを、クラスをキャッシュキーにしてキャッシュしている。
で、キャッシュヒットした場合の分岐は相当単純で、それ以外のコードが大きすぎることにより普通にコード生成を行なうとコンパイラがうまく最適化できないので、JITを行なう瞬間のインデックスを生成コードにインライン化もしつつ、キャッシュヒットしない分岐はVMにフォールバックするようなコードを生成すると、インスタンス変数のアクセスに無駄がなくなりかなり速くなる。
これがpreview3におけるOptcarrotの性能のブレークスルーで、Optcarrotは一番のボトルネックがインスタンス変数アクセスなので、主にこれが理由で良いベンチマーク結果を出せている。
検証はしてみたが未採用のもの
attr_reader のメソッド呼び出しのインライン化
attr_reader という、オブジェクトのインスタンス変数のreaderメソッドがあるのだが、これはRubyで書かれたメソッドやC拡張で定義されたメソッドとは違うタイプのメソッドになる。メソッド呼び出し処理はバイパスし、インスタンス変数を取り出すVMの処理をインライン化するのが最速になるが、これは一度trunkにコミットしたもののなんかバグが発覚したので諦めてrevertした状態になっている。理由はまだよくわかっていないが、Optcarrotには効かない最適化なので放置したままになっていた(!?)。Railsには効くかも。
最適化命令の型ごとの細分化と、動的な命令書き換え
現在のRubyのVMにおける動的な命令書き換えはたった一つで、VMの動作の様々な挙動をフックして突然任意の処理を差し込むためのTracePointという機能が有効化された時に全てのバイトコードをtraceされる命令に書き換えるというものだけがある。これにより、TracePointが有効になるまでは、traceが必要かどうかのチェックを完全にスキップするという最適化が実現されている。
この最適化は、例えば現在は + を最適化する opt_plus という命令がレシーバの型をInteger, Float, String, Arrayで順番にチェックしてディスパッチするのを、 opt_plus_Integer, opt_plus_Float, ... といった具合に細分化しておき、レシーバの型が変わったタイミングでバイトコードの命令も差し替えてしまうというもの。
これをやると、JITコンパイラはある1つの型に特化した分岐だけを生成できるようになり、その方がコンパイラが最適化しやすいだろう、というもの。
実装はここ https://github.com/k0kubun/ruby/pull/11 にあって、ちゃんと動くのだが、Optcarrotは速くならないのでやめた。多分だけど、最初にチェックしてるIntegerの型が一番使われてるからかなあと思う。
JITの生成コード読み込みのオーバーヘッド削減
現在のRubyのJITの問題点として、JITして生成した関数を呼び出す際にいくつかのメモリアクセスが必要でオーバーヘッドが大きいという話がある。実は今RailsアプリケーションでJITをすると大体のケースで遅くなってしまうのだが、こういう遅くなる問題の温床になる。
JITの生成コードは 「現在のスレッド」→「現在のコールスタックのフレーム→」「バイトコード」→「そのバイトコード用のJIT生成コード」という感じで読んでくるのだが、他の方法でメソッド呼び出しをする時は、メソッドキャッシュを持つstructの要素の一つに関数ポインタをセットしておき、メソッドキャッシュのinvalidate直後にその関数ポインタを呼ぶという処理になっているので、JITが生成する関数をこの場所にセットできるようなものにしてしまえば、独自の方法でコードを読んでくる必要がなくなると考えている。
何故そうしていないかというと今JITで生成している関数はそこで呼び出しても全く動かないからで、まずフォーマットを揃えてあげる必要がある。で、それをやってるのがこれ https://github.com/k0kubun/ruby/pull/8 なんだけど、これをやるとVMが遅くなってしまうので、まだ入っていない。VMが遅くなる理由もどうにかつきとめて直すなどのことが必要になる。
Rubyで定義されたメソッドの1段階インライン化
「メソッドディスパッチ処理のバイパス」ができていると前述したが、これができるということはメソッドの定義もその場にインライン化する準備ができているということで、まあやればできるレベルの話。もう大分古いコミットがベースになってしまっているが実装はここにある https://github.com/k0kubun/ruby/tree/mjit-inline-send-yield 。RubyKaigiでこれのデモをして https://speakerdeck.com/k0kubun/the-method-jit-compiler?slide=150 、"C language is dead"といった訳のわからないことを口走ったりした。
入れてないのは、コンパイル時間が伸びるのでOptcarrotのベンチマーク中コンパイルが間に合うメソッドの数が減りベンチマークスコアが落ちてしまうからで、(Optcarrotのベンチマークのためにも)ある程度インライン化するかどうかの取捨選択はいれたいと思っているものの、マージへの障壁はそれほどない。
ただし、現在の実装ではJITの対象にしたメソッドから直接呼び出されているメソッドしかインライン化していない。要するに再帰で実装されていないのである。
構想レベルのもの
複数段階のメソッドインライン化
直前に書いた1段階インライン化を拡張したもの。やればできるような気がするが、全てをインライン化するとコンパイル速度が大きく犠牲になるので、何をどこまでインライン化するかの戦略を慎重に考えないといけない。面倒なのでJavaがやってることをパクりたいのだけど、OpenJDKの実装に詳しくないので、関連した論文なり内部実装なりに詳しい人がいたら連絡をいただきたい。何もなければ自分で読むなりして調査するつもり。
あと、よく呼び出されているメソッドがJITの対象になった時、そのメソッド呼び出し自体がインライン化されることによる恩恵は大きいはずなので、JIT対象にしたメソッドが呼び出しているメソッドのインライン化だけでなく、JIT対象にしたメソッドを呼び出しているメソッドからJITした方が速そう、という話もある。が、そういう箇所は複数存在し得るはずで、そのうちどれをコンパイルするか、があまり自明ではなくなるのが難しいポイント。全部コンパイルしたらコンパイルに時間がかかるし、少ないと使われないかもしれないし。
longjmpを契機にしたdeoptimization
メソッドキャッシュのキャッシュミスなど、最適化の条件が無効になった場合JIT関数はその場でspやpcを復元した後returnし、VM実行にフォールバックするようになっている。これは最適化の条件が満たされている場合にチェックのコストがかかってしまうが、一応このJITにおけるDeoptimizationにあたる。
普通は、Deoptimizationというのはそういうチェックは最適化対象のコードではなくしておいて、別のメソッド呼び出し中に最適化の条件が無効になった場合に、どうにかそこで過去に最適化していたコードの実行をキャンセルするようになっている。今のRubyのJITはそうなってないが、呼び出し済のJIT関数は、program counterやstack pointerを復元可能な状態にした上で、その上のVMフレームとかまでlongjmpしてしまえば、その関数を呼び出していなかったかのような状態に無理矢理することができるので、いわゆる普通のDeoptimizationが可能になる。On-Stack Replacementの一種になる。
これをやるためには、メソッドの再定義などのタイミングでprogram counterやstack pointerを全てのJITフレームで復元可能にしつつ、またsetjmpが呼ばれている回数だけ何度もlongjmpが必要なことなどの対応が面倒で、まだやっていない。
あと、普通に分岐する変わりに、後からSEGVさせることが可能なtest命令をインラインアセンブラで仕込んでおいて、DeoptimizationしたいタイミングでSEGVさせてSEGVのシグナルハンドラで全てをどうにかするという話もChris Seatonが提案していて、TruffleRubyというGraalでJITするTruffleベースのRuby実装ではこれが行なわれているらしいのだが、なんというかCRubyのためだけにこの機構を実装してメンテするのはあまり割に合わない感じがするし、それより先に最適化すべき問題がいっぱいあると思っている。
Multi-tier JIT
HotSpot(JavaのJIT)にはC1(client)コンパイラとC2(server)コンパイラというのがあって、コンパイル時間は短いがそこまで最適化されないC1コンパイラと、コンパイル時間が長いが最大限最適化を行なうC2コンパイラを使い分けている。
今のRubyのJIT方法はややコンパイル時間が長い(100〜500msくらい)のでWebアプリケーションのような長く走らせるプログラムに向いているが、数秒で終わるような処理にはコンパイル時間が短いJITができると良いかもしれない、という話がある。
これは現在のRubyのJITの仕組みを提案・実装した人が取り組んでいるので、僕は特に何かをやるつもりがないが、その後の進捗次第では入るかもしれない。
Profile-guided optimization
GCCとかは -fprofile-generate を渡して生成したプロファイリング用バイナリを実行した後、その実行で生成されたプロファイリング情報も入力にして再度コンパイルすると、そのプロファイリング結果に基いた最適化を行なう、という機能がある。実際具体的に何が行なわれるのか詳しいわけではないが、よく実行されるコードを近くに集めて命令の読み込みがキャッシュされやすくしたり、よく使われるパスへの分岐コストが低くなるようなコード生成を行なっていると想像している。FacebookがHHVMでやってる奴も似たような話だと思う。
これをやろうとすると、Multi-tier JITと同じく、どのタイミングでどのメソッドをプロファイリング開始し、最終的なコンパイルを行なうかといった戦略を決めないといけないのが面倒でまだ始めていないが、広く効果が期待できそうな最適化なのでRuby 2.7で入れたいなあと思っている。
オブジェクトのスタック割付
mallocやfreeを普通に呼ぶとヒープの適当な場所を割り当て確保するためのオーバーヘッドがかかるのに対し、オブジェクトのためのメモリをスタックに確保すると、マシンのスタックポインタを移動させるだけでメモリの確保ができる上、同時に使用されやすい値と近い場所に置かれるのでキャッシュも効きやすいという話がある。
なので、メソッドの外側に返されることがないようなオブジェクトはJIT関数のフレームのスタックに割り当てられると速そうなわけだが、それをやるためにはまさにその「メソッドの外側に返されることがない」ことを解析しなければならない。これはエスケープ解析と呼ばれている。これの実装はちゃんとやろうとするとまあ大変なのとOptcarrotがベンチマーク中オブジェクトを作らないので効かなそうなこと等からまだやってないのだが、そこそこ効果が期待できる気がするので2.7開発のどこかのタイミングで検証くらいはしたいと思っている。