LuaJIT

LuaJIT 解析

More than 1 year has passed since last update.

動的型付けの言語のJITコンパイラの中でも魔術度が圧倒的に高いLuaJITの内部について調べました。


LuaJITの概要

LuaJITはこんな流れで実行されます

   処理     
   生成物     
ファイル       

Luaプログラムをパースし、バイトコードを生成する
バイトコード
字句解析 lj_lex.c 構文解析 lj_parse.c

バイトコードを実行する

バイトコードインタプリタ vm_x86.dasc

バイトコードのうちホットな場所を中間コードSSA IRに変換する

SSA IR  
トレースの制御をおこなう lj_trace.c 実際にトレースしてIRを生成する lj_record.c

SSA IRをアセンブラに変換する
機械語
アセンブラを生成する lj_asm.c

生成した高速なコードを実行する

とりあえず、最初のパースはだいたい普通の言語と同じだと思うから(実はちゃんと読んでない)

のでバイトコードを実行しているところから話を始めます1


バイトコードインタープリタ

LuaJITのバイトコードインタープリタはアセンブラで書いてあります。詳細は読者の宿題とします(読むのに挫折した)


トレーシング

LuaJITはTracing JITを採用しています。LuaJITのバイトコードインタープリタはトレーシングの開始になる命令(関数呼び出しとループ)の実行回数を数えておいて、必要に応じてSSA IRを生成するトレーシングを開始します。

トレーシング開始する判定を行うコードはこんな感じ(x86版)

|.macro hotcall, reg

| mov reg, PC
| shr reg, 1
| and reg, HOTCOUNT_PCMASK
| sub word [DISPATCH+reg+GG_DISP2HOT], HOTCOUNT_CALL
| jb ->vm_hotcall
|.endmacro

トレーシングを開始すると、コンパイルを実行と並行して行います。この生成しながら実行する処理の最上位部はこんな感じ。whileのところがそうです。(lj_trace.c)

/* A bytecode instruction is about to be executed. Record it. */

void lj_trace_ins(jit_State *J, const BCIns *pc)
{
/* Note: J->L must already be set. pc is the true bytecode PC here. */
J->pc = pc;
J->fn = curr_func(J->L);
J->pt = isluafunc(J->fn) ? funcproto(J->fn) : NULL;
while (lj_vm_cpcall(J->L, NULL, (void *)J, trace_state) != 0)
J->state = LJ_TRACE_ERR;
}

ここで、trace_stateの関数ポインタを渡しているが、この関数がJITコンパイラの状態を管理しています。コードはこんな感じ

/* State machine for the trace compiler. Protected callback. */

static TValue *trace_state(lua_State *L, lua_CFunction dummy, void *ud)
{
jit_State *J = (jit_State *)ud;
UNUSED(dummy);
do {
retry:
switch (J->state) {
case LJ_TRACE_START:
J->state = LJ_TRACE_RECORD; /* trace_start() may change state. */
trace_start(J);
lj_dispatch_update(J2G(J));
break;

case LJ_TRACE_RECORD:
trace_pendpatch(J, 0);
setvmstate(J2G(J), RECORD);
lj_vmevent_send_(L, RECORD,
/* Save/restore tmptv state for trace recorder. */
TValue savetv = J2G(J)->tmptv;
TValue savetv2 = J2G(J)->tmptv2;
setintV(L->top++, J->cur.traceno);
setfuncV(L, L->top++, J->fn);
setintV(L->top++, J->pt ? (int32_t)proto_bcpos(J->pt, J->pc) : -1);
setintV(L->top++, J->framedepth);
,
J2G(J)->tmptv = savetv;
J2G(J)->tmptv2 = savetv2;
);
lj_record_ins(J);
break;

case LJ_TRACE_END:
trace_pendpatch(J, 1);
J->loopref = 0;
if ((J->flags & JIT_F_OPT_LOOP) &&
中略
J->state = LJ_TRACE_ASM;
break;

case LJ_TRACE_ASM:
setvmstate(J2G(J), ASM);
lj_asm_trace(J, &J->cur);
trace_stop(J);
setvmstate(J2G(J), INTERP);
J->state = LJ_TRACE_IDLE;
lj_dispatch_update(J2G(J));
return NULL;

default: /* Trace aborted asynchronously. */
setintV(L->top++, (int32_t)LJ_TRERR_RECERR);
/* fallthrough */
case LJ_TRACE_ERR:
中略
J->state = LJ_TRACE_IDLE;
lj_dispatch_update(J2G(J));
return NULL;
}
} while (J->state > LJ_TRACE_RECORD);
return NULL;
}

コンパイラの状態はこんな感じ。

typedef enum {

LJ_TRACE_IDLE, /* Trace compiler idle. */
LJ_TRACE_ACTIVE = 0x10,
LJ_TRACE_RECORD, /* Bytecode recording active. */
LJ_TRACE_START, /* New trace started. */
LJ_TRACE_END, /* End of trace. */
LJ_TRACE_ASM, /* Assemble trace. */
LJ_TRACE_ERR /* Trace aborted with error. */
} TraceState;

LJ_TRACE_RECORDが、バイトコードからSSA IRを生成するモードです。何をやっているかは、lj_record_ins関数を読むといいでしょう。この関数で各バイトコード命令毎にSSA IRの生成処理にジャンプしています。

SSA IRを生成しながら適時オプティマイザーを呼び出しますが、これまた膨大で難しそうなので読んでいません。


SSA IR

SSA IRは最適化等をCPUに依存しない形で行うための中間コードです。64bit固定長で最適化された結果を保持するよう工夫されています。この命令セットは静的な型を持っています。推論やガードの結果などから型チェックのためのガードを減らす最適化が効率よく行えるようになっています。


SnapShot

Tracing JITでは全体をコンパイルするわけではなく、ループなど実行頻度の高い所だけをコンパイルするのでVMに戻る必要があります。しかし、VMに戻るためにはVMを整合性のとれた状態に戻す必要があります。コンパイルした機械語の状態(レジスタとかスタックの内容など)をVMの状態にマッピングするデータ構造をSnapShotといいます。SSA IRは複数個のSnapShotを保持することが出来ます。

SnapShotの定義はlj_jit.hに、SnapShotの操作に関する関数はlj_snap.cにあります。


機械語の生成

IRを生成したら、最適化を行い、機械語を生成します。SSA IRのレジスタ(16bitなので事実上無限個)からCPUのレジスタにマッピングする必要があります。この辺は大変活発に研究されています。LuaJITはreverse-linear-scan register allocationを独自のコストモデルと共に使っているそうですが、良くわからないです。linear-scan register allocationは比較的低コストでよいレジスタ割り当てが出来ると言うことでJITコンパイラでは結構メジャーなわけですが。


 まとめ

LuaJIT難しい (涙)





  1. LuaJITのパーサーは職人が1つ1つ手作りした温かみのあるコードなので、実は面白い可能性が結構ありそうです