Ruby 2.6にJITコンパイラをマージしました

The English version of this article is available here: medium.com

2/4(日)に、去年のRubyKaigiが終わった直後の新幹線で開発を始め10月に公開したJITコンパイラRubyのtrunk (2.6.0-dev) にマージし、昨日TD Tech Talk 2018で以下のような内容の発表をしました。

speakerdeck.com

まだそれほど速くできていないということもあり、私はTwitterでのみ共有して満足していたのですが、海外の方がいくつか記事を書いてくださいました。

とても丁寧に書かれているので、私の記事がわかりにくければそちらもどうぞ。ただ、いくつか補足したい点があったのと、また昨日の発表内でうまいこと説明できなかった点があるので、そのあたりに触れつつ今回マージしたJITコンパイラについての概要を記事にしようと思います。

"MJIT"とは何か

この見出しの内容はある程度知ってる方がいる気がするので、僕のYARV-MJITの話を聞いたことがある人は「RubyJITの現状」まで飛ばしてください。昨日の僕の発表を聞いてくださった方は、発表時に話せなかった「JITコンパイラの自動生成の裏側」まで飛ばしてください。

ここでは、RubyJITにまつわる最新の動向をダイジェストでお伝えします。

RubyVMのRTL命令化プロジェクト

これは2016年7月の話なのですが、Ruby 2.4でHashを高速化 (Feature #12142)したVladmir Makarovさん(以下Vlad)が、VM performance improvement proposal (Feature #12589)というチケットを作りました。

タイトルの通りVMのパフォーマンスを上げる目的の提案なのですが、RubyVMJava等と同じくスタックベースの命令で作られてるのに対し、その命令を全てレジスタベースのもの(以下RTL: Register Transfer Language)に置き換えようという内容になっています。

VladはGCCレジスタアロケータの作者であり、インストラクションスケジューラのメンテナでもあります。そのGCCは内部でレジスタベースの中間表現を使っているため、そこで得たノウハウを適用して高速化しようという目論みだったのでしょう。

全ての命令を書き変えるような大きな変更にも関わらず、make testが通る完成度のRTL命令の実装が公開されました。純粋なレジスタベースへの置き換えだけでなく、実行時にレシーバや引数のクラスに応じてバイトコードを動的に書き変える仕組みを用意して、クラスに応じた特化命令が使われるようにする変更も入っています。

一方この変更だけでは、一部のマイクロベンチマークではパフォーマンスが向上したものの、Optcarrotのようなある程度規模の大きいベンチマークではパフォーマンスが劣化してしまう状態でした。

RTL命令上で動くJITコンパイラMJIT

そのような問題を解決するため、2017年3月末、RTL命令をベースに動作するMJIT (MRI JIT)というJITコンパイラの実装が公開されました。

rtl_mjit_branchのページにRubyJITのための方針がいくつか検討されていますが、採用されたのはCのコードを生成してディスクに書き出し、Cのコンパイラを実行してそれをコンパイルしロードするというものでした。

f:id:k0kubun:20180216205959p:plain

RTL化の際にインライン化に有利な命令列が動的に生成されるようになっており、そのコードをJITすると、先ほどのOptcarrotでもかなり高速(Ruby 2.0の2〜3倍)に動作するようになりました。

一方で、VMの実装を大きく書き換えた副作用として、make testは通るもののmake test-all等のよりシビアなテストはJIT有効/無効に関わらず通らない状態になってしまい、Optcarrotもベンチマークモードでしか起動可能でなかったり、Railsも起動に失敗する状態でした。

YARV-MJITとは

このようにRTL命令化はある程度リスクのある変更だったのですが、私の会社ではクラウドサービスをRailsアプリケーションとして実装して提供しておりとても高い可用性が求められるため、想定外のバグが発覚するリスクの高い変更がRubyに行なわれてしまうと、バージョンアップがとても大変になるだろうという懸念がありました。

RubyKaigiでのVladの発表の後、私はどうにかしてリスクを軽減しながらRubyJITを導入できないか考えました。そこで思いついたのがYARV-MJITと私が呼んでいるアプローチで、MJITのJIT基盤をそのまま使いながら、RTL命令ではなく現在のスタックベースのYARV命令をベースにJITコンパイルを行なうというものでした。

このJITのアプローチはVMの実装を全く変えずに実現できるため、リリース後JITコンパイラにバグが発覚しても、JITコンパイラを無効にしてしまえばRubyが今まで通り動作することがほぼ確実となります。

リスクを抑えるためVMの命令セットに変更を加えないという制約をかけていたのもありMJITほどのパフォーマンスは出せなかったのですが、公開後様々な改善を減てYARV-MJITでもある程度MJITに近い速度が出るようになりました。

いくつかの未熟な最適化を取り除くと、YARV-MJITはある程度の速度を保ちながらもJIT有効な状態でRubyのテストが全て通るようになったため、2017年末にMJITのJIT基盤とYARV-MJITのマージを提案し、ポテンシャルバグやランダムなクラッシュの修正の後、2018年2月に無事マージされました。

RubyJITの現状

Ruby 2.6に向けて今回マージされたのは主に以下の2つのコンポーネントに分けられます。

VladのMJITはもともとJITコンパイラ部分も含めた呼称と考えられるため、そのうちJITの基盤となる部分のことを私はチケットやコミットメッセージ上ではMJIT Infrastructureと区別して呼んでいます。

が、たくさん名前を覚えるのは面倒ですし、ユーザーの皆さんからしたら中の実装に何が採用されたとかはどうでもいいかもしれないので、両方合わせてMJITと呼んでいただくのは別に構わないと思います。

これは昨日TD Tech Talk 2018で話した内容なのですが、現状それぞれどこまで実装が進んでいるのかという点について共有します。

JITコンパイラについて

ベンチマーク

パフォーマンス改善のための変更なので、いくらリリースのリスクを抑える目的があるとはいえ、当然パフォーマンスが向上しなければ導入するべきではないでしょう。なので、私のマシン(Intel 4.0GHz i7-4790K 8 Cores, 16GB memory, x86_64 Linux)でのベンチマーク結果をいくつか載せておきます。

Optcarrot

f:id:k0kubun:20180216210028p:plain

これは最新のOptcarrotベンチマークの結果です。YARV-MJITも一時期68fpsくらい出ていたんですが、そこからいくつか不完全な最適化を取り除いているためこのようなパフォーマンスになっています。

JITのマージのPull Requestでは最初63fpsのベンチマークを貼っていて、Squareさんの記事でもこれを引用されているように見えるのですが、そこから更に最適化を1つ一時的に外しているため、PRの本文に書いてあるようにマージの段階でベンチマーク結果が変わって58fpsになっています。マージ後r62388で59fpsくらいになりましたが、63fpsに届くために必要な最適化*1はまだ外したままです。そのうち直します。

RubyBenchの更新が最近止まっているため正確にどの変更の影響が大きいのか調べるのが難しい状況*2なのですが、最近shyouheiさんがVMで様々な改善を行なっている*3ためか、Ruby 2.6では実はJIT無しでもOptcarrotベンチマークは結構速くなっています。

JITがそれほどパフォーマンスに影響を与えられていないというのは私個人にとっては少し残念な話ですが、Ruby 2.6がJITなしでも速くなっているというのは良いニュースでしょう。

Benchmark in "Playing with ruby's new JIT"

VM実行に比べ現時点のJITの高速化の可能性は本当に上記のように数パーセントしかないのでしょうか?

Playing with ruby's new JIT: MJITの記事にはベンチマークスクリプトがあり、それをtrufflerubyの人が改良したものが以下のようなものになっています。

def calculate(a, b, n = 40_000_000)
  i = 0
  c = 0
  while i < n
    a = a * 16807 % 2147483647
    b = b * 48271 % 2147483647
    c += 1 if (a & 0xffff) == (b & 0xffff)
    i += 1
  end
  c
end

benchmark.ips do |x|
  x.iterations = 3
  x.report("calculate") do |times|
    calculate(65, 8921, 100_000)
  end
end

これを私のマシンでも実行してみました。

$ ruby -v
ruby 2.6.0dev (2018-02-15 trunk 62410) [x86_64-linux]
$ ruby bench.rb
Warming up --------------------------------------
    calculate    13.000  i/100ms
    calculate    13.000  i/100ms
    calculate    13.000  i/100ms
Calculating -------------------------------------
    calculate    1.800k (± 2.7%) i/s - 8.996k in   5.002504s
    calculate    1.785k (± 7.4%) i/s - 8.853k in   5.003616s
    calculate    1.802k (± 4.0%) i/s - 8.996k in   5.006199s
$ ruby --jit bench.rb
Warming up --------------------------------------
    calculate    13.000  i/100ms
    calculate    18.000  i/100ms
    calculate    27.000  i/100ms
Calculating -------------------------------------
    calculate    7.182k (± 9.1%) i/s - 35.397k in   5.000332s
    calculate    7.296k (± 2.9%) i/s - 36.450k in   5.001392s
    calculate    7.295k (± 3.1%) i/s - 36.450k in   5.002572s

1.802k → 7.296k ですから、大体4.0倍くらいですかね。Ruby 3x3達成しましたね。

これは多分JITの宣伝のために作られたような架空のワークロードだと思いますが、正しく現実世界の用途に向けJITをチューニングしていけば、現在の最適化戦略で大きな改善ができる可能性を示唆していると思います。

またこのベンチマークに関しても、手元のtrufflerubyで動かしたところ667.854k (trunkのVM実行の約370倍)になったので、我々もこのような最適化しやすいケースに関しては数百倍のオーダーを目指すべきだと認識しています。

その他のベンチマーク

JITコンパイラをマージした時のコミットのコミットメッセージに、Optcarrot以外にもマイクロベンチマークの結果(JIT offの状態に比べ2倍程度速くなっているのがいくつかあります)と、Railsでのベンチマーク結果を載せています。

Railsベンチマークに関してはNoah Gibbsさんのrails_ruby_benchを使用したいところなのですが、私の手元で動かない状態になっているので、それが直せるまでは現状そのベースとなっているdiscourseリポジトリベンチマークを使用しています。

そのRailsベンチマークだと速くなるどころか若干遅くなってしまうというのが現状であり、原因はまだ調査中なのですが、Sam SaffronさんがUnicornはforkするので子プロセスでJITが有効になっていないのではという指摘をされていて、実際無効にしているので、このあたり直しつつ遅くなっている原因の切り分けもやっていこうとしている状態です。

JITコンパイラの自動生成

JITコンパイラVMと同じ挙動をしなければいけないため、素朴に実装するとJITコンパイラVMの実装のコピペが多くなり、保守性が下がってしまいます。

RubyVMはもともとinsns.defという特殊なテンプレートフォーマットをパースしてVMを生成するERBに渡すことで生成されているのですが、insns.defからJITコンパイラも生成するようにできないかという提案をささださんからいただき、やってみたところ上手くいきました。

以下のようなmjit_compile.inc.erbというファイル(わかりやすいようにいろいろ省略)があり、

switch (insn) {
% (RubyVM::BareInstructions.to_a + RubyVM::OperandsUnifications.to_a).each do |insn|
  case BIN(<%= insn.name %>):
<%= render 'mjit_compile_insn', locals: { insn: insn } -%>
    break;
% end
}

命令に応じてswitch-caseで分岐しながら、_mjit_compile_insn_body.erbというファイルでinsns.defの中身をfprintfするコードを生成しています。

% expand_simple_macros.call(insn.expr.expr).each_line do |line|
%   if line =~ /\A\s+JUMP\((?<dest>[^)]+)\);\s+\z/
        /* JUMPの動的生成コード */
%   else
        fprintf(f, <%= to_cstr.call(line) %>);
%   end
% end

マクロをJIT向けに事前に変換し、またJUMP命令は特にcase-when文に関して簡単な変換が難しいため動的に生成するようになっています。それ以外のコードは基本的にはfprintfの羅列で、このmjit_compile.inc.erbからは以下のようなファイルが生成されます。

switch (insn) {
  case BIN(nop):
    fprintf(f, "{\n");
    {

        fprintf(f, "    reg_cfp->pc = original_body_iseq + %d;\n", pos);
        fprintf(f, "    reg_cfp->sp = (VALUE *)reg_cfp->bp + %d;\n", b->stack_size + 1);
        fprintf(f, "    {\n");
        fprintf(f, "        /* none */\n");
        fprintf(f, "    }\n");
        b->stack_size += attr_sp_inc_nop();
    }
    fprintf(f, "}\n");
    break;
  case BIN(getlocal):
    /* ... */
}

上記はMJITのワーカースレッドとして実行され、実行時に以下のようなコードをファイルに書き出します。

VALUE
_mjit0(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp)
{
    VALUE *stack = reg_cfp->sp;
    static const VALUE *const original_body_iseq = (VALUE *)0x5643d9a852a0;
    if (reg_cfp->pc != original_body_iseq) {
        return Qundef;
    }

label_0: /* nop */
{
    reg_cfp->pc = original_body_iseq + 1;
    reg_cfp->sp = (VALUE *)reg_cfp->bp + 2;
    {
        /* none */
    }
}

/* ... */

} /* end of _mjit0 */

rb_execution_context_tというのはスレッドのようなもので、rb_control_frame_tがコールスタックのうちの1つのフレームを表しているので、あるスレッドのあるフレームで実行できるようなインターフェースになっています。

このファイルをコンパイルしてdlopen, dlsymするとこれの関数ポインタが得られるので、そのポインタが既に存在したらVM実行のかわりにこれを呼び出すことでJIT実行が実現されます。

JITコンパイラの自動生成の裏側

ささださんくらいしか興味がないような気もするのですが、この自動生成のために必要だったことを紹介します。

次のメソッド呼び出し処理のマクロ定義分岐

RubyVMを1つ起動するとsetjmpという重い処理をする上マシンスタックを大きく消費するため、あまり何度も起動したくありません。従ってRubyVMは、Rubyのコードだけを処理している間は別のVMを起動せず使い回すように作られています。

これがRESTORE_REGS();,NEXT_INSN();という2つのマクロで実現されていたのですが、マクロが2つに分かれていると変換もしづらくとても都合が悪いため、これらをまとめて実行するEXEC_EC_CFP()というマクロを勝手に作り、JITの関数内ではそれがvm_exec()というVMを起動する関数を実行するようになっています。書いてて思ったけどここで先にmjit_exec()するようにした方が速いような。

vm_exec.h内で、MJIT_HEADERというマクロが定義されていたらマクロの定義自体が変わるように実装されています。

例外命令のマクロ定義分岐

Rubyの例外は基本的にはlongjmpを呼び出し大域脱出することで実現されているのですが、それがまあとても遅い処理なわけで、Rubyが内部実装の都合で使っているthrow命令は単にVM実行の関数(vm_exec_core)からreturnして、その関数をラップしている関数(vm_exec)で大域脱出先のISeqを線形に探索して実行する、というような実装になっています。

当然このコードはvm_execにラップされていないと動かないのですが、先ほどいったようにこいつはsetjmpするのでそれも呼びたくなく、基本的にmjit_execvm_execを経由せず呼び出しているため、throw命令に使われているTHROW_EXCEPTION()というマクロを、単にreturnするのではなくlongjmpするように変更しています。

これもvm_exec.h内で分岐しています。

例外ハンドラーからのISeq実行の無視

vm_execにラップされてないと動かないのと同じ問題が他の場所にもあり、vm_exec内の例外ハンドラからISeqの実行が始まる場合setjmpの実行がスキップされておりこれも動かないので、その場合そのISeqはJITで実行されないようになっています。

opt_case_dispatchの動的コンパイル

case-whenってありますよね。opt_case_dispatchオペランドにHashが入っていて、キーが単純で高速に分岐できる場合はそこからテーブル引きで分岐先のアドレスが取得されるようになっています。

で、それはVMの実装向けにgotoが走るようになっていてそのままでは動かないので、JITの実装向けにJITの関数内でのラベルにgoto先を変換しています。

これにはそのHashの中身を見る必要があり、これを見るためにHashの操作を行なう関数を使っていたらランダムにSEGVしていた(GVLのコントロールから外れているMJITのスレッドでRubyVMを利用すると、当然同時にメインのスレッドでも使っているので壊れてしまう)ので、Hashからstを取り出し、stを操作する関数を使うようにしています。

プログラムカウンターの移動

これはRubyのメソッドJITコンパイラの実装初心者(誰)がハマりがちな罠なのですが、一見不要そうに見えるプログラムカウンターの移動をサボると、突然例外で大域脱出した時に実行するべきISeqのlookupが意図通りではなくなってしまいます。なので、全ての命令の実行時に毎回プログラムカウンターを動かす必要があります。

スタックポインターの移動

もともとYARV-MJITはVMのスタック操作相当のものをJITされた関数のローカル変数上で行なうようになっており、それによってそのJITされた関数内でのスタック操作が他の関数から影響を受けないことを保証してコンパイラに最適化させるという意図を持っていました。値を適切なローカル変数に代入してVMのスタックの挙動を再現するため、JITコンパイル時にスタックのサイズを計算するコードが入っており、スタックポインタ相当の計算をコンパイル時に終わらせることができるようになっています。

が、これだと例外で大域脱出した瞬間そのローカル変数の値が失われてしまうという問題がありました。そのため、この速さを維持するためにはlongjmp前にlongjmpで失われる全てのフレームのローカル変数をそのフレームの外側から退避する必要がありました。

それの実装がまだできてないので現在はローカル変数のかわりにVMのスタックを使用しており、スタックポインタも毎回動かさなければならない状態になっており、これが最初のPull Requestの提案から遅くなってしまったポイントです。

この退避の実装を実装するためにローカル変数のアドレスを必ず取る必要が発生するので、もしかするとそのせいでコンパイラが最適化できなくなり、例外が絶対に発生しなさそうなISeqではそれをスキップした方がいいかなあ、とか考えています。

JIT基盤について

プラットフォームの対応状況

MJITのワーカースレッドがpthread直書きだったのを、YARV-MJITの開発中にRubyのスレッドの抽象化レイヤーを使って移植するのはマージ前に終わっていました。

マージ後、MJITのランタイムではなく、JITコンパイル時に使用するヘッダーのビルドが以下のようなプラットフォームで動かなくなるのを報告いただいたり、RubyCIで確認したりしました。

私の手元にはLinux, macOS, MinGW, 新しめのVisual Studioの環境があり、そこで動くことは確認していたのですが、Rubyが要求するポータビリティにはそれでは足りなかったようです。NetBSDSolarisVMをインストールしたり、hsbtさんにICCの環境を用意していただいたりしてデバッグし、どうにかビルドが通る状態に持ってきました。

その後、同期的にコンパイルを行なう--jit-waitというオプションを用いてRubyユニットテストJITのテストを追加し、gcc+Linuxという環境では基本的に動作しているのですが、clangとMinGWの環境でいくつかRubyCIの失敗が残っていて、これらはこれから修正していこうと思っている状態です。

また、Windows環境に入っているCのヘッダは、コンパイル時間を短縮するための変換をしようとするとエラーになる状態になっており、一旦それをやらないようにしているので、MinGWではJITは一応動いていますがコンパイル時間がかなり長いという問題があります。これはヘッダの変換のバグを直すしかないでしょう。

セキュリティ

JIT基盤に関してはnobuさんやusaさんがいろいろ修正をしてくださっていて、その中でnobuさんがセキュリティの脆弱性についても修正してくださっています。

MJITのワーカースレッドがCのファイルを書き出す時に/tmp/_ruby_mjit_p123u4みたいなファイルをfopenしていたのですが、まずファイルのパーミッションに対するコントロールがなかった(これはnobuさんからの指摘で認識していたのでそのうち直すつもりでマージしていた)のと、既存のファイルに書き始めると任意のパーミッションになってしまうため、ファイルを作成する時もそのパスに既にファイルがないことが保証されるコードになりました。

いまのところそれ以外の脆弱性は確認されていませんが、リリース前に精査する必要がありそうです。現時点で本番環境で試していただくのは、こちらとしては助かるのですが、正式リリースまではある程度自己責任でやっていただく形になってしまいます。

その他の補足

"Startup Time"について

Ruby's New JITには、"Startup Time"というセクションで以下のように説明されています。

Startup time is one thing to take into consideration when considering using the new JIT. Starting up Ruby with the JIT takes about six times (6x) longer.

重要な気付きだと思います。この理由は、--jitがついているとRubyの起動時にJIT用のprecompiled headerのコンパイルを開始し、またRubyの終了時にそのJITワーカースレッドが現在行なっているJITコンパイルが終わるのを待つため、最低でもCのファイル1回分のコンパイル時間がかかってしまいます。コンパイルが遅いコンパイラを使うとこの時間は伸びるわけです。

解決策は以下の2つあります。

  • プリコンパイルドヘッダをビルド時に生成する
  • JIT用スレッドをいきなりキャンセルする

プリコンパイルドヘッダは実際にはRubyのビルド時にはプリコンパイルされていません。これはコンパイラのバージョンが途中でアップグレードされプリコンパイルドヘッダのフォーマットが変わってしまうのではないかということを懸念して毎回コンパイルしていることによります。もしサポートしているコンパイラ全てでそのようなことがほとんどないことがわかればそれで良いでしょうが、現在は保守的な方向に倒しているためこうなっています。

JIT用のスレッドを任意の瞬間にキャンセルして、途中で生成されていた全てのファイルを上手いこと特定して削除できる方法があればより速く終了できるのですが、これはあまり安全な方法ではないでしょう。

--jit-ccはなくしました

記事では--jit-cc=gccが使われていたりするのですが、JIT時に使うヘッダの生成に使用したコンパイラ以外のコンパイラJITに使用すると割と多くのケースで壊れてしまい、それが壊れるのはRubyの責任ではないのでサポートしないことにしています。

--jit-ccは間に合せで残していたオプションなのですが、nobuさんがいろいろ改善してくださって、Rubyのビルドに使われたコンパイラ(JIT用のヘッダの生成にも使われる)をマクロで判定してコンパイラを特定するように変わりました。

従って、普段gccを使っているがclangでのパフォーマンスも確認してみたい、という人はclangでRubyをビルドし直すようにしてください。

謝辞

Ruby 2.6のリリースまで私のJITコンパイラが残っているかはまだわかりませんが(!?)、RubyJITが入るまでに様々な方に助けていただいたので、ここで以下の方にお礼を申し上げます。

  • MJITの基盤を発明したVladimir Makarovさん
  • 様々な場所でお世話になっているまつもとさん
  • YARV-MJITについて様々な相談に乗っていただいたささださん、mameさん
  • マージ前に多くのバグを報告・修正していただいたwanabeさん
  • 最初のMinGWサポートを書いたLars Kanisさん
  • ビルド環境でいつもお世話になっているhsbtさん
  • マージ後基盤の方を直してくださっているnobuさん、usaさん、znzさん、knuさん
  • コードを通じてJITコンパイラの生成回りで助けていただいたshyouheiさん

Ruby 2.6.0 preview1 そろそろ出ます

JITが入ったのでpreview1早めに出しませんかみたいな提案があり、成瀬さんがpreview1の準備を始めています。 その時点ではJITはそれほど速くなってないと思いますが、新しいものが好きな方はそこで少し遊んでみていただけると助かります。

*1:冒頭のスライドでLazy stack pointer motionと呼んでる奴

*2:私がRubyアソシエーションの開発助成金で行なっているベンチマークプロジェクトでも3月にそのような継続的ベンチマーク環境を用意しようとしているところなので、来月にはわかりそうです

*3:https://github.com/ruby/ruby/pull/1779 など