Rust+WebAssemblyのx86エミュレータにJITを育てる — 「1bitも変わらない」証明がバグを2匹仕留めた

Rust

ブラウザで32bit Linuxを起動する自作x86エミュレータ rustx86 に、テンプレートJITを育てている。前回はレジスタ間演算だけの骨格 (F1a) を通し、「JITを有効にしても命令数もシリアル出力もビット同一」を3つのISA (Apple M1 / x86_64×2) で証明した。

今回はJITの語彙をメモリアクセスとスタック操作に広げた。そして、その過程で決定性の審判がバグを2匹仕留めた。この記事はその謎解きの記録である。

前提: このJITの「絶対に譲らない約束」

このプロジェクトのJITには、速度より優先する約束が1つある。

JITを有効にしても、実行結果が1bitも変わらないこと。

具体的には、同じLinuxイメージをブートしたとき:

  • シェル到達までの命令数 (TSC) が完全一致する
  • シリアル出力が全文ビット同一である

これをCIの「決定性ゲート」(jit-check.mjs) が毎PR見張っている。エミュレータのJITは「だいたい合っている」が一番怖い。速くなったが挙動が微妙に違うのでは、最適化ではなく別のマシンになってしまう。

F1b: メモリアクセスをJITに入れる — フォールト脱出モデル

F1aがメモリアクセスを締め出していた理由は #PF (ページフォールト) である。実CPUの約束は「フォールトした命令は何も起きなかったことになる」— つまり命令の途中で失敗したら、レジスタもフラグも命令前の姿に巻き戻す必要がある。生成コードに巻き戻し機構を持ち込むと一気に複雑になる。

答えは巻き戻さない。失敗しそうなら、失敗する前に逃げるだった。

/// メモリロード 32bit。フォールトしそうなときは記録せず
/// 上位32bitで合図する。生成コードはそれを見て、
/// 状態を1つも変えずにブロックを脱出する。
#[no_mangle]
pub unsafe extern "C" fn rx86_jit_ld32(m: *const Machine, seg: i32, off: i32) -> i64 {
    let m = &*m;
    let la = m.cpu.lin(seg as usize, off as u32);
    match m.jit_try_read32(la) {
        Some(v) => v as i64,   // 成功: 値 (上位32bit=0)
        None => 1i64 << 32,    // 脱出の合図
    }
}

生成されるwasmコードはこの合図を見て、こう逃げる (擬似コード):

v = call ld32(machine, DS, eff_addr)
if (v >> 32) != 0:
    store(ip, block_head + offset_of_this_insn)  ; 生成時定数
    return k                                     ; 完全に実行し終えた命令数

脱出は「現命令の状態を1つも変える前」に起きるので、巻き戻すものが無い。あとはインタプリタがその命令を白紙からやり直し、#PFの記録も配送も従来経路がやる。

フォールト脱出モデル

この設計には便利な非対称がある。脱出は保守的でよい — 本当はフォールトしない場面で余計に脱出しても、インタプリタがやり直して同じ結果になる (遅くなるだけで意味は不変)。逆 (フォールトを見逃して実行) だけが許されない。おかげでページ跨ぎロードのような稀で面倒な形は、判定を作り込まず無条件脱出に倒せる。

ストアとRMW (add [mem], reg のような読んで書く形) も同じ思想で入れた。RMWは「書き込み権限のtranslateを最初に試す」のが肝で、x86のページングに書き込み専用ページは無い (書ければ必ず読める) ため、これが通ればフラグを汚した後に失敗する道が消える。

事件1: printkの時刻が±20µsずれる

ロードを語彙に入れてゲートを回したら、シリアル出力の全文一致が落ちた。差分を見ると:

interp: [    4.836823] VFS: Disk quotas dquot_6.6.0
jit:    [    4.836838] VFS: Disk quotas dquot_6.6.0

printkのタイムスタンプだけが±20µsずれている。本文は同じ。タイムスタンプはゲストがrdtscで読んだ時刻から作られるので、JIT有効時だけ、ゲストの見る時計が実行列とずれていることになる。

決定性は、それ自体が最強のデバッガである

ここで効いたのが決定性だ。このエミュレータは同じ入力なら全状態が毎回ビット同一に進む。ならば、インタプリタ実行とJIT実行を並走させて、スナップショット (CPU+装置+RAM全体) をバイト比較すれば、最初に食い違った瞬間を二分探索で特定できる

// 粗い刻みで食い違い区間を見つけたら、新しいペアを作り直して
// その区間の頭まで走らせ、刻みを1/100にして再走する
async function compare(base, step, cap) {
  const A = await mk(false);  // インタプリタ
  const B = await mk(true);   // JIT
  runTo(A, base); runTo(B, base);
  for (;;) {
    const t0 = A.emu.tsc();
    if (t0 >= cap) return { t0: null };
    runTo(A, t0 + step); runTo(B, t0 + step);
    const d = firstDiff(A.emu.save_state(), B.emu.save_state());
    if (d !== -1) return { t0, d };
  }
}

スナップショット並走の二分探索

10M命令刻み→100k→1k→100と絞ると、tsc=140,003,300から100命令以内で最初の食い違いが起きていた。スナップショットを解読すると:

interp: ip=c11e23d3  edx=c24eee38   (mov (%eax),%edx を実行した後)
jit:    ip=c11e23d1  edx=0          (実行する前)

同一tscで、JIT側だけipが1命令手前に居る。レジスタも「1命令ぶん過去」の値だ。つまりJIT側は時計だけが1命令ぶん先行している — どこかで実行していない命令の時刻を払っている。

犯人: ページ跨ぎ着地の前払い

このエミュレータの時計は「実行する命令の分を前払いする」規約で回っている。インタプリタの連結ループは、次の命令がページを跨ぐときは払わずに外へ返す (外側が払い直す)。ところがJITブロックの清算は、先にn命令ぶん (=ブロック内の残り + 着地の1命令の前払い) を払ってから跨ぎ判定でreturnしていた:

// 修正前: 払ってから跨ぎ判定 — 跨ぐたびに tsc が+1過払い
m.cpu.tsc += n;
m.tick_countdown -= n;
let new_lin = m.cpu.lin(CS, m.cpu.ip);
if new_lin >> 12 != lin >> 12 {
    return;  // 着地の前払いが宙に浮く。外側がもう一度払う
}
// 修正後: インタプリタと同じ順序 — 跨ぎ判定を先に、前払いは続行するときだけ
m.cpu.tsc += n - 1;            // ブロック内の残りだけ先に払う
let new_lin = m.cpu.lin(CS, m.cpu.ip);
if new_lin >> 12 != lin >> 12 {
    return;                    // 着地の時計は外側の担当
}
m.cpu.tsc += 1;                // 続行するときだけ前払い

面白いのはこのバグがF1aから潜んでいたことだ。F1aの語彙 (レジスタ間のみ) では該当ブロックが稀すぎて、ずれがprintkのµs解像度に届かず、ゲートをすり抜けていた。F1bでブロックが倍増して初めて±20µsの姿を現した。

修正後は、スナップショット全体がブート完了 (6億命令) まで完全一致。この並走比較はゲートより強い審判として tools/webtest/jit-lockstep.mjs に常設した。

事件2: カバレッジ1.2%の謎

語彙をスタック操作 (push/pop/call/ret/leave) まで広げても、速度は伸びなかった。原因を探るため「JITブロックの中で実行された命令数」を数えるカウンタを足したら、衝撃の数字が出た。

ブート6億命令のうち、JITの中で走ったのは7M命令 = 1.2%。

仮に生成コードが無限に速くても全体は1.2%しか縮まない。そして動的uop分布を測ると、語彙内の命令は既に81.2%あった。語彙は足りている。ブロックが張れていないのだ。犯人は4人いた:

犯人 何が起きていたか 対策
焼きの遅配 コンパイルが50M命令に1回しか届かず、ブートのフェーズ的な熱に間に合わない スライスを2Mに
ブロックの断片化 語彙外の命令 (18.8%) が点在し、走路が平均5命令の断片になる 語彙外の1命令を挟んで前後をブロック化 (タイル張り)
断片の再加熱 断片ごとに熱カウンタ1024回が要る 熱い頭から走路全体を一気に焼く (走路焼き)
入口の蒸発 直接マップのキャッシュにfillが15M回来て、据え付けたJITブロックの入口が衝突で剥がれ続けていた 永続台帳から fill のついでに埋め直す

最大の犯人は4人目だった。ブロック自体は生きているのに、そこへ入るためのポインタが平均数十万命令で消えていた。対策後、カバレッジは1.2% → 18.1%

正直な現在地

それでも壁時計はまだ微赤字である (インタプリタ13.1s vs JIT 13.7s)。平均ブロック長4.4命令で、メモリアクセスのヘルパはインタプリタと同じ仕事+呼び出し税を払う — ディスパッチ除去だけでは1命令あたりの収支が並ぶ。カバレッジの機構は完成したので、次は単価 (変換のインライン化、あるいはネイティブのCraneliftバックエンド) の番だ。

数字だけ見れば「JITはまだ勝っていない」。だが今回の収穫は速度ではない。

  • 決定性の審判は、µs未満で潜んでいた会計バグを二分探索一発で特定した
  • 4段の計測 (カバレッジ→uop分布→ブロック長→入場数) が、次に何をすべきかを毎回一意に決めた

「測ったものだけ信じる」を徹底すると、失敗した増分さえ次の税の在り処を教えてくれる。JITの旅はまだ続くが、審判団がいる限り道に迷うことはない。

リポジトリ

github.com/yoshiharu-ishii/rustx86 — JITの考え方と仕組みの図解は docs/jit.md にまとめてある。

コメント

タイトルとURLをコピーしました