AArch64でJITしてみる

AArch64で実行時にコード生成して動かしてみよう。いわゆるJITアセンブル・JITコンパイルというやつだ。

AArch64の機械語について

AArch64の機械語についての公式の情報源は「Arm Architecture Reference Manual」である。以下のリンクからダウンロードできる：

Arm Architecture Reference Manual Armv8, for A-profile architecture

機械語で関数を呼び出すには、呼び出し規約に関する情報も必要である。これは以下から入手できる：

Releases · ARM-software/abi-aa
- ABI for the Arm 64-bit Architecture > Procedure Call Standard for the Arm 64-bit Architecture

Appleのプラットフォーム向けにコードを書く場合は、Appleのマニュアルも参照する必要がある。

Writing ARM64 Code for Apple Platforms | Apple Developer Documentation

Appleのプラットフォームで実行時にコード生成をしたいなら、さらに次の資料も読んでおく必要がある：

Porting Just-In-Time Compilers to Apple Silicon | Apple Developer Documentation

Architecture Reference Manualを読むと、例えば関数から戻るための命令 RET は次のようにエンコードされることがわかる：

31-28	27-24	23-20	19-16	15-12	11-8	7-4	3-0
1101	0110	0101	1111	0000	00<Rnの上位2ビット>	<Rnの下位3ビット>0	0000

RET命令

ただし、Rnは戻り先のアドレスを格納したレジスターで、通常はX30（整数値の30がエンコードされる）である。

C言語で書くなら次のようなコードでこの機械語を生成できる：

0xD65F0000 | (Rn << 5)

二進数と十六進数の変換は、Mac標準添付の「計算機」アプリの「プログラマ」モードが便利だ。

レジスターがcaller-saveかcallee-saveか、といった情報はProcedure Call Standardに書かれている。ここに再掲しておく：

SP: The Stack Pointer
r30: The Link Register
r29: The Frame Pointer
r19…r28: Callee-saved registers
r18: The Platform Register
- Appleのプラットフォームではreservedとなる
r17, r16: intra-procedure-call temporary/scratch register
r9…r15: Temporary registers
r8: Indirect result location register
r0…r7: Parameter/result registers

レジスターrnは32ビット幅で扱うときはwnという名前で呼ばれ、64ビット幅で扱うときはxnという名前で呼ばれる。

整数のパラメーターはr0から順番にレジスターを使って渡される。返り値が整数一つの場合はr0が使用される。

では早速書いてみよう。

はじめてのJIT

「何もせずに RET する」関数を実行時に生成するコードは次のようになる：

#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/errno.h>
#include <sys/mman.h> // mmap, mprotect, munmap
#include <unistd.h> // getpagesize
#if defined(__APPLE__)
#include <libkern/OSCacheControl.h> // sys_icache_invalidate
#endif

int main(void)
{
    printf("page size = %d\n", getpagesize());
    size_t len = getpagesize();
    void *mem = mmap(NULL, len, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
    if (mem == MAP_FAILED) {
        int e = errno;
        printf("mem = %p (failed), errno = %d (%s)\n", mem, e, strerror(e));
        return 1;
    }
    printf("mem = %p (successful)\n", mem);
    uint32_t *instr = mem;
    *instr++ = 0xD65F0000 | (30 << 5); /* RET x30 */
    int ret = mprotect(mem, len, PROT_READ | PROT_EXEC);
    if (ret != 0) {
        int e = errno;
        printf("mprotect failed with errno = %d (%s)\n", e, strerror(e));
        return 1;
    }
#if defined(__APPLE__)
    sys_icache_invalidate(mem, (char *)instr - (char *)mem);
#else
    __builtin___clear_cache((void *)mem, (void *)instr);
#endif
    int (*f)(int) = (int (*)(int))mem;
    printf("f(42) = %d\n", f(42));
    printf("f(37) = %d\n", f(37));
    munmap(mem, len);
}

まずmmapでコード用のメモリーを確保している。「Porting Just-In-Time Compilers to Apple Silicon」に書かれているが、Apple Silicon Macでは同じメモリーに書き込み権限と実行権限の両方を同時に与えることはできない。これを（俗に？）W^X (write-xor-execute)と呼ぶ。なので最初は読み込み+書き込みのみの権限を与えている。

メモリーが確保できたら機械語を書き込む。ここでは RET の一命令を書き込んでいる。戻り先アドレスは呼び出し規約にしたがって x30 に格納されているものを使う。

次にmprotectで書き込み権限を外し、代わりに実行権限を与えている。

その次に、命令キャッシュをクリアしている。Macでは sys_icache_invalidate を呼び出し、それ以外ではGCC/Clangの組み込み関数である __builtin___clear_cache を呼び出す。x86系だとプロセッサーがこの辺上手いことやってくれるらしいが、Armではプログラマーが明示的に命令キャッシュをクリアする必要がある、らしい。ハードウェアに近いところを触っている感覚がしてワクワクする。

それができたらいよいよ関数の呼び出しだ。ここでは書き込んだコードを「一個のintを受け取って一個のintを返す関数」と見做している。この引数と返り値はいずれもw0を使って受け渡されるが、書き込んだコードではw0は触っていないため、「与えられたintをそのまま返却する関数」として動作するはずだ。

では実行してみよう。Apple Silicon Macでの実行結果：

page size = 16384
mem = 0x104bac000 (successful)
f(42) = 42
f(37) = 37

Linuxでの実行結果：

page size = 4096
mem = 0xffffb42a0000 (successful)
f(42) = 42
f(37) = 37

期待通り動作しているようだ。

なお、上記のコードはシステムがリトルエンディアンであることを仮定している。Arm的にはデータはリトルエンディアン・ビッグエンディアン両方いける（？）らしいが、命令列は必ずリトルエンディアンである必要がある（Architecture Reference ManualのB2.6.2）。

足し算してみる

次は足し算を書いてみよう。C言語で言うところの

int f(int x, int y) {
    return x + y;
}

だ。

AArch64でのADD命令は3種類ほどあるが、ここではADD (shifted register)を使う。これはC言語で書けば

d = n + (m << amount);

みたいなことをする命令で、amountを0にすれば通常の足し算として使える。

この命令をエンコードするCのコードは次のようになる：

    uint32_t sf = 0; // 0: 32 bit, 1: 64 bit
    uint32_t shift = 0; // 0: LSL, 1: LSR, 2: ASR, 3: reserved
    uint32_t Rm = 0; // w0
    uint32_t imm6 = 0;
    uint32_t Rn = 1; // w1
    uint32_t Rd = 0; // w0
    *instr++ = 0x0B000000 | (sf << 31) | (shift << 22) | (Rm << 16) | (imm6 << 10) | (Rn << 5) | Rd; /* ADD w0, w1, w0 */

w0とw1の和をw0に代入している。

完全なソースコードはこんな感じになる：

#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/errno.h>
#include <sys/mman.h> // mmap, mprotect, munmap
#include <unistd.h> // getpagesize
#if defined(__APPLE__)
#include <libkern/OSCacheControl.h> // sys_icache_invalidate
#endif

int main(void)
{
    printf("page size = %d\n", getpagesize());
    size_t len = getpagesize();
    void *mem = mmap(NULL, len, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
    if (mem == MAP_FAILED) {
        int e = errno;
        printf("mem = %p (failed), errno = %d (%s)\n", mem, e, strerror(e));
        return 1;
    }
    printf("mem = %p (successful)\n", mem);
    uint32_t *instr = mem;
    {
        uint32_t sf = 0; // 0: 32 bit, 1: 64 bit
        uint32_t shift = 0; // 0: LSL, 1: LSR, 2: ASR, 3: reserved
        uint32_t Rm = 0; // w0
        uint32_t imm6 = 0;
        uint32_t Rn = 1; // w1
        uint32_t Rd = 0; // w0
        *instr++ = 0x0B000000 | (sf << 31) | (shift << 22) | (Rm << 16) | (imm6 << 10) | (Rn << 5) | Rd; /* ADD w0, w1, w0 */
    }
    *instr++ = 0xD503201F; /* NOP */
    *instr++ = 0xd65f0000 | (30 << 5); /* RET x30 */
    int ret = mprotect(mem, len, PROT_READ | PROT_EXEC);
    if (ret != 0) {
        int e = errno;
        printf("mprotect failed with errno = %d (%s)\n", e, strerror(e));
        return 1;
    }
#if defined(__APPLE__)
    sys_icache_invalidate(mem, (char *)instr - (char *)mem);
#else
    __builtin___clear_cache((void *)mem, (void *)instr);
#endif
    int (*f)(int, int) = (int (*)(int, int))mem;
    printf("f(42, 11) = %d\n", f(42, 11));
    printf("f(37, -5) = %d\n", f(37, -5));
    munmap(mem, len);
}

ADDだけだと寂しいのでNOPも入れてみた。

Macでの実行結果：

page size = 16384
mem = 0x104ab8000 (successful)
f(42, 11) = 53
f(37, -5) = 32

Linuxでの実行結果：

page size = 4096
mem = 0xffff8ab75000 (successful)
f(42, 11) = 53
f(37, -5) = 32

うまくいったようだ。

便利ツール

簡単な算術命令をやっているうちは良いが、ポインター等が絡む命令はどういうアセンブリーを書けば良いのか（筆者のような）初心者には難しい。そこで活用したいのがCコンパイラーだ。-Sオプションでアセンブリーを出力させれば参考になる。あるいはobjdump -dも良いだろう。

出力させたコードがうまくいかないときは、一旦コードをファイル等に書き出して逆アセンブラーにかけてみると良いだろう。ググったらOnline Disassemblerという手軽に使える逆アセンブラーを見つけた。

あとは、デバッガー（gdb, lldb）だ。デバッガーがあると逆アセンブルもできるしレジスターの値もメモリーの値も見られるし、デバッガーはデバッグに便利という世界の真実に†目覚めて†しまう。

MAP_JITとpthread_jit_write_protect_np

先程のコードではApple SiliconのW^X対策としてコード書き込み後にmprotectで権限を変更していた。しかし、頻繁にコードを書き換える必要がある状況では、その都度mprotectを呼んでいては遅い。そこでAppleが提供している方法が、MAP_JITとpthread_jit_write_protect_npだ。

まず、mmapにMAP_JITオプションを渡すと権限としてPROT_READ | PROT_WRITE | PROT_EXECの3つを同時に指定できるようになる。

しかし、W^X制限は依然として有効で、書き込みと実行は同時にはできない。各スレッドには「書き込みモード」と「実行モード」のいずれかが割り当てられていて、そのスレッドからは許可された操作しかできない。

「書き込みモード」と「実行モード」を切り替えるAPIがpthread_jit_write_protect_npだ。引数に0を渡すとMAP_JITで確保したメモリーが書き込み可能・実行不能になり、引数に0でない値を渡すと書き込み不能・実行可能になる。

使用例は次のようになる：

#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/errno.h>
#include <sys/mman.h> // mmap, mprotect, munmap
#include <unistd.h> // getpagesize
#if defined(__APPLE__)
#include <pthread.h> // pthread_jit_write_protect_np
#include <libkern/OSCacheControl.h> // sys_icache_invalidate
#endif

int main(void)
{
    printf("page size = %d\n", getpagesize());
    size_t len = getpagesize();
    void *mem = mmap(NULL, len, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_ANON | MAP_PRIVATE
#if defined(__APPLE__)
                     | MAP_JIT
#endif
                     , -1, 0);
    if (mem == MAP_FAILED) {
        int e = errno;
        printf("mem = %p (failed), errno = %d (%s)\n", mem, e, strerror(e));
        return 1;
    }
    printf("mem = %p (successful)\n", mem);
#if defined(__APPLE__)
    pthread_jit_write_protect_np(0); // Make the memory writable
#endif
    uint32_t *instr = mem;
    {
        uint32_t sf = 0; // 0: 32 bit, 1: 64 bit
        uint32_t shift = 0; // 0: LSL, 1: LSR, 2: ASR, 3: reserved
        uint32_t Rm = 0; // w0
        uint32_t imm6 = 0;
        uint32_t Rn = 1; // w1
        uint32_t Rd = 0; // w0
        *instr++ = 0x0B000000 | (sf << 31) | (shift << 22) | (Rm << 16) | (imm6 << 10) | (Rn << 5) | Rd; /* ADD w0, w1, w0 */
    }
    *instr++ = 0xD503201F; /* NOP */
    *instr++ = 0xd65f0000 | (30 << 5); /* RET x30 */
#if defined(__APPLE__)
    pthread_jit_write_protect_np(1); // Make the memory executable
    sys_icache_invalidate(mem, (char *)instr - (char *)mem);
#else
    __builtin___clear_cache((void *)mem, (void *)instr);
#endif
    int (*f)(int, int) = (int (*)(int, int))mem;
    printf("f(42, 11) = %d\n", f(42, 11));
    printf("f(37, -5) = %d\n", f(37, -5));
    munmap(mem, len);
}

JITアセンブラーの活用

ここでは機械語をゴリゴリ書いていったが、大規模なJITコンパイラーを書く際はJITアセンブラーを使うことになるだろう。

x86系にはXbyakというC++で使えるJITアセンブラーがあるが、AArch64向けには富士通が作ったXbyak_aarch64というのがある。

https://github.com/fujitsu/xbyak_aarch64

こいつは一応Apple Siliconにも対応しているようだ。使用例は次のようになる：

#include <cstdio>
#include <xbyak_aarch64/xbyak_aarch64.h>

struct Code : Xbyak_aarch64::CodeGenerator
{
    Code()
    {
        add(w0, w1, w0);
        ret();
    }
};

int main()
{
    Code c;
    c.ready();
    auto f = c.getCode<int (*)(int, int)>();
    printf("%d\n", f(42, 37));
}

Xbyak_aarch64にはここではこれ以上深く突っ込まない。

その他のJITアセンブラーについては

https://github.com/wdv4758h/awesome-jit

に色々まとまっているようだ。

Universal MachineのJITコンパイラーに向けて

そもそも何で筆者がJITコンパイラーを書こうとしているかというと、前回の記事のVM (Universal Machine)を高速化したかったからだ。

ICFP Programming Contest 2006 “The Cult of Bound Variable” に挑戦してみる

というわけで、この数日で筆者が書き上げたJITコンパイラーがこれだ：

run_jit.c

JITしたコードからは文字の出力や何やらで外部の関数を呼び出す必要がある。そのためには

スタックポインターの操作
フレームポインターの操作
レジスターの退避……はcallee-save registerを使うことで省いた
ジャンプ先アドレスの埋め込み
- ジャンプ先の相対アドレスが±128MBに収まっていればBL命令に即値で埋め込めるが、一般の場合を考えると64ビットの値をMOVZ/MOVKで埋め込んでBRでジャンプするか、どこかのレジスターに保存しておいてジャンプするか、等の手段を取る必要がある。

などを考える必要がある。

JITの利点は、VMのレジスターをマシンのレジスターに対応させられる点だ。ここではw21からw28の8個のレジスターを使うことにした（いずれもcallee-save）。

Load Program命令ではVMの命令を単位とする指定されたPCにジャンプする必要がある。それぞれのPCに対応する機械語のアドレスを void *jumptable[] という配列に格納しておくことにした。Load Programで0番配列が書き換えられた場合は一旦JITされた関数を抜けてmain関数に戻って再度コンパイルする。

厄介なのは自己書き換えだ。0番配列が書き換えられた場合、JITコンパイル済みのコードも書き換える必要がある。書き換え元のコードが書き換え後のコードよりも長ければin-placeに書き換えができる（余った分はNOPで埋める）が、そうでない場合は別途メモリーを用意してそこへジャンプさせる必要がある。

で、一通り実装できた（と言いたいところだがsandmarkをLinuxで実行するとmallocがメモリー破壊を報告してきやがるのでどこかでメモリーが破壊されている可能性がある）のだが、遅い。どうやら自己書き換えが遅いのか？高速化のためにJITの実装を始めたのに、これでは示しがつかない！

動的な性質がなければ……と思うところだが、そもそも自己書き換えのない静的なコードであれば普通にAOTコンパイルすれば良いので、JITするということはすなわち動的なコードに立ち向かうということなのだろう。

高速化はともかく、一応JITするという目標は達成できたので、今回の記事はここまでとする。

雑記帳

人生やっていき