Ruby が YJIT でなんで速くなるのか？ Lazy Basic Block Versioning をサクッと理解してみた

technology

estie でソフトウェアエンジニアをしている徳永(@yTo_9)です。 estie では Ruby を書いたりTypeScriptを書いたりしています！

estie 夏のブログ祭りにかこつけて、せっかくなら普段は追わない部分だけど、気になっていたYJITなるものを深掘りしてみようと思い、「YJITがなぜRailsアプリケーションの高速化を実現できたのか」を調べてみたので紹介したいと思います。

「どうせ難しいんでしょ？」と思いながら調べてみたのですが、講演や論文の説明がわかりやすく、意外に概要を把握することは難しくありませんでした。

YJIT の核となっているのは Lazy Basic Block Versioning (LBBV) という手法で、これはRubyだけに限らず動的言語全般に適用可能な強力なアプローチであることがわかりました。

「あるタイプの条件分岐は、ほとんどの場合で片側しか実行されないため、実際には分岐を大幅にカットすることができる」という一見直感に反する話ですが、このブログを読んだ方がメカニズムのイメージを持って帰って頂けると嬉しいです。

実行時に変数の型は変化するか？

動的言語では実行時に、型検査・型推論が行われている

動的言語で演算子を書いた場合、処理系内では値の型検査とオーバーフロー検査が行われます。

プログラマーとしては 「 if で条件文を書いていないじゃん」と思うかもしれませんが、以下の例にあるようにi < n や s += i のような処理を入れるだけでも、 内部的にはi や n , s の型をチェックしたり、演算結果がオーバーフローしないかなどの条件分岐が何度も行われることになります。

仮にJSで書かれたこんなプログラムを考えてみましょう。

function sum(n) {
  for (var i=0, s=0; i<n; i++) {
    s += i;
  }
  return s;
}

図1: JavaScriptで0 ~ n-1 の総和(sum)を計算するコードの例

一見、条件分岐はほとんどないように見えますが、図2を見ていただけると、各行の演算の中では動的な型検査が走っており、その結果に応じて処理が分岐していく様子がイメージしてもらえるかと思います。（この図では… と書かれて省略されている部分にも分岐ロジックがあり、実際にはもっと複雑な分岐が広がっている）

図2: 左のプログラムを実行した際の内部的な処理の流れ(Control Flow Graph)。D (i < n) の処理をするために、B・Cで operand(i, n)の型を検査。H (s +=1) の処理をするために、F・Gでoperand(s, i)の型を検査している。https://arxiv.org/pdf/1411.0352v2.pdf の Figure.4 を引用

条件分岐の数に応じて倍々で制御の分岐フローが増えていきます。たったこれだけの処理でも、実際に出てくる条件分岐の膨大さをイメージしてもらえるでしょうか？

さて、ここからは話を簡単にするため、このプログラムの一部である i < n にのみ注目していきます。

図3: 図1 のうち for文の終了条件部分 (i < n) のみ考える

実行時の型検査・型推論はコストに合わない

動的言語の処理系ではそれぞれの変数の型検査の結果に応じた条件分岐を行います。JS処理系[]を表したこの概念図を見ると、ループ内で i < n のような処理を行うだけでも、型検査が何度も繰り返し走ることがわかります。(その上で n 回ループします)

図4: 図2における B → C → D の省略された分岐を拡大。動的な型推論ではループ全体でどの経路を通るかの推論に失敗しがちだったが、LBBVでは、実際には到達しない分岐についてはコンパイルしない先送り(Lazy)の方式を取ることで実行時に処理が到達する太字の部分のみのバージョン固有のコードをコンパイルすることに成功したという。

もし何らかの方法で、この型検査を省略することができればプログラムのあらゆる部分を高速化することができそうです。

静的型付け言語ではコンパイル時に型を推論することで、実行時の型検査を外すことができます。一方、RubyやJSのように「常にどんな型が入ってくるか実行時にならないとわからない」という前提がある動的言語では、型推論も実行時に行う必要があります。

実行時のパフォーマンスを上げるという目的上、込み入った型推論はできません。その上 eval のように実行してみないと型がわからない要素もあり、動的言語ではこの型推論によるアプローチは多くの場合に特定に失敗し、コストに見合った高速化を実現することが難しいとされていました。

LBBV の型検査削除へのアプローチ

今回紹介するLazy Basic Block Versioning(以下LBBVとする) では、「型推論を行わない」という逆転の発想によってこのジレンマを乗り越えました。以下ではこのアイデアを説明します。

実行時に判明した各変数の型ごとに Basic Block の型情報を確定させ、後段のブロックにもその情報を伝播させます（ここでいう「ブロック」とは、Control Flow Graph の各ノードに対応するものだと思ってもらえれば大丈夫です）。そのため各変数の型に応じたコードが都度生成されることになりますが、後段のブロックではその変数の型を確定情報として利用することができます。

ここからは少しJITの話がはいります。LBBVによって動的に収集した型情報をJITに利用させることで、処理系は型検査を省略した高速な実行コードを「少しずつ」生成していくことができます。これは投機的実行によりメソッド単位（ブロックから見るとはるかに大きい）でコンパイルしようとする従来のJITとは異なる戦略で、「なるべく不要なコンパイルをしない」ことでコンパイル時間の短縮とコード最適化をはかるという試みになっています。

図5:https://arxiv.org/pdf/1411.0352v2.pdf の Figure.6 を引用

ところで、こうやって「小さく」各分岐ごとにコンパイルをすすめるとして、実行中に変数の型が変わってしまったらどうなるでしょうか。変更された型情報に対して追加のコンパイルが必要になるため、2の冪乗で実行コードのコンパイル回数が膨れ上がり、巨大なオーバーヘッドがかかりそうに思えます。

実行時の型による分岐の90%以上で逆の分岐に到達しないプログラムばかりだった

ところが、現実のコードではほとんどの場合、一つの変数はそのライフタイムにおいて同一の型を持ち続けることが多いことがわかりました。

HiggsというOSSランタイムでJSのベンチマークを実行した場合のデータを見てみると、型検査の分岐は90％以上で片方の分岐しかたどらず、ブロックのコンパイルをその分岐に入るまで遅延した場合に2つ以上のバージョンのコードがコンパイルされたのは 6%に過ぎなかったとのことでした。

図6: Lazy Basic Block Versioning を行なった際に同一ブロックに対して作成されたバージョン数の比率。https://arxiv.org/pdf/1411.0352v2.pdf の Figure.10 を引用

これはどういうことでしょう？

考えてみると、変数は文法上再代入可能ですが、運用上で immutable であることも多いです。プロダクションコードでは、変数に後から異なる値を再代入することは良いパターンではなく、別の変数名をつけて持ち回ることの方が多いですよね。

その結果として、JITコードのコンパイル時間が短縮されるだけでなく、生成されるコードサイズにおいても大きなアドバンテージがあることがわかりました。論文中では Eager Basic Block Versioning を（すべての分岐に対してあり得るバージョン毎のコードをコンパイル)した場合、LBBVの数10倍のコードサイズになるケースが見受けられました。

まとめ

LBBV による高速化は、「組み合わせ上は存在するコードバージョンの多くが実行時にほとんど現れない」という現実のコードへの鋭い観察に基づいたシンプルなアプローチであり、あらゆる動的言語のパフォーマンスを向上させる汎用的な手法でした。

動的言語の書き味は変えないまま、型固有の最適化されたコードを生成することで、パフォーマンスが向上するのは非常に強力だなと思いました。生成するバージョン数の組み合わせ爆発が起こり得るため、ハードリミットを入れてはいるようですが、図6に見られるように現実世界のプログラムはかなり偏った特徴を持つというのは驚きでした。

Rubyにおいてはまだ最新バージョンに入ったばかりの仕組みですが、ベンチマークのみならず、Shopifyでは一部の本番環境でリリースされ、実際に高速化を達成したとのことでした。 estie pro でも YJIT オプションを有効化してみようと思ったので、実際に高速化が見られたらまたレポートします。

最後に

estie で働いていると、ありがたいことに Ruby Kaigi 登壇経験者やCRuby のコントリビューター、Rust の本を出版されている方など、言語自体にちゃんと向き合いながら活動している人に接する機会が増えました。その結果、自分もしばらく言語の処理系の話からは離れていましたが久しぶりに最新事情を追ってみようかと思い、今回 YJIT (もとい LBBV) を調べてみた次第です。

そんな estie での開発・仕事に興味を持った人はぜひ以下のリンクからカジュアル面談にお申し込みください。

hrmos.co

参考

© 2019- estie, inc.