From 21d7cc5f56e744267eeab57f44080ce318641470 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Sun, 16 Aug 2026 04:07:54 +0900 Subject: [PATCH 1/7] wasm: hand the guard's bridge cell address to the exit epilogue in a local A guard arm wrote its fail index into frame[0] and the shared epilogue read it back, subtracted the trace's fail-index base, scaled it by four and added the cell array base to reach the guard's bridge slot. The fail index is a constant at the arm, so the whole cell address is one too: the arm now computes it at emit time and leaves it in `bridge_slot_local`, and the epilogue loads the slot from there. `frame[0]` is still written, for the host round-trip that reads it. Finish and GuardAlwaysFails reach the same epilogue, so they set the local too; GuardAlwaysFails previously left an index there, which nothing read. Executed wasm ops, measured with `wasm_ops`: global_quasiimmut_invalidation -2.88%, the nested-loop microbenchmark -1.56%, fannkuch -0.63%, short_circuit_value_kept_stack -0.52%, raise_catch_loop -0.51%, fib_recursive -0.29%. A flat loop, whose steady state takes no guard exit, moves +0.00%. `compile_ms` on fannkuch reads 62.3 against 91.4. An earlier form of this inlined the whole dispatch, tail call included, into every guard arm. It cut more executed ops (gqi -6.71%) and lost to its own compile time: fannkuch's `compile_ms` went 75.2 to 181.9, against a +9.1% wall clock regression on a 1.6s fixture. pyre/check.py --backend wasm: 428/428. Assisted-by: Claude --- majit/majit-backend-wasm/src/codegen.rs | 142 +++++++++++++++++++----- 1 file changed, 112 insertions(+), 30 deletions(-) diff --git a/majit/majit-backend-wasm/src/codegen.rs b/majit/majit-backend-wasm/src/codegen.rs index 91b405b7148..3b0e0fdcbc8 100644 --- a/majit/majit-backend-wasm/src/codegen.rs +++ b/majit/majit-backend-wasm/src/codegen.rs @@ -2472,6 +2472,12 @@ fn build_function( debug_assert_eq!(ca_fi_local, ca_cfp_local + 1); debug_assert_eq!(alloc_scratch_local, bridge_slot_local + base_i32_locals); debug_assert_eq!(alloc_size_local, alloc_scratch_local + 1); + let guard_dispatch = BridgeDispatch { + cells_base, + fail_index_base, + bridge_slot_local, + enabled: bridge_dispatch, + }; let mut locals = Vec::new(); let mut start = 0; while start < value_types.types().len() { @@ -2823,6 +2829,7 @@ fn build_function( guard_idx, guard, block_exit_depth, + guard_dispatch, ); guard_idx += 1; fused_guard_at = Some(op_idx + 1); @@ -2951,6 +2958,9 @@ fn build_function( OpCode::Finish => { emit_guard_spill(&mut sink, constants, value_types, guard_idx, op); + if guard_dispatch.enabled { + emit_guard_bridge_dispatch(&mut sink, guard_idx, guard_dispatch); + } sink.br(block_exit_depth); guard_idx += 1; } @@ -2964,6 +2974,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -2975,6 +2986,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -2996,6 +3008,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3009,6 +3022,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3023,6 +3037,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3084,6 +3099,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3106,6 +3122,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3121,6 +3138,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3145,6 +3163,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); } guard_idx += 1; @@ -3172,6 +3191,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3193,6 +3213,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; // Success path: capture the caught exception into the result @@ -3770,6 +3791,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); } guard_idx += 1; @@ -3845,6 +3867,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3971,13 +3994,23 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } OpCode::GuardFutureCondition | OpCode::GuardAlwaysFails => { - // GuardAlwaysFails always exits. - sink.i32_const(guard_idx as i32); - sink.local_set(bridge_slot_local); + // GuardAlwaysFails always exits. This arm writes neither the + // fail args nor `frame[0]`, so it keeps branching to the shared + // epilogue — giving it the ordinary guard spill would change + // what the exit reports, not just how it gets there. The + // epilogue takes a cell address when bridge dispatch is on; + // otherwise retain the unused index write unchanged. + if guard_dispatch.enabled { + emit_guard_bridge_dispatch(&mut sink, guard_idx, guard_dispatch); + } else { + sink.i32_const(guard_idx as i32); + sink.local_set(bridge_slot_local); + } sink.br(block_exit_depth); guard_idx += 1; } @@ -5402,30 +5435,17 @@ fn build_function( sink.return_(); sink.end(); // end A $hot_exit - // Epilogue bridge dispatch. Control reaches here only - // after a guard `br`'d out of the exit block, having written its - // `fail_index` into `frame[0]`. Look up that guard's bridge slot in the - // shared cell array; if a bridge has been compiled (slot != 0), tail into - // it via `call_indirect` through the shared table — staying inside wasm — - // and return its result. Otherwise fall through to the host round-trip - // (return `frame_ptr`, the metainterp reads `frame[0]`). With every cell 0 - // (no bridge yet) this is inert and behavior is unchanged. + // Epilogue bridge dispatch for exits that branch out of the hot exit + // block. Their arms have already placed the exit's constant cell address + // in `bridge_slot_local`; load the table slot from that cell. If a bridge + // has been compiled (slot != 0), tail into it via `call_indirect` through + // the shared table and return its result. Otherwise fall through to the + // host round-trip (return `frame_ptr`, whose `frame[0]` was written by the + // ordinary guard or Finish arm). With every cell 0 (no bridge yet) this is + // inert and behavior is unchanged. if bridge_dispatch { - // slot = *(cells_base + (fail_index - fail_index_base) * 4) - // The cell array is local to this trace (one i32 per local guard); - // `frame[0]` carries the GLOBAL fail index, so subtract this trace's - // base back to a local cell index. - sink.i32_const(cells_base as i32); - sink.local_get(0); - sink.i64_load(mem64(0)); // frame[0] = fail_index - sink.i32_wrap_i64(); - if fail_index_base != 0 { - sink.i32_const(fail_index_base as i32); - sink.i32_sub(); - } - sink.i32_const(4); - sink.i32_mul(); - sink.i32_add(); + // slot = *(bridge_slot_local), where the local holds a cell address. + sink.local_get(bridge_slot_local); sink.i32_load(memarg(0, 2)); sink.local_tee(bridge_slot_local); sink.if_(BlockType::Empty); @@ -5833,6 +5853,14 @@ fn emit_array_addr( // ── Guard emission helpers ── +#[derive(Clone, Copy)] +struct BridgeDispatch { + cells_base: u32, + fail_index_base: u32, + bridge_slot_local: u32, + enabled: bool, +} + fn emit_guard_true( sink: &mut PeepSink<'_, '_>, constants: &indexmap::IndexMap, @@ -5840,6 +5868,7 @@ fn emit_guard_true( guard_idx: u32, op: &Op, block_exit_depth: u32, + dispatch: BridgeDispatch, ) { emit_resolve(sink, constants, value_types, op.arg(0).to_opref()); sink.i64_eqz(); @@ -5850,6 +5879,7 @@ fn emit_guard_true( guard_idx, op, block_exit_depth, + dispatch, ); } @@ -5860,6 +5890,7 @@ fn emit_guard_false( guard_idx: u32, op: &Op, block_exit_depth: u32, + dispatch: BridgeDispatch, ) { emit_resolve(sink, constants, value_types, op.arg(0).to_opref()); sink.i64_const(0); @@ -5871,6 +5902,7 @@ fn emit_guard_false( guard_idx, op, block_exit_depth, + dispatch, ); } @@ -5943,7 +5975,7 @@ fn next_ovf_guard(ops: &[Op], i: usize) -> Option<&Op> { matches!(next.opcode, OpCode::GuardNoOverflow | OpCode::GuardOverflow).then_some(next) } -/// Common guard exit: condition is on stack (i32), spill and branch on failure. +/// Common guard exit: condition is on stack (i32), spill and leave on failure. /// /// The spill belongs in this arm rather than in one shared exit handler after /// the trace. x86/assembler.py:835-846 `write_pending_failure_recoveries` can @@ -5961,7 +5993,9 @@ fn next_ovf_guard(ops: &[Op], i: usize) -> Option<&Op> { /// `block_exit_depth` is the statement-level depth of the enclosing exit /// `block` (preamble = 0, loop body = 1); the `+ 1` accounts for the `if` /// this opens. The stores run only on the failing edge, so the fallthrough -/// carries no frame traffic. +/// carries no frame traffic. With bridge dispatch enabled, the failing arm +/// writes its fail index to `frame[0]`, records its constant bridge-cell +/// address in a local, and branches to the shared dispatch epilogue. fn emit_guard_if_exit( sink: &mut PeepSink<'_, '_>, constants: &indexmap::IndexMap, @@ -5969,19 +6003,65 @@ fn emit_guard_if_exit( guard_idx: u32, op: &Op, block_exit_depth: u32, + dispatch: BridgeDispatch, ) { sink.if_(BlockType::Empty); - emit_guard_spill(sink, constants, value_types, guard_idx, op); - sink.br(block_exit_depth + 1); + emit_guard_exit( + sink, + constants, + value_types, + guard_idx, + op, + block_exit_depth + 1, + dispatch, + ); sink.end(); } +fn emit_guard_exit( + sink: &mut PeepSink<'_, '_>, + constants: &indexmap::IndexMap, + value_types: &ValueLocals, + guard_idx: u32, + op: &Op, + block_exit_depth: u32, + dispatch: BridgeDispatch, +) { + emit_guard_spill(sink, constants, value_types, guard_idx, op); + if dispatch.enabled { + emit_guard_bridge_dispatch(sink, guard_idx, dispatch); + } + sink.br(block_exit_depth); +} + +fn emit_guard_bridge_dispatch( + sink: &mut PeepSink<'_, '_>, + guard_idx: u32, + dispatch: BridgeDispatch, +) { + debug_assert!(guard_idx >= dispatch.fail_index_base); + let cell_addr = dispatch.cells_base + + (guard_idx - dispatch.fail_index_base) * std::mem::size_of::() as u32; + sink.i32_const(cell_addr as i32); + sink.local_set(dispatch.bridge_slot_local); +} + fn emit_guard_spill( sink: &mut PeepSink<'_, '_>, constants: &indexmap::IndexMap, value_types: &ValueLocals, guard_idx: u32, op: &Op, +) { + emit_guard_fail_args_spill(sink, constants, value_types, op); + emit_guard_fail_index_store(sink, guard_idx); +} + +fn emit_guard_fail_args_spill( + sink: &mut PeepSink<'_, '_>, + constants: &indexmap::IndexMap, + value_types: &ValueLocals, + op: &Op, ) { let fail_args: Vec = op .getfailargs() @@ -5994,7 +6074,9 @@ fn emit_guard_spill( emit_resolve(sink, constants, value_types, arg_ref); sink.i64_store(mem64(offset)); } +} +fn emit_guard_fail_index_store(sink: &mut PeepSink<'_, '_>, guard_idx: u32) { sink.local_get(0); sink.i64_const(guard_idx as i64); sink.i64_store(mem64(0)); From 722f2fa4c6e73026727d6d422b7efc12ce06baf8 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Sun, 16 Aug 2026 10:41:34 +0900 Subject: [PATCH 2/7] wasm: re-emit a compiled loop into its own table slot, and inline a bridge into it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `CompiledWasmLoop` retains the post-intern `ModuleBuildInputs` its module was built from, so `WasmBackend::reemit_loop` can rebuild it. The rebuild takes a fresh global fail-index base and a fresh guard-cell array, replays the bridge slots recorded for still-standalone bridges, re-bases the recorded descr ranges, and installs the result through a new `jit_replace_wasm` host binding that `table.set`s the loop's ORIGINAL slot — the slot is a baked `i32.const` immediate inside entry bridges and sibling loops' bridges, so a new slot would strand them. `intern_ref_constants` is never re-run: a second pass over already-rewritten ops finds no `ConstPtr`, yielding a zero GC-table base that reads linear memory at 0. The snapshot is kept only when a re-emission is armed, since it is live for the token's whole lifetime. The guard-cell array moves out of `build_wasm_module` to its callers, and `bridge_cells_base` / `num_guard_cells` become `Cell`s, because a rebuild's array has a different address and `JitCellToken.compiled` is a write-once `OnceLock`. `alloc_bridge_cells(0)` returns no array, keeping dispatch omitted for a guardless trace rather than handing out a zero-length allocation's non-null address. `ModuleBuildInputs` also carries inlined bridge regions. A region is emitted inside the loop as a `block` opened after `loop` and closed before the region's ops, so the failing guard reaches it with a `br` and its terminal JUMP takes the existing local-LABEL lowering — a parallel move between wasm locals. The guard arm moves its fail args straight into the region's inputarg locals instead of spilling them to the frame. A trial build decides eligibility, so `build_wasm_module` declines rather than asserting on a shape it cannot emit, such as a trace with no LABEL to branch back to. Both switches are read on the host and armed through guest exports; the guest has no environment, so `std::env::var` inside the backend never fires. Bridge diagnostics gain the re-emit and inline outcome counters. Two synthetic fixtures record one more `guard_failures` than their baseline with both switches off, and the cause is not yet identified. Sweeping `FROZEN_CHAIN_VALUE_SLOTS` over 64/96/128/192 moves which fixture is off by one without ever putting all three at their baseline, so the floor stays at 64. Assisted-by: Claude --- majit/majit-backend-wasm/js/jit_glue.js | 31 +- majit/majit-backend-wasm/src/codegen.rs | 503 +++++++++++++--- majit/majit-backend-wasm/src/failguard.rs | 28 +- majit/majit-backend-wasm/src/glue.rs | 23 + majit/majit-backend-wasm/src/lib.rs | 556 +++++++++++++++--- .../majit-backend-wasm/tests/codegen_test.rs | 281 +++++---- pyre/pyre-wasm-runner/src/main.rs | 114 +++- pyre/pyre-wasm-runner/src/wasmi_host.rs | 45 +- pyre/pyre-wasm/src/lib.rs | 29 + 9 files changed, 1304 insertions(+), 306 deletions(-) diff --git a/majit/majit-backend-wasm/js/jit_glue.js b/majit/majit-backend-wasm/js/jit_glue.js index 50456cae8cf..82460e90ef5 100644 --- a/majit/majit-backend-wasm/js/jit_glue.js +++ b/majit/majit-backend-wasm/js/jit_glue.js @@ -55,13 +55,16 @@ function jitCallTrampoline(framePtr, callAreaOfs = CALL_RESULT_OFS) { } export function jit_compile_wasm(bytesPtr, bytesLen) { + const trace = instantiateTrace(bytesPtr, bytesLen); + return registerTrace(trace); +} + +function instantiateTrace(bytesPtr, bytesLen) { if (!mainMemory) { throw new Error("jit_set_memory() must be called before jit_compile_wasm()"); } const bytes = new Uint8Array(mainMemory.buffer, bytesPtr, bytesLen).slice(); const module = new WebAssembly.Module(bytes); - const imports = { env: { memory: mainMemory } }; - // Check if the module needs jit_call import // (wasm-encoder adds it when trace has CALL ops) try { @@ -71,13 +74,33 @@ export function jit_compile_wasm(bytesPtr, bytesLen) { const instance = new WebAssembly.Instance(module, { env: { memory: mainMemory, jit_call: jitCallTrampoline, jit_call_compact: jitCallTrampoline, __indirect_function_table: mainTable } }); - return registerTrace(instance.exports.trace); + return instance.exports.trace; } catch (e) { // Retry without jit_call (for traces without CALL ops) const instance = new WebAssembly.Instance(module, { env: { memory: mainMemory } }); - return registerTrace(instance.exports.trace); + return instance.exports.trace; + } +} + +// Compile and instantiate a trace, then overwrite an existing shared-table +// slot. A caller already running the old function retains that invocation; +// later indirect calls use the replacement. +export function jit_replace_wasm(funcId, bytesPtr, bytesLen) { + try { + if (!funcTable[funcId]) { + return 0; + } + const trace = instantiateTrace(bytesPtr, bytesLen); + if (mainTable) { + mainTable.set(funcId, trace); + } + funcTable[funcId] = trace; + return funcId; + } catch (e) { + console.error('[jit_replace_wasm] failed:', e); + return 0; } } diff --git a/majit/majit-backend-wasm/src/codegen.rs b/majit/majit-backend-wasm/src/codegen.rs index 3b0e0fdcbc8..aea0e8124b0 100644 --- a/majit/majit-backend-wasm/src/codegen.rs +++ b/majit/majit-backend-wasm/src/codegen.rs @@ -1543,7 +1543,7 @@ pub struct GuardExit { /// The default sets `supports_guard_gc_type = false`, matching /// `AbstractCPU.supports_guard_gc_type` in `backend/model.py:21`; the /// codegen arms assert this flag before reading any other field. -#[derive(Default)] +#[derive(Clone, Default)] pub struct GuardGcTypeInfo { pub supports_guard_gc_type: bool, /// `get_translated_info_for_typeinfo()` = (base, shift, sizeof_ti). @@ -1893,6 +1893,12 @@ fn collect_guards_and_vars(inputargs: &[InputArg], ops: &[Op]) -> (Vec usize { + collect_guards_and_vars(inputargs, ops).0.len() +} + /// Dense wasm-local assignment and type lookup for each addressed SSA value. fn collect_value_types(inputargs: &[InputArg], ops: &[Op], num_vars: u32) -> ValueLocals { ValueLocals::collect(inputargs, ops, num_vars) @@ -1918,7 +1924,14 @@ fn collect_value_types(inputargs: &[InputArg], ops: &[Op], num_vars: u32) -> Val /// On native the trace is never executed, so the dispatch is omitted and no /// cells are needed — returning `(0, None)` keeps the emitted module /// byte-identical to the pre-chaining output and allocates nothing. -fn alloc_bridge_cells(num_guards: usize) -> (u32, Option>) { +pub fn alloc_bridge_cells(num_guards: usize) -> (u32, Option>) { + // `Box<[u32; 0]>` has a non-null dangling `as_mut_ptr()`. The pointer is + // not a dispatch table, so preserve the no-dispatch representation even + // on wasm where allocating that empty box would otherwise make the + // epilogue load an uninitialised bridge-slot local. + if num_guards == 0 { + return (0, None); + } #[cfg(target_arch = "wasm32")] { let mut cells = vec![0u32; num_guards].into_boxed_slice(); @@ -2014,6 +2027,7 @@ pub struct CaInlineParams { /// backend lowers as `malloc_cond`: load free, bump, compare top, call the /// slow path only on overflow). `None` keeps every allocation on the /// `wasm_jit_alloc` helper call. +#[derive(Clone)] pub struct NurseryAllocParams { /// Linear-memory address of the GC's `nursery_free` bump pointer. pub free_addr: u32, @@ -2043,54 +2057,187 @@ pub struct AllocHelpers { pub new_array_oldgen_fn_ptr: i64, } -type BuildWasmModuleOutput = (Vec, Vec, usize, u32, Option>); +type BuildWasmModuleOutput = (Vec, Vec, usize); + +/// Owned inputs for one wasm module build. A loop retains this after its +/// first build so it can emit the same trace again without revisiting mutable +/// backend state such as the constants pool or GC-reference interning pass. +pub struct ModuleBuildInputs { + pub inputargs: Vec, + /// These are the post-intern operations. Re-interning them would lose the + /// already allocated GC-table base encoded by `gc_table_base`. + pub ops: Vec, + /// Loop-closing bridge regions emitted inside this loop's wasm function. + /// Their value ids are already in the owning trace's global space; retain + /// them verbatim so every analysis and the generated locals see the same + /// identities as the bridge metadata. + pub inlined_bridges: Vec, + pub constants: indexmap::IndexMap, + pub vtable_offset: Option, + pub classptr_to_typeid: HashMap, + pub guard_gc_type_info: GuardGcTypeInfo, + pub alloc: AllocHelpers, + pub wb_fn_ptr: i64, + pub nursery: Option, + pub invalidated_flag_addr: u32, + pub gc_table_base: u32, + pub fail_index_base: u32, + pub bridge_cells_base: u32, + pub external_jump_slot: u32, + pub external_jump_key: u32, + pub frame: FrameGeometry, + pub ca: CaParams, +} + +pub struct InlinedBridge { + /// Per-trace fail index of the guard that enters this region. + pub source_fail_index: u32, + pub trace_id: u64, + pub inputargs: Vec, + pub ops: Vec, + /// Base of this already-interned region's GC table. Each region retains + /// its own roots; codegen selects it by the LoadFromGcTable producer. + pub gc_table_base: u32, +} + +/// Whether the exact operation stream emitted for `inputs` has a local loop +/// back-edge target. An inline bridge transfers with `br`, which can only +/// target the wasm `loop` opened for that LABEL. +pub fn merged_stream_has_loop_label(inputs: &ModuleBuildInputs) -> bool { + let mut ops = inputs.ops.clone(); + for bridge in &inputs.inlined_bridges { + ops.extend(bridge.ops.iter().cloned()); + } + find_loop_label_index(&ops).is_some_and(|label_idx| label_idx < inputs.ops.len()) +} + +impl Clone for InlinedBridge { + fn clone(&self) -> Self { + Self { + source_fail_index: self.source_fail_index, + trace_id: self.trace_id, + inputargs: self + .inputargs + .iter() + .map(InputArg::fresh_value_copy) + .collect(), + ops: self.ops.clone(), + gc_table_base: self.gc_table_base, + } + } +} + +impl Clone for ModuleBuildInputs { + fn clone(&self) -> Self { + Self { + inputargs: self + .inputargs + .iter() + .map(InputArg::fresh_value_copy) + .collect(), + ops: self.ops.clone(), + inlined_bridges: self.inlined_bridges.clone(), + constants: self.constants.clone(), + vtable_offset: self.vtable_offset, + classptr_to_typeid: self.classptr_to_typeid.clone(), + guard_gc_type_info: self.guard_gc_type_info.clone(), + alloc: self.alloc, + wb_fn_ptr: self.wb_fn_ptr, + nursery: self.nursery.clone(), + invalidated_flag_addr: self.invalidated_flag_addr, + gc_table_base: self.gc_table_base, + fail_index_base: self.fail_index_base, + bridge_cells_base: self.bridge_cells_base, + external_jump_slot: self.external_jump_slot, + external_jump_key: self.external_jump_key, + frame: self.frame, + ca: self.ca.clone(), + } + } +} /// Build a wasm module from majit IR. -#[expect( - clippy::too_many_arguments, - reason = "this is the single code-generation phase boundary and keeps each RPython backend input—IR, descriptors, GC state, frame geometry, and chaining state—explicit and independently auditable" -)] pub fn build_wasm_module( - inputargs: &[InputArg], - ops: &[Op], - constants: &indexmap::IndexMap, - vtable_offset: Option, - classptr_to_typeid: &HashMap, - guard_gc_type_info: &GuardGcTypeInfo, - alloc: AllocHelpers, - wb_fn_ptr: i64, - // Inline nursery-bump fast path for eligible `New`/`NewWithVtable` - // (see `NurseryAllocParams`); `None` keeps allocations on the helper. - nursery: Option<&NurseryAllocParams>, - // Address of the owning JitCellToken.invalidated AtomicBool in shared - // linear memory. GUARD_NOT_INVALIDATED reads this byte at runtime, like - // the native backends bake the same Arc allocation's address. - invalidated_flag_addr: u32, - // Base address of this trace's per-loop `GcTable` slot array in shared - // linear memory (`gcreftracer.py:9` `array_base_addr`), baked as the - // `LoadFromGcTable` base immediate exactly as the native backends bake - // it. `0` when the trace holds no reference constant. - gc_table_base: u32, - fail_index_base: u32, - // Table slot of the loop a JUMP-with-no-local-LABEL re-enters (a loop-closing - // bridge). `0` for a loop trace (its JUMP is a local back-edge `br`) and for a - // straight-line bridge (no JUMP). When set, the terminal external JUMP writes - // the loop's next inputargs into the frame and `return_call_indirect`s the - // loop's table slot — a wasm tail call, so the loop⇄bridge cycle runs at - // constant stack depth instead of growing one frame per iteration. - external_jump_slot: u32, - // 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, + inputs: &ModuleBuildInputs, ) -> Result { - let (mut guards, num_vars) = collect_guards_and_vars(inputargs, ops); + let ModuleBuildInputs { + inputargs, + ops, + inlined_bridges, + constants, + vtable_offset, + classptr_to_typeid, + guard_gc_type_info, + alloc, + wb_fn_ptr, + nursery, + invalidated_flag_addr, + gc_table_base, + fail_index_base, + bridge_cells_base, + external_jump_slot, + external_jump_key, + frame, + ca, + } = inputs; + // A bridge region has no function-entry loads, but its InputArgs and ops + // still need locals, liveness, homes, guard exits, and call signatures. + // Analyse the complete function as one stream while keeping `inputargs` + // below as the actual function-entry list. + // A normal module is emitted directly from its retained vectors. Keep + // that path allocation-free: code generation runs in the guest process, + // so transient merged-stream allocations can otherwise perturb the next + // collection boundary before any bridge is attached. + let mut merged_inputargs = Vec::new(); + let mut merged_ops = Vec::new(); + let mut gc_table_bases = HashMap::new(); + let (analysis_inputargs, analysis_ops): (&[InputArg], &[Op]) = if inlined_bridges.is_empty() { + (inputargs, ops) + } else { + merged_inputargs.extend(inputargs.iter().map(InputArg::fresh_value_copy)); + merged_ops.extend(ops.iter().cloned()); + for bridge in inlined_bridges { + merged_inputargs.extend(bridge.inputargs.iter().map(InputArg::fresh_value_copy)); + for op in &bridge.ops { + if op.opcode == OpCode::LoadFromGcTable { + gc_table_bases.insert(op.pos.get().raw(), bridge.gc_table_base); + } + } + merged_ops.extend(bridge.ops.iter().cloned()); + } + (&merged_inputargs, &merged_ops) + }; + let (mut guards, num_vars) = collect_guards_and_vars(analysis_inputargs, analysis_ops); + + // An inlined bridge branches back into the owner with wasm `br`. The + // merged stream must therefore contain the local LABEL that opens the + // wasm loop; a label-less cross-loop bridge has no in-function target. + if !inlined_bridges.is_empty() && !merged_stream_has_loop_label(inputs) { + return Err(BackendError::Unsupported( + "wasm backend: inlined bridge stream has no local loop LABEL".into(), + )); + } + for bridge in inlined_bridges { + if bridge.ops.is_empty() { + return Err(BackendError::Unsupported( + "wasm backend: inlined bridge stream has an empty region".into(), + )); + } + let source_guard = guards + .get(bridge.source_fail_index as usize) + .ok_or_else(|| { + BackendError::Unsupported( + "wasm backend: inlined bridge source guard is outside the owner stream".into(), + ) + })?; + let source_args = source_guard.fail_arg_refs.len(); + if source_args != bridge.inputargs.len() { + return Err(BackendError::Unsupported(format!( + "wasm backend: inlined bridge input arity {} differs from source guard arity {source_args}", + bridge.inputargs.len(), + ))); + } + } // Every trace's guard/finish exits draw their indices from ONE global // fail-index space (`failguard::FAIL_DESCR_REGISTRY`): a cross-trace chain @@ -2118,22 +2265,20 @@ pub fn build_wasm_module( // into its CA bridge, and a BRIDGE's own guards chain nested sub-bridges the // same way (a hot guard inside a chained bridge would otherwise round-trip // to the host forever). So any guarded trace wants dispatch cells. - let want_dispatch = !guards.is_empty(); - let (cells_base, cells_owner) = if want_dispatch { - alloc_bridge_cells(guards.len()) - } else { - (0, None) - }; - let bridge_dispatch = cells_base != 0; + let bridge_dispatch = *bridge_cells_base != 0; // Frame value slots (inputs at entry, fail-arg spills at guard exit) occupy // `[1, 1 + max(num inputs, max fail args))`. They precede the dispatch key, // Ref homes, and the always-present tail call area; a chained bridge must // fit the source token's frozen value-slot count before it can share that // frame. - let label_resume = LabelResumeData::collect(inputargs, ops); - let max_value_slots = normal_frame_value_slots(inputargs, ops) + label_resume.scalar_slots; + let label_resume = LabelResumeData::collect(&analysis_inputargs, &analysis_ops); + let max_value_slots = + normal_frame_value_slots(&analysis_inputargs, &analysis_ops) + label_resume.scalar_slots; if max_value_slots > frame.value_slots { + if !inlined_bridges.is_empty() { + super::record_inline_geometry(max_value_slots, frame.value_slots); + } return Err(BackendError::Unsupported(format!( "wasm backend: {max_value_slots} frame value slots exceed frozen frame layout \ ({})", @@ -2141,11 +2286,19 @@ pub fn build_wasm_module( ))); } - let value_types = collect_value_types(inputargs, ops, num_vars); - let ref_values = RefValues::collect(inputargs, ops); - let ref_homes = RefHomes::collect(inputargs, ops, ca.emit_ca, &label_resume.captured_refs); + let value_types = collect_value_types(&analysis_inputargs, &analysis_ops, num_vars); + let ref_values = RefValues::collect(&analysis_inputargs, &analysis_ops); + let ref_homes = RefHomes::collect( + &analysis_inputargs, + &analysis_ops, + ca.emit_ca, + &label_resume.captured_refs, + ); let num_ref_homes = ref_homes.len(); - if num_ref_homes > frame.ordinary_home_slots() || !label_resume.supported_by(frame) { + if num_ref_homes > frame.ordinary_home_slots() || !label_resume.supported_by(*frame) { + if !inlined_bridges.is_empty() { + super::record_inline_geometry(num_ref_homes, frame.ordinary_home_slots()); + } return Err(BackendError::Unsupported(format!( "wasm backend: {num_ref_homes} ordinary ref homes and {} LABEL ref captures exceed frozen frame layout ({}, {})", label_resume.ref_slots, @@ -2176,7 +2329,7 @@ pub fn build_wasm_module( // residual helpers, including the CA arm's inline fast path, use // `call_indirect` and need no import, although their frozen frame still // keeps the tail call area for future bridges. - let needs_call = has_trampoline_calls(inputargs, ops, ca.emit_ca); + let needs_call = has_trampoline_calls(&analysis_inputargs, &analysis_ops, ca.emit_ca); // In-module residual calls (`WASM_DIRECT_RESIDUAL_CALL`): the largest // eligible `(i64×n)->i64` arity in this trace — residual CALLs (word // result or word-ABI void) plus the `New*` / write-barrier helper @@ -2184,7 +2337,7 @@ pub fn build_wasm_module( // are none. Each distinct arity `0..=max` gets its own function type // (declared below) so those arms can `call_indirect` with a static type. let residual_max_arity = if WASM_DIRECT_RESIDUAL_CALL { - let scanned = ops + let scanned = analysis_ops .iter() .filter_map(|op| direct_helper_i64_arity(op, &ref_values)) .max(); @@ -2210,7 +2363,7 @@ pub fn build_wasm_module( // trace gets stable type indices while declaring each signature once. let mut float_residual_sigs = Vec::new(); if WASM_DIRECT_RESIDUAL_CALL { - for op in ops { + for op in analysis_ops { if let Some(sig) = residual_call_float_sig(op) && !float_residual_sigs.contains(&sig) { @@ -2222,7 +2375,10 @@ pub fn build_wasm_module( // i64- and f64-result types. As with the uniform i64 family, declaring // `0..=max` makes each type index a base plus the call arity. let true_void_residual_max_arity = if WASM_DIRECT_RESIDUAL_CALL { - ops.iter().filter_map(residual_call_void_true_arity).max() + analysis_ops + .iter() + .filter_map(residual_call_void_true_arity) + .max() } else { None }; @@ -2348,51 +2504,50 @@ pub fn build_wasm_module( let jit_call_idx = if needs_call { Some(0u32) } else { None }; let func = build_function( inputargs, - ops, + &analysis_inputargs, + &analysis_ops, + inlined_bridges, constants, num_vars, &value_types, jit_call_idx, - vtable_offset, + *vtable_offset, classptr_to_typeid, guard_gc_type_info, - alloc, - wb_fn_ptr, - nursery, + *alloc, + *wb_fn_ptr, + nursery.as_ref(), &ref_values, &ref_homes, &label_resume, - cells_base, + *bridge_cells_base, bridge_dispatch, - invalidated_flag_addr, - gc_table_base, - fail_index_base, - external_jump_slot, - external_jump_key, - frame, + *invalidated_flag_addr, + *gc_table_base, + &gc_table_bases, + *fail_index_base, + *external_jump_slot, + *external_jump_key, + *frame, residual_max_arity.map(|_| residual_type_base), &float_residual_type_indices, true_void_residual_max_arity.map(|_| true_void_residual_type_base), - ca, + ca.clone(), bridge_finish_fi, ca_helper_type_idx, )?; codes.function(&func); module.section(&codes); - Ok(( - module.finish(), - guards, - num_ref_homes, - cells_base, - cells_owner, - )) + Ok((module.finish(), guards, num_ref_homes)) } #[allow(clippy::too_many_arguments)] fn build_function( + entry_inputargs: &[InputArg], inputargs: &[InputArg], ops: &[Op], + inlined_bridges: &[InlinedBridge], constants: &indexmap::IndexMap, num_vars: u32, value_types: &ValueLocals, @@ -2410,6 +2565,7 @@ fn build_function( bridge_dispatch: bool, invalidated_flag_addr: u32, gc_table_base: u32, + gc_table_bases: &HashMap, fail_index_base: u32, external_jump_slot: u32, // Resume-at-LABEL dispatch key the terminal external JUMP writes before @@ -2472,11 +2628,26 @@ fn build_function( debug_assert_eq!(ca_fi_local, ca_cfp_local + 1); debug_assert_eq!(alloc_scratch_local, bridge_slot_local + base_i32_locals); debug_assert_eq!(alloc_size_local, alloc_scratch_local + 1); + let inline_guards: Vec> = inlined_bridges + .iter() + .enumerate() + .map(|(region, bridge)| InlineGuard { + guard_idx: fail_index_base + bridge.source_fail_index, + inputargs: &bridge.inputargs, + // Blocks open in reverse attach order, making region 0 innermost. + // This is the target depth at statement level; the guard `if` + // contributes the final +1 in `emit_guard_if_exit`. + branch_depth: region as u32, + }) + .collect(); let guard_dispatch = BridgeDispatch { cells_base, fail_index_base, bridge_slot_local, enabled: bridge_dispatch, + inline_guards: &inline_guards, + ref_homes, + frame, }; let mut locals = Vec::new(); let mut start = 0; @@ -2525,6 +2696,15 @@ fn build_function( // re-execute the complete loop body. let loop_label_idx = find_loop_label_index(ops); let has_loop = loop_label_idx.is_some(); + let bridge_op_count = inlined_bridges + .iter() + .map(|bridge| bridge.ops.len()) + .sum::(); + let bridge_start = ops.len().checked_sub(bridge_op_count).ok_or_else(|| { + BackendError::Unsupported( + "wasm backend: inlined bridge operations are not contained in the merged stream".into(), + ) + })?; // Def / last-use positions for the post-collection Ref reload filter. let liveness = HomeLiveness::collect(inputargs, ops); @@ -2643,7 +2823,7 @@ fn build_function( // bridge never scatters its frame-passed label values into the function // inputargs' home slots; those stay null-initialized (GC-safe) and the // resume loader sets the live label-arg homes. - for (k, ia) in inputargs.iter().enumerate() { + for (k, ia) in entry_inputargs.iter().enumerate() { let local_idx = value_types.local(ia.index); let offset = FRAME_SLOT_BASE + k as u64 * SLOT_SIZE; sink.local_get(0).i64_load(mem64(offset)); @@ -2670,6 +2850,7 @@ fn build_function( ref_homes, frame, gc_table_base, + gc_table_bases, *result, ); } @@ -2768,6 +2949,7 @@ fn build_function( ref_homes, frame, gc_table_base, + gc_table_bases, *result, ); } @@ -2776,8 +2958,49 @@ fn build_function( } if Some(op_idx) == loop_label_idx { sink.loop_(BlockType::Empty); + for _ in inlined_bridges.iter().rev() { + sink.block(BlockType::Empty); + } in_loop_body = true; } + // The loop's normal body ends with its JUMP, which branches around all + // regions. Closing one block before each attached region makes its + // body reachable only from the guard that branched to that block. + let mut started_bridge_regions = 0usize; + if in_loop_body && op_idx >= bridge_start { + let mut start = bridge_start; + for bridge in inlined_bridges { + if op_idx == start { + sink.end(); + } + if op_idx >= start { + started_bridge_regions += 1; + } + start += bridge.ops.len(); + } + } + // The blocks opened at the loop LABEL are closed once at the start of + // each appended region. Compute the remaining nesting directly from + // this operation's position, so a label-less stream cannot close a + // block that was never opened. + let open_bridge_blocks = if in_loop_body { + let remaining = inlined_bridges + .len() + .checked_sub(started_bridge_regions) + .ok_or_else(|| { + BackendError::Unsupported( + "wasm backend: inlined bridge region bookkeeping exceeded its open blocks" + .into(), + ) + })?; + u32::try_from(remaining).map_err(|_| { + BackendError::Unsupported( + "wasm backend: too many inlined bridge regions for wasm branch depth".into(), + ) + })? + } else { + 0 + }; // Depth (from statement level) of the enclosing `block` that guard // exits `br` to. Without `key_dispatch`: preamble = 0, loop body = 1 // (the `loop` sits between the body and block A). With `key_dispatch` @@ -2787,15 +3010,27 @@ fn build_function( // before the loop). Straight-line traces use the universal hot exit // block at depth 0. let block_exit_depth = match (has_loop, in_loop_body) { - (false, _) => 0u32, + (false, _) => { + if open_bridge_blocks != 0 { + return Err(BackendError::Unsupported( + "wasm backend: inlined bridge regions require a local loop LABEL".into(), + )); + } + 0u32 + } (true, false) => { + if open_bridge_blocks != 0 { + return Err(BackendError::Unsupported( + "wasm backend: inlined bridge region opened before its loop LABEL".into(), + )); + } if key_dispatch { 2 * (num_labels - labels_passed) as u32 } else { 0u32 } } - (true, true) => 1u32, + (true, true) => 1u32 + open_bridge_blocks, }; // The guard whose condition the previous op already pushed and tested. // `block_exit_depth` is unchanged across the pair: only a LABEL moves @@ -2953,7 +3188,7 @@ fn build_function( sink.i64_store(mem64(frame.home_slot_base + h as u64 * SLOT_SIZE)); } } - sink.br(0); + sink.br(open_bridge_blocks); } OpCode::Finish => { @@ -3287,6 +3522,7 @@ fn build_function( guard_idx, guard, block_exit_depth, + guard_dispatch, ); guard_idx += 1; fused_guard_at = Some(op_idx + 1); @@ -4102,8 +4338,9 @@ fn build_function( let vi = op.pos.get().raw(); if !OpRef::raw_is_constant(vi) && !hoisted_gc_table_loads.contains(&op.pos.get()) { let index = resolve_const_bits(constants, op.arg(0).to_opref()); - let slot = gc_table_base as u64 - + index as u64 * std::mem::size_of::() as u64; + let base = gc_table_bases.get(&vi).copied().unwrap_or(gc_table_base); + let slot = + base as u64 + index as u64 * std::mem::size_of::() as u64; sink.i32_const(slot as i32); sink.i32_load(MemArg { offset: 0, @@ -5659,6 +5896,7 @@ fn emit_seed_gc_table_ref( ref_homes: &RefHomes, frame: FrameGeometry, gc_table_base: u32, + gc_table_bases: &HashMap, result: OpRef, ) { let producer = ops @@ -5666,7 +5904,11 @@ fn emit_seed_gc_table_ref( .find(|op| op.pos.get() == result) .expect("hoisted GC-table result must have a producer"); let index = resolve_const_bits(constants, producer.arg(0).to_opref()); - let slot = gc_table_base as u64 + index as u64 * std::mem::size_of::() as u64; + let base = gc_table_bases + .get(&result.raw()) + .copied() + .unwrap_or(gc_table_base); + let slot = base as u64 + index as u64 * std::mem::size_of::() as u64; sink.i32_const(slot as i32); sink.i32_load(MemArg { offset: 0, @@ -5854,11 +6096,21 @@ fn emit_array_addr( // ── Guard emission helpers ── #[derive(Clone, Copy)] -struct BridgeDispatch { +struct InlineGuard<'a> { + guard_idx: u32, + inputargs: &'a [InputArg], + branch_depth: u32, +} + +#[derive(Clone, Copy)] +struct BridgeDispatch<'a> { cells_base: u32, fail_index_base: u32, bridge_slot_local: u32, enabled: bool, + inline_guards: &'a [InlineGuard<'a>], + ref_homes: &'a RefHomes, + frame: FrameGeometry, } fn emit_guard_true( @@ -5868,7 +6120,7 @@ fn emit_guard_true( guard_idx: u32, op: &Op, block_exit_depth: u32, - dispatch: BridgeDispatch, + dispatch: BridgeDispatch<'_>, ) { emit_resolve(sink, constants, value_types, op.arg(0).to_opref()); sink.i64_eqz(); @@ -5890,7 +6142,7 @@ fn emit_guard_false( guard_idx: u32, op: &Op, block_exit_depth: u32, - dispatch: BridgeDispatch, + dispatch: BridgeDispatch<'_>, ) { emit_resolve(sink, constants, value_types, op.arg(0).to_opref()); sink.i64_const(0); @@ -6003,7 +6255,7 @@ fn emit_guard_if_exit( guard_idx: u32, op: &Op, block_exit_depth: u32, - dispatch: BridgeDispatch, + dispatch: BridgeDispatch<'_>, ) { sink.if_(BlockType::Empty); emit_guard_exit( @@ -6025,8 +6277,27 @@ fn emit_guard_exit( guard_idx: u32, op: &Op, block_exit_depth: u32, - dispatch: BridgeDispatch, + dispatch: BridgeDispatch<'_>, ) { + if let Some(inline) = dispatch + .inline_guards + .iter() + .find(|g| g.guard_idx == guard_idx) + { + emit_guard_inline_bridge_move( + sink, + constants, + value_types, + dispatch.ref_homes, + dispatch.frame, + op, + inline.inputargs, + ); + // The target depth is measured outside this failing `if`; include the + // `if` itself before selecting the enclosing bridge block. + sink.br(inline.branch_depth + 1); + return; + } emit_guard_spill(sink, constants, value_types, guard_idx, op); if dispatch.enabled { emit_guard_bridge_dispatch(sink, guard_idx, dispatch); @@ -6034,10 +6305,50 @@ fn emit_guard_exit( sink.br(block_exit_depth); } +/// Transfer a failing guard directly into an inlined bridge. All sources are +/// pushed before any destination local is written, preserving parallel-move +/// semantics when fail arguments overlap bridge input locals. +fn emit_guard_inline_bridge_move( + sink: &mut PeepSink<'_, '_>, + constants: &indexmap::IndexMap, + value_types: &ValueLocals, + ref_homes: &RefHomes, + frame: FrameGeometry, + op: &Op, + inputargs: &[InputArg], +) { + let fail_args: Vec = op + .getfailargs() + .map(|args| args.iter().map(|arg| arg.to_opref()).collect()) + .unwrap_or_else(|| op.getarglist().iter().map(|arg| arg.to_opref()).collect()); + assert_eq!( + fail_args.len(), + inputargs.len(), + "guard and bridge input arity diverged" + ); + for (arg, input) in fail_args.iter().zip(inputargs) { + if value_types.ty(input.index) == ValType::F64 { + emit_resolve_f64(sink, constants, value_types, *arg); + } else { + emit_resolve(sink, constants, value_types, *arg); + } + } + for input in inputargs.iter().rev() { + sink.local_set(value_types.local(input.index)); + } + for input in inputargs { + if let Some(home) = ref_homes.home_id(input.index) { + sink.local_get(0); + sink.local_get(value_types.local(input.index)); + sink.i64_store(mem64(frame.home_slot_base + home as u64 * SLOT_SIZE)); + } + } +} + fn emit_guard_bridge_dispatch( sink: &mut PeepSink<'_, '_>, guard_idx: u32, - dispatch: BridgeDispatch, + dispatch: BridgeDispatch<'_>, ) { debug_assert!(guard_idx >= dispatch.fail_index_base); let cell_addr = dispatch.cells_base diff --git a/majit/majit-backend-wasm/src/failguard.rs b/majit/majit-backend-wasm/src/failguard.rs index 3d625a0cf47..a3442e562df 100644 --- a/majit/majit-backend-wasm/src/failguard.rs +++ b/majit/majit-backend-wasm/src/failguard.rs @@ -612,12 +612,12 @@ pub struct CompiledWasmLoop { /// epilogue reads `cells[fail_index]` and `compile_bridge` writes a bridge's /// table slot here. `0` when the trace has no in-module dispatch (native, or /// a guardless / straight-line trace). - pub bridge_cells_base: u32, + pub bridge_cells_base: Cell, /// Number of cells in the `bridge_cells_base` array = this loop's own guard /// count at compile time. A bridge attaches only to one of these original /// guards (`source_fail_index < num_guard_cells`); descrs appended past this /// range belong to already-chained bridges and have no cell of their own. - pub num_guard_cells: usize, + pub num_guard_cells: Cell, /// True when this is a peeled loop (`codegen::is_resumable_peeled`) — there /// is real work (a preamble = the unrolled first iteration) before the last /// `LABEL`, single- or multi-label. Such a loop carries the resume-at-LABEL @@ -657,14 +657,18 @@ pub struct CompiledWasmLoop { /// by the bridge's backend `trace_id` (see [`ChainedTraceMeta`]). Lets a /// guard INSIDE a chained bridge chain its own nested sub-bridge. pub chained_trace_meta: RefCell>, - /// Owns this loop's per-guard bridge-slot cell array so it is freed on - /// `Drop`; `bridge_cells_base` aliases its heap address (stable across the - /// struct move). `None` when the trace has no in-module dispatch. - pub _bridge_cells_owner: Option>, - /// Owns the cell arrays of every bridge chained onto this loop. A bridge - /// module lives as long as the source loop it attaches to, so its cells are - /// freed when this loop drops. Appended by `compile_bridge`. + /// Owns this loop's current cell array and every bridge cell array chained + /// onto it. A re-emission retains the old array for an already-running + /// module before switching its baked base to a new array. pub _bridge_owned_cells: RefCell>>, + /// Direct-loop guard index to bridge table slot. A re-emission replays + /// these slots into its fresh loop cell array. + pub bridge_slots: RefCell>, + /// Post-intern module inputs retained for a loop re-emission. Entry + /// bridges store `None` because they tail-call another loop. + pub reemit: RefCell>, + /// The environment-gated identity re-emission runs once per token. + pub reemitted: Cell, /// `(descr identity, table slot)` for every label published by a bridge /// chained onto this loop. The bridge module lives as long as its source /// loop, so `Drop` retracts entries that still name that bridge's slot. @@ -685,6 +689,12 @@ pub struct CompiledWasmLoop { pub ca_callers: RefCell>>, } +// Compiled loop metadata is transferred through the token's `Any + Send` +// holder, but all access to its IR snapshot and cell arrays is confined to the +// single wasm execution thread. The contained `RefCell`s enforce that runtime +// ownership model; moving the holder does not permit concurrent access. +unsafe impl Send for CompiledWasmLoop {} + impl CompiledWasmLoop { pub fn eager_func_handle(&self) -> u32 { self.func_handle.get() diff --git a/majit/majit-backend-wasm/src/glue.rs b/majit/majit-backend-wasm/src/glue.rs index 1baeebf2e51..f9c3c1d9524 100644 --- a/majit/majit-backend-wasm/src/glue.rs +++ b/majit/majit-backend-wasm/src/glue.rs @@ -37,6 +37,7 @@ mod imports { #[wasm_bindgen(raw_module = "./jit_glue.js")] unsafe extern "C" { pub(super) fn jit_compile_wasm(bytes_ptr: u32, bytes_len: u32) -> u32; + pub(super) fn jit_replace_wasm(func_id: u32, bytes_ptr: u32, bytes_len: u32) -> u32; pub(super) fn jit_execute_wasm(func_id: u32, frame_ptr: u32) -> u32; pub(super) fn jit_free_wasm(func_id: u32); } @@ -50,6 +51,7 @@ mod imports { #[link(wasm_import_module = "pyre_jit")] unsafe extern "C" { pub(super) fn jit_compile_wasm(bytes_ptr: u32, bytes_len: u32) -> u32; + pub(super) fn jit_replace_wasm(func_id: u32, bytes_ptr: u32, bytes_len: u32) -> u32; pub(super) fn jit_execute_wasm(func_id: u32, frame_ptr: u32) -> u32; pub(super) fn jit_free_wasm(func_id: u32); } @@ -65,6 +67,9 @@ mod imports { pub(super) unsafe fn jit_compile_wasm(_bytes_ptr: u32, _bytes_len: u32) -> u32 { panic!("{NO_BINDING}") } + pub(super) unsafe fn jit_replace_wasm(_func_id: u32, _bytes_ptr: u32, _bytes_len: u32) -> u32 { + panic!("{NO_BINDING}") + } pub(super) unsafe fn jit_execute_wasm(_func_id: u32, _frame_ptr: u32) -> u32 { panic!("{NO_BINDING}") } @@ -107,6 +112,24 @@ pub fn compile_module_cached(wasm_bytes: &[u8]) -> u32 { handle } +/// Replace the function stored in an existing trace slot. +/// +/// Replacements deliberately bypass `MODULE_CACHE`: a slot belongs to one +/// token, while a byte-identical later trace must not inherit that token's +/// table identity. +pub fn replace_module(func_id: u32, wasm_bytes: &[u8]) -> u32 { + let ptr = wasm_bytes.as_ptr() as u32; + let len = wasm_bytes.len() as u32; + #[cfg(feature = "web")] + { + imports::jit_replace_wasm(func_id, ptr, len) + } + #[cfg(not(feature = "web"))] + unsafe { + imports::jit_replace_wasm(func_id, ptr, len) + } +} + /// Execute a compiled JIT function with the given frame pointer. pub fn execute(func_id: u32, frame_ptr: u32) -> u32 { #[cfg(feature = "web")] diff --git a/majit/majit-backend-wasm/src/lib.rs b/majit/majit-backend-wasm/src/lib.rs index 2270af3ca08..f9f36391f46 100644 --- a/majit/majit-backend-wasm/src/lib.rs +++ b/majit/majit-backend-wasm/src/lib.rs @@ -15,8 +15,8 @@ mod glue; use std::cell::RefCell; use std::collections::HashMap; -use std::sync::Arc; -use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; /// Diagnostic-only `compile_bridge` outcome tallies, read out via the /// `pyre_jit_bridge_diag` guest export (the runner prints them at @@ -66,7 +66,77 @@ use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; /// re-bridged after its first bridge was outgrown); a count that tracks /// `BRIDGE_OK` says the epilogue dispatch is not taking the cell at all and /// every bridge after the first is dead weight. -pub static BRIDGE_DIAG: [AtomicU64; 30] = [const { AtomicU64::new(0) }; 30]; +/// 30 = a host-armed loop re-emission was attempted but failed; +/// 31 = it succeeded and the rebuilt module is installed in the loop's +/// original table slot. 31 is the only positive evidence that a re-emission +/// ran at all: a re-emission that silently never fires is indistinguishable +/// from one that fires and changes nothing. +/// 32 = a loop-closing bridge region was inlined; 33 = inlining declined +/// because the source guard belongs to an already chained trace; 34 = the +/// bridge is not loop-closing; 35 = the owner has no retained module inputs; +/// 36 = that guard already owns a region; 37 = the merged stream exceeds the +/// owner's frozen frame geometry; 38 = the bridge does not resume at the loop +/// header; 39 = the merged stream has no local loop LABEL for the wasm back +/// edge. 40-43 split a rejected inline trial into value-layout, +/// Ref-home-layout, missing-local-label, and other backend errors. +pub static BRIDGE_DIAG: [AtomicU64; 44] = [const { AtomicU64::new(0) }; 44]; + +/// The first three inline geometry failures, packed as `(needed, available)`. +/// They expose a frozen-layout shortage without changing the compile result. +static INLINE_GEOMETRY: [AtomicU64; 3] = [const { AtomicU64::new(0) }; 3]; +static INLINE_GEOMETRY_COUNT: AtomicU64 = AtomicU64::new(0); +static INLINE_TRIAL_ERRORS: Mutex> = Mutex::new(Vec::new()); + +pub(crate) fn record_inline_geometry(needed: usize, available: usize) { + let index = INLINE_GEOMETRY_COUNT.fetch_add(1, Ordering::Relaxed) as usize; + if let Some(slot) = INLINE_GEOMETRY.get(index) { + slot.store( + ((needed as u64) << 32) | available as u64, + Ordering::Relaxed, + ); + } +} + +/// Read a packed `(needed, available)` inline geometry failure. +pub fn inline_geometry_diag(index: usize) -> u64 { + INLINE_GEOMETRY + .get(index) + .map_or(0, |slot| slot.load(Ordering::Relaxed)) +} + +pub fn inline_trial_errors() -> String { + INLINE_TRIAL_ERRORS.lock().unwrap().join(" | ") +} + +fn record_inline_trial_error(error: &BackendError) { + let mut errors = INLINE_TRIAL_ERRORS.lock().unwrap(); + if errors.len() < 3 { + errors.push(error.to_string()); + } +} + +static REEMIT_ENABLED: AtomicBool = AtomicBool::new(false); +static INLINE_BRIDGE_ENABLED: AtomicBool = AtomicBool::new(false); + +/// Arm loop-module replacement from the host before guest execution starts. +pub fn reemit_enable() { + REEMIT_ENABLED.store(true, Ordering::Relaxed); +} + +fn reemit_enabled() -> bool { + REEMIT_ENABLED.load(Ordering::Relaxed) +} + +/// Arm loop-closing bridge inlining. Inlining rebuilds the owning loop, so it +/// also enables the replacement path. +pub fn inline_bridge_enable() { + INLINE_BRIDGE_ENABLED.store(true, Ordering::Relaxed); + reemit_enable(); +} + +fn inline_bridge_enabled() -> bool { + INLINE_BRIDGE_ENABLED.load(Ordering::Relaxed) +} /// Read a `BRIDGE_DIAG` tally (saturating index). Surfaced to the host through /// the `pyre_jit_bridge_diag` export in the `pyre-wasm` crate. @@ -114,6 +184,7 @@ fn diag_bump(i: usize) { // per compiled token. const FROZEN_CHAIN_VALUE_SLOTS: usize = 64; const FROZEN_CHAIN_REF_HOMES: usize = 128; +const FROZEN_CHAIN_LABEL_REF_SLOTS: usize = 2; /// An op whose result advances loop-carried state. A value produced inside the /// re-running region by arithmetic or by a heap load is fresh on each pass, so @@ -1636,6 +1707,164 @@ impl WasmBackend { } } } + + /// Rebuild a loop module and install it into its original shared-table + /// slot. The retained inputs are post-intern, so this does not allocate a + /// second GC reference table or change any reference-constant immediate. + #[allow(unreachable_code, unused_variables)] + pub fn reemit_loop(&mut self, token: &JitCellToken) -> Result<(), BackendError> { + let compiled = token + .compiled + .get() + .and_then(|c| c.downcast_ref::()) + .ok_or_else(|| { + BackendError::Unsupported("wasm backend: no compiled loop to re-emit".into()) + })?; + let Some(mut inputs) = compiled.reemit.borrow().as_ref().cloned() else { + return Err(BackendError::Unsupported( + "wasm backend: entry bridge is not re-emittable".into(), + )); + }; + let old_handle = compiled.eager_func_handle(); + if old_handle == 0 { + return Err(BackendError::Unsupported( + "wasm backend: unmaterialized loop is not re-emittable".into(), + )); + } + + inputs.fail_index_base = fail_descr_base(); + let merged_guard_count = codegen::guard_exit_count(&inputs.inputargs, &inputs.ops) + + inputs + .inlined_bridges + .iter() + .map(|region| codegen::guard_exit_count(®ion.inputargs, ®ion.ops)) + .sum::(); + let (new_cells_base, new_cells_owner) = codegen::alloc_bridge_cells(merged_guard_count); + inputs.bridge_cells_base = new_cells_base; + let (wasm_bytes, guard_exits, _) = codegen::build_wasm_module(&inputs)?; + let code_size = wasm_bytes.len(); + let own_guard_count = codegen::guard_exit_count(&inputs.inputargs, &inputs.ops); + let descrs: Vec> = guard_exits + .iter() + .enumerate() + .map(|(index, g)| { + let mut region_start = own_guard_count; + let trace_id = inputs + .inlined_bridges + .iter() + .find_map(|region| { + let count = codegen::guard_exit_count(®ion.inputargs, ®ion.ops); + let contains = (region_start..region_start + count).contains(&index); + region_start += count; + contains.then_some(region.trace_id) + }) + .unwrap_or(compiled.trace_id); + Arc::new(WasmFailDescr { + fail_index: g.fail_index, + trace_id, + fail_arg_types: g.fail_arg_types.clone(), + is_finish: g.is_finish, + meta_descr: g.meta_descr.clone(), + }) + }) + .collect(); + + #[cfg(all(target_arch = "wasm32", not(target_os = "wasi")))] + if glue::replace_module(old_handle, &wasm_bytes) != old_handle { + return Err(BackendError::Unsupported( + "wasm host rejected the re-emitted trace module".into(), + )); + } + #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))] + { + let _ = wasm_bytes; + return Err(BackendError::Unsupported( + "wasm backend: no host replacement binding".into(), + )); + } + + // The host has accepted the replacement, so its newly encoded global + // indices can now be made visible in the registry and local metadata. + // Both instances remain resident, so account for the replacement block + // in the same lifetime ledger as an ordinary compiled module. + let block = self.asm_memory_stats.record_block(code_size, code_size); + self.asm_memory_blocks.push(block); + // Keep still-standalone bridge descriptors after the rebuilt merged + // prefix. Adding regions grows that prefix, so every old positional + // range moves by exactly the difference in guard-cell counts. + let old_guard_count = compiled.num_guard_cells.get(); + let chained_descrs = compiled.fail_descrs.borrow()[old_guard_count..].to_vec(); + let mut replacement_descrs = descrs.clone(); + replacement_descrs.extend(chained_descrs); + *compiled.fail_descrs.borrow_mut() = replacement_descrs; + register_fail_descrs(&descrs); + let guard_growth = guard_exits.len().saturating_sub(old_guard_count); + if guard_growth != 0 { + for (_, _, start, _) in compiled.bridge_descr_ranges.borrow_mut().iter_mut() { + *start += guard_growth; + } + } + #[cfg(all(target_arch = "wasm32", not(target_os = "wasi")))] + if new_cells_base != 0 { + for (&fail_index, &bridge_slot) in compiled.bridge_slots.borrow().iter() { + let cell = (new_cells_base as usize + fail_index as usize * 4) as *mut u32; + unsafe { core::ptr::write(cell, bridge_slot) }; + } + } + if let Some(owner) = new_cells_owner { + compiled._bridge_owned_cells.borrow_mut().push(owner); + } + compiled.bridge_cells_base.set(new_cells_base); + compiled.num_guard_cells.set(guard_exits.len()); + { + let mut metas = compiled.chained_trace_meta.borrow_mut(); + let mut offset = own_guard_count; + for region in &inputs.inlined_bridges { + let count = codegen::guard_exit_count(®ion.inputargs, ®ion.ops); + let exits = &guard_exits[offset..offset + count]; + metas.insert( + region.trace_id, + ChainedTraceMeta { + cells_base: new_cells_base + offset as u32 * 4, + num_cells: count, + guard_fail_arg_advanced: guard_fail_args_advanced(®ion.ops, exits), + }, + ); + offset += count; + } + } + *compiled.reemit.borrow_mut() = Some(inputs.clone()); + + // LABEL targets bake only the stable table slot, so restamp them for + // this build. CA dispatch additionally carries the new finish index. + let _ = stamp_and_publish_label_targets( + old_handle, + compiled.frame, + &inputs.inputargs, + &inputs.ops, + ); + let loop_finish_fi = descrs + .iter() + .find(|descr| { + descr.is_finish + && !failguard::meta_descr_is_exit_frame_with_exception(&descr.meta_descr) + }) + .map(|descr| descr.fail_index) + .unwrap_or(failguard::WASM_CA_FINISH_FI_UNKNOWN); + ca_dispatch_publish( + token.number, + old_handle, + loop_finish_fi, + compiled as *const CompiledWasmLoop as usize as u32, + ); + if let Some(mut target) = call_assembler_target(token.number) { + target.func_handle = old_handle; + target.loop_finish_fi = loop_finish_fi; + target.compiled_ptr = compiled as *const CompiledWasmLoop as usize as u64; + publish_call_assembler_target(token.number, target); + } + Ok(()) + } } unsafe impl Send for WasmBackend {} @@ -2220,7 +2449,8 @@ impl majit_backend::Backend for WasmBackend { // geometry for both the loop and each nursery-allocated self callee. let raw_frame_value_slots = codegen::frame_value_slots(inputargs, ops); let raw_num_ref_homes = codegen::count_ref_homes(inputargs, ops); - let label_ref_slots = codegen::label_ref_capture_slots(inputargs, ops); + let label_ref_slots = + codegen::label_ref_capture_slots(inputargs, ops).max(FROZEN_CHAIN_LABEL_REF_SLOTS); // An entry bridge (`compile.py:1006-1022 ResumeFromInterpDescr`) is sent // to the backend through `compile_loop` like any loop, but it is not one: // it has no LABEL of its own and ends in a JUMP into an @@ -2307,50 +2537,51 @@ impl majit_backend::Backend for WasmBackend { // chain's `frame[0]` resolves regardless of which module wrote it // (`failguard::FAIL_DESCR_REGISTRY`). let fail_index_base = fail_descr_base(); - let (wasm_bytes, guard_exits, num_ref_homes, bridge_cells_base, bridge_cells_owner) = - codegen::build_wasm_module( - inputargs, - ops, - &self.constants, - self.vtable_offset, - &typeid_table, - &guard_gc_type_info, - alloc, - wb_fn_ptr, - nursery_alloc_params(ops).as_ref(), - Arc::as_ptr(&token.invalidated) as usize as u32, - gc_table_base, - fail_index_base, - // A real loop's JUMP is a local back-edge `br` and needs - // neither; an entry bridge tail-calls the target loop's table - // slot and resumes at the label `external_jump_key` selects. - entry_bridge_target.map_or(0, |t| t.func_handle), - entry_bridge_target.map_or(0, |t| t.key), - frame, - ca_targets.as_ref().map_or_else( - || codegen::CaParams { - // A loop can run on a nursery CA frame and must reload - // local 0 after a collection even when it emits no CA. - ca_reload_fn_ptr: wasm_jit_ca_reload_frame as *const () as usize as i64, - jf_top_addr: jf_top_addr(), - ..codegen::CaParams::default() - }, - |targets| codegen::CaParams { - emit_ca: true, - targets: ca_codegen_targets(targets), - deopt_helper_slot: ca_deopt_helper_slot(), - 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, - // Inline allocation is shared module state, so admit - // it only if the largest target frame is eligible. - inline: ca_inline_params(ca_max_frame_bytes(targets)), - jf_top_addr: jf_top_addr(), - }, - ), - )?; + let (bridge_cells_base, bridge_cells_owner) = + codegen::alloc_bridge_cells(codegen::guard_exit_count(inputargs, ops)); + let module_inputs = codegen::ModuleBuildInputs { + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + // Keep these rewritten operations exactly as intern_ref_constants + // produced them; their LoadFromGcTable immediates share this base. + ops: ops_owned.clone(), + inlined_bridges: Vec::new(), + constants: self.constants.clone(), + vtable_offset: self.vtable_offset, + classptr_to_typeid: typeid_table, + guard_gc_type_info, + alloc, + wb_fn_ptr, + nursery: nursery_alloc_params(ops), + invalidated_flag_addr: Arc::as_ptr(&token.invalidated) as usize as u32, + gc_table_base, + fail_index_base, + bridge_cells_base, + // A real loop's JUMP is a local back-edge `br`; an entry bridge + // tail-calls its target loop and is deliberately not re-emittable. + external_jump_slot: entry_bridge_target.map_or(0, |t| t.func_handle), + external_jump_key: entry_bridge_target.map_or(0, |t| t.key), + frame, + ca: ca_targets.as_ref().map_or_else( + || codegen::CaParams { + ca_reload_fn_ptr: wasm_jit_ca_reload_frame as *const () as usize as i64, + jf_top_addr: jf_top_addr(), + ..codegen::CaParams::default() + }, + |targets| codegen::CaParams { + emit_ca: true, + targets: ca_codegen_targets(targets), + deopt_helper_slot: ca_deopt_helper_slot(), + 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, + inline: ca_inline_params(ca_max_frame_bytes(targets)), + jf_top_addr: jf_top_addr(), + }, + ), + }; + let (wasm_bytes, guard_exits, num_ref_homes) = codegen::build_wasm_module(&module_inputs)?; // Build fail descriptors let fail_descrs: Vec> = guard_exits @@ -2469,15 +2700,25 @@ impl majit_backend::Backend for WasmBackend { max_output_slots, num_ref_homes, frame, - bridge_cells_base, - num_guard_cells: guard_exits.len(), + bridge_cells_base: std::cell::Cell::new(bridge_cells_base), + num_guard_cells: std::cell::Cell::new(guard_exits.len()), has_preamble, label_descrs, guard_fail_arg_advanced, bridge_descr_ranges: std::cell::RefCell::new(Vec::new()), 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()), + _bridge_owned_cells: std::cell::RefCell::new(bridge_cells_owner.into_iter().collect()), + bridge_slots: std::cell::RefCell::new(HashMap::new()), + // Retaining the snapshot costs long-lived heap for the token's + // whole lifetime, which moves when the collector next runs and so + // moves which iteration a back edge's eval-breaker guard bails on. + // Keep it only when a re-emission can actually consume it, so a run + // with the switches off allocates exactly what it did before. + reemit: std::cell::RefCell::new( + (entry_bridge_target.is_none() && (reemit_enabled() || inline_bridge_enabled())) + .then_some(module_inputs), + ), + reemitted: std::cell::Cell::new(false), bridge_owned_label_targets: std::cell::RefCell::new(Vec::new()), ca_active: std::cell::Cell::new(false), ca_terminal_declined: std::cell::Cell::new(false), @@ -2578,8 +2819,8 @@ impl majit_backend::Backend for WasmBackend { // round-tripping through the interpreter, the source loop's epilogue // `call_indirect`s the bridge in-module (see `codegen` epilogue). The // bridge runs in the SOURCE loop's reused frame: the guard spilled its - // fail args positionally into `frame[1..]`, exactly where the bridge's - // `build_function` reads its inputs (`inputargs[k].index == k`), so no + // fail args positionally into `frame[1..]`. `build_function` reads the + // positional slot `k`, independently of the bridge value id, so no // argument-recovery layout is needed — hence `caller_recovery_layout` // and `previous_tokens` are unused. let ops_owned: Vec = normalize_ops_for_codegen(inputargs, ops); @@ -2607,7 +2848,7 @@ impl majit_backend::Backend for WasmBackend { // Scalars read from the source loop up front, so the immutable borrow of // `original_token` is released before the `&mut self` codegen calls. - let (source_guard, source_func_handle, source_has_preamble, source_frame) = { + let (source_guard, source_func_handle, source_has_preamble, source_frame, is_direct) = { let source_loop = original_token .compiled .get() @@ -2627,8 +2868,8 @@ impl majit_backend::Backend for WasmBackend { let is_direct = source_trace_id == source_loop.trace_id; let guard = if is_direct { Some(( - source_loop.bridge_cells_base, - source_loop.num_guard_cells, + source_loop.bridge_cells_base.get(), + source_loop.num_guard_cells.get(), source_loop .guard_fail_arg_advanced .get(source_fail_index as usize) @@ -2656,6 +2897,7 @@ impl majit_backend::Backend for WasmBackend { source_loop.materialize_func_handle()?, source_loop.has_preamble, source_loop.frame, + is_direct, ) }; @@ -2837,6 +3079,128 @@ impl majit_backend::Backend for WasmBackend { } } + if inline_bridge_enabled() { + if !is_direct { + diag_bump(33); + } else if !bridge_is_loop_closing { + diag_bump(34); + } else if !resumes_at_loop_header { + diag_bump(38); + } else if let Some(mut candidate) = original_token + .compiled + .get() + .and_then(|c| c.downcast_ref::()) + .and_then(|loop_| loop_.reemit.borrow().as_ref().cloned()) + { + if candidate + .inlined_bridges + .iter() + .any(|r| r.source_fail_index == source_fail_index) + { + diag_bump(36); + } else if !codegen::merged_stream_has_loop_label(&candidate) { + diag_bump(39); + } else { + self.collect_constants_from_ops(ops); + candidate.inlined_bridges.push(codegen::InlinedBridge { + source_fail_index, + trace_id: self.trace_counter, + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + ops: ops_owned.clone(), + gc_table_base, + }); + let mut merged_ops = candidate.ops.clone(); + for region in &candidate.inlined_bridges { + merged_ops.extend(region.ops.iter().cloned()); + } + candidate.constants = self.constants.clone(); + candidate.classptr_to_typeid = self.collect_classptr_typeid_table(&merged_ops); + candidate.guard_gc_type_info = self.collect_guard_gc_type_info(&merged_ops); + candidate.nursery = nursery_alloc_params(&merged_ops); + match codegen::build_wasm_module(&candidate) { + Err(ref error @ BackendError::Unsupported(ref reason)) => { + record_inline_trial_error(error); + diag_bump(37); + if reason.contains("frame value slots exceed frozen frame layout") { + diag_bump(40); + } else if reason.contains("ordinary ref homes") { + diag_bump(41); + } else if reason + .contains("inlined bridge stream has no local loop LABEL") + { + diag_bump(42); + } else { + diag_bump(43); + } + } + Err(ref error @ BackendError::CompilationFailed(_)) => { + record_inline_trial_error(error); + diag_bump(37); + diag_bump(43); + } + Ok(_) => { + let source_loop = original_token + .compiled + .get() + .and_then(|c| c.downcast_ref::()) + .expect("source loop disappeared before inline install"); + // The local branch supersedes any previous direct-cell + // dispatch for this guard. Remove it before reemit so + // the fresh array cannot replay a contradictory slot. + let old_bridge_slot = source_loop + .bridge_slots + .borrow_mut() + .remove(&source_fail_index); + #[cfg(all(target_arch = "wasm32", not(target_os = "wasi")))] + if source_cells_base != 0 { + let cell = (source_cells_base as usize + + source_fail_index as usize * 4) + as *mut u32; + unsafe { core::ptr::write(cell, 0) }; + } + let old_inputs = source_loop.reemit.replace(Some(candidate)); + match self.reemit_loop(original_token) { + Ok(()) => { + self.trace_counter += 1; + if let Some(table) = gc_table { + Self::register_gc_table(original_token, table); + } + diag_bump(31); + diag_bump(32); + return Ok(AsmInfo { + code_addr: 0, + code_size: 0, + }); + } + Err(_) => { + source_loop.reemit.replace(old_inputs); + if let Some(slot) = old_bridge_slot { + source_loop + .bridge_slots + .borrow_mut() + .insert(source_fail_index, slot); + #[cfg(all( + target_arch = "wasm32", + not(target_os = "wasi") + ))] + if source_cells_base != 0 { + let cell = (source_cells_base as usize + + source_fail_index as usize * 4) + as *mut u32; + unsafe { core::ptr::write(cell, slot) }; + } + } + diag_bump(30); + } + } + } + } + } + } else { + diag_bump(35); + } + } + self.collect_constants_from_ops(ops); let trace_id = self.trace_counter; self.trace_counter += 1; @@ -2882,29 +3246,29 @@ impl majit_backend::Backend for WasmBackend { // invalidation starts valid; only a later invalidation may kill its // `GUARD_NOT_INVALIDATED` operations. let bridge_flag = original_token.mint_bridge_invalidation_flag(); - let (wasm_bytes, guard_exits, _num_ref_homes, bridge_cells_base, bridge_cells_owner) = - codegen::build_wasm_module( - inputargs, - ops, - &self.constants, - self.vtable_offset, - &typeid_table, - &guard_gc_type_info, - alloc, - wb_fn_ptr, - nursery_alloc_params(ops).as_ref(), - Arc::as_ptr(&bridge_flag) as usize as u32, - gc_table_base, - base, - // A loop-closing bridge's terminal JUMP re-enters the target - // loop (own or sibling, resolved via `LABEL_TARGETS`) through - // its table slot via a tail call, resuming at the label - // `external_jump_key` selects. - external_jump_slot, - external_jump_key, - source_frame, - ca_params, - )?; + let (bridge_cells_base, bridge_cells_owner) = + codegen::alloc_bridge_cells(codegen::guard_exit_count(inputargs, ops)); + let module_inputs = codegen::ModuleBuildInputs { + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + ops: ops_owned.clone(), + inlined_bridges: Vec::new(), + constants: self.constants.clone(), + vtable_offset: self.vtable_offset, + classptr_to_typeid: typeid_table, + guard_gc_type_info, + alloc, + wb_fn_ptr, + nursery: nursery_alloc_params(ops), + invalidated_flag_addr: Arc::as_ptr(&bridge_flag) as usize as u32, + gc_table_base, + fail_index_base: base, + bridge_cells_base, + external_jump_slot, + external_jump_key, + frame: source_frame, + ca: ca_params, + }; + let (wasm_bytes, guard_exits, _num_ref_homes) = codegen::build_wasm_module(&module_inputs)?; // Bridge exit descrs (fail_index already base-offset by build_wasm_module). let bridge_descrs: Vec> = guard_exits @@ -3050,6 +3414,21 @@ impl majit_backend::Backend for WasmBackend { unsafe { core::ptr::write(cell, bridge_slot); } + // Only retained module replacement needs to restore this cell + // after allocating a fresh dispatch array. Without replacement, + // the live cell is already the sole dispatch state. + if is_direct && reemit_enabled() { + if let Some(source_loop) = original_token + .compiled + .get() + .and_then(|c| c.downcast_ref::()) + { + source_loop + .bridge_slots + .borrow_mut() + .insert(source_fail_index, bridge_slot); + } + } } #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))] let _ = (source_cells_base, bridge_slot); @@ -3060,6 +3439,23 @@ impl majit_backend::Backend for WasmBackend { let block = self.asm_memory_stats.record_block(code_size, code_size); self.asm_memory_blocks.push(block); + // The first bridge installation is the identity re-emission probe. + // A failed probe leaves the old module installed and must not disrupt + // the bridge that just became reachable. + if is_direct && reemit_enabled() { + let should_reemit = original_token + .compiled + .get() + .and_then(|c| c.downcast_ref::()) + .is_some_and(|loop_| !loop_.reemitted.replace(true)); + if should_reemit { + match self.reemit_loop(original_token) { + Ok(()) => diag_bump(31), + Err(_) => diag_bump(30), + } + } + } + Ok(AsmInfo { code_addr: 0, code_size, diff --git a/majit/majit-backend-wasm/tests/codegen_test.rs b/majit/majit-backend-wasm/tests/codegen_test.rs index c907ef6a568..be7c5dd66f8 100644 --- a/majit/majit-backend-wasm/tests/codegen_test.rs +++ b/majit/majit-backend-wasm/tests/codegen_test.rs @@ -420,25 +420,28 @@ fn build_module_with_frame( gc_info: &codegen::GuardGcTypeInfo, frame: codegen::FrameGeometry, ) -> (Vec, Vec) { - let (bytes, guards, _, _, _) = codegen::build_wasm_module( - inputargs, - ops, - constants, + let inputs = codegen::ModuleBuildInputs { + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + ops: ops.iter().cloned().collect(), + inlined_bridges: Vec::new(), + constants: constants.clone(), vtable_offset, - &HashMap::new(), - gc_info, - codegen::AllocHelpers::default(), - 0, - None, // nursery - 0, // invalidated_flag_addr - 0, // gc_table_base - 0, // fail_index_base - 0, // external_jump_slot - 0, // external_jump_key + classptr_to_typeid: HashMap::new(), + guard_gc_type_info: gc_info.clone(), + alloc: codegen::AllocHelpers::default(), + wb_fn_ptr: 0, + nursery: None, + invalidated_flag_addr: 0, + gc_table_base: 0, + fail_index_base: 0, + bridge_cells_base: 0, + external_jump_slot: 0, + external_jump_key: 0, frame, - codegen::CaParams::default(), - ) - .expect("wasm codegen should succeed"); + ca: codegen::CaParams::default(), + }; + let (bytes, guards, _) = + codegen::build_wasm_module(&inputs).expect("wasm codegen should succeed"); (bytes, guards) } @@ -620,27 +623,29 @@ fn build_module_with_write_barrier_target( ops: &[Op], write_barrier_target: i64, ) -> Vec { - let (bytes, _, _, _, _) = codegen::build_wasm_module( - inputargs, - ops, - &indexmap::IndexMap::new(), - Some(0), - &HashMap::new(), - &codegen::GuardGcTypeInfo::default(), - codegen::AllocHelpers::default(), - write_barrier_target, - None, - 0, - 0, - 0, - 0, - 0, + let inputs = codegen::ModuleBuildInputs { + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + ops: ops.iter().cloned().collect(), + inlined_bridges: Vec::new(), + constants: indexmap::IndexMap::new(), + vtable_offset: Some(0), + classptr_to_typeid: HashMap::new(), + guard_gc_type_info: codegen::GuardGcTypeInfo::default(), + alloc: codegen::AllocHelpers::default(), + wb_fn_ptr: write_barrier_target, + nursery: None, + invalidated_flag_addr: 0, + gc_table_base: 0, + fail_index_base: 0, + bridge_cells_base: 0, + external_jump_slot: 0, + external_jump_key: 0, // The allocated trace keeps both Ref inputs live across New; reserve // their homes in the shared helper geometry. - codegen::FrameGeometry::compact(4, 2, 0), - codegen::CaParams::default(), - ) - .expect("wasm codegen should succeed"); + frame: codegen::FrameGeometry::compact(4, 2, 0), + ca: codegen::CaParams::default(), + }; + let (bytes, _, _) = codegen::build_wasm_module(&inputs).expect("wasm codegen should succeed"); bytes } @@ -1131,25 +1136,28 @@ fn test_cold_guard_recovery_preserves_nonzero_base_and_typed_bits() { let finish = Op::new(OpCode::Finish, &fail_args); finish.setfailargs(fail_args); let ops = [guard, finish]; - let (bytes, guards, _, _, _) = codegen::build_wasm_module( - &inputargs, - &ops, - &indexmap::IndexMap::new(), - Some(0), - &HashMap::new(), - &codegen::GuardGcTypeInfo::default(), - codegen::AllocHelpers::default(), - 0, - None, - 0, - 0, - FAIL_INDEX_BASE, - 0, - 0, - codegen::FrameGeometry::fixed(), - codegen::CaParams::default(), - ) - .expect("wasm codegen should succeed"); + let inputs = codegen::ModuleBuildInputs { + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + ops: ops.iter().cloned().collect(), + inlined_bridges: Vec::new(), + constants: indexmap::IndexMap::new(), + vtable_offset: Some(0), + classptr_to_typeid: HashMap::new(), + guard_gc_type_info: codegen::GuardGcTypeInfo::default(), + alloc: codegen::AllocHelpers::default(), + wb_fn_ptr: 0, + nursery: None, + invalidated_flag_addr: 0, + gc_table_base: 0, + fail_index_base: FAIL_INDEX_BASE, + bridge_cells_base: 0, + external_jump_slot: 0, + external_jump_key: 0, + frame: codegen::FrameGeometry::fixed(), + ca: codegen::CaParams::default(), + }; + let (bytes, guards, _) = + codegen::build_wasm_module(&inputs).expect("wasm codegen should succeed"); assert_eq!(guards[0].fail_index, FAIL_INDEX_BASE); let ref_bits = 0x1234_5678_i64; @@ -1218,6 +1226,49 @@ fn test_empty_trace() { assert!(guards[0].is_finish); } +#[test] +fn inlined_bridge_without_owner_loop_label_declines() { + let inputargs = vec![InputArg::from_type(Type::Int, 0)]; + let guard = make_guard( + OpCode::GuardTrue, + &[OpRef::input_arg_int(0)], + &[OpRef::input_arg_int(0)], + ); + let finish = Op::new(OpCode::Finish, &[rb(OpRef::input_arg_int(0))]); + let inputs = codegen::ModuleBuildInputs { + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + ops: vec![guard, finish], + inlined_bridges: vec![codegen::InlinedBridge { + source_fail_index: 0, + trace_id: 1, + inputargs: vec![InputArg::from_type(Type::Int, 1)], + ops: vec![Op::new(OpCode::Finish, &[])], + gc_table_base: 0, + }], + constants: indexmap::IndexMap::new(), + vtable_offset: Some(0), + classptr_to_typeid: HashMap::new(), + guard_gc_type_info: codegen::GuardGcTypeInfo::default(), + alloc: codegen::AllocHelpers::default(), + wb_fn_ptr: 0, + nursery: None, + invalidated_flag_addr: 0, + gc_table_base: 0, + fail_index_base: 0, + bridge_cells_base: 0, + external_jump_slot: 0, + external_jump_key: 0, + frame: codegen::FrameGeometry::fixed(), + ca: codegen::CaParams::default(), + }; + + let error = match codegen::build_wasm_module(&inputs) { + Ok(_) => panic!("a label-less owner cannot accept an inlined bridge"), + Err(error) => error, + }; + assert!(error.to_string().contains("no local loop LABEL")); +} + #[test] fn test_int_add_loop() { // Label(i, sum) -> IntAdd(sum, i) -> IntAdd(i, 1) -> IntLt(i, 100) @@ -1737,25 +1788,28 @@ fn test_guard_not_invalidated_loads_runtime_flag() { Op::new(OpCode::Finish, &[rb(OpRef::input_arg_int(0))]), ]; let constants = indexmap::IndexMap::new(); - let (bytes, guards, _, _, _) = codegen::build_wasm_module( - &inputargs, - &ops, - &constants, - None, - &HashMap::new(), - &codegen::GuardGcTypeInfo::default(), - codegen::AllocHelpers::default(), - 0, - None, - 0x1000, // invalidated_flag_addr - 0, // gc_table_base - 0, // fail_index_base - 0, // external_jump_slot - 0, // external_jump_key - codegen::FrameGeometry::fixed(), - codegen::CaParams::default(), - ) - .expect("wasm codegen should succeed"); + let inputs = codegen::ModuleBuildInputs { + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + ops: ops.iter().cloned().collect(), + inlined_bridges: Vec::new(), + constants: constants.clone(), + vtable_offset: None, + classptr_to_typeid: HashMap::new(), + guard_gc_type_info: codegen::GuardGcTypeInfo::default(), + alloc: codegen::AllocHelpers::default(), + wb_fn_ptr: 0, + nursery: None, + invalidated_flag_addr: 0x1000, + gc_table_base: 0, + fail_index_base: 0, + bridge_cells_base: 0, + external_jump_slot: 0, + external_jump_key: 0, + frame: codegen::FrameGeometry::fixed(), + ca: codegen::CaParams::default(), + }; + let (bytes, guards, _) = + codegen::build_wasm_module(&inputs).expect("wasm codegen should succeed"); validate_wasm(&bytes); assert_eq!(guards.len(), 2); @@ -2103,25 +2157,27 @@ fn loop_invariant_gc_table_load_stays_outside_non_collecting_loop() { Op::new(OpCode::Jump, &[rb(OpRef::int_op(3))]), ]; let gc_table_base = 4096; - let (bytes, _, _, _, _) = codegen::build_wasm_module( - &inputargs, - &ops, - &indexmap::IndexMap::new(), - Some(0), - &HashMap::new(), - &codegen::GuardGcTypeInfo::default(), - codegen::AllocHelpers::default(), - 0, - None, - 0, + let inputs = codegen::ModuleBuildInputs { + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + ops: ops.iter().cloned().collect(), + inlined_bridges: Vec::new(), + constants: indexmap::IndexMap::new(), + vtable_offset: Some(0), + classptr_to_typeid: HashMap::new(), + guard_gc_type_info: codegen::GuardGcTypeInfo::default(), + alloc: codegen::AllocHelpers::default(), + wb_fn_ptr: 0, + nursery: None, + invalidated_flag_addr: 0, gc_table_base, - 0, - 0, - 0, - codegen::FrameGeometry::fixed(), - codegen::CaParams::default(), - ) - .expect("wasm codegen should succeed"); + fail_index_base: 0, + bridge_cells_base: 0, + external_jump_slot: 0, + external_jump_key: 0, + frame: codegen::FrameGeometry::fixed(), + ca: codegen::CaParams::default(), + }; + let (bytes, _, _) = codegen::build_wasm_module(&inputs).expect("wasm codegen should succeed"); validate_wasm(&bytes); let mut control_stack = Vec::new(); @@ -2426,30 +2482,33 @@ fn test_non_moving_descr_allocates_through_the_oldgen_helper() { let finish = Op::new(OpCode::Finish, &[rb(OpRef::input_arg_int(0))]); finish.setfailargs(smallvec![rb(OpRef::input_arg_int(0))]); - let (bytes, _, _, _, _) = codegen::build_wasm_module( - &[InputArg::from_type(Type::Int, 0)], - &[new_op, new_array_op, finish], - &indexmap::IndexMap::new(), - Some(0), - &HashMap::new(), - &codegen::GuardGcTypeInfo::default(), - codegen::AllocHelpers { + let inputs = codegen::ModuleBuildInputs { + inputargs: vec![InputArg::from_type(Type::Int, 0)], + ops: vec![new_op, new_array_op, finish], + inlined_bridges: Vec::new(), + constants: indexmap::IndexMap::new(), + vtable_offset: Some(0), + classptr_to_typeid: HashMap::new(), + guard_gc_type_info: codegen::GuardGcTypeInfo::default(), + alloc: codegen::AllocHelpers { new_fn_ptr: NEW_FN, new_array_fn_ptr: NEW_ARRAY_FN, new_oldgen_fn_ptr: NEW_OLDGEN_FN, new_array_oldgen_fn_ptr: NEW_ARRAY_OLDGEN_FN, }, - 0, - None, // nursery: the inline bump is off, so only the helper choice shows - 0, - 0, - 0, - 0, - 0, - codegen::FrameGeometry::fixed(), - codegen::CaParams::default(), - ) - .expect("wasm codegen should succeed"); + wb_fn_ptr: 0, + nursery: None, // the inline bump is off, so only the helper choice shows + invalidated_flag_addr: 0, + gc_table_base: 0, + fail_index_base: 0, + bridge_cells_base: 0, + external_jump_slot: 0, + external_jump_key: 0, + frame: codegen::FrameGeometry::fixed(), + ca: codegen::CaParams::default(), + }; + let (bytes, _, _) = + codegen::build_wasm_module(&inputs).expect("wasm codegen should succeed"); validate_wasm(&bytes); const_immediates(&bytes) }; diff --git a/pyre/pyre-wasm-runner/src/main.rs b/pyre/pyre-wasm-runner/src/main.rs index 1c2907f144b..90d39f4e26e 100644 --- a/pyre/pyre-wasm-runner/src/main.rs +++ b/pyre/pyre-wasm-runner/src/main.rs @@ -536,6 +536,17 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { { arm.call(&mut store, ())?; } + if std::env::var_os("PYRE_WASM_REEMIT").is_some() + && let Ok(arm) = instance.get_typed_func::<(), ()>(&mut store, "pyre_jit_reemit_enable") + { + arm.call(&mut store, ())?; + } + if std::env::var_os("PYRE_WASM_INLINE_BRIDGE").is_some() + && let Ok(arm) = + instance.get_typed_func::<(), ()>(&mut store, "pyre_jit_inline_bridge_enable") + { + arm.call(&mut store, ())?; + } let src = source.as_bytes(); let len = src.len() as u32; @@ -688,6 +699,20 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { "cell_set", "cell_missing", "cell_rebridge", + "reemit_failed", + "reemit_ok", + "inline_ok", + "inline_decl_not_direct", + "inline_decl_not_loop_closing", + "inline_decl_not_reemittable", + "inline_decl_already_owned", + "inline_decl_frame", + "inline_decl_not_header", + "inline_decl_no_loop_label", + "inline_decl_value_layout", + "inline_decl_ref_layout", + "inline_decl_missing_label", + "inline_decl_other", ]; let mut parts = Vec::new(); for (i, lbl) in labels.iter().enumerate() { @@ -696,6 +721,20 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { } eprintln!("[jit-stats] bridge_diag {}", parts.join(" ")); } + if let Ok(geometry) = + instance.get_typed_func::(&mut store, "pyre_jit_inline_geometry_diag") + { + let mut parts = Vec::new(); + for i in 0..3 { + let packed = geometry.call(&mut store, i).unwrap_or(0); + if packed != 0 { + parts.push(format!("{}:{}/{}", i + 1, packed >> 32, packed as u32)); + } + } + if !parts.is_empty() { + eprintln!("[jit-stats] inline_geometry {}", parts.join(" ")); + } + } // Per-walk full-body-walk census (diagnostic). Prints the same record // the native backends print as `[fbw-census]` under PYRE_FBW_CENSUS, // which the guest cannot read. Slot layout in @@ -1079,6 +1118,21 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { // back through their own exports. Absent on a module predating them, in // which case the run keeps the old stdout-only, always-0 behaviour. let err_bytes = take_guest_stderr(&mut store, &instance, &memory)?; + if let Ok(errors) = + instance.get_typed_func::<(), u64>(&mut store, "pyre_jit_inline_trial_errors") + && let Ok(packed) = errors.call(&mut store, ()) + { + let (ptr, len) = ((packed >> 32) as u32, packed as u32); + if len != 0 { + let mut bytes = vec![0u8; len as usize]; + memory.read(&store, ptr as usize, &mut bytes)?; + dealloc.call(&mut store, (ptr, len))?; + eprintln!( + "[jit-stats] inline_trial_errors {}", + String::from_utf8_lossy(&bytes) + ); + } + } // `PYRE_FBW_DEBUG_ABORT` cannot select the walker's decline census in the // guest, and its `eprintln!` would reach nothing anyway, so the census // comes back through its own export and is printed here. Absent on a @@ -1252,6 +1306,20 @@ fn build_linker(engine: &Engine) -> Result> { }, )?; + linker.func_wrap( + "pyre_jit", + "jit_replace_wasm", + |mut caller: Caller<'_, Host>, func_id: u32, bytes_ptr: u32, bytes_len: u32| -> u32 { + match jit_replace(&mut caller, func_id, bytes_ptr, bytes_len) { + Ok(id) => id, + Err(e) => { + eprintln!("[jit_replace_wasm] {e:?}"); + 0 + } + } + }, + )?; + linker.func_wrap( "pyre_jit", "jit_execute_wasm", @@ -1406,7 +1474,11 @@ fn host_read( /// Compile and instantiate a JIT-emitted trace module, sharing the main /// module's linear memory and wiring the `jit_call` trampoline. -fn jit_compile(caller: &mut Caller<'_, Host>, bytes_ptr: u32, bytes_len: u32) -> Result { +fn jit_compile_trace( + caller: &mut Caller<'_, Host>, + bytes_ptr: u32, + bytes_len: u32, +) -> Result<(Table, Func)> { caller.data_mut().jit_compile_count += 1; let memory = caller .data() @@ -1488,16 +1560,50 @@ fn jit_compile(caller: &mut Caller<'_, Host>, bytes_ptr: u32, bytes_len: u32) -> .get_func(&mut *caller, "trace") .context("trace module is missing its `trace` export")?; + Ok((table, trace)) +} + +/// Compile and instantiate a trace, then append its export to the trace table. +fn jit_compile(caller: &mut Caller<'_, Host>, bytes_ptr: u32, bytes_len: u32) -> Result { + let (table, trace) = jit_compile_trace(caller, bytes_ptr, bytes_len)?; // Register the trace into the shared indirect function table so it is - // reachable by table index. `grow` returns the previous size, i.e. the - // index of the newly appended entry, which becomes this trace's id. + // reachable by table index. `grow` returns the newly appended slot. let slot = table .grow(&mut *caller, 1, Ref::Func(Some(trace))) .context("register trace into shared table")? as u32; - Ok(slot) } +/// Compile and instantiate a trace, then replace an existing trace-table slot. +fn jit_replace( + caller: &mut Caller<'_, Host>, + func_id: u32, + bytes_ptr: u32, + bytes_len: u32, +) -> Result { + if (func_id as u64) < caller.data().trace_base { + return Err(Error::msg(format!( + "jit_replace_wasm: id {func_id} is not a trace slot" + ))); + } + let (table, trace) = jit_compile_trace(caller, bytes_ptr, bytes_len)?; + if !matches!( + table.get(&mut *caller, func_id as u64), + Some(Ref::Func(Some(_))) + ) { + return Err(Error::msg(format!( + "jit_replace_wasm: id {func_id} is not a live trace" + ))); + } + // The Store retains every instantiated module. Replacing this table entry + // therefore leaves an old trace live when a non-tail indirect call still + // has one of its frames on the guest stack. + table + .set(&mut *caller, func_id as u64, Ref::Func(Some(trace))) + .context("replace trace in shared table")?; + Ok(func_id) +} + /// Run a previously compiled trace, returning its guard-exit index. fn jit_execute(caller: &mut Caller<'_, Host>, func_id: u32, frame_ptr: u32) -> Result { caller.data_mut().jit_execute_count += 1; diff --git a/pyre/pyre-wasm-runner/src/wasmi_host.rs b/pyre/pyre-wasm-runner/src/wasmi_host.rs index e7423234324..06b57eafb55 100644 --- a/pyre/pyre-wasm-runner/src/wasmi_host.rs +++ b/pyre/pyre-wasm-runner/src/wasmi_host.rs @@ -230,6 +230,22 @@ fn build_linker(engine: &Engine) -> Result, String> { ) .map_err(estr)?; + linker + .func_wrap( + "pyre_jit", + "jit_replace_wasm", + |mut caller: Caller<'_, Host>, func_id: u32, bytes_ptr: u32, bytes_len: u32| -> u32 { + match jit_replace(&mut caller, func_id, bytes_ptr, bytes_len) { + Ok(id) => id, + Err(e) => { + eprintln!("[jit_replace_wasm] {e}"); + 0 + } + } + }, + ) + .map_err(estr)?; + linker .func_wrap( "pyre_jit", @@ -381,11 +397,11 @@ fn host_read( /// Compile and instantiate a JIT-emitted trace module, sharing the main /// module's linear memory and wiring the `jit_call` trampoline. -fn jit_compile( +fn jit_compile_trace( caller: &mut Caller<'_, Host>, bytes_ptr: u32, bytes_len: u32, -) -> Result { +) -> Result { let memory = caller.data().memory.ok_or("main memory not initialized")?; let mut bytes = vec![0u8; bytes_len as usize]; @@ -460,6 +476,15 @@ fn jit_compile( .get_func(&*caller, "trace") .ok_or("trace module is missing its `trace` export")?; + Ok(trace) +} + +fn jit_compile( + caller: &mut Caller<'_, Host>, + bytes_ptr: u32, + bytes_len: u32, +) -> Result { + let trace = jit_compile_trace(caller, bytes_ptr, bytes_len)?; let host = caller.data_mut(); let id = host.next_id; host.next_id += 1; @@ -467,6 +492,22 @@ fn jit_compile( Ok(id) } +fn jit_replace( + caller: &mut Caller<'_, Host>, + func_id: u32, + bytes_ptr: u32, + bytes_len: u32, +) -> Result { + if !caller.data().traces.contains_key(&func_id) { + return Err(format!("jit_replace_wasm: unknown func id {func_id}")); + } + let trace = jit_compile_trace(caller, bytes_ptr, bytes_len)?; + // The map owns the old function until this overwrite. Existing calls keep + // their already-entered instance; later dispatches use this replacement. + caller.data_mut().traces.insert(func_id, trace); + Ok(func_id) +} + /// Run a previously compiled trace, returning its guard-exit index. fn jit_execute(caller: &mut Caller<'_, Host>, func_id: u32, frame_ptr: u32) -> Result { let trace = *caller diff --git a/pyre/pyre-wasm/src/lib.rs b/pyre/pyre-wasm/src/lib.rs index 6fbd1aa1957..c8c1a3d2900 100644 --- a/pyre/pyre-wasm/src/lib.rs +++ b/pyre/pyre-wasm/src/lib.rs @@ -351,6 +351,14 @@ pub extern "C" fn pyre_jit_bridge_diag(i: u32) -> u64 { majit_backend_wasm::bridge_diag(i as usize) } +/// Packed `(needed, available)` geometry for an inline-module trial that did +/// not fit its owner's frozen frame. The host formats this diagnostic only. +#[cfg(all(target_arch = "wasm32", feature = "wasm-host"))] +#[unsafe(no_mangle)] +pub extern "C" fn pyre_jit_inline_geometry_diag(i: u32) -> u64 { + majit_backend_wasm::inline_geometry_diag(i as usize) +} + /// Test-only control plane for the terminal-declined CALL_ASSEMBLER regression. /// Exported rather than imported so it cannot perturb table/function indices /// used by JIT-emitted modules. Zero disables it; see @@ -1141,6 +1149,20 @@ mod host_abi { majit_metainterp::guard_census_enable(); } + /// Arm loop-module replacement before tracing begins. The host owns the + /// environment; this guest export carries that choice into the backend. + #[unsafe(no_mangle)] + pub extern "C" fn pyre_jit_reemit_enable() { + majit_backend_wasm::reemit_enable(); + } + + /// Arm loop-closing bridge inlining. This also arms module replacement, + /// because an accepted region is installed by rebuilding its owner. + #[unsafe(no_mangle)] + pub extern "C" fn pyre_jit_inline_bridge_enable() { + majit_backend_wasm::inline_bridge_enable(); + } + /// The armed census as the same packed pair, for the host to print under /// `PYRE_WASM_GUARD_CENSUS`. Same `top` as `pyrex` prints natively, so the /// two lines compare directly. Reading, not draining. @@ -1149,6 +1171,13 @@ mod host_abi { pack_into_guest(majit_metainterp::guard_census_summary(12).into_bytes()) } + /// Trial-build errors recorded while deciding whether a bridge can be + /// merged into its loop module. The host surfaces this diagnostic text. + #[unsafe(no_mangle)] + pub extern "C" fn pyre_jit_inline_trial_errors() -> u64 { + pack_into_guest(majit_backend_wasm::inline_trial_errors().into_bytes()) + } + /// Status the last `pyre_run_python` ended with: `SystemExit`'s code, 1 for /// an uncaught exception or a `SyntaxError`, else 0. The host exits with /// it, as `pyrex` does with `targetpypystandalone.py:37 entry_point`'s From d0cdd8b022ad0dd57fa6cd81d3b52a8437537226 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Sun, 16 Aug 2026 16:36:21 +0900 Subject: [PATCH 3/7] wasm: pass a failing guard's fail args to its bridge as call parameters A bridge module's entry was `(i32) -> i32` and read its inputargs back out of positional frame slots that the failing guard had just written. Declare one entry type per arity, `(i32, i64 x n) -> i32`, bind param `k+1` to `inputargs[k]`, and carry f64 values as their raw bits so one type serves each arity. Ref-typed params are stored into the callee's own Ref homes at entry, since wasm locals are not scanned. A bridge's inputarg list is its guard's fail-arg list, so the arity and the wasm type index are constants where the guard is emitted. Each armed guard arm therefore loads its own bridge cell and, when it is nonzero, resolves its fail args straight onto the operand stack and tail-calls with its own constant type index; a zero cell keeps the spill and the return to the host. The shared epilogue no longer carries a bridge tail call. `PYRE_WASM_BRIDGE_PARAMS` now only turns this off (`0`, `false`, `off`), read on the host and applied through a guest export. Measured with `PYRE_WASM_JIT_STATS`, off vs on, executed ops: global_quasiimmut_invalidation 6,100,873,004 -> 5,573,151,872 short_circuit_value_kept_stack 10,425,576,281 -> 9,972,315,185 fannkuch 30,139,493,121 -> 27,755,382,974 fib_recursive 38,583,936,147 -> 38,140,657,358 `compile_ms` rises on all four (6.0 -> 8.6, 20.5 -> 27.6, 57.4 -> 77.6, 38.1 -> 53.2): each armed arm grows by `6 + k` instructions, and every arm carries the cold spill path whether or not a bridge ever attaches. `pyre/check.py --backend wasm --synthetic-only` is 415/415 both with the switch defaulted on and with `PYRE_WASM_BRIDGE_PARAMS=0`. Assisted-by: Claude --- majit/majit-backend-wasm/src/codegen.rs | 286 ++++++++++++++---- majit/majit-backend-wasm/src/failguard.rs | 10 + majit/majit-backend-wasm/src/lib.rs | 76 ++++- .../majit-backend-wasm/tests/codegen_test.rs | 14 + pyre/pyre-wasm-runner/src/main.rs | 14 + pyre/pyre-wasm/src/lib.rs | 8 + 6 files changed, 352 insertions(+), 56 deletions(-) diff --git a/majit/majit-backend-wasm/src/codegen.rs b/majit/majit-backend-wasm/src/codegen.rs index aea0e8124b0..80a1a38dd8a 100644 --- a/majit/majit-backend-wasm/src/codegen.rs +++ b/majit/majit-backend-wasm/src/codegen.rs @@ -34,6 +34,9 @@ const UMULHI_SCRATCH: u32 = 5; struct ValueLocals { by_id: Vec>, types: Vec, + /// First non-parameter local. Ordinary traces have one frame-pointer + /// parameter; parameter-entry bridges have that plus their fail values. + first_local: u32, } impl ValueLocals { @@ -57,7 +60,7 @@ impl ValueLocals { has_authoritative_type[i] |= authoritative; } - fn collect(inputargs: &[InputArg], ops: &[Op], num_vars: u32) -> Self { + fn collect(inputargs: &[InputArg], ops: &[Op], num_vars: u32, first_local: u32) -> Self { let mut by_id = vec![None; num_vars as usize]; let mut id_types = vec![ValType::I64; num_vars as usize]; let mut has_authoritative_type = vec![false; num_vars as usize]; @@ -133,11 +136,15 @@ impl ValueLocals { let mut types = Vec::new(); for (id, slot) in by_id.iter_mut().enumerate() { if slot.is_some() { - *slot = Some(types.len() as u32 + 1); + *slot = Some(types.len() as u32 + first_local); types.push(id_types[id]); } } - Self { by_id, types } + Self { + by_id, + types, + first_local, + } } fn local(&self, id: u32) -> u32 { @@ -149,7 +156,7 @@ impl ValueLocals { } fn ty(&self, id: u32) -> ValType { - self.types[(self.local(id) - 1) as usize] + self.types[(self.local(id) - self.first_local) as usize] } fn count(&self) -> u32 { @@ -159,6 +166,16 @@ impl ValueLocals { fn types(&self) -> &[ValType] { &self.types } + + /// Local index immediately after the dense value-local range. + fn end_local(&self) -> u32 { + self.first_local + self.count() + } + + /// Last dense value-local index, used as the base before scratch locals. + fn last_local(&self) -> u32 { + self.end_local() - 1 + } } /// Call area layout in the historical fixed frame geometry. @@ -1900,8 +1917,13 @@ pub fn guard_exit_count(inputargs: &[InputArg], ops: &[Op]) -> usize { } /// Dense wasm-local assignment and type lookup for each addressed SSA value. -fn collect_value_types(inputargs: &[InputArg], ops: &[Op], num_vars: u32) -> ValueLocals { - ValueLocals::collect(inputargs, ops, num_vars) +fn collect_value_types( + inputargs: &[InputArg], + ops: &[Op], + num_vars: u32, + first_local: u32, +) -> ValueLocals { + ValueLocals::collect(inputargs, ops, num_vars, first_local) } /// Assign each Ref-typed value (input arg or op result) a dense home-slot @@ -2083,6 +2105,12 @@ pub struct ModuleBuildInputs { pub gc_table_base: u32, pub fail_index_base: u32, pub bridge_cells_base: u32, + /// A bridge reached from an armed guard takes its fail values as `i64` + /// parameters after the frame pointer. Float bits use the same i64 carrier, + /// so a single function type per arity covers every failure signature. + pub bridge_entry_arity: Option, + /// Emit fixed-arity guard-to-bridge parameter tail-call arms for this module. + pub bridge_param_dispatch: bool, pub external_jump_slot: u32, pub external_jump_key: u32, pub frame: FrameGeometry, @@ -2148,6 +2176,8 @@ impl Clone for ModuleBuildInputs { gc_table_base: self.gc_table_base, fail_index_base: self.fail_index_base, bridge_cells_base: self.bridge_cells_base, + bridge_entry_arity: self.bridge_entry_arity, + bridge_param_dispatch: self.bridge_param_dispatch, external_jump_slot: self.external_jump_slot, external_jump_key: self.external_jump_key, frame: self.frame, @@ -2175,6 +2205,8 @@ pub fn build_wasm_module( gc_table_base, fail_index_base, bridge_cells_base, + bridge_entry_arity, + bridge_param_dispatch, external_jump_slot, external_jump_key, frame, @@ -2266,6 +2298,19 @@ pub fn build_wasm_module( // same way (a hot guard inside a chained bridge would otherwise round-trip // to the host forever). So any guarded trace wants dispatch cells. let bridge_dispatch = *bridge_cells_base != 0; + // All boundary values use an i64 carrier, including raw Float bits. That + // makes the call type depend only on arity while preserving f64 payloads. + let bridge_param_arities: Vec = if *bridge_param_dispatch && bridge_dispatch { + let mut arities: Vec = guards + .iter() + .map(|guard| guard.fail_arg_refs.len()) + .collect(); + arities.sort_unstable(); + arities.dedup(); + arities + } else { + Vec::new() + }; // Frame value slots (inputs at entry, fail-arg spills at guard exit) occupy // `[1, 1 + max(num inputs, max fail args))`. They precede the dispatch key, @@ -2286,7 +2331,21 @@ pub fn build_wasm_module( ))); } - let value_types = collect_value_types(&analysis_inputargs, &analysis_ops, num_vars); + let entry_param_count = 1 + bridge_entry_arity.unwrap_or(0) as u32; + if let Some(arity) = bridge_entry_arity + && *arity != inputargs.len() + { + return Err(BackendError::Unsupported(format!( + "wasm backend: bridge parameter arity {arity} differs from input arity {}", + inputargs.len(), + ))); + } + let value_types = collect_value_types( + &analysis_inputargs, + &analysis_ops, + num_vars, + entry_param_count, + ); let ref_values = RefValues::collect(&analysis_inputargs, &analysis_ops); let ref_homes = RefHomes::collect( &analysis_inputargs, @@ -2398,40 +2457,63 @@ pub fn build_wasm_module( // Type section let mut types = TypeSection::new(); - // Type 0: trace function (param i32) -> (result i32) + // Type 0 remains the loop/host entry signature. A bridge parameter entry + // receives a separate type so its terminal JUMP can still call a loop. types.ty().function(vec![ValType::I32], vec![ValType::I32]); - if needs_call { - // Type 1: `jit_call_compact(base, call_area_ofs)` trampoline. + let mut next_type_idx = 1u32; + let bridge_entry_type_idx = bridge_entry_arity.map(|arity| { + let idx = next_type_idx; + next_type_idx += 1; + types.ty().function( + std::iter::once(ValType::I32) + .chain(std::iter::repeat_n(ValType::I64, arity)) + .collect::>(), + vec![ValType::I32], + ); + idx + }); + let mut bridge_param_type_indices = indexmap::IndexMap::new(); + if let (Some(arity), Some(idx)) = (*bridge_entry_arity, bridge_entry_type_idx) { + bridge_param_type_indices.insert(arity, idx); + } + let jit_call_type_idx = if needs_call { + let idx = next_type_idx; + next_type_idx += 1; types .ty() .function(vec![ValType::I32, ValType::I32], vec![]); - } + Some(idx) + } else { + None + }; // Residual-call types follow: `(i64×n) -> i64` for arity `n`, indexed by // `residual_type_base + n`. `residual_type_base` = the count of types above. - let residual_type_base = 1 + needs_call as u32; + let residual_type_base = next_type_idx; if let Some(max) = residual_max_arity { for n in 0..=max { types .ty() .function(vec![ValType::I64; n], vec![ValType::I64]); } + next_type_idx += max as u32 + 1; } // CA deopt-helper type `(i64 frame_ptr, i64 compiled_ptr) -> i64`. The CA arm // `call_indirect`s `wasm_ca_resume_deopt` through it when a self-recursive // callee leaves its trace through a guard (a deopt). Declared after the // residual-call type family so its index is independent of which residual // arities the bridge happens to use. - let ca_helper_type_idx = residual_type_base + residual_max_arity.map_or(0, |m| m as u32 + 1); + let ca_helper_type_idx = next_type_idx; if ca.emit_ca { types .ty() .function(vec![ValType::I64, ValType::I64], vec![ValType::I64]); + next_type_idx += 1; } // Float residual types follow all pre-existing direct helper types. Their // parameter sequence comes from the call descr (`i64` for Int/Ref, `f64` // for Float) and their result is always `f64`; the emitter uses this map to // select the exact `call_indirect` type for each callee. - let float_residual_type_base = ca_helper_type_idx + ca.emit_ca as u32; + let float_residual_type_base = next_type_idx; let float_residual_type_indices = float_residual_sigs .iter() .cloned() @@ -2441,12 +2523,26 @@ pub fn build_wasm_module( for sig in float_residual_type_indices.keys() { types.ty().function(sig.clone(), vec![ValType::F64]); } - let true_void_residual_type_base = - float_residual_type_base + float_residual_type_indices.len() as u32; + next_type_idx += float_residual_type_indices.len() as u32; + let true_void_residual_type_base = next_type_idx; if let Some(max) = true_void_residual_max_arity { for n in 0..=max { types.ty().function(vec![ValType::I64; n], vec![]); } + next_type_idx += max as u32 + 1; + } + for arity in bridge_param_arities { + if bridge_param_type_indices.contains_key(&arity) { + continue; + } + bridge_param_type_indices.insert(arity, next_type_idx); + next_type_idx += 1; + types.ty().function( + std::iter::once(ValType::I32) + .chain(std::iter::repeat_n(ValType::I64, arity)) + .collect::>(), + vec![ValType::I32], + ); } module.section(&types); @@ -2465,7 +2561,11 @@ pub fn build_wasm_module( ); if needs_call { // Import jit_call trampoline as function index 0 - imports.import("env", "jit_call_compact", EntityType::Function(1)); + imports.import( + "env", + "jit_call_compact", + EntityType::Function(jit_call_type_idx.expect("jit_call type")), + ); } if needs_table { // Import the host's shared indirect function table as table index 0. @@ -2490,7 +2590,7 @@ pub fn build_wasm_module( // Function section let mut functions = FunctionSection::new(); - functions.function(0); // type 0 + functions.function(bridge_entry_type_idx.unwrap_or(0)); module.section(&functions); // Export section: trace function index depends on whether we imported jit_call @@ -2522,6 +2622,8 @@ pub fn build_wasm_module( &label_resume, *bridge_cells_base, bridge_dispatch, + *bridge_entry_arity, + &bridge_param_type_indices, *invalidated_flag_addr, *gc_table_base, &gc_table_bases, @@ -2563,6 +2665,8 @@ fn build_function( label_resume: &LabelResumeData, cells_base: u32, bridge_dispatch: bool, + bridge_entry_arity: Option, + bridge_param_type_indices: &indexmap::IndexMap, invalidated_flag_addr: u32, gc_table_base: u32, gc_table_bases: &HashMap, @@ -2602,26 +2706,27 @@ fn build_function( // while `WASM_DIRECT_RESIDUAL_CALL` is enabled). Its `jit_call` fallback // branches are retained solely for the direct-family-disabled baseline. debug_assert!(!ca.emit_ca || residual_type_base.is_some()); - let num_value_locals = value_types.count(); + let value_locals_end = value_types.end_local(); // Value locals occupy the dense local range beginning at 1; reserve // `UMULHI_SCRATCH` i64 locals past them for the `UintMulHigh` // 32-bit-split expansion, plus one i64 local for the pending overflow flag. - // One i32 local past those holds the bridge table slot for the epilogue - // `call_indirect` dispatch (unused when `!bridge_dispatch`). - let ovf_flag_local = num_value_locals + UMULHI_SCRATCH + 1; - let bridge_slot_local = num_value_locals + UMULHI_SCRATCH + 2; + // One i32 local past those holds a bridge table slot while a guard arm + // performs its direct indirect tail call (or while the frame-entry + // dispatcher is enabled without parameter entries). + let ovf_flag_local = value_locals_end + UMULHI_SCRATCH; + let bridge_slot_local = ovf_flag_local + 1; // The self-recursive CALL_ASSEMBLER arm needs two more i32 scratch locals: // `ca_cfp_local` (the current callee frame pointer) and `ca_fi_local` (the // returned frame[0] fail index). Reserve them only under `emit_ca` so a // flag-off module keeps exactly one i32 local (byte-identical). - let ca_cfp_local = num_value_locals + UMULHI_SCRATCH + 3; - let ca_fi_local = num_value_locals + UMULHI_SCRATCH + 4; + let ca_cfp_local = bridge_slot_local + 1; + let ca_fi_local = ca_cfp_local + 1; // Extra i32 scratches when the inline nursery-bump fast path is armed: // one holds the loaded `nursery_free` across the bump/commit sequence; // runtime varsize array allocation also needs one for the computed // total/new-free word. - let base_i32_locals: u32 = if ca.emit_ca { 3 } else { 1 }; - let alloc_scratch_local = num_value_locals + UMULHI_SCRATCH + 2 + base_i32_locals; + let base_i32_locals: u32 = 1 + if ca.emit_ca { 2 } else { 0 }; + let alloc_scratch_local = bridge_slot_local + base_i32_locals; let alloc_size_local = alloc_scratch_local + 1; debug_assert_eq!(bridge_slot_local, ovf_flag_local + 1); debug_assert_eq!(ca_cfp_local, bridge_slot_local + 1); @@ -2645,6 +2750,7 @@ fn build_function( fail_index_base, bridge_slot_local, enabled: bridge_dispatch, + param_type_indices: bridge_param_type_indices, inline_guards: &inline_guards, ref_homes, frame, @@ -2825,10 +2931,19 @@ fn build_function( // resume loader sets the live label-arg homes. for (k, ia) in entry_inputargs.iter().enumerate() { let local_idx = value_types.local(ia.index); - let offset = FRAME_SLOT_BASE + k as u64 * SLOT_SIZE; - sink.local_get(0).i64_load(mem64(offset)); - if value_types.ty(ia.index) == ValType::F64 { - sink.f64_reinterpret_i64(); + if bridge_entry_arity.is_some() { + // Parameter entries carry raw i64 words after frame_ptr. Float + // values use their IEEE bit pattern, matching the guard boundary. + sink.local_get(k as u32 + 1); + if value_types.ty(ia.index) == ValType::F64 { + sink.f64_reinterpret_i64(); + } + } else { + let offset = FRAME_SLOT_BASE + k as u64 * SLOT_SIZE; + sink.local_get(0).i64_load(mem64(offset)); + if value_types.ty(ia.index) == ValType::F64 { + sink.f64_reinterpret_i64(); + } } sink.local_set(local_idx); if let Some(h) = ref_homes.home_id(ia.index) { @@ -3192,11 +3307,15 @@ fn build_function( } OpCode::Finish => { - emit_guard_spill(&mut sink, constants, value_types, guard_idx, op); - if guard_dispatch.enabled { - emit_guard_bridge_dispatch(&mut sink, guard_idx, guard_dispatch); - } - sink.br(block_exit_depth); + emit_guard_exit( + &mut sink, + constants, + value_types, + guard_idx, + op, + block_exit_depth, + guard_dispatch, + ); guard_idx += 1; } @@ -3484,9 +3603,13 @@ fn build_function( // High 64 bits of the unsigned 64×64→128 product. The optimizer // emits this for division/modulo-by-constant strength reduction; // wasm has no mul-high instruction, so expand via 32-bit split. - OpCode::UintMulHigh => { - emit_umulhi(&mut sink, constants, value_types, op, value_types.count()) - } + OpCode::UintMulHigh => emit_umulhi( + &mut sink, + constants, + value_types, + op, + value_types.last_local(), + ), // Overflow variants: compute result + overflow flag OpCode::IntAddOvf | OpCode::IntSubOvf | OpCode::IntMulOvf => { @@ -3507,7 +3630,7 @@ fn build_function( value_types, op, binop, - value_types.count(), + value_types.last_local(), ovf_flag_local, fused_guard.map(|guard| guard.opcode), ) { @@ -4241,7 +4364,16 @@ fn build_function( // what the exit reports, not just how it gets there. The // epilogue takes a cell address when bridge dispatch is on; // otherwise retain the unused index write unchanged. - if guard_dispatch.enabled { + if !guard_dispatch.param_type_indices.is_empty() { + emit_guard_param_tail_call( + &mut sink, + constants, + value_types, + guard_idx, + op, + guard_dispatch, + ); + } else if guard_dispatch.enabled { emit_guard_bridge_dispatch(&mut sink, guard_idx, guard_dispatch); } else { sink.i32_const(guard_idx as i32); @@ -5672,15 +5804,11 @@ fn build_function( sink.return_(); sink.end(); // end A $hot_exit - // Epilogue bridge dispatch for exits that branch out of the hot exit - // block. Their arms have already placed the exit's constant cell address - // in `bridge_slot_local`; load the table slot from that cell. If a bridge - // has been compiled (slot != 0), tail into it via `call_indirect` through - // the shared table and return its result. Otherwise fall through to the - // host round-trip (return `frame_ptr`, whose `frame[0]` was written by the - // ordinary guard or Finish arm). With every cell 0 (no bridge yet) this is - // inert and behavior is unchanged. - if bridge_dispatch { + // Frame-entry bridge dispatch for exits that branch out of the hot exit + // block. Parameter-entry bridges tail-call from their own guard arm: that + // arm knows the fixed failure arity and therefore the fixed wasm type. + // The shared epilogue remains only for the established frame-entry form. + if bridge_dispatch && bridge_param_type_indices.is_empty() { // slot = *(bridge_slot_local), where the local holds a cell address. sink.local_get(bridge_slot_local); sink.i32_load(memarg(0, 2)); @@ -6108,6 +6236,9 @@ struct BridgeDispatch<'a> { fail_index_base: u32, bridge_slot_local: u32, enabled: bool, + /// `arity -> indirect-call type` for armed parameter dispatch. Every + /// signature carries values as i64, including Float bit patterns. + param_type_indices: &'a indexmap::IndexMap, inline_guards: &'a [InlineGuard<'a>], ref_homes: &'a RefHomes, frame: FrameGeometry, @@ -6298,13 +6429,62 @@ fn emit_guard_exit( sink.br(inline.branch_depth + 1); return; } - emit_guard_spill(sink, constants, value_types, guard_idx, op); - if dispatch.enabled { - emit_guard_bridge_dispatch(sink, guard_idx, dispatch); + if dispatch.param_type_indices.is_empty() { + emit_guard_spill(sink, constants, value_types, guard_idx, op); + if dispatch.enabled { + emit_guard_bridge_dispatch(sink, guard_idx, dispatch); + } + } else { + emit_guard_param_tail_call(sink, constants, value_types, guard_idx, op, dispatch); + // A missing cell keeps the historical recovery path. It is deliberately + // after the cell test so a bridge crossing performs no frame spill. + emit_guard_spill(sink, constants, value_types, guard_idx, op); } sink.br(block_exit_depth); } +/// Tail-call this guard's bridge directly when its cell is armed. The guard's +/// failure list fixes both the values and the wasm function type, so this path +/// needs neither an arity tag nor staging locals. +fn emit_guard_param_tail_call( + sink: &mut PeepSink<'_, '_>, + constants: &indexmap::IndexMap, + value_types: &ValueLocals, + guard_idx: u32, + op: &Op, + dispatch: BridgeDispatch<'_>, +) { + let fail_args: Vec = op + .getfailargs() + .map(|args| args.iter().map(|arg| arg.to_opref()).collect()) + .unwrap_or_else(|| op.getarglist().iter().map(|arg| arg.to_opref()).collect()); + let arity = fail_args.len(); + let type_idx = *dispatch + .param_type_indices + .get(&arity) + .expect("parameter dispatch type missing for guard fail arity"); + debug_assert!(dispatch.enabled); + debug_assert!(guard_idx >= dispatch.fail_index_base); + let cell_addr = dispatch.cells_base + + (guard_idx - dispatch.fail_index_base) * std::mem::size_of::() as u32; + sink.i32_const(cell_addr as i32); + sink.i32_load(memarg(0, 2)); + sink.local_tee(dispatch.bridge_slot_local); + sink.if_(BlockType::Empty); + sink.local_get(0); + for arg in fail_args { + if arg.ty() == Some(Type::Float) { + emit_resolve_f64(sink, constants, value_types, arg); + sink.i64_reinterpret_f64(); + } else { + emit_resolve(sink, constants, value_types, arg); + } + } + sink.local_get(dispatch.bridge_slot_local); + sink.return_call_indirect(0, type_idx); + sink.end(); +} + /// Transfer a failing guard directly into an inlined bridge. All sources are /// pushed before any destination local is written, preserving parallel-move /// semantics when fail arguments overlap bridge input locals. diff --git a/majit/majit-backend-wasm/src/failguard.rs b/majit/majit-backend-wasm/src/failguard.rs index a3442e562df..1669c5d1c0f 100644 --- a/majit/majit-backend-wasm/src/failguard.rs +++ b/majit/majit-backend-wasm/src/failguard.rs @@ -568,6 +568,11 @@ pub struct ChainedTraceMeta { /// Per-guard, per-fail-arg induction-advance flags /// (`CompiledWasmLoop::guard_fail_arg_advanced` analog). pub guard_fail_arg_advanced: Vec>, + /// Number of values each guard transfers to a bridge. A parameter entry + /// is admitted only when this agrees with the bridge's input list. + pub guard_fail_arg_counts: Vec, + /// Whether this trace's guard epilogue has typed parameter dispatch arms. + pub bridge_param_dispatch: bool, } /// Compiled wasm loop metadata, stored in `JitCellToken.compiled`. @@ -642,6 +647,11 @@ pub struct CompiledWasmLoop { /// `compile_bridge`'s livelock check: a loop-closing bridge that JUMPs /// such a fail arg verbatim still advances the chained cycle. pub guard_fail_arg_advanced: Vec>, + /// Number of fail arguments for every guard/finish exit in this trace. + pub guard_fail_arg_counts: Vec, + /// Whether this module transfers a compiled bridge's fail arguments as + /// wasm call parameters instead of reloading their positional frame slots. + pub bridge_param_dispatch: bool, /// `(source_trace_id, source_fail_index, start, count)` ranges into /// `fail_descrs` for each chained bridge `compile_bridge` appended (lib.rs /// extend site). Lets `compiled_bridge_fail_descr_layouts` / diff --git a/majit/majit-backend-wasm/src/lib.rs b/majit/majit-backend-wasm/src/lib.rs index f9f36391f46..61d6be6b9f3 100644 --- a/majit/majit-backend-wasm/src/lib.rs +++ b/majit/majit-backend-wasm/src/lib.rs @@ -78,8 +78,11 @@ use std::sync::{Arc, Mutex}; /// owner's frozen frame geometry; 38 = the bridge does not resume at the loop /// header; 39 = the merged stream has no local loop LABEL for the wasm back /// edge. 40-43 split a rejected inline trial into value-layout, -/// Ref-home-layout, missing-local-label, and other backend errors. -pub static BRIDGE_DIAG: [AtomicU64; 44] = [const { AtomicU64::new(0) }; 44]; +/// Ref-home-layout, missing-local-label, and other backend errors. 44 = a +/// bridge compiled with a parameter entry; 45 = parameter entry declined +/// because the source module has frame-only dispatch; 46 = parameter entry +/// declined because the source guard and bridge input arities disagree. +pub static BRIDGE_DIAG: [AtomicU64; 47] = [const { AtomicU64::new(0) }; 47]; /// The first three inline geometry failures, packed as `(needed, available)`. /// They expose a frozen-layout shortage without changing the compile result. @@ -117,6 +120,7 @@ fn record_inline_trial_error(error: &BackendError) { static REEMIT_ENABLED: AtomicBool = AtomicBool::new(false); static INLINE_BRIDGE_ENABLED: AtomicBool = AtomicBool::new(false); +static BRIDGE_PARAMS_ENABLED: AtomicBool = AtomicBool::new(true); /// Arm loop-module replacement from the host before guest execution starts. pub fn reemit_enable() { @@ -138,6 +142,18 @@ fn inline_bridge_enabled() -> bool { INLINE_BRIDGE_ENABLED.load(Ordering::Relaxed) } +/// Disable guard-to-bridge value parameters from the host before guest +/// execution. By default, a generated guard keeps the ordinary frame recovery +/// state for the uncompiled case, then passes its live failure values directly +/// once a bridge table slot is present. +pub fn bridge_params_disable() { + BRIDGE_PARAMS_ENABLED.store(false, Ordering::Relaxed); +} + +fn bridge_params_enabled() -> bool { + BRIDGE_PARAMS_ENABLED.load(Ordering::Relaxed) +} + /// Read a `BRIDGE_DIAG` tally (saturating index). Surfaced to the host through /// the `pyre_jit_bridge_diag` export in the `pyre-wasm` crate. pub fn bridge_diag(i: usize) -> u64 { @@ -1828,6 +1844,11 @@ impl WasmBackend { cells_base: new_cells_base + offset as u32 * 4, num_cells: count, guard_fail_arg_advanced: guard_fail_args_advanced(®ion.ops, exits), + guard_fail_arg_counts: exits + .iter() + .map(|guard| guard.fail_arg_refs.len()) + .collect(), + bridge_param_dispatch: inputs.bridge_param_dispatch, }, ); offset += count; @@ -2556,6 +2577,8 @@ impl majit_backend::Backend for WasmBackend { gc_table_base, fail_index_base, bridge_cells_base, + bridge_entry_arity: None, + bridge_param_dispatch: bridge_params_enabled(), // A real loop's JUMP is a local back-edge `br`; an entry bridge // tail-calls its target loop and is deliberately not re-emittable. external_jump_slot: entry_bridge_target.map_or(0, |t| t.func_handle), @@ -2705,6 +2728,11 @@ impl majit_backend::Backend for WasmBackend { has_preamble, label_descrs, guard_fail_arg_advanced, + guard_fail_arg_counts: guard_exits + .iter() + .map(|guard| guard.fail_arg_refs.len()) + .collect(), + bridge_param_dispatch: bridge_params_enabled(), bridge_descr_ranges: std::cell::RefCell::new(Vec::new()), chained_trace_meta: std::cell::RefCell::new(std::collections::HashMap::new()), _bridge_owned_cells: std::cell::RefCell::new(bridge_cells_owner.into_iter().collect()), @@ -2875,6 +2903,11 @@ impl majit_backend::Backend for WasmBackend { .get(source_fail_index as usize) .cloned() .unwrap_or_default(), + source_loop + .guard_fail_arg_counts + .get(source_fail_index as usize) + .copied(), + source_loop.bridge_param_dispatch, )) } else { source_loop @@ -2889,6 +2922,10 @@ impl majit_backend::Backend for WasmBackend { .get(source_fail_index as usize) .cloned() .unwrap_or_default(), + m.guard_fail_arg_counts + .get(source_fail_index as usize) + .copied(), + m.bridge_param_dispatch, ) }) }; @@ -2906,7 +2943,13 @@ impl majit_backend::Backend for WasmBackend { // that trace's array. A foreign descr has no cell to flip; decline so // the metainterp keeps the correct interpreter fallback rather than // installing an unreachable bridge module. - let Some((source_cells_base, source_num_cells, source_fail_arg_advanced)) = source_guard + let Some(( + source_cells_base, + source_num_cells, + source_fail_arg_advanced, + source_fail_arg_count, + source_bridge_param_dispatch, + )) = source_guard else { diag_bump(3); // declined: source guard's trace is not chained here return Err(BackendError::Unsupported( @@ -2919,6 +2962,23 @@ impl majit_backend::Backend for WasmBackend { "wasm backend: bridge source guard index has no dispatch cell".into(), )); } + let bridge_entry_arity = if bridge_params_enabled() { + if !source_bridge_param_dispatch { + diag_bump(45); + return Err(BackendError::Unsupported( + "wasm backend: source guard has no parameter bridge dispatch".into(), + )); + } + if source_fail_arg_count != Some(inputargs.len()) { + diag_bump(46); + return Err(BackendError::Unsupported( + "wasm backend: guard and bridge input arities differ".into(), + )); + } + Some(inputargs.len()) + } else { + None + }; let allow_ca = ca_candidate; if let Some(reason) = wasm_unsupported_trace_reason(ops, allow_ca) { diag_bump(1); // declined: CALL_ASSEMBLER @@ -3263,6 +3323,8 @@ impl majit_backend::Backend for WasmBackend { gc_table_base, fail_index_base: base, bridge_cells_base, + bridge_entry_arity, + bridge_param_dispatch: bridge_params_enabled(), external_jump_slot, external_jump_key, frame: source_frame, @@ -3313,6 +3375,9 @@ impl majit_backend::Backend for WasmBackend { Self::register_gc_table(original_token, table); } diag_bump(5); // bridge compiled — chained in-module + if bridge_entry_arity.is_some() { + diag_bump(44); // bridge compiled with a parameter entry + } // x86/assembler.py:706 publishes the target tokens defined by an // accepted bridge. `codegen::is_resumable_peeled` and @@ -3363,6 +3428,11 @@ impl majit_backend::Backend for WasmBackend { cells_base: bridge_cells_base, num_cells: guard_exits.len(), guard_fail_arg_advanced: guard_fail_args_advanced(ops, &guard_exits), + guard_fail_arg_counts: guard_exits + .iter() + .map(|guard| guard.fail_arg_refs.len()) + .collect(), + bridge_param_dispatch: bridge_params_enabled(), }, ); // The bridge module lives as long as this source loop, so hand its diff --git a/majit/majit-backend-wasm/tests/codegen_test.rs b/majit/majit-backend-wasm/tests/codegen_test.rs index be7c5dd66f8..1c00a9702e1 100644 --- a/majit/majit-backend-wasm/tests/codegen_test.rs +++ b/majit/majit-backend-wasm/tests/codegen_test.rs @@ -435,6 +435,8 @@ fn build_module_with_frame( gc_table_base: 0, fail_index_base: 0, bridge_cells_base: 0, + bridge_entry_arity: None, + bridge_param_dispatch: false, external_jump_slot: 0, external_jump_key: 0, frame, @@ -638,6 +640,8 @@ fn build_module_with_write_barrier_target( gc_table_base: 0, fail_index_base: 0, bridge_cells_base: 0, + bridge_entry_arity: None, + bridge_param_dispatch: false, external_jump_slot: 0, external_jump_key: 0, // The allocated trace keeps both Ref inputs live across New; reserve @@ -1151,6 +1155,8 @@ fn test_cold_guard_recovery_preserves_nonzero_base_and_typed_bits() { gc_table_base: 0, fail_index_base: FAIL_INDEX_BASE, bridge_cells_base: 0, + bridge_entry_arity: None, + bridge_param_dispatch: false, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), @@ -1256,6 +1262,8 @@ fn inlined_bridge_without_owner_loop_label_declines() { gc_table_base: 0, fail_index_base: 0, bridge_cells_base: 0, + bridge_entry_arity: None, + bridge_param_dispatch: false, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), @@ -1803,6 +1811,8 @@ fn test_guard_not_invalidated_loads_runtime_flag() { gc_table_base: 0, fail_index_base: 0, bridge_cells_base: 0, + bridge_entry_arity: None, + bridge_param_dispatch: false, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), @@ -2172,6 +2182,8 @@ fn loop_invariant_gc_table_load_stays_outside_non_collecting_loop() { gc_table_base, fail_index_base: 0, bridge_cells_base: 0, + bridge_entry_arity: None, + bridge_param_dispatch: false, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), @@ -2502,6 +2514,8 @@ fn test_non_moving_descr_allocates_through_the_oldgen_helper() { gc_table_base: 0, fail_index_base: 0, bridge_cells_base: 0, + bridge_entry_arity: None, + bridge_param_dispatch: false, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), diff --git a/pyre/pyre-wasm-runner/src/main.rs b/pyre/pyre-wasm-runner/src/main.rs index 90d39f4e26e..df7875d8b6c 100644 --- a/pyre/pyre-wasm-runner/src/main.rs +++ b/pyre/pyre-wasm-runner/src/main.rs @@ -547,6 +547,17 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { { arm.call(&mut store, ())?; } + // Parameter bridge entries are the default. The guest has no environment, + // so an explicit host-side opt-out must travel through this export before + // tracing begins. + if std::env::var_os("PYRE_WASM_BRIDGE_PARAMS").is_some_and(|value| { + matches!(value.to_str().map(str::trim), Some("0" | "false" | "off")) + }) + && let Ok(arm) = + instance.get_typed_func::<(), ()>(&mut store, "pyre_jit_bridge_params_disable") + { + arm.call(&mut store, ())?; + } let src = source.as_bytes(); let len = src.len() as u32; @@ -713,6 +724,9 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { "inline_decl_ref_layout", "inline_decl_missing_label", "inline_decl_other", + "bridge_param_ok", + "bridge_param_decl_source_frame", + "bridge_param_decl_arity", ]; let mut parts = Vec::new(); for (i, lbl) in labels.iter().enumerate() { diff --git a/pyre/pyre-wasm/src/lib.rs b/pyre/pyre-wasm/src/lib.rs index c8c1a3d2900..c46107a0cca 100644 --- a/pyre/pyre-wasm/src/lib.rs +++ b/pyre/pyre-wasm/src/lib.rs @@ -1163,6 +1163,14 @@ mod host_abi { majit_backend_wasm::inline_bridge_enable(); } + /// Disable the default guard-to-bridge parameter entries. The host owns + /// the environment, so this call carries its explicit opt-out into the + /// guest before tracing begins. + #[unsafe(no_mangle)] + pub extern "C" fn pyre_jit_bridge_params_disable() { + majit_backend_wasm::bridge_params_disable(); + } + /// The armed census as the same packed pair, for the host to print under /// `PYRE_WASM_GUARD_CENSUS`. Same `top` as `pyrex` prints natively, so the /// two lines compare directly. Reading, not draining. From 0efae530dfc82b7987d7f05a350b6ca3433d4e9a Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Mon, 17 Aug 2026 00:05:22 +0900 Subject: [PATCH 4/7] wasm: do not publish LABEL targets on a bridge whose entry takes parameters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `compile_bridge` publishes an accepted bridge's own LABELs keyed on that bridge's table slot. A later trace whose cross-module terminal JUMP resolves to one of those targets emits `return_call_indirect(0, 0)` — type 0, which is `(i32) -> i32`. Since a bridge entry can declare `(i32, i64 x n) -> i32`, the two can disagree, and an indirect call whose declared type does not match the callee's traps at run time. Suppress publication when `bridge_entry_arity` is `Some(n)` with `n > 0`. `Some(0)` keeps publishing: `call_indirect` compares signatures structurally, and a zero-arity parameter entry is `(i32) -> i32` even though codegen mints it a separate type index. The emitted-module test asserts that equality. `bridge_param_label_suppressed` counts the suppression. Over the four graded fixtures plus `const_arg_call_resume` and `retrace_accumulator_type_flip` it fires once, on the last of those. `pyre/check.py --backend wasm --synthetic-only` is 415/415. Assisted-by: Claude --- majit/majit-backend-wasm/src/lib.rs | 61 ++++++++++++------- .../majit-backend-wasm/tests/codegen_test.rs | 52 ++++++++++++++++ pyre/pyre-wasm-runner/src/main.rs | 1 + 3 files changed, 92 insertions(+), 22 deletions(-) diff --git a/majit/majit-backend-wasm/src/lib.rs b/majit/majit-backend-wasm/src/lib.rs index 61d6be6b9f3..76568fe7941 100644 --- a/majit/majit-backend-wasm/src/lib.rs +++ b/majit/majit-backend-wasm/src/lib.rs @@ -81,8 +81,9 @@ use std::sync::{Arc, Mutex}; /// Ref-home-layout, missing-local-label, and other backend errors. 44 = a /// bridge compiled with a parameter entry; 45 = parameter entry declined /// because the source module has frame-only dispatch; 46 = parameter entry -/// declined because the source guard and bridge input arities disagree. -pub static BRIDGE_DIAG: [AtomicU64; 47] = [const { AtomicU64::new(0) }; 47]; +/// declined because the source guard and bridge input arities disagree; 47 = +/// LABEL publication suppressed because the bridge entry has nonzero parameters. +pub static BRIDGE_DIAG: [AtomicU64; 48] = [const { AtomicU64::new(0) }; 48]; /// The first three inline geometry failures, packed as `(needed, available)`. /// They expose a frozen-layout shortage without changing the compile result. @@ -319,6 +320,7 @@ fn stamp_and_publish_label_targets( frame: codegen::FrameGeometry, inputargs: &[InputArg], ops: &[Op], + bridge_entry_arity: Option, ) -> (Vec, Vec) { // Stamp each LABEL's loop-target descr with its ordinal (0, 1, 2, …) so a // loop-closing bridge can recover which label its terminal JUMP targets: @@ -361,6 +363,10 @@ fn stamp_and_publish_label_targets( let label_num_args = codegen::label_arg_counts(ops); let label_resume_info = codegen::label_resume_info(inputargs, ops, frame); let mut published_descrs = Vec::new(); + // A parameter entry with no fail values remains structurally `(i32) -> + // i32`, so type-0 indirect calls may enter it. Only a nonzero parameter + // entry is incompatible with published LABEL targets. + let suppress_publication = matches!(bridge_entry_arity, Some(arity) if arity > 0); // Publish this loop's enterable labels so a loop-closing bridge from // ANY loop can chain into them in-module (jump-to-existing-trace). A @@ -380,20 +386,24 @@ fn stamp_and_publish_label_targets( if id == 0 { continue; } - diag_bump(19); - publish_label_target( - id, - LabelTarget { - func_handle, - key: j as u32 + 1, - num_args: label_num_args[j], - resume_safe: label_resume_info[j].0, - requires_own_frame: label_resume_info[j].1, - is_last_label: j == header, - frame, - }, - ); - published_descrs.push(id); + if suppress_publication { + diag_bump(47); + } else { + diag_bump(19); + publish_label_target( + id, + LabelTarget { + func_handle, + key: j as u32 + 1, + num_args: label_num_args[j], + resume_safe: label_resume_info[j].0, + requires_own_frame: label_resume_info[j].1, + is_last_label: j == header, + frame, + }, + ); + published_descrs.push(id); + } } } else { // A LABEL with real work before it is not reachable through the plain @@ -415,7 +425,9 @@ fn stamp_and_publish_label_targets( if !publishable && !label_descrs.is_empty() { diag_bump(21); } - if publishable { + if publishable && suppress_publication { + diag_bump(47); + } else if publishable { let id = label_descrs[0]; diag_bump(20); publish_label_target( @@ -1863,6 +1875,7 @@ impl WasmBackend { compiled.frame, &inputs.inputargs, &inputs.ops, + inputs.bridge_entry_arity, ); let loop_finish_fi = descrs .iter() @@ -2707,7 +2720,8 @@ impl majit_backend::Backend for WasmBackend { // the last LABEL. Computed through the same predicate codegen's wrapper // gates on, so the recorded field and the emitted wrapper cannot drift. let has_preamble = codegen::is_resumable_peeled(ops); - let (label_descrs, _) = stamp_and_publish_label_targets(func_handle, frame, inputargs, ops); + let (label_descrs, _) = + stamp_and_publish_label_targets(func_handle, frame, inputargs, ops, None); // Per-guard, per-fail-arg induction-advance flags for // `compile_bridge`'s livelock check (see `guard_fail_args_advanced`). let guard_fail_arg_advanced = guard_fail_args_advanced(ops, &guard_exits); @@ -3390,8 +3404,13 @@ impl majit_backend::Backend for WasmBackend { // LABEL; the // existing `first_label_at_entry` / arity guard correctly leaves that // label unpublished, because key 0 would re-run the work before it. - let (_, published_label_descrs) = - stamp_and_publish_label_targets(bridge_slot, source_frame, inputargs, ops); + let (_, published_label_descrs) = stamp_and_publish_label_targets( + bridge_slot, + source_frame, + inputargs, + ops, + bridge_entry_arity, + ); { let source_loop = original_token @@ -3475,8 +3494,6 @@ impl majit_backend::Backend for WasmBackend { } #[cfg(all(target_arch = "wasm32", not(target_os = "wasi")))] if source_cells_base != 0 && bridge_slot != 0 { - // cells[source_fail_index] = bridge_slot — the loop epilogue now - // tails into this bridge instead of returning to the host. let cell = (source_cells_base as usize + source_fail_index as usize * 4) as *mut u32; if unsafe { core::ptr::read(cell) } != 0 { diag_bump(29); // this guard already had a reachable bridge diff --git a/majit/majit-backend-wasm/tests/codegen_test.rs b/majit/majit-backend-wasm/tests/codegen_test.rs index 1c00a9702e1..6bc2031ac67 100644 --- a/majit/majit-backend-wasm/tests/codegen_test.rs +++ b/majit/majit-backend-wasm/tests/codegen_test.rs @@ -1510,6 +1510,58 @@ fn function_type( panic!("module has no type section"); } +fn entry_function_type_index(bytes: &[u8]) -> usize { + for payload in wasmparser::Parser::new(0).parse_all(bytes) { + if let wasmparser::Payload::FunctionSection(functions) = payload.unwrap() { + return functions + .into_iter() + .next() + .expect("module has an entry function") + .expect("entry function type is valid") as usize; + } + } + panic!("module has no function section"); +} + +#[test] +fn zero_arity_parameter_entry_is_structurally_type_zero() { + // A published LABEL target is entered through type 0, `(i32) -> i32`. + // The separate type index emitted for a zero-arity parameter bridge must + // retain that same structural signature. + let inputargs = Vec::new(); + let ops = vec![Op::new(OpCode::Label, &[]), Op::new(OpCode::Finish, &[])]; + let inputs = codegen::ModuleBuildInputs { + inputargs, + ops, + inlined_bridges: Vec::new(), + constants: indexmap::IndexMap::new(), + vtable_offset: Some(0), + classptr_to_typeid: HashMap::new(), + guard_gc_type_info: codegen::GuardGcTypeInfo::default(), + alloc: codegen::AllocHelpers::default(), + wb_fn_ptr: 0, + nursery: None, + invalidated_flag_addr: 0, + gc_table_base: 0, + fail_index_base: 0, + bridge_cells_base: 4, + bridge_entry_arity: Some(0), + bridge_param_dispatch: true, + external_jump_slot: 0, + external_jump_key: 0, + frame: codegen::FrameGeometry::fixed(), + ca: codegen::CaParams::default(), + }; + let (bytes, _, _) = codegen::build_wasm_module(&inputs).unwrap(); + + validate_wasm(&bytes); + assert_eq!( + function_type(&bytes, entry_function_type_index(&bytes)), + function_type(&bytes, 0), + "a published LABEL target's module entry must be structurally `(i32) -> i32`" + ); +} + #[test] fn test_nullary_true_void_call_uses_indirect_call_without_drop() { let inputargs = vec![InputArg::from_type(Type::Int, 0)]; diff --git a/pyre/pyre-wasm-runner/src/main.rs b/pyre/pyre-wasm-runner/src/main.rs index df7875d8b6c..4998ebd4e0e 100644 --- a/pyre/pyre-wasm-runner/src/main.rs +++ b/pyre/pyre-wasm-runner/src/main.rs @@ -727,6 +727,7 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { "bridge_param_ok", "bridge_param_decl_source_frame", "bridge_param_decl_arity", + "bridge_param_label_suppressed", ]; let mut parts = Vec::new(); for (i, lbl) in labels.iter().enumerate() { From d270418e8e77c7c96e56612f56a2a26d82d50c66 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Mon, 17 Aug 2026 01:23:58 +0900 Subject: [PATCH 5/7] wasm: count trace entries per resume key `executes` counts host entries only, so it is blind to the in-guest `return_call_indirect` traffic between a loop and its bridges. Nothing said how often a trace module is entered, or through which `br_table` key, which is what an executed-op budget has to be divided by. Add a per-trace, per-key counter in guest memory, armed from the host: `PYRE_WASM_TRACE_ENTRY_CENSUS` calls `pyre_jit_trace_entry_census_enable`, and `pyre_jit_trace_entry_census` reads the table back. The counter instructions, their locals and the census globals are emitted only when the census is armed at compile time; the `local.tee`/`local.get` pair that keeps the dispatch key for them is inside the same condition, because otherwise it costs two fuel on every entry into a peeled module. Disarmed `global_quasiimmut_invalidation` reads 5,575,086,653 against 5,575,060,318 for the same tree without the census. That fixture, N=10,193,192, reads per outer iteration in steady state: the loop entered 3x at key 2, one bridge 3x at key 0, another 1x at key 0. Its executed-op slope, differenced at N=200,000 and 400,000, is 497.0002, which those counts and the taken-path lengths account for exactly. `pyre/check.py --backend wasm --synthetic-only` is 415/415. Assisted-by: Claude --- majit/majit-backend-wasm/src/codegen.rs | 117 +++++++++++++++++- majit/majit-backend-wasm/src/lib.rs | 102 +++++++++++++++ .../majit-backend-wasm/tests/codegen_test.rs | 8 ++ pyre/pyre-wasm-runner/src/main.rs | 54 ++++++-- pyre/pyre-wasm/src/lib.rs | 16 +++ 5 files changed, 284 insertions(+), 13 deletions(-) diff --git a/majit/majit-backend-wasm/src/codegen.rs b/majit/majit-backend-wasm/src/codegen.rs index 80a1a38dd8a..d1acaa79f38 100644 --- a/majit/majit-backend-wasm/src/codegen.rs +++ b/majit/majit-backend-wasm/src/codegen.rs @@ -17,9 +17,9 @@ use majit_backend::BackendError; use majit_gc::header::{GcHeader, TYPE_ID_MASK}; use majit_ir::{InputArg, Op, OpCode, OpRef, Type}; use wasm_encoder::{ - BlockType, CodeSection, EntityType, ExportKind, ExportSection, Function, FunctionSection, - ImportSection, InstructionSink, MemArg, MemoryType, Module, RefType, TableType, TypeSection, - ValType, + BlockType, CodeSection, ConstExpr, EntityType, ExportKind, ExportSection, Function, + FunctionSection, GlobalSection, GlobalType, ImportSection, InstructionSink, MemArg, MemoryType, + Module, RefType, TableType, TypeSection, ValType, }; /// Frame slot byte offset: slot[i] is at frame_ptr + 8 + i * 8. @@ -2111,6 +2111,9 @@ pub struct ModuleBuildInputs { pub bridge_entry_arity: Option, /// Emit fixed-arity guard-to-bridge parameter tail-call arms for this module. pub bridge_param_dispatch: bool, + /// Guest-memory counters baked into an armed trace-entry census module. + /// `None` keeps the generated module byte-identical to the normal path. + pub trace_entry_census: Option, pub external_jump_slot: u32, pub external_jump_key: u32, pub frame: FrameGeometry, @@ -2178,6 +2181,7 @@ impl Clone for ModuleBuildInputs { bridge_cells_base: self.bridge_cells_base, bridge_entry_arity: self.bridge_entry_arity, bridge_param_dispatch: self.bridge_param_dispatch, + trace_entry_census: self.trace_entry_census, external_jump_slot: self.external_jump_slot, external_jump_key: self.external_jump_key, frame: self.frame, @@ -2207,6 +2211,7 @@ pub fn build_wasm_module( bridge_cells_base, bridge_entry_arity, bridge_param_dispatch, + trace_entry_census, external_jump_slot, external_jump_key, frame, @@ -2593,10 +2598,29 @@ pub fn build_wasm_module( functions.function(bridge_entry_type_idx.unwrap_or(0)); module.section(&functions); + // Only armed modules carry this global. The runner reads it after + // instantiation to give `PYRE_WASM_DUMP_ALL_TRACES` the same trace id the + // census reports; it is omitted entirely from ordinary trace modules. + if let Some(census) = trace_entry_census { + let mut globals = GlobalSection::new(); + globals.global( + GlobalType { + val_type: ValType::I64, + mutable: false, + shared: false, + }, + &ConstExpr::i64_const(census.trace_id as i64), + ); + module.section(&globals); + } + // Export section: trace function index depends on whether we imported jit_call let trace_func_idx = if needs_call { 1 } else { 0 }; let mut exports = ExportSection::new(); exports.export("trace", ExportKind::Func, trace_func_idx); + if trace_entry_census.is_some() { + exports.export("trace_entry_census_id", ExportKind::Global, 0); + } module.section(&exports); // Code section @@ -2637,6 +2661,7 @@ pub fn build_wasm_module( ca.clone(), bridge_finish_fi, ca_helper_type_idx, + *trace_entry_census, )?; codes.function(&func); module.section(&codes); @@ -2701,6 +2726,7 @@ fn build_function( // module type section when `ca.emit_ca`. The CA arm uses it to `call_indirect` // `ca.deopt_helper_slot` for a deopted callee. ca_helper_type_idx: u32, + trace_entry_census: Option, ) -> Result { // The CA arm requires residual types (the setup above forces arity >= 2 // while `WASM_DIRECT_RESIDUAL_CALL` is enabled). Its `jit_call` fallback @@ -2728,6 +2754,17 @@ fn build_function( let base_i32_locals: u32 = 1 + if ca.emit_ca { 2 } else { 0 }; let alloc_scratch_local = bridge_slot_local + base_i32_locals; let alloc_size_local = alloc_scratch_local + 1; + // A keyed census must preserve the raw dispatch value until `br_table`. + // Its counter-address scratch cannot share `bridge_slot_local`, because + // the latter would replace the selector with a guest-memory address. + let trace_entry_key_local = bridge_slot_local + + base_i32_locals + + if nursery.is_some() || ca.inline.is_some() { + 2 + } else { + 0 + }; + let trace_entry_needs_key_local = trace_entry_census.is_some() && is_resumable_peeled(ops); debug_assert_eq!(bridge_slot_local, ovf_flag_local + 1); debug_assert_eq!(ca_cfp_local, bridge_slot_local + 1); debug_assert_eq!(ca_fi_local, ca_cfp_local + 1); @@ -2777,7 +2814,8 @@ fn build_function( 2 } else { 0 - }, + } + + u32::from(trace_entry_needs_key_local), ValType::I32, )); let mut func = Function::new(locals); @@ -2876,6 +2914,24 @@ fn build_function( sink.local_get(0); sink.i64_load(mem64(frame.dispatch_key_ofs)); sink.i32_wrap_i64(); + // Without a census the key is already where `br_table` wants it. Only + // the census needs it a second time, so only the census pays to keep a + // copy: a `tee`/`get` pair here costs every entry into a peeled module. + if let Some(census) = trace_entry_census { + let dispatch_key_local = if trace_entry_needs_key_local { + trace_entry_key_local + } else { + bridge_slot_local + }; + sink.local_tee(dispatch_key_local); + emit_trace_entry_census( + &mut sink, + census, + bridge_slot_local, + Some(dispatch_key_local), + ); + sink.local_get(dispatch_key_local); + } // 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 // the default target D (the entry path). @@ -2884,6 +2940,8 @@ fn build_function( .collect(); sink.br_table(br_targets, 0); sink.end(); // end D $dispatch — key-0 entry path continues here + } else if let Some(census) = trace_entry_census { + emit_trace_entry_census(&mut sink, census, bridge_slot_local, None); } // Fresh entry owns key 0 and must clear both the trace's ordinary homes @@ -5866,6 +5924,57 @@ pub fn resumable_label_count(ops: &[Op]) -> usize { .count() } +/// Number of entry-dispatch keys an armed trace module can observe. Ordinary +/// traces have only the fresh-entry bucket; a resumable peeled loop has key 0 +/// plus one bucket for each `br_table` resume arm. +pub fn entry_dispatch_key_count(ops: &[Op]) -> usize { + if is_resumable_peeled(ops) { + resumable_label_count(ops) + 1 + } else { + 1 + } +} + +/// Increment one module's guest-memory entry counter. `key_local` holds the +/// i32 value consumed by the entry `br_table`; out-of-range values retain that +/// table's normal default-to-fresh-entry behaviour but do not index beyond the +/// fixed counter array. +fn emit_trace_entry_census( + sink: &mut PeepSink<'_, '_>, + census: crate::TraceEntryCensusStorage, + scratch_local: u32, + key_local: Option, +) { + if let Some(key_local) = key_local { + sink.local_get(key_local); + sink.i32_const(census.key_count as i32); + sink.i32_lt_u(); + sink.if_(BlockType::Empty); + sink.i32_const(census.base as i32); + sink.local_get(key_local); + sink.i32_const(std::mem::size_of::() as i32); + sink.i32_mul(); + sink.i32_add(); + sink.local_set(scratch_local); + sink.local_get(scratch_local); + sink.local_get(scratch_local); + sink.i64_load(mem64(0)); + sink.i64_const(1); + sink.i64_add(); + sink.i64_store(mem64(0)); + sink.end(); + } else { + sink.i32_const(census.base as i32); + sink.local_set(scratch_local); + sink.local_get(scratch_local); + sink.local_get(scratch_local); + sink.i64_load(mem64(0)); + sink.i64_const(1); + sink.i64_add(); + sink.i64_store(mem64(0)); + } +} + /// The single-label subset of `is_resumable_peeled`: exactly one LABEL. /// No longer consulted by the bridge accept-condition (which resolves the /// JUMP's target label by descr identity uniformly); kept as a shape diff --git a/majit/majit-backend-wasm/src/lib.rs b/majit/majit-backend-wasm/src/lib.rs index 76568fe7941..184b8c368a7 100644 --- a/majit/majit-backend-wasm/src/lib.rs +++ b/majit/majit-backend-wasm/src/lib.rs @@ -122,6 +122,102 @@ fn record_inline_trial_error(error: &BackendError) { static REEMIT_ENABLED: AtomicBool = AtomicBool::new(false); static INLINE_BRIDGE_ENABLED: AtomicBool = AtomicBool::new(false); static BRIDGE_PARAMS_ENABLED: AtomicBool = AtomicBool::new(true); +static TRACE_ENTRY_CENSUS_FORCED: AtomicBool = AtomicBool::new(false); + +/// One compiled trace's guest-memory entry counters. The generated module +/// updates `counts[key]` directly, so this owner must outlive every module +/// that bakes its base address. +struct TraceEntryCensus { + trace_id: u64, + counts: Box<[u64]>, +} + +/// The census deliberately has no per-entry Rust callback: a module writes +/// this guest-memory storage itself. The runner reads it only after Python +/// exits, when no trace is executing. +static TRACE_ENTRY_CENSUS: Mutex> = Mutex::new(Vec::new()); + +/// Baked into an armed module. `trace_id` is the backend's monotonic trace id, +/// which stays attached to a loop when its module is re-emitted. +#[derive(Clone, Copy)] +pub struct TraceEntryCensusStorage { + pub trace_id: u64, + pub base: u32, + pub key_count: u32, +} + +/// Arm trace-entry instrumentation before the guest starts compiling traces. +/// Native runs select the same facility with `MAJIT_TRACE_ENTRY_CENSUS`; wasm +/// has no environment, so its host calls this function through pyre-wasm. +pub fn trace_entry_census_enable() { + TRACE_ENTRY_CENSUS_FORCED.store(true, Ordering::Relaxed); +} + +fn trace_entry_census_enabled() -> bool { + if TRACE_ENTRY_CENSUS_FORCED.load(Ordering::Relaxed) { + return true; + } + #[cfg(not(target_arch = "wasm32"))] + { + static ON: std::sync::OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| std::env::var_os("MAJIT_TRACE_ENTRY_CENSUS").is_some()) + } + #[cfg(target_arch = "wasm32")] + { + false + } +} + +/// Allocate the one counter array that an armed physical trace module uses. +/// Re-emission clones the stored descriptor, preserving both the trace id and +/// the counters rather than assigning the replacement a second identity. +fn alloc_trace_entry_census(trace_id: u64, key_count: usize) -> Option { + if !trace_entry_census_enabled() { + return None; + } + #[cfg(target_arch = "wasm32")] + { + let mut counts = vec![0u64; key_count].into_boxed_slice(); + let base = counts.as_mut_ptr() as usize as u32; + TRACE_ENTRY_CENSUS + .lock() + .unwrap() + .push(TraceEntryCensus { trace_id, counts }); + Some(TraceEntryCensusStorage { + trace_id, + base, + key_count: key_count as u32, + }) + } + #[cfg(not(target_arch = "wasm32"))] + { + let _ = (trace_id, key_count); + None + } +} + +/// Greppable, stable host readout of the guest-written entry counters. +pub fn trace_entry_census_summary() -> String { + let census = TRACE_ENTRY_CENSUS.lock().unwrap(); + let mut total = 0u64; + let mut report = String::new(); + for trace in census.iter() { + for (key, count) in trace.counts.iter().enumerate() { + // Trace modules update this memory directly, outside Rust's alias + // analysis; volatile makes the post-run host read explicit. + let count = unsafe { core::ptr::read_volatile(count) }; + if count != 0 { + total = total.saturating_add(count); + report.push_str(&format!( + "[trace-entry-census] trace_id={} key={key} entries={count}\n", + trace.trace_id + )); + } + } + } + report.push_str(&format!("[trace-entry-census] total={total}\n")); + report +} /// Arm loop-module replacement from the host before guest execution starts. pub fn reemit_enable() { @@ -2559,6 +2655,8 @@ impl majit_backend::Backend for WasmBackend { self.collect_constants_from_ops(ops); let trace_id = self.trace_counter; self.trace_counter += 1; + let trace_entry_census = + alloc_trace_entry_census(trace_id, codegen::entry_dispatch_key_count(ops)); let typeid_table = self.collect_classptr_typeid_table(ops); let guard_gc_type_info = self.collect_guard_gc_type_info(ops); @@ -2592,6 +2690,7 @@ impl majit_backend::Backend for WasmBackend { bridge_cells_base, bridge_entry_arity: None, bridge_param_dispatch: bridge_params_enabled(), + trace_entry_census, // A real loop's JUMP is a local back-edge `br`; an entry bridge // tail-calls its target loop and is deliberately not re-emittable. external_jump_slot: entry_bridge_target.map_or(0, |t| t.func_handle), @@ -3278,6 +3377,8 @@ impl majit_backend::Backend for WasmBackend { self.collect_constants_from_ops(ops); let trace_id = self.trace_counter; self.trace_counter += 1; + let trace_entry_census = + alloc_trace_entry_census(trace_id, codegen::entry_dispatch_key_count(ops)); let typeid_table = self.collect_classptr_typeid_table(ops); let guard_gc_type_info = self.collect_guard_gc_type_info(ops); @@ -3339,6 +3440,7 @@ impl majit_backend::Backend for WasmBackend { bridge_cells_base, bridge_entry_arity, bridge_param_dispatch: bridge_params_enabled(), + trace_entry_census, external_jump_slot, external_jump_key, frame: source_frame, diff --git a/majit/majit-backend-wasm/tests/codegen_test.rs b/majit/majit-backend-wasm/tests/codegen_test.rs index 6bc2031ac67..6af265d826a 100644 --- a/majit/majit-backend-wasm/tests/codegen_test.rs +++ b/majit/majit-backend-wasm/tests/codegen_test.rs @@ -437,6 +437,7 @@ fn build_module_with_frame( bridge_cells_base: 0, bridge_entry_arity: None, bridge_param_dispatch: false, + trace_entry_census: None, external_jump_slot: 0, external_jump_key: 0, frame, @@ -642,6 +643,7 @@ fn build_module_with_write_barrier_target( bridge_cells_base: 0, bridge_entry_arity: None, bridge_param_dispatch: false, + trace_entry_census: None, external_jump_slot: 0, external_jump_key: 0, // The allocated trace keeps both Ref inputs live across New; reserve @@ -1157,6 +1159,7 @@ fn test_cold_guard_recovery_preserves_nonzero_base_and_typed_bits() { bridge_cells_base: 0, bridge_entry_arity: None, bridge_param_dispatch: false, + trace_entry_census: None, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), @@ -1264,6 +1267,7 @@ fn inlined_bridge_without_owner_loop_label_declines() { bridge_cells_base: 0, bridge_entry_arity: None, bridge_param_dispatch: false, + trace_entry_census: None, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), @@ -1547,6 +1551,7 @@ fn zero_arity_parameter_entry_is_structurally_type_zero() { bridge_cells_base: 4, bridge_entry_arity: Some(0), bridge_param_dispatch: true, + trace_entry_census: None, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), @@ -1865,6 +1870,7 @@ fn test_guard_not_invalidated_loads_runtime_flag() { bridge_cells_base: 0, bridge_entry_arity: None, bridge_param_dispatch: false, + trace_entry_census: None, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), @@ -2236,6 +2242,7 @@ fn loop_invariant_gc_table_load_stays_outside_non_collecting_loop() { bridge_cells_base: 0, bridge_entry_arity: None, bridge_param_dispatch: false, + trace_entry_census: None, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), @@ -2568,6 +2575,7 @@ fn test_non_moving_descr_allocates_through_the_oldgen_helper() { bridge_cells_base: 0, bridge_entry_arity: None, bridge_param_dispatch: false, + trace_entry_census: None, external_jump_slot: 0, external_jump_key: 0, frame: codegen::FrameGeometry::fixed(), diff --git a/pyre/pyre-wasm-runner/src/main.rs b/pyre/pyre-wasm-runner/src/main.rs index 4998ebd4e0e..b4da1a4ec60 100644 --- a/pyre/pyre-wasm-runner/src/main.rs +++ b/pyre/pyre-wasm-runner/src/main.rs @@ -127,6 +127,8 @@ struct Host { /// read back through the `pyre_fbw_census` export, and `MAJIT_GUARD_CENSUS`'s /// per-guard deopt census is `PYRE_WASM_GUARD_CENSUS`, armed through /// `pyre_jit_guard_census_enable` and read through `pyre_jit_guard_census`. +/// `PYRE_WASM_TRACE_ENTRY_CENSUS` similarly arms the emitted-module entry +/// census before tracing starts. /// /// Exempt: the names this runner interprets host-side (`PYRE_WASM_*`, /// `PYRE_STDLIB`, `MAJIT_STATS`) and `check.py`'s own `PYRE_CHECK_*` @@ -536,6 +538,12 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { { arm.call(&mut store, ())?; } + if std::env::var_os("PYRE_WASM_TRACE_ENTRY_CENSUS").is_some() + && let Ok(arm) = + instance.get_typed_func::<(), ()>(&mut store, "pyre_jit_trace_entry_census_enable") + { + arm.call(&mut store, ())?; + } if std::env::var_os("PYRE_WASM_REEMIT").is_some() && let Ok(arm) = instance.get_typed_func::<(), ()>(&mut store, "pyre_jit_reemit_enable") { @@ -550,9 +558,8 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { // Parameter bridge entries are the default. The guest has no environment, // so an explicit host-side opt-out must travel through this export before // tracing begins. - if std::env::var_os("PYRE_WASM_BRIDGE_PARAMS").is_some_and(|value| { - matches!(value.to_str().map(str::trim), Some("0" | "false" | "off")) - }) + if std::env::var_os("PYRE_WASM_BRIDGE_PARAMS") + .is_some_and(|value| matches!(value.to_str().map(str::trim), Some("0" | "false" | "off"))) && let Ok(arm) = instance.get_typed_func::<(), ()>(&mut store, "pyre_jit_bridge_params_disable") { @@ -1197,6 +1204,27 @@ fn run(module_path: &Path, source: &str, script: &Path) -> Result { Err(_) => eprintln!("[jit-stats] guard_census=unexported"), } } + if std::env::var_os("PYRE_WASM_TRACE_ENTRY_CENSUS").is_some() { + match instance.get_typed_func::<(), u64>(&mut store, "pyre_jit_trace_entry_census") { + Ok(census) => { + let census_result: Result<()> = (|| { + let packed = census.call(&mut store, ())?; + let (ptr, clen) = ((packed >> 32) as u32, (packed & 0xffff_ffff) as u32); + if clen != 0 { + let mut bytes = vec![0u8; clen as usize]; + memory.read(&store, ptr as usize, &mut bytes)?; + dealloc.call(&mut store, (ptr, clen))?; + eprint!("{}", String::from_utf8_lossy(&bytes)); + } + Ok(()) + })(); + if let Err(err) = census_result { + eprintln!("pyre-wasm-runner: trace entry census failed: {err}"); + } + } + Err(_) => eprintln!("[trace-entry-census] unexported"), + } + } let exit_code = instance .get_typed_func::<(), i32>(&mut store, "pyre_exit_code") .and_then(|f| f.call(&mut store, ())) @@ -1507,12 +1535,6 @@ fn jit_compile_trace( .context("read trace module bytes")?; let engine = caller.engine().clone(); - if std::env::var_os("PYRE_WASM_DUMP_ALL_TRACES").is_some() { - match wasmprinter::print_bytes(&bytes) { - Ok(wat) => eprintln!("=== trace module ({} bytes) ===\n{wat}", bytes.len()), - Err(pe) => eprintln!("[jit_compile_wasm] wat print failed: {pe}"), - } - } let compile_start = std::time::Instant::now(); let module_result = Module::new(&engine, &bytes); caller.data_mut().jit_compile_time_ns += compile_start.elapsed().as_nanos(); @@ -1571,6 +1593,20 @@ fn jit_compile_trace( let instance = Instance::new(&mut *caller, &module, &externs).context("instantiate trace module")?; + if std::env::var_os("PYRE_WASM_DUMP_ALL_TRACES").is_some() { + let trace_id = instance + .get_global(&mut *caller, "trace_entry_census_id") + .and_then(|global| global.get(&mut *caller).i64()) + .map(|id| format!(" trace_id={id}")) + .unwrap_or_default(); + match wasmprinter::print_bytes(&bytes) { + Ok(wat) => eprintln!( + "=== trace module{trace_id} ({} bytes) ===\n{wat}", + bytes.len() + ), + Err(pe) => eprintln!("[jit_compile_wasm] wat print failed: {pe}"), + } + } let trace = instance .get_func(&mut *caller, "trace") .context("trace module is missing its `trace` export")?; diff --git a/pyre/pyre-wasm/src/lib.rs b/pyre/pyre-wasm/src/lib.rs index c46107a0cca..a1948d8dc31 100644 --- a/pyre/pyre-wasm/src/lib.rs +++ b/pyre/pyre-wasm/src/lib.rs @@ -1149,6 +1149,15 @@ mod host_abi { majit_metainterp::guard_census_enable(); } + /// Arm the per-trace, per-entry-resume-key census before tracing starts. + /// The runner owns `PYRE_WASM_TRACE_ENTRY_CENSUS`; the wasm guest cannot + /// read its environment, and arming before compilation keeps the counter + /// instructions absent from every unarmed trace module. + #[unsafe(no_mangle)] + pub extern "C" fn pyre_jit_trace_entry_census_enable() { + majit_backend_wasm::trace_entry_census_enable(); + } + /// Arm loop-module replacement before tracing begins. The host owns the /// environment; this guest export carries that choice into the backend. #[unsafe(no_mangle)] @@ -1179,6 +1188,13 @@ mod host_abi { pack_into_guest(majit_metainterp::guard_census_summary(12).into_bytes()) } + /// Guest-memory trace-entry census readout. The runner prints every line + /// verbatim, preserving one greppable record per nonzero `(trace_id, key)`. + #[unsafe(no_mangle)] + pub extern "C" fn pyre_jit_trace_entry_census() -> u64 { + pack_into_guest(majit_backend_wasm::trace_entry_census_summary().into_bytes()) + } + /// Trial-build errors recorded while deciding whether a bridge can be /// merged into its loop module. The host surfaces this diagnostic text. #[unsafe(no_mangle)] From f41ec9b8e8af4faf415caed745ff08b5b856d917 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Mon, 17 Aug 2026 02:39:33 +0900 Subject: [PATCH 6/7] docs: add triage rows for the four gates the wasm commits introduced MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `gate_triage_complete` fails when a `PYRE_*` or `MAJIT_*` name is read from the environment with no entry in the matching triage file. `PYRE_WASM_BRIDGE_PARAMS` goes to §6a as a live default-ON brake; `PYRE_WASM_INLINE_BRIDGE` and `PYRE_WASM_REEMIT` get a §6a2 for default-OFF experiments; `PYRE_WASM_TRACE_ENTRY_CENSUS` joins the §6c diagnostics list and `MAJIT_TRACE_ENTRY_CENSUS` gets a majit entry. Assisted-by: Claude --- majit/gate-triage.md | 7 +++++++ pyre/gate-triage.md | 24 ++++++++++++++++++------ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/majit/gate-triage.md b/majit/gate-triage.md index 3f1415a6b82..6753d3d2592 100644 --- a/majit/gate-triage.md +++ b/majit/gate-triage.md @@ -371,6 +371,13 @@ Each entry records its reader, purpose, and retirement condition. `UNRECORDED` m - What it does: **UNRECORDED** — no doc comment at the read site. - Retirement condition: **UNRECORDED** — owed by this gate's owner. +### `MAJIT_TRACE_ENTRY_CENSUS` + +- Read sites: 1 — `majit/majit-backend-wasm/src/lib.rs` +- Accessor: `trace_entry_census_enabled()`; the wasm guest has no environment, so a host arms the same facility through `trace_entry_census_force()` +- What it does: Counts entries into each emitted trace module per resume key, so a steady state can be attributed to the module and dispatch key it re-enters. +- Retirement condition: Remove when the wasm trace-crossing epic closes and per-key entry counts are no longer the way that budget is attributed. + ### `MAJIT_VERIFY` - Read sites: 1 — `majit/majit-backend-cranelift/src/compiler.rs` diff --git a/pyre/gate-triage.md b/pyre/gate-triage.md index abe2717971d..a150c7688c2 100644 --- a/pyre/gate-triage.md +++ b/pyre/gate-triage.md @@ -947,15 +947,26 @@ Polarity below follows this file's rule, with one correction it needed: an **OFF**. Three diagnostics (`PYRE_DESCR_SPELLING_GATE`, `PYRE_GC_DIAG`, `PYRE_MC_DIAG`) read as ON under the unqualified rule and are OFF in fact. -### §6a — Live default-ON (4): the removal targets +### §6a — Live default-ON (5): the removal targets | gate | what is ON by default | retire when | |---|---|---| | PYRE_JD1 | the jd1 compiled-loop experiment (`eval.rs jd1_experiment_enabled`); `PYRE_NO_JD1` or `PYRE_JD1=0` turns it off, and no-JIT implies off | the jd1 experiment concludes | | PYRE_JD1_NO_ENTER | entering the compiled jd1 loop directly rather than leaving the drain to the interpreter caller | with `PYRE_JD1` | | PYRE_WALKABORT_OFF | the non-carrier walk-abort leg (`trace.rs walk_abort_leg_enabled`) | kept deliberately: the leg commits irrevocably once the blackhole runs, so it is the one-binary A/B for the bug class it sits in | +| PYRE_WASM_BRIDGE_PARAMS | a wasm guard passing its fail args to the bridge as call parameters (`lib.rs bridge_params_enabled`); `=0`/`false`/`off` restores the jitframe spill crossing | the wasm trace-crossing epic closes; until then it is the one-binary A/B for the crossing shape | | PYRE_WASM_FULL_TEARDOWN | skipping the ~0.2s wasm engine teardown at exit; setting it restores the drops for leak diagnostics | when teardown stops being the dominant fixed startup tax | +### §6a2 — Default-OFF experiments (2): the wasm re-emission A/Bs + +Both are measured to lose today and are kept as the switched-off arm of a +one-binary comparison, not as latent defaults. + +| gate | what turning it ON does | retire when | +|---|---|---| +| PYRE_WASM_INLINE_BRIDGE | merges a bridge's ops into the loop module that guards into it, so `guard → bridge → loop` becomes a `br` | the wasm trace-crossing epic closes, or the shape is measured to win | +| PYRE_WASM_REEMIT | re-emits a compiled loop's wasm module into its own table slot, which is what lets an inlined bridge reach live code | with `PYRE_WASM_INLINE_BRIDGE` | + ### §6b — VALUE knobs (12): config, not gates `PYRE_FBW_MULTIFRAME_DEPTH`, `PYRE_FBW_NO_SPECIALIZE`, `PYRE_JD1_THRESHOLD`, @@ -974,7 +985,7 @@ the folds it selects, not before them. `PYRE_FBW_SPEC_CENSUS` in §6c is its read-only half: the per-fold consulted/fired tallies. -### §6c — Default-OFF diagnostics, censuses and probes (65): keep, cost nothing +### §6c — Default-OFF diagnostics, censuses and probes (66): keep, cost nothing Each is inert unless set, so none is a removal target by this file's already-ON criterion. They are listed so they cannot be missed again. @@ -1004,7 +1015,8 @@ already-ON criterion. They are listed so they cannot be missed again. `PYRE_VSTACK_EXACT_AUDIT`, `PYRE_VSTACK_KEEP_REORDER`, `PYRE_VSTACK_NO_EXACT`, `PYRE_WASM_DUMP_BAD_TRACE`, `PYRE_WASM_EXEC_TRACE`, `PYRE_WASM_FBW_CENSUS`, `PYRE_WASM_GUARD_CENSUS`, `PYRE_WASM_JIT_STATS`, `PYRE_WASM_CALL_HIST`, -`PYRE_WASM_NO_CACHE`, `PYRE_WASM_STARTUP_TRACE`. +`PYRE_WASM_NO_CACHE`, `PYRE_WASM_STARTUP_TRACE`, +`PYRE_WASM_TRACE_ENTRY_CENSUS`. `PYRE_ALLOCSITES` enables stack attribution in the standalone `allocsites` example; it is unset by default. Its `AFTER`, `BUDGET`, `EVERY`, and `ROWS` @@ -1052,8 +1064,8 @@ input. | retired (§1 + §1b + §1c + §1d parity pass) | 5 + 4 + 17 + 1 | | not gates (identifiers) | 12 | | dead (no read site) | 10 | -| live default-ON, kept until epic closes | 6 | -| diagnostics (OFF) | ~34 | -| default-OFF experiments | 0 | +| live default-ON, kept until epic closes | 7 | +| diagnostics (OFF) | ~35 | +| default-OFF experiments | 2 | | config / value / master | ~17 | | test harness | 1 | From bdf4fe5dd1406fe5ee636c0c0fc1f63f41d2e95e Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Mon, 17 Aug 2026 02:41:19 +0900 Subject: [PATCH 7/7] majit: read a bridge's close target off the procedure token `compile_trace_inner` took the JUMP descr from the `compiled_loops` side table while `has_compiled_targets`, consulted for the same decision, answered from the procedure token. pyjitpl.py:3005-3007 reads the token once and hands that same object to `compile_trace`, so both halves now come from `warm_state.get_procedure_token`, which applies the warmstate.py:191-196 invalidation filter. `first_target_token` is the token-side accessor for the head of `target_tokens`. Measured on `synth/global_quasiimmut_invalidation`: the per-trace, per-resume-key entry census is byte-identical before and after, and `check.py --backend wasm --synthetic-only` is 415/415. Assisted-by: Claude --- majit/majit-backend/src/lib.rs | 11 +++++++++++ majit/majit-metainterp/src/pyjitpl.rs | 17 +++++++++++++---- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/majit/majit-backend/src/lib.rs b/majit/majit-backend/src/lib.rs index e91254b3b89..210e94feab5 100644 --- a/majit/majit-backend/src/lib.rs +++ b/majit/majit-backend/src/lib.rs @@ -1553,6 +1553,17 @@ impl JitCellToken { !self.target_tokens.lock().is_empty() } + /// The head of `token.target_tokens` — the descr `compile.py:290` + /// seeds the list with, and the one `pyjitpl.py:3007` closes a bridge + /// onto once `has_compiled_targets` has admitted the token. Reading it + /// from the token rather than from a side table keeps the target and the + /// gate that admitted it the same object, which is what makes the + /// `warmstate.py:191-196` invalidation filter cover both. + #[inline] + pub fn first_target_token(&self) -> Option { + self.target_tokens.lock().first().cloned() + } + /// `compile.py:286-296` / `:312-323` — append a freshly minted /// TargetToken's descr to `token.target_tokens`. Idempotent on /// `Arc::ptr_eq` so retrace paths that reuse `prior_front_target_tokens` diff --git a/majit/majit-metainterp/src/pyjitpl.rs b/majit/majit-metainterp/src/pyjitpl.rs index 886ab7a5d4d..c6d969ba67f 100644 --- a/majit/majit-metainterp/src/pyjitpl.rs +++ b/majit/majit-metainterp/src/pyjitpl.rs @@ -7663,11 +7663,20 @@ impl MetaInterp { if let Some(descr) = finish_descr { ctx.finish(finish_args, descr); } else { + // pyjitpl.py:3005-3007 reads the procedure token once, through + // the `warmstate.py:191-196` invalidation filter, and hands that + // same object to `compile_trace` — the loop a bridge closes onto + // is the loop `has_compiled_targets` admitted. `has_compiled_targets` + // here (`:7564`) already answers from the token, so taking the + // descr from the `compiled_loops` side table made one decision read + // two sources: the side table applies no invalidation filter, and + // `jitdriver.rs:6288` keeps an invalidated loop's target tokens on + // purpose. Reading both halves off the token is what makes the + // filter cover the target as well as the gate. let jump_descr = self - .compiled_loops - .get(&green_key) - .and_then(|compiled| compiled.front_target_tokens.first()) - .map(|target_token| target_token.as_jump_target_descr()); + .warm_state + .get_procedure_token(green_key) + .and_then(|token| token.first_target_token()); let Some(jump_descr) = jump_descr else { if crate::majit_log_enabled() { eprintln!(