diff --git a/majit/majit-backend-wasm/src/codegen.rs b/majit/majit-backend-wasm/src/codegen.rs index 5a5ad1643e9..bf249e32204 100644 --- a/majit/majit-backend-wasm/src/codegen.rs +++ b/majit/majit-backend-wasm/src/codegen.rs @@ -792,8 +792,9 @@ impl RefHomes { ops: &[Op], include_ca_collects: bool, forced_refs: &[OpRef], + regions: &[InlinedRegionSpan], ) -> Self { - let liveness = HomeLiveness::collect(inputargs, ops); + let liveness = HomeLiveness::collect_with_regions(inputargs, ops, regions); let collect_positions = collecting_call_positions(ops, include_ca_collects); let ref_values = RefValues::collect(inputargs, ops); let mut by_id = Vec::new(); @@ -895,8 +896,45 @@ struct LabelResumeData { ref_slots: usize, } +/// Where one inlined bridge region starts in a merged analysis stream, and +/// which value ids carry that region's own live-ins. +struct InlinedRegionSpan { + ops_start: usize, + inputarg_ids: Vec, +} + +impl InlinedRegionSpan { + /// The regions occupy the tail of the merged stream in `inlined_bridges` + /// order, so their starts run back from the end of `ops`. `bridges` must be + /// the rebased copies the merged stream was built from, so the recorded ids + /// are the ids that stream reads. + fn collect(ops_len: usize, bridges: &[InlinedBridge]) -> Vec { + let mut start = + ops_len.saturating_sub(bridges.iter().map(|bridge| bridge.ops.len()).sum::()); + bridges + .iter() + .map(|bridge| { + let span = Self { + ops_start: start, + inputarg_ids: bridge.inputargs.iter().map(|ia| ia.index).collect(), + }; + start += bridge.ops.len(); + span + }) + .collect() + } +} + impl LabelResumeData { fn collect(inputargs: &[InputArg], ops: &[Op]) -> Self { + Self::collect_with_regions(inputargs, ops, &[]) + } + + fn collect_with_regions( + inputargs: &[InputArg], + ops: &[Op], + regions: &[InlinedRegionSpan], + ) -> Self { let (_, num_vars) = collect_guards_and_vars(inputargs, ops); let ref_values = RefValues::collect(inputargs, ops); let normal_value_slots = normal_frame_value_slots(inputargs, ops); @@ -963,6 +1001,24 @@ impl LabelResumeData { *v = true; } } + // An appended region's live-ins reach it only through the + // guard-fail branch that is the region's sole predecessor, and that + // branch assigns them. Nothing the entry dispatch can land on + // reaches a region's first read without passing it, so those ids + // are dead until written here. Treating them as live would reserve + // one frozen-frame slot per region live-in at every resumable + // label, and the resume loader would reload a value the guard + // overwrites before anything reads it. + for region in regions { + if region.ops_start <= label_pos { + continue; + } + for &id in ®ion.inputarg_ids { + if let Some(v) = available.get_mut(id as usize) { + *v = true; + } + } + } let mut missing = Vec::new(); let mut bad = false; @@ -1087,7 +1143,7 @@ pub fn count_ref_homes(inputargs: &[InputArg], ops: &[Op]) -> usize { // This pre-sizing query is used for CA bridges before `CaParams` exists, so // count CALL_ASSEMBLER as a collecting position to match CA codegen. let resume = LabelResumeData::collect(inputargs, ops); - RefHomes::collect(inputargs, ops, true, &resume.captured_refs).len() + RefHomes::collect(inputargs, ops, true, &resume.captured_refs, &[]).len() } /// Number of high GC-rooted homes reserved exclusively for LABEL live-ins. @@ -1336,7 +1392,11 @@ struct HomeLiveness { } impl HomeLiveness { - fn collect(inputargs: &[InputArg], ops: &[Op]) -> Self { + fn collect_with_regions( + inputargs: &[InputArg], + ops: &[Op], + regions: &[InlinedRegionSpan], + ) -> Self { let mut n = inputargs .iter() .map(|ia| ia.index as usize + 1) @@ -1379,6 +1439,22 @@ impl HomeLiveness { } } } + // An appended region's live-ins are written by the guard-fail branch + // that is the region's sole predecessor, and that branch jumps straight + // into the region. Their entry in the merged input-arg list would + // otherwise date them to trace entry, making them live across every + // collecting call in the owner's body: each would take a Ref home and + // be reloaded there on every iteration, for a value nothing in the + // owner reads. Date them to the region instead, so they stay live + // across the region's own collect points and nowhere else. + for region in regions { + let defined_at = region.ops_start as i32 - 1; + for &id in ®ion.inputarg_ids { + if let Some(d) = def_pos.get_mut(id as usize) { + *d = defined_at; + } + } + } Self { def_pos, last_use } } @@ -2151,9 +2227,9 @@ pub struct ModuleBuildInputs { /// 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. + /// Each is retained in its own trace's numbering; `build_wasm_module` + /// rebases them onto a private id range before merging, because the owner + /// and every region number their values independently from zero. pub inlined_bridges: Vec, pub constants: indexmap::IndexMap, pub vtable_offset: Option, @@ -2190,6 +2266,11 @@ pub struct InlinedBridge { /// 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, + /// The constant pool registered for this region's own trace. A pool is + /// per-trace (`Backend::set_constants_pool` names the next compile), and + /// its value-id keys — the folded values that have no producing op — are + /// in that trace's numbering, so the merge rebases them with the region. + pub constants: indexmap::IndexMap, } /// Whether the exact operation stream emitted for `inputs` has a local loop @@ -2215,6 +2296,7 @@ impl Clone for InlinedBridge { .collect(), ops: self.ops.clone(), gc_table_base: self.gc_table_base, + constants: self.constants.clone(), } } } @@ -2251,6 +2333,109 @@ impl Clone for ModuleBuildInputs { } } +/// One past the highest value id `inputargs`/`ops` define or read. Mirrors the +/// `max_var` half of `collect_guards_and_vars` without its guard collection, +/// which stamps per-value counters onto guard descrs and must run once only. +fn value_id_end(inputargs: &[InputArg], ops: &[Op]) -> u32 { + let mut end: u32 = 0; + let widen = |r: OpRef, end: &mut u32| { + if r != OpRef::NONE && !r.is_constant() && r.raw() + 1 > *end { + *end = r.raw() + 1; + } + }; + for ia in inputargs { + if ia.index + 1 > end { + end = ia.index + 1; + } + } + for op in ops { + widen(op.pos.get(), &mut end); + for a in op.getarglist().iter() { + widen(a.to_opref(), &mut end); + } + if let Some(fa) = op.getfailargs() { + for a in fa.iter() { + widen(a.to_opref(), &mut end); + } + } + } + end +} + +/// Move every value id a region defines or reads up by `offset`, returning the +/// rebased region and the width of the id range it now occupies. +/// +/// The owner trace and each region are separately recorded traces, so both +/// number their values from zero and their ids overlap. A region is entered by +/// `local.set`ting the id each of its input args carries +/// (`emit_guard_inline_bridge_move`) and leaves through the loop header, so an +/// id it shares with an owner value that is live across the back edge +/// overwrites that value for every following iteration. Rebasing onto a +/// disjoint range is what makes the merged stream's single local namespace +/// sound. +/// +/// `TempVar` ids live in a reserved high strip and constants in their own +/// namespace; neither indexes a value local, so both pass through unchanged. +fn rebase_region_value_ids(bridge: &InlinedBridge, offset: u32) -> (InlinedBridge, u32) { + use majit_ir::operand::Operand; + + let shift = |r: OpRef| -> OpRef { + if r.is_none() || r.is_constant() || r.is_temp_var() { + r + } else { + r.with_raw(r.raw() + offset) + } + }; + + let width = value_id_end(&bridge.inputargs, &bridge.ops); + let inputargs: Vec = bridge + .inputargs + .iter() + .map(|ia| InputArg::from_type(ia.tp, ia.index + offset)) + .collect(); + // `Op::clone` gives the copy its own arg/failarg slots, but the operands in + // them keep pointing at the region's original producers, whose `pos` this + // must not touch — the region is retained for the next re-emission. So each + // moved reference is rebound to a synthetic producer carrying the new id. + let ops: Vec = bridge.ops.to_vec(); + for op in &ops { + op.pos.set(shift(op.pos.get())); + for (i, arg) in op.getarglist().iter().enumerate() { + let before = arg.to_opref(); + let after = shift(before); + if after != before { + op.setarg(i, Operand::bound_from_opref(after)); + } + } + if let Some(mut fail_args) = op.getfailargs() { + let mut moved = false; + for slot in fail_args.iter_mut() { + let before = slot.to_opref(); + let after = shift(before); + if after != before { + *slot = Operand::bound_from_opref(after); + moved = true; + } + } + if moved { + op.setfailargs(fail_args); + } + } + } + + ( + InlinedBridge { + source_fail_index: bridge.source_fail_index, + trace_id: bridge.trace_id, + inputargs, + ops, + gc_table_base: bridge.gc_table_base, + constants: bridge.constants.clone(), + }, + width, + ) +} + /// Build a wasm module from majit IR. pub fn build_wasm_module( inputs: &ModuleBuildInputs, @@ -2289,12 +2474,37 @@ pub fn build_wasm_module( let mut merged_inputargs = Vec::new(); let mut merged_ops = Vec::new(); let mut gc_table_bases = HashMap::new(); + let mut rebased_bridges: Vec = Vec::new(); + let mut rebased_constants = indexmap::IndexMap::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()); + // The merged stream has one local namespace, so every region has to be + // moved off the ids the owner and the earlier regions already use. + rebased_constants = constants.clone(); + let mut next_value_id = value_id_end(inputargs, ops); for bridge in inlined_bridges { + let (bridge, width) = rebase_region_value_ids(bridge, next_value_id); + // The pool is keyed by value position for a folded value with no + // producing op, so rebasing the region's ids moved its reads off + // its own entries. Replay that window at the offset, and drop a + // key another trace left inside it, or `unbound_pool_const_seeds` + // either declines a resolvable value or seeds an unrelated one's + // bits. Keys outside the window are left alone: rewriting them + // would overwrite the entries the owner's own operations read. + for id in 0..width { + match bridge.constants.get(&id) { + Some(&bits) => { + rebased_constants.insert(id + next_value_id, bits); + } + None => { + rebased_constants.shift_remove(&(id + next_value_id)); + } + } + } + next_value_id += width; merged_inputargs.extend(bridge.inputargs.iter().map(InputArg::fresh_value_copy)); for op in &bridge.ops { if op.opcode == OpCode::LoadFromGcTable { @@ -2302,9 +2512,23 @@ pub fn build_wasm_module( } } merged_ops.extend(bridge.ops.iter().cloned()); + rebased_bridges.push(bridge); } (&merged_inputargs, &merged_ops) }; + // Guard-entry moves and region emission must name the rebased ids, not the + // ids the retained regions still carry. + let emitted_bridges: &[InlinedBridge] = if inlined_bridges.is_empty() { + inlined_bridges + } else { + &rebased_bridges + }; + let region_spans = InlinedRegionSpan::collect(analysis_ops.len(), emitted_bridges); + let constants = if inlined_bridges.is_empty() { + constants + } else { + &rebased_constants + }; 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 @@ -2383,7 +2607,8 @@ pub fn build_wasm_module( // 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(&analysis_inputargs, &analysis_ops); + let label_resume = + LabelResumeData::collect_with_regions(&analysis_inputargs, &analysis_ops, ®ion_spans); let max_value_slots = normal_frame_value_slots(&analysis_inputargs, &analysis_ops) + label_resume.scalar_slots; if max_value_slots > frame.value_slots { @@ -2422,6 +2647,7 @@ pub fn build_wasm_module( &analysis_ops, ca.emit_ca, &label_resume.captured_refs, + ®ion_spans, ); let num_ref_homes = ref_homes.len(); let shortage = if num_ref_homes > frame.ordinary_home_slots() { @@ -2716,7 +2942,7 @@ pub fn build_wasm_module( inputargs, &analysis_inputargs, &analysis_ops, - inlined_bridges, + emitted_bridges, constants, num_vars, &value_types, @@ -2936,8 +3162,11 @@ fn build_function( ) })?; - // Def / last-use positions for the post-collection Ref reload filter. - let liveness = HomeLiveness::collect(inputargs, ops); + // Def / last-use positions for the post-collection Ref reload filter. The + // spans must match the ones `RefHomes` was built from, or a home would be + // reserved and never reloaded (or the reverse). + let region_spans = InlinedRegionSpan::collect(ops.len(), inlined_bridges); + let liveness = HomeLiveness::collect_with_regions(inputargs, ops, ®ion_spans); // `LOAD_FROM_GC_TABLE` is the backend form of a ConstPtr. Native PyPy // keeps such loop-invariant references in their allocated location across @@ -3039,8 +3268,13 @@ fn build_function( // immediately and nothing between the two allocates, so no collection can // read the slot while it is stale. Homes no input fills keep their clear // because store-on-def writes them only later. + // The loop below fills `entry_inputargs`, not every arg of the merged + // stream: an appended region's live-ins are stored by the guard-fail branch + // that reaches the region, which is nowhere near this entry. Marking those + // homes filled here would skip their clear and leave the collector reading + // an uninitialised slot. let mut input_filled_home = vec![false; ref_homes.len()]; - for ia in inputargs { + for ia in entry_inputargs { if let Some(h) = ref_homes.home_id(ia.index) { input_filled_home[h as usize] = true; } diff --git a/majit/majit-backend-wasm/src/failguard.rs b/majit/majit-backend-wasm/src/failguard.rs index 8a5099699dd..6208905d378 100644 --- a/majit/majit-backend-wasm/src/failguard.rs +++ b/majit/majit-backend-wasm/src/failguard.rs @@ -671,9 +671,19 @@ pub struct CompiledWasmLoop { /// 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. + /// Direct-loop guard index to bridge table slot. `patch_jump_for_descr` + /// rewrites the guard's own jump to reach a newly attached bridge; a wasm + /// module is immutable once compiled, so the branch instead reads a slot + /// out of a mutable cell array, and these are the writes a re-emission has + /// to replay into its fresh array. pub bridge_slots: RefCell>, + /// The same, for a guard that lives inside a trace chained onto this loop, + /// keyed by `(owning trace_id, per-trace fail index)`. A standalone chained + /// bridge keeps its cells in its own module's array, which survives; a + /// region merged into this loop does not, because a re-emission reallocates + /// the loop array its guards are carved out of. Replayed once the rebuilt + /// `chained_trace_meta` names the new bases. + pub chained_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>, diff --git a/majit/majit-backend-wasm/src/lib.rs b/majit/majit-backend-wasm/src/lib.rs index 7bdc2969ff2..60a09a7559c 100644 --- a/majit/majit-backend-wasm/src/lib.rs +++ b/majit/majit-backend-wasm/src/lib.rs @@ -1987,6 +1987,24 @@ impl WasmBackend { for region in &inputs.inlined_bridges { let count = codegen::guard_exit_count(®ion.inputargs, ®ion.ops); let exits = &guard_exits[offset..offset + count]; + // This region's guards are carved out of the array that was + // just reallocated, so every bridge already chained onto one of + // them has lost its dispatch entry. Unreplayed, that guard + // deopts to the tracer on every failure and retraces a bridge + // it can never reach. + #[cfg(all(target_arch = "wasm32", not(target_os = "wasi")))] + if new_cells_base != 0 { + for (&(trace_id, fail_index), &bridge_slot) in + compiled.chained_bridge_slots.borrow().iter() + { + if trace_id != region.trace_id || fail_index as usize >= count { + continue; + } + let cell = (new_cells_base as usize + (offset + fail_index as usize) * 4) + as *mut u32; + unsafe { core::ptr::write(cell, bridge_slot) }; + } + } metas.insert( region.trace_id, ChainedTraceMeta { @@ -2889,6 +2907,7 @@ impl majit_backend::Backend for WasmBackend { chained_trace_meta: std::cell::RefCell::new(std::collections::HashMap::new()), _bridge_owned_cells: std::cell::RefCell::new(bridge_cells_owner.into_iter().collect()), bridge_slots: std::cell::RefCell::new(HashMap::new()), + chained_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. @@ -3320,6 +3339,7 @@ impl majit_backend::Backend for WasmBackend { inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), ops: ops_owned.clone(), gc_table_base, + constants: self.constants.clone(), }); let mut merged_ops = candidate.ops.clone(); for region in &candidate.inlined_bridges { @@ -3646,16 +3666,23 @@ impl majit_backend::Backend for WasmBackend { } // Retained module replacement and loop-closing bridge inlining // restore this cell after allocating a fresh dispatch array. - if is_direct && (reemit_enabled() || inline_bridge_enabled()) { + if reemit_enabled() || inline_bridge_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); + if is_direct { + source_loop + .bridge_slots + .borrow_mut() + .insert(source_fail_index, bridge_slot); + } else { + source_loop + .chained_bridge_slots + .borrow_mut() + .insert((source_trace_id, source_fail_index), bridge_slot); + } } } } diff --git a/majit/majit-backend-wasm/tests/codegen_test.rs b/majit/majit-backend-wasm/tests/codegen_test.rs index ef0b76395ad..9e7d8842011 100644 --- a/majit/majit-backend-wasm/tests/codegen_test.rs +++ b/majit/majit-backend-wasm/tests/codegen_test.rs @@ -1276,6 +1276,7 @@ fn inlined_bridge_without_owner_loop_label_declines() { inputargs: vec![InputArg::from_type(Type::Int, 1)], ops: vec![Op::new(OpCode::Finish, &[])], gc_table_base: 0, + constants: indexmap::IndexMap::new(), }], constants: indexmap::IndexMap::new(), vtable_offset: Some(0), @@ -1304,6 +1305,127 @@ fn inlined_bridge_without_owner_loop_label_declines() { assert!(error.to_string().contains("no local loop LABEL")); } +/// A region's value ids are its own trace's, so they collide with the owner's. +/// The merged stream has one local namespace, so an unrebased collision makes +/// the region's entry moves land in locals the owner still holds live across +/// the back edge. Where a region's numbering happens to start must therefore +/// not be observable in the emitted code. +#[test] +fn inlined_bridge_emission_is_independent_of_the_regions_own_numbering() { + fn owner_ops() -> Vec { + vec![ + // Defined before the LABEL and read after it, so it is live across + // the back edge and is restored only on preamble/resume entry. + make_op( + OpCode::IntAdd, + &[OpRef::input_arg_int(0), OpRef::input_arg_int(1)], + OpRef::int_op(2), + ), + Op::new( + OpCode::Label, + &[rb(OpRef::input_arg_int(0)), rb(OpRef::input_arg_int(1))], + ), + make_op( + OpCode::IntAdd, + &[OpRef::input_arg_int(0), OpRef::const_int(1)], + OpRef::int_op(3), + ), + make_op( + OpCode::IntLt, + &[OpRef::int_op(3), OpRef::const_int(10)], + OpRef::int_op(4), + ), + make_guard( + OpCode::GuardTrue, + &[OpRef::int_op(4)], + &[OpRef::int_op(3), OpRef::input_arg_int(1)], + ), + make_op( + OpCode::IntAdd, + &[OpRef::input_arg_int(1), OpRef::int_op(2)], + OpRef::int_op(5), + ), + make_op( + OpCode::IntLt, + &[OpRef::int_op(5), OpRef::const_int(1000)], + OpRef::int_op(6), + ), + make_guard( + OpCode::GuardTrue, + &[OpRef::int_op(6)], + &[OpRef::int_op(3), OpRef::int_op(5)], + ), + Op::new(OpCode::Jump, &[rb(OpRef::int_op(3)), rb(OpRef::int_op(5))]), + ] + } + + // `base` picks where the region numbers its own values. `base = 2` makes + // its first input arg share an id with the owner's loop-invariant + // `int_op(2)`; `base = 40` clears every owner id. + fn build(base: u32) -> Vec { + let inputargs = vec![ + InputArg::from_type(Type::Int, 0), + InputArg::from_type(Type::Int, 1), + ]; + let region_ops = vec![ + make_op( + OpCode::IntAdd, + &[OpRef::input_arg_int(base), OpRef::input_arg_int(base + 1)], + OpRef::int_op(base + 2), + ), + Op::new( + OpCode::Jump, + &[ + rb(OpRef::int_op(base + 2)), + rb(OpRef::input_arg_int(base + 1)), + ], + ), + ]; + let inputs = codegen::ModuleBuildInputs { + inputargs: inputargs.iter().map(InputArg::fresh_value_copy).collect(), + ops: owner_ops(), + inlined_bridges: vec![codegen::InlinedBridge { + source_fail_index: 1, + trace_id: 7, + inputargs: vec![ + InputArg::from_type(Type::Int, base), + InputArg::from_type(Type::Int, base + 1), + ], + ops: region_ops, + gc_table_base: 0, + constants: indexmap::IndexMap::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: 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(), + }; + codegen::build_wasm_module(&inputs) + .expect("a loop-closing region merges into its owner") + .0 + } + + let colliding = build(2); + let disjoint = build(40); + validate_wasm(&colliding); + validate_wasm(&disjoint); + assert_eq!(colliding, disjoint); +} + #[test] fn test_int_add_loop() { // Label(i, sum) -> IntAdd(sum, i) -> IntAdd(i, 1) -> IntLt(i, 100)