JITコンパイルをサポートした新たなRuby処理系をRustとアセンブリで書いている話

この記事は[言語実装 Advent Calendar 2022](https://qiita.com/advent-calendar/2022/lang_dev) の24日目の記事です。

## はじめに

　monochromeです。今回はmonorubyという名前の高速なRuby処理系の紹介をします。

　ここ何年か、わたしはruruby[^ruruby]というRubyのインタプリタをRustで書いていて、Advent Calendarでもいくつか紹介記事を書きました。

[^ruruby]:https://github.com/sisshiki1969/ruruby
 
- [Rustでつくる（つくれるかもしれない）Ruby](https://qiita.com/sisshiki1969/items/3d25aa81a376eee2e7c2)
- [ruruby: RustでつくっているRuby](https://qiita.com/sisshiki1969/items/4d76e69545ca1c26ed48)
- [Ruby処理系自作入門](https://qiita.com/sisshiki1969/items/6f1852142f9f351e941a)
- [rurubyのガベージコレクタとアロケータ](https://qiita.com/sisshiki1969/items/a3de541d3d5b9eec0500)

　rurubyはRubyの言語機能のかなりの部分をカバーするところまで開発が進んだのですが、CRuby[^cruby]と比較するとやはり遅いのが自分としては不満でした。そこでJITコンパイル[^JIT]をサポートする高速な処理系を一から作ろう、ということで開発を始めたのがmonorubyです。対応するアーキテクチャは後述の理由でポータビリティに乏しいためx86-64・linuxのみ[^arch]サポートしています。

[^cruby]:Rubyが言語の名前でCRubyはCで書かれたRubyの処理系実装の一つにして標準実装でもある。正式名称ではない気がするが、本稿では具体的な処理系を指すときはCRubyと呼ぶ。
[^JIT]:プログラム実行時に、必要に応じて必要な部分を機械語にコンパイルする形式。
[^arch]:他のアーキは永遠にサポートできる気がしない…

 https://github.com/sisshiki1969/monoruby

　で、各種リテラルや制御構文、クラスとメソッドの定義、メソッド呼び出しにブロック、ローカル変数・インスタンス変数・グローバル変数を実装してそこそこRuby処理系っぽいところになったところでベンチマークを取った結果が下記です。

|                |3.2.0-preview3|     --yjit|    monoruby|monoruby --no-jit|
|:---------------|-------------:|----------:|-----------:|----------------:|
|quicksort       |      124.557k|   316.856k|    346.593k|         191.412k|
|app_fib         |         3.566|     14.378|      18.625|            4.859|
|tarai           |         2.922|     13.015|      14.117|            3.733|
|so_mandelbrot   |         0.626|      0.978|      16.751|            0.737|
|so_nbody        |         0.913|      1.507|       5.771|            0.749|
|app_aobench     |         0.028|      0.045|       0.101|            0.026|

　CRuby(3.2.0-preview3)およびCRubyの新しいJITコンパイラ・モジュールであるYJITを有効にした場合と、monoruby（JITあり・なし）の単位時間当たりの実行回数を[benchmark-driver.gem](https://qiita.com/k0kubun/items/82196240d3a7f987d1be) を用いて比較しています。数字が大きい方が速いです。ベンチマークの種類によりますが、YJITと比較しても高速、特に浮動小数点演算を多数実行するso_mandelbrotやso_nbody、app_aobenchでは数倍高速となっています。インタプリタもCRuby（JITなし）と遜色ない速度になっています。

　以下、設計と実装について紹介していきますが、その前に必要な基礎知識をおさらいしておきます。

## 基礎知識

　JITコンパイラは実行時に対象言語のプログラムを機械語へコンパイルするモジュールですが、一般的な実装ではインタプリタと並存して動きます。プログラムは抽象構文木を経由してバイトコード[^bytecode]に翻訳され、最初はインタプリタがバイトコードを実行しつつ最適化コンパイルに必要な情報を収集し、プログラムの一部分（メソッドやループ、特定の実行パスなど）が一定回数実行されるとJITコンパイラが起動して対象のバイトコードから機械語（JITコード）を生成し、以後はJITコードを実行するようになります[^1]。具体的な構成は実装により様々なのですが、例として下記にmonorubyの実行パイプラインを示します。実行される種々のコードが灰色の箱、オレンジの箱はコードを変換・実行する主体を表しています。

![無題のプレゼンテーション.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/332404/bcce3f8f-8aba-4760-2e68-1d0d3dc62287.png)

　一般に動的言語で高速で小さなJITコードを生成するにはいくつかの前提条件（制約）の元でコンパイルを行う必要があります。具体的にいうと「ローカル変数`x`は整数である」「この`a + b`という加算の左辺`a`の型は整数で右辺`b`の型は浮動小数点数である」「定数`A`の値は文字列"Ruby"である[^const]」などの情報です。これらの情報は実行時にならないと分からない[^ai]ため、インタプリタがバイトコードを実行しながらオブジェクトの型やメソッドの情報、定数の値などの情報（プロファイル）を収集する必要があります。動的言語の性質上、コンパイルに用いた制約が常に満たされるとは限らない[^correct]ため、制約が崩れた場合には実行継続が不可能となり、その場でインタプリタへ戻って（deoptimization）、新たなプロファイルを元にコンパイルし直します。
　このようにJITコンパイラの動作は複雑で、その性能には生成される機械語の品質やコンパイル速度の他、プロファイル情報を収集・保存するコストやdeoptimizationのコストも影響します。一般に高度な最適化を行うほどコンパイルに時間がかかり、またdeoptimizationが発生するリスクも高くなります。

[^1]:あくまで一般論であり、様々な実装があるので念のため。
[^bytecode]:ソースコードから処理系が生成する、インタプリタが解釈実行するための中間コード。雰囲気としては仮想的なCPUが実行する仮想的な機械語的なもの。「仮想的なCPU」は一般に仮想マシン、VMなどと呼ばれる。Java仮想マシン用の中間コードがバイトコードと呼ばれていたので、他の言語でも仮想マシン形式インタプリタ用の中間コードを慣用的にバイトコードと呼んでいる（ようである）。
[^ai]:抽象実行 abstract interpretation とかを使って静的に頑張って解析する方法もある。@miura1729 さんのRuby to C translator [MMC](https://www.youtube.com/watch?v=kr2RXLoiLNA)とか Julia (https://github.com/JuliaLang/julia/blob/master/base/compiler/types.jl) とか。
[^const]:ここで「定数の値は実行しなくてもわかるんじゃないのか？」という疑問を持つのが自然であるが、Rubyの定数は実行時に書き換え可能、しかも値を取得するには名前と値のマップのリンクを探索する必要がありコストが高い。なので前回得た値とタイムスタンプをセットにしてキャッシュしておくことで高速化する。
[^correct]:例えばメソッドが再定義された、定数の値が書き換えられた、など。

## 下準備

　処理系の話に行く前に機械語生成の具体的なやり方について説明します。Rustには既に機械語を生成する優れたライブラリがいくつかありますが、今回は自分の処理系用に使いやすいようにカスタマイズしたいので、monoasmというライブラリを自前で用意しています。これはx86-64専用で、proc macroというRustのマクロ機能を使ってアセンブリっぽい記法で実行時に機械語を吐くコードを書けます。

https://github.com/sisshiki1969/monoasm

　下記はmonoasmの使用例で、メソッド呼び出しで引き渡す引数をスタック上にセットする関数です。monoasm!マクロの内部にアセンブリ（っぽいもの）を記述しておくとコンパイル時に「機械語を生成してヒープに格納するRustコード」に変換され、実行時に実際の機械語が生成されます。

- 分岐や関数呼び出しの飛び先ラベルはあらかじめRustの変数として定義しておきます。
- 即値やレジスタ間接のオフセット指定などでは（）の中にRustの式を書けます。
- JITコンパイラでは複数のコードページを切り替えて出力できると便利なので、その機能もつけてあります。

```rust
impl Codegen {
    fn vm_set_arguments(&mut self, args: SlotId, len: u16, has_splat: bool) {
        // set arguments
        if len != 0 {
            let splat = self.splat;
            monoasm!(self.jit,
                lea r8, [rsp - (16 + OFFSET_ARG0)];
            );
            for i in 0..len {
                let next = self.jit.label();
                let reg = args + i;
                monoasm! {self.jit,
                    movq rax, [rbp - (conv(reg))];
                }
                if has_splat {
                    let no_splat = self.jit.label();
                    monoasm! {self.jit,
                        testq rax, 0b111;
                        jne  no_splat;
                        cmpw [rax + 2], (ObjKind::SPLAT);
                        jne  no_splat;
                        call splat;
                        jmp  next;
                    no_splat:
                    }
                }
                monoasm! {self.jit,
                    movq [r8], rax;
                    subq r8, 8;
                next:
                }
            }
        }
    }
}
```

　なお、パーサとアロケータ（ガベージコレクタ）はrurubyのものをそのまま流用しています。

## 設計

　monorubyの設計にあたり、いくつかの目標を立てました。

1. 構造がシンプルであること
2. インタプリタもそれなりに高速であること
3. JITコンパイルが高速であること
4. JITコード・インタプリタ間の遷移が高速に行えること

　これらの目標を考慮してバイトコード、スタックのレイアウト、メソッド呼び出しのABIの設計を行なっています。

#### バイトコード

　monorubyのインタプリタはレジスタマシンで、バイトコードは16バイト[^byte]の固定長です。特徴としては、命令コードやオペランドの他にインタプリタが収集するプロファイル情報もバイトコード内に書き込むようになっていることでしょうか[^profile]。例えばフィボナッチ関数のメソッドのバイトコードは下記のような感じになっています。％で示されているのはスタック上のスロット（レジスタと呼びます[^register]）番号で、%0は`self`です。先頭に`+`がついている命令は基本ブロックの開始を示します。

[^byte]:何がバイトコードなのかよく分からなくなってくるが、ネット上の風説では本家Javaバイトコードはオペコードが１バイトなのでバイトコードと呼ばれたらしい。monorubyもオペコードは１バイトなので問題なし。
[^profile]:これは手で作った機械語コードから読み取る／書き込むのが簡単かつ高速になるので良い方法ではあるが、マルチスレッド化すると色々と問題が生じる予感がする。難しいことは考えないことにする。
[^register]:Rubyオブジェクトを格納する8バイトのスタックフレーム内の領域、ということなのだがCPUのレジスタと一緒になってややこしい。ちなみに「スタック」という言葉もRubyのスタックとCPUが使うスタックが混ざってややこしい。

```text
+:00000 init_method reg:5 arg:1 pos:1 req:1 block:0 stack_offset:5
 :00001 _%2 = %1 < 3: i16
 :00002 condnotbr _%2 =>:00005
 :00003 %2 = 1: i32
 :00004 ret %2
+:00005 %3 = %1 - 1: i16
 :00006 %2 = %0.call fib(%3; 1)
 :00009 %4 = %1 - 2: i16
 :00010 %3 = %0.call fib(%4; 1)
 :00013 %2 = %2 + %3
 :00014 ret %2
```

　この中の`:00013 %2 = %2 + %3`はレジスタ同士を加算してレジスタに戻す命令（`ADD_RR`）で、`:00006 %2 = %0.call fib(%3; 1)`はレシーバが%0（`self`）、メソッド名が`fib`、%3から始まって１個の引数を与えてメソッド呼び出しを行い、結果を%2へ格納するメソッド呼び出し命令（`METHOD_CALL`）です。下記にこれらの命令のメモリレイアウトを示します。

![無題のプレゼンテーション (1).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/332404/88aac89c-7c93-ff01-ec98-77f38649fc18.png)

　基本的には前半8バイトが命令コード（赤い部分）＋オペランド（黄色）、後半8バイトがプロファイル情報（青）になっています。`ADD_RR`命令は命令コード（0xaa）で始まり、前半部には結果を格納するレジスタ番号（`%dst`）と左辺（`%lhs`）・右辺（`%rhs`）のレジスタ番号、後半は左辺・右辺の値のクラスIDが格納されます。
　その下にある`METHOD_CALL`は16バイトの命令３個分を使うメソッド呼び出し命令です。メソッド名、レシーバ`%rcv`・引数`%args`とその個数`len`・戻り値を返すレジスタ`%dst`などのオペランドの他、プロファイル情報としてレシーバのクラス`ClassId(rcv)`、関数ポインタ`CodePtr`やメタ情報（関数IDなど）、バイトコードのアドレス`PC`がキャッシュされています。これらの情報はJITコンパイラだけでなく、インタプリタ自身もインライン・メソッドキャッシュとして使用します。

#### スタックフレーム構造の単一化

　CRubyにはRuby上でのメソッド呼び出しに使用されるスタックがあり、メソッドを呼び出すとスタックを伸ばして制御情報やローカル変数領域を確保し（スタックフレーム）、メソッドから帰るときにはスタックを巻き戻します。一方、CRubyはＣ言語で書かれていますので、処理系内部の関数呼び出しや、Ｃで書かれたメソッドが呼ばれたときにはＣのスタック（CPUが使うマシンスタック）を消費します。要はRubyスタックとマシンスタックの２つがあり、これらは基本的には別ものですが、Rubyで書かれたメソッド→Cで書かれたメソッド→Rubyで書かれたメソッドのように処理が進むと両方のスタックを消費することになり、処理が煩雑になる原因になっていました。
　そこで、monorubyでは全てをマシンスタックにまとめる設計としました。これにより例外発生時などのスタックの巻き戻し処理やコルーチン（RubyではFiberクラス）の管理が単純化されます[^stack]。また、インタプリタやJITコード、Rustで書かれたメソッドも全て同一のスタックレイアウトを共有しているため、ループの途中でJITコンパイラを走らせてJITコードへ遷移したり(on stack replacement)、任意の場所でJITコードからインタプリタへ戻る（deoptimization）のも単にジャンプするだけ[^deopt]で済むので低コストです。

[^stack]:CRubyでは謎の黒魔術を使って実装されている。単純化される、と断言しているがmonorubyではraiseはあるがまだrescueできないし、Fiberもまだ実装していないので本当かどうかは未検証。
[^deopt]:実際には後述するCPUの浮動小数点数レジスタに移動されている変数があるので、これらをスタックに書き戻す必要がある。

#### メソッド呼び出しABIの単一化

　また、CRubyもrurubyもメソッドを呼び出すときにはメソッドの実体がRubyで書かれたもの（＝バイトコード）なのか、ＣなりRustなりで書かれた関数なのか、またはインスタンス変数のアクセサ[^accessor]なのかで条件分岐して別々の処理を行っていました。monorubyのインタプリタではこれらすべてが同じABIになっている（＝必要があれば処理系がラッパーを自動的に生成する）ので、呼び出し側は単に制御情報や引数をスタックにセットして指定のアドレスをcallすればよい仕組みになっています。インタプリタとJITコードも同一のABIで動きますので呼び出し側が考える必要はありません。
　下記にスタックの構造を示します。詳細は省きますが、`cfp`を辿っていくことで呼び出し元のメソッドを、`outer`を辿っていくことで現在のブロックの外側のブロックを辿れます。

[^accessor]:Rubyではオブジェクト`x`のインスタンス変数`@var`へのアクセスを`x.var = 42`のように書ける。（クラス定義内でアクセスを許可する必要がある）

 ```text
        +-------------+
 +0x08  | return addr |　　　　戻り先アドレス
        +-------------+
  rbp-> |  prev rbp   |　　　　前のスタックフレームの　rbp
        +-------------+
 -0x08  |  prev cfp   |　　　　前のフレームの　cfp　アドレス
        +-------------+
 -0x10  |    outer    |　　　　外側のフレームへのリンク
        +-------------+ 
 -0x18  |    meta     |　　　　メタ情報。このメソッドの関数IDなど。
        +-------------+ 
 -0x20  |    block    |　　　　このメソッドに渡されたブロックの情報
        +-------------+ 
 -0x28  |    self     |　　　　selfの値
        +-------------+ 
 -0x30  |    reg0     |　　　　引数、ローカル変数、一時的に使用される変数など
        +-------------+ 
        |      :      |
        +-------------+
  rsp-> |   reg(n-1)  |
        +-------------+
        |      :      |
 ```

#### インタプリタ

　上記のようにマシンスタックレイアウトやメソッド呼び出しABIを単一化するにはインタプリタが直接マシンスタックを触る必要があり、これをRustで書くのは困難です。また、monorubyのインタプリタは高速化のためにdirect threading[^DT]という手法を用いており、これもRustで実現するのは困難[^gcc]です。このため、monorubyではインタプリタをアセンブリで記述し、処理系起動時に動的に生成することにしました。もちろん複雑な処理はRustでヘルパー関数を書いて呼び出しています。保守性・移植性を完全に犠牲にしていますが、LuaJIT[^LuaJIT]という前例があります。
　その他、メソッドキャッシュや定数キャッシュなど、高速化の仕組みは一通り備えています。

[^gcc]:CRubyでは（使える環境では）GCC拡張を使っている。
[^DT]:下記の@κeenさんの記事がわかりやすいです。<br> https://keens.github.io/blog/2015/05/29/daiikkyuuraberuwomotanaigengoniokerudirect_threaded_vmnojissou/
[^LuaJIT]:https://luajit.org/
LuaJITは様々なCPU・OS上で動くことになっているが、インタプリタは作者が一人でCPUごとにアセンブリで書き下している。頭がおかしい。

## JITコンパイラの実装

　monorubyのJITコンパイラはメソッド単位・ループ単位の２つのモードを持ち、メソッドに関してはメソッド自身がカウンタを持っており、一定回数呼ばれるとJITコンパイラを起動して自分自身をコンパイルし、以後はJITコードへ飛ぶようにエントリを書き換えます。ループはインタプリタが回数をカウントして一定回数で同様にJITコンパイルを行い、そのままJITコードに遷移します。
　以下、実装上の工夫をいくつか挙げていきます。

#### Floatクラスの四則演算の高速化

　CRubyのインタプリタではFloatクラス（浮動小数点数クラス）の演算にあたって

1. 両辺のRubyオブジェクトをスタックからロードしてf64に変換（unboxing）
2. f64の演算
3. 結果をRubyオブジェクトに変換してスタックへ戻す（boxing）

　という操作が必要であり、本来の演算(2)よりも(1)(3)のオーバーヘッドが相対的に大きく、高速化の余地が大きいと考えました。monorubyのJITコンパイラはプロファイル情報からf64（Cのdouble）として計算に使用されるレジスタをCPUの浮動小数点数レジスタ（xmmレジスタ）に割り当て、簡易な制御フロー解析を行って不要な(1)(3)を削除するようにしました。通常コンパイラが行うレジスタ割り付けと基本的には同じですが、１つの変数をRubyオブジェクトとして使用する場合とf64として使用する場合の両方があるので、変換コストが最小限になるように工夫しています。
　結果として、冒頭のベンチマークで挙げたように浮動小数点演算を繰り返し実行するso_mandelbrotやso_nbody、app_aobenchで顕著な高速化が得られています[^float]。

[^float]:Rubyの浮動小数点数演算が速くなって嬉しい人がどのくらいいるかは不明。そういうことは考えない。

#### インスタンス変数アクセスの最適化

　Rubyのインスタンス変数は、オブジェクトに紐づけられた名前の付いたプロパティ（値はRubyオブジェクト）です。素朴にハッシュマップなどで実装すると遅いので、monorubyではオブジェクトには可変長配列（下図のVarTable）を持たせて値を格納しています。一方、クラス情報の構造体にインスタンス変数名とインデックスの対応マップ（下図のInstanceVarMap）を持たせてあります。これだとInstanceVarMapを変数名をキーにして引いてインデックスを得て、さらにVarTableにインデックスでアクセスすることになるので二度手間ですが、インデックスをキャッシュしておくことで２回目以降は高速になります。（以上は下図の上半分、クラスCの説明）
　さらに、Objectクラスの派生クラスではRubyオブジェクトのメモリ表現（８バイト）に空きがあるので、この中にインスタンス変数の値を埋め込んでしまうことができます。これを用いるとさらにアクセスが高速化する上に可変長配列をアロケートする必要がなくなるのと、JITコードではアセンブリで直書き（mov命令一発）できるので超高速になります[^instance]。（下図の下半分、クラスD）

[^instance]:CRubyにも同様の最適化がなされており、さらに最近Shapeという仕組みが入った。

![2022アドカレ.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/332404/e66fc314-e0ef-2608-7e07-3b06bddfa7e9.png)

以下、ベンチマークです。

|                     |3.2.0-preview3|     --yjit|       monoruby| monoruby --no-jit|
|:--------------------|-------------:|----------:|--------------:|-----------------:|
|vm_ivar_get          |        11.958|     26.649|         67.922|            12.395|
|vm_ivar_set          |      123.794M|   118.631M|         1.530G|           62.200M|
|vm_attr_ivar         |       57.361M|    57.916M|        68.746M|           48.746M|
|vm_attr_ivar_set     |       51.297M|    50.350M|       305.655M|           36.901M|

#### 小さなメソッドのインライン化

　数命令の機械語で済むような単純な計算のみを行うメソッドは直接JITコードの中に埋め込んでしまうことで高速化を図っています。メソッド呼び出し・復帰のコストを省けるのが大きいですが、特に浮動小数点数関係の単純な演算メソッド（整数から浮動小数点数への変換`Integer#to_f`や機械語１命令で計算できる平方根計算`Math.#sqrt`など）ではRubyオブジェクトとf64の変換も削減できて大きな高速化が得られます。`Math.#cos`など、正味の処理自体に一定時間かかる演算では相対的に恩恵は少なくなります。

|             |3.2.0-preview3|        --yjit|       monoruby|monoruby --no-jit|
|:------------|-------------:|-------------:|--------------:|----------------:|
|to_f         |       13.813M|       13.805M|       320.955M|          18.606M|
|math_sqrt    |        9.860M|        9.670M|        51.387M|          11.530M|
|math_sin     |        5.902M|        6.175M|        14.793M|           6.148M|
|math_cos     |        6.068M|        6.047M|        16.028M|           6.228M|

#### プロファイルのない命令のJITコンパイル

　インタプリタが実行時にプロファイルを収集するわけですが、JITコンパイラが起動した時点で全ての命令が実行済み（＝プロファイル収集済み）とは限りません。プロファイルのない状態でコンパイルすることは一応可能ですが、最適化の程度は当然劣ります。
　例えばフィボナッチ関数のスクリプト

```ruby
def fib n
  if n < 3
    1
  else
    fib(n-1) + fib(n-2)  # <- ココ
  end
end
```

を考えてみると、５行目の１個目の`fib()`で自分自身が再帰的に呼ばれるため、n<3になるまでは２個目の`fib()`は実行されず、プロファイルも収集されません。この状態でコンパイルすると遅いコードが生成されます。もっと一般的なケースでは条件分岐でほとんど分岐しない枝でこういう状態が起きます。
　monorubyのJITコンパイラでは、プロファイルのないバイトコードに対してはカウンタとともにインタプリタへ戻るコードを生成しておき[^trace]、一定回数以上通過した時点で（つまりプロファイルが収集された時点で）メソッドないしループ全体を再度コンパイルするようになっています。次回`fib()`が呼ばれる時には、より最適化された新しいJITコードが呼ばれるわけです。

[^trace]:monorubyは一応 method JITを謳っているが、こういう所は tracing JITっぽい。YJITもまだ一度も分岐していない枝ではコンパイルを止めてスタブコードだけを仕込んでおいて、分岐した時点でプロファイルを収集しつつその先のコンパイルを行うようになっている。

## 今後の展望

　ここまでで力尽きたので、Rubyの言語機能を充実させたい、とかフロー解析をもっとちゃんとしたものにして高度な最適化を適用できるようにしたい、とか呟きつつ、メリークリスマス。みなさま良いお年を。