diff --git a/majit/majit-backend-wasm/js/jit_glue.js b/majit/majit-backend-wasm/js/jit_glue.js index 1ac791ccfb8..50456cae8cf 100644 --- a/majit/majit-backend-wasm/js/jit_glue.js +++ b/majit/majit-backend-wasm/js/jit_glue.js @@ -25,16 +25,16 @@ export function jit_set_table(table) { // Call trampoline — invoked by generated wasm when it encounters a CALL op. // Reads func_ptr + args from the frame's call area, performs the call via // the main module's indirect function table, writes result back. -function jitCallTrampoline(framePtr) { +function jitCallTrampoline(framePtr, callAreaOfs = CALL_RESULT_OFS) { const view = new DataView(mainMemory.buffer); - const funcPtrLo = view.getUint32(framePtr + CALL_FUNC_OFS, true); - const numArgs = Number(view.getBigInt64(framePtr + CALL_NARGS_OFS, true)); + const funcPtrLo = view.getUint32(framePtr + callAreaOfs + 8, true); + const numArgs = Number(view.getBigInt64(framePtr + callAreaOfs + 16, true)); // Read args from call area const args = []; for (let i = 0; i < numArgs; i++) { // On wasm32, values are i32 (pointers/ints). Read low 32 bits of each i64 slot. - args.push(view.getInt32(framePtr + CALL_ARGS_OFS + i * 8, true)); + args.push(view.getInt32(framePtr + callAreaOfs + 24 + i * 8, true)); } // Call via the main module's function table @@ -50,8 +50,8 @@ function jitCallTrampoline(framePtr) { } // Write result to call area (as i64: low 32 bits = result, high 32 bits = 0) - view.setInt32(framePtr + CALL_RESULT_OFS, result, true); - view.setInt32(framePtr + CALL_RESULT_OFS + 4, 0, true); + view.setInt32(framePtr + callAreaOfs, result, true); + view.setInt32(framePtr + callAreaOfs + 4, 0, true); } export function jit_compile_wasm(bytesPtr, bytesLen) { @@ -69,7 +69,7 @@ export function jit_compile_wasm(bytesPtr, bytesLen) { // chaining; the module imports it only when it has CALL ops. Extra // entries in the import object are ignored when not declared. const instance = new WebAssembly.Instance(module, { - env: { memory: mainMemory, jit_call: jitCallTrampoline, __indirect_function_table: mainTable } + env: { memory: mainMemory, jit_call: jitCallTrampoline, jit_call_compact: jitCallTrampoline, __indirect_function_table: mainTable } }); return registerTrace(instance.exports.trace); } catch (e) { diff --git a/majit/majit-backend-wasm/src/codegen.rs b/majit/majit-backend-wasm/src/codegen.rs index 7c4adec45de..dea0058d281 100644 --- a/majit/majit-backend-wasm/src/codegen.rs +++ b/majit/majit-backend-wasm/src/codegen.rs @@ -48,6 +48,75 @@ const CALL_ARGS_OFS: u64 = 2024; /// Minimum frame allocation size in bytes to accommodate the call area. pub const MIN_FRAME_BYTES: usize = 2024 + 16 * 8; // 16 max call args +/// Per-token layout of a wasm execution frame. Host entry frames retain the +/// historical [`MIN_FRAME_BYTES`] allocation floor, but generated code and CA +/// nursery frames use these offsets. This mirrors `jitframe.py`'s +/// `JITFRAME_FIXED_SIZE + frame_depth`: a frame carries the slots its token +/// actually needs, rather than a backend-wide slot floor. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct FrameGeometry { + /// Number of data slots before the call trampoline (including frame[0]). + pub value_slots: usize, + /// Byte offset of the call trampoline result word. + pub call_result_ofs: u64, + pub call_func_ofs: u64, + pub call_nargs_ofs: u64, + pub call_args_ofs: u64, + /// Byte offset of the resume-at-LABEL key. + pub dispatch_key_ofs: u64, + /// Byte offset of Ref-home zero. + pub home_slot_base: u64, + /// Number of Ref-home slots the layout reserves. + pub home_slots: usize, + /// Bytes in the JitFrame item area required by this layout. + pub frame_bytes: u32, +} + +impl FrameGeometry { + const CALL_AREA_SLOTS: usize = 3 + 16; // result, function, nargs, args + + /// Historical fixed geometry, used by direct codegen tests and by callers + /// that deliberately need the arena-compatible layout. + pub const fn fixed() -> Self { + Self { + value_slots: MIN_FRAME_BYTES / 8, + call_result_ofs: CALL_RESULT_OFS, + call_func_ofs: CALL_FUNC_OFS, + call_nargs_ofs: CALL_NARGS_OFS, + call_args_ofs: CALL_ARGS_OFS, + dispatch_key_ofs: DISPATCH_KEY_OFS, + home_slot_base: HOME_SLOT_BASE, + home_slots: 0, + frame_bytes: (MIN_FRAME_BYTES + SLOT_SIZE as usize) as u32, + } + } + + /// Compact geometry for one token. `value_slots` includes frame[0], so + /// the call area begins immediately after the greatest positional input or + /// fail-arg slot used by that token. + pub fn compact(value_slots: usize, home_slots: usize) -> Self { + let value_slots = value_slots.max(1); + let call_result_ofs = (value_slots as u64) * SLOT_SIZE; + let call_func_ofs = call_result_ofs + SLOT_SIZE; + let call_nargs_ofs = call_func_ofs + SLOT_SIZE; + let call_args_ofs = call_nargs_ofs + SLOT_SIZE; + let dispatch_key_ofs = call_result_ofs + Self::CALL_AREA_SLOTS as u64 * SLOT_SIZE; + let home_slot_base = dispatch_key_ofs + SLOT_SIZE; + let frame_bytes = home_slot_base + home_slots as u64 * SLOT_SIZE; + Self { + value_slots, + call_result_ofs, + call_func_ofs, + call_nargs_ofs, + call_args_ofs, + dispatch_key_ofs, + home_slot_base, + home_slots, + frame_bytes: frame_bytes as u32, + } + } +} + /// Byte offset of the Ref-home region within the frame. Each Ref value that is /// live across a collecting call is given a dedicated home slot here: it is /// null-initialized at trace entry and written on every definition @@ -87,6 +156,17 @@ fn memarg(offset: u64, align: u32) -> MemArg { } } +/// Invoke the residual-call trampoline. The historical import receives only a +/// frame pointer and therefore reads the fixed call area; compact frames use a +/// second import carrying their call-area base. The old trampoline remains +/// unchanged for fixed-layout frames. +fn emit_jit_call(sink: &mut InstructionSink<'_>, jit_call_idx: u32, frame: FrameGeometry) { + if frame.call_result_ofs != CALL_RESULT_OFS { + sink.i32_const(frame.call_result_ofs as i32); + } + sink.call(jit_call_idx); +} + /// Emit a width-correct integer load. The element address (i32) must be on /// the stack; the result is an i64, sign- or zero-extended from `size` /// bytes. Word-sized fields are 4 bytes on wasm32 (`isize`/`usize`/pointer), @@ -270,6 +350,7 @@ impl RefHomes { fn collect(inputargs: &[InputArg], ops: &[Op], include_ca_collects: bool) -> Self { let liveness = HomeLiveness::collect(inputargs, ops); let collect_positions = collecting_call_positions(ops, include_ca_collects); + let ref_values = RefValues::collect(inputargs, ops); let mut by_id = Vec::new(); let mut next = 0u32; for ia in inputargs { @@ -287,6 +368,20 @@ impl RefHomes { Self::assign(&mut by_id, &mut next, r.raw()); } } + if include_ca_collects { + // The CA arm allocates its callee frame before it resolves this + // CallAssemblerR's arguments. Those Ref operands are used at (not + // after) this op, so ordinary `live_across` deliberately excludes + // them; they nevertheless need homes through the prior allocation. + for op in ops.iter().filter(|op| op.opcode == OpCode::CallAssemblerR) { + for arg in op.getarglist() { + let arg = arg.to_opref(); + if ref_values.contains(arg) { + Self::assign(&mut by_id, &mut next, arg.raw()); + } + } + } + } RefHomes { by_id, len: next as usize, @@ -335,6 +430,19 @@ pub fn count_ref_homes(inputargs: &[InputArg], ops: &[Op]) -> usize { RefHomes::collect(inputargs, ops, true).len() } +/// Positional frame slots required for a token's inputs and guard spills. +/// Slot zero is the fail index; the returned count therefore also gives the +/// first free slot for the call trampoline. +pub fn frame_value_slots(inputargs: &[InputArg], ops: &[Op]) -> usize { + let (guards, _) = collect_guards_and_vars(inputargs, ops); + let max_fail_args = guards + .iter() + .map(|g| g.fail_arg_refs.len()) + .max() + .unwrap_or(0); + 1 + max_fail_args.max(inputargs.len()) +} + /// Argument index of the stored value for a GC ref-storing op. `SetfieldRaw` / /// `SetarrayitemRaw` store into non-GC memory and never need a write barrier, /// so only the `*Gc` variants are listed (rewrite.py only routes `SETFIELD_GC` @@ -385,6 +493,7 @@ fn emit_write_barrier( residual_type_base: Option, wb_fn_ptr: i64, base_ref: OpRef, + frame: FrameGeometry, ) { if let Some(base) = residual_type_base { // Header word is a u64 at `obj - GcHeader::SIZE` with the flags in @@ -419,19 +528,19 @@ fn emit_write_barrier( // func_ptr = wasm_jit_write_barrier sink.local_get(0); sink.i64_const(wb_fn_ptr); - sink.i64_store(mem64(CALL_FUNC_OFS)); + sink.i64_store(mem64(frame.call_func_ofs)); // num_args = 1 (the trampoline reflects arity from the wasm signature; // written for protocol symmetry with the alloc/call paths) sink.local_get(0); sink.i64_const(1); - sink.i64_store(mem64(CALL_NARGS_OFS)); + sink.i64_store(mem64(frame.call_nargs_ofs)); // arg0 = base object pointer sink.local_get(0); emit_resolve(sink, constants, base_ref); - sink.i64_store(mem64(CALL_ARGS_OFS)); + sink.i64_store(mem64(frame.call_args_ofs)); // call trampoline; void result ignored sink.local_get(0); - sink.call(jit_call); + emit_jit_call(sink, jit_call, frame); } /// Per-value def / last-use op positions over the trace, used to filter the @@ -512,19 +621,18 @@ impl HomeLiveness { } /// Static collecting-call positions whose gcmap-visible homes may be forwarded. -/// These are exactly the sites that emit post-call reloads from homes: -/// `New`/`NewWithVtable`, `NewArray`/`NewArrayClear`, and the CA -/// `CallAssemblerR` arm when enabled. General residual calls deliberately stay -/// out of this set because their host-side allocations use the no-collect hook -/// and the codegen emits no reload for them. +/// Alongside `New*`, conservatively include residual calls: an eligible direct +/// residual target may allocate or force, so it needs the same post-call home +/// reload as the allocation helpers. fn collecting_call_positions(ops: &[Op], include_ca_collects: bool) -> Vec { ops.iter() .enumerate() .filter_map(|(i, op)| { - matches!( - op.opcode, - OpCode::New | OpCode::NewWithVtable | OpCode::NewArray | OpCode::NewArrayClear - ) + (op.opcode.is_call() + || matches!( + op.opcode, + OpCode::New | OpCode::NewWithVtable | OpCode::NewArray | OpCode::NewArrayClear + )) .then_some(i) .or_else(|| (include_ca_collects && op.opcode == OpCode::CallAssemblerR).then_some(i)) }) @@ -546,6 +654,7 @@ fn emit_reload_refs_from_homes( liveness: &HomeLiveness, at_op: usize, skip_raw: Option, + frame: FrameGeometry, ) { // `iter` yields id order, so the emitted module is reproducible without a // sort; each reload is independent (home and local storage are disjoint). @@ -554,11 +663,53 @@ fn emit_reload_refs_from_homes( continue; } sink.local_get(0); - sink.i64_load(mem64(HOME_SLOT_BASE + h as u64 * SLOT_SIZE)); + sink.i64_load(mem64(frame.home_slot_base + h as u64 * SLOT_SIZE)); sink.local_set(1 + raw); } } +/// RPython `_reload_frame_if_necessary` (x86 `assembler.py:1369`) for wasm +/// trace bodies: a collecting direct call may have forwarded the running +/// JitFrame, while wasm local 0 still holds its old ITEMS base. +fn emit_reload_frame_if_necessary( + sink: &mut InstructionSink<'_>, + residual_type_base: Option, + ca_reload_fn_ptr: i64, +) { + if let Some(base) = residual_type_base { + sink.i32_const(ca_reload_fn_ptr as i32); + sink.call_indirect(0, base); + sink.i32_wrap_i64(); + sink.local_set(0); + } else { + // The trampoline path still assumes a non-moving frame: its scratch writes use local 0. + } +} + +/// Reload the Ref operands which the CA arm resolves only after its collecting +/// callee-frame allocation. Unlike ordinary post-call reloads, these are live +/// *at* the CALL_ASSEMBLER op, not after it. +fn emit_reload_ca_input_refs_from_homes( + sink: &mut InstructionSink<'_>, + ref_homes: &RefHomes, + ref_values: &RefValues, + op: &Op, + frame: FrameGeometry, +) { + for arg in op.getarglist() { + let arg = arg.to_opref(); + if !ref_values.contains(arg) { + continue; + } + let Some(home) = ref_homes.home(arg) else { + continue; + }; + sink.local_get(0); + sink.i64_load(mem64(frame.home_slot_base + home as u64 * SLOT_SIZE)); + sink.local_set(1 + arg.raw()); + } +} + /// llsupport/gc.py:563 GcLLDescr_framework /// .get_typeid_from_classptr_if_gcremovetypeptr(classptr) /// Looks up the materialized table populated by the runner from the @@ -859,7 +1010,7 @@ pub struct CaParams { pub source_compiled_ptr: u64, /// `__indirect_function_table` slot (`fn as usize`) of /// `lib.rs::wasm_jit_ca_alloc_frame`, which allocates each callee frame as - /// a GC-managed old-gen `JitFrame` (push_jf-rooted, traced by its own + /// a young nursery GC-managed `JitFrame` (push_jf-rooted, traced by its own /// per-frame gcmap). `call_indirect`ed in-module through the residual /// `(i64,i64)->i64` type when declared, else via the `jit_call` trampoline. pub ca_alloc_fn_ptr: i64, @@ -867,6 +1018,14 @@ pub struct CaParams { /// called on CA-arm exit to pop the callee frame off the jitframe shadow /// stack (strict LIFO). pub ca_pop_fn_ptr: i64, + /// `__indirect_function_table` slot of `lib.rs::wasm_jit_ca_reload_frame`, + /// called after the recursive call to recover this level's possibly-moved + /// nursery frame from the jitframe shadow stack. + pub ca_reload_fn_ptr: i64, + /// `__indirect_function_table` slot of + /// `lib.rs::wasm_jit_ca_reload_caller_frame`, called while the callee is + /// still pushed to recover this invocation's possibly-moved local-0 frame. + pub ca_reload_caller_fn_ptr: i64, /// Leaked per-bridge `jf_gcmap` (`lib.rs::build_callee_gcmap`) marking the /// callee frame's CA input + home Ref slots; baked into each frame's /// `jf_gcmap` field at alloc time. @@ -917,6 +1076,10 @@ pub fn build_wasm_module( // Resume-at-LABEL dispatch key for the terminal external JUMP (`target // label ordinal + 1`, or 0 for a non-peeled target); see `build_function`. external_jump_key: u32, + // Frozen layout of the frame this module executes on. A chained bridge + // receives its source token's layout; a loop receives its own compact + // layout at first compilation. + frame: FrameGeometry, // Self-recursive CALL_ASSEMBLER arm parameters (`PYRE_WASM_CA`); `emit_ca` // off keeps the module byte-identical. ca: CaParams, @@ -969,16 +1132,23 @@ pub fn build_wasm_module( .max() .unwrap_or(0); let max_value_slots = 1 + max_fail_args.max(inputargs.len()); - if max_value_slots as u64 > CALL_AREA_FIRST_SLOT { + if max_value_slots > frame.value_slots { return Err(BackendError::Unsupported(format!( - "wasm backend: {max_value_slots} frame value slots reach the call area \ - (limit {CALL_AREA_FIRST_SLOT})" + "wasm backend: {max_value_slots} frame value slots exceed frozen frame layout \ + ({})", + frame.value_slots, ))); } let ref_values = RefValues::collect(inputargs, ops); let ref_homes = RefHomes::collect(inputargs, ops, ca.emit_ca); let num_ref_homes = ref_homes.len(); + if num_ref_homes > frame.home_slots { + return Err(BackendError::Unsupported(format!( + "wasm backend: {num_ref_homes} ref homes exceed frozen frame layout ({})", + frame.home_slots, + ))); + } // Self-recursive CALL_ASSEMBLER arm (`PYRE_WASM_CA`): `bridge_finish_fi` is // THIS bridge's own DoneWithThisFrame index (the recursive return), which the @@ -992,16 +1162,9 @@ pub fn build_wasm_module( .find(|g| g.is_finish) .map(|g| g.fail_index) .unwrap_or(0); - let ca = if ca.emit_ca { - let bridge_base_slots = (MIN_FRAME_BYTES / 8).max(max_value_slots); - let bridge_frame_bytes = ((bridge_base_slots + 1 + num_ref_homes) * 8) as u32; - CaParams { - callee_frame_bytes: ca.callee_frame_bytes.max(bridge_frame_bytes) + 128, - ..ca - } - } else { - ca - }; + // CA frames execute the source loop and this bridge on the same frozen + // geometry. `compile_bridge` rejects a bridge that needs more slots, so + // no global floor or speculative slack is needed here. // A ref-storing store needs the `jit_call` import for its write barrier, // even when the trace has no `New*`/CALL of its own. The CA arm needs it too @@ -1026,10 +1189,16 @@ pub fn build_wasm_module( .filter_map(|op| direct_helper_i64_arity(op, &ref_values)) .max(); if ca.emit_ca { - // The CA arm's frame helpers (`wasm_jit_ca_alloc_frame(frame_bytes, - // gcmap_ptr)` / `wasm_jit_ca_pop_frame(frame_base)`) lower through - // this same `(i64×n)->i64` family; make sure arity 2 is declared. + // The CA arm's frame helpers (`wasm_jit_ca_reload_frame()`, + // `wasm_jit_ca_pop_frame(frame_base)`, and + // `wasm_jit_ca_alloc_frame(frame_bytes, gcmap_ptr)`) lower through + // this same `(i64×n)->i64` family; make sure arity 2 is declared, + // which declares the full 0..=2 range including reload's arity 0. Some(scanned.map_or(2, |m| m.max(2))) + } else if ca.ca_reload_fn_ptr != 0 { + // Every trace body can reload its own frame after a collecting + // direct call, even though only bridges emit the CA arm. + Some(scanned.map_or(0, |m| m.max(0))) } else { scanned } @@ -1044,8 +1213,14 @@ pub fn build_wasm_module( // Type 0: trace function (param i32) -> (result i32) types.ty().function(vec![ValType::I32], vec![ValType::I32]); if needs_call { - // Type 1: jit_call trampoline (param i32) -> () - types.ty().function(vec![ValType::I32], vec![]); + // Type 1: fixed `jit_call(frame)` or compact + // `jit_call_compact(frame, call_area_ofs)` trampoline. + let params = if frame.call_result_ofs == CALL_RESULT_OFS { + vec![ValType::I32] + } else { + vec![ValType::I32, ValType::I32] + }; + types.ty().function(params, vec![]); } // Residual-call types follow: `(i64×n) -> i64` for arity `n`, indexed by // `residual_type_base + n`. `residual_type_base` = the count of types above. @@ -1085,7 +1260,15 @@ pub fn build_wasm_module( ); if needs_call { // Import jit_call trampoline as function index 0 - imports.import("env", "jit_call", EntityType::Function(1)); + imports.import( + "env", + if frame.call_result_ofs == CALL_RESULT_OFS { + "jit_call" + } else { + "jit_call_compact" + }, + EntityType::Function(1), + ); } if needs_table { // Import the host's shared indirect function table as table index 0. @@ -1142,6 +1325,7 @@ pub fn build_wasm_module( fail_index_base, external_jump_slot, external_jump_key, + frame, residual_max_arity.map(|_| residual_type_base), ca, bridge_finish_fi, @@ -1184,6 +1368,7 @@ fn build_function( // target's entry `br_table` lands on that label's resume loader. `0` when // the target is not peeled (no dispatch reads the slot). external_jump_key: u32, + frame: FrameGeometry, // Base wasm type index of the `(i64×n)->i64` residual-call types (type // `residual_type_base + n` for arity `n`), or `None` when the trace has no // eligible residual call / `New*` / write barrier, so those arms always @@ -1233,7 +1418,7 @@ fn build_function( for h in 0..ref_homes.len() as u64 { sink.local_get(0); sink.i64_const(0); - sink.i64_store(mem64(HOME_SLOT_BASE + h * SLOT_SIZE)); + sink.i64_store(mem64(frame.home_slot_base + h * SLOT_SIZE)); } // A peeled loop arrives as `[preamble..][LABEL][body..][JUMP]`: the @@ -1294,7 +1479,7 @@ fn build_function( } sink.block(BlockType::Empty); // D $dispatch sink.local_get(0); - sink.i64_load(mem64(DISPATCH_KEY_OFS)); + sink.i64_load(mem64(frame.dispatch_key_ofs)); sink.i32_wrap_i64(); // Depths at this point, innermost first: D=0, then (C_j, B_j) pairs // with C_j at 2j+1. Entry j+1 of the table targets C_j; entry 0 and @@ -1327,7 +1512,7 @@ fn build_function( if let Some(h) = ref_homes.home_id(ia.index) { sink.local_get(0); sink.local_get(local_idx); - sink.i64_store(mem64(HOME_SLOT_BASE + h as u64 * SLOT_SIZE)); + sink.i64_store(mem64(frame.home_slot_base + h as u64 * SLOT_SIZE)); } } @@ -1369,7 +1554,7 @@ fn build_function( if let Some(h) = ref_homes.home(*la) { sink.local_get(0); sink.local_get(1 + la.raw()); - sink.i64_store(mem64(HOME_SLOT_BASE + h as u64 * SLOT_SIZE)); + sink.i64_store(mem64(frame.home_slot_base + h as u64 * SLOT_SIZE)); } } sink.end(); // end B_j $past_loader @@ -1432,7 +1617,7 @@ fn build_function( // dispatch and ignores the slot (`external_jump_key` 0). sink.local_get(0); // frame_ptr sink.i64_const(external_jump_key as i64); // dispatch key - sink.i64_store(mem64(DISPATCH_KEY_OFS)); + sink.i64_store(mem64(frame.dispatch_key_ofs)); sink.local_get(0); // frame_ptr argument to the loop sink.i32_const(external_jump_slot as i32); // table slot sink.return_call_indirect(0, 0); // table 0, type 0: (i32) -> i32 @@ -1464,7 +1649,7 @@ fn build_function( if let Some(h) = ref_homes.home(*la) { sink.local_get(0); sink.local_get(1 + la.raw()); - sink.i64_store(mem64(HOME_SLOT_BASE + h as u64 * SLOT_SIZE)); + sink.i64_store(mem64(frame.home_slot_base + h as u64 * SLOT_SIZE)); } } sink.br(0); @@ -1872,6 +2057,7 @@ fn build_function( residual_type_base, wb_fn_ptr, base, + frame, ); } emit_resolve(&mut sink, constants, op.arg(0).to_opref()); // struct ptr @@ -1961,6 +2147,7 @@ fn build_function( residual_type_base, wb_fn_ptr, base, + frame, ); } emit_array_addr(&mut sink, constants, op); @@ -2011,6 +2198,7 @@ fn build_function( residual_type_base, wb_fn_ptr, base, + frame, ); } emit_resolve(&mut sink, constants, op.arg(0).to_opref()); @@ -2357,18 +2545,18 @@ fn build_function( let vi = op.pos.get().raw(); sink.local_get(0); emit_resolve(&mut sink, constants, op.arg(0).to_opref()); // length - sink.i64_store(mem64(CALL_ARGS_OFS)); + sink.i64_store(mem64(frame.call_args_ofs)); sink.local_get(0); sink.i64_const(0); // func_ptr = 0 signals "newstr" to host - sink.i64_store(mem64(CALL_FUNC_OFS)); + sink.i64_store(mem64(frame.call_func_ofs)); sink.local_get(0); sink.i64_const(1); - sink.i64_store(mem64(CALL_NARGS_OFS)); + sink.i64_store(mem64(frame.call_nargs_ofs)); sink.local_get(0); - sink.call(jit_call); + emit_jit_call(&mut sink, jit_call, frame); if !OpRef::raw_is_constant(vi) { sink.local_get(0); - sink.i64_load(mem64(CALL_RESULT_OFS)); + sink.i64_load(mem64(frame.call_result_ofs)); sink.local_set(1 + vi); } } @@ -2416,8 +2604,8 @@ fn build_function( // Lower `vi = CallAssemblerR(frame, ec)` into an in-module // `call_indirect` into the SOURCE loop (self-recursion) instead of a // host round-trip. A fresh callee frame is allocated as a real - // GC-managed `JitFrame` (old-gen ⇒ non-moving; push_jf-rooted on the - // jitframe shadow stack; traced by its OWN per-frame gcmap covering + // GC-managed nursery `JitFrame` (push_jf-rooted on the jitframe + // shadow stack; traced by its OWN per-frame gcmap covering // its input + home Ref slots), the two red inputs are written to its // input slots, the loop runs on it (recursing through this same arm // for deeper levels), then the result Ref is read back from output @@ -2436,10 +2624,11 @@ fn build_function( let self_slot = external_jump_slot as i32; // Allocate the callee frame as a GC JitFrame - // (`wasm_jit_ca_alloc_frame(frame_bytes, gcmap_ptr)` — a plain - // `(i64,i64)->i64` table entry that never itself collects, so - // it lowers like an eligible residual call when the type family - // is declared; otherwise via the jit_call trampoline). + // (`wasm_jit_ca_alloc_frame(frame_bytes, gcmap_ptr)` — a + // collecting `(i64,i64)->i64` table entry whose caller's Refs + // are rooted in frame homes, so it lowers like an eligible + // residual call when the type family is declared; otherwise via + // the jit_call trampoline). // `ca_cfp_local = frame_base + FIRST_ITEM_OFFSET` is the // bespoke-layout frame pointer — every `mem64(OFS)` below is // relative to it, exactly as the source loop reads its local 0. @@ -2453,30 +2642,43 @@ fn build_function( jit_call_idx.expect("CA arm needs jit_call for the frame trampolines"); sink.local_get(0); sink.i64_const(ca.ca_alloc_fn_ptr); - sink.i64_store(mem64(CALL_FUNC_OFS)); + sink.i64_store(mem64(frame.call_func_ofs)); sink.local_get(0); sink.i64_const(2); - sink.i64_store(mem64(CALL_NARGS_OFS)); + sink.i64_store(mem64(frame.call_nargs_ofs)); sink.local_get(0); sink.i64_const(ca.callee_frame_bytes as i64); - sink.i64_store(mem64(CALL_ARGS_OFS)); + sink.i64_store(mem64(frame.call_args_ofs)); sink.local_get(0); sink.i64_const(ca.callee_gcmap_ptr); - sink.i64_store(mem64(CALL_ARGS_OFS + SLOT_SIZE)); + sink.i64_store(mem64(frame.call_args_ofs + SLOT_SIZE)); sink.local_get(0); - sink.call(jit_call); + emit_jit_call(&mut sink, jit_call, frame); sink.local_get(0); - sink.i64_load(mem64(CALL_RESULT_OFS)); + sink.i64_load(mem64(frame.call_result_ofs)); } sink.i32_wrap_i64(); sink.i32_const(majit_backend::jitframe::FIRST_ITEM_OFFSET as i32); sink.i32_add(); sink.local_set(ca_cfp_local); + // The collecting callee allocation ran while this invocation's + // own frame was the shadow-stack top. Now that the callee is + // pushed, reload local 0 from the entry beneath it before + // resolving inputs through local-0-relative homes. The + // trampoline path intentionally keeps the A0-era assumption: + // its scratch writes themselves dereference stale local 0. + if let Some(base) = residual_type_base { + sink.i32_const(ca.ca_reload_caller_fn_ptr as i32); + sink.call_indirect(0, base + 0); + sink.i32_wrap_i64(); + sink.local_set(0); + } + emit_reload_ca_input_refs_from_homes(&mut sink, ref_homes, ref_values, op, frame); // dispatch key = 0: run the loop from its entry (preamble), not a // LABEL resume — this is a fresh call. sink.local_get(ca_cfp_local); sink.i64_const(0); - sink.i64_store(mem64(DISPATCH_KEY_OFS)); + sink.i64_store(mem64(frame.dispatch_key_ofs)); // inputs: F'[1] = arg0 (callee frame), F'[2] = arg1 (ec). sink.local_get(ca_cfp_local); emit_resolve(&mut sink, constants, op.arg(0).to_opref()); @@ -2490,6 +2692,29 @@ fn build_function( sink.i32_const(self_slot); sink.call_indirect(0, 0); sink.drop(); + // The recursive call may minor-collect and move this nursery + // callee frame. Deeper levels have already popped, so the + // jitframe shadow-stack top is this level's frame; reload its + // ITEMS base before reading F'[0] or F'[1]. + if let Some(base) = residual_type_base { + sink.i32_const(ca.ca_reload_fn_ptr as i32); + sink.call_indirect(0, base + 0); + } else { + let jit_call = + jit_call_idx.expect("CA arm needs jit_call for the frame trampolines"); + sink.local_get(0); + sink.i64_const(ca.ca_reload_fn_ptr); + sink.i64_store(mem64(frame.call_func_ofs)); + sink.local_get(0); + sink.i64_const(0); + sink.i64_store(mem64(frame.call_nargs_ofs)); + sink.local_get(0); + emit_jit_call(&mut sink, jit_call, frame); + sink.local_get(0); + sink.i64_load(mem64(frame.call_result_ofs)); + } + sink.i32_wrap_i64(); + sink.local_set(ca_cfp_local); // F'[0] is the callee's exit `fail_index`. The base-case loop // finish or this bridge's own recursive finish is a clean // DoneWithThisFrame — the result is already in the callee output @@ -2522,6 +2747,18 @@ fn build_function( // call_indirect(table_index, type_index): the shared table is 0. sink.call_indirect(0, ca_helper_type_idx); sink.end(); + // The recursive call or deopt helper may have collected and + // moved this invocation's own frame. Reload it before the pop + // trampoline and post-call home loads address local 0. As above, + // the trampoline-only configuration retains its A0-era stale- + // local-0 limitation because its scratch writes cannot reload it + // safely. + if let Some(base) = residual_type_base { + sink.i32_const(ca.ca_reload_caller_fn_ptr as i32); + sink.call_indirect(0, base + 0); + sink.i32_wrap_i64(); + sink.local_set(0); + } // store-on-def homes the result Ref (from whichever branch). if !OpRef::raw_is_constant(vi) { sink.local_set(1 + vi); @@ -2543,16 +2780,16 @@ fn build_function( jit_call_idx.expect("CA arm needs jit_call for the frame trampolines"); sink.local_get(0); sink.i64_const(ca.ca_pop_fn_ptr); - sink.i64_store(mem64(CALL_FUNC_OFS)); + sink.i64_store(mem64(frame.call_func_ofs)); sink.local_get(0); sink.i64_const(1); - sink.i64_store(mem64(CALL_NARGS_OFS)); + sink.i64_store(mem64(frame.call_nargs_ofs)); sink.local_get(0); sink.local_get(ca_cfp_local); sink.i64_extend_i32_u(); - sink.i64_store(mem64(CALL_ARGS_OFS)); + sink.i64_store(mem64(frame.call_args_ofs)); sink.local_get(0); - sink.call(jit_call); + emit_jit_call(&mut sink, jit_call, frame); } // The callee recursion minor-collected; this bridge's other live // Ref locals are now stale. Reload them from the forwarded homes. @@ -2560,7 +2797,8 @@ fn build_function( // and its home is not written until the store-on-def below, so a // reload would clobber it with the home's pre-call (stale) value. let skip = (!OpRef::raw_is_constant(vi)).then_some(vi); - emit_reload_refs_from_homes(&mut sink, ref_homes, &liveness, op_idx, skip); + emit_reload_frame_if_necessary(&mut sink, residual_type_base, ca.ca_reload_fn_ptr); + emit_reload_refs_from_homes(&mut sink, ref_homes, &liveness, op_idx, skip, frame); } // ── CALL operations (via trampoline) ── @@ -2595,9 +2833,9 @@ fn build_function( // `call_indirect` the callee's table slot with a static // `(i64×n)->i64` type. The residual ABI is uniformly i64 for // Int/Ref args+result, so args/result move on the wasm stack with - // no marshalling and no call-area traffic. The callee uses the - // *no-collect* nursery hook (like the trampoline path), so no Ref - // reload is needed. Falls back below when ineligible. + // no marshalling and no call-area traffic. A direct target may + // collect or force, so reload local 0 and its live Ref homes on + // return. Falls back below when ineligible. if let (Some(base), Some(nargs)) = (residual_type_base, residual_call_i64_arity(op)) { let call_args = &op.getarglist()[1..]; @@ -2614,6 +2852,19 @@ fn build_function( } else { sink.drop(); // value-producing call whose result is unused } + emit_reload_frame_if_necessary( + &mut sink, + residual_type_base, + ca.ca_reload_fn_ptr, + ); + emit_reload_refs_from_homes( + &mut sink, + ref_homes, + &liveness, + op_idx, + (!OpRef::raw_is_constant(vi)).then_some(vi), + frame, + ); // store-on-def (end of loop) homes a Ref result, so the // direct path must NOT `continue` past it. } else if let (Some(base), Some(nargs)) = @@ -2630,6 +2881,14 @@ fn build_function( sink.i32_wrap_i64(); sink.call_indirect(0, base + nargs as u32); sink.drop(); + emit_reload_frame_if_necessary( + &mut sink, + residual_type_base, + ca.ca_reload_fn_ptr, + ); + emit_reload_refs_from_homes( + &mut sink, ref_homes, &liveness, op_idx, None, frame, + ); } else { let jit_call = jit_call_idx.expect("CALL op present but jit_call not imported"); @@ -2640,23 +2899,23 @@ fn build_function( // Store func_ptr to call area sink.local_get(0); emit_resolve(&mut sink, constants, func_ptr_ref); - sink.i64_store(mem64(CALL_FUNC_OFS)); + sink.i64_store(mem64(frame.call_func_ofs)); // Store num_args sink.local_get(0); sink.i64_const(call_args.len() as i64); - sink.i64_store(mem64(CALL_NARGS_OFS)); + sink.i64_store(mem64(frame.call_nargs_ofs)); // Store each arg for (i, arg) in call_args.iter().enumerate() { sink.local_get(0); emit_resolve(&mut sink, constants, arg.to_opref()); - sink.i64_store(mem64(CALL_ARGS_OFS + i as u64 * SLOT_SIZE)); + sink.i64_store(mem64(frame.call_args_ofs + i as u64 * SLOT_SIZE)); } // Call trampoline sink.local_get(0); - sink.call(jit_call); + emit_jit_call(&mut sink, jit_call, frame); // Read result (for non-void calls) let is_void = matches!( @@ -2670,14 +2929,9 @@ fn build_function( ); if !OpRef::raw_is_constant(vi) && !is_void { sink.local_get(0); - sink.i64_load(mem64(CALL_RESULT_OFS)); + sink.i64_load(mem64(frame.call_result_ofs)); sink.local_set(1 + vi); } - // No reload after a residual call: the interpreter's host-side - // allocations use the *no-collect* nursery hook (their callers - // hold unrooted raw pointers), so a residual callee never moves - // objects. Only `New*` (which uses the collecting allocator) - // needs a reload. } } @@ -2742,12 +2996,18 @@ fn build_function( sink.i64_const(size); sink.i32_const(alloc_fn_ptr as i32); sink.call_indirect(0, base + 2); + emit_reload_frame_if_necessary( + &mut sink, + residual_type_base, + ca.ca_reload_fn_ptr, + ); emit_reload_refs_from_homes( &mut sink, ref_homes, &liveness, op_idx, (!OpRef::raw_is_constant(vi)).then_some(vi), + frame, ); sink.else_(); // Commit: *nursery_free = free + total. @@ -2798,27 +3058,27 @@ fn build_function( // func_ptr = wasm_jit_alloc sink.local_get(0); sink.i64_const(alloc_fn_ptr); - sink.i64_store(mem64(CALL_FUNC_OFS)); + sink.i64_store(mem64(frame.call_func_ofs)); // num_args = 2 sink.local_get(0); sink.i64_const(2); - sink.i64_store(mem64(CALL_NARGS_OFS)); + sink.i64_store(mem64(frame.call_nargs_ofs)); // arg0 = type_id sink.local_get(0); sink.i64_const(type_id); - sink.i64_store(mem64(CALL_ARGS_OFS)); + sink.i64_store(mem64(frame.call_args_ofs)); // arg1 = size sink.local_get(0); sink.i64_const(size); - sink.i64_store(mem64(CALL_ARGS_OFS + SLOT_SIZE)); + sink.i64_store(mem64(frame.call_args_ofs + SLOT_SIZE)); // call trampoline sink.local_get(0); - sink.call(jit_call); + emit_jit_call(&mut sink, jit_call, frame); if !OpRef::raw_is_constant(vi) { // result pointer sink.local_get(0); - sink.i64_load(mem64(CALL_RESULT_OFS)); + sink.i64_load(mem64(frame.call_result_ofs)); sink.local_set(1 + vi); } } @@ -2851,7 +3111,14 @@ fn build_function( // nothing). if residual_type_base.is_none() || inline_nursery.is_none() { let skip = (!OpRef::raw_is_constant(vi)).then_some(vi); - emit_reload_refs_from_homes(&mut sink, ref_homes, &liveness, op_idx, skip); + emit_reload_frame_if_necessary( + &mut sink, + residual_type_base, + ca.ca_reload_fn_ptr, + ); + emit_reload_refs_from_homes( + &mut sink, ref_homes, &liveness, op_idx, skip, frame, + ); } } OpCode::NewArray | OpCode::NewArrayClear => { @@ -2947,12 +3214,18 @@ fn build_function( sink.i64_const(len_offset); sink.i32_const(alloc_array_fn_ptr as i32); sink.call_indirect(0, base + 5); + emit_reload_frame_if_necessary( + &mut sink, + residual_type_base, + ca.ca_reload_fn_ptr, + ); emit_reload_refs_from_homes( &mut sink, ref_homes, &liveness, op_idx, (!OpRef::raw_is_constant(vi)).then_some(vi), + frame, ); sink.else_(); // Commit: *nursery_free = free + total. @@ -3009,12 +3282,18 @@ fn build_function( sink.i64_const(len_offset); sink.i32_const(alloc_array_fn_ptr as i32); sink.call_indirect(0, base + 5); + emit_reload_frame_if_necessary( + &mut sink, + residual_type_base, + ca.ca_reload_fn_ptr, + ); emit_reload_refs_from_homes( &mut sink, ref_homes, &liveness, op_idx, (!OpRef::raw_is_constant(vi)).then_some(vi), + frame, ); sink.else_(); // total = round_up_8(max(header + base + item * length, @@ -3066,12 +3345,18 @@ fn build_function( sink.i64_const(len_offset); sink.i32_const(alloc_array_fn_ptr as i32); sink.call_indirect(0, base + 5); + emit_reload_frame_if_necessary( + &mut sink, + residual_type_base, + ca.ca_reload_fn_ptr, + ); emit_reload_refs_from_homes( &mut sink, ref_homes, &liveness, op_idx, (!OpRef::raw_is_constant(vi)).then_some(vi), + frame, ); sink.else_(); // Commit: *nursery_free = new_free. @@ -3134,38 +3419,38 @@ fn build_function( // func_ptr = wasm_jit_alloc_array sink.local_get(0); sink.i64_const(alloc_array_fn_ptr); - sink.i64_store(mem64(CALL_FUNC_OFS)); + sink.i64_store(mem64(frame.call_func_ofs)); // num_args = 5 sink.local_get(0); sink.i64_const(5); - sink.i64_store(mem64(CALL_NARGS_OFS)); + sink.i64_store(mem64(frame.call_nargs_ofs)); // arg0 = type_id sink.local_get(0); sink.i64_const(type_id); - sink.i64_store(mem64(CALL_ARGS_OFS)); + sink.i64_store(mem64(frame.call_args_ofs)); // arg1 = base_size sink.local_get(0); sink.i64_const(base_size); - sink.i64_store(mem64(CALL_ARGS_OFS + SLOT_SIZE)); + sink.i64_store(mem64(frame.call_args_ofs + SLOT_SIZE)); // arg2 = item_size sink.local_get(0); sink.i64_const(item_size); - sink.i64_store(mem64(CALL_ARGS_OFS + 2 * SLOT_SIZE)); + sink.i64_store(mem64(frame.call_args_ofs + 2 * SLOT_SIZE)); // arg3 = length (op.arg(0)) sink.local_get(0); emit_resolve(&mut sink, constants, op.arg(0).to_opref()); - sink.i64_store(mem64(CALL_ARGS_OFS + 3 * SLOT_SIZE)); + sink.i64_store(mem64(frame.call_args_ofs + 3 * SLOT_SIZE)); // arg4 = len_offset sink.local_get(0); sink.i64_const(len_offset); - sink.i64_store(mem64(CALL_ARGS_OFS + 4 * SLOT_SIZE)); + sink.i64_store(mem64(frame.call_args_ofs + 4 * SLOT_SIZE)); // call trampoline sink.local_get(0); - sink.call(jit_call); + emit_jit_call(&mut sink, jit_call, frame); if !OpRef::raw_is_constant(vi) { sink.local_get(0); - sink.i64_load(mem64(CALL_RESULT_OFS)); + sink.i64_load(mem64(frame.call_result_ofs)); sink.local_set(1 + vi); } } @@ -3175,7 +3460,14 @@ fn build_function( || (inline_nursery_total.is_none() && inline_nursery_varsize.is_none()) { let skip = (!OpRef::raw_is_constant(vi)).then_some(vi); - emit_reload_refs_from_homes(&mut sink, ref_homes, &liveness, op_idx, skip); + emit_reload_frame_if_necessary( + &mut sink, + residual_type_base, + ca.ca_reload_fn_ptr, + ); + emit_reload_refs_from_homes( + &mut sink, ref_homes, &liveness, op_idx, skip, frame, + ); } } @@ -3273,7 +3565,7 @@ fn build_function( if let Some(h) = ref_homes.home(result) { sink.local_get(0); sink.local_get(1 + result.raw()); - sink.i64_store(mem64(HOME_SLOT_BASE + h as u64 * SLOT_SIZE)); + sink.i64_store(mem64(frame.home_slot_base + h as u64 * SLOT_SIZE)); } } diff --git a/majit/majit-backend-wasm/src/failguard.rs b/majit/majit-backend-wasm/src/failguard.rs index d935bac03a4..ce20315bcec 100644 --- a/majit/majit-backend-wasm/src/failguard.rs +++ b/majit/majit-backend-wasm/src/failguard.rs @@ -85,26 +85,12 @@ pub struct LabelTarget { /// livelock advance-check applies; earlier labels execute the peeled /// segment, which advances the state by itself. pub is_last_label: bool, - /// Ref-home slot count of the owning loop, for the chain-soundness check: - /// a chained trace runs in the frame `execute_token` sized for the loop - /// the chain ENTERED through, so a tail-call may only target a loop whose - /// Ref-home region fits `max(source loop's homes, FRAME_REF_HOME_FLOOR)` - /// — inductively, every loop in a chain then fits the entry frame (which - /// is sized to at least the floor). Ref homes are the ONLY variable frame - /// requirement: value slots are bounded by `codegen`'s - /// `CALL_AREA_FIRST_SLOT` decline, below the constant - /// `MIN_FRAME_BYTES / 8` value region every host frame carries. - pub num_ref_homes: usize, + /// Frozen frame geometry of the target token. A tail-call can only reuse + /// a frame when its offsets agree exactly, not merely when its allocation + /// is large enough. + pub frame: crate::codegen::FrameGeometry, } -/// Minimum Ref-home slot count `execute_token` sizes every host frame for -/// while bridge chaining is enabled (see `LabelTarget::num_ref_homes`). -/// Chains between traces whose home counts stay at or under this floor need -/// no source-vs-target comparison at all, which is what lifts most frame-fit -/// bridge declines. Costs `8` bytes + one GC-root registration per slot per -/// `execute_token` call. -pub const FRAME_REF_HOME_FLOOR: usize = 64; - /// Global `frame[0]` fail-index space. /// /// Cross-trace chaining (`LABEL_TARGETS`) means the module that last wrote @@ -230,6 +216,9 @@ pub struct CompiledWasmLoop { /// region (`codegen::HOME_SLOT_BASE`). `execute_token` sizes the host /// frame to include this region and registers each home slot as a GC root. pub num_ref_homes: usize, + /// Geometry frozen when this token was first compiled. Every bridge + /// chained onto it is emitted against this exact layout. + pub frame: crate::codegen::FrameGeometry, /// Base address (shared linear memory) of this loop's per-guard bridge-slot /// cell array — one i32 per `fail_index`, `0` = no bridge. The trace's /// epilogue reads `cells[fail_index]` and `compile_bridge` writes a bridge's @@ -288,15 +277,6 @@ pub struct CompiledWasmLoop { /// module lives as long as the source loop it attaches to, so its cells are /// freed when this loop drops. Appended by `compile_bridge`. pub _bridge_owned_cells: RefCell>>, - /// Max `num_ref_homes` over the self-recursive `CallAssemblerR` bridges - /// (`PYRE_WASM_CA`) chained onto this loop, or 0 when there are none. Such a - /// bridge runs in the host entry frame `F0` for the outermost call, so - /// `execute_token` must size `F0` (and register its GC roots) for the LARGER - /// of the loop's own homes and this — the bridge's home writes would - /// otherwise overflow a loop-sized `F0`. Set by `compile_bridge` when it - /// accepts a CA bridge; `Cell` because the source token is shared (`&`) and - /// the wasm host is single-threaded. - pub ca_bridge_ref_homes: Cell, /// Set when `compile_bridge` accepts a self-recursive `CallAssemblerR` /// bridge (`PYRE_WASM_CA`) for this loop. While set, `compile_bridge` /// declines chaining any FURTHER bridge into this recursion (the guard diff --git a/majit/majit-backend-wasm/src/lib.rs b/majit/majit-backend-wasm/src/lib.rs index 3e699413be9..e2bcaded1a8 100644 --- a/majit/majit-backend-wasm/src/lib.rs +++ b/majit/majit-backend-wasm/src/lib.rs @@ -50,6 +50,13 @@ fn diag_bump(i: usize) { BRIDGE_DIAG[i].fetch_add(1, Ordering::Relaxed); } +// A source token is compiled before a later guard may become a CA bridge. +// Freeze modest room for that bridge at first compilation; a later trace that +// exceeds either bound is declined rather than changing the live frame's +// offsets. Fib's bridge needs about 20 positional slots and 8 homes. +const FROZEN_CHAIN_VALUE_SLOTS: usize = 32; +const FROZEN_CHAIN_REF_HOMES: usize = 16; + /// An arithmetic op whose result advances a loop-carried numeric value (the /// `IntAdd`/`IntSub`/… and float-arithmetic block plus the overflow-checked /// variants and the unary `IntNeg`/`IntInvert`). Excludes copies (`SameAs*`), @@ -449,118 +456,16 @@ pub extern "C" fn wasm_jit_write_barrier(obj: i64) -> i64 { 0 } -/// Plain unsynchronized static cell for CA runtime-helper state. These -/// helpers only ever execute inside the wasm32 guest, which is -/// single-threaded (native builds compile them for the codegen unit tests -/// but never call them — `execute_token` requires a wasm host, see the -/// `func_handle` placeholder arm in `compile_loop`). There is no -/// thread-locality to model, so `thread_local!` here was asserting a -/// per-thread identity the state never had; same invariant as the arena -/// globals in pyre-jit's `call_jit`. -/// -/// Safety contract for [`GuestCell::get_mut`]: single-threaded caller and no -/// re-entrant second reference to the same cell while one is live. -struct GuestCell(std::cell::UnsafeCell); -unsafe impl Sync for GuestCell {} -impl GuestCell { - const fn new(value: T) -> Self { - Self(std::cell::UnsafeCell::new(value)) - } - /// See the safety contract on [`GuestCell`]. - #[allow(clippy::mut_from_ref)] - unsafe fn get_mut(&self) -> &mut T { - unsafe { &mut *self.0.get() } - } -} - -/// Per-gcmap cache of the contiguous marked-item runs to re-zero on CA -/// frame reuse: `(gcmap_ptr, &[(first_item, item_count)])`. One gcmap is -/// built per CA bridge and leaked for the program's life, so a linear -/// scan over a handful of entries resolves in one pointer compare and -/// the derived runs can be leaked alongside it. -static CA_GCMAP_ZERO_RUNS: GuestCell> = - GuestCell::new(Vec::new()); - -/// CA callee frames as `(frame_addr, alloc_capacity_bytes)`, live-prefix -/// + pooled-suffix in one Vec: `entries[..live]` are the live frames in -/// recursion order (mirroring the CA entries on the jf shadow stack) and -/// `entries[live..]` is the LIFO reuse pool with its top at -/// `entries[live]`. Strict CA recursion order means a returning frame is -/// retired IN PLACE by decrementing `live` and the next call revives it -/// by incrementing `live` back, so alloc and pop each touch this single -/// static once and steady-state recursion performs no allocator calls at -/// all. Entries stay registered as libc jitframes and are never freed -/// (bounded by peak recursion depth). -static CA_FRAMES: GuestCell = GuestCell::new(CaFrames { - entries: Vec::new(), - live: 0, -}); - -/// Backing store for [`CA_FRAMES`]: `entries[..live]` live, `entries[live..]` -/// pooled (top at `entries[live]`). -struct CaFrames { - entries: Vec<(usize, usize)>, - live: usize, -} - -/// Contiguous runs of gcmap-marked jf_frame items for a CA callee gcmap, as -/// `(first_item, item_count)` pairs — the slots [`wasm_jit_ca_alloc_frame`] -/// must re-zero when reusing a pooled frame. Derived once per gcmap (built by -/// [`build_callee_gcmap`], leaked per bridge) and cached: the marked items are -/// the CA input slots and the Ref-home region, i.e. two dense runs, so frame -/// reuse re-zeroes them with a couple of bulk `memory.fill`s instead of either -/// a whole-frame memset or a per-bit scatter loop (both measurably slower). -fn ca_gcmap_zero_runs(gcmap_ptr: usize) -> &'static [(usize, usize)] { - let cache = unsafe { CA_GCMAP_ZERO_RUNS.get_mut() }; - if let Some(&(_, runs)) = cache.iter().find(|&&(p, _)| p == gcmap_ptr) { - return runs; - } - // Marked items sit one per 8-byte slot at Signed (4-byte on wasm32) - // item granularity, i.e. stride 2 — never bit-contiguous. Merging - // across small gaps turns each dense region (inputs, homes) into one - // run; zeroing the unmarked gap items is sound (a strict subset of - // the whole-frame zero this replaces). - const MERGE_SLACK: usize = 16; - let mut runs: Vec<(usize, usize)> = Vec::new(); - unsafe { - let gcmap = gcmap_ptr as *const usize; - let num_words = *gcmap; - let bits_per_word = usize::BITS as usize; - for w in 0..num_words { - let mut word = *gcmap.add(1 + w); - while word != 0 { - let bit = word.trailing_zeros() as usize; - word &= word - 1; - let item = w * bits_per_word + bit; - match runs.last_mut() { - Some((first, count)) if item <= *first + *count + MERGE_SLACK => { - *count = item - *first + 1; - } - _ => runs.push((item, 1)), - } - } - } - } - let leaked: &'static [(usize, usize)] = Box::leak(runs.into_boxed_slice()); - cache.push((gcmap_ptr, leaked)); - leaked -} - /// Self-recursive CALL_ASSEMBLER (`PYRE_WASM_CA`) callee-frame allocation -/// helper. Allocates the callee's execution frame as a **libc-jitframe** -/// (malloc memory, like dynasm's `execute_token` calloc frames — registered -/// via `register_libc_jitframe` so the collector's jf-root walk traces its -/// gcmap-marked Ref slots), initializes its header + per-frame `jf_gcmap` -/// (covering the callee's input + home Ref slots), and pushes it on the -/// jitframe shadow stack so a mid-recursion collection forwards those Refs -/// via the registered libc-jitframe tracer. Malloc (not nursery/old-gen) -/// keeps the frame non-moving AND off the GC's `bytes_made_old_since_cycle` -/// accounting — a per-call old-gen frame made every recursion level look -/// like heap growth and drove back-to-back major collections (fib: 1704 -/// majors for one bench run). Frames are pooled LIFO on pop and re-zeroed on -/// reuse, so steady-state recursion performs no allocator calls at all. -/// Returns the frame base (codegen adds `FIRST_ITEM_OFFSET` for the -/// bespoke-layout frame pointer), or 0 on allocation failure. +/// helper. Allocates the callee's execution frame as a young nursery +/// GC-managed `JitFrame`, mirroring rewrite.py's nursery frame allocation: +/// steady recursive frames die young, while only frames alive across a +/// collection are promoted. The frame is traced through the jitframe type id's +/// custom trace using its per-frame `jf_gcmap`, rooted by pushing it on the +/// jitframe shadow stack, and reloaded after the recursive call because a +/// nursery frame may move. Returns the frame base (codegen adds +/// `FIRST_ITEM_OFFSET` for the bespoke-layout frame pointer), or 0 on +/// allocation failure. /// /// Each callee frame self-describes through its own per-frame gcmap, so /// mixed-geometry frames from distinct CA bridges are each forwarded by their @@ -569,105 +474,81 @@ fn ca_gcmap_zero_runs(gcmap_ptr: usize) -> &'static [(usize, usize)] { pub extern "C" fn wasm_jit_ca_alloc_frame(frame_bytes: i64, gcmap_ptr: i64) -> i64 { use majit_backend::jitframe::JitFrame; let depth = frame_bytes as usize / std::mem::size_of::(); - let alloc_size = JitFrame::alloc_size(depth); - // Reuse the pool top (`entries[live]`) when it is large enough (`>=` also - // covers a smaller bridge nesting inside a larger one) by reviving it in - // place — bump `live`. A too-small top is left in place: the fresh frame - // is inserted at `live`, becoming the new live top with the old pool top - // right below it in LIFO order. - // `(frame_addr, needs_rezero)`: a revived pool frame carries the previous - // run's bytes; a fresh `alloc_zeroed` frame does not. - let got = { - let f = unsafe { CA_FRAMES.get_mut() }; - if let Some(&(addr, _)) = f.entries.get(f.live).filter(|&&(_, cap)| cap >= alloc_size) { - f.live += 1; - Some((addr, true)) - } else { - let layout = std::alloc::Layout::from_size_align(alloc_size, 16) - .expect("CA frame layout overflow"); - let p = unsafe { std::alloc::alloc_zeroed(layout) }; - if p.is_null() { - None - } else { - majit_gc::shadow_stack::register_libc_jitframe(p as usize); - let live = f.live; - f.entries.insert(live, (p as usize, alloc_size)); - f.live += 1; - Some((p as usize, false)) - } - } - }; - let addr = match got { - Some((addr, true)) => { - // Re-zero only what a reused frame must present as zeroed: - // * the JitFrame header — `JitFrame::init` writes only - // jf_frame_info + the jf_frame length and relies on the other - // header fields (jf_descr, jf_guard_exc, jf_forward, ...) - // being zero (jitframe.py:48-52 "other fields are zero from - // malloc"), and - // * the gcmap-marked items — stale Refs from a previous run - // must not reach the jf-root tracer, and the trace's - // store-on-def home discipline assumes null-initialized homes. - // Every unmarked jf_frame slot is trace data the compiled code - // writes before any read (inputs by the CA arm, outputs/spills by - // the trace itself), so its stale bytes are unobservable. The - // marked items form contiguous runs (the input slots and the - // Ref-home region), pre-derived per gcmap by - // [`ca_gcmap_zero_runs`], so this is a couple of small - // `memory.fill`s instead of the whole `alloc_size` (which was - // ~35% of a recursive CALL_ASSEMBLER call). - use majit_backend::jitframe::{FIRST_ITEM_OFFSET, JITFRAME_FIXED_SIZE, SIGN_SIZE}; - unsafe { - std::ptr::write_bytes(addr as *mut u8, 0, JITFRAME_FIXED_SIZE); - let items = (addr + FIRST_ITEM_OFFSET) as *mut u8; - for &(first, count) in ca_gcmap_zero_runs(gcmap_ptr as usize) { - std::ptr::write_bytes(items.add(first * SIGN_SIZE), 0, count * SIGN_SIZE); - } - } - addr - } - Some((addr, false)) => addr, - None => return 0, - }; - let jf = addr as *mut JitFrame; + // Slice A1: collecting nursery allocation, matching rewrite.py's + // `gen_malloc_nursery_varsize_frame`. The caller frame remains rooted at + // the shadow-stack top during a collection; wasm reloads it from there + // after this call, then this freshly allocated callee is pushed below its + // own execution. Steady recursive frames die young; only frames that live + // through a collection are promoted instead of inflating the old-gen major + // collection threshold on every call. + let jf_ref = WASM_ACTIVE_GC.with(|cell| match cell.borrow_mut().as_deref_mut() { + Some(gc) => gc.alloc_nursery_typed(wasm_jitframe_tid(), JitFrame::alloc_size(depth)), + None => GcRef(0), + }); + if jf_ref.0 == 0 { + return 0; + } + let jf = jf_ref.0 as *mut JitFrame; unsafe { JitFrame::init(jf, std::ptr::null(), depth); (*jf).jf_gcmap = gcmap_ptr as *const u8; } - majit_gc::shadow_stack::push_jf(GcRef(addr)); - addr as i64 + majit_gc::shadow_stack::push_jf(jf_ref); + jf_ref.0 as i64 } /// Companion to [`wasm_jit_ca_alloc_frame`]: pop the top jitframe shadow-stack -/// entry on CA-arm exit (the callee frame just ran to finish/deopt) and retire -/// the frame into the reuse pool — in place, by decrementing the [`CA_FRAMES`] -/// live watermark. The CA recursion is strict LIFO — each level pushes one -/// frame before its `call_indirect` and pops after, and a deopt resume runs on -/// the host's own shadow stack — so removing the top entry releases exactly +/// entry on CA-arm exit. The CA recursion is strict LIFO — each level pushes +/// one frame before its `call_indirect` and pops after, and a deopt resume runs +/// on the host's own shadow stack — so removing the top entry releases exactly /// this callee's frame. pub extern "C" fn wasm_jit_ca_pop_frame(_frame_base: i64) -> i64 { majit_gc::shadow_stack::pop_jf_top(); - let f = unsafe { CA_FRAMES.get_mut() }; - f.live = f.live.saturating_sub(1); 0 } -/// Build the per-frame `jf_gcmap` for a CA callee frame: mark the CA input slots -/// (`v64` + `ec`, at `FRAME_SLOT_BASE`) AND the surviving home slots (Refs live -/// across collecting calls), in the `JitFrame`'s Signed-granular item indexing -/// (see [`build_home_gcmap`] for the wasm32 layout). Unlike the host-entry frame -/// F0 (homes only), a callee frame keeps its virtualizable `v64` in an input slot -/// (never homed by the loop), so the input slots are roots too. Returned buffer -/// is leaked by the caller (one per bridge) and lives for the program's life. -fn build_callee_gcmap(input_count: usize, home_count: usize) -> Box<[usize]> { +/// Reload the current CA callee frame pointer after a recursive call. The GC +/// may have moved the callee frame during the recursive call; `jf_top_ptr()` +/// reads the forwarded base from the jitframe shadow-stack slot. At this point +/// this recursion level's frame is the top — deeper levels have already popped. +/// Analog of `_reload_frame_if_necessary`; returns the ITEMS base held in the +/// CA arm's `ca_cfp_local`. +pub extern "C" fn wasm_jit_ca_reload_frame() -> i64 { + majit_gc::shadow_stack::jf_top_ptr().0 as i64 + + majit_backend::jitframe::FIRST_ITEM_OFFSET as i64 +} + +/// Reload the CA caller's frame pointer after the callee-frame allocation. +/// The allocation occurs before the callee is pushed, so while the callee is +/// live the caller remains one entry below the shadow-stack top. Returns that +/// caller's ITEMS base for local 0. +pub extern "C" fn wasm_jit_ca_reload_caller_frame() -> i64 { + majit_gc::shadow_stack::jf_under_top_ptr().0 as i64 + + majit_backend::jitframe::FIRST_ITEM_OFFSET as i64 +} + +/// Build the per-frame `jf_gcmap` for a CA callee frame: mark the input slots +/// (at `FRAME_SLOT_BASE`) and the home slots (at `HOME_SLOT_BASE`), in the +/// `JitFrame`'s Signed-granular item indexing (see [`build_home_gcmap`] for the +/// wasm32 layout). The collector's `is_nursery_object_start` gate skips any +/// marked slot that does not hold a live nursery object base, so a slot holding +/// a scalar or an already-promoted Ref is traced harmlessly. +/// +/// Returned buffer is leaked by the caller (one per bridge) and lives for the +/// program's life. +fn build_callee_gcmap( + input_types: &[majit_ir::Type], + frame: codegen::FrameGeometry, +) -> Box<[usize]> { let sign = std::mem::size_of::(); let bits_per_word = std::mem::size_of::() * 8; - let mut indices: Vec = Vec::with_capacity(input_count + home_count); + let input_count = input_types.len(); + let mut indices: Vec = Vec::with_capacity(input_count + frame.home_slots); for i in 0..input_count { indices.push((codegen::FRAME_SLOT_BASE as usize + i * 8) / sign); } - for h in 0..home_count { - indices.push((codegen::HOME_SLOT_BASE as usize + h * 8) / sign); + for h in 0..frame.home_slots { + indices.push((frame.home_slot_base as usize + h * 8) / sign); } let max_index = indices.iter().copied().max().unwrap_or(0); let num_words = max_index / bits_per_word + 1; @@ -767,8 +648,7 @@ pub fn set_wasm_jitframe_tid(id: u32) { WASM_JITFRAME_TID.store(id, std::sync::atomic::Ordering::Relaxed); } -// Only read on the wasm32 execute_token path (CA frame allocs use libc -// jitframes and no longer consume the tid). +// Only read on the wasm32 execute_token path and by CA callee-frame allocation. #[cfg_attr(not(target_arch = "wasm32"), allow(dead_code))] fn wasm_jitframe_tid() -> u32 { WASM_JITFRAME_TID.load(std::sync::atomic::Ordering::Relaxed) @@ -793,19 +673,19 @@ fn wasm_jitframe_tid() -> u32 { reason = "native test builds compile the wasm backend without running wasm frame entry" ) )] -fn build_home_gcmap(home_count: usize) -> Box<[usize]> { +fn build_home_gcmap(frame: codegen::FrameGeometry) -> Box<[usize]> { let sign = std::mem::size_of::(); let bits_per_word = std::mem::size_of::() * 8; - if home_count == 0 { + if frame.home_slots == 0 { // One empty data word: a non-null jf_gcmap that traces nothing. return vec![1usize, 0usize].into_boxed_slice(); } - let last_index = (codegen::HOME_SLOT_BASE as usize + (home_count - 1) * 8) / sign; + let last_index = (frame.home_slot_base as usize + (frame.home_slots - 1) * 8) / sign; let num_words = last_index / bits_per_word + 1; let mut buf = vec![0usize; 1 + num_words]; buf[0] = num_words; - for h in 0..home_count { - let index = (codegen::HOME_SLOT_BASE as usize + h * 8) / sign; + for h in 0..frame.home_slots { + let index = (frame.home_slot_base as usize + h * 8) / sign; buf[1 + index / bits_per_word] |= 1usize << (index % bits_per_word); } buf.into_boxed_slice() @@ -1268,6 +1148,14 @@ impl majit_backend::Backend for WasmBackend { let alloc_fn_ptr = wasm_jit_alloc as *const () as usize as i64; let alloc_array_fn_ptr = wasm_jit_alloc_array as *const () as usize as i64; let wb_fn_ptr = wasm_jit_write_barrier as *const () as usize as i64; + // Freeze this token's generated frame layout at first compilation. + // `jitframe.py` sizes native JitFrames as `JITFRAME_FIXED_SIZE + + // frame_depth`; wasm CA frames follow the same per-token depth instead + // of inheriting the host arena's call-area floor. + let frame = codegen::FrameGeometry::compact( + codegen::frame_value_slots(inputargs, ops).max(FROZEN_CHAIN_VALUE_SLOTS), + codegen::count_ref_homes(inputargs, ops).max(FROZEN_CHAIN_REF_HOMES), + ); // Exit indices come from the global fail-index space so a cross-trace // chain's `frame[0]` resolves regardless of which module wrote it // (`failguard::FAIL_DESCR_REGISTRY`). @@ -1287,7 +1175,13 @@ impl majit_backend::Backend for WasmBackend { fail_index_base, 0, // external_jump_slot: a loop's JUMP is a local back-edge `br` 0, // external_jump_key: unused without an external JUMP - codegen::CaParams::default(), // a loop never emits the CA arm + frame, + codegen::CaParams { + // Loops do not emit the CA arm, but they can run on a + // nursery CA frame and must reload local 0 after a collect. + ca_reload_fn_ptr: wasm_jit_ca_reload_frame as *const () as usize as i64, + ..codegen::CaParams::default() + }, )?; // Build fail descriptors @@ -1412,7 +1306,7 @@ impl majit_backend::Backend for WasmBackend { num_args: label_num_args[j], resume_safe: label_resume_safe[j], is_last_label: j == last, - num_ref_homes, + frame, }, ); } @@ -1429,7 +1323,7 @@ impl majit_backend::Backend for WasmBackend { // an entry re-run lands at the header without any // advancing segment — the livelock check applies. is_last_label: true, - num_ref_homes, + frame, }, ); } @@ -1443,6 +1337,7 @@ impl majit_backend::Backend for WasmBackend { num_inputs: inputargs.len(), max_output_slots, num_ref_homes, + frame, bridge_cells_base, num_guard_cells: guard_exits.len(), has_preamble, @@ -1452,7 +1347,6 @@ impl majit_backend::Backend for WasmBackend { chained_trace_meta: std::cell::RefCell::new(std::collections::HashMap::new()), _bridge_cells_owner: bridge_cells_owner, _bridge_owned_cells: std::cell::RefCell::new(Vec::new()), - ca_bridge_ref_homes: std::cell::Cell::new(0), ca_active: std::cell::Cell::new(false), }; @@ -1552,11 +1446,10 @@ impl majit_backend::Backend for WasmBackend { let ( source_guard, source_is_direct, - source_num_ref_homes, source_func_handle, source_has_preamble, - source_max_output_slots, - source_num_inputs, + source_frame, + source_input_types, source_loop_finish_fi, source_compiled_ptr, source_ca_active, @@ -1618,11 +1511,10 @@ impl majit_backend::Backend for WasmBackend { ( guard, is_direct, - source_loop.num_ref_homes, source_loop.func_handle, source_loop.has_preamble, - source_loop.max_output_slots, - source_loop.num_inputs, + source_loop.frame, + source_loop.input_types.clone(), loop_finish_fi, // Address of the source loop's metadata, baked into the CA arm // (opaque cookie in the deopt-helper ABI; `frame[0]` resolution @@ -1669,6 +1561,25 @@ impl majit_backend::Backend for WasmBackend { )); } + // A chained bridge executes in the source token's *same* frame. Its + // offsets are frozen when that token is compiled, so accept it only if + // its positional spill region and Ref-home region fit exactly within + // that layout. Declining here preserves the normal blackhole fallback; + // it is never safe to grow an already-allocated CA frame underneath a + // later bridge. + let bridge_value_slots = codegen::frame_value_slots(inputargs, ops); + let bridge_ref_homes = codegen::count_ref_homes(inputargs, ops); + if bridge_value_slots > source_frame.value_slots + || bridge_ref_homes > source_frame.home_slots + { + diag_bump(4); + return Err(BackendError::Unsupported(format!( + "wasm backend: bridge frame needs values={bridge_value_slots}, homes={bridge_ref_homes}; \ + source frozen layout has values={}, homes={}", + source_frame.value_slots, source_frame.home_slots, + ))); + } + // A loop-closing bridge (terminal JUMP, no local LABEL) re-enters the // source loop through `source_func_handle` — the function entry. For a // peeled source loop, entering at the function entry re-runs the preamble @@ -1712,15 +1623,8 @@ impl majit_backend::Backend for WasmBackend { // reads exactly that many positional frame slots), the label's args // are not the complete live set of the target trace's remainder // (`resume_safe` — resuming there would read a null local), or the - // target loop's Ref-home region exceeds what the chain's entry frame - // is guaranteed to carry (`max(source homes, FRAME_REF_HOME_FLOOR)` — - // the frame `execute_token` sized for the loop the chain entered - // through; requiring target ≤ that bound keeps every hop within the - // entry frame by induction, since the entry frame itself is sized to - // at least the floor). Ref homes are the only variable requirement: - // value slots are bounded by codegen's `CALL_AREA_FIRST_SLOT` decline, - // below the constant `MIN_FRAME_BYTES / 8` value region every host - // frame carries. A declined guard falls back to blackhole resume and + // target loop's frozen frame geometry differs from the source frame. + // A declined guard falls back to blackhole resume and // `declined_bridge_guards` stops the metainterp re-tracing it. let mut external_jump_key: u32 = 0; let mut external_jump_slot: u32 = source_func_handle; @@ -1750,11 +1654,8 @@ impl majit_backend::Backend for WasmBackend { diag_bump(9); // label args not the full live set false } - Some(t) - if t.num_ref_homes - > source_num_ref_homes.max(failguard::FRAME_REF_HOME_FLOOR) => - { - diag_bump(4); // target Ref homes exceed the entry frame's bound + Some(t) if t.frame != source_frame => { + diag_bump(4); // target uses different frozen frame offsets false } Some(t) => { @@ -1869,68 +1770,39 @@ impl majit_backend::Backend for WasmBackend { let alloc_array_fn_ptr = wasm_jit_alloc_array as *const () as usize as i64; let wb_fn_ptr = wasm_jit_write_barrier as *const () as usize as i64; - // Self-recursive CALL_ASSEMBLER (PYRE_WASM_CA): the CA arm bump-allocates - // a fresh callee frame per recursive `call_indirect` into the source - // loop. Size it for the source loop's frame layout (base_slots + the - // dispatch-key slot + surviving ref-home region, mirroring - // `execute_token`); - // `build_wasm_module` widens it to also fit THIS bridge, which reuses the - // same frame when the loop's guard-exit chains back into it. - // This bridge materializes the recursive callee frame and homes any - // bridge Refs live across collecting calls in the SAME arena frame the - // source loop runs on, store-on-def'ing at its OWN dense home indices. A - // self-recursive fib bridge may reserve more surviving homes than the - // source loop, so the arena frame and the GC walker must cover the WIDER - // of the two: otherwise a bridge home (index >= `source_num_ref_homes`) - // lands past the frame's walked region and a minor collection - // mid-recursion reclaims it, leaving a later deopt to read zeroed nursery - // memory. `count_ref_homes` matches the CA-enabled `num_ref_homes` - // `build_wasm_module` returns below. - // - // The recursion can also chain into NESTED bridges (and sibling loops via - // loop-closing tail calls) while running ON a CA callee frame — and those - // were accepted against the `FRAME_REF_HOME_FLOOR` bound `execute_token` - // guarantees for HOST frames. The callee frame must give the same - // guarantee, or a chained bridge homes/reads Ref slots past the frame's - // sized (and gcmap-walked) region — wrong-value corruption (suite - // `recursion_memo_branch` / `generator_tree_recursion`). Mirror - // `execute_token`'s `chain_floor`. - let chain_floor = failguard::FRAME_REF_HOME_FLOOR; - let ca_ref_homes = if allow_ca { - source_num_ref_homes - .max(codegen::count_ref_homes(inputargs, ops)) - .max(chain_floor) - } else { - source_num_ref_homes - }; + // Self-recursive CALL_ASSEMBLER (PYRE_WASM_CA): the CA arm allocates a + // fresh callee using the source token's frozen geometry. The earlier + // frame-fit decline guarantees this bridge uses those same offsets. let ca_params = if allow_ca { - let min_slots = (codegen::MIN_FRAME_BYTES / 8) as u32; - let src_base_slots = - min_slots.max(1 + source_max_output_slots.max(source_num_inputs) as u32); - let src_frame_slots = src_base_slots + 1 + ca_ref_homes as u32; - // Per-bridge callee-frame gcmap (input + home Ref slots), leaked to - // live for the program — each callee frame's `jf_gcmap` points at it. + // Per-bridge callee-frame gcmap (real Ref inputs + home Ref slots), + // leaked to live for the program — each callee frame's `jf_gcmap` + // points at it. let callee_gcmap_ptr = - Box::leak(build_callee_gcmap(source_num_inputs as usize, ca_ref_homes)).as_ptr() - as i64; + Box::leak(build_callee_gcmap(&source_input_types, source_frame)).as_ptr() as i64; codegen::CaParams { emit_ca: true, - callee_frame_bytes: src_frame_slots * 8, + callee_frame_bytes: source_frame.frame_bytes, loop_finish_fi: source_loop_finish_fi, deopt_helper_slot: ca_deopt_helper_slot(), source_compiled_ptr, ca_alloc_fn_ptr: wasm_jit_ca_alloc_frame as *const () as usize as i64, ca_pop_fn_ptr: wasm_jit_ca_pop_frame as *const () as usize as i64, + ca_reload_fn_ptr: wasm_jit_ca_reload_frame as *const () as usize as i64, + ca_reload_caller_fn_ptr: wasm_jit_ca_reload_caller_frame as *const () as usize + as i64, callee_gcmap_ptr, } } else { - codegen::CaParams::default() + codegen::CaParams { + ca_reload_fn_ptr: wasm_jit_ca_reload_frame as *const () as usize as i64, + ..codegen::CaParams::default() + } }; // This bridge's exit indices come from the global fail-index space, // like every trace's (`failguard::FAIL_DESCR_REGISTRY`). let base = fail_descr_base(); - let (wasm_bytes, guard_exits, num_ref_homes, bridge_cells_base, bridge_cells_owner) = + let (wasm_bytes, guard_exits, _num_ref_homes, bridge_cells_base, bridge_cells_owner) = codegen::build_wasm_module( inputargs, ops, @@ -1949,35 +1821,10 @@ impl majit_backend::Backend for WasmBackend { // `external_jump_key` selects. external_jump_slot, external_jump_key, + source_frame, ca_params, )?; - // The bridge runs in the chain's entry frame, so it must not address - // more Ref-home slots than that frame is guaranteed to carry — - // `max(source loop's homes, FRAME_REF_HOME_FLOOR)`, the same inductive - // bound as the JUMP-target check above. If it would, decline: the host - // round-trip path allocates a frame sized for the bridge. - // The bridge's value/output slots need no separate bound check: - // `build_wasm_module` already declined this bridge (above, via `?`) if - // its value slots reach `CALL_AREA_FIRST_SLOT` (codegen.rs), and - // `execute_token` floors the frame at `MIN_FRAME_BYTES/8` slots — which - // exceeds `CALL_AREA_FIRST_SLOT` — before the Ref-home region, so every - // in-bounds value-slot write lands strictly below both the call area and - // the home roots regardless of how the bridge's exit arity compares to - // the source loop's `max_output_slots`. - // A self-recursive CALL_ASSEMBLER bridge (`allow_ca`) is exempt: its - // recursive calls run in arena callee frames sized for it - // (`callee_frame_bytes`), and the outermost call runs in the host entry - // frame `F0`, which `execute_token` widens via `ca_bridge_ref_homes` - // (set below). So its home writes never overflow. - if !allow_ca && num_ref_homes > source_num_ref_homes.max(failguard::FRAME_REF_HOME_FLOOR) { - diag_bump(4); // declined: ref-home overflow - return Err(BackendError::Unsupported(format!( - "wasm backend: bridge needs {num_ref_homes} ref homes, entry frame bound is \ - max({source_num_ref_homes}, floor)" - ))); - } - // Bridge exit descrs (fail_index already base-offset by build_wasm_module). let bridge_descrs: Vec> = guard_exits .iter() @@ -2057,13 +1904,7 @@ impl majit_backend::Backend for WasmBackend { if let Some(owner) = bridge_cells_owner { source_loop._bridge_owned_cells.borrow_mut().push(owner); } - // A self-recursive CALL_ASSEMBLER bridge runs its outermost call in - // the loop's host entry frame `F0`. Record its (possibly larger) - // home count so `execute_token` sizes `F0` and registers GC roots for - // it, not just the loop's own homes. if allow_ca { - let prev = source_loop.ca_bridge_ref_homes.get(); - source_loop.ca_bridge_ref_homes.set(prev.max(num_ref_homes)); // Freeze this recursion to the CA mechanism: no further bridge // chains here (see the decline above the codegen call). source_loop.ca_active.set(true); @@ -2297,22 +2138,15 @@ impl majit_backend::Backend for WasmBackend { // and the Ref-home region (now at HOME_SLOT_BASE = MIN_FRAME_BYTES + 8). // `vec![0i64]` zeroes it, so a fresh host entry reads key 0 (preamble). // - // A self-recursive CALL_ASSEMBLER bridge runs its outermost call in this - // frame and may home more collection-live Refs than the loop, so size - // for the LARGER of the two (`ca_bridge_ref_homes`); the extra slots are - // zeroed and GC-rooted below exactly like the loop's own homes. - // A cross-trace tail call can land in a loop or bridge homing more Refs - // than this one, so also size at least `FRAME_REF_HOME_FLOOR` — the bound - // `compile_bridge`'s frame-fit accept checks rely on. - let chain_floor = failguard::FRAME_REF_HOME_FLOOR; - let eff_ref_homes = compiled - .num_ref_homes - .max(compiled.ca_bridge_ref_homes.get()) - .max(chain_floor); - let frame_size = base_slots + 1 + eff_ref_homes; + // The host/arena allocation deliberately retains its historical floor: + // any token may run here. Generated code, however, addresses only the + // compact offsets frozen on `compiled.frame`; all chained bridges use + // that same geometry or are declined at compile time. + let _ = base_slots; + let frame_size = min_slots.max((compiled.frame.frame_bytes as usize).div_ceil(8)); #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))] { - let _ = (frame_size, eff_ref_homes, args); + let _ = (frame_size, args); panic!("wasm backend execute_token requires a wasm host"); } #[cfg(all(target_arch = "wasm32", not(target_os = "wasi")))] @@ -2351,7 +2185,7 @@ impl majit_backend::Backend for WasmBackend { // Per-loop gcmap over the surviving Ref-home region. Held in this // stack frame (jf_gcmap points at it) until the outputs are read // after the trace returns. - let gcmap = build_home_gcmap(eff_ref_homes); + let gcmap = build_home_gcmap(compiled.frame); unsafe { (*jf).jf_gcmap = gcmap.as_ptr() as *const u8 }; let items_base = jf as usize + FIRST_ITEM_OFFSET; @@ -2412,13 +2246,13 @@ impl majit_backend::Backend for WasmBackend { }; } let frame_ptr = frame.as_mut_ptr() as usize as u32; - let home_base = codegen::HOME_SLOT_BASE as usize / 8; - for h in 0..eff_ref_homes { + let home_base = compiled.frame.home_slot_base as usize / 8; + for h in 0..compiled.frame.home_slots { let slot = unsafe { frame.as_mut_ptr().add(home_base + h) } as *mut GcRef; unsafe { wasm_gc_add_root(slot) }; } glue::execute(compiled.func_handle, frame_ptr); - for h in 0..eff_ref_homes { + for h in 0..compiled.frame.home_slots { let slot = unsafe { frame.as_mut_ptr().add(home_base + h) } as *mut GcRef; wasm_gc_remove_root(slot); } diff --git a/majit/majit-backend-wasm/tests/codegen_test.rs b/majit/majit-backend-wasm/tests/codegen_test.rs index df1f20b6a22..946180e924c 100644 --- a/majit/majit-backend-wasm/tests/codegen_test.rs +++ b/majit/majit-backend-wasm/tests/codegen_test.rs @@ -59,6 +59,7 @@ fn build_module( 0, // fail_index_base 0, // external_jump_slot 0, // external_jump_key + codegen::FrameGeometry::fixed(), codegen::CaParams::default(), ) .expect("wasm codegen should succeed"); diff --git a/majit/majit-gc/src/shadow_stack.rs b/majit/majit-gc/src/shadow_stack.rs index 0f99c44c74e..674114ef48e 100644 --- a/majit/majit-gc/src/shadow_stack.rs +++ b/majit/majit-gc/src/shadow_stack.rs @@ -499,6 +499,29 @@ pub fn jf_top_ptr() -> GcRef { }) } +/// Read the jf_ptr one entry below the top jitframe shadow-stack entry. +/// +/// While a CALL_ASSEMBLER callee is pushed, this is its caller's frame. A +/// collecting allocation performed before the callee push can update this +/// shadow-stack slot, so compiled wasm reloads its own frame from here before +/// addressing local-0-relative frame homes. +pub fn jf_under_top_ptr() -> GcRef { + JF_ROOT_STACK.with(|stack| { + let mut stack = stack.borrow_mut(); + stack.ensure_init(); + unsafe { + let top = stack.top.get(); + if top <= stack.base + 2 * WORD { + return GcRef::NULL; + } + // Each entry is [is_minor, jf_ptr]. The caller's jf_ptr is three + // words below top: top[-1] is this callee, top[-3] its caller. + let jf_ptr_addr = (top - 3 * WORD) as *const usize; + GcRef(*jf_ptr_addr) + } + }) +} + // ── Blackhole register bank shadow stack ──────────────────────── // // blackhole.py:840 BlackholeInterpreter.registers_r parity: diff --git a/majit/majit-macros/src/jit_interp/codegen_state.rs b/majit/majit-macros/src/jit_interp/codegen_state.rs index 70a051f7a5f..c304826e092 100644 --- a/majit/majit-macros/src/jit_interp/codegen_state.rs +++ b/majit/majit-macros/src/jit_interp/codegen_state.rs @@ -1220,10 +1220,10 @@ fn generate_state_fields_jit_state(config: &JitInterpConfig, func: &ItemFn) -> T // Value-routing types in `extract_live` order: int scalars, // int array elements, then per virt-array the identity ptr // (Ref) + length (Int), then the appended ref scalars (Ref). - // The ptr slot MUST be Ref so the folded loop-invariant - // `&state` identity numbers as a ref const (TAGCONST), not an - // int const — the resume reader decodes it through `decode_ref` - // in both the vable section and the frame ref-liveness. + // The ptr slot MUST be Ref so the live `&state` identity is a + // Ref failarg (TAGBOX), which the resume reader decodes through + // `decode_ref` in both the vable section and the frame + // ref-liveness. let mut types: Vec = Vec::new(); for _ in 0..#num_scalars { types.push(majit_ir::Type::Int); @@ -1407,6 +1407,19 @@ fn generate_state_fields_jit_state(config: &JitInterpConfig, func: &ItemFn) -> T Some(self as *const Self as *mut u8) } + fn blackhole_virtualizable_identity( + &self, + _meta: &Self::Meta, + _virtualizable: &str, + _info: &majit_metainterp::virtualizable::VirtualizableInfo, + ) -> Option<*mut u8> { + // This is intentionally distinct from the PyFrame path: + // `state` is the current host-stack object, not a movable GC + // object whose trace-time address may be baked into resume + // data. At deopt, re-derive its identity from this call. + Some(self as *const Self as *mut u8) + } + fn export_virtualizable_boxes( &self, _meta: &Self::Meta, @@ -2157,7 +2170,6 @@ fn generate_state_fields_jit_state(config: &JitInterpConfig, func: &ItemFn) -> T __all_liveness: &[u8], __virtualizable_boxes: &[majit_ir::OpRef], __virtualref_boxes: &[(majit_ir::OpRef, usize)], - __identity_const: Option, ) -> Option { use majit_metainterp::JitCodeSym as _; if frames.frames.is_empty() { @@ -2195,7 +2207,6 @@ fn generate_state_fields_jit_state(config: &JitInterpConfig, func: &ItemFn) -> T false, __virtualizable_boxes, __virtualref_boxes, - __identity_const, Some((sym.int_identity_slots_base(), sym.int_identity_slots_end())), ); let __root = &mut frames.frames[0]; diff --git a/majit/majit-metainterp/src/blackhole.rs b/majit/majit-metainterp/src/blackhole.rs index 2215a105206..23fda889348 100644 --- a/majit/majit-metainterp/src/blackhole.rs +++ b/majit/majit-metainterp/src/blackhole.rs @@ -2424,6 +2424,7 @@ pub fn resume_in_blackhole( None, // vrefinfo None, // vinfo None, // ginfo + None, // virtualizable_identity_override &null_alloc, ); diff --git a/majit/majit-metainterp/src/jit_state.rs b/majit/majit-metainterp/src/jit_state.rs index 653913feeb5..ea2cc9653f9 100644 --- a/majit/majit-metainterp/src/jit_state.rs +++ b/majit/majit-metainterp/src/jit_state.rs @@ -504,6 +504,24 @@ pub trait JitState: Sized { None } + /// Per-call virtualizable identity for blackhole resume when the host + /// keeps its virtualizable in stack storage rather than in a JIT failarg. + /// + /// RPython's virtualizable is a red input, so resume data always obtains + /// this value from its TAGBOX. The state-field macro JIT instead has a + /// host-stack `&state`: its loop-invariant identity can be folded out of + /// the backend's failargs. That host opts in here and supplies the current + /// call's address at deopt entry. Heap virtualizables deliberately keep + /// the default so their TAGBOX remains the source of truth. + fn blackhole_virtualizable_identity( + &self, + _meta: &Self::Meta, + _virtualizable: &str, + _info: &VirtualizableInfo, + ) -> Option<*mut u8> { + None + } + fn virtualizable_array_lengths( &self, _meta: &Self::Meta, @@ -608,7 +626,6 @@ pub trait JitState: Sized { _all_liveness: &[u8], _virtualizable_boxes: &[OpRef], _virtualref_boxes: &[(OpRef, usize)], - _identity_const: Option, ) -> Option { None } diff --git a/majit/majit-metainterp/src/jitdriver.rs b/majit/majit-metainterp/src/jitdriver.rs index dab0df9ff15..ae876d32c7a 100644 --- a/majit/majit-metainterp/src/jitdriver.rs +++ b/majit/majit-metainterp/src/jitdriver.rs @@ -1713,14 +1713,13 @@ impl JitDriver { // active trace ctx before borrowing `self.meta.framestack` // mutably below. Cloning is acceptable because this is // the segmented-loop force path (slow path). - let (vable_boxes, vref_boxes, identity_const) = self + let (vable_boxes, vref_boxes) = self .meta .trace_ctx() .map(|ctx| { ( ctx.virtualizable_boxes.clone().unwrap_or_default(), ctx.virtualref_boxes.clone(), - ctx.state_field_identity_const(), ) }) .unwrap_or_default(); @@ -1732,7 +1731,6 @@ impl JitDriver { &all_liveness, &vable_boxes, &vref_boxes, - identity_const, ) }); let mut current_live = self @@ -2798,6 +2796,15 @@ impl JitDriver { self.meta_interp().staticdata.op_rvmprof_code, ); let all_liveness = self.meta_interp().staticdata.liveness_info.as_slice(); + // The state-field macro's `&state` is host-stack storage, so + // its identity may be folded out of the failing frame. Ask + // only an explicit host opt-in for the current call's address; + // heap virtualizables retain the live resume TAGBOX path. + let vable_identity_override = self.meta.virtualizable_info().and_then(|info| { + state + .blackhole_virtualizable_identity(&compiled_meta, &info.name, info) + .map(|ptr| ptr as i64) + }); let bh = crate::resume::blackhole_from_resumedata( &mut *bh_builder, &resolve_jitcode, @@ -2819,6 +2826,7 @@ impl JitDriver { .virtualizable_info() .map(|a| a.as_ref() as &dyn crate::resume::VirtualizableInfo), None, // ginfo + vable_identity_override, allocator, ); if let Some((mut bh, _vable_ptr)) = bh { @@ -4885,6 +4893,13 @@ impl JitDriver { self.meta_interp().staticdata.op_rvmprof_code, ); let all_liveness = self.meta_interp().staticdata.liveness_info.as_slice(); + // See `back_edge_internal`: only the macro state-field host + // opts in to this per-call host-stack identity source. + let vable_identity_override = self.meta.virtualizable_info().and_then(|info| { + state + .blackhole_virtualizable_identity(&meta, &info.name, info) + .map(|ptr| ptr as i64) + }); let bh = crate::resume::blackhole_from_resumedata( &mut *bh_builder, &resolve_jitcode, @@ -4906,6 +4921,7 @@ impl JitDriver { .virtualizable_info() .map(|a| a.as_ref() as &dyn crate::resume::VirtualizableInfo), None, // ginfo + vable_identity_override, allocator, ); if let Some((mut bh, _vable_ptr)) = bh { diff --git a/majit/majit-metainterp/src/pyjitpl/dispatch.rs b/majit/majit-metainterp/src/pyjitpl/dispatch.rs index 07f8638c825..beb8b7bb314 100644 --- a/majit/majit-metainterp/src/pyjitpl/dispatch.rs +++ b/majit/majit-metainterp/src/pyjitpl/dispatch.rs @@ -1119,7 +1119,6 @@ where // two per live vref). let virtualizable_snapshot = ctx.virtualizable_boxes.clone().unwrap_or_default(); let virtualref_snapshot = ctx.virtualref_boxes.clone(); - let identity_const = ctx.state_field_identity_const(); let snapshot = build_state_field_snapshot( self.frames, op_live, @@ -1127,7 +1126,6 @@ where after_residual_call, &virtualizable_snapshot, &virtualref_snapshot, - identity_const, Some((sym.int_identity_slots_base(), sym.int_identity_slots_end())), ); for idx in 0..n { @@ -7595,7 +7593,6 @@ pub fn build_state_field_snapshot( after_residual_call: bool, virtualizable_boxes: &[OpRef], virtualref_boxes: &[(OpRef, usize)], - identity_const: Option, inline_int_identity: Option<(usize, usize)>, ) -> crate::recorder::Snapshot { let frame_count = frames.frames.len(); @@ -7659,7 +7656,7 @@ pub fn build_state_field_snapshot( // requiring `OpRef::ty()` to be `Some` rather than silently dropping // misshapen entries (which would shrink the snapshot relative to // upstream and desync the resume reader). - let vable_boxes_snap = build_vable_snapshot_boxes(virtualizable_boxes, identity_const); + let vable_boxes_snap = build_vable_snapshot_boxes(virtualizable_boxes); let vref_boxes_snap = build_vref_snapshot_boxes(virtualref_boxes); crate::recorder::Snapshot { frames: snapshot_frames, @@ -7678,31 +7675,19 @@ pub fn build_state_field_snapshot( /// snapshot relative to upstream. pub fn build_vable_snapshot_boxes( virtualizable_boxes: &[OpRef], - identity_const: Option, ) -> Vec { let mut vable_boxes_snap: Vec = Vec::new(); if !virtualizable_boxes.is_empty() { let last = virtualizable_boxes.last().copied().unwrap(); - // The identity is encoded identity-FIRST. For the state-field JIT the - // `&state` identity is a loop-invariant constant the backend drops from - // live registers, so its deadframe slot decodes null at resume; encode - // it as a `Ref` constant (rd_consts TAGCONST) so `consume_vable_info`'s - // `next_ref()` returns the real pointer, matching RPython `_encode`'s - // `isinstance(box, Const)` arm (opencoder.py:603-640). `None` keeps the - // genuinely-live box (heap-object virtualizables like PyFrame, whose - // identity flows through the trace and may be forced/mutated). - match identity_const { - Some(ptr) => vable_boxes_snap.push(crate::recorder::SnapshotTagged::Const( - ptr, - majit_ir::Type::Ref, - )), - None => { - let last_ty = last - .ty() - .expect("build_vable_snapshot_boxes: virtualizable identity must be typed"); - vable_boxes_snap.push(crate::recorder::SnapshotTagged::Box(last, last_ty)); - } - } + // `pyjitpl.py:3319`: the virtualizable identity is the live red input + // box at `virtualizable_boxes[-1]`. It must remain a TAGBOX failarg, + // never a trace-time Ref constant: resume.py:1566-1578 reads the + // current failing jitframe slot, whose GC map keeps a moving-GC pointer + // updated in place. + let last_ty = last + .ty() + .expect("build_vable_snapshot_boxes: virtualizable identity must be typed"); + vable_boxes_snap.push(crate::recorder::SnapshotTagged::Box(last, last_ty)); for opref in &virtualizable_boxes[..virtualizable_boxes.len() - 1] { let ty = opref .ty() @@ -9745,7 +9730,6 @@ mod tests { &[], &[], None, - None, ); assert_eq!(snapshot.frames.len(), 1); @@ -9807,7 +9791,6 @@ mod tests { &[], &[], None, - None, ); let f = &snapshot.frames[0]; @@ -9872,7 +9855,6 @@ mod tests { &[], &[], None, - None, ); assert_eq!(snapshot.frames.len(), 2); let root_frame = &snapshot.frames[0]; @@ -9931,7 +9913,6 @@ mod tests { &[], &[], None, - None, ); let f = &snapshot.frames[0]; @@ -9945,31 +9926,26 @@ mod tests { } #[test] - fn build_vable_snapshot_boxes_encodes_identity_const_as_ref_const() { - // State-field JIT: the loop-invariant `&state` identity (at `[-1]`) is - // folded out of live registers, so it is supplied as a concrete pointer - // and encoded identity-FIRST as a `Ref` constant. `consume_vable_info` - // then reads the real pointer via `next_ref()` (resume.py:1404) with no - // `LIVE_VABLE_PTR` recovery. Non-identity entries keep their `Box` tag. + fn build_vable_snapshot_boxes_encodes_identity_as_live_ref_box() { + // `pyjitpl.py:3319` / `resume.py:192-221`: the identity at `[-1]` is + // identity-FIRST and always a live Box, so resume.py:1566-1578 reads + // the current failing jitframe's ref failarg slot. let other = majit_ir::OpRef::int_op(3); let identity = majit_ir::OpRef::ref_op(7); - let snap = build_vable_snapshot_boxes(&[other, identity], Some(0xdead_beef)); + let snap = build_vable_snapshot_boxes(&[other, identity]); assert_eq!( snap, vec![ - crate::recorder::SnapshotTagged::Const(0xdead_beef, majit_ir::Type::Ref), + crate::recorder::SnapshotTagged::Box(identity, majit_ir::Type::Ref), crate::recorder::SnapshotTagged::Box(other, majit_ir::Type::Int), ] ); } #[test] - fn build_vable_snapshot_boxes_keeps_live_identity_as_box_when_no_const() { - // Heap-object virtualizables (e.g. PyFrame) pass `None`: the - // genuinely-live identity box stays a `Box` snapshot entry because its - // pointer is present in the deadframe and decodes non-null. + fn build_vable_snapshot_boxes_keeps_single_live_identity_as_box() { let identity = majit_ir::OpRef::ref_op(7); - let snap = build_vable_snapshot_boxes(&[identity], None); + let snap = build_vable_snapshot_boxes(&[identity]); assert_eq!( snap, vec![crate::recorder::SnapshotTagged::Box( diff --git a/majit/majit-metainterp/src/resume.rs b/majit/majit-metainterp/src/resume.rs index e7f906dfcce..0ac016b0350 100644 --- a/majit/majit-metainterp/src/resume.rs +++ b/majit/majit-metainterp/src/resume.rs @@ -5263,6 +5263,7 @@ mod tests { None, None, None, + None, &NullAllocator, ) .expect("runtime-only jitcode should still resume"); @@ -5302,7 +5303,7 @@ mod tests { let mut reader = ResumeDataDirectReader::new(&rd_numb, &[], &[], &[], None, None, &NullAllocator); - reader.consume_vref_and_vable(None, Some(&TestVirtualizableInfo), None); + reader.consume_vref_and_vable(None, Some(&TestVirtualizableInfo), None, None); } #[test] @@ -6852,7 +6853,12 @@ impl<'a> ResumeDataDirectReader<'a> { } /// resume.py:1399 consume_vable_info - pub fn consume_vable_info(&mut self, vinfo: &dyn VirtualizableInfo, vable_size: i32) { + pub fn consume_vable_info( + &mut self, + vinfo: &dyn VirtualizableInfo, + vable_size: i32, + identity_override: Option, + ) { // resume.py:1403 assert!(vable_size > 0); // The vable section is encoded identity-FIRST: the snapshot writer @@ -6863,12 +6869,15 @@ impl<'a> ResumeDataDirectReader<'a> { // payload the remaining `vable_size - 1`, read sequentially. // resume.py:1404 virtualizable = self.next_ref() // - // The state-field JIT's `&state` identity is a loop-invariant the - // backend drops from live registers; `build_vable_snapshot_boxes` - // encodes it into the resume snapshot as a `Ref` constant, so it decodes - // to the real pointer here with no thread-local recovery — matching - // resume.py:1404, which reads the identity solely from resume data. - let virtualizable = self.next_ref(); + // Consume the encoded identity even when a host supplies an override: + // it occupies one resume-data item and keeps the reader aligned for + // the field payload. Heap virtualizables (PyFrame) use this live + // TAGBOX exactly as RPython does. The state-field macro JIT opts in + // to `identity_override` because its host-stack `&state` is folded out + // of backend failargs; at deopt it must use the current call's address, + // never a trace-time pointer or an unrelated deadframe slot. + let encoded_identity = self.next_ref(); + let virtualizable = identity_override.unwrap_or(encoded_identity); self.virtualizable_ptr = virtualizable; // resume.py:1406: assert vinfo.get_total_size(virtualizable) == vable_size - 1 let expected = vinfo.get_total_size(virtualizable) as i32; @@ -6893,6 +6902,7 @@ impl<'a> ResumeDataDirectReader<'a> { vrefinfo: Option<&dyn VRefInfo>, vinfo: Option<&dyn VirtualizableInfo>, ginfo: Option<&dyn GreenfieldInfo>, + identity_override: Option, ) { // resume.py:1425 let vable_size = self.resumecodereader.next_item(); @@ -6900,7 +6910,7 @@ impl<'a> ResumeDataDirectReader<'a> { if self.resume_after_guard_not_forced != 2 { // resume.py:1427-1428 if let Some(vi) = vinfo { - self.consume_vable_info(vi, vable_size); + self.consume_vable_info(vi, vable_size, identity_override); } // resume.py:1429-1430 if ginfo.is_some() { @@ -7274,6 +7284,7 @@ pub fn blackhole_from_resumedata<'a>( vrefinfo: Option<&dyn VRefInfo>, vinfo: Option<&dyn VirtualizableInfo>, ginfo: Option<&dyn GreenfieldInfo>, + virtualizable_identity_override: Option, allocator: &'a dyn BlackholeAllocator, ) -> Option<(BlackholeInterpreter, i64)> { // resume.py:1315-1327 The initialization is stack-critical code: it @@ -7328,7 +7339,7 @@ pub fn blackhole_from_resumedata<'a>( } // resume.py:1325 - resumereader.consume_vref_and_vable(vrefinfo, vinfo, ginfo); + resumereader.consume_vref_and_vable(vrefinfo, vinfo, ginfo, virtualizable_identity_override); drop(_cc_guard); // resume.py:1404: virtualizable pointer read by consume_vable_info. @@ -7432,7 +7443,7 @@ pub fn force_from_resumedata<'a>( resumereader.prepare(rd_virtuals, rd_guard_pendingfields); resumereader.handling_async_forcing(); // resume.py:1350 - resumereader.consume_vref_and_vable(vrefinfo, vinfo, ginfo); + resumereader.consume_vref_and_vable(vrefinfo, vinfo, ginfo, None); // resume.py:1351: return resumereader.force_all_virtuals() let (ptrs, ints) = resumereader.force_all_virtuals(); (ptrs.to_vec(), ints.to_vec()) diff --git a/majit/majit-metainterp/src/trace_ctx.rs b/majit/majit-metainterp/src/trace_ctx.rs index 56543ed651c..81492c4fc91 100644 --- a/majit/majit-metainterp/src/trace_ctx.rs +++ b/majit/majit-metainterp/src/trace_ctx.rs @@ -2225,33 +2225,11 @@ impl TraceCtx { Vec, ) { let vable_slice: &[OpRef] = self.virtualizable_boxes.as_deref().unwrap_or(&[]); - let vable_boxes = crate::pyjitpl::build_vable_snapshot_boxes( - vable_slice, - self.state_field_identity_const(), - ); + let vable_boxes = crate::pyjitpl::build_vable_snapshot_boxes(vable_slice); let vref_boxes = crate::pyjitpl::build_vref_snapshot_boxes(&self.virtualref_boxes); (vable_boxes, vref_boxes) } - /// Concrete `&state` pointer to encode as the resume-snapshot identity for - /// the state-field JIT, whose loop-invariant identity is folded out of the - /// live registers and otherwise decodes null at resume. Encoding it as a - /// `Ref` constant lets `consume_vable_info` read the real pointer from - /// resume data, matching `resume.py:1404`. Returns `None` for heap-object - /// virtualizables (e.g. `PyFrame`), whose identity is a genuinely-live box - /// that must stay a `Box` snapshot entry, and when no concrete shadow is - /// available (bridge-entry rebuild leaves `virtualizable_values` unset). - pub(crate) fn state_field_identity_const(&self) -> Option { - let info = self.virtualizable_info.as_ref()?; - if !info.elements_carried_via_shadow() { - return None; - } - match self.standard_virtualizable_concrete()? { - Value::Ref(r) if r.0 != 0 => Some(r.0 as i64), - _ => None, - } - } - /// Concrete shadow of the standard virtualizable — the raw heap pointer /// `standard_virtualizable_box` refers to. Parallels /// `MetaInterp.virtualizable_boxes[-1].getref_base()` at runtime; pyre diff --git a/pyre/pyre-jit/src/call_jit.rs b/pyre/pyre-jit/src/call_jit.rs index 4a5cdd9491e..2a159b795a9 100644 --- a/pyre/pyre-jit/src/call_jit.rs +++ b/pyre/pyre-jit/src/call_jit.rs @@ -1669,6 +1669,7 @@ pub fn blackhole_resume_via_rd_numb( Some(vrefinfo_dyn), // resume.py:1314 metainterp_sd.virtualref_info Some(vinfo_dyn), // resume.py:1312 self.jitdriver_sd.virtualizable_info None, // resume.py:1316 greenfield_info unused in pyre + None, // heap PyFrame identity remains the live TAGBOX &allocator, ) }); @@ -5500,7 +5501,7 @@ pub fn cranelift_resumedata_deopt( reader.prepare(rd_virtuals_slice, rd_pendingfields); let vinfo_dyn: &dyn resume::VirtualizableInfo = driver_vinfo.as_ref(); let vrefinfo_dyn: &dyn resume::VRefInfo = driver.meta_interp().virtualref_info(); - reader.consume_vref_and_vable(Some(vrefinfo_dyn), Some(vinfo_dyn), None); + reader.consume_vref_and_vable(Some(vrefinfo_dyn), Some(vinfo_dyn), None, None); // 7. resume.py:1339 jitcodes[jitcode_pos] lookup — same shape as // blackhole_resume_via_rd_numb's resolve_jitcode (line 1891), diff --git a/pyre/pyre-wasm-runner/src/main.rs b/pyre/pyre-wasm-runner/src/main.rs index aade1c365b2..eb99d7344be 100644 --- a/pyre/pyre-wasm-runner/src/main.rs +++ b/pyre/pyre-wasm-runner/src/main.rs @@ -33,13 +33,11 @@ use wasmtime::{ // Frame call-area offsets — must match `majit-backend-wasm/src/codegen.rs`. // Shared with `wasmi_host`, which mirrors the same call-area protocol. pub(crate) const CALL_RESULT_OFS: usize = 2000; -pub(crate) const CALL_FUNC_OFS: usize = 2008; // The arg count also lives at offset 2016, but the trampoline derives arity // (and the exact value types) from the resolved function's wasm signature // instead, which is authoritative on wasm32. Kept for layout documentation. #[allow(dead_code)] pub(crate) const CALL_NARGS_OFS: usize = 2016; -pub(crate) const CALL_ARGS_OFS: usize = 2024; pub(crate) const DEFAULT_MODULE: &str = "target/wasm32-unknown-unknown/release/pyre_wasm.wasm"; @@ -567,7 +565,7 @@ fn build_linker(engine: &Engine) -> Result> { "pyre_jit", "jit_call_host", |mut caller: Caller<'_, Host>, frame_ptr: u32| { - if let Err(e) = jit_call_trampoline(&mut caller, frame_ptr) { + if let Err(e) = jit_call_trampoline(&mut caller, frame_ptr, CALL_RESULT_OFS as u32) { eprintln!("[jit_call_host] {e:?}"); } }, @@ -718,11 +716,22 @@ fn jit_compile(caller: &mut Caller<'_, Host>, bytes_ptr: u32, bytes_len: u32) -> let jit_call = Func::wrap( &mut *caller, |mut inner: Caller<'_, Host>, frame_ptr: i32| { - if let Err(e) = jit_call_trampoline(&mut inner, frame_ptr as u32) { + if let Err(e) = + jit_call_trampoline(&mut inner, frame_ptr as u32, CALL_RESULT_OFS as u32) + { eprintln!("[jit_call] {e:?}"); } }, ); + let jit_call_compact = Func::wrap( + &mut *caller, + |mut inner: Caller<'_, Host>, frame_ptr: i32, call_area_ofs: i32| { + if let Err(e) = jit_call_trampoline(&mut inner, frame_ptr as u32, call_area_ofs as u32) + { + eprintln!("[jit_call_compact] {e:?}"); + } + }, + ); // Supply imports in the module's declared order. let mut externs: Vec = Vec::new(); @@ -730,6 +739,7 @@ fn jit_compile(caller: &mut Caller<'_, Host>, bytes_ptr: u32, bytes_len: u32) -> match (import.module(), import.name()) { ("env", "memory") => externs.push(Extern::Memory(memory)), ("env", "jit_call") => externs.push(Extern::Func(jit_call)), + ("env", "jit_call_compact") => externs.push(Extern::Func(jit_call_compact)), ("env", "__indirect_function_table") => externs.push(Extern::Table(table)), (m, n) => { return Err(Error::msg(format!( @@ -796,25 +806,29 @@ fn jit_execute(caller: &mut Caller<'_, Host>, func_id: u32, frame_ptr: u32) -> R } /// Dispatch a residual call requested by a running trace. -fn jit_call_trampoline(caller: &mut Caller<'_, Host>, frame_ptr: u32) -> Result<()> { +fn jit_call_trampoline( + caller: &mut Caller<'_, Host>, + frame_ptr: u32, + call_area_ofs: u32, +) -> Result<()> { caller.data_mut().jit_call_count += 1; let memory = caller.data().memory.context("memory")?; let table = caller.data().table.context("table")?; - let frame = frame_ptr as usize; + let call_area = frame_ptr as usize + call_area_ofs as usize; - let func_ptr = read_u32(&memory, &*caller, frame + CALL_FUNC_OFS); + let func_ptr = read_u32(&memory, &*caller, call_area + 8); // `func_ptr == 0` is the "newstr" sentinel; without a host string // allocator (matching the browser glue's null table slot) it yields 0. if func_ptr == 0 { - write_i64(&memory, &mut *caller, frame + CALL_RESULT_OFS, 0)?; + write_i64(&memory, &mut *caller, call_area, 0)?; return Ok(()); } let func = match table.get(&mut *caller, func_ptr as u64) { Some(Ref::Func(Some(f))) => f, _ => { - write_i64(&memory, &mut *caller, frame + CALL_RESULT_OFS, 0)?; + write_i64(&memory, &mut *caller, call_area, 0)?; return Ok(()); } }; @@ -823,7 +837,7 @@ fn jit_call_trampoline(caller: &mut Caller<'_, Host>, frame_ptr: u32) -> Result< let params: Vec = ty.params().collect(); let mut args: Vec = Vec::with_capacity(params.len()); for (i, pty) in params.iter().enumerate() { - let raw = read_i64(&memory, &*caller, frame + CALL_ARGS_OFS + i * 8); + let raw = read_i64(&memory, &*caller, call_area + 24 + i * 8); args.push(match pty { ValType::I32 => Val::I32(raw as i32), ValType::I64 => Val::I64(raw), @@ -852,7 +866,7 @@ fn jit_call_trampoline(caller: &mut Caller<'_, Host>, frame_ptr: u32) -> Result< // reported as a zero result rather than aborting the whole run. if let Err(e) = func.call(&mut *caller, &args, &mut results) { eprintln!("[jit_call] residual target trapped: {e:?}"); - write_i64(&memory, &mut *caller, frame + CALL_RESULT_OFS, 0)?; + write_i64(&memory, &mut *caller, call_area, 0)?; return Ok(()); } @@ -863,7 +877,7 @@ fn jit_call_trampoline(caller: &mut Caller<'_, Host>, frame_ptr: u32) -> Result< Some(Val::F32(x)) => (*x as u64) as i64, _ => 0, }; - write_i64(&memory, &mut *caller, frame + CALL_RESULT_OFS, result)?; + write_i64(&memory, &mut *caller, call_area, result)?; Ok(()) } diff --git a/pyre/pyre-wasm-runner/src/wasmi_host.rs b/pyre/pyre-wasm-runner/src/wasmi_host.rs index 88f49711b2b..56677547f1f 100644 --- a/pyre/pyre-wasm-runner/src/wasmi_host.rs +++ b/pyre/pyre-wasm-runner/src/wasmi_host.rs @@ -17,7 +17,7 @@ use wasmi::{ Module, Ref, Store, Table, Val, ValType, }; -use crate::{CALL_ARGS_OFS, CALL_FUNC_OFS, CALL_RESULT_OFS}; +use crate::CALL_RESULT_OFS; /// Per-store host state, mirroring the wasmtime path's `Host`. wasmi needs the /// engine handle stored too, because trace modules are compiled from inside an @@ -222,7 +222,8 @@ fn build_linker(engine: &Engine) -> Result, String> { "pyre_jit", "jit_call_host", |mut caller: Caller<'_, Host>, frame_ptr: u32| { - if let Err(e) = jit_call_trampoline(&mut caller, frame_ptr) { + if let Err(e) = jit_call_trampoline(&mut caller, frame_ptr, CALL_RESULT_OFS as u32) + { eprintln!("[jit_call_host] {e}"); } }, @@ -382,11 +383,22 @@ fn jit_compile( let jit_call = Func::wrap( &mut *caller, |mut inner: Caller<'_, Host>, frame_ptr: i32| { - if let Err(e) = jit_call_trampoline(&mut inner, frame_ptr as u32) { + if let Err(e) = + jit_call_trampoline(&mut inner, frame_ptr as u32, CALL_RESULT_OFS as u32) + { eprintln!("[jit_call] {e}"); } }, ); + let jit_call_compact = Func::wrap( + &mut *caller, + |mut inner: Caller<'_, Host>, frame_ptr: i32, call_area_ofs: i32| { + if let Err(e) = jit_call_trampoline(&mut inner, frame_ptr as u32, call_area_ofs as u32) + { + eprintln!("[jit_call_compact] {e}"); + } + }, + ); // Supply imports by name; a trace that imports only `env.memory` simply // leaves the defined `env.jit_call` unused. @@ -397,6 +409,9 @@ fn jit_compile( linker .define("env", "jit_call", Extern::Func(jit_call)) .map_err(estr)?; + linker + .define("env", "jit_call_compact", Extern::Func(jit_call_compact)) + .map_err(estr)?; let instance = linker .instantiate_and_start(&mut *caller, &module) .map_err(|e| format!("instantiate trace module: {e}"))?; @@ -435,17 +450,21 @@ fn jit_execute(caller: &mut Caller<'_, Host>, func_id: u32, frame_ptr: u32) -> R } /// Dispatch a residual call requested by a running trace. -fn jit_call_trampoline(caller: &mut Caller<'_, Host>, frame_ptr: u32) -> Result<(), String> { +fn jit_call_trampoline( + caller: &mut Caller<'_, Host>, + frame_ptr: u32, + call_area_ofs: u32, +) -> Result<(), String> { let memory = caller.data().memory.ok_or("memory")?; let table = caller.data().table.ok_or("table")?; - let frame = frame_ptr as usize; + let call_area = frame_ptr as usize + call_area_ofs as usize; - let func_ptr = read_u32(&memory, &*caller, frame + CALL_FUNC_OFS); + let func_ptr = read_u32(&memory, &*caller, call_area + 8); // `func_ptr == 0` is the "newstr" sentinel; without a host string // allocator it yields 0. if func_ptr == 0 { - write_i64(&memory, &mut *caller, frame + CALL_RESULT_OFS, 0)?; + write_i64(&memory, &mut *caller, call_area, 0)?; return Ok(()); } @@ -453,12 +472,12 @@ fn jit_call_trampoline(caller: &mut Caller<'_, Host>, frame_ptr: u32) -> Result< Some(Val::FuncRef(fr)) => match func_of(&fr) { Some(f) => f, None => { - write_i64(&memory, &mut *caller, frame + CALL_RESULT_OFS, 0)?; + write_i64(&memory, &mut *caller, call_area, 0)?; return Ok(()); } }, _ => { - write_i64(&memory, &mut *caller, frame + CALL_RESULT_OFS, 0)?; + write_i64(&memory, &mut *caller, call_area, 0)?; return Ok(()); } }; @@ -467,7 +486,7 @@ fn jit_call_trampoline(caller: &mut Caller<'_, Host>, frame_ptr: u32) -> Result< let params: Vec = ty.params().to_vec(); let mut args: Vec = Vec::with_capacity(params.len()); for (i, pty) in params.iter().enumerate() { - let raw = read_i64(&memory, &*caller, frame + CALL_ARGS_OFS + i * 8); + let raw = read_i64(&memory, &*caller, call_area + 24 + i * 8); args.push(match *pty { ValType::I32 => Val::I32(raw as i32), ValType::I64 => Val::I64(raw), @@ -495,7 +514,7 @@ fn jit_call_trampoline(caller: &mut Caller<'_, Host>, frame_ptr: u32) -> Result< // reported as a zero result rather than aborting the whole run. if let Err(e) = func.call(&mut *caller, &args, &mut results) { eprintln!("[jit_call] residual target trapped: {e}"); - write_i64(&memory, &mut *caller, frame + CALL_RESULT_OFS, 0)?; + write_i64(&memory, &mut *caller, call_area, 0)?; return Ok(()); } @@ -506,7 +525,7 @@ fn jit_call_trampoline(caller: &mut Caller<'_, Host>, frame_ptr: u32) -> Result< Some(Val::F32(x)) => (x.to_bits() as u64) as i64, _ => 0, }; - write_i64(&memory, &mut *caller, frame + CALL_RESULT_OFS, result)?; + write_i64(&memory, &mut *caller, call_area, result)?; Ok(()) }