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

の続き。AArch64向けJITがインタープリターに勝ったので記念カキコ（死語）。

今回書いたコードはこれ：

Apple M1でsandmark.umzを実行した時の時間：

  • 通常のVM: 13秒
  • Direct threaded VM: 11.5秒
  • JIT-compiling VM: 6.5秒

配列の確保と解放の処理を改善したので、インタープリターの方もタイムが縮まっている。

Raspberry Pi 4でも試してみた：

  • 通常のVM: 2分26秒
  • Direct threaded VM: 2分24秒
  • JIT-compiling VM: 1分22秒

JITの方が半分近い所要時間で済んでいる。大勝利だ。

JITコンパイラーの概要

JITで作られる関数のシグニチャーは

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

とする。arraysは配列のIDからアドレスへの対応表で、jumptableはVMの命令位置から機械語のコード位置への対応表、registersはVMのレジスターの初期値、initial_pcは最初のProgram Counterである。返り値は、実行を中断することになった機械語のアドレスである。この返り値でjumptableを二分探索すれば中断した時点のPCが復元できる。

この関数を組み立てるのはcompile関数である。ここではまずmmapで読み書き可能なメモリーを確保する。

レジスターの使い方は次のようにする：

レジスター用途
x19arrays
x20jumptable
w21R0
w22R1
w23R2
w24R3
w25R4
w26R5
w27R6
w28R7
その他引数の受け渡し・テンポラリー

r19からr28の10個はcallee-saveなので、関数呼び出し（um_modify_0, um_alloc, um_free, um_putchar, um_getchar）の際に退避する必要がない。その代わり、最初に関数に入る時に元の値をスタックに保存し、関数から抜ける時に復元する必要がある。

関数のプロローグでは、まずフレームポインター・リンクレジスターの値を退避する。次にcallee-saveレジスターを退避し、x2の値（引数のregisters）もスタックに保存する。

その次に使うレジスター10個を初期化する。引数のarraysとjumptableはx19, x20に保存し、w21からw28は引数のregistersから読み出す。

それが終わったら、jumptable[initial_pc]にジャンプする。C言語のGCC拡張で書くと goto *jumptable[initial_pc] である。

エピローグはプロローグの直後に配置した。

エピローグへはBL (Branch with Link)命令で飛んでくる想定で、x30にリターンアドレスが保存されている。これ（から4を引いた値）をx0に移して返り値とする。

その次はスタックに保存したregistersからuint32_t registers[8]のアドレスを取ってきて、VMのレジスターを保存する。

callee-saveなレジスター10個とx29, x30を復元してRETする。

それぞれのVM命令について見ていこう。オペランドとなるレジスターはWa, Wb, Wcで表現する。XarraysはX19, XjumptableはX20の意味とし、XtmpはX0を使用する。

  • 0: Conditional Move
    • CMP Wc, #0
    • CSEL Wa, Wa, Wb
  • 1: Array Index
    • LDR Xtmp, [Xarrays, Wb, UXTW #3]
    • ADD Xtmp, Xtmp, Wc, UXTW #2
    • LDR Wa, [Xtmp, #4]
    • Cコンパイラーに出力させたコードを参考にした。
  • 2: Array Amendment
    • LDR Xtmp, [Xarrays, Wa, UXTW #3]
    • ADD Xtmp, Xtmp, Wb, UXTW #2
    • LDR W2, [Xtmp, #4]
      • um_modify_0に渡すために変更前の値を保存しておく
    • STR Wc, [Xtmp, #4]!
    • 以下、0番配列を変更する場合の処理
    • CBNZ Wa, L_next
    • MOV W0, Wb
    • MOV W1, Wc
    • BL um_modify_0
      • um_modify_0関数を(Wb, Wc, 元の値)で呼び出す
    • L_next: （次の命令）
    • um_modify_0関数については後で詳しく述べる。
  • 3: Addition
    • ADD Wa, Wb, Wc
    • シンプル。
  • 4: Multiplication
    • MUL Wa, Wb, Wc
  • 5: Division
    • UDIV Wa, Wb, Wc
  • 6: Not-And
    • Wb=Wcの場合
      • 単なるbitwise negation
      • MVN Wa, Wb
    • Wb≠Wcの場合
      • AND Wa, Wb, Wc
      • MVN Wa, Wa
    • NotがつかないAndやOrを実行する場合は「Not-Andの結果の後に続けてNot-Andを実行する」というような形になるが、それを一命令にまとめる（macro-op fusionみたいなやつ）とジャンプ命令との兼ね合いが面倒なことになりそうなのでここでは採用しない。
  • 7: Halt
    • JITコンパイルされたコードを抜ける。
    • BL L_epilogue
  • 8: Allocation
    • 外部の関数um_allocを呼び出す。um_allocからは新しい配列のIDと、arraysの新しい値が返ってくる（C言語では構造体を値で返すように見えるが、機械語レベルではレジスターR0, R1に展開される）。
    • MOV W0, Wc
    • BL um_alloc
    • MOV Wb, W0
    • MOV Xarrays, X1
  • 9: Abandonment
    • 関数um_freeを呼び出す。
    • MOV W0, Wc
    • BL um_free
  • 10: Output
    • 関数um_putcharを呼び出す。
    • MOV W0, Wc
    • BL um_putchar
  • 11: Input
    • 関数um_getcharを呼び出す。
    • BL um_getchar
    • MOV Wc, W0
  • 12: Load Program
    • 0番配列を置き換える場合：JITコンパイルされたコードを抜ける。
    • 置き換えない場合：goto *jumptable[new_pc] に相当する命令を実行する。
    • CBNZ Wb, L_reload
    • LDR Xtmp, [Xjumptable, Wc, UXTW #3]
    • BR X0
    • L_reload:
    • BL L_epilogue
  • 13: Orthography
    • 即値が16ビットに収まる場合（65535以下）：
      • MOVZ Wa, #value
    • 即値が16ビットに収まらない場合：
      • MOVZ Wa #(value & 0xFFFF)
      • MOVK Wa, #(value >> 16), LSL #16

main関数では、コンパイル結果の関数を適当な引数で呼び出す。返り値からコンパイル済みコードを中断した命令が分かるので、それを解釈してやる。コンパイル済みコードを中断する命令は典型的にはHaltかLoad Programだが、「0番配列の書き換え」の対象となった命令は「コンパイル済みコードを中断する命令」に書き換えられる（後述）ため、任意の命令が入っている可能性がある。

コンパイル済みコードを中断した命令がLoad Programで0番配列の置き換えが行われた場合は、改めてコンパイルを行う。

VMの命令の中で一番厄介なのは配列の書き換えで、0番配列を書き換えるケースを高速で処理しないとJITといえどインタープリターに勝てない。

筆者の当初の実装では、書き換えられた部分に相当する機械語を改めて生成していたが、そうすると書き換えのたびにpthread_jit_write_protect_npを呼び出すことになり遅い。そこで、筆者は「0番配列で書き換えられた部分は、実際にはコードとしては滅多に実行されない」「0番配列の同じ位置に読み書きされることが多い」（0番配列の特定の位置が、コードではなくデータとして使用されている）という仮定を置いた。そして、書き換えの対象となった部分は問答無用で「コンパイル済みコードを中断する（実際の処理はmain関数内のインタープリターで行う）」命令に書き換えることにした。

そうすると、「0番配列の同じ位置を2回以上書き換える場合、機械語がすでに『コンパイル済みコードを中断する命令』になっていればそれ以上の書き換えは必要ない」ということになる。つまり、同じ位置の2回目以降の書き換えではpthread_jit_write_protect_npを呼ばなくて良いようになる。

こうして筆者のJITコンパイルするVM実装は、単純なインタープリターの倍くらい速くなった。めでたしめでたし。

JIT所感

JIT自体は意外とすんなりとできたが、最初インタープリターよりも遅かったのがショックだった。ただ漫然とJITコンパイルすればいいというものではなく、どこがボトルネックか考えながらJITコンパイルする必要がある。

今回のVMはレジスターマシンで命令も機械語に近いものが多く、割とJITコンパイルしやすい方だったのではないかと思う。スタックマシンではレジスター割り付けを考える必要がありそうだし、スクリプト言語のように動的な要素が多い場合はJITコンパイルはもっと複雑になりそうだ。

今回はAArch64がターゲットでJITアセンブラーを使わずにゴリゴリと書いていったが、x86系は命令のエンコードが複雑そう（偏見）だし、x86系をターゲットとするJITコンパイラーを書くとしたら何らかのJITアセンブラーを使いたいところである。

まあなんだかんだ言って「VMをJITで高速化する」という目標は達成できたので今夜は枕を高くして寝られそうだ。

次からはいよいよICFPC2006の課題にチャレンジすることになる。

このエントリーをはてなブックマークに追加

「Universal MachineのJITコンパイラーを書いた」への2件のフィードバック

  1. ピンバック: AArch64アセンブリーを出力するBrainfuckコンパイラーを書いてみる | 雑記帳

  2. ピンバック: Universal MachineのJITコンパイラーを書いた／x86_64編 | 雑記帳

コメントは停止中です。