fabi/xenia-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[iterate-4A] jit: inline load fast-path (skip trampoline) — JIT reaches interpreter parity

Browse Source 
Loads no longer call a trampoline on the common path. When memory exposes a flat
mapping and the address is proven non-MMIO and committed, the compiled block
loads directly from `membase + ea` (byte-swapped) with zero calls; otherwise it
falls back to the read trampoline.

Mechanism:
  * MemoryAccess::fast_mem() -> Option<FastMem { membase, page_table, mmio_mask,
    mmio_value }>. Default None; GuestMemory returns Some (flat 4 GiB mapping,
    page-table pointer, MMIO aperture pair). Wrappers that intercept accesses
    (the recompiler's speculative OverlayMemory) inherit None, so the fast path
    is disabled under the diff harness and their interception is preserved.
  * MemEnv carries membase/page_table/mmio_mask/mmio_value; MemEnv::new(mem)
    fills them (or nulls when fast_mem() is None).
  * emit_inline_load emits a 5-block diamond: membase==0 -> slow (disabled);
    (ea & mmio_mask) == mmio_value -> slow (maybe MMIO); page_table[ea>>12]
    COMMIT bit (49) clear -> slow (unmapped; the mapped check is mandatory —
    unmapped pages are PROT_NONE and a raw load would fault); else fast load
    membase+ea, bswap, zero-extend. Wired all integer loads (lbz/lhz/lwz/ld +
    x-forms) and FP loads (lfs/lfsx reuse the diamond then bitcast/fpromote;
    lfd/lfdx bit-copy). Stores and FP-punt still trampoline.

Validation (three ways, since the diff harness disables the fast path):
  * unit test jit_inline_load_matches_interpreter — real GuestMemory, both the
    mapped (fast) and unmapped (slow) address bit-match the interpreter.
  * e2e XENIA_JIT=1 boot+movie plays, clean exit.
  * diff regression: checked 149.2M blocks, MISMATCHES=0 (validates the
    unchanged slow/ALU/branch/FP logic).

Measured (2e9 instr, block-exec MIPS): interpreter 98.2 / JIT trampoline-loads
93.8 / JIT inline-loads 96.5 — inline loads recover the trampoline overhead,
bringing the JIT to parity (~1.7% slower, within noise). Beating the interpreter
needs the levers it can't do: block-linking (30% dispatch), host registers.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
MechaCat02
2026-07-07 19:44:56 +02:00
parent d91311c486
commit da0509b4ac
5 changed files with 292 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

270
crates/xenia-cpu/src/jit.rs
View File

@@ -22,7 +22,9 @@
use core::mem::offset_of;
use cranelift_codegen::ir::condcodes::IntCC;
use cranelift_codegen::ir::{types, AbiParam, FuncRef, InstBuilder, MemFlags, Signature, Value};
use cranelift_codegen::ir::{
types, AbiParam, BlockArg, FuncRef, InstBuilder, MemFlags, Signature, Value,
};
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{default_libcall_names, FuncId, Linkage, Module};
@@ -38,10 +40,46 @@ use xenia_memory::MemoryAccess;
/// blocks today; the rest exist for when branch/sc/trap ops are covered).
pub const RET_CONTINUE: u32 = 0;
/// Thin environment handed to memory trampolines: wraps the live `MemoryAccess`.
/// Environment handed to compiled blocks. Carries the live `MemoryAccess` (used
/// by the trampoline slow path) plus the raw pointers/constants a load needs to
/// take the **inline fast path** (skip the trampoline call entirely). When
/// `membase` is null the fast path is disabled and every access goes through the
/// trampoline — this is how the recompiler's speculative overlay forces its
/// interception (its `fast_mem()` returns `None`).
///
/// `#[repr(C)]` with a stable field order: the JIT reads `membase`/`page_table`/
/// `mmio_mask`/`mmio_value` from fixed `offset_of!` offsets.
#[repr(C)]
pub struct MemEnv<'a> {
pub mem: &'a dyn MemoryAccess,
pub membase: *mut u8,
pub page_table: *const u64,
pub mmio_mask: u32,
pub mmio_value: u32,
}
impl<'a> MemEnv<'a> {
/// Build the env, pulling the inline fast-path pointers from `mem` if it
/// exposes them ([`MemoryAccess::fast_mem`]). Memories that intercept
/// accesses return `None`, leaving `membase` null (fast path disabled).
pub fn new(mem: &'a dyn MemoryAccess) -> Self {
match mem.fast_mem() {
Some(f) => Self {
mem,
membase: f.membase,
page_table: f.page_table,
mmio_mask: f.mmio_mask,
mmio_value: f.mmio_value,
},
None => Self {
mem,
membase: std::ptr::null_mut(),
page_table: std::ptr::null(),
mmio_mask: 0,
mmio_value: 1, // (addr & 0) == 1 is never true; membase gate wins anyway
},
}
}
}
/// A compiled block: a native entry point valid for as long as the owning
@@ -621,6 +659,126 @@ fn call_store(b: &mut FunctionBuilder, ec: &EmitCtx, fref: FuncRef, ea: Value, v
b.ins().call(fref, &[ec.ctxp, ec.memenv, ea, val]);
}
/// Load an `I64` field from the `MemEnv` at `offset`.
#[inline]
fn load_env_i64(b: &mut FunctionBuilder, ec: &EmitCtx, offset: usize) -> Value {
b.ins().load(types::I64, MemFlags::trusted(), ec.memenv, off(offset))
}
/// Native (host-endian) guest load of `size` bytes at `hptr`, byte-swapped to
/// the guest's big-endian value and zero-extended to `I64`. `notrap` because the
/// caller has already proven the page is committed.
fn emit_native_load(b: &mut FunctionBuilder, hptr: Value, size: u8) -> Value {
let mut mf = MemFlags::new();
mf.set_notrap();
match size {
1 => {
let v = b.ins().load(types::I8, mf, hptr, 0);
b.ins().uextend(types::I64, v)
}
2 => {
let v = b.ins().load(types::I16, mf, hptr, 0);
let s = b.ins().bswap(v);
b.ins().uextend(types::I64, s)
}
4 => {
let v = b.ins().load(types::I32, mf, hptr, 0);
let s = b.ins().bswap(v);
b.ins().uextend(types::I64, s)
}
8 => {
let v = b.ins().load(types::I64, mf, hptr, 0);
b.ins().bswap(v)
}
_ => unreachable!("bad load size {size}"),
}
}
/// Slow-path load: call the read trampoline and zero-extend to `I64`.
fn emit_slow_load(
b: &mut FunctionBuilder,
ec: &EmitCtx,
slow_fref: FuncRef,
ea: Value,
size: u8,
) -> Value {
let v = call_load(b, ec, slow_fref, ea);
if size == 8 {
v // read64 already returns I64
} else {
b.ins().uextend(types::I64, v)
}
}
/// Emit a load with an inline fast path. When the memory exposes a flat mapping
/// (`membase != 0`) and the address is proven **non-MMIO and committed**, load
/// directly from `membase + ea` (byte-swapped) with no call; otherwise call the
/// read trampoline (`slow_fref`). Returns the value zero-extended to `I64`.
///
/// The mapped check is mandatory — unmapped guest pages are `PROT_NONE`, so a
/// raw load there would fault. When `membase == 0` (the overlay/mock path)
/// everything goes slow, preserving interception.
fn emit_inline_load(
b: &mut FunctionBuilder,
ec: &EmitCtx,
ea: Value,
size: u8,
slow_fref: FuncRef,
) -> Value {
let membase = load_env_i64(b, ec, offset_of!(MemEnv, membase));
let chk_mmio = b.create_block();
let chk_map = b.create_block();
let fast = b.create_block();
let slow = b.create_block();
let merge = b.create_block();
let result = b.append_block_param(merge, types::I64);
// membase == 0 → fast path disabled (overlay/mock).
let dis = b.ins().icmp_imm(IntCC::Equal, membase, 0);
b.ins().brif(dis, slow, &[], chk_mmio, &[]);
// (ea & mmio_mask) == mmio_value ⇒ maybe MMIO ⇒ slow.
b.switch_to_block(chk_mmio);
b.seal_block(chk_mmio);
let mask = b.ins().load(types::I32, MemFlags::trusted(), ec.memenv, off(offset_of!(MemEnv, mmio_mask)));
let value = b.ins().load(types::I32, MemFlags::trusted(), ec.memenv, off(offset_of!(MemEnv, mmio_value)));
let masked = b.ins().band(ea, mask);
let is_mmio = b.ins().icmp(IntCC::Equal, masked, value);
b.ins().brif(is_mmio, slow, &[], chk_map, &[]);
// page_table[ea >> 12] COMMIT bit (bit 49) set ⇒ mapped ⇒ fast.
b.switch_to_block(chk_map);
b.seal_block(chk_map);
let pt = load_env_i64(b, ec, offset_of!(MemEnv, page_table));
let page = b.ins().ushr_imm(ea, 12);
let page64 = b.ins().uextend(types::I64, page);
let poff = b.ins().ishl_imm(page64, 3);
let entry_ptr = b.ins().iadd(pt, poff);
let raw = b.ins().load(types::I64, MemFlags::trusted(), entry_ptr, 0);
let commit = b.ins().band_imm(raw, 1i64 << 49);
let committed = b.ins().icmp_imm(IntCC::NotEqual, commit, 0);
b.ins().brif(committed, fast, &[], slow, &[]);
// Fast: raw load from membase + ea.
b.switch_to_block(fast);
b.seal_block(fast);
let ea64 = b.ins().uextend(types::I64, ea);
let hptr = b.ins().iadd(membase, ea64);
let fv = emit_native_load(b, hptr, size);
b.ins().jump(merge, &[BlockArg::Value(fv)]);
// Slow: trampoline through MemoryAccess.
b.switch_to_block(slow);
b.seal_block(slow);
let sv = emit_slow_load(b, ec, slow_fref, ea, size);
b.ins().jump(merge, &[BlockArg::Value(sv)]);
b.seal_block(merge);
b.switch_to_block(merge);
result
}
/// Emit IR for one covered instruction. Must mirror [`covered`] and the exact
/// semantics of the interpreter's `execute()` (validated by the diff harness).
fn emit_op(b: &mut FunctionBuilder, ec: &EmitCtx, instr: &DecodedInstr) {
@@ -881,48 +1039,42 @@ fn emit_op(b: &mut FunctionBuilder, ec: &EmitCtx, instr: &DecodedInstr) {
// ---- integer loads (zero-extended into the 64-bit GPR). ----
PpcOpcode::lbz => {
let ea = ea_d(b, ec, instr.ra(), instr.d());
let v = call_load(b, ec, ec.tr.read8, ea);
let z = b.ins().uextend(types::I64, v);
store_gpr(b, ctxp, instr.rd(), z);
let v = emit_inline_load(b, ec, ea, 1, ec.tr.read8);
store_gpr(b, ctxp, instr.rd(), v);
}
PpcOpcode::lbzx => {
let ea = ea_x(b, ec, instr.ra(), instr.rb());
let v = call_load(b, ec, ec.tr.read8, ea);
let z = b.ins().uextend(types::I64, v);
store_gpr(b, ctxp, instr.rd(), z);
let v = emit_inline_load(b, ec, ea, 1, ec.tr.read8);
store_gpr(b, ctxp, instr.rd(), v);
}
PpcOpcode::lhz => {
let ea = ea_d(b, ec, instr.ra(), instr.d());
let v = call_load(b, ec, ec.tr.read16, ea);
let z = b.ins().uextend(types::I64, v);
store_gpr(b, ctxp, instr.rd(), z);
let v = emit_inline_load(b, ec, ea, 2, ec.tr.read16);
store_gpr(b, ctxp, instr.rd(), v);
}
PpcOpcode::lhzx => {
let ea = ea_x(b, ec, instr.ra(), instr.rb());
let v = call_load(b, ec, ec.tr.read16, ea);
let z = b.ins().uextend(types::I64, v);
store_gpr(b, ctxp, instr.rd(), z);
let v = emit_inline_load(b, ec, ea, 2, ec.tr.read16);
store_gpr(b, ctxp, instr.rd(), v);
}
PpcOpcode::lwz => {
let ea = ea_d(b, ec, instr.ra(), instr.d());
let v = call_load(b, ec, ec.tr.read32, ea);
let z = b.ins().uextend(types::I64, v);
store_gpr(b, ctxp, instr.rd(), z);
let v = emit_inline_load(b, ec, ea, 4, ec.tr.read32);
store_gpr(b, ctxp, instr.rd(), v);
}
PpcOpcode::lwzx => {
let ea = ea_x(b, ec, instr.ra(), instr.rb());
let v = call_load(b, ec, ec.tr.read32, ea);
let z = b.ins().uextend(types::I64, v);
store_gpr(b, ctxp, instr.rd(), z);
let v = emit_inline_load(b, ec, ea, 4, ec.tr.read32);
store_gpr(b, ctxp, instr.rd(), v);
}
PpcOpcode::ld => {
let ea = ea_d(b, ec, instr.ra(), instr.ds());
let v = call_load(b, ec, ec.tr.read64, ea);
let v = emit_inline_load(b, ec, ea, 8, ec.tr.read64);
store_gpr(b, ctxp, instr.rd(), v);
}
PpcOpcode::ldx => {
let ea = ea_x(b, ec, instr.ra(), instr.rb());
let v = call_load(b, ec, ec.tr.read64, ea);
let v = emit_inline_load(b, ec, ea, 8, ec.tr.read64);
store_gpr(b, ctxp, instr.rd(), v);
}
@@ -971,14 +1123,16 @@ fn emit_op(b: &mut FunctionBuilder, ec: &EmitCtx, instr: &DecodedInstr) {
// ---- FP loads. lfs/lfsx: load single, widen to the f64 FPR. ----
PpcOpcode::lfs => {
let ea = ea_d(b, ec, instr.ra(), instr.d());
let bits = call_load(b, ec, ec.tr.read32, ea);
let bits64 = emit_inline_load(b, ec, ea, 4, ec.tr.read32);
let bits = b.ins().ireduce(types::I32, bits64);
let f32v = b.ins().bitcast(types::F32, MemFlags::new(), bits);
let f64v = b.ins().fpromote(types::F64, f32v);
b.ins().store(MemFlags::trusted(), f64v, ctxp, fpr_off(instr.rd()));
}
PpcOpcode::lfsx => {
let ea = ea_x(b, ec, instr.ra(), instr.rb());
let bits = call_load(b, ec, ec.tr.read32, ea);
let bits64 = emit_inline_load(b, ec, ea, 4, ec.tr.read32);
let bits = b.ins().ireduce(types::I32, bits64);
let f32v = b.ins().bitcast(types::F32, MemFlags::new(), bits);
let f64v = b.ins().fpromote(types::F64, f32v);
b.ins().store(MemFlags::trusted(), f64v, ctxp, fpr_off(instr.rd()));
@@ -986,12 +1140,12 @@ fn emit_op(b: &mut FunctionBuilder, ec: &EmitCtx, instr: &DecodedInstr) {
// lfd/lfdx: 64-bit — a pure bit copy into the FPR.
PpcOpcode::lfd => {
let ea = ea_d(b, ec, instr.ra(), instr.d());
let bits = call_load(b, ec, ec.tr.read64, ea);
let bits = emit_inline_load(b, ec, ea, 8, ec.tr.read64);
b.ins().store(MemFlags::trusted(), bits, ctxp, fpr_off(instr.rd()));
}
PpcOpcode::lfdx => {
let ea = ea_x(b, ec, instr.ra(), instr.rb());
let bits = call_load(b, ec, ec.tr.read64, ea);
let bits = emit_inline_load(b, ec, ea, 8, ec.tr.read64);
b.ins().store(MemFlags::trusted(), bits, ctxp, fpr_off(instr.rd()));
}
@@ -1161,7 +1315,7 @@ mod tests {
let f = jit.compile(&block).expect("compile");
let mut jit_ctx = PpcContext::new();
jit_ctx.pc = 0x1000;
let env = MemEnv { mem: &mem };
let env = MemEnv::new(&mem);
let r = f(&mut jit_ctx as *mut _, &env as *const _);
assert_eq!(r, RET_CONTINUE);
@@ -1170,4 +1324,66 @@ mod tests {
assert_eq!(jit_ctx.cycle_count, ref_ctx.cycle_count, "cycle mismatch");
assert_eq!(jit_ctx.timebase, ref_ctx.timebase, "timebase mismatch");
}
// lwz rD, D(rA) = (32<<26)|(rD<<21)|(rA<<16)|(D & 0xFFFF)
fn lwz(rd: u32, ra: u32, d: u16) -> u32 {
(32 << 26) | (rd << 21) | (ra << 16) | d as u32
}
fn b_self() -> u32 {
18 << 26 // b . (LI=0) — covered terminator
}
/// Exercises the inline load fast path against a real `GuestMemory` (the
/// diff harness can't — its overlay disables the fast path), asserting the
/// JIT-loaded value bit-matches the interpreter for both a mapped fast-path
/// address and an unmapped slow-path address.
#[test]
fn jit_inline_load_matches_interpreter() {
use xenia_memory::page_table::MemoryProtect;
use xenia_memory::GuestMemory;
let mem = GuestMemory::new().expect("reserve 4GB");
// Commit a data page and a code page.
mem.alloc(0x40000, 0x1000, MemoryProtect::READ | MemoryProtect::WRITE).unwrap();
mem.alloc(0x41000, 0x1000, MemoryProtect::READ | MemoryProtect::WRITE).unwrap();
mem.write_u32(0x40010, 0xDEAD_BEEF); // big-endian store
// Code at 0x41000: lwz r3, 0x10(r4) ; b .
mem.write_u32(0x41000, lwz(3, 4, 0x10));
mem.write_u32(0x41004, b_self());
let mut bc = BlockCache::new();
let block = bc.lookup_or_build(0x41000, &mem).clone();
assert!(block_covered(&block), "lwz+b block should be covered");
let mut jit = Jit::new().expect("jit init");
let f = jit.compile(&block).expect("compile");
let env = MemEnv::new(&mem);
assert!(!env.membase.is_null(), "GuestMemory must expose fast_mem");
// Mapped address → fast path.
let mut ref_ctx = PpcContext::new();
ref_ctx.pc = 0x41000;
ref_ctx.gpr[4] = 0x40000;
let _ = crate::interpreter::step_block(&mut ref_ctx, &mem, &block);
let mut jit_ctx = PpcContext::new();
jit_ctx.pc = 0x41000;
jit_ctx.gpr[4] = 0x40000;
f(&mut jit_ctx as *mut _, &env as *const _);
assert_eq!(jit_ctx.gpr[3], 0xDEAD_BEEF, "fast-path load value");
assert_eq!(jit_ctx.gpr[3], ref_ctx.gpr[3], "fast-path vs interpreter");
// Unmapped address → slow path returns 0 (no fault).
let mut ref2 = PpcContext::new();
ref2.pc = 0x41000;
ref2.gpr[4] = 0x9000_0000; // unmapped
let _ = crate::interpreter::step_block(&mut ref2, &mem, &block);
let mut jit2 = PpcContext::new();
jit2.pc = 0x41000;
jit2.gpr[4] = 0x9000_0000;
f(&mut jit2 as *mut _, &env as *const _);
assert_eq!(jit2.gpr[3], ref2.gpr[3], "slow-path (unmapped) vs interpreter");
}
}

2
crates/xenia-cpu/src/recompiler.rs
View File

@@ -150,7 +150,7 @@ pub fn run_block(
) -> StepResult {
if let Some(f) = jit.get_or_compile(block) {
JIT_COMPILED_RUN.fetch_add(1, Ordering::Relaxed);
let env = MemEnv { mem };
let env = MemEnv::new(mem);
let raw = f(ctx as *mut PpcContext, &env as *const MemEnv);
// Covered blocks are straight-line (branch/sc/trap/db16cyc are all
// uncovered), so a compiled block always runs to completion and

34
crates/xenia-memory/src/access.rs
View File

@@ -1,3 +1,29 @@
/// Raw pointers/constants a JIT needs to inline the load fast-path (skip the
/// trait-object call). Only the real [`crate::heap::GuestMemory`] exposes this;
/// wrappers that intercept accesses (e.g. the recompiler's speculative overlay)
/// return `None` from [`MemoryAccess::fast_mem`] so the JIT falls back to the
/// trait-object slow path and their interception is preserved.
///
/// The flat mapping means a guest address `a` translates to `membase + a`. A
/// load is safe to inline only when `a` is **not** MMIO and its page is
/// committed: `(a & mmio_mask) != mmio_value` AND page-table entry
/// `page_table[a >> 12]` has the COMMIT bit (bit 49). Unmapped pages are
/// `PROT_NONE`, so the committed check must precede the raw load.
#[derive(Clone, Copy)]
pub struct FastMem {
/// Base of the flat 4 GiB guest mapping. `membase + guest_addr` is the host
/// address. Never null for a real mapping (used by the JIT as the
/// "inline enabled" sentinel).
pub membase: *mut u8,
/// Base of the per-page allocation table (`AtomicU64` entries, reinterpreted
/// as `u64`). Entry `page_table[addr >> 12]` bit 49 = COMMIT.
pub page_table: *const u64,
/// MMIO fast-reject mask/value: `(addr & mask) == value` is the *necessary*
/// condition for `addr` to be MMIO, so `!=` proves non-MMIO.
pub mmio_mask: u32,
pub mmio_value: u32,
}
/// Trait for all guest memory access. Every load/store goes through this,
/// enabling MMIO checking and debugger observation on every access.
/// This is the key abstraction that eliminates the need for MMIO exception handlers.
@@ -58,6 +84,14 @@ pub trait MemoryAccess {
false
}
/// Raw pointers/constants for a JIT to inline the load fast-path, or `None`
/// to force the trait-object slow path. Default `None` (mock memories and
/// interception wrappers have no flat mapping to expose). `GuestMemory`
/// overrides it. See [`FastMem`].
fn fast_mem(&self) -> Option<FastMem> {
None
}
/// Get a direct host pointer for the given guest address.
/// Returns None if the address is invalid or in an MMIO region.
fn translate(&self, addr: u32) -> Option<*const u8>;

13
crates/xenia-memory/src/heap.rs
View File

@@ -509,6 +509,19 @@ impl MemoryAccess for GuestMemory {
self.find_mmio(addr).is_some()
}
fn fast_mem(&self) -> Option<crate::access::FastMem> {
Some(crate::access::FastMem {
membase: self.membase,
// AtomicU64 is #[repr(transparent)] over u64, so the element pointer
// reinterprets directly. The JIT reads entries with a plain load —
// sound because alloc publishes each entry before the guest can
// access the page (single-thread lockstep; x86 loads are acquire).
page_table: self.page_table.as_ptr() as *const u64,
mmio_mask: self.mmio_aperture_mask,
mmio_value: self.mmio_aperture_value,
})
}
// Tier-3 perf: `#[inline]` on the hot read/write paths lets LLVM
// fold the MMIO + mapping checks into the interpreter's load/store
// handlers, hoisting the "not-MMIO, mapped" branch out of the loop

2
crates/xenia-memory/src/lib.rs
View File

@@ -7,7 +7,7 @@ mod platform;
use thiserror::Error;
pub use access::MemoryAccess;
pub use access::{FastMem, MemoryAccess};
pub use heap::{set_writer_ctx, GuestMemory, HeapType};
pub use mmio::MmioRegion;
pub use page_table::PageEntry;