Tech Racho エンジニアの「?」を「!」に。
  • 開発

Rubyの新しいJIT「MJIT」で早速遊んでみた(翻訳)

概要

原著者の許諾を得て翻訳・公開いたします。

画像は英語記事からの引用です。

Rubyの新しいJIT「MJIT」で早速遊んでみた(翻訳)

Just-in-timeは幻想である: Albert Einstein
元ネタ: Time is an illusion

今週、Takashi Kokubun (@k0kubun)がMRI Rubyに最初の実装をmergeしました。

commitの解説にあるように、MRI RubyのJITはまだ登場したばかりです。まだRailsを高速化できる状態ではなく、現状では初期プロトタイプの一部が前より遅くなっているものもありますが、それでもJITは登場しました。

私はJITの登場に興奮しました。Rubyは以前から遅いという評判がありました(本当に遅いと言えるのは一部にとどまりますが)。JITはこの点を解決できる可能性があります。

JITは実験段階ですが、私は試さずにはいられませんでした。

MRI trunkをcloneしてビルドし、~/.rubiesディレクトリにインストールしてみました。

$ git clone https://github.com/ruby/ruby
$ cd ruby
$ autoconf
$ ./configure --prefix=$HOME/.rubies/ruby-2.6.0-dev
$ make
$ make install

続いて新しいRubyに切り替えました。

$ chruby 2.6
$ ruby -v
ruby 2.6.0dev (2018-02-05 trunk 62211) [x86_64-linux]

Advent of Code 2017day 15で使った私のソリューションでテストしました(以下がネタバレ)。

def calculate(a, b, n = 40_000_000)
  n.times.count do
    a = a * 16807 % 2147483647
    b = b * 48271 % 2147483647

    (a & 0xffff) == (b & 0xffff)
  end
end

raise unless calculate(65, 8921) == 588
p result: calculate(699, 124)

このコードはRuby JIT向けの理想的な候補となるはずです。コストの高い他のメソッドを呼び出さずに単に演算を実行しており、ほぼインタプリタの速度にのみ縛られるはずだからです。

$ time ruby --disable-gems 15a.rb
{:result=>600}
  8.30s user
  0.00s system
   100% cpu
  8.306 total
$ time ruby --disable-gems --jit 15a.rb
{:result=>600}
  6.42s user
  0.03s system
   105% cpu
  6.132 total

動きました!JITをオンにしただけで8.3秒から6.1秒に短縮されています。

MJITの内部

まだよくわかっているわけではありませんが、MJITのアプローチはある意味で型破りです。mjit.cのコメントから引用します。

MJITの実装ではCコンパイラ(GCCとLLVM Clang)を広範囲に渡って用いている。これらのコンパイラには、ISEQから生成されたCコードを渡す。これら「産業向け」Cコンパイラの(コンパイル)速度は、通常のJITエンジンよりも遅いが、コンパイル速度よりも、これらのCコンパイラを用いて生成されたコードのパフォーマンスの方が優先度が高い。

MJITはRubyのYARVバイトコードのブロックの1つを受け取り、基本的にはそれをインラインバージョンのCコードに変換し、これが解釈時に実行されることになります。

ここはある意味で他のJITの動作と変わりません。他のJITは実行時にバイトコードをマシンコードにコンパイルしますが、他のJITがこのように既製のCコンパイラを直接使っているのかどうかについてはわかりません。

私はこのやり方が好きです。

この方法でとにかく動いていますし、UNIXらしい素敵な方法だと思います。何よりJITがとても調べやすくなります。

MJITの動作を詳しく見てみましょう。

$ ruby --disable-gems --jit --jit-verbose=2 --jit-save-temps 15a.rb
Starting process: gcc -O2 [...] /tmp/_ruby_mjitp18966u0.c -o /tmp/_ruby_mjitp18966u0.so
JIT success (133.1ms): block in calculate@15a.rb:2 -> /tmp/_ruby_mjitp18966u0.c
{:result=>600}
$ ls -l /tmp/_ruby_mjitp18966u0.*
-rw-r--r-- 1 jhawthorn users  8765 Feb  5 22:18 /tmp/_ruby_mjitp18966u0.c
-rwxr-xr-x 1 jhawthorn users 64384 Feb  5 22:18 /tmp/_ruby_mjitp18966u0.so

MJITは、calculate関数の内部ループから「ホットな」ブロックを取り出し、以下を行います。

  • Cコードに変換する
  • 変換したCコードを/tmp/_ruby_mjitp18966u0.cに書き出す
  • GCC(またはClang)を用いてコンパイルし、/tmp/_ruby_mjitp18966u0.soに出力する
  • 出力された共有ライブラリを動的に読み込んで実行する

補足メモ: このJITの実行時には、Rubyコードがシングルスレッドの場合であってもCPUが105%に達します。2番目のスレッドはJITコンパイラを実行します。余分な5%の部分に該当するのは、生成されるそうしたスレッドやGCCプロセスであり、これらは別のCPUコアで動作するようです。うまい方法です!

生成されたCコードを見てみましょう(完全なコードはgistを参照)。

#include "/tmp/_mjit_hp18966u0.h"
/* block in calculate@15a.rb:2 */

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

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

label_1: /* getlocal_WC_1 */
{
    MAYBE_UNUSED(VALUE) val;
    MAYBE_UNUSED(lindex_t) idx;
    MAYBE_UNUSED(rb_num_t) level;
    level = 1;
    idx = (lindex_t)0x4;
    reg_cfp->pc = (VALUE *)0x561e21fc1118;
    reg_cfp->sp = reg_cfp->bp + 1;
    {
        val = *(vm_get_ep(GET_EP(), level) - idx);
        RB_DEBUG_COUNTER_INC(lvar_get);
        (void)RB_DEBUG_COUNTER_INC_IF(lvar_get_dynamic, level > 0);
    }
    stack[0] = val;
}

label_3: /* putobject */
{
    MAYBE_UNUSED(VALUE) val;
    val = (VALUE)0x834f;
    reg_cfp->pc = (VALUE *)0x561e21fc1128;
    reg_cfp->sp = reg_cfp->bp + 2;
    {
        /* */
    }
    stack[1] = val;
}

label_5: /* opt_mult */
{
    MAYBE_UNUSED(CALL_CACHE) cc;
    MAYBE_UNUSED(CALL_INFO) ci;
    MAYBE_UNUSED(VALUE) obj, recv, val;
    ci = (CALL_INFO)0x561e21fc13a0;
    cc = (CALL_CACHE)0x561e21fc1420;
    recv = stack[0];
    obj = stack[1];
    reg_cfp->pc = (VALUE *)0x561e21fc1138;
    reg_cfp->sp = reg_cfp->bp + 3;
    {
        val = vm_opt_mult(recv, obj);

        if (val == Qundef) {
            return Qundef; /* cancel JIT */
        }
    }
    stack[0] = val;
}

/* あと300行ぐらい続く */

コードのコメントは、コンパイルされているYARVインストラクションに対応しています。各インストラクションはスタック、プログラムカウンタ、およびRubyの残りのVMを、インタプリタのvm_exec_coreとまったく同じ方法で操作します。

最初の乗算で行われるすべての内容を以下のセクションに示します。

  • ローカル変数aを取得する(スタックにpushされている)
  • 被乗数16807をスタックにpushする(object_id 0x834fで表されている)
  • スタック上の2つの値を引数としてvm_opt_multを呼び出す

ここで、Cコンパイラがいかにうまくやっているかが伺える点が1つあります。Cコンパイラは、ここで呼び出されるRubyの内部メソッド(vm_get_epvm_opt_multなど)をできるかぎりインライン化します。これによって、スタックや他のメモリへの代入が不要であると推測できる場合や、最後に書き込まれた値を代入するだけでよい場合に、代入を回避します。したがって、現在のシンプルな実装であっても理にかなった動作になるはずです。

RubyコードをJIT向けに変更するとさらに速くなるか?

今の段階であらゆるRubyコードをJIT向けに書き換えるのは時期尚早でしょう。MJITは数か月もすればまるで違うものになるかもしれません。

ただしあくまでお楽しみとして、MJITでちょっと遊んでみようと思います。

私の元のRubyコードで使っている#timesメソッドや#countメソッドはCで書かれており(訳注: Rubyのネイティブメソッド)、それぞれが独立したメソッド呼び出しなので、JIT化された内部ブロックで最適化されていません。Rubyらしい書き方から少し外れてみることで、ループの外側でもJITが効くようになります。

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
  • JITなしの場合:
$ time ruby --disable-gems 15a.rb
{:result=>600}
ruby --disable-gems 15a.rb
  6.80s user
  0.00s system
  99% cpu
  6.803 total
  • JITありの場合:
$ time ruby --jit --jit-verbose=1 15a.rb
{:result=>600}
  6.85s user
  0.02s system
   103% cpu
  6.862 total

何が起きたかおわかりでしょうか。

(JITありのコードは)JITなしの場合より遅いだけでなく、元のコードよりも遅くなってしまいました。whileループはJITなしの場合でより高速であるにもかかわらずです。

MJITはこのメソッドを最適化対象として認識していませんでした。最適化するメソッドの決定には、かなり素朴なヒューリスティクスが使われています(このあたりは今後変更されるかもしれません)。

$ ruby --help
...
MJIT options (experimental):
  --jit-min-calls=num
                  Number of calls to trigger JIT (for testing, default: 5)

MJITは、呼び出し回数が5回を超えた関数を最適化しますが、さきほどのコードはcalculateを2回しか呼び出していません。元のコードで内部ブロックが最適化されのは、呼び出しが数百万回行われていたためです。

Webサーバーのように長期間実行されるプロセスでは最終的に最適化されるでしょう。私は、この5回という値は小さすぎるのでもっと大きくすべきではないかと考えています。

あくまでお遊びとしてですが、ちょっと細工を加えて、私のメソッドがコンパイルされるよう私たちの愛をMJITに伝えましょう。

# HACK THE PLANET!
# calculateを(反復回数を減らして)5回呼び出すことで
# MJITがここをホットなブロックと認識して最適化するようにする
5.times { calculate(0,0,1) }

raise unless calculate(65, 8921) == 588
p calculate(699, 124)

次の問題は、MJITが非同期に別のスレッドで動作することです。just-in-timeコンパイルは、calculateが最初に5回実行された後で開始されますが、calculateへの本当の呼び出しが最初に行われるまでにコンパイルが完了しません。

これではjust-in-timeではなくjust-too-late(惜しくも間に合わない)です。

これを回避するための--jit-waitというコマンドラインオプションがあります。このオプションをオンにするとMJITが同期的に振る舞い、次に進む前にコンパイルが完了するようになります。ほとんどの場合パフォーマンスは向上するどころか下がってしまいますが、このシンプルなスクリプトにはちょうどうってつけです。

$ time ruby --disable-gems --jit --jit-wait --jit-verbose=1 15a.rb
JIT success (59.8ms): block in <main>@15a.rb:14 -> /tmp/_ruby_mjitp12605u0.c
JIT success (202.6ms): calculate@15a.rb:1 -> /tmp/_ruby_mjitp12605u1.c
{:result=>600}
  3.93s user
  0.03s system
   100% cpu
  3.957 total

これを実際のコードで使うのは時期尚早ですが、Rubyのパフォーマンスがどんなふうになっているかをざっくり見る分にはよいかと思います。

MJITが今後どんなふうに発展するかがとても楽しみです。

もっと詳しく知りたい方へ

関連記事

Ruby 3 JITの最新情報: 現状と今後(翻訳)


CONTACT

TechRachoでは、パートナーシップをご検討いただける方からの
ご連絡をお待ちしております。ぜひお気軽にご意見・ご相談ください。