RubyはJust-In-Time（JIT）コンパイラという機能を備えており、これを有効化すると実行時に機械語を生成して様々な最適化が行なわれ、実行が高速になります。Ruby 3.3にはYJITとRJITという2つのJITコンパイラがあり、デフォルトでは無効になっていますが、それぞれ--yjitと--rjitで有効化できます。

この回では、Ruby 3.3でYJITの性能特性が変化した点や、YJITに新たに追加された便利な機能、またRJITはどのように使うものであるかについて解説します。

YJIT

YJITは現在様々な企業の本番環境で利用されており、最大40%の高速化が報告されています。RubyはJITコンパイラなしでもWebアプリケーションなどの実装では十分な性能の出る言語ですが、アプリケーションの実装を書き換えずにYJITを有効化するだけで高速になるわけですから、これを活用しない手はありません。現在開発中のRails 7.2ではYJITがデフォルトで有効化されるようになり、今ではRubyはYJITを使って実行するのが業界標準になりつつあると言えるでしょう。

JITコンパイラは実行速度を改善する一方で、機械語やその管理のためにより多くのメモリを消費するというトレードオフがあります。これはJavaやJavaScriptといった他の言語の処理系にも一般に当てはまる問題です。しかしRubyは並列化のために複数プロセスが使われることの多い言語ですから、YJITが消費するメモリの量は並列化に使用するプロセスの数分倍になってしまいます。そのため、Ruby 3.2ではYJITがメモリ管理の点で運用の難しさが指摘されることもありました。

Ruby 3.3のYJITでは、高速化のためにRuby 3.2より多くメモリを消費するようになる変更もあった一方で、それを相殺するべく大幅な省メモリ化が行なわれた上、メモリ管理のための様々な機能が追加されました。本節では、YJITが扱うメモリを管理する上で重要となるRuby 3.3の機能や変更点を中心に解説します。

より多くのコードを生成

Ruby 3.3のYJITでは様々な最適化が行われました。たとえば次のような最適化が挙げられます。

• VMのスタック操作に新たにレジスタ割付が実装された
• 本番環境やベンチマークでよく使われる様々なCメソッドに特別な最適化が実装された
• 簡単な定義のRubyメソッドがインライン化できるようになった

これにより、より多くのタイプのメソッド呼び出しやインスタンス変数アクセスが高速化されるようになりました。

よく多くの種類のコードがコンパイルできるようになった副作用として、様々なアプリケーションにおいて、Ruby 3.3 YJITが生成するコードのサイズはRuby 3.2 YJITより大きくなりました。これはなぜかというと、YJITはコンパイルに対応していないRubyコードに遭遇すると諦めてインタプリタにフォールバックするようになっており、Ruby 3.2 YJITではコンパイルを諦めていた場所の多くがRuby 3.3 YJITではコード生成できるようになったからです。

これによるYJITが使用するメモリサイズへのインパクトはアプリケーションによります。たとえばRailsのベンチマークでは、Ruby 3.2の時点では93%程度のRuby VM命令がYJITにより実行されていたのが、Ruby 3.3では99%がYJIT実行になっています。Rubyを3.2から3.3にアップグレードする際、この変更のメモリ使用量へのインパクトが気になる場合は、ステージング環境などを使って、YJITの生成コードのサイズを示すRubyVM::YJIT.runtime_stats[:code_region_size]の変化を監視しておくと良いかもしれません。

メタデータの省メモリ化

YJITが消費するメモリには、code_region_sizeでカウントされる「YJITが生成したコード」に加えて、yjit_alloc_sizeでカウントされる「生成コードの管理に使うメタデータ」があります。YJITのすべての生成コードはメタデータを持つため、yjit_alloc_sizeは一般にcode_region_sizeに比例する傾向にあります。

Ruby 3.3ではcode_region_sizeが増えるような変更が行われ、また新たなレジスタ割付のためのメタデータが新たに管理されるようになったため、yjit_alloc_sizeも増えてしまうことが考えられます。しかし、Ruby 3.3ではメタデータの構造が3.2よりコンパクトになり、生成コードあたりのメタデータのサイズはむしろ3.2より小さくなりました。

そのため、ベンチマークによっては、code_region_sizeとyjit_alloc_sizeの合計はRuby 3.2よりRuby 3.3のほうが小さくなっています。RubyVM::YJIT.runtime_stats[:yjit_alloc_size]はRuby 3.2では特殊なビルドをしないと使えませんでしたが、Ruby 3.3ではこの統計情報はデフォルトで提供されるようになりました。code_region_sizeとyjit_alloc_sizeの両方を監視しておくと、YJITのメモリ使用量のインパクトが理解しやすくなります。

Code GCの無効化と--yjit-exec-mem-size

YJITが生成するコード量は--yjit-exec-mem-sizeでコントロールできます（デフォルト64MiB⁠）⁠。生成コードのサイズが--yjit-exec-mem-sizeに達すると、YJITはデフォルトで以下のような動きをします。

• Ruby 3.2：すべての生成コードを破棄し、以降呼ばれたメソッドをコンパイルし直す。
• Ruby 3.3：新たにメソッドをコンパイルしなくなる。未コンパイルのメソッドはインタプリタ実行される。

Ruby 3.2のこの挙動をCode GCと呼んでいます。数時間に一回Code GCが走る程度なら大した性能影響はないのですが、--yjit-exec-mem-sizeが小さすぎると、頻繁にコンパイルし直すコストによってアプリがむしろ遅くなる場合があります。

こういった問題にヒットしにくくなるよう、Code GCはデフォルトで無効になりました。これの最大の利点は気軽に--yjit-exec-mem-sizeが下げられるようになった点で、code_region_sizeの変化を参考にしつつ、--yjit-exec-mem-size=32のような設定を使うのが現実的な選択肢になりました。

参考までに、筆者が所属しているShopify社内最大のアプリケーションであるモノリスでは、デフォルトである--yjit-exec-mem-size=64でそれより大きな値とほぼ同じ性能が出ました。--yjit-exec-mem-size=32では性能は多少劣化しましたが、そこまで大きな差ではありませんでした。このモノリスは大変大きなアプリケーションですから、読者の環境ではこれより小さな設定で十分な性能が出るかもしれません。

RubyVM::YJIT.enableを使った省メモリ化

YJITは、Ruby 3.2まではコマンドライン引数--yjitや環境変数RUBY_YJIT_ENABLE=1で有効化するしかありませんでした。しかしRuby 3.3ではRubyコード内でRubyVM::YJIT.enableを呼び出すだけでYJITが有効化できるようになりました。

これを使ってYJITの起動を遅延させることで、アプリ初期化後は使われないコードのコンパイルを避けメモリ消費量を削減できます。Railsのイニシャライザから使うだけでも効果はありますが、理想的にはUnicornのafter_forkやPumaのafter_worker_forkから呼び出すのが望ましいです。

なお、--yjit-exec-mem-sizeなどのチューニングオプションを指定するだけでも起動時にYJITが有効化されるため、その場合も遅延起動するには--yjit-disableを明示する必要があります。

RJIT

RJITはRubyのJITコンパイラの最適化実験を行なう環境を提供することを目的としたJITコンパイラ基盤で、Rubyで書かれたJITコンパイラを提供しています。Ruby 3.2にはYJITの他にMJITという別のJITコンパイラがありましたが、RJITはMJITを置き換える形でRuby 3.3に新たに登場しました。本節では、RJITについて最低限知っておくべき知識と、JITコンパイラ基盤の使い方について解説します。

RJITがMJITを置き換えた理由

MJITはRuby 2.6でマージされたRuby最初のJITコンパイラで、Ruby 3.0をRuby 2.0の3倍高速にする「Ruby 3x3」目標の達成で活躍しました。しかし、Ruby 3.1で2つ目のJITコンパイラであるYJITが登場し、Railsなど数々のベンチマークでMJITの性能を圧倒したため、YJITの登場後にMJITを本番環境で使う人はいませんでした。

YJITが追加されてからも、Ruby 3.2 MJITは以下の点でYJITにはないメリットを持っていました。

• Rubyで書かれているため、Rustで書かれているYJITに比べ最適化アイデアの実験が容易である。
• MJITをモンキーパッチすることで、MJITではない独自のJITコンパイラを実装する基盤にできる。

Ruby 3.3 YJITに実装されたいくつかの最適化は元々MJITやRJITにあったものですし、MJITの基盤を使って独自のJITコンパイラを開発しているコミッタたちはそこで得た経験をYJITのデザインの議論に繋げたりしているため、実際にユーザーが使うYJIT以外にも実験用にJITコンパイラ開発基盤を維持しておくことは、間接的にYJITの性能改善に貢献しています。JavaバーチャルマシンにもJavaでJITコンパイラを記述するための機能が備わっていますが、Rubyにもそのような環境を維持することでRubyの性能改善の機会を最大化したいと筆者は考えていました。

一方、MJITは実行時にCコンパイラを起動する非常にユニークなデザインであったことから、それを可能にするための特殊な実装が随所で必要で、時にはYJITの改善も困難にするような保守性の問題がありました。RJITは、その問題を解決しつつ上記のメリットを残すべく作られました。RJITはYJITをRustからRubyに書き直したようなデザインになっており、そのためRuby VMに必要な基盤がYJITと共通化されることでシンプルかつ高速になり、またYJITに活かす最適化アイデアの実験により適した実装になりました。

RJITとYJITの使い分け

RJITはYJITの開発を助けることを主目的とした実験的なプロジェクトであり、本番環境ではRJITではなくYJITを使うことが推奨されます。RJITにYJITにはない最適化が実装されることもありますが、それがベンチマークに有意差をもたらす場合は基本的にその最適化はRubyのリリース前にYJITに移植されています。そのため、リリースされるRuby上では常にYJITがRJITより速いことが期待されます。

RJITは現在x86_64アーキテクチャのUnix環境のみサポートしています。最適化の実験を容易にする目的から、RJITサポート対象の環境でRubyをビルドするとRJITサポートがデフォルトで有効になります。一方、YJITはビルド時にRustコンパイラが存在していないとYJITサポートは有効になりません。そのため、Rustに依存しない分RJITの方がビルドは容易ですが、性能面でも安定性でもRJITはYJITに劣ります。よってYJITサポートが有効でないRubyをビルドしてしまっている場合、そのままRJITを使うのではなくYJITサポートつきでRubyをビルドし直すことが推奨されます。

RJITで独自のJITコンパイラを実装する方法

RJITはそれ自体がJITコンパイラを提供しており、これはRJITの基盤が十分な機能を維持していることのテストに役立ちますが、基本的にはそれとは全く異なるJITコンパイラに差し替えて使うことが想定されています。RJIT基盤上に実装されたJITコンパイラとしては、たとえばtenderjitやhawthjitなどがあります。

ここでは、RJITの基盤を使って独自のJITコンパイラを作って動かす方法を紹介します。Rubyを--rjitつきで起動すると、あるメソッドの呼び出し回数が10回に達した時、そのメソッドを引数に取りRubyVM::RJIT::Compiler#compileが呼ばれることによってJITコンパイルが行なわれます。そのため、このメソッドをモンキーパッチで上書きすることによってJITコンパイラを全く別のものに差し替えることができます。

具体的には、以下のようにします。

module RubyVM::RJIT
  Compiler.prepend(Module.new {
    def compile(iseq, cfp)
      cb = CodeBlock.new(mem_block: C.mmap(4 * 1024), mem_size: 4 * 1024)

      asm = Assembler.new
      # Pop cfp: ec->cfp = cfp + 1 (rdi is EC, rsi is CFP)
      asm.lea(:rax, [:rsi, C.rb_control_frame_t.size])
      asm.mov([:rdi, C.rb_execution_context_t.offsetof(:cfp)], :rax)
      # return true
      asm.mov(:rax, Qtrue)
      asm.ret

      iseq.body.jit_entry = cb.write(asm)
    end
  })

  enable
end if RubyVM::RJIT.enabled?

def compiled = false
p compiled #=> true

JITコンパイラを差し替える最中にRJITが有効になっていると、差し替えの最中にRJITが少しコンパイルを行なってしまうリスクがあります。Rubyを--rjit-disableつきで起動することでRJIT無効の状態でRubyを立ち上げ、JITコンパイラを差し替えた後に上記コードのようにRubyVM::RJIT.enableを呼び出すことでそれを防ぐことができます。

上記のコードを--rjit-disable --rjit-call-threshold=1で実行すると、#compiledの初回の呼び出しで自前のJITコンパイラが起動されます。これは本来falseを出力するスクリプトですが、このJITコンパイラは何をコンパイルしてもtrueを返すように実装されているため、trueを出力します。

ここで使われているRubyVM::RJIT::Assemblerの使い方や具体的なJITコンパイラの記述方法はruby-jit-challengeというリポジトリで解説しています。興味のある方はこちらをご覧ください。

まとめ

Ruby 3.3ではYJITのメモリ使用量を抑えるための様々な改善が行なわれたほか、新たなJITコンパイラであるRJITが追加されました。Ruby 3.4ではYJITでの複雑なメソッドのインライン化など高度な最適化の実現が計画されています。Rubyの今後の速度向上にご期待ください。