「60分でできる!JITコンパイラ自作入門」
(0)
- JITコンパイラはインタプリタを高速にするために使われる定番の技術です。
- たいして難しい技術ではないのですが、多くの人は難しいと漠然と思っているようです。
- だから解説します。
- 以下でも示しますが、普通のインタプリタをJITコンパイラに切り替えるだけで10倍以上高速化できます。
- JITコンパイラ化する以外で10倍も速くするのは相当に困難なので、とてもおすすめの方法なのです。
- タイトルでは雑に60分とか書きましたが、ざっと理解するだけなら30分もかからないかもしれません。
(1) 64bit-インタプリタ編 [61行]
- まず、以下のような命令セットを想定して、インタプリタとJITコンパイラを作り比べてみます。
Sub_AI(SP,16); // SP = SP - 16;
Mov_RI(R0,0); // R0 = 0;
Sto_RMD(R0,SP,0); // [SP+0] = R0;
Sto_RMD(R0,SP,8); // [SP+8] = R0;
Lb_I(1); // L_1:
Lod_RMD(R0,SP,8); // R0 = [SP+8];
Lod_RMD(R1,SP,0); // R1 = [SP+0];
Add_RR(R0,R1); // R0 += R1;
Sto_RMD(R0,SP,0); // [SP+0] = R0;
Lod_RMD(R0,SP,8); // R0 = [SP+8];
Add_RI(R0,1); // R0 += 1;
Sto_RMD(R0,SP,8); // [SP+8] = R0;
CmpJlt_RII(R0,100000000,1); // if (R0 < 100000000) goto L_1;
Lod_RMD(R0,SP,0); // R0 = [SP+0];
Add_AI(SP,16); // SP = SP + 16;
End(); // R0に計算結果が入っているので、それを表示して終了;
// ちなみに上記のコードは下記のソースコードをkccにコンパイルさせて生成したものです.
// int s, i; s = 0;
// for (i = 0; i < 100000000; i++) { s = s + i; }
// つまり、0から99999999までの和を計算しています.
| Mov_RI(Rx,const); | 整数レジスタ(Rx)に定数を代入 | Rx = const; |
| Lod_RMD(Rx,Ax,disp); | メモリから値を読んで整数レジスタに入れる | Rx = [Ax + disp]; |
| Sto_RMD(Rx,Ax,disp); | 整数レジスタの値をメモリに書き込む | [Ax + disp] = Rx; |
| Add_RR(Rx,Ry); | 整数レジスタ同士の加算 | Rx += Ry; |
| Add_RI(Rx,const); | 整数レジスタに定数を加算 | Rx += const; |
| Add_AI(Ax,const); | アドレスレジスタ(Ax)に定数を加算 | Ax += const; |
| Sub_AI(Ax,const); | アドレスレジスタに定数を減算 | Ax -= const; |
| Lb_I(x); | 分岐先指定用のラベル定義命令(ラベルには名前は付けられず番号で区別する) | L_x: |
| CmpJlt_RII(Rx,const,y); | 条件分岐命令 | if (Rx < const) goto L_y; |
| End(); | R0の値を表示して終了 | |
- この10個の命令があれば上記は動きます。
- SPはA6レジスタの別名だということにします。
#include <stdio.h>
#include <time.h>
#include <stdint.h>
#include <inttypes.h>
struct Code { int64_t op, a, b, c; };
enum { Mov_RI=1, Lod_RMD, Sto_RMD, Add_RR, Add_RI, Add_AI, Sub_AI, Lb_I, CmpJlt_RII, End };
enum { R0=0, R1=1, SP=6 };
int main()
{
static struct Code t[17] = {
{ Sub_AI, SP, 16, -0 }, // -0は命令としては意味を持ってないパラメータ.
{ Mov_RI, R0, 0, -0 },
{ Sto_RMD, R0, SP, 0 },
{ Sto_RMD, R0, SP, 8 },
{ Lb_I, 1, -0, -0 },
{ Lod_RMD, R0, SP, 8 },
{ Lod_RMD, R1, SP, 0 },
{ Add_RR, R0, R1, -0 },
{ Sto_RMD, R0, SP, 0 },
{ Lod_RMD, R0, SP, 8 },
{ Add_RI, R0, 1, -0 },
{ Sto_RMD, R0, SP, 8 },
{ CmpJlt_RII, R0, 100000000, 1 },
{ Lod_RMD, R0, SP, 0 },
{ Add_AI, SP, 16, -0 },
{ End, -0, -0, -0 },
{ 0, -0, -0, -0 }
};
int64_t lb[256], r[8], a[8], pc;
char m[256]; // メモリ.
// Lb_I(x)命令を探して位置をlb[]に格納する.
for (pc = 0; t[pc].op > 0; pc++) {
if (t[pc].op == Lb_I) { lb[t[pc].a] = pc + 1; }
// 高速化のために、ラベル命令の次の命令の位置を登録している.
}
// 実行開始.
pc = 0; a[6] = 256;
clock_t tm0 = clock();
for (;;) {
struct Code *tp = &t[pc]; int64_t *mp;
switch (tp->op) {
case Mov_RI: r[tp->a] = tp->b; pc++; continue;
case Lod_RMD: mp = (int64_t *) (m + a[tp->b] + tp->c); r[tp->a] = *mp; pc++; continue;
case Sto_RMD: mp = (int64_t *) (m + a[tp->b] + tp->c); *mp = r[tp->a]; pc++; continue;
case Add_RR: r[tp->a] += r[tp->b]; pc++; continue;
case Add_RI: r[tp->a] += tp->b; pc++; continue;
case Add_AI: a[tp->a] += tp->b; pc++; continue;
case Sub_AI: a[tp->a] -= tp->b; pc++; continue;
case Lb_I: pc++; continue;
case CmpJlt_RII: if (r[tp->a] < tp->b) { pc = lb[tp->c]; } else { pc++; } continue;
case End: goto fin;
}
}
fin:
printf("R0=%" PRId64 " (%.3f[sec])\n", r[0], (clock() - tm0) / (double) CLOCKS_PER_SEC);
return 0;
}
- 実行すると R0=4999999950000000 と表示されます。
(2) 64bit-JITコンパイラ編(x64用のコード生成) [69行]
- 適当に作ればこんな感じです。
- Lod_RMD/Sto_RMD命令では、第二パラメータにSPしか指定しない想定で作って手抜きをしています。
- Sub_AI/Add_AI命令でも第一パラメータにSPしか指定しない想定で作っています。
#include <stdio.h>
#include <time.h>
#include <stdint.h>
#include <inttypes.h>
struct Code { int64_t op, a, b, c; };
enum { Mov_RI=1, Lod_RMD, Sto_RMD, Add_RR, Add_RI, Add_AI, Sub_AI, Lb_I, CmpJlt_RII, End };
enum { R0=0, R1=1, SP=6 };
#include <windows.h>
void *mallocRWX(int siz) { return VirtualAlloc(0, siz, MEM_COMMIT, PAGE_EXECUTE_READWRITE); } // [参考] https://essen.osask.jp/?a21_txt02 の(4).
void set32(char *p, int64_t i) { p[0] = i & 0xff; p[1] = (i >> 8) & 0xff; p[2] = (i >> 16) & 0xff; p[3] = (i >> 24) & 0xff; }
void set64(char *p, int64_t i) { set32(p, i); set32(p + 4, i >> 32); }
int main()
{
static struct Code t[17] = {
{ Sub_AI, SP, 16, -0 }, // -0は命令としては意味を持ってないパラメータ.
{ Mov_RI, R0, 0, -0 },
{ Sto_RMD, R0, SP, 0 },
{ Sto_RMD, R0, SP, 8 },
{ Lb_I, 1, -0, -0 },
{ Lod_RMD, R0, SP, 8 },
{ Lod_RMD, R1, SP, 0 },
{ Add_RR, R0, R1, -0 },
{ Sto_RMD, R0, SP, 0 },
{ Lod_RMD, R0, SP, 8 },
{ Add_RI, R0, 1, -0 },
{ Sto_RMD, R0, SP, 8 },
{ CmpJlt_RII, R0, 100000000, 1 },
{ Lod_RMD, R0, SP, 0 },
{ Add_AI, SP, 16, -0 },
{ End, -0, -0, -0 },
{ 0, -0, -0, -0 }
};
int64_t lb[256], r[8], pc, pass, rip;
char *tx = mallocRWX(65536); // とりあえず64KB.
// 方針: R0はRAX, R1はRCXに割り当てる. SPはRSPに割り当てる. それ以外のレジスタ指定はないと仮定(簡単のために).
// JITコンパイル開始.
for (pass = 0; pass < 2; pass++) { // あえて2回まわす.
rip = 0;
for (pc = 0; t[pc].op > 0; pc++) {
struct Code *tp = &t[pc];
switch (tp->op) {
case Mov_RI: tx[rip + 0] = 0x48; tx[rip + 1] = 0xb8 + tp->a; set64(&tx[rip + 2], tp->b); rip += 10; continue;
case Lod_RMD: tx[rip + 0] = 0x48; tx[rip + 1] = 0x8b; tx[rip + 2] = 0x84 + tp->a * 8; tx[rip + 3] = 0x24; set32(&tx[rip + 4], tp->c); rip += 8; continue; // tp->b == SP のみ想定.
case Sto_RMD: tx[rip + 0] = 0x48; tx[rip + 1] = 0x89; tx[rip + 2] = 0x84 + tp->a * 8; tx[rip + 3] = 0x24; set32(&tx[rip + 4], tp->c); rip += 8; continue; // tp->b == SP のみ想定.
case Add_RR: tx[rip + 0] = 0x48; tx[rip + 1] = 0x03; tx[rip + 2] = 0xc0 + tp->a * 8 + tp->b; rip += 3; continue;
case Add_RI: tx[rip + 0] = 0x48; tx[rip + 1] = 0xba; set64(&tx[rip + 2], tp->b); tx[rip + 10] = 0x48; tx[rip + 11] = 0x03; tx[rip + 12] = 0xc2 + tp->a * 8; rip += 13; continue;
case Add_AI: tx[rip + 0] = 0x48; tx[rip + 1] = 0x81; tx[rip + 2] = 0xc4 + 0 * 8; set32(&tx[rip + 3], tp->b); rip += 7; continue; // tp->a == SP のみ想定.
case Sub_AI: tx[rip + 0] = 0x48; tx[rip + 1] = 0x81; tx[rip + 2] = 0xc4 + 5 * 8; set32(&tx[rip + 3], tp->b); rip += 7; continue; // tp->a == SP のみ想定.
case Lb_I: lb[tp->a] = rip; continue;
case CmpJlt_RII: tx[rip + 0] = 0x48; tx[rip + 1] = 0xba; set64(&tx[rip + 2], tp->b); tx[rip + 10] = 0x48; tx[rip + 11] = 0x3b; tx[rip + 12] = 0xc2 + tp->a * 8;
tx[rip + 13] = 0x0f; tx[rip + 14] = 0x8c; set32(&tx[rip + 15], lb[tp->c] - (rip + 19)); rip += 19; continue;
case End: tx[rip + 0] = 0xc3; rip++; continue;
}
}
}
// int i; for (i = 0; i < rip; i++) printf("%02x ", ((unsigned char *) tx)[i]); // デバッグ用.
// 実行開始.
int64_t (*f)() = (void *) tx;
clock_t tm0 = clock();
r[0] = f();
printf("R0=%" PRId64 " (%.3f[sec])\n", r[0], (clock() - tm0) / (double) CLOCKS_PER_SEC);
return 0;
}
- 当然ですが、実行すると R0=4999999950000000 と表示されます。
- 変数passを使って、コード生成を2回やっています。これはほぼ無駄なのですが、しかし1回目のループが終わると変数lb[]の値がすべて確定して、それで2回目のループが回ると、正しいコードが生成できるようになっています。
- 機械語を生成している部分は、まあ普通に見て意味不明だろうと思いますが、しかしプログラムの長さとしては、インタプリタ編とそれほど違いはないということをわかってほしいです。
- そして10倍以上速いのです。
(3) 今回使った機械語
48 bx [8バイト] : RAX/RCX/RDXに定数を代入. x=8:RAX, x=9:RCX, x=a:RDX.
48 03 xx : レジスタをレジスタに加える演算( a += b; ). xx=c0+b+a*8.
48 81 xx [4バイト] : レジスタと32bit整数定数の演算. 定数が32bitで収まると分かっているのならこれを使った方が速い. xxの部分で対象レジスタや演算タイプを指定する.
48 3b xx : レジスタ同士の比較.
0f 8c [4バイト] : 条件分岐(8cならJL)
c3 : return;
48 8b xx 24 [disp] : レジスタ=[RSP+disp]; xx=84+a*8.
48 89 xx 24 [disp] : [RSP+disp]=レジスタ; xx=84+a*8.
- 参考にしたサイト(もし基礎的なことから勉強したいなら、以下のサイトは向いてないです。)
(4) レジスタだけで計算するようにすれば、インタプリタでもJITコンパイラでも速くなります(2倍速くらい)。
Mov_RI(R0,0);
Mov_RI(R1,0);
Lb_I(1); // L_1:
Add_RR(R0,R1); // R0 += R1;
Add_RI(R1,1); // R1 += 1;
CmpJlt_RII(R1,100000000,1); // if (R1 < 100000000) goto L_1;
End(); // R0に計算結果が入っているので、それを表示して終了;
(5) 理解を深めるためのQ&A
- [Q1]プログラムを見やすくするために以下の関数を導入してほしいです。
void set2(char *p, char a, char b) { p[0] = a; p[1] = b; }
void set3(char *p, char a, char b, char c) { p[0] = a; p[1] = b; p[2] = c; }
void set4(char *p, char a, char b, char c, char d) { p[0] = a; p[1] = b; p[2] = c; p[3] = d; }
void set5(char *p, char a, char b, char c, char d, char e) { p[0] = a; p[1] = b; p[2] = c; p[3] = d; p[4] = e; }
- [A1]なるほど。それを使うとこうなります。
switch (tp->op) {
case Mov_RI: set2(&tx[rip + 0], 0x48, 0xb8 + tp->a); set64(&tx[rip + 2], tp->b); rip += 10; continue;
case Lod_RMD: set4(&tx[rip + 0], 0x48, 0x8b, 0x84 + tp->a * 8, 0x24); set32(&tx[rip + 4], tp->c); rip += 8; continue; // tp->b == SP のみ想定.
case Sto_RMD: set4(&tx[rip + 0], 0x48, 0x89, 0x84 + tp->a * 8, 0x24); set32(&tx[rip + 4], tp->c); rip += 8; continue; // tp->b == SP のみ想定.
case Add_RR: set3(&tx[rip + 0], 0x48, 0x03, 0xc0 + tp->a * 8 + tp->b); rip += 3; continue;
case Add_RI: set2(&tx[rip + 0], 0x48, 0xba); set64(&tx[rip + 2], tp->b); set3(&tx[rip + 10], 0x48, 0x03, 0xc2 + tp->a * 8); rip += 13; continue;
case Add_AI: set3(&tx[rip + 0], 0x48, 0x81, 0xc4 + 0 * 8); set32(&tx[rip + 3], tp->b); rip += 7; continue; // tp->a == SP のみ想定.
case Sub_AI: set3(&tx[rip + 0], 0x48, 0x81, 0xc4 + 5 * 8); set32(&tx[rip + 3], tp->b); rip += 7; continue; // tp->a == SP のみ想定.
case Lb_I: lb[tp->a] = rip; continue;
case CmpJlt_RII: set2(&tx[rip + 0], 0x48, 0xba); set64(&tx[rip + 2], tp->b); set5(&tx[rip + 10], 0x48, 0x3b, 0xc2 + tp->a * 8, 0x0f, 0x8c);
set32(&tx[rip + 15], lb[tp->c] - (rip + 19)); rip += 19; continue;
case End: tx[rip + 0] = 0xc3; rip++; continue;
}
- たしかにだいぶ短くなりました。
- [Q2]この例では仮想的な機械語をインタプリタ実行したりJITコンパイル実行していますが、もっと高級な命令のインタプリタやJITコンパイラは作れないのですか?
- [A2]その気になればもちろん誰でも作れますが、そうすると処理も増えてプログラムが複雑になってきます。私はここで、インタプリタもJITコンパイラもほぼ同等の規模で書けるということを簡潔に示したかっただけなので、やりたい人はチャレンジしましょう!
- [Q3]このプログラム例はWindows用になっていると思うのだけど、他のOSではどうしたらいいのかな?
- [Q4]このプログラムでは、Add_RI命令をRDX経由で実現しているけど、このプログラムのように32bitで収まる定数で演算するときは、RDXを使わない機械語を出力したほうが速くなるのでは?
- [A4]はい、そうです。じゃあそのように書き直してみますかー。
case Add_RI:
if (-0x80000000LL <= tp->b && tp->b <= 0x7fffffffLL) {
set3(&tx[rip + 0], 0x48, 0x81, 0xc0 + 0 * 8 + tp->a); set32(&tx[rip + 3], tp->b); rip += 7;
} else {
set2(&tx[rip + 0], 0x48, 0xba); set64(&tx[rip + 2], tp->b); set3(&tx[rip + 10], 0x48, 0x03, 0xc2 + tp->a * 8); rip += 13;
}
continue;
- どのくらい速くなったかなーと思いましたが、手元の環境では速さに違いはないようです。x64が優秀で、RDX経由でも遅くならないように工夫しているようです。
- もっとプログラムが大きくなってコードキャッシュの大きさを意識するようになったら、短い機械語を出すことは有利に働いてくると思います。
- [Q5]このJITコンパイラがシンプルなのって、RAX/RCX/RDXしか使わないように工夫して、しかも生成した機械語からは他の関数を呼ばないようにしているからだと思うのだけど、あってますか?
- [A5]そうですね、x64のABIはいろいろややこしいので、この制約を外すとだんだん複雑になってきます。でもそれは「JITコンパイラだから複雑になる」のではなく、x64のABIが複雑だから複雑になるということなので、そこは分けて考えてほしいです。・・・そのことをわかってもらうために、あえてRAX/RCX/RDXしか使わないようにしていました。
- [Q6]なんか思ったよりもずいぶん短いのだけど、本当にこれでいいの?
- [A6]いやインタプリタやJITコンパイラなんて結局この程度のものです。知らないとなんか凄そうに思えるわけですが、実際はそれほどでもないのです。今はまだ加算しかないですが、減算や乗算など他の演算を足していくことは容易にできますし、条件分岐だって他のパターンを追加すればいいのです。コピペでどんどんできます。私はkccのインタプリタ部とJITコンパイラ部をあわせて数時間で書きましたが、この短さを思えば「まあ確かにそれくらいでできるんだろうな」って思いますよね?
- JITコンパイラが69行って思うかもしれないですけど、テスト用のデモコードで19行使っているので、JITコンパイラは50行程度です。