RubyKaigi 2024 - YJIT Makes Rails 1.7x Faster (Day3)

ogp

YJIT Makes Rails 1.7x Faster

こんにちは、2024年に新卒で入社し、ジョブハウスでバックエンドエンジニアをしているnozomemeinです。
本記事では、3日目のTakashi Kokubun(@k0kubun)さんによるセッション、YJIT Makes Rails 1.7x Faster について紹介させていただきます。

YJITとは？

講演の内容に入る前に、YJITの歴史と特徴についてまとめます。

YJITの概要

YJITは「Yet Another Ruby JIT」の略で、Shopifyで開発され、Ruby 3.1から導入されたJITコンパイラです。

Ruby2.6から、MJITと呼ばれるJITコンパイラが存在していましたが、YJITはそれをさらに進化させたものです。

YJITはJITコード（中間表現 or IR）を使用し、機械語に変換して実行します。

この方法により、パフォーマンスの大幅な向上が期待できるとともに、中間表現を介して機械語に直接変換できるので、実行環境を問わないことが特徴です。

YJITの歴史

Ruby 1.8では、パーサが出力するAST（抽象構文木）を直接Cのプログラムが実行する形式でした。

Ruby 1.9でYARVが導入され、ASTをYARVがYARVバイトコード（ISeq）にコンパイルし、これをYARV仮想マシンが実行するようになりました。

Ruby 2.6ではMJITが追加され、実行時に一部の処理をC言語にコンパイルすることで高速化を図りましたが、Railsアプリケーションなどを使った本番環境では、まだまだ十分に使える状態ではありませんでした。

そこで、Ruby 3.1でYJITが登場し、さらに高性能な実行環境が提供されるようになりました。YJITではYARVのバイトコードのうち、変換可能なコードを部分的に機械語に置き換えることで高速化を図ります。

YJITの登場により、本番環境でも十分に使えるJITコンパイラが整備され、Railsアプリケーションでも高いパフォーマンスを発揮することが可能となりました。

Ruby 3.2からはYJITの実装がCからRustに置き換えられ、YJITのメモリ使用量をより削減できるようになったとともに、YJITが機械語にコンパイルする前に、中間表現としてJITコードを挟むようになったので、x86_64だけでなく、ARM64にも対応できるようになりました。

Ruby 3.3からは、Rubyで書かれたRJITがKokubunさんの手によって開発され、MJITを置き換える形で登場しています。

github.com

YJITの技術的特徴

Lazy Basic Block Versioning

YJITの核となる技術は「Lazy Basic Block Versioning（LBBV）」です。

Shopify YJITチームのメンバーであるMaxime Chevalier-Boisvert(@maximecb)さんによって考案され、彼女の博士時代の研究がもとになっているそうです。(原著論文)

LBBVは、実行時に型情報を収集し、それをもとにJITコンパイラにフィードバックします。これにより、不必要な型検査を省略でき、より効率的にコンパイルすることが可能になります。

また、LBBVでは型情報の収集を事前に行うのではなく、型情報が条件分岐で確定してから収集する形式(可能な限り遅く、lazyに実行する戦略)をとります。これにより、JITコンパイラによって生成されるコードが不必要に増えることを防いでくれます。

メモリ効率の向上

YJITはメモリ効率にも優れています。

従来のJITコンパイラはピークパフォーマンスに注目していましたが、メモリ使用量が多かったり、パフォーマンスが出るまでに時間がかかるという問題点がありました。

現実問題として、本番環境ではサーバープロセスやキャッシュにメモリを使うため、JITに割り当てられるメモリ量が限られています。

また、デプロイ頻度が多いと、立ち上がり時間がボトルネックになるという問題がShopifyでも指摘されていたそうです。

そこで、YJITはウォームアップ時間やメモリ使用量にも配慮された設計になっています。これにより、実際の本番環境でも高いパフォーマンスを発揮できます。

YJITのアップデート

RubyKaigi 2024で紹介された、Ruby 3.3からのYJITの主なアップデートは以下の通りです。

Method Call Fallbacks
Exception Handlers
Register Allocator for Stack Values
Method Inlining for Ruby/C

また講演発表資料はこちらからご覧いただけます。
speakerdeck.com

Method Call Fallbacks

これまでのYJITでは、YARVバイトコードをJITコードにコンパイルする際に、対応するJITコードがない場合、それ以降のJITコンパイル処理を停止し、残りのYARVバイトコードはコンパイルされずにYARV仮想マシンによって実行されていました。

結果として、後続のYARVバイトコードが仮にJITコードにコンパイル可能であってもスキップされていました。これではせっかく残りのバイトコードがJITコンパイル可能であっても、YJITの機能を活かせなくなってしまいます。

この問題を解決するために、Ruby 3.3からは対応するJITコードがない場合でも、JITコード側からYARVのC関数を呼び出せるようにすることで（フォールバックを可能にすることで）、コンパイル処理が中断されなくなりました。

このおかげで、よりYJITの機能を活かした実行が可能になり、パフォーマンスが改善されるようになりました。

Exception Handlers

CRubyでは、breakやreturn文は内部的にはYARVバイトコードでthrow命令として扱われています（returnもbreakも内部表現では同じというのは面白いですよね）。

throw命令を実行すると、制御フレームスタックを順に辿って補足テーブルを探し、補足テーブルでキャッチされると、補足テーブルが持っているcontインデックスを使って戻り先の命令シーケンスにjumpするような処理がCRuby内部では行われています。

// iseq.h (Ruby 3.3.1)
// 補足テーブルの実体
struct iseq_catch_table_entry {
    enum rb_catch_type type;
    rb_iseq_t *iseq;

    unsigned int start;
    unsigned int end;
    unsigned int cont;
    unsigned int sp;
};

これまでは、このthrow命令がJITコードにコンパイルできませんでしたが、Ruby 3.3からはJITコードにコンパイルできるようになりました。

このおかげで、「Method Call Fallbacks」のセクションでも述べたように、JITコンパイル処理の最適化を行うことができるようになりました。

Register Allocator for Stack Values

これまでのYJITでは、putobjectのようなYARVバイトコードをコンパイルして実行した場合、VMスタックを内部的に表現する際にメモリに直接書き込んでいました。

しかし、メモリに直接書き込むよりもレジスタに直接書き込む方が圧倒的に速いです。レジスタへの書き込みが1サイクルに対して、メモリは80サイクル程度かかるため、オーダーがおよそ2桁違います。

そこで、Ruby 3.3からはレジスタに直接書き込むようになりました。

これにより、メモリアクセスのオーバーヘッドが削減され、パフォーマンスが向上しました。

さらに、Ruby 3.4からはローカル変数もレジスタに書き込むように開発が進められています。これにより、さらなるパフォーマンス改善が期待でき、非常に楽しみです。

Method Inlining for Ruby/C

Ruby 3.3では、RubyおよびCメソッドのインライン化が進みました。

これにより、特定のメソッド呼び出しのパフォーマンスが向上しました。

Rubyメソッドのインライン化

シングルラインメソッドや単純な戻り値を返すメソッドがインライン化されました(e.g. 1, 2, 3..., true, false, nil, :symbol)。

これにより、JITコンパイルによって出力される機械語が数十行から一行に集約されるので、メソッド呼び出しのオーバーヘッドが減少し、特にRailsのblank?やpresent?などの基本的なメソッドがインライン化されることで、パフォーマンスが向上しました。

Ruby 3.4からローカル変数やselfもinline化できるように進めているようです。

C関数メソッドのインライン化

Cで書かれた関数の一部もインライン化され、アセンブリのバイトコードに直接対応するようになりました。 これにより、C関数の実行が高速化されましたが、例外処理などの複雑なケースへの対応が求められるようになったそうです。

下記に示しているのは、String#empty?メソッドの最適化の例です。事前に対応するメソッドを登録しておいて、コード生成関数を呼び出すことで、メソッドの最適化を行っています。

// String#empty?の例
// yjit/src/codegen.rs (Ruby 3.3.1)

// メソッドに対応するコード生成関数を登録しておく
pub fn yjit_reg_method_codegen_fns() {
  unsafe {
    // 略
    yjit_reg_method(rb_cString, "empty?", jit_rb_str_empty_p);
    // 略
  }
}

// コード生成関数
// Codegen for rb_str_empty_p()
fn jit_rb_str_empty_p(
    _jit: &mut JITState,
    asm: &mut Assembler,
    _ocb: &mut OutlinedCb,
    _ci: *const rb_callinfo,
    _cme: *const rb_callable_method_entry_t,
    _block: Option<BlockHandler>,
    _argc: i32,
    _known_recv_class: Option<VALUE>,
) -> bool {
    let recv_opnd = asm.stack_pop(1);

    asm_comment!(asm, "get string length");
    let str_len_opnd = Opnd::mem(
        std::os::raw::c_long::BITS as u8,
        asm.load(recv_opnd),
        RUBY_OFFSET_RSTRING_LEN as i32,
    );

    asm.cmp(str_len_opnd, Opnd::UImm(0));
    let string_empty = asm.csel_e(Qtrue.into(), Qfalse.into());
    let out_opnd = asm.stack_push(Type::UnknownImm);
    asm.mov(out_opnd, string_empty);

    return true;
}