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/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 91b405b7148..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. @@ -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. @@ -1543,7 +1560,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,9 +1910,20 @@ 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) +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 @@ -1918,7 +1946,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 +2049,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 +2079,202 @@ 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, + /// 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, + /// 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, + 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, + 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, + 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, + bridge_entry_arity, + bridge_param_dispatch, + trace_entry_census, + 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 +2302,33 @@ 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()) + 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 { - (0, None) + Vec::new() }; - let bridge_dispatch = 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 +2336,33 @@ 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 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, + &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 +2393,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 +2401,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 +2427,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 +2439,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 }; @@ -2242,40 +2462,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() @@ -2285,12 +2528,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); @@ -2309,7 +2566,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. @@ -2334,13 +2595,32 @@ 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); + // 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 @@ -2348,51 +2628,53 @@ 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, + *bridge_entry_arity, + &bridge_param_type_indices, + *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, + *trace_entry_census, )?; 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, @@ -2408,8 +2690,11 @@ 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, fail_index_base: u32, external_jump_slot: u32, // Resume-at-LABEL dispatch key the terminal external JUMP writes before @@ -2441,37 +2726,72 @@ 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 // 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; + // 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); 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, + param_type_indices: bridge_param_type_indices, + inline_guards: &inline_guards, + ref_homes, + frame, + }; let mut locals = Vec::new(); let mut start = 0; while start < value_types.types().len() { @@ -2494,7 +2814,8 @@ fn build_function( 2 } else { 0 - }, + } + + u32::from(trace_entry_needs_key_local), ValType::I32, )); let mut func = Function::new(locals); @@ -2519,6 +2840,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); @@ -2584,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). @@ -2592,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 @@ -2637,12 +2987,21 @@ 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)); - 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) { @@ -2664,6 +3023,7 @@ fn build_function( ref_homes, frame, gc_table_base, + gc_table_bases, *result, ); } @@ -2762,6 +3122,7 @@ fn build_function( ref_homes, frame, gc_table_base, + gc_table_bases, *result, ); } @@ -2770,8 +3131,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` @@ -2781,15 +3183,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 @@ -2823,6 +3237,7 @@ fn build_function( guard_idx, guard, block_exit_depth, + guard_dispatch, ); guard_idx += 1; fused_guard_at = Some(op_idx + 1); @@ -2946,12 +3361,19 @@ 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 => { - emit_guard_spill(&mut sink, constants, value_types, guard_idx, op); - sink.br(block_exit_depth); + emit_guard_exit( + &mut sink, + constants, + value_types, + guard_idx, + op, + block_exit_depth, + guard_dispatch, + ); guard_idx += 1; } @@ -2964,6 +3386,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -2975,6 +3398,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -2996,6 +3420,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3009,6 +3434,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3023,6 +3449,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3084,6 +3511,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3106,6 +3534,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3121,6 +3550,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3145,6 +3575,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); } guard_idx += 1; @@ -3172,6 +3603,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3193,6 +3625,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; // Success path: capture the caught exception into the result @@ -3228,9 +3661,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 => { @@ -3251,7 +3688,7 @@ fn build_function( value_types, op, binop, - value_types.count(), + value_types.last_local(), ovf_flag_local, fused_guard.map(|guard| guard.opcode), ) { @@ -3266,6 +3703,7 @@ fn build_function( guard_idx, guard, block_exit_depth, + guard_dispatch, ); guard_idx += 1; fused_guard_at = Some(op_idx + 1); @@ -3770,6 +4208,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); } guard_idx += 1; @@ -3845,6 +4284,7 @@ fn build_function( guard_idx, op, block_exit_depth, + guard_dispatch, ); guard_idx += 1; } @@ -3971,13 +4411,32 @@ 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.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); + sink.local_set(bridge_slot_local); + } sink.br(block_exit_depth); guard_idx += 1; } @@ -4069,8 +4528,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, @@ -5402,30 +5862,13 @@ 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. - 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(); + // 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)); sink.local_tee(bridge_slot_local); sink.if_(BlockType::Empty); @@ -5481,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 @@ -5639,6 +6133,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 @@ -5646,7 +6141,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, @@ -5833,6 +6332,27 @@ fn emit_array_addr( // ── Guard emission helpers ── +#[derive(Clone, Copy)] +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, + /// `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, +} + fn emit_guard_true( sink: &mut PeepSink<'_, '_>, constants: &indexmap::IndexMap, @@ -5840,6 +6360,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 +6371,7 @@ fn emit_guard_true( guard_idx, op, block_exit_depth, + dispatch, ); } @@ -5860,6 +6382,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 +6394,7 @@ fn emit_guard_false( guard_idx, op, block_exit_depth, + dispatch, ); } @@ -5943,7 +6467,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 +6485,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 +6495,173 @@ fn emit_guard_if_exit( guard_idx: u32, op: &Op, block_exit_depth: u32, + dispatch: BridgeDispatch<'_>, +) { + sink.if_(BlockType::Empty); + 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<'_>, +) { + 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; + } + 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); - emit_guard_spill(sink, constants, value_types, guard_idx, op); - sink.br(block_exit_depth + 1); + 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. +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<'_>, +) { + 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 +6674,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)); diff --git a/majit/majit-backend-wasm/src/failguard.rs b/majit/majit-backend-wasm/src/failguard.rs index 3d625a0cf47..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`. @@ -612,12 +617,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 @@ -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` / @@ -657,14 +667,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 +699,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..184b8c368a7 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,190 @@ 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. 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; 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. +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); +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() { + 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) +} + +/// 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. @@ -114,6 +297,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 @@ -232,6 +416,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: @@ -274,6 +459,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 @@ -293,20 +482,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 @@ -328,7 +521,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( @@ -1636,6 +1831,170 @@ 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), + guard_fail_arg_counts: exits + .iter() + .map(|guard| guard.fail_arg_refs.len()) + .collect(), + bridge_param_dispatch: inputs.bridge_param_dispatch, + }, + ); + 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, + inputs.bridge_entry_arity, + ); + 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 +2579,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 @@ -2295,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); @@ -2307,50 +2669,54 @@ 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, + 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), + 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 @@ -2453,7 +2819,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); @@ -2469,15 +2836,30 @@ 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, + 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_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 +2960,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 +2989,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,13 +3009,18 @@ 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) .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 @@ -2648,6 +3035,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, ) }) }; @@ -2656,6 +3047,7 @@ impl majit_backend::Backend for WasmBackend { source_loop.materialize_func_handle()?, source_loop.has_preamble, source_loop.frame, + is_direct, ) }; @@ -2664,7 +3056,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( @@ -2677,6 +3075,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 @@ -2837,9 +3252,133 @@ 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; + 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); @@ -2882,29 +3421,32 @@ 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, + bridge_entry_arity, + bridge_param_dispatch: bridge_params_enabled(), + trace_entry_census, + 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 @@ -2949,6 +3491,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 @@ -2961,8 +3506,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 @@ -2999,6 +3549,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 @@ -3041,8 +3596,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 @@ -3050,6 +3603,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 +3628,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..6af265d826a 100644 --- a/majit/majit-backend-wasm/tests/codegen_test.rs +++ b/majit/majit-backend-wasm/tests/codegen_test.rs @@ -420,25 +420,31 @@ 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, + bridge_entry_arity: None, + bridge_param_dispatch: false, + trace_entry_census: None, + 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 +626,32 @@ 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, + 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 // 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 +1142,31 @@ 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, + bridge_entry_arity: None, + bridge_param_dispatch: false, + trace_entry_census: None, + 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 +1235,52 @@ 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, + bridge_entry_arity: None, + bridge_param_dispatch: false, + trace_entry_census: None, + 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) @@ -1451,6 +1514,59 @@ 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, + trace_entry_census: None, + 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)]; @@ -1737,25 +1853,31 @@ 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, + bridge_entry_arity: None, + bridge_param_dispatch: false, + trace_entry_census: None, + 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 +2225,30 @@ 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, + bridge_entry_arity: None, + bridge_param_dispatch: false, + trace_entry_census: None, + 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 +2553,36 @@ 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, + bridge_entry_arity: None, + bridge_param_dispatch: false, + trace_entry_census: None, + 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/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!( 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 | diff --git a/pyre/pyre-wasm-runner/src/main.rs b/pyre/pyre-wasm-runner/src/main.rs index 1c2907f144b..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,33 @@ 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") + { + 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, ())?; + } + // 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; @@ -688,6 +717,24 @@ 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", + "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() { @@ -696,6 +743,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 +1140,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 @@ -1128,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, ())) @@ -1252,6 +1349,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 +1517,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() @@ -1420,12 +1535,6 @@ fn jit_compile(caller: &mut Caller<'_, Host>, bytes_ptr: u32, bytes_len: u32) -> .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(); @@ -1484,20 +1593,68 @@ fn jit_compile(caller: &mut Caller<'_, Host>, bytes_ptr: u32, bytes_len: u32) -> 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")?; + 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..a1948d8dc31 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,37 @@ 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)] + 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(); + } + + /// 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. @@ -1149,6 +1188,20 @@ 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)] + 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