Universal MachineのJITコンパイラーを書いた／x86_64編

前にこういう記事を書きました：

Universal MachineのJITコンパイラーを書いた

かいつまんで説明しておくと、Universal Machineというのは2006年のICFP Programming Contestで挑戦者が実装する必要のあるVMです。これは速ければ速いほどいいので、JITコンパイルに挑戦したくなります。2年前の記事ではAArch64向けのJITコンパイラーを書きました（私の家の最速マシンがM1 Macだったので）。

時は流れ、私の家の最速のマシンはM1 MacからRyzen搭載ミニPCへ交代しました。なので、x86_64版のJITコンパイラーも書きたくなりました。

今回対応したのはUnix系で、Windows版は未実装です。

x86機械語

私としてはx86のアセンブリー言語はある程度触ったことがありますが、今回やりたいのはJITコンパイルです。JITコンパイルでは as コマンド的なアセンブラーには頼れません（コマンド呼び出しのオーバーヘッドが許容できる場合はともかく）。普通はXbyakのようなライブラリーとして使えるJITアセンブラーを使うか、自力で機械語を書くか、ということになります。

今回必要な命令はそこまで多くはないと思われること、実行時に命令列を書き換える都合上命令の長さ（バイト数）も意識したいことを踏まえて、今回は自力で機械語を書くことにしました。ただ、x86の機械語は複雑で、ソースコード中にガリガリビット演算で書き込んでいくと見通しが悪そうなので、それぞれの命令に対応する機械語を書き込む処理をCの関数として用意することにしました。簡易的なJITアセンブラーと呼べるかもしれません。

例えば、 PUSH r64 型の機械語を uint8_t *instr に書き込む関数は

uint8_t *push_r64(uint8_t *instr, enum reg64 r);

で、機械語を1バイトまたは2バイト書き込み、その分進んだ新しい instr の値を返します。

もっと複雑なものだと、 MOV r64, QWORD PTR [base + scale*index + disp] 型の機械語を書き込む関数は

uint8_t *mov_r64_qword_ptr(uint8_t *instr, enum reg64 dst, enum reg64 base, enum address_scale scale, enum reg64 index, int32_t disp);

となります。

enum reg64 は汎用レジスターを表す型で、

enum reg64 {
    rax,
    rcx,
    rdx,
    rbx,
    rsp,
    rbp,
    rsi,
    rdi,
    r8,
    r9,
    r10,
    r11,
    r12,
    r13,
    r14,
    r15
};

と定義しています。32ビットの汎用レジスターを指す enum reg32 も定義しています。なんで最初の4つがA, C, D, Bの順番なのかは知りません。

肝心のx86機械語のエンコードは、Intel SDMのそれぞれの命令のところを見ても最初は正直よくわかりませんでした。Armとは大違いです。なので、ググって出てくるページも参考にしました。

ググって出てくるオンラインのアセンブラーも参考にしました。

苦労の甲斐あってIntel SDMの読み解き方もある程度わかるようになってきたので、x86の機械語についてはまたそのうち記事を書きたいです。

コードを実行時に生成するためのおまじない

Apple Silicon MacだとW^X制限とか命令キャッシュのflushとかの「おまじない」がありましたが、x86だとそういうのは基本的に必要ありません。Macでも PROT_READ | PROT_WRITE | PROT_EXEC を同時に指定できますし、命令とデータは同一のキャッシュに載るらしいです。

ただ、だからと言ってGCC/Clangの __builtin___clear_cache が完全な無意味かというとそうでもなくて、ある種のコンパイラーの最適化を抑止する役割があるらしいです。memory barrierってやつです。参考：

https://stackoverflow.com/questions/35741814/how-does-builtin-clear-cache-work#comment85964322_35741869

レジスターの使い方

アセンブリー言語や機械語でプログラミングする上で一番楽しいのは、レジスターの使い方を考える部分だと思います。逆に、スタックの使い方を考えるのは色々制約があったりスタックポインターの値を考慮しないといけなかったりして面倒くさいです。

Universal Machineには8個のレジスターがあります。インタープリターに対するJITコンパイラーの優位点は、VMのレジスターをCPUのレジスターに割り当てることができることです。もちろん、32ビットx86のように汎用レジスターが少ないとそれは難しいでしょうが、x86_64には16個の汎用レジスターがあるので、なんとかいけそうです。

汎用レジスターが16個あると言っても、いくつかは役割が決まっていたり（rsp, rbp）、関数呼び出しの引数・返り値で使う必要があったり（rdi, rsi, rdx, rcx, rax, rdx）、整数の乗除算で特別な役割があったりします（rax, rdx）。なので、VMのレジスターはそれらを避けるように割り当てたいです。

また、VMのレジスター8個以外にも、レジスターに載せて持っておきたい値が2つほどあります（配列の配列、program counterと機械語を結びつけてジャンプに使う配列）。これらの割り当てを考える必要があります。

今回考えたレジスターの使い方は以下の通りです：

rax: 関数呼び出しの最初の返り値／整数の乗除算／テンポラリー
rcx: 関数呼び出しの第4引数／ジャンプテーブル
rdx: 関数呼び出しの第3引数／関数呼び出しの第2返り値／整数の乗除算（上位）
rbx: 配列（callee-save）
rsp
rbp
rsi: 関数呼び出しの第2引数
rdi: 関数呼び出しの第1引数
r8: VMのレジスター0
r9: VMのレジスター1
r10: VMのレジスター2
r11: VMのレジスター3
r12: VMのレジスター4（callee-save）
r13: VMのレジスター5（callee-save）
r14: VMのレジスター6（callee-save）
r15: VMのレジスター7（callee-save）

callee-saveのレジスターはJITコンパイルした関数に入る時に退避して、それ以外の必要なレジスターはコンパイルした関数から補助関数を呼び出す時に退避します。

JITコンパイルで作られる関数の概要

JITコンパイルで作る関数のプロトタイプは

uint32_t func(struct array **arrays, void **jumptable, uint32_t registers[8], uint32_t initial_pc);
/* ただし
struct array {
    uint32_t length;
    uint32_t data[];
};
*/

です。

arrays は配列のID（32ビット整数）から配列の実体を対応させる配列です。jumptable はprogram counter pc にその命令に対応する機械語のアドレス jumptable[pc] を対応させる配列で、ジャンプとコード書き換えに使います。registers はVMのレジスターの値の受け渡しに使います。initial_pc は最初に実行する命令のprogram counterです。

関数の返り値は、JITで実行できず中断した命令のprogram counterです。AArch64版では機械語のアドレスを返していましたが、今回は素直にprogram counterとしました。中断した命令はC言語のswitch文で実装されたVMにより実行します。「JITで実行しない命令」というのは、HALT命令や、自己書き換えの対象となった命令を含みます。

関数のプロローグやエピローグは、入力に依存しないので、古典的なアセンブラーでアセンブルしてCの配列としてコードに埋め込み、memcpy で貼り付けています。

各命令のコンパイル

各opcodeに対応するx86_64の命令列を見ていきます。アドレス計算に使う部分以外では、A, B, Cは32ビットレジスターを指します（r8d – r15d）。

まず、JITコンパイルされたコードを脱出してインタープリターに処理を任せる場合は、現在のprogram counterを pc として

return pc;

となります。アセンブリー言語で書けば

mov eax, <pc>; imm32
jmp L_epilogue; rel32

となります。レジスターを復帰させる作業があるので、 ret ではなく、それ用のエピローグにジャンプします。この2命令は機械語で10バイトなので、（書き換えを考慮して）他の命令が10バイト未満だった場合はNOPを埋めて1命令あたり10バイト以上となるようにします。

0 Conditional Moveは

if (C != 0) {
    A = B;
}

みたいなことをします。アセンブリー言語で書けば

#if A == B
    ; no-op
#else
    test C, C
    cmovne A, B
#endif

となります。

1 Array Indexは

A = arrays[B]->data[C];

みたいなことをします。アセンブリー言語で書けば

mov rax, qword ptr [rbx + 8*B]; rbx = arrays
mov A, dword ptr [rax + 4*C + 4]

となります。

2 Array Amendmentは

arrays[A]->data[B] = C;
if (B == 0) {
    um_modify_0(B, C);
}

みたいなことをします。um_modify_0 はJITコンパイルされたコードにパッチを当てる（「JITコンパイルされたコードを書き換えてインタープリターへ脱出するコード」に置き換える）関数です。プロトタイプは

void um_modify_0(uint32_t b, uint32_t c);

です。

アセンブリー言語で書けば

mov rax, qword ptr [rbx, 8*A]; rbx = arrays
mov dword ptr [rax + 4*B + 4], C;
test A, A
jne L_next
mov edi, B
mov esi, C
call um_modify_0; rel32
L_next:

となります。ただし、JITコンパイルでできたコードとCで書かれた um_modify_0 関数の相対アドレスが32ビットに収まる保証がないのと、毎回レジスターを退避・復帰させるコードを書くのがだるいので、実際にはJITコンパイルで「レジスターを退避した上で本物の um_modify_0 を呼び出す関数」を作ってそれを呼び出すようにします。

3 Additionは

A = B + C;

みたいなことをします。オーバーフローはwrap aroundします。アセンブリー言語で書けば

lea A, [B + C]

となります。x86は add 命令というのもありますが、lea 命令の方が3オペランドとして使えて便利そうです。ただしbaseが r13 の場合は特別扱いが必要で、それのせいで微妙なバグを埋め込んで悩みました。

4 Multiplicationは

A = B * C;

みたいなことをします。オーバーフローはwrap aroundします。アセンブリー言語で書けば

mov eax, B
mul C
mov A, eax

となります。

5 Divisionは

A = B / C;

みたいなことをします。符号なしです。アセンブリー言語で書けば

xor edx, edx
mov eax, B
div C
mov A, eax

となります。

6 Not-Andは

A = ~(B & C);

みたいなことをします。アセンブリー言語で書けば

#if B == C
    #if A != B
        mov A, B
    #endif
    not A
#elif A == B
    and A, C
    not A
#elif A == C
    and A, B
    not A
#else
    mov A, B
    and A, C
    not A
#endif

となります。

7 Haltは

exit(0);

みたいなことをしますが、ここでは処理をインタープリターに任せるのでJITされたコードからは脱出します。

8 Allocationは

uint32_t id = fresh();
arrays[id] = calloc(C + 1, sizeof(uint32_t));
arrays[id]->length = C;
B = id;

みたいなことをします。ここではCで実装された関数 um_alloc を呼び出します。um_alloc は

struct alloc_result {
    uint32_t identifier;
    struct array *arrays;
};
struct alloc_result um_alloc(uint32_t capacity);

という風な関数で、構造体の返り値はレジスター eax, rdx で返却されます。

アセンブリー言語で書けば

mov edi, C
call um_alloc
mov B, eax
mov rbx, rdx; rbx = arrays

となります。関数を呼び出す時の注意点は um_modify_0 と同様です。

9 Abandonmentは

free(arrays[C]);

みたいなことをします。ここではCで実装された関数 um_free を呼び出します。um_free は

void um_free(uint32_t id);

という風な関数です。

アセンブリー言語で書けば

mov edi, C
call um_free

となります。

10 Outputは

putchar(C);

みたいなことをします。ここではCで実装された関数 um_putchar を呼び出します。um_putchar は

void um_putchar(uint32_t c);

という風な関数です。

アセンブリー言語で書けば

mov edi, C
call um_putchar

となります。

11 Inputは

C = getchar();

みたいなことをします。ここではCで実装された関数 um_getchar を呼び出します。um_getchar は

uint32_t um_getchar(void);

という風な関数です。

アセンブリー言語で書けば

call um_getchar
mov C, eax

となります。

12 Load Programは

if (B != 0) {
    arrays[0] = dup(array[B]);
}
pc = C;

みたいなことをします。Bが0なら単なるジャンプです。そうでない場合は、JITコンパイルをやり直す必要があるので、JITコンパイルされたコードを脱出してインタープリターで処理します。

アセンブリー言語で書けば

test B, B
jne L_reload
jmp qword ptr [rcx + 8*C]; rcx = jumptable
L_reload:
mov eax, <pc>
jmp L_epilogue

となります。

13 Orthographyは

A = imm;

みたいなことをします。アセンブリー言語で書けば

mov A, imm; imm32

です。

プロローグとエピローグ

JITコンパイルで作る関数のプロローグとエピローグも説明しておきます。まず、関数のプロトタイプは以下の通りです：

uint32_t func(struct array **arrays, void **jumptable, uint32_t registers[8], uint32_t initial_pc);

プロローグは次のようになります：

; rdi = arrays
; rsi = jumptable
; rdx = registers
; rcx = initial_pc

push rbp
mov rbp, rsp

; callee-saveなレジスターを退避する
push r15
push r14
push r13
push r12
push rbx

; ローカル変数を保存するための追加の領域を確保する
sub rsp, 40; 必要な領域は32バイトだが、スタックポインターの16バイトアラインのために40確保する
mov rbx, rdi; arraysをセット

; VMのレジスターの初期値をセットする
mov r8d, dword ptr [rdx]
mov r9d, dword ptr [rdx + 4]
; 中略
mov r15d, dword ptr [rdx + 28]

; registersとjumptableをスタックに保存しておく
mov qword ptr [rsp + 24], rdx; registers
mov qword ptr [rsp + 16], rsi; jumptable

; 実行したいprogram counterに相当する機械語のアドレスを計算する
mov rax, qword ptr [rsi + 8*rcx]

mov rcx, rsi; jumptableをセット

; 実行したい命令に飛ぶ
jmp rax

エピローグは次の通りです。

; eax = 実行を中断した（インタープリターで実行するべき）program counter
; これがこのまま返り値となる

; VMのレジスターを保存する
mov rdx, qword ptr [rsp + 24]; registers
mov dword ptr [rdx], r8d
mov dword ptr [rdx + 4], r9d
; 中略
mov dword ptr [rdx + 28], r15d

add rsp, 40
pop rbx
pop r12
pop r13
pop r14
pop r15
pop rbp
ret

補助関数を呼び出す部分は次のようになります。

sub rsp, 8; スタックには8バイトのreturn addressが積まれているので、16バイトアラインを維持するために必要

; 関数呼び出しで破壊される（caller-saveな）レジスターを退避
mov dword ptr [rsp + 16], r8d
mov dword ptr [rsp + 20], r9d
mov dword ptr [rsp + 24], r10d
mov dword ptr [rsp + 28], r11d

; 実際の関数を呼び出す。アドレスは実行時に得られた64ビットのものを使う。
mov rax, <address>; imm64
call rax

; レジスターを復帰する
mov r8d, dword ptr [rsp + 16]
mov r9d, dword ptr [rsp + 20]
mov r10d, dword ptr [rsp + 24]
mov r11d, dword ptr [rsp + 28]
mov rcx, qword ptr [rsp + 32]

add rsp, 8
ret

高速なself-modification

何度か書いたかもしれませんが、VMのバイトコードが格納される0番配列はVMの命令によって変更される可能性があります。その際の処理を行うのが um_modify_0 関数です。UMのプログラム（sandmark.umz など）は0番配列の書き換えを頻繁に行うので、um_modify_0 関数は高速に実行したいです。

高速な実行のために、「0番配列の書き換えの対象となる位置はデータとしてのみ使用され、コードとしては滅多に（あるいは全く）実行されない」という仮定を置きます（AArch64版と同様）。そして、最初の書き込みで機械語を「インタープリターへ脱出する」ものに置き換え、2回目以降の書き込みでは何もしないようにします。

AArch64版では機械語を見にいくことで「機械語を置き換え済みか否か」を判定していましたが、x86_64の機械語はバイト単位で、比較が大変なので、専用の配列を新たに導入します。これで、um_modify_0 関数は次のようになります：

static bool *patched; // 0番配列と同じ長さの配列
void um_modify_0(uint32_t b, uint32_t c)
{
    if (patched[b]) {
        // 2回目以降の書き込みである；機械語は書き換え済み
        return;
    }
    patched[b] = true;

    // 機械語を書き換えに行く
    uint8_t *target = jumptable[b];
    uint8_t *instr = target;
    instr = mov_r32_imm32(instr, eax, b); // mov eax, <b>; 5 bytes
    instr = jmp_rel32(instr, L_epilogue - (target + 10)); // jmp L_epilogue; 5 bytes
    assert(instr == target + 10);
    __builtin___clear_cache((void *)target, (void *)(target + 10));
}

ベンチマーク

sandmark.umzの実行時間を載せておきます。hyperfineを使って10回実行しました。

Ryzen 9 7940HS / Ubuntu (WSL)

まずは最近買ったミニPCです。インタープリターはGCCでコンパイルしています。WSLで動かしていますが、そのうちWindowsネイティブ版も作りたいです。

インタープリター（素朴）

$ hyperfine "./run ./sandmark.umz"
Benchmark 1: ./run ./sandmark.umz
  Time (mean ± σ):      9.287 s ±  0.171 s    [User: 9.287 s, System: 0.000 s]
  Range (min … max):    9.106 s …  9.636 s    10 runs

インタープリター（direct threaded）

$ hyperfine "./run_directthreaded ./sandmark.umz"
Benchmark 1: ./run_directthreaded ./sandmark.umz
  Time (mean ± σ):      7.733 s ±  0.080 s    [User: 7.731 s, System: 0.001 s]
  Range (min … max):    7.660 s …  7.927 s    10 runs

JIT

$ hyperfine "./run_jit ../sandmark.umz"
Benchmark 1: ./run_jit ../sandmark.umz
  Time (mean ± σ):      3.622 s ±  0.151 s    [User: 3.620 s, System: 0.001 s]
  Range (min … max):    3.469 s …  4.012 s    10 runs

JITによってインタープリター（direct threaded）の倍くらい早くなっています。良いですね。

2 GHz クアッドコアIntel Core i5 (MacBook Pro 2020) / macOS

次に、2020年に買った最後のIntel MacBookです。当世のマシンと比べると性能が骨董品です。インタプリターはClangでコンパイルしています。

インタープリター（素朴）

$ hyperfine "./run sandmark.umz"
Benchmark 1: ./run sandmark.umz
  Time (mean ± σ):     18.433 s ±  0.201 s    [User: 18.015 s, System: 0.309 s]
  Range (min … max):   18.246 s … 18.800 s    10 runs

インタープリター（direct threaded）

$ hyperfine "./run_directthreaded sandmark.umz"
Benchmark 1: ./run_directthreaded sandmark.umz
  Time (mean ± σ):     17.619 s ±  0.925 s    [User: 16.999 s, System: 0.363 s]
  Range (min … max):   16.690 s … 19.162 s    10 runs

JIT

$ hyperfine "x86_64/run_jit sandmark.umz"
Benchmark 1: x86_64/run_jit sandmark.umz
  Time (mean ± σ):      6.958 s ±  0.095 s    [User: 6.683 s, System: 0.202 s]
  Range (min … max):    6.851 s …  7.201 s    10 runs

JITによってインタープリター（direct threaded）の2.5倍くらい速くなっています。とても良いですね。

M1 Mac mini (2020) / macOS

今回作ったわけではありませんが、性能の比較対象としてM1 Mac miniも測っておきます。インタプリターはClangでコンパイルしています。

インタープリター（素朴）

$ hyperfine "./run sandmark.umz"
Benchmark 1: ./run sandmark.umz
  Time (mean ± σ):      9.738 s ±  0.171 s    [User: 9.590 s, System: 0.120 s]
  Range (min … max):    9.549 s … 10.028 s    10 runs

インタープリター（direct threaded）

$ hyperfine "./run_directthreaded sandmark.umz"
Benchmark 1: ./run_directthreaded sandmark.umz
  Time (mean ± σ):      7.829 s ±  0.050 s    [User: 7.700 s, System: 0.116 s]
  Range (min … max):    7.751 s …  7.895 s    10 runs

JIT (AArch64)

$ hyperfine "aarch64/run_jit sandmark.umz"
Benchmark 1: aarch64/run_jit sandmark.umz
  Time (mean ± σ):      3.346 s ±  0.006 s    [User: 3.255 s, System: 0.087 s]
  Range (min … max):    3.337 s …  3.353 s    10 runs

Apple M1は総合的な性能ではRyzenに負けていると思っていたのですが、JITの結果はRyzenに勝っています。インタープリターでは順当にRyzenの方が上なので、x86_64版JITコンパイラーがCPUの性能を引き出せていないのだと考えられます。

今後やりたいこと

今後やりたいのは

高速化
Windows版
x86_64機械語入門・JIT入門の執筆

です。

まず、2023年に買ったRyzenマシンでの実行結果が2020年のApple M1に負けているのは納得がいきません。x86_64版JITをもうちょっとチューニングできないか考えたいです。

次に、今回使った呼び出し規約はUnixで使われるやつで、macOSやLinuxでは通用しますが、Windowsでは通用しません。なのでWindows版も作りたいです。

この他、今回の知見を活かしてx86_64機械語入門・JIT入門みたいな記事を書きたいです。類似の記事はすでにたくさんありそうなので、自分用の備忘録としての意味合いが大きいかもしれません。

雑記帳

人生やっていき

Universal MachineのJITコンパイラーを書いた／x86_64編

x86機械語

コードを実行時に生成するためのおまじない

レジスターの使い方

JITコンパイルで作られる関数の概要

各命令のコンパイル

プロローグとエピローグ

高速なself-modification

ベンチマーク

Ryzen 9 7940HS / Ubuntu (WSL)

2 GHz クアッドコアIntel Core i5 (MacBook Pro 2020) / macOS

M1 Mac mini (2020) / macOS

今後やりたいこと