Rust製x86エミュレータに手書きAArch64テンプレートJITを載せて、インタプリタに勝つまで【rustx86 JIT編 #1】

Rustでx86エミュレータ rustx86 を書いている。ブラウザ(WebAssembly)でもネイティブでも同じcoreが動き、実物のLinuxが起動してgccでCコンパイルまでできる。ネットワーク編・ディスク編に続いて、今回からJIT編。インタプリタを捨てずに、その上へテンプレートJITを重ねて、初めてインタプリタに勝つまでの記録である。

結論から書くと、Linuxブート970M命令の定規で インタプリタ9.5秒 vs JIT 8.2秒(-14%、119 MIPS)、gccコンパイル込みの総合時間でもJITが最速になった。ここに至るまでに、生成した機械語が速いかどうかより先に戦うべき敵が3つあった。税、正しさ、密度である。

前提: 100 MIPSのインタプリタと、一度死んだJIT

rustx86のインタプリタは決定的に動く。同じディスクイメージなら起動までの命令数はビット単位で同じ、TSC(ゲストの時計)は命令数そのもの。この決定性はテストの土台(命令数指紋・golden trace・実カーネルのUnicornロックステップ)なので、JITを載せても1bitも壊せないという縛りがある。

実はJITは一度死んでいる。Cranelift(Rust製コンパイラバックエンド)で作ったF1cは、相対では-9〜13%まで来たのに、Craneliftをリンクするだけでバイナリが太ってインタプリタ自体が+25%遅くなり、絶対速度で負けた。コードフットプリントは税である。そこで今回のF1dは方針を変えた:

• バックエンドは dynasmrt でAArch64機械語を手書き(リンク税ほぼゼロ)
• 生成コードはテンプレート式。レジスタ割付も最適化パスもない。1命令 = 決まった雛形

// mov r32, imm32 の雛形。ゲストレジスタはメモリ上のMachine構造体に居る
JitOp::MovRI { dst, imm } => {
    mov_imm32(&mut a, 10, imm);                 // w10 = imm
    store_reg(&mut a, l, machine, 10, dst);     // str w10, [x19, #regs+dst*4]
}

敵1: 税 — 受け口はチェーン入口にだけ置く

JITの入口(「ここに焼けたブロックがあるか?」の問い合わせ)をどこに置くかが最初の設計判断だった。F1c時代、インタプリタのデコードキャッシュEntryにスロットを畳んだら、Entryが8バイト太っただけでインタプリタが-22%した。ホットループには1バイトも足せない。

F1dの受け口はブロック頭(チェーン入口とtaken分岐の着地)だけ。毎命令のループには分岐1つも足さない。

pub struct JitHook {
    /// pa に焼けたブロックがあり世代が一致し予算内なら実行して命令数を返す。
    /// 0 = 入らなかった (未焼き・世代落ち・予算不足)
    pub try_enter: fn(pa: u32, gen: u32, budget: u32) -> u32,
}

budget = min(チェーン残り+1, tick_countdown) がミソで、ブロック実行中にデバイスtickが起きないことを入場時に保証する。割り込みの受付位置が毎命令実行と同一になり、決定性が構造で守られる。設計を工夫すれば、JITは決定性の敵ではない。

もうひとつの税はブロックの探索。焼けたブロックはHashMapではなく直接マップ64K〜2Mスロット(taken分岐の着地は約7命令に1回来る — HashMapのハッシュ代だけで+1.9秒を実測済み)。

敵2: 正しさ — 指紋が緑でも穴はある

決定性の門番は「指紋」である。ブート970M命令を走らせ、総命令数とシリアル出力のFNVハッシュがJIT on/offで完全一致したら緑。

指紋: 命令数=970000000 シリアルFNV=e841cce816fe3eb7   ← on/off両方でこの値

ところがこの門番が3日間緑のまま、gccコースでJITだけがデッドハルトする事件が起きた。犯人は自己書き換え検出の穴。コードページには page_has_code ビットが立ち、ストアのたびに世代を進めて古い写しを捨てる仕組みだが、このビットを立てるのはインタプリタのデコード充填だけだった。JITのbakeしか踏まないページはビットが立たず、virtioの記述子域をゲストが書き換えても誰も気づかない — 完了割り込みが迷子になってハルトした。

修正はbake時に同じ義務を果たす1行。教訓は「門番はコースの数だけ強くなる」で、以後はブート指紋(jboot)に加えてgccコンパイル窓の指紋(jcmd)も毎PRの関所になった。再発防止テストには罠まで書き残した:

// 全opがJIT語彙のSMCループ。8bitストア版にしないこと —
// あちらは語彙外でインタプリタに落ち、fillがページを守って偽緑になる
let code = [
    0xB9, 0x05, 0, 0, 0,        // mov ecx, 5   ← 書き換えの標的
    // ...
    0xC7, 0x05, ...,            // mov dword [0x2001], 7 (JIT内ストア)
    0xEB, 0xE2,                 // jmp ループ
];

敵3: 密度 — カバレッジを上げても速くならない日

語彙(JITが焼ける命令の種類)を広げればカバレッジが上がり、カバレッジが上がれば速くなる — と思っていた時期が私にもあった。

cc1(gccの中身)は8bit演算だらけで、当初は語彙外に当たってブロックが平均1.8命令に断片化していた。そこで8bit語彙を全部実装した。AArch64はリトルエンディアンなので、8bitレジスタはバイト番地の直接アクセスで書ける(ALはregs[0]の第0バイト、AHは第1バイト)。

/// AL/CL/DL/BL = 各regsの第0バイト、AH/CH/DH/BH = 第1バイト
fn reg8_addr(l: &JitLayout, r8: u8) -> usize {
    l.regs + (r8 as usize & 3) * 4 + usize::from(r8 >= 4)
}

ここで事件その2。語彙を広げた途端、gcc窓が44 MIPSまで暴落した(インタプリタは70)。観測カウンタを見ると焼き358,932回に対して据付76,484 — スロット衝突で追い出されたブロックが何度も焼き直される「再焼きの嵐」だった。同期エミットは1回数十µs、それが28万回。スロットを256K→2Mに広げて鎮火した。「焼ける頭が増える最適化」はスロット表の再設計とセットである。

さらに、ブロックごとに損益を数えるローリング再審も入れた。大きく焼けてもtaken分岐で2命令で退出するブロックは、入場の固定費が勝つ。

b.enters += 1;
b.execd += k;
if b.enters == 256 {
    if b.execd < 256 * 3 {
        slot.block = None;   // 平均3命令未満 — 負の印に降格、インタプリタに返す
    } else {
        (b.enters, b.execd) = (0, 0);  // 再審は繰り返す (局面は変わる)
    }
}

一度きりの審査ではダメだった。ブート期に好成績で「殿堂入り」したブロックが、コンパイル定常状態では分岐早退ばかりになる。局面は変わるものとして扱う。

それでも窓は-15%。決定打はセンサスだった。collectを止めたuop(ブロックの断片化犯)を度数で出すと:

grp5 (間接call/jmp・push rm・inc rm)   38.2%
mov moffs (絶対番地mov)                26.7%
grp3 (test [mem],imm)                  20.1%
setcc                                   7.8%

推測でなく実測で的を決め、上位4系を埋めた。moffsに至っては新しい雛形すら要らず、「base/index無しのメモリ形」に畳むだけで既存語彙が受けてくれた。……が、カバレッジが58.5%に上がっても壁時計は1msも動かなかった。増えた分の生成コードが、ちょうどインタプリタと同速だったのである。敵はカバレッジではなく、生成コード1命令あたりの密度だった。

密度の2梃子: 死んだフラグを書かない、条件はホストフラグで

テンプレートJITの1命令は重い。ALU 1個につきゲストレジスタのロード/ストアに加えて、遅延フラグの材料(op/幅/a/b/cin/r)を6ストアする。条件分岐のたびに条件判定ヘルパを関数呼び出しする(x86コードは約7命令に1回分岐する)。

梃子その1は死んだ材料を書かないこと。xor eax,eax; test ... のように、次の命令が読む前にフラグを完全上書きするなら、前の6ストアは死んでいる。エミット時に隣を見て省く。

梃子その2が本命の条件判定のインライン化。cmp の直後の jcc は、材料(a, b)がまだホストレジスタに生きている。ならばARMのフラグで直接判定すればいい:

5 | 7 => {  // SUB/CMP: subs後のARM条件はx86 jccと1対1
    if w == 2 {
        dynasm!(a; .arch aarch64; cmp w10, w11);
    } else {
        // 8bitは上位バイトに寄せて桁を合わせる (C/V/N/Zが8bit演算と一致)
        dynasm!(a; .arch aarch64; lsl w4, w10, 24; lsl w5, w11, 24; cmp w4, w5);
    }
}
// x86 cc → ARM cond: E→eq NE→ne B→lo AE→hs BE→ls A→hi L→lt GE→ge …
4 => dynasm!(a; .arch aarch64; cset w0, eq),

これは危ない橋でもある。rustx86には「意味論の原本は1つ」という掟があり、フラグの意味をJITに二重実装してはいけない。踏み込むなら守りを置く — 全kind×全16条件×両幅×境界値オペランドで、alu; setcc をJIT on/offの両方で実行して6フラグまで突き合わせる全数照合テストを先に書いた。

そしてこのテストが初回実行で赤くなった。

kind=0 cc=2 wide=true a=0x0 b=0x0: interp=(0, …) jit=(1, …)

ADD後のキャリー判定(jb)が逆である。x86のCFは、ADDではARMのCと同じ極性だが、SUB系では逆(ARMは借りを!Cで表す)。subs用の写像 B→lo をADDに流用した写し間違いで、全数照合がなければブートは緑のまま(カーネルはadd直後のjbをほぼ踏まない)、いつかgccが黙って壊れた計算をするところだった。ついでにARMの単一条件では表せない組(ADD後のBE/A = CF|ZFの合成)も見つかり、そこは正直にヘルパへ落とすことにした。

おまけの罠: この全数照合テスト、最初は全部緑なのに何も検証していなかった。テスト用ブロックが2命令しかなく、最小ブロック長4の門前払いでJITに一度も入っていなかったのである。パディングを足し、jit_instrs > 0 の空回り検出を入れて塞いだ。テストはまず「本当に対象を踏んでいるか」を疑うべし。

結果

冷間・交互A/B(offとonを交互に走らせて差だけ読む)3周:

決定性はビット単位で無傷のまま(指紋2コース+全数照合+SMCテストが毎PRの関所)。ブラウザ版にもJITのON/OFFセレクタを付けたので、同じマシンで切り替えて遊べる。ちなみにwasm版では「WebAssembly.Moduleのコンパイル固定費0.5ms×36,000ブロック=+19秒」という別の税(instantiate地獄)を踏んでおり、4096ブロックを1モジュールに束ねるバッチ据え付けで解消した — 単価6.3µsの操作も、Nで来ると別の敵になる。

残る宿題はgcc窓の-2%。正体は毎命令のゲストレジスタのロード/ストアで、ブロック内でホストレジスタに常駐させれば消える — ただしヘルパ呼び出し境界のflush規律という大工事になる。台帳に寝かせて、実測が呼んだら取りに行く。

(rustx86のコードは GitHub にある。ネットワーク編・ディスク編もどうぞ)

次回: QEMUの主砲「TB chaining」が不発だった話 — 真の敵はtick粒度だった【rustx86 JIT編 #2】