対象のコード

独自言語ですが、まあ説明は不要でしょう。
再帰関数になっているfib関数を高速化します。

fun fib(n) do
  if (n == 0) do
    return 1;
  end
  if (n == 1) do
    return 1;
  end

  return fib(n - 1) + fib(n - 2);
end

インタープリターを作る

まずは普通にインタープリターを作ります。今回はZig言語を使ってみることにしました。
(Zig歴がまだ3日くらいの頃に書いているので、コードは拙いです)
中身はただの、LexerとParserとEvaluatorを組み合わせただけのものです。

そして書いたインタープリターのパフォーマンスを計測してみます。hyperfineを使ってみました。

該当のcommit

# fib(20)
hyperfine --warmup 3 "zig build run -Doptimize=ReleaseSafe"
Benchmark 1: zig build run -Doptimize=ReleaseSafe
  Time (mean ± σ):      2.584 s ±  0.095 s    [User: 0.060 s, System: 2.536 s]
  Range (min … max):    2.414 s …  2.729 s    10 runs

概ね2.5秒くらいかかっているようです。
比較のために、Rustで同じコードを書いて計測してみます。

# fib(20)
hyperfine --warmup 3 "cargo run --release"
Benchmark 1: cargo run --release
  Time (mean ± σ):     185.6 ms ±   3.3 ms    [User: 31.7 ms, System: 9.3 ms]
  Range (min … max):   180.2 ms … 192.3 ms    15 runs

Rustの方が10倍以上早いですね。
圧倒的です。

VMへのコンパイルを行う

なぜ先ほどのインタープリターはそこまで遅いのでしょうか。
それはASTを舐めているからで、これを直列の命令列にすることで高速化が図れます。

内部的に関数が呼ばれている回数を記録しておき、一定回数を超えたらVM上の命令列にコンパイルしてそれをキャッシュし、以降のその関数呼び出しではVM上の実行に切り替えるという手段を取ります。

命令列は単純なスタックマシンとし、簡単な演算とpush,pop等のスタックの操作とcall,retあたりを実装しておきます。また、ここではcall,retを実装するにあたり内部的な呼び出し規約を作っておきました。
VMの実行は命令列と、スタックと、base pointerを入れるレジスタと、ローカル変数を名前で解決できるようにする辞書を渡して実行を開始します。

該当のcommit

# fib(20)
hyperfine --warmup 3 "zig build run -Doptimize=ReleaseSafe"
Benchmark 1: zig build run -Doptimize=ReleaseSafe
  Time (mean ± σ):     375.8 ms ±   3.5 ms    [User: 347.4 ms, System: 45.3 ms]
  Range (min … max):   372.5 ms … 381.8 ms    10 runs

かなりの高速化ですね。6~7倍くらい早くなりました。

少しの最適化 + Arm64上のネイティブ実行

さてここからArm64上でネイティブに実行できるようにしていきます。
マシン語に落とす前には2つほど解決しておかなければならない問題があります。1つはローカル変数の解決で、もう1つはラベルの解決です。

ローカル変数は、get_local("hoge") のような命令を、名前ではなくアドレスか何かで解決するようにすることです。私は今回はベースポインタをレジスタに入れる規約のもとでVMを動かしているので、ベースポインタからの相対オフセットに変換し get_local(-3) のような命令に変換する機構をいれました。
ラベルの解決は単純に命令列のindexを直接指定しました。これでjump命令を内部的にただのip(インストラクションポインタ)の代入に置き換えられます。

それができたら、いよいよArm64のマシン語をバリバリ書くフェーズになります。
mmapでメモリ領域を確保し、ここにArm64のマシン語を書き込んで最後に関数ポインタにして呼び出すような仕組みにしました。
JITコンパイラには以下の資料を大変参考にさせていただきました。

(↓なんとZigでしかもArm64という、奇跡的に構成が同じだったため参考にできる部分が多かったです)

https://github.com/srijan-paul/tinyjit

(↓Arm64初心者なので、入門に良かったです)

(↓アセンブリをマシン語に変換したり、逆をするのに使いました)

https://armconverter.com/

JITコンパイラは内部的にはVMの命令列を受け取って、関数ポインタを返すような構成です。
ここで関数ポインタは上記のVMのexecuteを行う関数とほぼ同じ構成をしており、スタックへのポインタ、ベースポインタを入れるレジスタへのポインタを受け取るような関数です。

このようにすることでVMとインターフェースがほぼ同じになるため、挙動が変わらないことをテストで簡単に確認できます。また、直接スタックを操作する機会が減ったり呼び出し規約で考えることも減るので、どちらかというとCPUアーキテクチャ依存な知識が減るので少し実装はしやすくなるかと思います。
(とはいえVM側で呼び出し規約を定めたりはどうしても必要なので、どっちが良いかという話はあります。私はVMの設計が初めてではないのでこっちの方が楽でしたが)

これができたら、先ほどのVM実行の部分をネイティブ実行の方にすり替えてまた計測してみます。

該当のcommit

hyperfine --warmup 3 "zig build run -Doptimize=ReleaseSafe"
Benchmark 1: zig build run -Doptimize=ReleaseSafe
  Time (mean ± σ):      54.7 ms ±   1.4 ms    [User: 35.1 ms, System: 36.7 ms]
  Range (min … max):    52.4 ms …  59.0 ms    50 runs

なんと7倍くらい早くなりました。
最初の実装から比べるとほぼ50倍なので、素晴らしいですね！

fib(30)以降の話

さて、参考用に出したRustの実行結果をみておやと思った人もいるかもしれません。
Rustより速いとは何事か？！と思うのですが、実際はそこまで簡単な話ではありません。

fibの引数を増やしてみると、以下のようになります。

Rust

# fib(30)
hyperfine --warmup 3 "cargo run --release"
Benchmark 1: cargo run --release
  Time (mean ± σ):     186.1 ms ±   1.6 ms    [User: 34.1 ms, System: 8.7 ms]
  Range (min … max):   182.1 ms … 188.9 ms    16 runs

# fib(35)
hyperfine --warmup 3 "cargo run --release"
Benchmark 1: cargo run --release
  Time (mean ± σ):     208.4 ms ±   3.1 ms    [User: 59.0 ms, System: 7.5 ms]
  Range (min … max):   202.4 ms … 213.9 ms    14 runs

# fib(40)
hyperfine --warmup 3 "cargo run --release"
Benchmark 1: cargo run --release
  Time (mean ± σ):     496.4 ms ±   4.8 ms    [User: 315.4 ms, System: 9.1 ms]
  Range (min … max):   488.4 ms … 500.9 ms    10 runs

今回の処理系

# fib(30)
hyperfine --warmup 3 "zig build run -Doptimize=ReleaseSafe"
Benchmark 1: zig build run -Doptimize=ReleaseSafe
  Time (mean ± σ):     219.2 ms ±   2.1 ms    [User: 195.4 ms, System: 40.2 ms]
  Range (min … max):   216.7 ms … 224.6 ms    13 runs

# fib(35)
hyperfine --warmup 3 "zig build run -Doptimize=ReleaseSafe"
Benchmark 1: zig build run -Doptimize=ReleaseSafe
  Time (mean ± σ):      1.860 s ±  0.004 s    [User: 1.814 s, System: 0.061 s]
  Range (min … max):    1.854 s …  1.866 s    10 runs

# fib(40)
30秒くらいかかっていてhyperfineは諦めました。

ということで、引数を増やすと加速的に遅くなります。

そもそも今回のfibは再帰関数で再帰が深くなると加速的にスタックを消費するようなコードなので遅くなるのは当然なのですが、Rustの方はそれがかなり抑えられているように見えます。この辺はRust(かLLVMかも)の最適化の頑張りによるものかなと思います。
ということで言語同士のパフォーマンスは単純に比較できるようなものではありません。とはいえ、JITコンパイルを(うまく)すると爆速になるという体験ができました。

ハマったところ

一部ハマったところとしては、関数(というかサブルーチン)の最初でリンクレジスタ(x30)の値をスタックに退避し、RETの直前でそれを読み出ししてあげる必要があることです。
これはBL命令を発行してRETで後で戻ってくるために、BL命令時にリンクレジスタに今いる位置を保存してしまうために、2回目のRET時に戻り先がわからなくなってしまうためです。またArm64ではスタックのアラインメントを16バイトごとにしないといけないらしく、stp命令を使うなどをしました。

これらの話は以下のページの「サブルーチン呼び出しの例」に記載がありますが、最初やっていた時はあまり内容を理解していなかったのでcall命令周りが正しく動作せず困りました。

終わりに

今回は関数ごとにネイティブコードにコンパイルするようにしましたが、命令列のトレースをとってtracing JITにするのをやってみたいな、と思っています。

またかなりどうでもいいですが、JITコンパイルって名前的には実行時コンパイラで、実行時にVMに変換するのもJITコンパイルじゃないの？と思ったのでArm64上で動くネイティブコードに変換するのだけをJITコンパイルって言うのなんか不思議な感じだなあと思いました。でも一般的にはJITコンパイルって言うとネイティブコードの生成までを指すことが多い気がします。

Arm64どころかアセンブリ自体も初心者なので最初はセグフォがいっぱい出たりしてかなり苦労しましたが、どうにか動くようになって良かったです。

Discussion