昨日、RubyのJITの性能改善のためのパッチを入れた。

github.com

JITすればするほどRailsが遅くなる問題

Rubyの次期バージョンである2.6には、バイトコードをCのコードに変換した後、gcc/clangでコンパイルして.soファイルにしdlopenすることで生成コードのロードを行なう、MJITと呼ばれるJITコンパイラが入っているのだが、マージしたころのツイートにも書いていた通り、Railsで使うとより多くのメソッドがJITされるほど遅くなってしまうという問題があった。

結果、"MJIT slows down Rails applications"というチケットが報告されることとなり、昨日までの5か月の間閉じることができなかった。

元の構成

対策を始める前のMJITは大雑把に言うとこういう感じだった。メソッド1つごとに1つの.soファイルが作られ、ロードされる。

無制限にロードしまくるわけではなく、--jit-max-cacheオプションで指定した数(デフォルトでは1000)までしか生成コードを維持しないようになっており、JITされた数が--jit-max-cacheに到達すると、「呼び出し回数が少なく、かつ現在呼び出し中でないメソッド」と「メソッドがGC済のメソッド」向けにdlcloseを行なってから、他の呼び出し回数の多いメソッドのJITを開始する。

遅くなる理由

JITのためにgccやclangが走っている最中は、そこにリソースが取られるからかある程度遅くなってしまうのだが、今回報告されたチケットの計測方法では計測中ほとんどコンパイルが走らない状態になっていた。

いくつかのマイクロベンチマークや、ピュアRubyのNESエミュレータでの性能を計測するOptcarrotというベンチマークではJITした方が明らかに速いのだが、先のチケットの計測方法だと遅くなってしまう。この理由は最初は不可解だったが、Optcarrotで普通にベンチを取ると20〜30メソッドくらいしかJITされないのに対し、このRailsでの計測は4000〜5000メソッドがJITされているという大きな違いがあった。

そもそも生成コードの最適化がRailsのコードに対して全然効いてなさそうなのも問題なのだが、最適化の余地が全くないようなただnil*1を返すだけのメソッドをたくさん定義して呼び出してみると、定義して呼び出すメソッドの数が多いほど遅くなることが発見された。

perfで計測してみると、遅くなっているのはicacheにヒットせずストールする時間が長くなっているのが原因のようで、それはメソッドごとに.soをdlopenしていることで生成コードが2MBおきに配置されてしまっておりメモリ局所性が悪いことが原因と結論づけた。*2

どうやって解決するか

解決策1: ELFオブジェクトを直接ロード

僕がこれに関する発表をRubyKaigiで行なってすぐ、shinhさんがELFオブジェクトを自力でロードしてくるパッチを作ってくださっていた。試してみると、ロードにかかる時間を遅くすることなく、40個くらいのメソッドを呼び出してもJIT無効相当の速度が出ていた。

一方で、shinhさん自身がブログで解説しているが、これを採用するとなると以下のような懸念点があった。

• ELFを使うOSでしか動かない
• (現状のパッチだと)ロードしたコードのデバッグ情報がgdbで出ない
• ローダ自体の保守やデバッグが大変そう
• 生成コードのメモリ管理が大変 (1GB固定アロケートか、遅くなるモードのみ実装されている)

そのためこれは直接採用はせず、以下の手法の評価にのみ利用した。

解決策2: 全てのメソッドを持つsoファイルを作ってdlopen

別々の.soになっているから問題が起きるわけなので、何らかの方法でコンパイル対象のメソッドが全て入ったsoを生成し、全てのコードをそこからロードしてくれば良いという話になる。コンパイル対象のメソッドの数が適当にハードコードした数に達したらまとめてコンパイルしてロードするようにしてみたら、実際速かった。

考慮したポイント

しかし、4000メソッド(計測に使われているRailsアプリのエンドポイントを叩いて放置するとコンパイルされる数)くらいをまとめてコンパイルすると普通に数分かかったりするので、この最適化をどのタイミングでどう実現するかは全く自明でない。

その戦略を考えるのが結構大変だったので、RubyKaigiの時点では定期的にまとめてコンパイルするだけのスレッドを新たに立てるつもりだったが、実装が複雑になるのでMJIT用のワーカースレッドは増やさないことにした。

その上で、短期間に少量のメソッドをコンパイルするOptcarrotでのパフォーマンスを維持しながら、数分かけて大量のメソッドをコンパイルするRailsでの性能を改善するため、以下の要素を考慮して設計することにした。

1. 初めてコンパイル+ロードされるまでの時間
2. soをまとめた後コンパイル済メソッドそれぞれのためにdlsym + ロックを取って生成コードを差し替える時間
3. 生成コードの数が溢れた時にメインスレッドで生成コードをdlcloseする時間
4. まとまったsoのコードに差し替えられるまでの時間
5. メモリ使用量
6. Ruby実行中 /tmp に持ち続けるファイルの数とサイズ

より上にある奴がより優先度が高く、下の方は(どうせ大した量使わないのもあり)どうでもいいと思っている。

最終的な構成

上記の1のためにワーカーはメソッドを1つずつコンパイルし、4のためにそのコンパイル結果の.oを /tmp に残し続けることにし、一方で一つ.oが増える度に一つのsoにしてロードし直してると2が線形に重くなって厳しいので4は多少犠牲にして一定回数おきにだけまとめてロードすることにして、そうするとある程度3や5が小さくて済み、6に関してはどうにかしたくなったら複数の.oファイル達を1つの.oにまとめれば良いだろう、という方針で作り始めた。

で、速度に関してはその方針でいいとして、メモリ使用量を考慮すると上記の図の"Sometimes"にあたる頻度がちょっと難しい。現状の実装では、呼び出されている最中の生成コードをそのフレームの外からVM実行に置き換えるOSR*3を実装できていないので、あるタイミングで1つの.soにまとめて生成したコードがどのフレームで使われているかを新たに管理するようにしないと、使われなくなったコードの破棄*4ができないのだが、それをやるとコードが結構複雑になる上に結局メモリ使用量も増えてしまう問題があり、あまりやりたくない。

それをサボる場合、頻度を上げれば上げるほどメモリリーク的な挙動になる*5わけなので、とりあえずキューイングされた全てのメソッドがコンパイルされた時 か--jit-max-cacheに達した時 だけ一つの.soにまとめてロードする処理をやる状態でコミットした。

OSRは他の最適化にどの道必要なので、長期的にはOSRを実装して任意のタイミングで古い生成コードを全て破棄できる状態にしようと思っている。

ベンチマーク

チケットの報告に使われているDiscourseというRailsアプリの、ウォームアップ*6後のGET / リクエスト100回で、以下のものを計測した。

• trunk: r64082
• trunkJIT: r64082 w/ --jit
• single-soJIT: r64082 + 今回のパッチ w/ --jit
• objfcnJIT:shinhさんのobjfcnをr64082からrebaseしたもの

レスポンスタイム(ms)

速度増加の割合

小さい値の方が良く、太字 が速くなっている箇所。

50%ileと60%ileは微妙だが、1000リクエストする計測でやり直すと50%ileや60%ileのtrunkとの差が1msとかだけになる*7ので、微妙に遅くなるか運が良いとちょっと速いという状態まで改善した。

感想

objfcnに比べ遜色ない効果が出せているし、Optcarrotも仕組み上今回の変更ではベンチマーク結果に影響はないので、生成コードのメモリ局所性の問題に関してはうまく解決できたと思う。Railsで遅くはならない状態にできたので、今度は速くしていくのをがんばりたい。