From 6fd2e25c9e90cbe6da2c8e5a28469b97f01c5a17 Mon Sep 17 00:00:00 2001 From: sbillig Date: Mon, 6 Jul 2026 22:15:58 -0700 Subject: [PATCH 01/14] Slim the stackalloc Allocator contract Delete the never-called traverse_edge; rename read/write/read_br_table_case to pre_inst/post_inst/br_table_case and drop the ignored operand params; take &dyn Allocator instead of &mut dyn; rewrite FinalAlloc action lists once at construction so lookups return borrows; factor the spill-storage store-action mapping into SpillStorage::store_action. --- crates/codegen/src/isa/evm/emit/alloc.rs | 54 ++++++++------ crates/codegen/src/isa/evm/emit/insn.rs | 64 ++++++++-------- crates/codegen/src/isa/evm/emit/layout.rs | 4 +- crates/codegen/src/isa/evm/emit/mod.rs | 13 +--- .../codegen/src/isa/evm/machine/lazy_frame.rs | 34 ++++----- crates/codegen/src/isa/evm/tests.rs | 37 ++++------ crates/codegen/src/machinst/lower.rs | 8 +- crates/codegen/src/stackalloc/mod.rs | 19 ++--- .../codegen/src/stackalloc/stackify/alloc.rs | 73 +++++++++++++------ crates/codegen/src/stackalloc/stackify/mod.rs | 4 +- .../src/stackalloc/stackify/planner/mod.rs | 16 +++- 11 files changed, 175 insertions(+), 151 deletions(-) diff --git a/crates/codegen/src/isa/evm/emit/alloc.rs b/crates/codegen/src/isa/evm/emit/alloc.rs index 8f3e2a130..f9238b977 100644 --- a/crates/codegen/src/isa/evm/emit/alloc.rs +++ b/crates/codegen/src/isa/evm/emit/alloc.rs @@ -1,4 +1,4 @@ -use sonatina_ir::{BlockId, Function, I256, Immediate, InstId}; +use sonatina_ir::{Function, I256, Immediate, InstId}; use crate::stackalloc::{Action, Actions, Allocator, StackifyAlloc}; @@ -11,11 +11,34 @@ pub(crate) struct FinalAlloc { impl FinalAlloc { pub(crate) fn new(inner: StackifyAlloc, mem_plan: MachineFuncPlan) -> Self { - let alloc = Self { inner, mem_plan }; + let mut alloc = Self { inner, mem_plan }; + // Validation inspects the virtual object ids that the rewrite consumes, so it must run + // before rewriting. alloc.validate_object_actions(); + alloc.rewrite_inner_actions(); alloc } + /// Rewrite every stored action list once, at construction, into its final form (virtual + /// object/local addresses resolved, per-call save/restore injected). The rewrite depends only + /// on `mem_plan`, so doing it eagerly lets the `Allocator` methods return borrows and avoids + /// re-rewriting the same lists on every lookup. + fn rewrite_inner_actions(&mut self) { + let mut inner = std::mem::take(&mut self.inner); + // `rewrite_action_lists` visits allocated entries only; make sure every call with a + // preserve plan is visited so save/restore injection (and its invariant checks) applies + // even if stackify left the call's action lists untouched. + for &inst in self.mem_plan.call_preserve.keys() { + inner.touch_inst_actions(inst); + } + inner.rewrite_action_lists( + |inst, actions| self.rewrite_actions(self.inject_call_save_pre(inst, actions)), + |inst, actions| self.rewrite_actions(self.inject_call_save_post(inst, actions)), + |actions| self.rewrite_actions(actions), + ); + self.inner = inner; + } + fn abs_addr_for_word(&self, word_off: u32) -> u32 { self.mem_plan.abs_addr_for_word(word_off) } @@ -149,12 +172,7 @@ impl FinalAlloc { } } - fn inject_call_save_pre( - &self, - inst: InstId, - _operand_count: usize, - actions: Actions, - ) -> Actions { + fn inject_call_save_pre(&self, inst: InstId, actions: Actions) -> Actions { let Some(plan) = self.mem_plan.call_preserve.get(&inst) else { return actions; }; @@ -259,24 +277,16 @@ impl Allocator for FinalAlloc { self.rewrite_actions(self.inner.enter_function(function)) } - fn read(&self, inst: InstId, vals: &[sonatina_ir::ValueId]) -> Actions { - let actions = self.inner.read(inst, vals); - let actions = self.inject_call_save_pre(inst, vals.len(), actions); - self.rewrite_actions(actions) - } - - fn read_br_table_case(&self, inst: InstId, case_index: usize) -> Actions { - self.rewrite_actions(self.inner.read_br_table_case(inst, case_index)) + fn pre_inst(&self, inst: InstId) -> &Actions { + self.inner.pre_inst(inst) } - fn write(&self, inst: InstId, vals: &[sonatina_ir::ValueId]) -> Actions { - let actions = self.inner.write(inst, vals); - let actions = self.inject_call_save_post(inst, actions); - self.rewrite_actions(actions) + fn br_table_case(&self, inst: InstId, case_index: usize) -> &Actions { + self.inner.br_table_case(inst, case_index) } - fn traverse_edge(&self, from: BlockId, to: BlockId) -> Actions { - self.rewrite_actions(self.inner.traverse_edge(from, to)) + fn post_inst(&self, inst: InstId) -> &Actions { + self.inner.post_inst(inst) } } diff --git a/crates/codegen/src/isa/evm/emit/insn.rs b/crates/codegen/src/isa/evm/emit/insn.rs index 9207fb18e..cba320739 100644 --- a/crates/codegen/src/isa/evm/emit/insn.rs +++ b/crates/codegen/src/isa/evm/emit/insn.rs @@ -1,7 +1,7 @@ use rustc_hash::FxHashSet; use smallvec::SmallVec; use sonatina_ir::{ - InstId, InstSetExt, U256, ValueId, + InstId, InstSetExt, U256, inst::evm::machine_inst_set::EvmMachineInstKind, isa::{Isa, evm::EvmMachine}, }; @@ -25,12 +25,7 @@ use super::{ }; impl EvmMachineFunctionLowering<'_> { - pub(crate) fn lower_insn( - &self, - ctx: &mut Lower, - alloc: &mut dyn Allocator, - insn: InstId, - ) { + pub(crate) fn lower_insn(&self, ctx: &mut Lower, alloc: &dyn Allocator, insn: InstId) { if self.is_elided_block(ctx.insn_block(insn)) { return; } @@ -42,17 +37,16 @@ impl EvmMachineFunctionLowering<'_> { let emit_post_actions = |ctx: &mut Lower, actions: &[Action]| { self.emit_actions_for_site(ctx, actions, frame_layout, FrameSite::PostInst(insn)) }; - let results: SmallVec<[ValueId; 4]> = ctx.insn_results(insn).iter().copied().collect(); let args = ctx.insn_data(insn).collect_values(); let machine = EvmMachine::new(ctx.module.triple); let data = machine.inst_set().resolve_inst(ctx.insn_data(insn)); let basic_op = |ctx: &mut Lower, ops: &[OpCode]| { - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); for op in ops { ctx.push(*op); } - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); }; match &data { @@ -74,7 +68,7 @@ impl EvmMachineFunctionLowering<'_> { EvmMachineInstKind::Xor(_) => basic_op(ctx, &[OpCode::XOR]), EvmMachineInstKind::Jump(jump) => { let dest = self.canonical_block_target(*jump.dest()); - emit_pre_actions(ctx, &alloc.read(insn, &[])); + emit_pre_actions(ctx, alloc.pre_inst(insn)); if !ctx.is_next_block(dest) { let push_op = ctx.push(OpCode::PUSH1); @@ -86,7 +80,7 @@ impl EvmMachineFunctionLowering<'_> { let nz_dest = self.canonical_block_target(*br.nz_dest()); let z_dest = self.canonical_block_target(*br.z_dest()); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); if nz_dest == z_dest { ctx.push(OpCode::POP); if !ctx.is_next_block(nz_dest) { @@ -109,7 +103,7 @@ impl EvmMachineFunctionLowering<'_> { } EvmMachineInstKind::Phi(_) => {} EvmMachineInstKind::Unreachable(_) => { - emit_pre_actions(ctx, &alloc.read(insn, &[])); + emit_pre_actions(ctx, alloc.pre_inst(insn)); ctx.push(OpCode::INVALID); } EvmMachineInstKind::BrTable(br) => { @@ -131,7 +125,7 @@ impl EvmMachineFunctionLowering<'_> { let dest = self.canonical_block_target(*dest); self.emit_actions_for_site( ctx, - &alloc.read_br_table_case(insn, case_idx), + alloc.br_table_case(insn, case_idx), frame_layout, FrameSite::PreInst(insn), ); @@ -149,7 +143,7 @@ impl EvmMachineFunctionLowering<'_> { } EvmMachineInstKind::Call(call) => { let callee = *call.callee(); - let mut actions = alloc.read(insn, &args); + let mut actions = alloc.pre_inst(insn).clone(); let cont_pos = actions .iter() @@ -197,7 +191,7 @@ impl EvmMachineFunctionLowering<'_> { let jumpdest_op = ctx.push(OpCode::JUMPDEST); ctx.add_label_reference(push_callback, Label::Insn(jumpdest_op)); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } else { self.emit_actions_for_site( ctx, @@ -220,7 +214,7 @@ impl EvmMachineFunctionLowering<'_> { } } EvmMachineInstKind::Return(_) => { - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); if !self.has_lazy_frame_lowering() && let Some(frame_layout) = frame_layout { @@ -243,14 +237,14 @@ impl EvmMachineFunctionLowering<'_> { bits < 256, "full-width saturating add must be legalized earlier" ); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); emit_narrow_unsigned_saturating_binary( ctx, OpCode::ADD, bits, low_bits_mask(bits).unwrap(), ); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } EvmMachineInstKind::EvmSaddsat(sat) => { let bits = scalar_bit_width(*sat.ty(), ctx.module) @@ -259,9 +253,9 @@ impl EvmMachineFunctionLowering<'_> { bits < 256, "full-width saturating add must be legalized earlier" ); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); emit_narrow_signed_saturating_binary(ctx, OpCode::ADD, bits); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } EvmMachineInstKind::EvmUsubsat(sat) => { let bits = scalar_bit_width(*sat.ty(), ctx.module) @@ -270,9 +264,9 @@ impl EvmMachineFunctionLowering<'_> { bits < 256, "full-width saturating sub must be legalized earlier" ); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); emit_narrow_unsigned_saturating_binary(ctx, OpCode::SUB, bits, U256::zero()); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } EvmMachineInstKind::EvmSsubsat(sat) => { let bits = scalar_bit_width(*sat.ty(), ctx.module) @@ -281,9 +275,9 @@ impl EvmMachineFunctionLowering<'_> { bits < 256, "full-width saturating sub must be legalized earlier" ); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); emit_narrow_signed_saturating_binary(ctx, OpCode::SUB, bits); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } EvmMachineInstKind::EvmUmulsat(sat) => { let bits = scalar_bit_width(*sat.ty(), ctx.module) @@ -292,14 +286,14 @@ impl EvmMachineFunctionLowering<'_> { bits < 256, "full-width saturating mul must be legalized earlier" ); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); emit_narrow_unsigned_saturating_binary( ctx, OpCode::MUL, bits, low_bits_mask(bits).unwrap(), ); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } EvmMachineInstKind::EvmSmulsat(sat) => { let bits = scalar_bit_width(*sat.ty(), ctx.module) @@ -308,9 +302,9 @@ impl EvmMachineFunctionLowering<'_> { bits < 256, "full-width saturating mul must be legalized earlier" ); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); emit_narrow_signed_saturating_binary(ctx, OpCode::MUL, bits); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } EvmMachineInstKind::EvmUmod(_) => basic_op(ctx, &[OpCode::MOD]), EvmMachineInstKind::EvmSmod(_) => basic_op(ctx, &[OpCode::SMOD]), @@ -373,14 +367,14 @@ impl EvmMachineFunctionLowering<'_> { EvmMachineInstKind::EvmSelfDestruct(_) => basic_op(ctx, &[OpCode::SELFDESTRUCT]), EvmMachineInstKind::GetFunctionPtr(get_fn) => { let func = *get_fn.func(); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); ctx.push_jump_target(OpCode::PUSH1, Label::Function(func)); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } EvmMachineInstKind::EvmInvalid(_) => basic_op(ctx, &[OpCode::INVALID]), EvmMachineInstKind::SymAddr(sym_addr) => { let sym = sym_addr.sym().clone(); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); ctx.push_sym_fixup( OpCode::PUSH0, SymFixup { @@ -388,11 +382,11 @@ impl EvmMachineFunctionLowering<'_> { sym, }, ); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } EvmMachineInstKind::SymSize(sym_size) => { let sym = sym_size.sym().clone(); - emit_pre_actions(ctx, &alloc.read(insn, &args)); + emit_pre_actions(ctx, alloc.pre_inst(insn)); ctx.push_sym_fixup( OpCode::PUSH0, SymFixup { @@ -400,7 +394,7 @@ impl EvmMachineFunctionLowering<'_> { sym, }, ); - emit_post_actions(ctx, &alloc.write(insn, results.as_slice())); + emit_post_actions(ctx, alloc.post_inst(insn)); } } } diff --git a/crates/codegen/src/isa/evm/emit/layout.rs b/crates/codegen/src/isa/evm/emit/layout.rs index a3e55c2bf..2e20674a8 100644 --- a/crates/codegen/src/isa/evm/emit/layout.rs +++ b/crates/codegen/src/isa/evm/emit/layout.rs @@ -152,8 +152,8 @@ pub(crate) fn compute_late_block_alias_plan( }; if Some(block) == entry - || !alloc.read(term, &[]).is_empty() - || !alloc.write(term, &[]).is_empty() + || !alloc.pre_inst(term).is_empty() + || !alloc.post_inst(term).is_empty() || lazy_frame_mentions_trampoline_site(frame_summary, block, term) { continue; diff --git a/crates/codegen/src/isa/evm/emit/mod.rs b/crates/codegen/src/isa/evm/emit/mod.rs index b397c345d..a83f54fcc 100644 --- a/crates/codegen/src/isa/evm/emit/mod.rs +++ b/crates/codegen/src/isa/evm/emit/mod.rs @@ -204,14 +204,14 @@ impl<'a> EvmMachineFunctionLowering<'a> { let _span = trace_span!("sonatina.codegen.evm.lower_prepared_machine_function.lower").entered(); module.func_store.view(func, |function| { - let mut alloc = FinalAlloc::new( + let alloc = FinalAlloc::new( self.function_plan.alloc.clone(), self.function_plan.mem_plan.clone(), ); let lower = Lower::new(&module.ctx, function, &emitted_block_order); lower .lower( - &mut alloc, + &alloc, |ctx, alloc, block| self.enter_block(ctx, alloc, block), |ctx, alloc, function| self.enter_function(ctx, alloc, function), |ctx, alloc, insn| self.lower_insn(ctx, alloc, insn), @@ -265,12 +265,7 @@ impl<'a> EvmMachineFunctionLowering<'a> { }) } - fn enter_function( - &self, - ctx: &mut Lower, - alloc: &mut dyn Allocator, - function: &Function, - ) { + fn enter_function(&self, ctx: &mut Lower, alloc: &dyn Allocator, function: &Function) { let frame_layout = self.frame_layout(); let actions = alloc.enter_function(function); debug_assert!( @@ -294,7 +289,7 @@ impl<'a> EvmMachineFunctionLowering<'a> { } } - fn enter_block(&self, ctx: &mut Lower, _: &mut dyn Allocator, block: BlockId) { + fn enter_block(&self, ctx: &mut Lower, _: &dyn Allocator, block: BlockId) { if self.is_elided_block(block) { return; } diff --git a/crates/codegen/src/isa/evm/machine/lazy_frame.rs b/crates/codegen/src/isa/evm/machine/lazy_frame.rs index 734b268f8..bb6df8979 100644 --- a/crates/codegen/src/isa/evm/machine/lazy_frame.rs +++ b/crates/codegen/src/isa/evm/machine/lazy_frame.rs @@ -321,13 +321,12 @@ fn compute_active_pre_insts( for inst in function.layout.iter_inst(block) { let data = machine_isa.inst_set().resolve_inst(function.dfg.inst(inst)); - let args = function.dfg.inst(inst).collect_values(); apply_site_state(plan, FrameSite::PreInst(inst), &mut active); match data { EvmMachineInstKind::Call(_) => { if let Some((prefix, suffix, prefix_len)) = - split_call_actions(alloc.read(inst, &args)) + split_call_actions(alloc.pre_inst(inst).clone()) { apply_actions_state( plan, @@ -347,7 +346,7 @@ fn compute_active_pre_insts( apply_actions_state( plan, FrameSite::PreInst(inst), - &alloc.read(inst, &args), + alloc.pre_inst(inst), 0, &mut active, ); @@ -355,8 +354,8 @@ fn compute_active_pre_insts( } EvmMachineInstKind::BrTable(br) => { for (case_idx, _) in br.table().iter().enumerate() { - let actions = alloc.read_br_table_case(inst, case_idx); - if fold_stack_actions(&actions).iter().any(|action| { + let actions = alloc.br_table_case(inst, case_idx); + if fold_stack_actions(actions).iter().any(|action| { matches!( action, Action::MemLoadFrameSlot(_) | Action::MemStoreFrameSlot(_) @@ -369,7 +368,7 @@ fn compute_active_pre_insts( _ => apply_actions_state( plan, FrameSite::PreInst(inst), - &alloc.read(inst, &args), + alloc.pre_inst(inst), 0, &mut active, ), @@ -385,7 +384,7 @@ fn compute_active_pre_insts( apply_actions_state( plan, FrameSite::PostInst(inst), - &alloc.write(inst, function.dfg.inst_results(inst)), + alloc.post_inst(inst), 0, &mut active, ); @@ -427,10 +426,10 @@ fn validate_lazy_frame_activity( apply_site_state(plan, FrameSite::BlockEntry(block), &mut active); for inst in function.layout.iter_inst(block) { - let args = function.dfg.inst(inst).collect_values(); apply_site_state(plan, FrameSite::PreInst(inst), &mut active); - if let Some((prefix, suffix, prefix_len)) = split_call_actions(alloc.read(inst, &args)) + if let Some((prefix, suffix, prefix_len)) = + split_call_actions(alloc.pre_inst(inst).clone()) { apply_actions_state(plan, FrameSite::PreInst(inst), &prefix, 0, &mut active); apply_actions_state( @@ -444,7 +443,7 @@ fn validate_lazy_frame_activity( apply_actions_state( plan, FrameSite::PreInst(inst), - &alloc.read(inst, &args), + alloc.pre_inst(inst), 0, &mut active, ); @@ -457,7 +456,7 @@ fn validate_lazy_frame_activity( apply_actions_state( plan, FrameSite::PostInst(inst), - &alloc.write(inst, function.dfg.inst_results(inst)), + alloc.post_inst(inst), 0, &mut active, ); @@ -563,12 +562,11 @@ fn collect_dep_points( for block in function.layout.iter_block() { for inst in function.layout.iter_inst(block) { let data = machine_isa.inst_set().resolve_inst(function.dfg.inst(inst)); - let args = function.dfg.inst(inst).collect_values(); match &data { EvmMachineInstKind::Call(_) => { if let Some((prefix, suffix, prefix_len)) = - split_call_actions(alloc.read(inst, &args)) + split_call_actions(alloc.pre_inst(inst).clone()) { collect_action_dep_points( &mut out, @@ -595,15 +593,15 @@ fn collect_dep_points( &order, block, FrameSite::PreInst(inst), - &alloc.read(inst, &args), + alloc.pre_inst(inst), 0, ); } } EvmMachineInstKind::BrTable(br) => { for (case_idx, _) in br.table().iter().enumerate() { - let actions = alloc.read_br_table_case(inst, case_idx); - if fold_stack_actions(&actions).iter().any(|action| { + let actions = alloc.br_table_case(inst, case_idx); + if fold_stack_actions(actions).iter().any(|action| { matches!( action, Action::MemLoadFrameSlot(_) | Action::MemStoreFrameSlot(_) @@ -619,7 +617,7 @@ fn collect_dep_points( &order, block, FrameSite::PreInst(inst), - &alloc.read(inst, &args), + alloc.pre_inst(inst), 0, ), } @@ -644,7 +642,7 @@ fn collect_dep_points( &order, block, FrameSite::PostInst(inst), - &alloc.write(inst, function.dfg.inst_results(inst)), + alloc.post_inst(inst), 0, ); } diff --git a/crates/codegen/src/isa/evm/tests.rs b/crates/codegen/src/isa/evm/tests.rs index 237600b5f..c627385a2 100644 --- a/crates/codegen/src/isa/evm/tests.rs +++ b/crates/codegen/src/isa/evm/tests.rs @@ -1754,24 +1754,19 @@ block0: }); } - let (call_inst, call_args) = parsed.module.func_store.view(caller, |function| { + let call_inst = parsed.module.func_store.view(caller, |function| { function .layout .iter_block() .flat_map(|block| function.layout.iter_inst(block)) - .find_map(|inst| { - function - .dfg - .cast_call(inst) - .map(|call| (inst, call.args().clone())) - }) + .find_map(|inst| function.dfg.cast_call(inst).map(|_| inst)) .expect("missing call inst") }); let actions = stack_allocs .get(&caller) .expect("missing caller analysis") - .read(call_inst, &call_args); + .pre_inst(call_inst); assert!( !actions .iter() @@ -3000,7 +2995,7 @@ block0: let (shadow_obj, runs) = (&save_plan.shadow_obj, &save_plan.runs); assert!(!runs.is_empty(), "expected at least one saved run"); - let actions = alloc.read(call_inst, &call_args); + let actions = alloc.pre_inst(call_inst); let cont_pos = actions .iter() .position(|a| matches!(a, Action::PushContinuationOffset)) @@ -3107,21 +3102,17 @@ block2: } let caller = names["caller"]; - let (call_inst, call_results): (InstId, SmallVec<[ValueId; 8]>) = - parsed.module.func_store.view(caller, |function| { - for block in function.layout.iter_block() { - for inst in function.layout.iter_inst(block) { - if function.dfg.call_info(inst).is_none() { - continue; - } - return ( - inst, - function.dfg.inst_results(inst).iter().copied().collect(), - ); + let call_inst: InstId = parsed.module.func_store.view(caller, |function| { + for block in function.layout.iter_block() { + for inst in function.layout.iter_inst(block) { + if function.dfg.call_info(inst).is_none() { + continue; } + return inst; } - panic!("missing call inst"); - }); + } + panic!("missing call inst"); + }); let stack_alloc = stack_allocs .remove(&caller) @@ -3151,7 +3142,7 @@ block2: let (shadow_obj, runs) = (&save_plan.shadow_obj, &save_plan.runs); assert!(!runs.is_empty(), "expected at least one saved run"); - let actions = alloc.write(call_inst, &call_results); + let actions = alloc.post_inst(call_inst); let mut expected = Actions::new(); let shadow_loc = alloc.obj_loc_for_id(*shadow_obj); for run in runs.iter().rev() { diff --git a/crates/codegen/src/machinst/lower.rs b/crates/codegen/src/machinst/lower.rs index 508ee71fc..f10f55a2a 100644 --- a/crates/codegen/src/machinst/lower.rs +++ b/crates/codegen/src/machinst/lower.rs @@ -346,10 +346,10 @@ impl<'a, Op: Default> Lower<'a, Op> { pub fn lower( mut self, - alloc: &mut dyn Allocator, - mut enter_block: impl FnMut(&mut Self, &mut dyn Allocator, BlockId), - mut enter_function: impl FnMut(&mut Self, &mut dyn Allocator, &Function), - mut lower_insn: impl FnMut(&mut Self, &mut dyn Allocator, InstId), + alloc: &dyn Allocator, + mut enter_block: impl FnMut(&mut Self, &dyn Allocator, BlockId), + mut enter_function: impl FnMut(&mut Self, &dyn Allocator, &Function), + mut lower_insn: impl FnMut(&mut Self, &dyn Allocator, InstId), ) -> CodegenResult> { let function = self.function; let entry = function.layout.entry_block(); diff --git a/crates/codegen/src/stackalloc/mod.rs b/crates/codegen/src/stackalloc/mod.rs index 3585b301a..2e0c7a4ac 100644 --- a/crates/codegen/src/stackalloc/mod.rs +++ b/crates/codegen/src/stackalloc/mod.rs @@ -1,7 +1,7 @@ use crate::bitset::BitSet; use cranelift_entity::SecondaryMap; use smallvec::SmallVec; -use sonatina_ir::{BlockId, Function, Immediate, InstId, ValueId}; +use sonatina_ir::{Function, Immediate, InstId, ValueId}; use crate::isa::evm::static_arena_alloc::StackObjId; @@ -14,18 +14,15 @@ pub use stackify::{StackifyAlloc, StackifyBuilder, StackifySearchProfile}; pub type Actions = SmallVec<[Action; 2]>; pub trait Allocator { + /// Actions to run at function entry (function-argument spill stores). fn enter_function(&self, function: &Function) -> Actions; - // xxx rename these to make it clear that these are pre- and post-insn operations - /// Return the actions required to place `vals` on the stack, - /// in the specified order. I.e. the first `Value` in `vals` - /// will be on the top of the stack. - fn read(&self, inst: InstId, vals: &[ValueId]) -> Actions; - /// Return the actions required for the `case_index`th `br_table` compare in IR order. - fn read_br_table_case(&self, inst: InstId, case_index: usize) -> Actions; - fn write(&self, inst: InstId, vals: &[ValueId]) -> Actions; - - fn traverse_edge(&self, from: BlockId, to: BlockId) -> Actions; + /// Actions to run immediately before `inst`'s opcode(s). + fn pre_inst(&self, inst: InstId) -> &Actions; + /// Actions to run immediately after `inst`'s opcode(s). + fn post_inst(&self, inst: InstId) -> &Actions; + /// Actions preparing the `case_index`th `br_table` compare, in IR case order. + fn br_table_case(&self, inst: InstId, case_index: usize) -> &Actions; } pub(crate) fn canonicalize_value_alias( diff --git a/crates/codegen/src/stackalloc/stackify/alloc.rs b/crates/codegen/src/stackalloc/stackify/alloc.rs index 214e2abd3..62edffe92 100644 --- a/crates/codegen/src/stackalloc/stackify/alloc.rs +++ b/crates/codegen/src/stackalloc/stackify/alloc.rs @@ -1,6 +1,6 @@ use cranelift_entity::SecondaryMap; use rustc_hash::FxHashMap; -use sonatina_ir::{BlockId, Function, InstId, ValueId}; +use sonatina_ir::{Function, InstId, ValueId}; use crate::{ analysis::memory_access::ExactLocalAddr, @@ -15,6 +15,18 @@ pub(crate) enum SpillStorage { ExactLocal(ExactLocalAddr), } +impl SpillStorage { + /// The action that stores a value from the stack top into this storage. + /// `None` for storage materialized without a store (exact local addresses). + pub(super) fn store_action(self) -> Option { + match self { + SpillStorage::Scratch(slot) => Some(Action::MemStoreAbs(slot * 32)), + SpillStorage::Object(obj) => Some(Action::MemStoreObj(obj)), + SpillStorage::ExactLocal(_) => None, + } + } +} + #[derive(Clone, Default)] pub struct StackifyAlloc { pub(super) pre_actions: SecondaryMap, @@ -134,6 +146,34 @@ impl StackifyAlloc { } }); } + /// Ensure `inst`'s pre/post action entries exist (empty if never planned), so that map-wide + /// transforms like [`Self::rewrite_action_lists`] visit them. + pub(crate) fn touch_inst_actions(&mut self, inst: InstId) { + let _ = &mut self.pre_actions[inst]; + let _ = &mut self.post_actions[inst]; + } + + /// Replace every stored pre/post/`br_table` action list in place with the result of the + /// given transforms (each consumes the old list and returns the rewritten one). + pub(crate) fn rewrite_action_lists( + &mut self, + mut pre: impl FnMut(InstId, Actions) -> Actions, + mut post: impl FnMut(InstId, Actions) -> Actions, + mut br_case: impl FnMut(Actions) -> Actions, + ) { + for (inst, actions) in self.pre_actions.iter_mut() { + *actions = pre(inst, std::mem::take(actions)); + } + for (inst, actions) in self.post_actions.iter_mut() { + *actions = post(inst, std::mem::take(actions)); + } + for cases in self.brtable_actions.values_mut() { + for actions in cases.iter_mut() { + *actions = br_case(std::mem::take(actions)); + } + } + } + pub(crate) fn remap_stack_objects(&mut self, remap: &FxHashMap) { fn remap_actions(actions: &mut Actions, remap: &FxHashMap) { for action in actions { @@ -182,29 +222,24 @@ impl Allocator for StackifyAlloc { idx < super::DUP_MAX, "function arg depth exceeds DUP16 reach" ); - match self.storage_for_value(arg) { - Some(SpillStorage::Scratch(slot)) => { - act.push(Action::StackDup(idx as u8)); - act.push(Action::MemStoreAbs(slot * 32)); - } - Some(SpillStorage::Object(obj)) => { - act.push(Action::StackDup(idx as u8)); - act.push(Action::MemStoreObj(obj)); - } - Some(SpillStorage::ExactLocal(_)) | None => {} + if let Some(store) = self + .storage_for_value(arg) + .and_then(SpillStorage::store_action) + { + act.push(Action::StackDup(idx as u8)); + act.push(store); } } act } - fn read(&self, inst: InstId, _vals: &[ValueId]) -> Actions { - self.pre_actions[inst].clone() + fn pre_inst(&self, inst: InstId) -> &Actions { + &self.pre_actions[inst] } - fn read_br_table_case(&self, inst: InstId, case_index: usize) -> Actions { + fn br_table_case(&self, inst: InstId, case_index: usize) -> &Actions { self.brtable_actions[inst] .get(case_index) - .cloned() .unwrap_or_else(|| { panic!( "missing br_table case actions for inst {} case {}", @@ -214,11 +249,7 @@ impl Allocator for StackifyAlloc { }) } - fn write(&self, inst: InstId, _vals: &[ValueId]) -> Actions { - self.post_actions[inst].clone() - } - - fn traverse_edge(&self, _from: BlockId, _to: BlockId) -> Actions { - Actions::new() + fn post_inst(&self, inst: InstId) -> &Actions { + &self.post_actions[inst] } } diff --git a/crates/codegen/src/stackalloc/stackify/mod.rs b/crates/codegen/src/stackalloc/stackify/mod.rs index 8d7e74e8c..4c073a9bf 100644 --- a/crates/codegen/src/stackalloc/stackify/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/mod.rs @@ -216,8 +216,8 @@ block3: ); assert_eq!(alloc.brtable_actions[term].len(), 2); - let first = alloc.read_br_table_case(term, 0); - let second = alloc.read_br_table_case(term, 1); + let first = alloc.br_table_case(term, 0); + let second = alloc.br_table_case(term, 1); assert!( !first.is_empty(), "expected first br_table case to include compare preparation" diff --git a/crates/codegen/src/stackalloc/stackify/planner/mod.rs b/crates/codegen/src/stackalloc/stackify/planner/mod.rs index abac1cf87..a56c948f4 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/mod.rs @@ -17,6 +17,7 @@ use sonatina_ir::ValueId; use super::{ StackifyContext, + alloc::SpillStorage, slots::{FreeSlotPools, SpillSlotInterference, SpillSlotPools}, spill::{SpillDiscovery, SpillSet}, sym_stack::SymStack, @@ -123,15 +124,22 @@ impl<'a> MemPlan<'a> { &mut self.free_slots.scratch, Some(self.scratch_spill_slots), ) { - actions.push(Action::MemStoreAbs(slot * 32)); + actions.push( + SpillStorage::Scratch(slot) + .store_action() + .expect("scratch storage has a store action"), + ); return; } self.request_object_storage(v); } - actions.push(Action::MemStoreObj( - self.spill_obj[v].expect("spilled value missing stack object id"), - )); + let obj = self.spill_obj[v].expect("spilled value missing stack object id"); + actions.push( + SpillStorage::Object(obj) + .store_action() + .expect("object storage has a store action"), + ); } fn load_frame_slot_or_placeholder(&mut self, v: ValueId) -> Action { From e51a882380cc552e4bf385596ae8f06872add6a6 Mon Sep 17 00:00:00 2001 From: sbillig Date: Mon, 6 Jul 2026 22:30:10 -0700 Subject: [PATCH 02/14] Make spill_storage the single stored truth in StackifyAlloc Delete the spill_obj and scratch_slot_of_value projection fields; derive them via accessors (spill_obj, scratch_slot, object_spills) with a single set_spill_object mutation path. The provisional per-iteration object-id map moves out of StackifyAlloc into the planning driver. The map-drift assertions in validate_spill_storage become unrepresentable. --- crates/codegen/src/isa/evm/backend.rs | 2 +- .../src/isa/evm/machine/final_spills.rs | 13 +-- .../codegen/src/isa/evm/static_arena_alloc.rs | 32 +++--- .../codegen/src/stackalloc/stackify/alloc.rs | 102 +++++++----------- .../src/stackalloc/stackify/builder.rs | 45 ++++---- .../src/stackalloc/stackify/iteration.rs | 16 ++- crates/codegen/src/stackalloc/stackify/mod.rs | 10 +- 7 files changed, 91 insertions(+), 129 deletions(-) diff --git a/crates/codegen/src/isa/evm/backend.rs b/crates/codegen/src/isa/evm/backend.rs index f544f45ac..37c49f305 100644 --- a/crates/codegen/src/isa/evm/backend.rs +++ b/crates/codegen/src/isa/evm/backend.rs @@ -265,7 +265,7 @@ impl EvmBackend { if let Some(alloc) = prepared.function_plan(func).map(|plan| &plan.alloc) { module.func_store.view(func, |function| { for v in function.dfg.value_ids() { - let Some(slot) = alloc.scratch_slot_of_value[v] else { + let Some(slot) = alloc.scratch_slot(v) else { continue; }; scratch_spills.push((v, slot)); diff --git a/crates/codegen/src/isa/evm/machine/final_spills.rs b/crates/codegen/src/isa/evm/machine/final_spills.rs index d7a8ff65d..3aee72de5 100644 --- a/crates/codegen/src/isa/evm/machine/final_spills.rs +++ b/crates/codegen/src/isa/evm/machine/final_spills.rs @@ -496,7 +496,7 @@ pub(crate) fn allocate_final_spills( alloc.remap_stack_objects(&remap); for (value, old_obj) in spills.spilled_values { let new_obj = remap[&old_obj]; - alloc.spill_obj[value] = Some(new_obj); + alloc.set_spill_object(value, new_obj); mem_plan.spill_obj[value] = Some(new_obj); } alloc.validate_spill_storage(); @@ -511,16 +511,7 @@ pub(crate) fn allocate_final_spills( } fn final_spilled_values(alloc: &StackifyAlloc) -> Vec<(ValueId, StackObjId)> { - alloc - .spill_obj - .iter() - .filter_map(|(value, obj)| { - let obj = (*obj)?; - alloc.scratch_slot_of_value[value] - .is_none() - .then_some((value, obj)) - }) - .collect() + alloc.object_spills().collect() } fn spill_count(len: usize) -> u32 { diff --git a/crates/codegen/src/isa/evm/static_arena_alloc.rs b/crates/codegen/src/isa/evm/static_arena_alloc.rs index 68d0f71b4..afc59d4f3 100644 --- a/crates/codegen/src/isa/evm/static_arena_alloc.rs +++ b/crates/codegen/src/isa/evm/static_arena_alloc.rs @@ -333,14 +333,9 @@ fn compute_func_stack_objects_from_input( let mut spilled_values: BitSet = BitSet::default(); if let Some(alloc) = analysis.stackify_alloc { - for (v, obj) in alloc.spill_obj.iter() { - if alloc.scratch_slot_of_value[v].is_some() { - continue; - } - if obj.is_some() { - spilled_values.insert(v); - spill_obj[v] = *obj; - } + for (v, obj) in alloc.object_spills() { + spilled_values.insert(v); + spill_obj[v] = Some(obj); } } @@ -360,10 +355,8 @@ fn compute_func_stack_objects_from_input( let mut next_id: u32 = analysis.stackify_alloc.map_or(0, |alloc| { alloc - .spill_obj - .values() - .filter_map(|o| *o) - .map(|id| id.as_u32()) + .object_spills() + .map(|(_, id)| id.as_u32()) .max() .map_or(0, |n| n.checked_add(1).expect("stack object id overflow")) }); @@ -859,7 +852,10 @@ mod tests { .value(func_ref, "v5") .expect("aliased pointer value exists"), ); - let spill_obj = analysis.alloc.spill_obj[spill_value].expect("spill object exists"); + let spill_obj = analysis + .alloc + .spill_obj(spill_value) + .expect("spill object exists"); let call = stack .call_sites @@ -945,8 +941,10 @@ block3: let (parsed, func_ref, analysis, stack) = analyze_function(&src, "loop_spill", 16); let spilled = analysis.canonicalize_value(parsed.debug.value(func_ref, "v18").expect("v18 exists")); - let spill_obj = - analysis.alloc.spill_obj[spilled].expect("loop-carried phi spill object exists"); + let spill_obj = analysis + .alloc + .spill_obj(spilled) + .expect("loop-carried phi spill object exists"); let region = &stack.obj_facts[&spill_obj].region; assert!( region @@ -1081,7 +1079,7 @@ block3: .value(func_ref, &name) .expect("branch value exists"), ); - let obj = analysis.alloc.spill_obj[value]?; + let obj = analysis.alloc.spill_obj(value)?; Some(&stack.obj_facts[&obj].region) }) .expect("left branch should produce a spilled phi operand value"); @@ -1094,7 +1092,7 @@ block3: .value(func_ref, &name) .expect("branch value exists"), ); - let obj = analysis.alloc.spill_obj[value]?; + let obj = analysis.alloc.spill_obj(value)?; Some(&stack.obj_facts[&obj].region) }) .expect("right branch should produce a spilled phi operand value"); diff --git a/crates/codegen/src/stackalloc/stackify/alloc.rs b/crates/codegen/src/stackalloc/stackify/alloc.rs index 62edffe92..54228b80d 100644 --- a/crates/codegen/src/stackalloc/stackify/alloc.rs +++ b/crates/codegen/src/stackalloc/stackify/alloc.rs @@ -35,29 +35,59 @@ pub struct StackifyAlloc { /// `br_table` lowering uses per-case action sequences stored in IR case order. pub(super) brtable_actions: SecondaryMap>, + /// Finalized storage for every spilled value. Single source of truth; the + /// object/scratch projections below are derived from it on demand. pub(crate) spill_storage: SecondaryMap>, - pub(crate) spill_obj: SecondaryMap>, - pub(crate) scratch_slot_of_value: SecondaryMap>, pub(crate) exact_local_addr: SecondaryMap>, } impl StackifyAlloc { #[cfg(test)] pub(crate) fn set_object_spill_for_test(&mut self, value: ValueId, obj: StackObjId) { - self.spill_obj[value] = Some(obj); - self.spill_storage[value] = Some(SpillStorage::Object(obj)); + self.set_spill_object(value, obj); } pub(crate) fn uses_scratch_spills(&self) -> bool { - self.scratch_slot_of_value + self.spill_storage .values() - .any(|slot| slot.is_some()) + .any(|storage| matches!(storage, Some(SpillStorage::Scratch(_)))) } pub(crate) fn storage_for_value(&self, value: ValueId) -> Option { self.spill_storage[value] } + #[cfg(test)] + pub(crate) fn spill_obj(&self, value: ValueId) -> Option { + match self.spill_storage[value] { + Some(SpillStorage::Object(obj)) => Some(obj), + _ => None, + } + } + + pub(crate) fn scratch_slot(&self, value: ValueId) -> Option { + match self.spill_storage[value] { + Some(SpillStorage::Scratch(slot)) => Some(slot), + _ => None, + } + } + + /// Iterate `(value, obj)` pairs for every object-stored spill. + pub(crate) fn object_spills(&self) -> impl Iterator + '_ { + self.spill_storage.iter().filter_map(|(value, storage)| { + if let Some(SpillStorage::Object(obj)) = storage { + Some((value, *obj)) + } else { + None + } + }) + } + + /// The one mutation downstream needs (final spill re-homing). + pub(crate) fn set_spill_object(&mut self, value: ValueId, obj: StackObjId) { + self.spill_storage[value] = Some(SpillStorage::Object(obj)); + } + pub(crate) fn for_each_action(&self, mut f: impl FnMut(&Action)) { for actions in self.pre_actions.values() { for action in actions { @@ -79,61 +109,6 @@ impl StackifyAlloc { } pub(crate) fn validate_spill_storage(&self) { - for (value, storage) in self.spill_storage.iter() { - match storage { - Some(SpillStorage::Scratch(slot)) => { - assert_eq!( - self.scratch_slot_of_value[value], - Some(*slot), - "scratch storage map drift for value {}", - value.as_u32() - ); - assert_eq!( - self.spill_obj[value], - None, - "scratch-spilled value {} must not have object storage", - value.as_u32() - ); - } - Some(SpillStorage::Object(obj)) => { - assert_eq!( - self.spill_obj[value], - Some(*obj), - "object storage map drift for value {}", - value.as_u32() - ); - assert_eq!( - self.scratch_slot_of_value[value], - None, - "object-spilled value {} must not have scratch storage", - value.as_u32() - ); - } - Some(SpillStorage::ExactLocal(exact)) => { - assert_eq!( - self.exact_local_addr[value], - Some(*exact), - "exact local storage map drift for value {}", - value.as_u32() - ); - } - None => { - assert_eq!( - self.scratch_slot_of_value[value], - None, - "unspilled value {} must not have scratch storage", - value.as_u32() - ); - assert_eq!( - self.spill_obj[value], - None, - "unspilled value {} must not have object storage", - value.as_u32() - ); - } - } - } - self.for_each_action(|action| { if let Action::MemLoadObj(id) | Action::MemStoreObj(id) = action { assert!( @@ -206,11 +181,6 @@ impl StackifyAlloc { *obj = *new_obj; } } - for obj in self.spill_obj.values_mut().flatten() { - if let Some(new_obj) = remap.get(obj) { - *obj = *new_obj; - } - } } } diff --git a/crates/codegen/src/stackalloc/stackify/builder.rs b/crates/codegen/src/stackalloc/stackify/builder.rs index fd3647245..629316113 100644 --- a/crates/codegen/src/stackalloc/stackify/builder.rs +++ b/crates/codegen/src/stackalloc/stackify/builder.rs @@ -347,14 +347,15 @@ impl<'a> StackifyBuilder<'a> { let checkpoint = observer.checkpoint(); let mut slots: SpillSlotPools = SpillSlotPools::default(); - let (mut alloc, spill_requests, object_spill_requests) = Self::plan_iteration( - &ctx, - observer, - SpillSet::new(&spill_set), - &forced_object_spills, - &mut slots, - &mut search_scratch, - ); + let (mut alloc, spill_obj, spill_requests, object_spill_requests) = + Self::plan_iteration( + &ctx, + observer, + SpillSet::new(&spill_set), + &forced_object_spills, + &mut slots, + &mut search_scratch, + ); let spill_stable = spill_requests.is_subset(&spill_set); let object_spills_stable = object_spill_requests.is_subset(&forced_object_spills); @@ -365,6 +366,7 @@ impl<'a> StackifyBuilder<'a> { &forced_object_spills, &mut slots, &mut alloc, + &spill_obj, ); alloc.validate_spill_storage(); return alloc; @@ -383,7 +385,12 @@ impl<'a> StackifyBuilder<'a> { forced_object_spills: &BitSet, slots: &mut SpillSlotPools, search_scratch: &mut NormalizeSearchScratch, - ) -> (StackifyAlloc, BitSet, BitSet) { + ) -> ( + StackifyAlloc, + SecondaryMap>, + BitSet, + BitSet, + ) { let mut object_spill_requests: BitSet = BitSet::default(); let mut arg_free_slots: FreeSlotPools = FreeSlotPools::default(); for &arg in ctx.func.arg_values.iter() { @@ -417,8 +424,6 @@ impl<'a> StackifyBuilder<'a> { post_actions: SecondaryMap::new(), brtable_actions: SecondaryMap::new(), spill_storage: SecondaryMap::new(), - spill_obj, - scratch_slot_of_value: SecondaryMap::new(), exact_local_addr: ctx.exact_local_addr.clone(), }; @@ -442,6 +447,7 @@ impl<'a> StackifyBuilder<'a> { &terminal_chain_blocks, &interfaces.carry_in, &mut alloc, + &spill_obj, &mut spill_requests, &mut object_spill_requests, forced_object_spills, @@ -451,7 +457,7 @@ impl<'a> StackifyBuilder<'a> { ); planner.plan_blocks(); - (alloc, spill_requests, object_spill_requests) + (alloc, spill_obj, spill_requests, object_spill_requests) } fn finalize_spill_storage( @@ -460,19 +466,12 @@ impl<'a> StackifyBuilder<'a> { forced_object_spills: &BitSet, slots: &mut SpillSlotPools, alloc: &mut StackifyAlloc, + spill_obj: &SecondaryMap>, ) { let scratch_slots = slots.scratch.take_slot_map(); - let raw_spill_obj = alloc.spill_obj.clone(); let mut spill_storage: SecondaryMap> = SecondaryMap::new(); - let mut spill_obj: SecondaryMap< - ValueId, - Option, - > = SecondaryMap::new(); - let mut scratch_slot_of_value: SecondaryMap> = SecondaryMap::new(); for value in ctx.func.dfg.value_ids() { let _ = &mut spill_storage[value]; - let _ = &mut spill_obj[value]; - let _ = &mut scratch_slot_of_value[value]; } for value in spill.bitset().iter() { @@ -482,12 +481,10 @@ impl<'a> StackifyBuilder<'a> { || ctx.scratch_live_values.contains(value) || forced_object_spills.contains(value) { - let obj = raw_spill_obj[value].expect("object spill missing stack object id"); + let obj = spill_obj[value].expect("object spill missing stack object id"); spill_storage[value] = Some(SpillStorage::Object(obj)); - spill_obj[value] = Some(obj); } else if let Some(slot) = scratch_slots[value] { spill_storage[value] = Some(SpillStorage::Scratch(slot)); - scratch_slot_of_value[value] = Some(slot); } else { panic!( "spilled value {} has no stable scratch slot or object storage", @@ -497,8 +494,6 @@ impl<'a> StackifyBuilder<'a> { } alloc.spill_storage = spill_storage; - alloc.spill_obj = spill_obj; - alloc.scratch_slot_of_value = scratch_slot_of_value; } } diff --git a/crates/codegen/src/stackalloc/stackify/iteration.rs b/crates/codegen/src/stackalloc/stackify/iteration.rs index 10bb4c216..5714d82d6 100644 --- a/crates/codegen/src/stackalloc/stackify/iteration.rs +++ b/crates/codegen/src/stackalloc/stackify/iteration.rs @@ -27,6 +27,9 @@ pub(super) struct IterationPlanner<'a, 'ctx, O: StackifyObserver> { terminal_chain_blocks: &'a SecondaryMap, carry_in: &'a SecondaryMap>, alloc: &'a mut StackifyAlloc, + /// Provisional per-iteration object-id assignment (`assign_spill_obj_ids`), read by + /// `MemPlan` during planning before storage is finalized. + spill_obj: &'a SecondaryMap>, spill_requests: &'a mut BitSet, object_spill_requests: &'a mut BitSet, forced_object_spills: &'a BitSet, @@ -67,6 +70,10 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { terminal_chain_blocks: &'a SecondaryMap, carry_in: &'a SecondaryMap>, alloc: &'a mut StackifyAlloc, + spill_obj: &'a SecondaryMap< + ValueId, + Option, + >, spill_requests: &'a mut BitSet, object_spill_requests: &'a mut BitSet, forced_object_spills: &'a BitSet, @@ -82,6 +89,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { terminal_chain_blocks, carry_in, alloc, + spill_obj, spill_requests, object_spill_requests, forced_object_spills, @@ -104,7 +112,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { self.spill, &mut *self.spill_requests, self.ctx, - &self.alloc.spill_obj, + self.spill_obj, &self.alloc.exact_local_addr, self.object_spill_requests, self.forced_object_spills, @@ -313,7 +321,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { self.spill, &mut *self.spill_requests, self.ctx, - &self.alloc.spill_obj, + self.spill_obj, &self.alloc.exact_local_addr, self.object_spill_requests, self.forced_object_spills, @@ -334,7 +342,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { self.spill, &mut *self.spill_requests, self.ctx, - &self.alloc.spill_obj, + self.spill_obj, &self.alloc.exact_local_addr, self.object_spill_requests, self.forced_object_spills, @@ -363,7 +371,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { self.spill, &mut *self.spill_requests, self.ctx, - &self.alloc.spill_obj, + self.spill_obj, &self.alloc.exact_local_addr, self.object_spill_requests, self.forced_object_spills, diff --git a/crates/codegen/src/stackalloc/stackify/mod.rs b/crates/codegen/src/stackalloc/stackify/mod.rs index 4c073a9bf..c7401ca61 100644 --- a/crates/codegen/src/stackalloc/stackify/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/mod.rs @@ -116,8 +116,8 @@ mod tests { .with_scratch_spills(2) .compute(); - for (v, slot) in alloc.scratch_slot_of_value.iter() { - if slot.is_some() { + for v in function.dfg.value_ids() { + if alloc.scratch_slot(v).is_some() { assert!( !scratch_live_values.contains(v), "scratch spill used for a scratch-live value" @@ -128,7 +128,7 @@ mod tests { assert!( scratch_live_values .iter() - .any(|v| alloc.spill_obj[v].is_some()), + .any(|v| alloc.spill_obj(v).is_some()), "expected at least one scratch-live value to spill to a stack object" ); }); @@ -297,7 +297,7 @@ block2: alloc.pre_actions[jump_inst] ); assert!( - alloc.spill_obj[v1].is_some(), + alloc.spill_obj(v1).is_some(), "expected dropped entry arg to become spill-reloadable on the backedge" ); }); @@ -386,7 +386,7 @@ block3: let alloc = StackifyBuilder::new(function, &cfg, &dom, &liveness, 16).compute(); assert!( - alloc.spill_obj[deep_phi].is_some(), + alloc.spill_obj(deep_phi).is_some(), "expected deepest phi to spill so merge repair must store it" ); assert!( From 58fb0446d8468a9eaa1977c655eb354e2dde7a6d Mon Sep 17 00:00:00 2001 From: sbillig Date: Mon, 6 Jul 2026 22:45:44 -0700 Subject: [PATCH 03/14] Stackify small cleanups - Replace three confirmed-unreachable branches with debug asserts: the dead-value defensive loop in clean_dead_stack_prefix, the dominated-pred arm in choose_transfer, and the plan_block prologue fallback; clear_inst_actions becomes an emptiness assertion. - Express StackifySearchProfile knobs as SearchBudgets const tables. - terminal_chain_blocks becomes a BitSet. - Drop the inst collect in run_block_sim (iterate the layout directly). - Share the storage-eligibility predicate between MemPlan and the builder's entry-arg pre-pass. - Document the SWAP-chain cost of the stable SymStack ops; drop the imm_materialization_code_len alias. --- .../src/stackalloc/stackify/block_sim.rs | 8 +- .../src/stackalloc/stackify/builder.rs | 132 +++++++++++------- .../src/stackalloc/stackify/iteration.rs | 53 +++---- .../src/stackalloc/stackify/planner/mod.rs | 24 +++- .../stackalloc/stackify/planner/normalize.rs | 4 +- .../stackify/planner/normalize_search.rs | 28 ++-- .../stackify/planner/operand_prep.rs | 8 +- .../src/stackalloc/stackify/sym_stack.rs | 5 + .../src/stackalloc/stackify/templates.rs | 16 ++- .../src/stackalloc/stackify/terminal_chain.rs | 11 +- 10 files changed, 178 insertions(+), 111 deletions(-) diff --git a/crates/codegen/src/stackalloc/stackify/block_sim.rs b/crates/codegen/src/stackalloc/stackify/block_sim.rs index 4606d4cee..be81299ce 100644 --- a/crates/codegen/src/stackalloc/stackify/block_sim.rs +++ b/crates/codegen/src/stackalloc/stackify/block_sim.rs @@ -154,13 +154,15 @@ pub(super) fn run_block_sim( ) -> BlockSimState { let empty_last_use: BitSet = BitSet::default(); - let insts: Vec<_> = planner.ctx().func.layout.iter_inst(state.block).collect(); - for inst in insts { + // `func` is a shared `&Function` copied out of the planner, so the layout iterator does not + // borrow the planner and is free to coexist with the `&mut planner` calls in the loop body. + let func = planner.ctx().func; + for inst in func.layout.iter_inst(state.block) { if planner.ctx().func.dfg.is_phi(inst) { continue; } - planner.clear_inst_actions(inst); + planner.debug_assert_inst_actions_empty(inst); let is_call = planner.ctx().func.dfg.is_call(inst); let call_has_stack_continuation = is_call && planner.call_uses_stack_continuation(inst); diff --git a/crates/codegen/src/stackalloc/stackify/builder.rs b/crates/codegen/src/stackalloc/stackify/builder.rs index 629316113..7662ed7e2 100644 --- a/crates/codegen/src/stackalloc/stackify/builder.rs +++ b/crates/codegen/src/stackalloc/stackify/builder.rs @@ -3,6 +3,7 @@ use crate::{ bitset::BitSet, cfg_scc::CfgSccAnalysis, domtree::DomTree, + isa::evm::immediate_materialization_code_len, liveness::Liveness, stackalloc::normalize_value_alias_map, }; @@ -13,8 +14,8 @@ use sonatina_ir::{BlockId, Function, I256, ValueId, cfg::ControlFlowGraph}; use super::{ alloc::{SpillStorage, StackifyAlloc}, - iteration::{IterationPlanner, imm_materialization_code_len, operand_order_for_stackify}, - planner::NormalizeSearchScratch, + iteration::{IterationPlanner, operand_order_for_stackify}, + planner::{NormalizeSearchScratch, must_use_object_storage}, slots::{FreeSlotPools, SpillSlotInterference, SpillSlotPools}, spill::SpillSet, sym_stack::SymStack, @@ -41,62 +42,84 @@ pub enum StackifySearchProfile { Exact, } -impl StackifySearchProfile { - pub(super) fn use_exact_normalize(self) -> bool { - matches!(self, Self::GreedyWide | Self::Exact) - } - - pub(super) fn use_exact_operand_prep(self) -> bool { - matches!(self, Self::GreedyWide | Self::Exact) - } +/// Operand-prep beam search depth slack: either a fixed amount or "as deep as `SWAP*` reach". +#[derive(Clone, Copy, Debug)] +pub(super) enum BeamSlack { + Fixed(usize), + SwapMax, +} - pub(super) fn exact_expansions(self) -> usize { +impl BeamSlack { + pub(super) fn resolve(self, swap_max: usize) -> usize { match self { - Self::Fast => 0, - Self::GreedyWide => 1_000, - Self::Exact => 50_000, + BeamSlack::Fixed(n) => n, + BeamSlack::SwapMax => swap_max, } } +} - pub(super) fn operand_prep_exact_expansions(self) -> usize { - match self { - Self::Fast => 0, - Self::GreedyWide => 250, - Self::Exact => 50_000, - } - } +/// Search-effort knobs for a `StackifySearchProfile`, expressed as data instead of code. +pub(super) struct SearchBudgets { + pub(super) exact_normalize: bool, + pub(super) exact_operand_prep: bool, + pub(super) exact_expansions: usize, + pub(super) operand_prep_exact_expansions: usize, + /// `(with incumbent, without incumbent)`. + normalize_max_states: (usize, usize), + pub(super) operand_prep_max_states: usize, + pub(super) operand_prep_beam_width: usize, + pub(super) operand_prep_beam_depth_slack: BeamSlack, +} - pub(super) fn normalize_max_states(self, have_incumbent: bool) -> usize { - match (self, have_incumbent) { - (Self::Fast, _) => 0, - (Self::GreedyWide, true) => 25_000, - (Self::GreedyWide, false) => 50_000, - (Self::Exact, true) => 200_000, - (Self::Exact, false) => 500_000, +impl SearchBudgets { + pub(super) fn normalize_max_states(&self, have_incumbent: bool) -> usize { + if have_incumbent { + self.normalize_max_states.0 + } else { + self.normalize_max_states.1 } } +} - pub(super) fn operand_prep_max_states(self) -> usize { - match self { - Self::Fast => 0, - Self::GreedyWide => 25_000, - Self::Exact => 400_000, - } - } +const FAST_BUDGETS: SearchBudgets = SearchBudgets { + exact_normalize: false, + exact_operand_prep: false, + exact_expansions: 0, + operand_prep_exact_expansions: 0, + normalize_max_states: (0, 0), + operand_prep_max_states: 0, + operand_prep_beam_width: 16, + operand_prep_beam_depth_slack: BeamSlack::Fixed(4), +}; - pub(super) fn operand_prep_beam_width(self) -> usize { - match self { - Self::Fast => 16, - Self::GreedyWide => 64, - Self::Exact => 192, - } - } +const GREEDY_WIDE_BUDGETS: SearchBudgets = SearchBudgets { + exact_normalize: true, + exact_operand_prep: true, + exact_expansions: 1_000, + operand_prep_exact_expansions: 250, + normalize_max_states: (25_000, 50_000), + operand_prep_max_states: 25_000, + operand_prep_beam_width: 64, + operand_prep_beam_depth_slack: BeamSlack::SwapMax, +}; + +const EXACT_BUDGETS: SearchBudgets = SearchBudgets { + exact_normalize: true, + exact_operand_prep: true, + exact_expansions: 50_000, + operand_prep_exact_expansions: 50_000, + normalize_max_states: (200_000, 500_000), + operand_prep_max_states: 400_000, + operand_prep_beam_width: 192, + operand_prep_beam_depth_slack: BeamSlack::SwapMax, +}; - pub(super) fn operand_prep_beam_depth_slack(self, swap_max: usize) -> usize { +impl StackifySearchProfile { + pub(super) fn budgets(self) -> &'static SearchBudgets { match self { - Self::Fast => 4, - Self::GreedyWide => swap_max, - Self::Exact => swap_max, + Self::Fast => &FAST_BUDGETS, + Self::GreedyWide => &GREEDY_WIDE_BUDGETS, + Self::Exact => &EXACT_BUDGETS, } } } @@ -397,9 +420,12 @@ impl<'a> StackifyBuilder<'a> { let arg = ctx.canonicalize_value(arg); if let Some(spilled) = spill.spilled(arg) && ctx.exact_local_addr[arg].is_none() - && ctx.scratch_spill_slots != 0 - && !ctx.scratch_live_values.contains(arg) - && !forced_object_spills.contains(arg) + && !must_use_object_storage( + ctx.scratch_spill_slots, + &ctx.scratch_live_values, + forced_object_spills, + arg, + ) && slots .scratch .try_ensure_slot( @@ -524,7 +550,7 @@ fn compute_hot_stack_cached_immediates( let entry = counts.entry(imm.as_i256()).or_insert((0, 0)); entry.0 += 1; - entry.1 = entry.1.max(imm_materialization_code_len(imm)); + entry.1 = entry.1.max(immediate_materialization_code_len(imm)); } } @@ -689,7 +715,7 @@ block0: let mut large_pushes = 0; let mut dup_count = 0; alloc.for_each_action(|action| match action { - Action::Push(imm) if super::imm_materialization_code_len(*imm) >= 17 => { + Action::Push(imm) if super::immediate_materialization_code_len(*imm) >= 17 => { large_pushes += 1; } Action::StackDup(_) => { @@ -741,7 +767,7 @@ block0: let mut push2_count = 0; let mut dup_count = 0; alloc.for_each_action(|action| match action { - Action::Push(imm) if super::imm_materialization_code_len(*imm) == 3 => { + Action::Push(imm) if super::immediate_materialization_code_len(*imm) == 3 => { push2_count += 1; } Action::StackDup(_) => { @@ -796,7 +822,7 @@ block0: let mut large_pushes = 0; let mut dup_count = 0; alloc.for_each_action(|action| match action { - Action::Push(imm) if super::imm_materialization_code_len(*imm) >= 17 => { + Action::Push(imm) if super::immediate_materialization_code_len(*imm) >= 17 => { large_pushes += 1; } Action::StackDup(_) => { diff --git a/crates/codegen/src/stackalloc/stackify/iteration.rs b/crates/codegen/src/stackalloc/stackify/iteration.rs index 5714d82d6..af5aeb909 100644 --- a/crates/codegen/src/stackalloc/stackify/iteration.rs +++ b/crates/codegen/src/stackalloc/stackify/iteration.rs @@ -1,6 +1,6 @@ use cranelift_entity::SecondaryMap; use smallvec::SmallVec; -use sonatina_ir::{BlockId, Function, I256, Immediate, InstId, ValueId}; +use sonatina_ir::{BlockId, Function, I256, InstId, ValueId}; use std::collections::BTreeMap; use crate::{bitset::BitSet, isa::evm::immediate_materialization_code_len, stackalloc::Actions}; @@ -24,7 +24,7 @@ pub(super) struct IterationPlanner<'a, 'ctx, O: StackifyObserver> { spill: SpillSet<'a>, slots: &'a mut SpillSlotPools, templates: &'a mut SecondaryMap, - terminal_chain_blocks: &'a SecondaryMap, + terminal_chain_blocks: &'a BitSet, carry_in: &'a SecondaryMap>, alloc: &'a mut StackifyAlloc, /// Provisional per-iteration object-id assignment (`assign_spill_obj_ids`), read by @@ -67,7 +67,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { spill: SpillSet<'a>, slots: &'a mut SpillSlotPools, templates: &'a mut SecondaryMap, - terminal_chain_blocks: &'a SecondaryMap, + terminal_chain_blocks: &'a BitSet, carry_in: &'a SecondaryMap>, alloc: &'a mut StackifyAlloc, spill_obj: &'a SecondaryMap< @@ -146,7 +146,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { self.resolve_pending_edges(block); let inherited = self.inherited_stack.remove(&block); - if self.terminal_chain_blocks[block] { + if self.terminal_chain_blocks.contains(block) { self.freeze_template(block, TransferOrder::new()); } else if let Some((_pred, stack)) = inherited.as_ref() { self.freeze_template_from_stack(block, stack); @@ -158,7 +158,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { .on_block_header(self.ctx.func, block, &self.templates[block]); self.planned_blocks.insert(block); - let stack = if self.terminal_chain_blocks[block] { + let stack = if self.terminal_chain_blocks.contains(block) { SymStack::opaque_prefix_empty(self.ctx.has_internal_return) } else if let Some((pred, mut inh)) = inherited { // Dynamic entry stack (single predecessor). @@ -202,13 +202,12 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { let state = BlockSimState::with_live_sets(block, stack, free_slots, prologue, live_sets); let state = run_block_sim(self, state); - // If the block had no lowered instructions, inject prologue into the terminator. - if !state.prologue.is_empty() - && !state.injected_prologue - && let Some(term) = self.ctx.func.layout.last_inst_of(block) - { - self.alloc.pre_actions[term].extend_from_slice(&state.prologue); - } + // `on_inst_start` runs for every non-phi instruction (including the terminator that every + // block has), so a non-empty prologue is always injected during `run_block_sim`. + debug_assert!( + state.prologue.is_empty() || state.injected_prologue, + "prologue was not injected during the block walk" + ); } fn resolve_pending_edges(&mut self, block: BlockId) { @@ -286,10 +285,15 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { &self.slots.scratch } - pub(super) fn clear_inst_actions(&mut self, inst: InstId) { - self.alloc.pre_actions[inst].clear(); - self.alloc.post_actions[inst].clear(); - self.alloc.brtable_actions[inst].clear(); + pub(super) fn debug_assert_inst_actions_empty(&self, inst: InstId) { + // Within one fixed-point iteration each inst is visited once on a fresh `StackifyAlloc`, + // so its action buffers are always still empty when the walk reaches it. + debug_assert!( + self.alloc.pre_actions[inst].is_empty() + && self.alloc.post_actions[inst].is_empty() + && self.alloc.brtable_actions[inst].is_empty(), + "inst action buffers are not empty at the start of the block walk" + ); } pub(super) fn pre_actions_len(&self, inst: InstId) -> usize { @@ -456,7 +460,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { dest: BlockId, action_start: usize, ) { - if self.terminal_chain_blocks[dest] { + if self.terminal_chain_blocks.contains(dest) { } else if self.ctx.cfg.pred_num_of(dest) > 1 && dest != self.ctx.entry && !self.planned_blocks.contains(dest) @@ -659,10 +663,11 @@ pub(super) fn clean_dead_stack_prefix( let Some(StackItem::Value(top)) = stack.top() else { break; }; - if !live_future.contains(*top) && !live_out.contains(*top) { - // Should have been handled by `pop_dead_tops`, but keep looping defensively. - continue; - } + // `pop_dead_tops` only stops on a live value (or a non-`Value`, handled above). + debug_assert!( + live_future.contains(*top) || live_out.contains(*top), + "pop_dead_tops left a dead value on top" + ); let is_dead = |v: ValueId| !live_future.contains(v) && !live_out.contains(v); let dead_run = stack @@ -902,9 +907,5 @@ fn is_evictable_imm(func: &Function, v: ValueId) -> bool { let Some(imm) = func.dfg.value_imm(v) else { return false; }; - imm_materialization_code_len(imm) <= MAX_MATERIALIZATION_BYTES -} - -pub(crate) fn imm_materialization_code_len(imm: Immediate) -> usize { - immediate_materialization_code_len(imm) + immediate_materialization_code_len(imm) <= MAX_MATERIALIZATION_BYTES } diff --git a/crates/codegen/src/stackalloc/stackify/planner/mod.rs b/crates/codegen/src/stackalloc/stackify/planner/mod.rs index a56c948f4..0fa8087b3 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/mod.rs @@ -23,6 +23,21 @@ use super::{ sym_stack::SymStack, }; +/// Storage-eligibility policy, owned in one place and shared by `MemPlan` and the entry-arg +/// scratch pre-pass in the builder. +/// +/// A spilled value must use arena (object) storage rather than a reusable scratch slot when there +/// are no scratch slots, when it is live across a scratch clobber, or when a previous fixed-point +/// iteration already forced object storage for it. +pub(super) fn must_use_object_storage( + scratch_spill_slots: u32, + scratch_live_values: &BitSet, + forced_object_spills: &BitSet, + v: ValueId, +) -> bool { + scratch_spill_slots == 0 || scratch_live_values.contains(v) || forced_object_spills.contains(v) +} + #[derive(Clone)] pub(super) struct MemPlanSnapshot { free_slots: FreeSlotPools, @@ -95,9 +110,12 @@ impl<'a> MemPlan<'a> { } fn must_use_object_storage(&self, v: ValueId) -> bool { - self.scratch_spill_slots == 0 - || self.scratch_live_values.contains(v) - || self.forced_object_spills.contains(v) + must_use_object_storage( + self.scratch_spill_slots, + self.scratch_live_values, + self.forced_object_spills, + v, + ) } fn request_object_storage(&mut self, v: ValueId) { diff --git a/crates/codegen/src/stackalloc/stackify/planner/normalize.rs b/crates/codegen/src/stackalloc/stackify/planner/normalize.rs index 87b796dd3..3d406a12e 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/normalize.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/normalize.rs @@ -86,7 +86,7 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { } fn try_normalize_to_exact(&mut self, desired: &[ValueId]) -> bool { - if !self.ctx.search_profile.use_exact_normalize() { + if !self.ctx.search_profile.budgets().exact_normalize { return false; } @@ -107,7 +107,7 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { swap_max: self.ctx.reach.swap_max, // Bound intermediate length to the stack reach window (+ optional slack). max_len: self.ctx.reach.swap_max, - max_expansions: self.ctx.search_profile.exact_expansions(), + max_expansions: self.ctx.search_profile.budgets().exact_expansions, }; let mut plan = diff --git a/crates/codegen/src/stackalloc/stackify/planner/normalize_search.rs b/crates/codegen/src/stackalloc/stackify/planner/normalize_search.rs index 5213d9b6a..c596b45f5 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/normalize_search.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/normalize_search.rs @@ -1194,8 +1194,8 @@ pub(super) fn solve_optimal_normalize_plan( push0_cost, dup_goal_keys_only: false, max_states: [ - ctx.search_profile.normalize_max_states(false), - ctx.search_profile.normalize_max_states(true), + ctx.search_profile.budgets().normalize_max_states(false), + ctx.search_profile.budgets().normalize_max_states(true), ], }; @@ -1467,8 +1467,8 @@ pub(super) fn solve_optimal_repair_prefix_plan( push0_cost, dup_goal_keys_only: true, max_states: [ - ctx.search_profile.normalize_max_states(false), - ctx.search_profile.normalize_max_states(true), + ctx.search_profile.budgets().normalize_max_states(false), + ctx.search_profile.budgets().normalize_max_states(true), ], }; @@ -2580,9 +2580,11 @@ pub(super) fn solve_greedy_operand_prep_plan( cost, cfg, linear_cost, - ctx.search_profile.operand_prep_beam_width(), + ctx.search_profile.budgets().operand_prep_beam_width, ctx.search_profile - .operand_prep_beam_depth_slack(cfg.swap_max), + .budgets() + .operand_prep_beam_depth_slack + .resolve(cfg.swap_max), ) .filter(|(_, greedy_cost)| *greedy_cost < linear_cost) .unwrap_or((linear_steps, linear_cost)); @@ -2699,9 +2701,11 @@ pub(super) fn solve_optimal_operand_prep_plan( cost, cfg, upper_bound, - ctx.search_profile.operand_prep_beam_width(), + ctx.search_profile.budgets().operand_prep_beam_width, ctx.search_profile - .operand_prep_beam_depth_slack(cfg.swap_max), + .budgets() + .operand_prep_beam_depth_slack + .resolve(cfg.swap_max), ) && greedy_cost < upper_bound { if debug { @@ -2759,7 +2763,7 @@ pub(super) fn solve_optimal_operand_prep_plan( surplus_last_use_penalty, push0_kid, push0_cost, - max_states: ctx.search_profile.operand_prep_max_states(), + max_states: ctx.search_profile.budgets().operand_prep_max_states, }; let run = run_bounded_astar( @@ -5496,9 +5500,11 @@ func public %f() { &cost, search_cfg, linear_cost, - ctx.search_profile.operand_prep_beam_width(), + ctx.search_profile.budgets().operand_prep_beam_width, ctx.search_profile - .operand_prep_beam_depth_slack(search_cfg.swap_max), + .budgets() + .operand_prep_beam_depth_slack + .resolve(search_cfg.swap_max), ) .expect("expected greedy incumbent"); assert!( diff --git a/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs b/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs index f5c88529b..cfc13d876 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs @@ -338,7 +338,7 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { return Some(plan); } - if !self.ctx.search_profile.use_exact_operand_prep() { + if !self.ctx.search_profile.budgets().exact_operand_prep { return solve_greedy_operand_prep_plan( self.ctx, self.stack, @@ -805,7 +805,11 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { dup_max: self.ctx.reach.dup_max, swap_max: self.ctx.reach.swap_max, max_len, - max_expansions: self.ctx.search_profile.operand_prep_exact_expansions(), + max_expansions: self + .ctx + .search_profile + .budgets() + .operand_prep_exact_expansions, } } diff --git a/crates/codegen/src/stackalloc/stackify/sym_stack.rs b/crates/codegen/src/stackalloc/stackify/sym_stack.rs index 285b73c3a..351457c5c 100644 --- a/crates/codegen/src/stackalloc/stackify/sym_stack.rs +++ b/crates/codegen/src/stackalloc/stackify/sym_stack.rs @@ -165,6 +165,8 @@ impl SymStack { } /// Delete `stack[depth-1]` (1-indexed), preserving the relative order of the remaining items. + /// + /// Emits a `SWAP` chain, so its cost is O(depth). pub(super) fn stable_delete_at_depth(&mut self, depth: usize, actions: &mut Actions) { assert!( (1..=super::SWAP_WINDOW_MAX).contains(&depth), @@ -203,6 +205,9 @@ impl SymStack { self.items.swap(0, depth); } + /// Rotate `stack[pos]` up to the top, preserving the relative order of the items below it. + /// + /// Emits a `SWAP` chain, so its cost is O(pos). pub(super) fn stable_rotate_to_top(&mut self, pos: usize, actions: &mut Actions) { if pos == 0 { return; diff --git a/crates/codegen/src/stackalloc/stackify/templates.rs b/crates/codegen/src/stackalloc/stackify/templates.rs index fde7a006f..77e57c128 100644 --- a/crates/codegen/src/stackalloc/stackify/templates.rs +++ b/crates/codegen/src/stackalloc/stackify/templates.rs @@ -112,13 +112,15 @@ pub(super) fn choose_transfer( return first.clone(); } - if let Some((_pred, cand)) = candidates - .iter() - .filter(|(pred, _)| ctx.dom.dominates(block, *pred)) - .min_by_key(|(pred, _)| pred.as_u32()) - { - return cand.clone(); - } + // Candidates are predecessors already planned before `block` in dominator-tree RPO. A + // dominator always precedes its dominated nodes in RPO, so an already-planned predecessor + // can never be dominated by `block`; a `dominates(block, pred)` candidate is impossible here. + debug_assert!( + candidates + .iter() + .all(|(pred, _)| !ctx.dom.dominates(block, *pred)), + "merge predecessor dominated by its own merge block" + ); candidates .iter() diff --git a/crates/codegen/src/stackalloc/stackify/terminal_chain.rs b/crates/codegen/src/stackalloc/stackify/terminal_chain.rs index 534565a5a..a68082aa4 100644 --- a/crates/codegen/src/stackalloc/stackify/terminal_chain.rs +++ b/crates/codegen/src/stackalloc/stackify/terminal_chain.rs @@ -11,8 +11,8 @@ use super::{builder::StackifyContext, templates::BlockInterfaces}; pub(super) fn compute_terminal_chain_blocks( ctx: &StackifyContext<'_>, interfaces: &BlockInterfaces, -) -> SecondaryMap { - let mut terminal_chain_blocks = SecondaryMap::new(); +) -> BitSet { + let mut terminal_chain_blocks = BitSet::default(); for &block in ctx.dom.rpo().iter().rev() { if !ctx.dom.is_reachable(block) @@ -35,7 +35,7 @@ pub(super) fn compute_terminal_chain_blocks( let is = ctx.func.inst_set(); let inst = ctx.func.dfg.inst(term); - terminal_chain_blocks[block] = + let is_terminal = <&evm::EvmRevert as InstDowncast>::downcast(is, inst).is_some() || <&evm::EvmReturn as InstDowncast>::downcast(is, inst).is_some() || <&evm::EvmStop as InstDowncast>::downcast(is, inst).is_some() @@ -49,9 +49,12 @@ pub(super) fn compute_terminal_chain_blocks( .dests() .iter() .copied() - .all(|dest| terminal_chain_blocks[dest]), + .all(|dest| terminal_chain_blocks.contains(dest)), } }); + if is_terminal { + terminal_chain_blocks.insert(block); + } } terminal_chain_blocks From f8d871795150f714f27dcb7676d8d849a4b2a604 Mon Sep 17 00:00:00 2001 From: sbillig Date: Mon, 6 Jul 2026 23:19:30 -0700 Subject: [PATCH 04/14] Close the StackifyEdgeSplitter pipeline gap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The machine pipeline ran only CriticalEdgeSplitter, so a multiway self-loop on the entry block tripped a debug assert (or silently dropped the edge fixup in release). Run StackifyEdgeSplitter in prepare_machine_stackify_analysis instead, and: - Fix a latent CfgEditor::split_edge bug: inserting the split block before an entry-block destination silently made it the new entry, recreating the multiway-edge-to-planned-block problem. - Narrow the splitter from all in-cycle multiway edges to retreating ones (dom-RPO rank[to] <= rank[from]) — exactly the planner's assert condition — so existing corpus bytecode is byte-identical. - Move stackify_edge.rs to stackalloc/edge_split.rs; document the precondition; upgrade the planner guards to release asserts. - Add an end-to-end regression test for the entry self-loop shape. --- crates/codegen/src/cfg_edit.rs | 15 ++- crates/codegen/src/isa/evm/machine/prepare.rs | 9 +- .../codegen/src/isa/evm/static_arena_alloc.rs | 5 +- crates/codegen/src/isa/evm/tests.rs | 47 +++++++- crates/codegen/src/lib.rs | 1 - crates/codegen/src/stackalloc/edge_split.rs | 111 ++++++++++++++++++ crates/codegen/src/stackalloc/mod.rs | 2 + .../src/stackalloc/stackify/iteration.rs | 10 +- crates/codegen/src/stackalloc/stackify/mod.rs | 8 +- crates/codegen/src/stackify_edge.rs | 88 -------------- 10 files changed, 188 insertions(+), 108 deletions(-) create mode 100644 crates/codegen/src/stackalloc/edge_split.rs delete mode 100644 crates/codegen/src/stackify_edge.rs diff --git a/crates/codegen/src/cfg_edit.rs b/crates/codegen/src/cfg_edit.rs index 9ebc18058..30350057e 100644 --- a/crates/codegen/src/cfg_edit.rs +++ b/crates/codegen/src/cfg_edit.rs @@ -347,9 +347,18 @@ impl<'f> CfgEditor<'f> { ); let mid = self.func.dfg.make_block(); - let mut cursor = InstInserter::at_location(CursorLocation::BlockTop(to)); - cursor.insert_block_before(self.func, mid); - cursor.set_location(CursorLocation::BlockTop(mid)); + if self.func.layout.entry_block() == Some(to) { + // Splitting an edge whose destination is the entry block (e.g. a multiway self-loop + // on the entry): inserting `mid` immediately before `to` would make `mid` the new + // entry, because `Layout::insert_block_before` reassigns `entry_block` when `before` + // has no predecessor in the layout. Keep `to` as the entry by placing `mid` after it; + // layout order past the entry is only a fallthrough hint and is recomputed by later + // analyses. + self.func.layout.insert_block_after(mid, to); + } else { + self.func.layout.insert_block_before(mid, to); + } + let mut cursor = InstInserter::at_location(CursorLocation::BlockTop(mid)); cursor.append_inst_data(self.func, Jump::new(self.func.dfg.inst_set().jump(), to)); self.func.dfg.rewrite_branch_edges_to_block(term, to, mid); diff --git a/crates/codegen/src/isa/evm/machine/prepare.rs b/crates/codegen/src/isa/evm/machine/prepare.rs index b1ceaa5f2..520a062f8 100644 --- a/crates/codegen/src/isa/evm/machine/prepare.rs +++ b/crates/codegen/src/isa/evm/machine/prepare.rs @@ -4,13 +4,12 @@ use sonatina_ir::{Module, cfg::ControlFlowGraph, isa::evm::EvmMachine, module::F use tracing::{debug_span, trace_span}; use crate::{ - critical_edge::CriticalEdgeSplitter, domtree::DomTree, liveness::{InstLiveness, Liveness}, module_analysis::CallGraphSchedule, stackalloc::{ HOT_IMMEDIATE_SIZE_MIN_BLOCK_USES, HOT_IMMEDIATE_SIZE_MIN_MATERIALIZATION_BYTES, - StackifyBuilder, + StackifyBuilder, StackifyEdgeSplitter, }, }; @@ -101,8 +100,10 @@ fn prepare_machine_stackify_analysis( let mut cfg = ControlFlowGraph::new(); cfg.compute(function); - let mut splitter = CriticalEdgeSplitter::new(); - splitter.run(function, &mut cfg); + // `StackifyEdgeSplitter` subsumes critical-edge splitting and additionally splits in-cycle + // multiway edges, establishing stackify's edge preconditions (see `stackalloc::stackify`). + // Split first, then compute domtree/liveness over the post-split CFG below. + StackifyEdgeSplitter::run(function, &mut cfg); let mut dom = DomTree::new(); dom.compute(&cfg); diff --git a/crates/codegen/src/isa/evm/static_arena_alloc.rs b/crates/codegen/src/isa/evm/static_arena_alloc.rs index afc59d4f3..6613e04a4 100644 --- a/crates/codegen/src/isa/evm/static_arena_alloc.rs +++ b/crates/codegen/src/isa/evm/static_arena_alloc.rs @@ -739,11 +739,10 @@ fn render_alloca_escapes( mod tests { use super::*; use crate::{ - critical_edge::CriticalEdgeSplitter, domtree::DomTree, isa::evm::{EvmBackend, canonicalize_alias_value}, liveness::Liveness, - stackalloc::StackifyBuilder, + stackalloc::{StackifyBuilder, StackifyEdgeSplitter}, }; use sonatina_parser::{ParsedModule, parse_module}; use sonatina_triple::{Architecture, EvmVersion, OperatingSystem, TargetTriple, Vendor}; @@ -803,7 +802,7 @@ mod tests { parsed.module.func_store.modify(func_ref, |function| { let mut cfg = sonatina_ir::cfg::ControlFlowGraph::new(); cfg.compute(function); - CriticalEdgeSplitter::new().run(function, &mut cfg); + StackifyEdgeSplitter::run(function, &mut cfg); let mut liveness = Liveness::new(); liveness.compute(function, &cfg); diff --git a/crates/codegen/src/isa/evm/tests.rs b/crates/codegen/src/isa/evm/tests.rs index c627385a2..67a7be082 100644 --- a/crates/codegen/src/isa/evm/tests.rs +++ b/crates/codegen/src/isa/evm/tests.rs @@ -10,7 +10,9 @@ use crate::{ }, object::{CompileOptions, SymbolId, link::link_section}, optim::pipeline::Pipeline, - stackalloc::{Action, Actions, Allocator, StackifyAlloc, StackifyBuilder}, + stackalloc::{ + Action, Actions, Allocator, StackifyAlloc, StackifyBuilder, StackifyEdgeSplitter, + }, }; use rustc_hash::{FxHashMap, FxHashSet}; use smallvec::{SmallVec, smallvec}; @@ -135,7 +137,7 @@ fn compute_test_stackify_alloc(function: &mut Function) -> StackifyAlloc { let mut cfg = ControlFlowGraph::new(); cfg.compute(function); - CriticalEdgeSplitter::new().run(function, &mut cfg); + StackifyEdgeSplitter::run(function, &mut cfg); let mut liveness = Liveness::new(); liveness.compute(function, &cfg); @@ -2766,6 +2768,47 @@ object @Contract { ); } +#[test] +fn machine_pipeline_compiles_entry_self_loop() { + // Regression for the `StackifyEdgeSplitter` pipeline gap: a multiway self-loop on the entry + // block (`br v0 block0 block1`) is an in-cycle multiway edge that is *not* critical (block0's + // only predecessor is itself), so plain critical-edge splitting leaves it intact. Stackify + // then plans a branch edge back to the already-planned entry block, which trips the guard in + // `on_branch_edge` (previously a debug assert / silent miscompile). The machine pipeline runs + // `StackifyEdgeSplitter`, which splits the edge and makes the shape compile. + let parsed = parse_module( + r#" +target = "evm-ethereum-osaka" + +func public %f(v0.i1, v1.i256) { +block0: + mstore v1 v1 i256; + br v0 block0 block1; + +block1: + evm_return 0.i256 0.i256; +} + +object @Contract { + section runtime { + entry %f; + } +} +"#, + ) + .unwrap(); + + let func = parsed.module.funcs()[0]; + let backend = test_backend(); + let prepared = backend + .prepare_section(work_module(&parsed.module, &[func])) + .expect("prepare should succeed for an entry self-loop"); + + backend + .lower_function(&prepared, func) + .expect("lowering an entry self-loop should succeed"); +} + #[test] fn prepared_section_is_reusable_across_lower_calls() { let parsed = parse_module( diff --git a/crates/codegen/src/lib.rs b/crates/codegen/src/lib.rs index 6654fbe13..08b5d2365 100644 --- a/crates/codegen/src/lib.rs +++ b/crates/codegen/src/lib.rs @@ -17,6 +17,5 @@ pub mod optim; pub mod post_domtree; pub mod range_analysis; pub mod stackalloc; -pub mod stackify_edge; pub mod transform; pub(crate) mod type_rewrite; diff --git a/crates/codegen/src/stackalloc/edge_split.rs b/crates/codegen/src/stackalloc/edge_split.rs new file mode 100644 index 000000000..309950742 --- /dev/null +++ b/crates/codegen/src/stackalloc/edge_split.rs @@ -0,0 +1,111 @@ +use std::collections::BTreeSet; + +use cranelift_entity::SecondaryMap; +use sonatina_ir::{BlockId, ControlFlowGraph, Function}; + +use crate::{ + cfg_edit::{CfgEditor, CleanupMode}, + critical_edge::CriticalEdgeSplitter, + domtree::DomTree, +}; + +/// Establishes the edge preconditions of the stackify allocator (`stackalloc::stackify`). +/// +/// It runs [`CriticalEdgeSplitter`] and, in addition, splits every *multiway* edge whose target +/// is planned no later than the branching block itself. +/// +/// Stackify plans blocks in dominator-tree RPO, threading a symbolic stack forward, and marks a +/// block planned before simulating its terminator. A multiway terminator whose target is a block +/// already planned at that point (a retreating edge in planning order: a self-loop, or a backedge +/// into a dominating block such as a multiway self-loop on the entry block) would have to fix up +/// the stack into an already-planned block, which the planner cannot express — it either asserts +/// or silently drops the fixup. Splitting the edge inserts a single-jump block that carries the +/// fixup instead. +/// +/// Forward multiway edges — the common in-loop conditional branch, whose target is planned after +/// the branch — are handled by the planner directly and are deliberately left intact, so block +/// layout (and emitted bytecode) is unchanged for functions without a retreating multiway edge. +pub struct StackifyEdgeSplitter; + +impl StackifyEdgeSplitter { + pub fn run(func: &mut Function, cfg: &mut ControlFlowGraph) { + CriticalEdgeSplitter::new().run(func, cfg); + + // Rank reachable blocks by stackify's planning order (dominator-tree RPO). + let mut dom = DomTree::new(); + dom.compute(cfg); + let mut plan_rank: SecondaryMap> = SecondaryMap::default(); + for (rank, &block) in dom.rpo().iter().enumerate() { + plan_rank[block] = Some(rank as u32); + } + + let mut edges = BTreeSet::<(BlockId, BlockId)>::new(); + for from in func.layout.iter_block() { + if cfg.succ_num_of(from) < 2 { + continue; + } + let Some(from_rank) = plan_rank[from] else { + continue; + }; + for &to in cfg.succs_of(from) { + // Retreating edge: `to` is planned before (backedge) or together with (self-loop) + // `from`, so `to` is already planned when `from`'s terminator is simulated. + if plan_rank[to].is_some_and(|to_rank| to_rank <= from_rank) { + edges.insert((from, to)); + } + } + } + + if edges.is_empty() { + return; + } + + let mut editor = CfgEditor::new(func, CleanupMode::Strict); + for (from, to) in edges { + editor.split_edge(from, to); + } + cfg.compute(editor.func()); + } +} + +#[cfg(test)] +mod tests { + use sonatina_ir::cfg::ControlFlowGraph; + use sonatina_parser::parse_module; + + use super::StackifyEdgeSplitter; + + #[test] + fn splits_multiway_self_loop_edge() { + const SRC: &str = r#" +target = "evm-ethereum-osaka" + +func public %f(v0.i1) { +block0: + br v0 block0 block1; + +block1: + return; +} +"#; + + let parsed = parse_module(SRC).expect("module parses"); + let func = parsed.module.funcs()[0]; + + parsed.module.func_store.modify(func, |function| { + let mut cfg = ControlFlowGraph::new(); + cfg.compute(function); + let entry = cfg.entry().expect("missing entry"); + assert!(cfg.succs_of(entry).any(|&succ| succ == entry)); + + StackifyEdgeSplitter::run(function, &mut cfg); + + // The split block must not displace the entry: stackify plans the entry first, so a + // split block inserted before it would reintroduce a multiway edge into an + // already-planned block. + assert_eq!(cfg.entry(), Some(entry)); + assert!(!cfg.succs_of(entry).any(|&succ| succ == entry)); + assert!(cfg.preds_of(entry).any(|&pred| cfg.succ_num_of(pred) == 1)); + }); + } +} diff --git a/crates/codegen/src/stackalloc/mod.rs b/crates/codegen/src/stackalloc/mod.rs index 2e0c7a4ac..0c81b637c 100644 --- a/crates/codegen/src/stackalloc/mod.rs +++ b/crates/codegen/src/stackalloc/mod.rs @@ -5,7 +5,9 @@ use sonatina_ir::{Function, Immediate, InstId, ValueId}; use crate::isa::evm::static_arena_alloc::StackObjId; +mod edge_split; mod stackify; +pub use edge_split::StackifyEdgeSplitter; pub(crate) use stackify::{ HOT_IMMEDIATE_SIZE_MIN_BLOCK_USES, HOT_IMMEDIATE_SIZE_MIN_MATERIALIZATION_BYTES, StackifyTrace, }; diff --git a/crates/codegen/src/stackalloc/stackify/iteration.rs b/crates/codegen/src/stackalloc/stackify/iteration.rs index af5aeb909..677ecfc74 100644 --- a/crates/codegen/src/stackalloc/stackify/iteration.rs +++ b/crates/codegen/src/stackalloc/stackify/iteration.rs @@ -517,9 +517,10 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { succ: BlockId, stack: SymStack, ) { - debug_assert!( + assert!( !self.planned_blocks.contains(succ), - "branch edge to planned block {succ:?} requires a split edge" + "multiway branch edge to already-planned block {succ:?}: run StackifyEdgeSplitter \ + before stackify to split in-cycle multiway edges" ); debug_assert_eq!( self.ctx.cfg.pred_num_of(succ), @@ -537,9 +538,10 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { succ: BlockId, stack: SymStack, ) { - debug_assert!( + assert!( !self.planned_blocks.contains(succ), - "br_table edge to planned block {succ:?} requires a split edge" + "multiway br_table edge to already-planned block {succ:?}: run StackifyEdgeSplitter \ + before stackify to split in-cycle multiway edges" ); debug_assert_eq!( self.ctx.cfg.pred_num_of(succ), diff --git a/crates/codegen/src/stackalloc/stackify/mod.rs b/crates/codegen/src/stackalloc/stackify/mod.rs index c7401ca61..e1c317cf8 100644 --- a/crates/codegen/src/stackalloc/stackify/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/mod.rs @@ -20,7 +20,10 @@ //! edge stores can be emitted directly without first staging every phi source on the stack. //! //! Notes specific to this codebase: -//! - Critical edges must be split before running this allocator. +//! - Run `StackifyEdgeSplitter` before this allocator: it establishes both split +//! preconditions, splitting critical edges *and* every multiway edge whose target is already +//! planned when the branch is reached (a self-loop or backedge, e.g. a multiway self-loop on +//! the entry block). Splitting only critical edges is not enough. //! - Internal calls rely on an implicit return address value on the EVM stack. //! The allocator models this as a special stack item barrier to avoid popping //! into the caller's preserved stack segment. @@ -63,8 +66,7 @@ mod tests { analysis::func_behavior::analyze_module, domtree::DomTree, liveness::{InstLiveness, Liveness}, - stackalloc::{Action, Allocator, canonicalize_value_alias}, - stackify_edge::StackifyEdgeSplitter, + stackalloc::{Action, Allocator, StackifyEdgeSplitter, canonicalize_value_alias}, }; use cranelift_entity::SecondaryMap; use sonatina_ir::{ diff --git a/crates/codegen/src/stackify_edge.rs b/crates/codegen/src/stackify_edge.rs deleted file mode 100644 index c2a50534e..000000000 --- a/crates/codegen/src/stackify_edge.rs +++ /dev/null @@ -1,88 +0,0 @@ -use std::collections::BTreeSet; - -use sonatina_ir::{BlockId, ControlFlowGraph, Function}; - -use crate::{ - cfg_edit::{CfgEditor, CleanupMode}, - cfg_scc::CfgSccAnalysis, - critical_edge::CriticalEdgeSplitter, -}; - -pub struct StackifyEdgeSplitter; - -impl StackifyEdgeSplitter { - pub fn run(func: &mut Function, cfg: &mut ControlFlowGraph) { - CriticalEdgeSplitter::new().run(func, cfg); - - let mut scc = CfgSccAnalysis::new(); - scc.compute(cfg); - - let mut edges = BTreeSet::<(BlockId, BlockId)>::new(); - for from in func.layout.iter_block() { - if cfg.succ_num_of(from) < 2 { - continue; - } - - let Some(from_scc) = scc.scc_of(from) else { - continue; - }; - if !scc.scc_data(from_scc).is_cycle { - continue; - } - - for &to in cfg.succs_of(from) { - if scc.scc_of(to) == Some(from_scc) { - edges.insert((from, to)); - } - } - } - - if edges.is_empty() { - return; - } - - let mut editor = CfgEditor::new(func, CleanupMode::Strict); - for (from, to) in edges { - editor.split_edge(from, to); - } - cfg.compute(editor.func()); - } -} - -#[cfg(test)] -mod tests { - use sonatina_ir::cfg::ControlFlowGraph; - use sonatina_parser::parse_module; - - use super::StackifyEdgeSplitter; - - #[test] - fn splits_multiway_self_loop_edge() { - const SRC: &str = r#" -target = "evm-ethereum-osaka" - -func public %f(v0.i1) { -block0: - br v0 block0 block1; - -block1: - return; -} -"#; - - let parsed = parse_module(SRC).expect("module parses"); - let func = parsed.module.funcs()[0]; - - parsed.module.func_store.modify(func, |function| { - let mut cfg = ControlFlowGraph::new(); - cfg.compute(function); - let entry = cfg.entry().expect("missing entry"); - assert!(cfg.succs_of(entry).any(|&succ| succ == entry)); - - StackifyEdgeSplitter::run(function, &mut cfg); - - assert!(!cfg.succs_of(entry).any(|&succ| succ == entry)); - assert!(cfg.preds_of(entry).any(|&pred| cfg.succ_num_of(pred) == 1)); - }); - } -} From 74d971a29d20bf47b65af72c262de98c210f2d7c Mon Sep 17 00:00:00 2001 From: sbillig Date: Mon, 6 Jul 2026 23:35:19 -0700 Subject: [PATCH 05/14] Replace stackify trace placeholders with an event log StackifyTrace records structured TraceEvents instead of interleaving placeholder markers into one string and substituting at render time. Rendering is a single forward walk; checkpoint/rollback is a vec truncate; deferred-exit actions fill in by event index; object-id remapping rewrites raw Action payloads in events. The observer trait loses both associated types. Rendered output is byte-identical. --- .../src/stackalloc/stackify/iteration.rs | 12 +- .../codegen/src/stackalloc/stackify/trace.rs | 274 ++++++++---------- 2 files changed, 127 insertions(+), 159 deletions(-) diff --git a/crates/codegen/src/stackalloc/stackify/iteration.rs b/crates/codegen/src/stackalloc/stackify/iteration.rs index 677ecfc74..c50f6ee7c 100644 --- a/crates/codegen/src/stackalloc/stackify/iteration.rs +++ b/crates/codegen/src/stackalloc/stackify/iteration.rs @@ -34,19 +34,21 @@ pub(super) struct IterationPlanner<'a, 'ctx, O: StackifyObserver> { object_spill_requests: &'a mut BitSet, forced_object_spills: &'a BitSet, inherited_stack: BTreeMap, - pending_edges: BTreeMap>>, + pending_edges: BTreeMap>, planned_blocks: BitSet, search_scratch: &'a mut NormalizeSearchScratch, observer: &'a mut O, } -struct PendingEdge { +struct PendingEdge { pred: BlockId, inst: InstId, stack: SymStack, free_slots: FreeSlotPools, action_start: usize, - deferred_exit: D, + /// Index of this edge's `DeferredExit` event in the observer's trace (0 for `NullObserver`), + /// used to backfill the exit fixup actions once the merge template is resolved. + trace_token: usize, } #[derive(Clone, Copy)] @@ -250,7 +252,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { |planner| planner.plan_edge_fixup_to_template(&tmpl, edge.pred, block), ); self.observer.on_deferred_exit_actions( - edge.deferred_exit, + edge.trace_token, &self.alloc.pre_actions[edge.inst][edge.action_start..], ); } @@ -478,7 +480,7 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { stack: state.stack.clone(), free_slots: state.free_slots.clone(), action_start, - deferred_exit: self.observer.on_deferred_inst_jump(inst, dest), + trace_token: self.observer.on_deferred_inst_jump(inst, dest), }); return; } else if self.ctx.cfg.pred_num_of(dest) == 1 diff --git a/crates/codegen/src/stackalloc/stackify/trace.rs b/crates/codegen/src/stackalloc/stackify/trace.rs index fdb6e2e1a..41404814c 100644 --- a/crates/codegen/src/stackalloc/stackify/trace.rs +++ b/crates/codegen/src/stackalloc/stackify/trace.rs @@ -27,11 +27,8 @@ use super::{ /// `StackifyBuilder::compute_with_trace` is allowed to run multiple fixed-point iterations. /// Observers must support `checkpoint`/`rollback` so unsuccessful iterations can be discarded. pub(super) trait StackifyObserver { - type Checkpoint: Copy; - type DeferredExit: Copy + Default; - - fn checkpoint(&mut self) -> Self::Checkpoint; - fn rollback(&mut self, checkpoint: Self::Checkpoint); + fn checkpoint(&mut self) -> usize; + fn rollback(&mut self, checkpoint: usize); fn on_block_header(&mut self, _func: &Function, _block: BlockId, _template: &BlockTemplate) {} @@ -78,12 +75,12 @@ pub(super) trait StackifyObserver { fn on_inst_jump(&mut self, _inst: InstId, _dest: BlockId) {} - fn on_deferred_inst_jump(&mut self, inst: InstId, dest: BlockId) -> Self::DeferredExit { + fn on_deferred_inst_jump(&mut self, inst: InstId, dest: BlockId) -> usize { self.on_inst_jump(inst, dest); - Self::DeferredExit::default() + 0 } - fn on_deferred_exit_actions(&mut self, _deferred: Self::DeferredExit, _actions: &[Action]) {} + fn on_deferred_exit_actions(&mut self, _deferred: usize, _actions: &[Action]) {} fn on_inst_br(&mut self, _func: &Function, _inst: InstId, _cond: ValueId, _dests: &[BlockId]) {} @@ -95,12 +92,11 @@ pub(super) trait StackifyObserver { pub(super) struct NullObserver; impl StackifyObserver for NullObserver { - type Checkpoint = (); - type DeferredExit = (); - - fn checkpoint(&mut self) -> Self::Checkpoint {} + fn checkpoint(&mut self) -> usize { + 0 + } - fn rollback(&mut self, _checkpoint: Self::Checkpoint) {} + fn rollback(&mut self, _checkpoint: usize) {} } /// A snapshot-oriented trace collector for stackify planning. @@ -110,26 +106,30 @@ impl StackifyObserver for NullObserver { /// such as cleanup, prelude, exit normalization, and internal returns. #[derive(Default)] pub(crate) struct StackifyTrace { - out: String, - action_chunks: Vec, - inst_chunks: Vec, - deferred_exit_chunks: Vec, -} - -struct ActionChunk { - placeholder: String, - actions: crate::stackalloc::Actions, -} - -struct InstChunk { - placeholder: String, - inst: InstId, + events: Vec, } -struct DeferredExitChunk { - placeholder: String, - dest: BlockId, - actions: crate::stackalloc::Actions, +/// One recorded trace line (or line group). Everything that can be formatted at record time is +/// pre-rendered into a `Line`; the two things that cannot are kept structured: +/// - instruction comments need a `FuncWriteCtx` only available at render time (`InstHeader`); +/// - action payloads (`Actions`/`DeferredExit`) are stored raw so downstream object-id remapping +/// can rewrite them before render, and deferred-exit actions are filled in later by index. +enum TraceEvent { + /// Fully pre-rendered text, including any trailing newline. + Line(String), + /// `// ` line (rendered at the end) followed by `stack_line` (pre-rendered). + InstHeader { inst: InstId, stack_line: String }, + /// An action group rendered as `{prefix}{actions}\n`. Only recorded for non-empty groups. + Actions { + prefix: String, + actions: crate::stackalloc::Actions, + }, + /// A deferred exit whose `actions` are filled in when its edge resolves; renders as + /// ` exit({dest:?}): {actions}\n`, or nothing when no actions were added. + DeferredExit { + dest: BlockId, + actions: crate::stackalloc::Actions, + }, } impl StackifyTrace { @@ -144,11 +144,13 @@ impl StackifyTrace { } } - for chunk in &mut self.action_chunks { - remap_actions(&mut chunk.actions, remap); - } - for chunk in &mut self.deferred_exit_chunks { - remap_actions(&mut chunk.actions, remap); + for event in &mut self.events { + match event { + TraceEvent::Actions { actions, .. } | TraceEvent::DeferredExit { actions, .. } => { + remap_actions(actions, remap); + } + TraceEvent::Line(_) | TraceEvent::InstHeader { .. } => {} + } } } @@ -197,94 +199,56 @@ impl StackifyTrace { } let _ = writeln!(&mut out, "trace:"); - let mut trace = self.out; - for chunk in self.action_chunks { - let formatted = fmt_actions(&chunk.actions); - trace = trace.replace(&chunk.placeholder, &formatted); - } - for chunk in self.deferred_exit_chunks { - let formatted = if chunk.actions.is_empty() { - String::new() - } else { - format!( - " exit({:?}): {}\n", - chunk.dest, - fmt_actions(&chunk.actions) - ) - }; - trace = trace.replace(&chunk.placeholder, &formatted); - } - for chunk in self.inst_chunks { - let comment = fmt_inst_comment(func, chunk.inst); - trace = trace.replace(&chunk.placeholder, &comment); + for event in &self.events { + match event { + TraceEvent::Line(line) => out.push_str(line), + TraceEvent::InstHeader { inst, stack_line } => { + let _ = writeln!(&mut out, " // {}", fmt_inst_comment(func, *inst)); + out.push_str(stack_line); + } + TraceEvent::Actions { prefix, actions } => { + let _ = writeln!(&mut out, "{prefix}{}", fmt_actions(actions)); + } + TraceEvent::DeferredExit { dest, actions } => { + if !actions.is_empty() { + let _ = + writeln!(&mut out, " exit({dest:?}): {}", fmt_actions(actions)); + } + } + } } - out.push_str(&trace); out } - fn push_actions_placeholder(&mut self, actions: &[Action]) -> String { - let idx = self.action_chunks.len(); - let placeholder = format!("@@ACTIONS:{idx}@@"); + fn push_line(&mut self, line: String) { + self.events.push(TraceEvent::Line(line)); + } + + fn push_actions(&mut self, prefix: String, actions: &[Action]) { let mut stored = crate::stackalloc::Actions::new(); stored.extend_from_slice(actions); - self.action_chunks.push(ActionChunk { - placeholder: placeholder.clone(), + self.events.push(TraceEvent::Actions { + prefix, actions: stored, }); - placeholder - } - - fn push_inst_placeholder(&mut self, inst: InstId) -> String { - let idx = self.inst_chunks.len(); - let placeholder = format!("@@INST:{idx}@@"); - self.inst_chunks.push(InstChunk { - placeholder: placeholder.clone(), - inst, - }); - placeholder - } - - fn push_deferred_exit_placeholder(&mut self, dest: BlockId) -> usize { - let idx = self.deferred_exit_chunks.len(); - let placeholder = format!("@@DEFERRED_EXIT:{idx}@@"); - self.deferred_exit_chunks.push(DeferredExitChunk { - placeholder: placeholder.clone(), - dest, - actions: crate::stackalloc::Actions::new(), - }); - let _ = write!(&mut self.out, "{placeholder}"); - idx } } impl StackifyObserver for StackifyTrace { - type Checkpoint = (usize, usize, usize, usize); - type DeferredExit = usize; - - fn checkpoint(&mut self) -> Self::Checkpoint { - ( - self.out.len(), - self.action_chunks.len(), - self.inst_chunks.len(), - self.deferred_exit_chunks.len(), - ) + fn checkpoint(&mut self) -> usize { + self.events.len() } - fn rollback(&mut self, checkpoint: Self::Checkpoint) { - let (out_len, action_chunk_len, inst_chunk_len, deferred_exit_chunk_len) = checkpoint; - self.out.truncate(out_len); - self.action_chunks.truncate(action_chunk_len); - self.inst_chunks.truncate(inst_chunk_len); - self.deferred_exit_chunks.truncate(deferred_exit_chunk_len); + fn rollback(&mut self, checkpoint: usize) { + self.events.truncate(checkpoint); } fn on_block_header(&mut self, func: &Function, block: BlockId, template: &BlockTemplate) { - let _ = writeln!( - &mut self.out, - " {block:?} P={} T={}", + self.push_line(format!( + " {block:?} P={} T={}\n", fmt_values(func, &template.params), fmt_values(func, template.transfer()) - ); + )); } fn on_block_inherited( @@ -296,19 +260,17 @@ impl StackifyObserver for StackifyTrace { live_future: &BitSet, live_out: &BitSet, ) { - let _ = writeln!( - &mut self.out, - " inherited from {pred:?}: {}", + self.push_line(format!( + " inherited from {pred:?}: {}\n", fmt_stack(func, inherited_stack, live_future, live_out) - ); + )); } fn on_block_prologue(&mut self, actions: &[Action]) { if actions.is_empty() { return; } - let placeholder = self.push_actions_placeholder(actions); - let _ = writeln!(&mut self.out, " prologue: {placeholder}"); + self.push_actions(" prologue: ".to_string(), actions); } fn on_inst_start( @@ -320,34 +282,29 @@ impl StackifyObserver for StackifyTrace { live_out: &BitSet, last_use: &BitSet, ) { - let comment = self.push_inst_placeholder(inst); - let _ = writeln!(&mut self.out, " // {comment}"); let stack_start = fmt_stack(func, stack, live_future, live_out); let last_use_list: Vec = last_use.iter().collect(); - if last_use_list.is_empty() { - let _ = writeln!(&mut self.out, " - stack={stack_start}"); + let stack_line = if last_use_list.is_empty() { + format!(" - stack={stack_start}\n") } else { - let _ = writeln!( - &mut self.out, - " - stack={stack_start}, last_use={}", + format!( + " - stack={stack_start}, last_use={}\n", fmt_values(func, &last_use_list) - ); - } + ) + }; + self.events + .push(TraceEvent::InstHeader { inst, stack_line }); } fn on_inst_actions(&mut self, label: &'static str, actions: &[Action], dest: Option) { if actions.is_empty() { return; } - let placeholder = self.push_actions_placeholder(actions); - match (label, dest) { - ("exit", Some(dest)) => { - let _ = writeln!(&mut self.out, " exit({dest:?}): {placeholder}"); - } - _ => { - let _ = writeln!(&mut self.out, " {label}: {placeholder}"); - } - } + let prefix = match (label, dest) { + ("exit", Some(dest)) => format!(" exit({dest:?}): "), + _ => format!(" {label}: "), + }; + self.push_actions(prefix, actions); } fn on_inst_normal( @@ -358,59 +315,68 @@ impl StackifyObserver for StackifyTrace { results: &[ValueId], ) { let op_name = func.dfg.inst(inst).as_text(); - let _ = write!( - &mut self.out, - " {op_name} {}", - fmt_trace_operands(func, inst, args) - ); + let mut line = format!(" {op_name} {}", fmt_trace_operands(func, inst, args)); match results { [] => {} [result] => { - let _ = write!(&mut self.out, " -> {}", fmt_value(func, *result)); + let _ = write!(&mut line, " -> {}", fmt_value(func, *result)); } _ => { - let _ = write!(&mut self.out, " -> {}", fmt_values(func, results)); + let _ = write!(&mut line, " -> {}", fmt_values(func, results)); } } - let _ = writeln!(&mut self.out); + line.push('\n'); + self.push_line(line); } fn on_inst_jump(&mut self, _inst: InstId, dest: BlockId) { - let _ = writeln!(&mut self.out, " jump -> {dest:?}"); + self.push_line(format!(" jump -> {dest:?}\n")); } - fn on_deferred_inst_jump(&mut self, _inst: InstId, dest: BlockId) -> Self::DeferredExit { - let token = self.push_deferred_exit_placeholder(dest); - self.on_inst_jump(_inst, dest); + fn on_deferred_inst_jump(&mut self, inst: InstId, dest: BlockId) -> usize { + let token = self.events.len(); + self.events.push(TraceEvent::DeferredExit { + dest, + actions: crate::stackalloc::Actions::new(), + }); + self.on_inst_jump(inst, dest); token } - fn on_deferred_exit_actions(&mut self, deferred: Self::DeferredExit, actions: &[Action]) { - self.deferred_exit_chunks[deferred] - .actions - .extend_from_slice(actions); + fn on_deferred_exit_actions(&mut self, deferred: usize, actions: &[Action]) { + let TraceEvent::DeferredExit { + actions: stored, .. + } = &mut self.events[deferred] + else { + debug_assert!( + false, + "deferred exit token does not point at a DeferredExit event" + ); + return; + }; + stored.extend_from_slice(actions); } fn on_inst_br(&mut self, func: &Function, inst: InstId, cond: ValueId, dests: &[BlockId]) { let op_name = func.dfg.inst(inst).as_text(); - let _ = writeln!( - &mut self.out, - " {op_name} {} -> {dests:?}", + self.push_line(format!( + " {op_name} {} -> {dests:?}\n", fmt_values(func, &[cond]) - ); + )); } fn on_inst_br_table(&mut self, _inst: InstId) { - let _ = writeln!(&mut self.out, " br_table"); + self.push_line(" br_table\n".to_string()); } fn on_inst_return(&mut self, func: &Function, inst: InstId, rets: &[ValueId]) { let op_name = func.dfg.inst(inst).as_text(); - if rets.is_empty() { - let _ = writeln!(&mut self.out, " {op_name} []"); + let line = if rets.is_empty() { + format!(" {op_name} []\n") } else { - let _ = writeln!(&mut self.out, " {op_name} {}", fmt_values(func, rets)); - } + format!(" {op_name} {}\n", fmt_values(func, rets)) + }; + self.push_line(line); } } From 2d5550a633e358c7970e59ce4ef35397abcbffcc Mon Sep 17 00:00:00 2001 From: sbillig Date: Mon, 6 Jul 2026 23:51:20 -0700 Subject: [PATCH 06/14] Stop hoisting br_table base pre-actions into case 0 Emit and the lazy-frame analysis read a br_table's pre-actions via pre_inst like every other instruction; brtable_actions now hold only per-case compare preparation. This removes the take/prefix threading through PlannerActionSink. The lazy-frame planner bails to the always-active frame when br_table base pre-actions touch the frame, since emit replays their action indexes once per case. --- crates/codegen/src/isa/evm/emit/insn.rs | 2 + .../codegen/src/isa/evm/machine/lazy_frame.rs | 187 ++++++++++++++++-- .../src/stackalloc/stackify/block_sim.rs | 16 +- .../src/stackalloc/stackify/iteration.rs | 15 +- crates/codegen/src/stackalloc/stackify/mod.rs | 15 +- 5 files changed, 193 insertions(+), 42 deletions(-) diff --git a/crates/codegen/src/isa/evm/emit/insn.rs b/crates/codegen/src/isa/evm/emit/insn.rs index cba320739..5d0797e37 100644 --- a/crates/codegen/src/isa/evm/emit/insn.rs +++ b/crates/codegen/src/isa/evm/emit/insn.rs @@ -121,6 +121,8 @@ impl EvmMachineFunctionLowering<'_> { "br_table has duplicate scrutinee values" ); + emit_pre_actions(ctx, alloc.pre_inst(insn)); + for (case_idx, (_, dest)) in table.iter().enumerate() { let dest = self.canonical_block_target(*dest); self.emit_actions_for_site( diff --git a/crates/codegen/src/isa/evm/machine/lazy_frame.rs b/crates/codegen/src/isa/evm/machine/lazy_frame.rs index bb6df8979..5248042d7 100644 --- a/crates/codegen/src/isa/evm/machine/lazy_frame.rs +++ b/crates/codegen/src/isa/evm/machine/lazy_frame.rs @@ -353,14 +353,24 @@ fn compute_active_pre_insts( } } EvmMachineInstKind::BrTable(br) => { + // Emit replays `FrameSite::PreInst` from action index 0 for the base + // pre-actions *and* again for every case's compare prep, so an injection + // point planned inside the base actions would fire once per case. Bail to + // the always-active frame whenever the base actions touch the frame, exactly + // as for frame actions inside a case. + if actions_touch_frame(alloc.pre_inst(inst)) { + return None; + } + apply_actions_state( + plan, + FrameSite::PreInst(inst), + alloc.pre_inst(inst), + 0, + &mut active, + ); for (case_idx, _) in br.table().iter().enumerate() { let actions = alloc.br_table_case(inst, case_idx); - if fold_stack_actions(actions).iter().any(|action| { - matches!( - action, - Action::MemLoadFrameSlot(_) | Action::MemStoreFrameSlot(_) - ) - }) { + if actions_touch_frame(actions) { return None; } } @@ -599,14 +609,23 @@ fn collect_dep_points( } } EvmMachineInstKind::BrTable(br) => { + // See `compute_active_pre_insts`: base pre-actions that touch the frame make + // the lazy plan invalid, because emit replays their action indexes per case. + if actions_touch_frame(alloc.pre_inst(inst)) { + return None; + } + collect_action_dep_points( + &mut out, + &mut seen, + &order, + block, + FrameSite::PreInst(inst), + alloc.pre_inst(inst), + 0, + ); for (case_idx, _) in br.table().iter().enumerate() { let actions = alloc.br_table_case(inst, case_idx); - if fold_stack_actions(actions).iter().any(|action| { - matches!( - action, - Action::MemLoadFrameSlot(_) | Action::MemStoreFrameSlot(_) - ) - }) { + if actions_touch_frame(actions) { return None; } } @@ -849,6 +868,17 @@ fn collect_action_dep_points( } } +fn actions_touch_frame(actions: &[Action]) -> bool { + fold_stack_actions(actions).iter().any(|action| { + matches!( + action, + Action::MemLoadFrameSlot(_) + | Action::MemStoreFrameSlot(_) + | Action::PushFrameAddr { .. } + ) + }) +} + fn push_dep_point( out: &mut Vec, seen: &mut FxHashSet, @@ -1038,3 +1068,136 @@ fn scalar_bit_width(ty: Type, module: &sonatina_ir::module::ModuleCtx) -> Option }; Some(bits) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::stackalloc::Actions; + use cranelift_entity::SecondaryMap; + use sonatina_parser::parse_module; + + #[derive(Default)] + struct TestAlloc { + enter: Actions, + pre: SecondaryMap, + post: SecondaryMap, + cases: SecondaryMap>, + } + + impl TestAlloc { + fn for_function(function: &Function) -> Self { + let mut alloc = Self::default(); + for block in function.layout.iter_block() { + for inst in function.layout.iter_inst(block) { + let _ = &mut alloc.pre[inst]; + let _ = &mut alloc.post[inst]; + } + } + alloc + } + } + + impl Allocator for TestAlloc { + fn enter_function(&self, _function: &Function) -> Actions { + self.enter.clone() + } + + fn pre_inst(&self, inst: InstId) -> &Actions { + &self.pre[inst] + } + + fn post_inst(&self, inst: InstId) -> &Actions { + &self.post[inst] + } + + fn br_table_case(&self, inst: InstId, case_index: usize) -> &Actions { + &self.cases[inst][case_index] + } + } + + #[test] + fn br_table_frame_actions_abort_lazy_frame_dependency_collection() { + const SRC: &str = r#" +target = "evm-ethereum-osaka" + +func public %dispatch(v0.i256) -> i256 { +block0: + v1.i256 = add v0 1.i256; + br_table v0 block3 (0.i256 block1) (11.i256 block2); + +block1: + return v1; + +block2: + return v1; + +block3: + return v1; +} +"#; + + let parsed = parse_module(SRC).expect("module parses"); + let func_ref = parsed.debug.func_order[0]; + + parsed.module.func_store.view(func_ref, |function| { + let term = function + .layout + .iter_block() + .flat_map(|block| function.layout.iter_inst(block)) + .find(|&inst| function.dfg.cast_br_table(inst).is_some()) + .expect("missing br_table terminator"); + let case_count = function + .dfg + .cast_br_table(term) + .expect("br_table terminator should downcast") + .table() + .len(); + let root = parsed.debug.value(func_ref, "v1").expect("v1 exists"); + let root_def = function + .dfg + .value_inst(root) + .expect("v1 should be instruction-defined"); + let mut roots = MachineFrameRoots::default(); + roots.root_def_insts.insert(root_def); + roots.rooted_values.insert(root); + let mut cfg = ControlFlowGraph::default(); + cfg.compute(function); + + let mut clean_alloc = TestAlloc::for_function(function); + clean_alloc.cases[term] = vec![Actions::new(); case_count]; + let clean_dep_points = collect_dep_points(function, &cfg, &clean_alloc, &roots) + .expect("control br_table should allow lazy-frame dependency collection"); + assert!( + !clean_dep_points.is_empty(), + "control br_table should collect lazy-frame dependency points" + ); + + let mut base_action_alloc = TestAlloc::for_function(function); + base_action_alloc.pre[term].push(Action::MemLoadFrameSlot(0)); + base_action_alloc.cases[term] = vec![Actions::new(); case_count]; + assert!( + collect_dep_points(function, &cfg, &base_action_alloc, &roots).is_none(), + "frame-touching br_table base actions must abort dependency collection" + ); + + let mut case_action_alloc = TestAlloc::for_function(function); + case_action_alloc.cases[term] = vec![Actions::new(); case_count]; + case_action_alloc.cases[term][1].push(Action::MemStoreFrameSlot(0)); + assert!( + collect_dep_points(function, &cfg, &case_action_alloc, &roots).is_none(), + "frame-touching br_table case actions must abort dependency collection" + ); + + let mut case_address_alloc = TestAlloc::for_function(function); + case_address_alloc.cases[term] = vec![Actions::new(); case_count]; + case_address_alloc.cases[term][0].push(Action::PushFrameAddr { + offset_words: 0, + extra_bytes: 0, + }); + assert!( + collect_dep_points(function, &cfg, &case_address_alloc, &roots).is_none(), + "frame-address br_table case actions must abort dependency collection" + ); + }); + } +} diff --git a/crates/codegen/src/stackalloc/stackify/block_sim.rs b/crates/codegen/src/stackalloc/stackify/block_sim.rs index be81299ce..61c423407 100644 --- a/crates/codegen/src/stackalloc/stackify/block_sim.rs +++ b/crates/codegen/src/stackalloc/stackify/block_sim.rs @@ -80,14 +80,10 @@ impl BlockSimState { } } -pub(super) enum PlannerActionSink<'a> { +pub(super) enum PlannerActionSink { Pre(InstId), Post(InstId), - BrTableCase { - inst: InstId, - case_idx: usize, - prefix: Option<&'a Actions>, - }, + BrTableCase { inst: InstId, case_idx: usize }, } enum TerminatorInfo { @@ -319,20 +315,14 @@ pub(super) fn run_block_sim( ); }); - let base_actions = planner.take_pre_actions_for_br_table(inst); let (case_stacks, default_stack) = plan_br_table_compare_chain( &table, &state.stack, |case_idx, case_val, case_stack| { - let prefix = (case_idx == 0).then_some(&base_actions); planner.with_planner( case_stack, &mut state.free_slots, - PlannerActionSink::BrTableCase { - inst, - case_idx, - prefix, - }, + PlannerActionSink::BrTableCase { inst, case_idx }, |planner| { let consume_last_use = BitSet::::default(); let mut compare_args = smallvec::smallvec![scrutinee, case_val]; diff --git a/crates/codegen/src/stackalloc/stackify/iteration.rs b/crates/codegen/src/stackalloc/stackify/iteration.rs index c50f6ee7c..39c5ad3d4 100644 --- a/crates/codegen/src/stackalloc/stackify/iteration.rs +++ b/crates/codegen/src/stackalloc/stackify/iteration.rs @@ -310,15 +310,11 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { f(&mut self.alloc.pre_actions[inst]) } - pub(super) fn take_pre_actions_for_br_table(&mut self, inst: InstId) -> Actions { - std::mem::take(&mut self.alloc.pre_actions[inst]) - } - pub(super) fn with_planner( &mut self, stack: &mut SymStack, free_slots: &mut FreeSlotPools, - sink: PlannerActionSink<'_>, + sink: PlannerActionSink, f: impl FnOnce(&mut Planner<'_, '_>) -> R, ) -> R { match sink { @@ -364,15 +360,8 @@ impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { ); f(&mut planner) } - PlannerActionSink::BrTableCase { - inst, - case_idx, - prefix, - } => { + PlannerActionSink::BrTableCase { inst, case_idx } => { let mut actions = Actions::new(); - if let Some(prefix) = prefix { - actions.extend_from_slice(prefix); - } let mem = planner::MemPlan::new( self.spill, &mut *self.spill_requests, diff --git a/crates/codegen/src/stackalloc/stackify/mod.rs b/crates/codegen/src/stackalloc/stackify/mod.rs index e1c317cf8..30d200f76 100644 --- a/crates/codegen/src/stackalloc/stackify/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/mod.rs @@ -212,20 +212,27 @@ block3: .with_value_aliases(&value_aliases) .compute(); + // The br_table terminator's accumulated pre-actions (dead-prefix cleanup, + // reachability rescue, prologue injection) are read from `pre_actions[term]` like + // every other instruction and are no longer hoisted into case 0. For this function + // that set is empty; the stored case lists therefore hold only per-case compare + // preparation, in IR case order. assert!( alloc.pre_actions[term].is_empty(), - "br_table pre-actions should be hoisted into case actions" + "br_table pre-actions are read from pre_inst, not hoisted into case 0" ); assert_eq!(alloc.brtable_actions[term].len(), 2); let first = alloc.br_table_case(term, 0); let second = alloc.br_table_case(term, 1); + // Case 0 is exactly its compare preparation (materializing the case scrutinee), with + // no hoisted cleanup/rescue prefix in front of it. assert!( - !first.is_empty(), - "expected first br_table case to include compare preparation" + matches!(first.as_slice(), [Action::Push(_)]), + "case 0 should hold only compare preparation, not a hoisted base-action prefix" ); assert!( - !second.is_empty(), + matches!(second.first(), Some(Action::Push(_))), "expected later br_table case to include compare preparation" ); }); From 4e257be21b158ca49746669ab08ef7e5123e21dd Mon Sep 17 00:00:00 2001 From: sbillig Date: Tue, 7 Jul 2026 00:03:49 -0700 Subject: [PATCH 07/14] Encapsulate the internal-call ABI in one planner seam Planner::prepare_internal_call now owns the operand-rotation + continuation-push + single-SWAP sequence with the ABI documented in one place; the rotation site and SymStack helpers cross-reference it. remove_call_ret_addr's linear marker scan becomes pop_call_ret_addr, asserting the continuation is on top (it is, by construction). --- .../src/stackalloc/stackify/block_sim.rs | 19 ++++----- .../src/stackalloc/stackify/iteration.rs | 1 + .../stackify/planner/control_flow.rs | 29 +++++++++++++- .../src/stackalloc/stackify/sym_stack.rs | 39 +++++++++++-------- 4 files changed, 60 insertions(+), 28 deletions(-) diff --git a/crates/codegen/src/stackalloc/stackify/block_sim.rs b/crates/codegen/src/stackalloc/stackify/block_sim.rs index 61c423407..1dfc33bc5 100644 --- a/crates/codegen/src/stackalloc/stackify/block_sim.rs +++ b/crates/codegen/src/stackalloc/stackify/block_sim.rs @@ -367,19 +367,16 @@ pub(super) fn run_block_sim( &mut stack, &mut state.free_slots, PlannerActionSink::Pre(inst), - |planner| planner.prepare_operands_for_inst(inst, &mut args, last_use, &cache_preserve), + |planner| { + if call_has_stack_continuation { + planner.prepare_internal_call(inst, &mut args, last_use, &cache_preserve); + } else { + planner.prepare_operands_for_inst(inst, &mut args, last_use, &cache_preserve); + } + }, ); state.stack = stack; - if call_has_stack_continuation { - planner.with_pre_actions(inst, |actions| { - state.stack.push_call_continuation(actions); - state - .stack - .position_call_ret_below_operands(args.len(), actions); - }); - } - planner.on_pre_actions(inst, after_cleanup_len); planner.on_normal_inst(inst, &args, &results); @@ -397,7 +394,7 @@ pub(super) fn run_block_sim( state.stack.pop_n_operands(args.len()); if call_has_stack_continuation { - state.stack.remove_call_ret_addr(); + state.stack.pop_call_ret_addr(); } for &res in results.iter().rev() { diff --git a/crates/codegen/src/stackalloc/stackify/iteration.rs b/crates/codegen/src/stackalloc/stackify/iteration.rs index 39c5ad3d4..b57cfe508 100644 --- a/crates/codegen/src/stackalloc/stackify/iteration.rs +++ b/crates/codegen/src/stackalloc/stackify/iteration.rs @@ -702,6 +702,7 @@ pub(super) fn operand_order_for_stackify( .collect(); if call_has_local_return(func, inst) && !args.is_empty() { + // Rotation pairs with `Planner::prepare_internal_call`; see the ABI note there. args.as_mut_slice().rotate_left(1); } diff --git a/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs b/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs index ea35368b0..811eb76f7 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs @@ -1,7 +1,7 @@ use smallvec::SmallVec; use sonatina_ir::{BlockId, InstId, ValueId}; -use crate::stackalloc::Action; +use crate::{bitset::BitSet, stackalloc::Action}; use super::{ super::{ @@ -105,6 +105,33 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { self.mem.emit_store_for_spilled_value(phi_res, self.actions); } + /// Prepare the operands and return continuation for an internal `call`. + /// + /// EVM internal-call ABI: at the `JUMP` into the callee the stack must read + /// `[arg0, arg1, …, argN-1, cont, ]`, where `cont` is the + /// return-continuation address the callee jumps back to. This method arranges exactly that + /// shape, in three pre-coordinated steps so a single `SWAP` finishes it: + /// + /// 1. `args` arrive already rotated left by one (see `operand_order_for_stackify`), so operand + /// preparation leaves `[arg1, …, argN-1, arg0, ]` on top. + /// 2. `push_call_continuation` pushes `cont`: `[cont, arg1, …, argN-1, arg0, …]`. + /// 3. `position_call_ret_below_operands(argc)` swaps `cont` with the bottom operand `arg0`, + /// yielding `[arg0, arg1, …, argN-1, cont, …]` — ABI order with the continuation directly + /// below the args. When `argc == 0` the continuation is already on top and the swap is + /// skipped. + pub(in super::super) fn prepare_internal_call( + &mut self, + inst: InstId, + args: &mut SmallVec<[ValueId; 8]>, + consume_last_use: &BitSet, + cache_preserve: &BitSet, + ) { + self.prepare_operands_for_inst(inst, args, consume_last_use, cache_preserve); + self.stack.push_call_continuation(self.actions); + self.stack + .position_call_ret_below_operands(args.len(), self.actions); + } + pub fn plan_internal_return(&mut self, inst: InstId) { let ret_vals: SmallVec<[ValueId; 16]> = self .ctx diff --git a/crates/codegen/src/stackalloc/stackify/sym_stack.rs b/crates/codegen/src/stackalloc/stackify/sym_stack.rs index 351457c5c..fc566b01a 100644 --- a/crates/codegen/src/stackalloc/stackify/sym_stack.rs +++ b/crates/codegen/src/stackalloc/stackify/sym_stack.rs @@ -224,17 +224,20 @@ impl SymStack { self.pushed_above_barrier(); } + /// Push the internal-call return continuation onto the stack. + /// + /// One step of the internal-call ABI setup; see `Planner::prepare_internal_call` for the + /// full callee-entry stack shape and how the continuation is repositioned afterwards. pub(super) fn push_call_continuation(&mut self, actions: &mut Actions) { actions.push(Action::PushContinuationOffset); self.push_call_ret_addr(); } - /// Ensure the call continuation address sits immediately below the `operand_count` - /// call operands at the top of the stack. + /// Move the call continuation from the top of the stack to just below the `operand_count` + /// call operands, using a single `SWAP`. /// - /// This matches the EVM backend's internal-call ABI, where the callee expects its return - /// address (`FuncRetAddr`) to sit below its arguments and above any caller values that must - /// survive the call. + /// Relies on the operands having been prepared as a left-rotation of callee ABI order; see + /// `Planner::prepare_internal_call` for the internal-call ABI this completes. pub(super) fn position_call_ret_below_operands( &mut self, operand_count: usize, @@ -258,9 +261,8 @@ impl SymStack { "call operand count exceeds stack depth" ); - // During call operand preparation we arrange the operands as a left-rotation of the - // callee ABI order. With that setup, a single swap with the bottom operand moves the - // continuation behind the operands and restores the ABI operand order. + // One swap with the bottom operand moves the continuation behind the operands and + // restores ABI operand order (the operands were pre-rotated for exactly this). self.swap(operand_count, actions); debug_assert_eq!( @@ -270,14 +272,19 @@ impl SymStack { ); } - pub(super) fn remove_call_ret_addr(&mut self) { - let Some(pos) = self.items.iter().position(|i| *i == StackItem::CallRetAddr) else { - panic!("expected StackItem::CallRetAddr") - }; - self.items.remove(pos); - if self.func_ret_index.is_some_and(|idx| pos < idx) { - self.popped_above_barrier(); - } + /// Pop the internal-call return continuation marker, which must be on top. + /// + /// After a call's operands are popped the continuation is by construction the top item (it was + /// positioned at depth `argc` by `position_call_ret_below_operands`, or left on top when + /// `argc == 0`). See `Planner::prepare_internal_call` for the full ABI sequence. + pub(super) fn pop_call_ret_addr(&mut self) { + assert_eq!( + self.items.front(), + Some(&StackItem::CallRetAddr), + "call return address must be on top of the stack after popping call operands" + ); + self.items.pop_front(); + self.popped_above_barrier(); } pub(super) fn pop_operand(&mut self) { From 7cb6b2244b4d3da4136707cbc138719d667c4738 Mon Sep 17 00:00:00 2001 From: sbillig Date: Tue, 7 Jul 2026 00:54:06 -0700 Subject: [PATCH 08/14] Consolidate the stackify block walk into BlockPlanner - UseTracker (uses.rs) owns per-block use counting, cached-immediate tracking, and last-use/preserve queries; BlockLiveSets is gone. - rescue.rs owns dead-prefix cleanup and reachability-rescue heuristics; the ReachabilityValues wrapper is gone. - MemState groups the driver's memory-planning borrows; MemPlan::new drops from nine args to four and with_planner constructs it once. - BlockPlanner (block.rs, was block_sim.rs) owns the per-block walk with one method per terminator kind; observer calls live at one layer; edge bookkeeping stays on the driver as record_*_edge. All four take_stack dances are gone. - IterationPlanner is renamed FunctionPlanner and lives in driver.rs. --- .../codegen/src/stackalloc/stackify/block.rs | 491 ++++++++++ .../src/stackalloc/stackify/block_sim.rs | 422 -------- .../src/stackalloc/stackify/builder.rs | 20 +- .../codegen/src/stackalloc/stackify/driver.rs | 413 ++++++++ .../src/stackalloc/stackify/iteration.rs | 905 ------------------ crates/codegen/src/stackalloc/stackify/mod.rs | 6 +- .../stackify/planner/control_flow.rs | 42 +- .../src/stackalloc/stackify/planner/mod.rs | 36 +- .../stackify/planner/normalize_search.rs | 42 +- .../stackify/planner/operand_prep.rs | 268 +++--- .../codegen/src/stackalloc/stackify/rescue.rs | 217 +++++ .../codegen/src/stackalloc/stackify/uses.rs | 185 ++++ 12 files changed, 1505 insertions(+), 1542 deletions(-) create mode 100644 crates/codegen/src/stackalloc/stackify/block.rs delete mode 100644 crates/codegen/src/stackalloc/stackify/block_sim.rs create mode 100644 crates/codegen/src/stackalloc/stackify/driver.rs delete mode 100644 crates/codegen/src/stackalloc/stackify/iteration.rs create mode 100644 crates/codegen/src/stackalloc/stackify/rescue.rs create mode 100644 crates/codegen/src/stackalloc/stackify/uses.rs diff --git a/crates/codegen/src/stackalloc/stackify/block.rs b/crates/codegen/src/stackalloc/stackify/block.rs new file mode 100644 index 000000000..39a38f245 --- /dev/null +++ b/crates/codegen/src/stackalloc/stackify/block.rs @@ -0,0 +1,491 @@ +use std::ops::ControlFlow; + +use cranelift_entity::SecondaryMap; +use smallvec::SmallVec; +use sonatina_ir::{BlockId, Function, InstId, ValueId, inst::control_flow::BranchKind}; + +use crate::{bitset::BitSet, stackalloc::Actions}; + +use super::{ + br_table::plan_br_table_compare_chain, + builder::StackifyContext, + driver::FunctionPlanner, + rescue::{clean_dead_stack_prefix, improve_reachability_before_operands}, + slots::FreeSlotPools, + sym_stack::SymStack, + trace::StackifyObserver, + uses::UseTracker, +}; + +pub(super) struct BlockSimState { + pub(super) block: BlockId, + pub(super) free_slots: FreeSlotPools, + pub(super) prologue: Actions, + pub(super) injected_prologue: bool, + pub(super) uses: UseTracker, + pub(super) stack: SymStack, +} + +impl BlockSimState { + pub(super) fn new( + block: BlockId, + stack: SymStack, + free_slots: FreeSlotPools, + prologue: Actions, + uses: UseTracker, + ) -> Self { + Self { + block, + free_slots, + prologue, + injected_prologue: false, + uses, + stack, + } + } +} + +pub(super) enum PlannerActionSink { + Pre(InstId), + Post(InstId), + BrTableCase { inst: InstId, case_idx: usize }, +} + +enum TerminatorInfo { + Jump(BlockId), + Br { + cond: ValueId, + dests: SmallVec<[BlockId; 2]>, + }, + BrTable { + scrutinee: ValueId, + table: Vec<(ValueId, BlockId)>, + default: Option, + }, +} + +fn terminator_info(ctx: &StackifyContext<'_>, inst: InstId) -> Option { + let branch = ctx.func.dfg.branch_info(inst)?; + match branch.branch_kind() { + BranchKind::Jump(jump) => Some(TerminatorInfo::Jump(*jump.dest())), + BranchKind::Br(br) => Some(TerminatorInfo::Br { + cond: ctx.canonicalize_value(*br.cond()), + dests: branch.dests(), + }), + BranchKind::BrTable(table) => Some(TerminatorInfo::BrTable { + scrutinee: ctx.canonicalize_value(*table.scrutinee()), + table: table.table().to_vec(), + default: *table.default(), + }), + } +} + +/// The per-block instruction walk: threads a `SymStack` through one block, emitting actions and +/// trace events, and dispatching each terminator. Owns the block's `BlockSimState` and borrows the +/// driver ([`FunctionPlanner`]) for the planner/memory seams and cross-block edge bookkeeping. +pub(super) struct BlockPlanner<'d, 'a, 'ctx, O: StackifyObserver> { + driver: &'d mut FunctionPlanner<'a, 'ctx, O>, + state: BlockSimState, +} + +impl<'d, 'a, 'ctx, O: StackifyObserver> BlockPlanner<'d, 'a, 'ctx, O> { + pub(super) fn new(driver: &'d mut FunctionPlanner<'a, 'ctx, O>, state: BlockSimState) -> Self { + Self { driver, state } + } + + pub(super) fn run(mut self) -> BlockSimState { + // `func` is a shared `&Function` copied out of the driver, so the layout iterator does not + // borrow the driver and is free to coexist with the `&mut self` calls in the loop body. + let func = self.driver.ctx().func; + let block = self.state.block; + for inst in func.layout.iter_inst(block) { + if func.dfg.is_phi(inst) { + continue; + } + if self.plan_inst(inst).is_break() { + break; + } + } + self.state + } + + /// Plan one non-phi instruction: dead-prefix cleanup + reachability rescue, then dispatch to + /// the terminator-specific or normal-instruction planner. Returns `Break` for terminators (the + /// block walk stops) and `Continue` otherwise. + fn plan_inst(&mut self, inst: InstId) -> ControlFlow<()> { + self.driver.debug_assert_inst_actions_empty(inst); + + let is_call = self.driver.ctx().func.dfg.is_call(inst); + let call_has_stack_continuation = + is_call && call_has_local_return(self.driver.ctx().func, inst); + let terminator = terminator_info(self.driver.ctx(), inst); + let is_return = self.driver.ctx().func.dfg.is_return(inst); + let is_normal = terminator.is_none() && !is_return; + let skip_cleanup = skip_pre_exit_cleanup(self.driver.ctx().func, inst); + + let empty_last_use: BitSet = BitSet::default(); + let mut args = SmallVec::<[ValueId; 8]>::new(); + let mut consume_last_use: BitSet = BitSet::default(); + if is_normal { + args = operand_order_for_stackify( + self.driver.ctx().func, + inst, + &self.driver.ctx().value_aliases, + ); + consume_last_use = self.state.uses.last_uses_in(self.driver.ctx(), &args); + } + let cache_preserve = if is_normal { + self.state.uses.cache_preserve_in(self.driver.ctx(), &args) + } else { + BitSet::default() + }; + let last_use = if is_normal { + &consume_last_use + } else { + &empty_last_use + }; + + self.inject_prologue_once(inst); + { + let func = self.driver.ctx().func; + let stack = &self.state.stack; + let live_future = self.state.uses.live_future(); + let live_out = self.state.uses.live_out(); + self.driver.observer().on_inst_start( + func, + inst, + stack, + live_future, + live_out, + last_use, + ); + } + + let before_cleanup_len = self.driver.pre_actions_len(inst); + if !skip_cleanup { + let cleanup_live_future = self + .state + .uses + .cleanup_live(self.driver.ctx(), &self.state.stack); + let live_out = self.state.uses.live_out(); + let reach = self.driver.ctx().reach; + self.driver.with_pre_actions(inst, |actions| { + clean_dead_stack_prefix( + reach, + &mut self.state.stack, + &cleanup_live_future, + live_out, + actions, + ); + }); + } + if is_normal && !skip_cleanup { + self.rescue_reachability(inst, &args); + } + let after_cleanup_len = self.driver.pre_actions_len(inst); + { + let (observer, alloc) = self.driver.trace_actions(); + observer.on_inst_actions( + "cleanup", + &alloc.pre_actions[inst][before_cleanup_len..after_cleanup_len], + None, + ); + } + + if let Some(terminator) = terminator { + match terminator { + TerminatorInfo::Jump(dest) => self.plan_jump(inst, dest, after_cleanup_len), + TerminatorInfo::Br { cond, dests } => { + self.plan_br(inst, cond, &dests, after_cleanup_len) + } + TerminatorInfo::BrTable { + scrutinee, + table, + default, + } => self.plan_br_table(inst, scrutinee, table, default), + } + return ControlFlow::Break(()); + } + + if is_return { + self.plan_return(inst, after_cleanup_len); + return ControlFlow::Break(()); + } + + self.plan_normal_inst( + inst, + args, + last_use, + &cache_preserve, + call_has_stack_continuation, + after_cleanup_len, + ); + ControlFlow::Continue(()) + } + + /// Inject the block prologue into the first non-phi instruction's pre-actions, exactly once. + /// + /// Every block has a terminator that runs through here, so a non-empty prologue is always + /// injected by the end of the walk (asserted by `plan_block` in the driver). + fn inject_prologue_once(&mut self, inst: InstId) { + if self.state.prologue.is_empty() || self.state.injected_prologue { + return; + } + let prologue = &self.state.prologue; + self.driver + .with_pre_actions(inst, |actions| actions.extend_from_slice(prologue)); + self.state.injected_prologue = true; + } + + /// Reachability rescue for `operands` into `inst`'s pre-actions. Called by the normal path and + /// the `Br`/`BrTable` arms with their own operand list (`args`, `[cond]`, `[scrutinee]`), each + /// at the point where the rescue must run to preserve emitted action order. + fn rescue_reachability(&mut self, inst: InstId, operands: &[ValueId]) { + let func = self.driver.ctx().func; + let reach = self.driver.ctx().reach; + self.driver.with_pre_actions(inst, |actions| { + improve_reachability_before_operands( + func, + operands, + reach, + &mut self.state.stack, + &self.state.uses, + actions, + ); + }); + } + + fn plan_jump(&mut self, inst: InstId, dest: BlockId, action_start: usize) { + // The pending-edge case emits its exit + jump trace inside `record_jump_edge` (the deferred + // token must be captured synchronously with the pending edge); the other cases emit here. + if self + .driver + .record_jump_edge(&mut self.state, inst, dest, action_start) + { + return; + } + let (observer, alloc) = self.driver.trace_actions(); + observer.on_inst_actions("exit", &alloc.pre_actions[inst][action_start..], Some(dest)); + observer.on_inst_jump(inst, dest); + } + + fn plan_br( + &mut self, + inst: InstId, + cond: ValueId, + dests: &[BlockId], + after_cleanup_len: usize, + ) { + let consume_last_use = self.state.uses.last_uses_in(self.driver.ctx(), &[cond]); + + self.rescue_reachability(inst, &[cond]); + self.driver.with_planner( + &mut self.state.stack, + &mut self.state.free_slots, + PlannerActionSink::Pre(inst), + |planner| { + planner.prepare_operands(&[cond], &consume_last_use, &BitSet::default()); + }, + ); + + { + let (observer, alloc) = self.driver.trace_actions(); + observer.on_inst_actions("pre", &alloc.pre_actions[inst][after_cleanup_len..], None); + } + { + let func = self.driver.ctx().func; + self.driver.observer().on_inst_br(func, inst, cond, dests); + } + + let mut post_branch_stack = self.state.stack.clone(); + post_branch_stack.pop_operand(); + + for succ in dests.iter().copied() { + self.driver + .record_branch_edge(&mut self.state, succ, post_branch_stack.clone()); + } + } + + fn plan_br_table( + &mut self, + inst: InstId, + scrutinee: ValueId, + table: Vec<(ValueId, BlockId)>, + default: Option, + ) { + self.rescue_reachability(inst, &[scrutinee]); + + let (case_stacks, default_stack) = plan_br_table_compare_chain( + &table, + &self.state.stack, + |case_idx, case_val, case_stack| { + self.driver.with_planner( + case_stack, + &mut self.state.free_slots, + PlannerActionSink::BrTableCase { inst, case_idx }, + |planner| { + let consume_last_use = BitSet::::default(); + let mut compare_args = smallvec::smallvec![scrutinee, case_val]; + planner.prepare_operands_for_commutative_pair( + &mut compare_args, + &consume_last_use, + &BitSet::default(), + ); + }, + ); + }, + ); + + for case in case_stacks { + self.driver + .record_br_table_edge(&mut self.state, case.dest, case.post_compare_stack); + } + if let Some(default) = default { + self.driver + .record_br_table_edge(&mut self.state, default, default_stack); + } + + self.driver.observer().on_inst_br_table(inst); + } + + fn plan_return(&mut self, inst: InstId, after_cleanup_len: usize) { + self.driver.with_planner( + &mut self.state.stack, + &mut self.state.free_slots, + PlannerActionSink::Pre(inst), + |planner| planner.plan_internal_return(inst), + ); + + { + let (observer, alloc) = self.driver.trace_actions(); + observer.on_inst_actions( + "return", + &alloc.pre_actions[inst][after_cleanup_len..], + None, + ); + } + let func = self.driver.ctx().func; + let ret_vals: SmallVec<[ValueId; 16]> = func + .dfg + .return_args(inst) + .map(|args| args.iter().copied().collect()) + .unwrap_or_default(); + self.driver + .observer() + .on_inst_return(func, inst, ret_vals.as_slice()); + } + + fn plan_normal_inst( + &mut self, + inst: InstId, + mut args: SmallVec<[ValueId; 8]>, + last_use: &BitSet, + cache_preserve: &BitSet, + call_has_stack_continuation: bool, + after_cleanup_len: usize, + ) { + let results: SmallVec<[ValueId; 4]> = self + .driver + .ctx() + .func + .dfg + .inst_results(inst) + .iter() + .map(|&v| self.driver.ctx().canonicalize_value(v)) + .collect(); + + self.driver.with_planner( + &mut self.state.stack, + &mut self.state.free_slots, + PlannerActionSink::Pre(inst), + |planner| { + if call_has_stack_continuation { + planner.prepare_internal_call(inst, &mut args, last_use, cache_preserve); + } else { + planner.prepare_operands_for_inst(inst, &mut args, last_use, cache_preserve); + } + }, + ); + + { + let (observer, alloc) = self.driver.trace_actions(); + observer.on_inst_actions("pre", &alloc.pre_actions[inst][after_cleanup_len..], None); + } + { + let func = self.driver.ctx().func; + self.driver + .observer() + .on_inst_normal(func, inst, &args, &results); + } + + self.state.uses.consume( + self.driver.ctx(), + &args, + self.driver.scratch_slots(), + &mut self.state.free_slots.scratch, + ); + + self.state.stack.pop_n_operands(args.len()); + + if call_has_stack_continuation { + self.state.stack.pop_call_ret_addr(); + } + + for &res in results.iter().rev() { + self.state.stack.push_value(res); + } + + for (depth, &res) in results.iter().enumerate() { + if self.state.uses.is_dead(res) { + continue; + } + self.driver.with_planner( + &mut self.state.stack, + &mut self.state.free_slots, + PlannerActionSink::Post(inst), + |planner| planner.emit_store_if_spilled_at_depth(res, depth), + ); + } + + { + let (observer, alloc) = self.driver.trace_actions(); + observer.on_inst_actions("post", &alloc.post_actions[inst], None); + } + } +} + +pub(super) fn skip_pre_exit_cleanup(func: &Function, inst: InstId) -> bool { + func.dfg.is_exit(inst) && !func.dfg.is_return(inst) +} + +/// The alias-canonicalized operand order stackify plans against. For internal calls (those that +/// may return to the caller) the operands are rotated left by one so that a single `SWAP` after +/// pushing the return continuation restores callee ABI order — pairing with +/// `Planner::prepare_internal_call`; see the ABI note there. +pub(super) fn operand_order_for_stackify( + func: &Function, + inst: InstId, + value_aliases: &SecondaryMap>, +) -> SmallVec<[ValueId; 8]> { + let mut args: SmallVec<[ValueId; 8]> = func + .dfg + .inst(inst) + .collect_values() + .into_iter() + .map(|v| value_aliases[v].unwrap_or(v)) + .collect(); + + if call_has_local_return(func, inst) && !args.is_empty() { + args.as_mut_slice().rotate_left(1); + } + + args +} + +fn call_has_local_return(func: &Function, inst: InstId) -> bool { + func.dfg.call_info(inst).is_some_and(|call| { + func.ctx() + .func_effects(call.callee()) + .may_return_to_caller() + }) +} diff --git a/crates/codegen/src/stackalloc/stackify/block_sim.rs b/crates/codegen/src/stackalloc/stackify/block_sim.rs deleted file mode 100644 index 1dfc33bc5..000000000 --- a/crates/codegen/src/stackalloc/stackify/block_sim.rs +++ /dev/null @@ -1,422 +0,0 @@ -use std::collections::BTreeMap; - -use smallvec::SmallVec; -use sonatina_ir::{BlockId, I256, InstId, ValueId, inst::control_flow::BranchKind}; - -use crate::{bitset::BitSet, stackalloc::Actions}; - -use super::{ - br_table::plan_br_table_compare_chain, - builder::StackifyContext, - iteration::{ - IterationPlanner, ReachabilityValues, cached_immediate_preserve_values_in_inst, - clean_dead_stack_prefix, consume_cached_immediate_uses, consume_operand_uses, - count_block_uses, improve_reachability_before_operands, last_use_values_in_inst, - operand_order_for_stackify, skip_pre_exit_cleanup, - }, - slots::FreeSlotPools, - sym_stack::{StackItem, SymStack}, - trace::StackifyObserver, -}; - -pub(super) struct BlockSimState { - pub(super) block: BlockId, - pub(super) free_slots: FreeSlotPools, - pub(super) prologue: Actions, - pub(super) injected_prologue: bool, - remaining_uses: BTreeMap, - cached_remaining_uses: BTreeMap, - live_future: BitSet, - live_out: BitSet, - pub(super) stack: SymStack, -} - -pub(super) struct BlockLiveSets { - pub(super) remaining_uses: BTreeMap, - pub(super) live_future: BitSet, - pub(super) live_out: BitSet, - pub(super) cached_remaining_uses: BTreeMap, -} - -impl BlockSimState { - pub(super) fn block_live_sets(ctx: &StackifyContext<'_>, block: BlockId) -> BlockLiveSets { - let (remaining_uses, live_future, cached_remaining_uses) = count_block_uses(ctx, block); - let mut live_out = ctx.liveness.block_live_outs(block).clone(); - live_out.union_with(&ctx.phi_out_sources[block]); - BlockLiveSets { - remaining_uses, - live_future, - live_out, - cached_remaining_uses, - } - } - - pub(super) fn with_live_sets( - block: BlockId, - stack: SymStack, - free_slots: FreeSlotPools, - prologue: Actions, - live_sets: BlockLiveSets, - ) -> Self { - Self { - block, - free_slots, - prologue, - injected_prologue: false, - remaining_uses: live_sets.remaining_uses, - cached_remaining_uses: live_sets.cached_remaining_uses, - live_future: live_sets.live_future, - live_out: live_sets.live_out, - stack, - } - } - - pub(super) fn live_future(&self) -> &BitSet { - &self.live_future - } - - pub(super) fn live_out(&self) -> &BitSet { - &self.live_out - } -} - -pub(super) enum PlannerActionSink { - Pre(InstId), - Post(InstId), - BrTableCase { inst: InstId, case_idx: usize }, -} - -enum TerminatorInfo { - Jump(BlockId), - Br { - cond: ValueId, - dests: SmallVec<[BlockId; 2]>, - }, - BrTable { - scrutinee: ValueId, - table: Vec<(ValueId, BlockId)>, - default: Option, - }, -} - -fn take_stack(stack: &mut SymStack, has_internal_return: bool) -> SymStack { - std::mem::replace(stack, SymStack::opaque_prefix_empty(has_internal_return)) -} - -fn terminator_info(ctx: &StackifyContext<'_>, inst: InstId) -> Option { - let branch = ctx.func.dfg.branch_info(inst)?; - match branch.branch_kind() { - BranchKind::Jump(jump) => Some(TerminatorInfo::Jump(*jump.dest())), - BranchKind::Br(br) => Some(TerminatorInfo::Br { - cond: ctx.canonicalize_value(*br.cond()), - dests: branch.dests(), - }), - BranchKind::BrTable(table) => Some(TerminatorInfo::BrTable { - scrutinee: ctx.canonicalize_value(*table.scrutinee()), - table: table.table().to_vec(), - default: *table.default(), - }), - } -} - -fn cleanup_live_future_with_cached( - ctx: &StackifyContext<'_>, - stack: &SymStack, - live_future: &BitSet, - cached_remaining_uses: &BTreeMap, -) -> BitSet { - let mut live = live_future.clone(); - if cached_remaining_uses.is_empty() { - return live; - } - - for item in stack.iter() { - if let StackItem::Value(value) = item - && ctx.stack_caches_immediate(*value) - && let Some(imm) = ctx.func.dfg.value_imm(*value) - && cached_remaining_uses - .get(&imm.as_i256()) - .is_some_and(|count| *count != 0) - { - live.insert(*value); - } - } - live -} - -pub(super) fn run_block_sim( - planner: &mut IterationPlanner<'_, '_, O>, - mut state: BlockSimState, -) -> BlockSimState { - let empty_last_use: BitSet = BitSet::default(); - - // `func` is a shared `&Function` copied out of the planner, so the layout iterator does not - // borrow the planner and is free to coexist with the `&mut planner` calls in the loop body. - let func = planner.ctx().func; - for inst in func.layout.iter_inst(state.block) { - if planner.ctx().func.dfg.is_phi(inst) { - continue; - } - - planner.debug_assert_inst_actions_empty(inst); - - let is_call = planner.ctx().func.dfg.is_call(inst); - let call_has_stack_continuation = is_call && planner.call_uses_stack_continuation(inst); - let terminator = terminator_info(planner.ctx(), inst); - let is_normal = terminator.is_none() && !planner.ctx().func.dfg.is_return(inst); - let skip_cleanup = skip_pre_exit_cleanup(planner.ctx().func, inst); - - let mut args = SmallVec::<[ValueId; 8]>::new(); - let mut consume_last_use: BitSet = BitSet::default(); - if is_normal { - args = - operand_order_for_stackify(planner.ctx().func, inst, &planner.ctx().value_aliases); - consume_last_use = last_use_values_in_inst( - planner.ctx(), - &args, - &state.remaining_uses, - &state.live_out, - ); - } - let cache_preserve = if is_normal { - cached_immediate_preserve_values_in_inst( - planner.ctx(), - &args, - &state.cached_remaining_uses, - ) - } else { - BitSet::default() - }; - let last_use = if is_normal { - &consume_last_use - } else { - &empty_last_use - }; - planner.on_inst_start(&mut state, inst, last_use); - - let results: SmallVec<[ValueId; 4]> = planner - .ctx() - .func - .dfg - .inst_results(inst) - .iter() - .map(|&v| planner.ctx().canonicalize_value(v)) - .collect(); - let before_cleanup_len = planner.pre_actions_len(inst); - if !skip_cleanup { - let cleanup_live_future = cleanup_live_future_with_cached( - planner.ctx(), - &state.stack, - &state.live_future, - &state.cached_remaining_uses, - ); - let reach = planner.ctx().reach; - planner.with_pre_actions(inst, |actions| { - clean_dead_stack_prefix( - reach, - &mut state.stack, - &cleanup_live_future, - &state.live_out, - actions, - ); - }); - } - - if is_normal && !skip_cleanup { - let func = planner.ctx().func; - let values = ReachabilityValues { func }; - let reach = planner.ctx().reach; - planner.with_pre_actions(inst, |actions| { - improve_reachability_before_operands( - values, - &args, - reach, - &mut state.stack, - &state.live_future, - &state.live_out, - actions, - ); - }); - } - let after_cleanup_len = planner.pre_actions_len(inst); - planner.on_cleanup_actions(inst, before_cleanup_len, after_cleanup_len); - - if let Some(terminator) = terminator { - match terminator { - TerminatorInfo::Jump(dest) => { - planner.on_jump(&mut state, inst, dest, after_cleanup_len); - return state; - } - TerminatorInfo::Br { cond, dests } => { - let consume_last_use = last_use_values_in_inst( - planner.ctx(), - &[cond], - &state.remaining_uses, - &state.live_out, - ); - - let func = planner.ctx().func; - let values = ReachabilityValues { func }; - let reach = planner.ctx().reach; - planner.with_pre_actions(inst, |actions| { - improve_reachability_before_operands( - values, - &[cond], - reach, - &mut state.stack, - &state.live_future, - &state.live_out, - actions, - ); - }); - let mut stack = take_stack(&mut state.stack, planner.ctx().has_internal_return); - planner.with_planner( - &mut stack, - &mut state.free_slots, - PlannerActionSink::Pre(inst), - |planner| { - planner.prepare_operands( - &[cond], - &consume_last_use, - &BitSet::default(), - ); - }, - ); - state.stack = stack; - - planner.on_pre_actions(inst, after_cleanup_len); - planner.on_branch(inst, cond, dests.as_slice()); - - let mut post_branch_stack = state.stack.clone(); - post_branch_stack.pop_operand(); - - for succ in dests.iter().copied() { - planner.on_branch_edge(&mut state, succ, post_branch_stack.clone()); - } - return state; - } - TerminatorInfo::BrTable { - scrutinee, - table, - default, - } => { - let func = planner.ctx().func; - let values = ReachabilityValues { func }; - let reach = planner.ctx().reach; - planner.with_pre_actions(inst, |actions| { - improve_reachability_before_operands( - values, - &[scrutinee], - reach, - &mut state.stack, - &state.live_future, - &state.live_out, - actions, - ); - }); - - let (case_stacks, default_stack) = plan_br_table_compare_chain( - &table, - &state.stack, - |case_idx, case_val, case_stack| { - planner.with_planner( - case_stack, - &mut state.free_slots, - PlannerActionSink::BrTableCase { inst, case_idx }, - |planner| { - let consume_last_use = BitSet::::default(); - let mut compare_args = smallvec::smallvec![scrutinee, case_val]; - planner.prepare_operands_for_commutative_pair( - &mut compare_args, - &consume_last_use, - &BitSet::default(), - ); - }, - ); - }, - ); - - for case in case_stacks { - planner.on_br_table_edge(&mut state, case.dest, case.post_compare_stack); - } - if let Some(default) = default { - planner.on_br_table_edge(&mut state, default, default_stack); - } - - planner.on_br_table(inst); - return state; - } - } - } - - if planner.ctx().func.dfg.is_return(inst) { - let mut stack = take_stack(&mut state.stack, planner.ctx().has_internal_return); - planner.with_planner( - &mut stack, - &mut state.free_slots, - PlannerActionSink::Pre(inst), - |planner| planner.plan_internal_return(inst), - ); - state.stack = stack; - planner.on_return(inst, after_cleanup_len); - return state; - } - - let mut stack = take_stack(&mut state.stack, planner.ctx().has_internal_return); - planner.with_planner( - &mut stack, - &mut state.free_slots, - PlannerActionSink::Pre(inst), - |planner| { - if call_has_stack_continuation { - planner.prepare_internal_call(inst, &mut args, last_use, &cache_preserve); - } else { - planner.prepare_operands_for_inst(inst, &mut args, last_use, &cache_preserve); - } - }, - ); - state.stack = stack; - - planner.on_pre_actions(inst, after_cleanup_len); - planner.on_normal_inst(inst, &args, &results); - - consume_operand_uses( - planner.ctx(), - &args, - &mut state.remaining_uses, - &mut state.live_future, - &state.live_out, - planner.scratch_slots(), - &mut state.free_slots.scratch, - ); - consume_cached_immediate_uses(planner.ctx(), &args, &mut state.cached_remaining_uses); - - state.stack.pop_n_operands(args.len()); - - if call_has_stack_continuation { - state.stack.pop_call_ret_addr(); - } - - for &res in results.iter().rev() { - state.stack.push_value(res); - } - - for (depth, &res) in results.iter().enumerate() { - if !state.live_future.contains(res) && !state.live_out.contains(res) { - continue; - } - let mut stack = take_stack(&mut state.stack, planner.ctx().has_internal_return); - planner.with_planner( - &mut stack, - &mut state.free_slots, - PlannerActionSink::Post(inst), - |planner| planner.emit_store_if_spilled_at_depth(res, depth), - ); - state.stack = stack; - } - - planner.on_post_actions(inst); - } - - state -} diff --git a/crates/codegen/src/stackalloc/stackify/builder.rs b/crates/codegen/src/stackalloc/stackify/builder.rs index 7662ed7e2..9adaee3ca 100644 --- a/crates/codegen/src/stackalloc/stackify/builder.rs +++ b/crates/codegen/src/stackalloc/stackify/builder.rs @@ -14,8 +14,9 @@ use sonatina_ir::{BlockId, Function, I256, ValueId, cfg::ControlFlowGraph}; use super::{ alloc::{SpillStorage, StackifyAlloc}, - iteration::{IterationPlanner, operand_order_for_stackify}, - planner::{NormalizeSearchScratch, must_use_object_storage}, + block::operand_order_for_stackify, + driver::FunctionPlanner, + planner::{MemState, NormalizeSearchScratch, must_use_object_storage}, slots::{FreeSlotPools, SpillSlotInterference, SpillSlotPools}, spill::SpillSet, sym_stack::SymStack, @@ -465,18 +466,21 @@ impl<'a> StackifyBuilder<'a> { } inherited_stack.insert(ctx.entry, (ctx.entry, entry_stack)); - let mut planner = IterationPlanner::new( - ctx, + let mem = MemState { spill, + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills, slots, + }; + let mut planner = FunctionPlanner::new( + ctx, + mem, &mut templates, &terminal_chain_blocks, &interfaces.carry_in, &mut alloc, - &spill_obj, - &mut spill_requests, - &mut object_spill_requests, - forced_object_spills, inherited_stack, search_scratch, observer, diff --git a/crates/codegen/src/stackalloc/stackify/driver.rs b/crates/codegen/src/stackalloc/stackify/driver.rs new file mode 100644 index 000000000..52995f576 --- /dev/null +++ b/crates/codegen/src/stackalloc/stackify/driver.rs @@ -0,0 +1,413 @@ +use cranelift_entity::SecondaryMap; +use sonatina_ir::{BlockId, InstId, ValueId}; +use std::collections::BTreeMap; + +use crate::{bitset::BitSet, stackalloc::Actions}; + +use super::{ + alloc::StackifyAlloc, + block::{BlockPlanner, BlockSimState, PlannerActionSink}, + builder::StackifyContext, + planner::{self, MemState, NormalizeSearchScratch, Planner}, + slots::{FreeSlotPools, SlotPool}, + sym_stack::SymStack, + templates::{ + BlockTemplate, TransferOrder, canonical_transfer_order, choose_transfer, project_transfer, + }, + trace::StackifyObserver, + uses::UseTracker, +}; + +pub(super) struct FunctionPlanner<'a, 'ctx, O: StackifyObserver> { + ctx: &'a StackifyContext<'ctx>, + /// Memory-planning state (spill set, provisional object ids, spill/object requests, slot + /// pools) handed to each `MemPlan` in one reborrow. + mem: MemState<'a>, + templates: &'a mut SecondaryMap, + terminal_chain_blocks: &'a BitSet, + carry_in: &'a SecondaryMap>, + alloc: &'a mut StackifyAlloc, + inherited_stack: BTreeMap, + pending_edges: BTreeMap>, + planned_blocks: BitSet, + search_scratch: &'a mut NormalizeSearchScratch, + observer: &'a mut O, +} + +struct PendingEdge { + pred: BlockId, + inst: InstId, + stack: SymStack, + free_slots: FreeSlotPools, + action_start: usize, + /// Index of this edge's `DeferredExit` event in the observer's trace (0 for `NullObserver`), + /// used to backfill the exit fixup actions once the merge template is resolved. + trace_token: usize, +} + +impl<'a, 'ctx, O: StackifyObserver> FunctionPlanner<'a, 'ctx, O> { + #[allow(clippy::too_many_arguments)] + pub(super) fn new( + ctx: &'a StackifyContext<'ctx>, + mem: MemState<'a>, + templates: &'a mut SecondaryMap, + terminal_chain_blocks: &'a BitSet, + carry_in: &'a SecondaryMap>, + alloc: &'a mut StackifyAlloc, + inherited_stack: BTreeMap, + search_scratch: &'a mut NormalizeSearchScratch, + observer: &'a mut O, + ) -> Self { + Self { + ctx, + mem, + templates, + terminal_chain_blocks, + carry_in, + alloc, + inherited_stack, + pending_edges: BTreeMap::new(), + planned_blocks: BitSet::default(), + search_scratch, + observer, + } + } + + fn with_actions_planner( + &mut self, + stack: &mut SymStack, + actions: &mut Actions, + free_slots: &mut FreeSlotPools, + f: impl FnOnce(&mut Planner) -> R, + ) -> R { + let mem = planner::MemPlan::new( + &mut self.mem, + self.ctx, + &self.alloc.exact_local_addr, + free_slots, + ); + let mut planner = Planner::new(self.ctx, stack, actions, mem, &mut *self.search_scratch); + f(&mut planner) + } + + pub(super) fn plan_blocks(&mut self) { + for &block in self.ctx.dom.rpo() { + if block != self.ctx.entry && !self.ctx.dom.is_reachable(block) { + continue; + } + + self.plan_block(block); + } + + debug_assert!( + self.pending_edges.is_empty(), + "unresolved stackify edges remain" + ); + } + + fn plan_block(&mut self, block: BlockId) { + let mut free_slots: FreeSlotPools = FreeSlotPools::default(); + let mut prologue: Actions = Actions::new(); + + let uses = UseTracker::for_block(self.ctx, block); + self.resolve_pending_edges(block); + + let inherited = self.inherited_stack.remove(&block); + if self.terminal_chain_blocks.contains(block) { + self.freeze_template(block, TransferOrder::new()); + } else if let Some((_pred, stack)) = inherited.as_ref() { + self.freeze_template_from_stack(block, stack); + } else if block != self.ctx.entry { + self.freeze_template_canonical(block); + } + + self.observer + .on_block_header(self.ctx.func, block, &self.templates[block]); + self.planned_blocks.insert(block); + + let stack = if self.terminal_chain_blocks.contains(block) { + SymStack::opaque_prefix_empty(self.ctx.has_internal_return) + } else if let Some((pred, mut inh)) = inherited { + // Dynamic entry stack (single predecessor). + if block != self.ctx.entry { + debug_assert_eq!( + self.ctx.cfg.pred_num_of(block), + 1, + "inherited stack implies single-predecessor block" + ); + self.observer.on_block_inherited( + self.ctx.func, + block, + pred, + &inh, + uses.live_future(), + uses.live_out(), + ); + let has_phi_params = !self.ctx.phi_results[block].is_empty(); + + // Single-predecessor blocks without phis do not need exact entry + // template normalization. Keeping the inherited stack avoids pointless bottom + // reshuffling that can cascade into SWAP/POP churn. + if has_phi_params { + let tmpl = self.templates[block].clone(); + self.with_actions_planner( + &mut inh, + &mut prologue, + &mut free_slots, + |planner| { + planner.plan_edge_fixup_to_template(&tmpl, pred, block); + }, + ); + } + self.observer.on_block_prologue(&prologue); + } + inh + } else { + SymStack::from_template(&self.templates[block], self.ctx.has_internal_return) + }; + + let state = BlockSimState::new(block, stack, free_slots, prologue, uses); + let state = BlockPlanner::new(self, state).run(); + + // The block walk injects the prologue on every non-phi instruction (including the + // terminator that every block has), so a non-empty prologue is always injected. + debug_assert!( + state.prologue.is_empty() || state.injected_prologue, + "prologue was not injected during the block walk" + ); + } + + fn resolve_pending_edges(&mut self, block: BlockId) { + let Some(mut edges) = self.pending_edges.remove(&block) else { + return; + }; + debug_assert!( + !self.inherited_stack.contains_key(&block), + "pending merge edges cannot also inherit one stack" + ); + debug_assert!( + !self.planned_blocks.contains(block), + "pending edge target already planned" + ); + + let projected: Vec<(BlockId, TransferOrder)> = edges + .iter() + .map(|edge| { + ( + edge.pred, + project_transfer(&edge.stack, &self.carry_in[block]), + ) + }) + .collect(); + if !projected.is_empty() { + self.freeze_template(block, choose_transfer(self.ctx, block, &projected)); + } + + let tmpl = self.templates[block].clone(); + for edge in edges.iter_mut() { + debug_assert_eq!( + self.alloc.pre_actions[edge.inst].len(), + edge.action_start, + "deferred edge action list changed before resolution" + ); + self.with_planner( + &mut edge.stack, + &mut edge.free_slots, + PlannerActionSink::Pre(edge.inst), + |planner| planner.plan_edge_fixup_to_template(&tmpl, edge.pred, block), + ); + self.observer.on_deferred_exit_actions( + edge.trace_token, + &self.alloc.pre_actions[edge.inst][edge.action_start..], + ); + } + } + + fn freeze_template_from_stack(&mut self, block: BlockId, stack: &SymStack) { + self.freeze_template(block, project_transfer(stack, &self.carry_in[block])); + } + + fn freeze_template_canonical(&mut self, block: BlockId) { + let transfer = canonical_transfer_order( + &self.carry_in[block], + &self.ctx.dom_depth, + &self.ctx.def_info, + ); + self.freeze_template(block, transfer); + } + + fn freeze_template(&mut self, block: BlockId, transfer: TransferOrder) { + self.templates[block].freeze_transfer(transfer); + } + + pub(super) fn ctx(&self) -> &StackifyContext<'ctx> { + self.ctx + } + + pub(super) fn scratch_slots(&self) -> &SlotPool { + &self.mem.slots.scratch + } + + pub(super) fn debug_assert_inst_actions_empty(&self, inst: InstId) { + // Within one fixed-point iteration each inst is visited once on a fresh `StackifyAlloc`, + // so its action buffers are always still empty when the walk reaches it. + debug_assert!( + self.alloc.pre_actions[inst].is_empty() + && self.alloc.post_actions[inst].is_empty() + && self.alloc.brtable_actions[inst].is_empty(), + "inst action buffers are not empty at the start of the block walk" + ); + } + + pub(super) fn pre_actions_len(&self, inst: InstId) -> usize { + self.alloc.pre_actions[inst].len() + } + + pub(super) fn with_pre_actions( + &mut self, + inst: InstId, + f: impl FnOnce(&mut Actions) -> R, + ) -> R { + f(&mut self.alloc.pre_actions[inst]) + } + + pub(super) fn with_planner( + &mut self, + stack: &mut SymStack, + free_slots: &mut FreeSlotPools, + sink: PlannerActionSink, + f: impl FnOnce(&mut Planner<'_, '_>) -> R, + ) -> R { + // Resolve the action buffer first: `Pre`/`Post` index into `alloc` (a field disjoint from + // `alloc.exact_local_addr`, which `MemPlan` borrows), while `BrTableCase` accumulates into + // a local buffer pushed onto `brtable_actions` after planning. Then construct the + // `MemPlan`/`Planner` once for all three. + let mut brtable_buf = Actions::new(); + let actions: &mut Actions = match sink { + PlannerActionSink::Pre(inst) => &mut self.alloc.pre_actions[inst], + PlannerActionSink::Post(inst) => &mut self.alloc.post_actions[inst], + PlannerActionSink::BrTableCase { .. } => &mut brtable_buf, + }; + let mem = planner::MemPlan::new( + &mut self.mem, + self.ctx, + &self.alloc.exact_local_addr, + free_slots, + ); + let result = { + let mut planner = + Planner::new(self.ctx, stack, actions, mem, &mut *self.search_scratch); + f(&mut planner) + }; + if let PlannerActionSink::BrTableCase { inst, case_idx } = sink { + debug_assert_eq!(self.alloc.brtable_actions[inst].len(), case_idx); + self.alloc.brtable_actions[inst].push(brtable_buf); + } + result + } + + pub(super) fn observer(&mut self) -> &mut O { + &mut *self.observer + } + + /// Split borrow of the observer and the (read-only) allocation, so the block walk can emit an + /// action-group trace event over a slice of `alloc` while notifying the observer. + pub(super) fn trace_actions(&mut self) -> (&mut O, &StackifyAlloc) { + (&mut *self.observer, &*self.alloc) + } + + /// Record a `jump` edge to `dest`, applying immediate normalization / inheritance / pending + /// deferral as appropriate. Returns `true` when the edge was deferred (a pending merge edge): + /// its exit + jump trace has already been emitted here (synchronously with capturing the + /// deferred trace token), so the caller emits nothing further. Returns `false` for a resolved + /// edge, so the caller emits the exit-normalization actions and the jump trace. + pub(super) fn record_jump_edge( + &mut self, + state: &mut BlockSimState, + inst: InstId, + dest: BlockId, + action_start: usize, + ) -> bool { + if self.terminal_chain_blocks.contains(dest) { + } else if self.ctx.cfg.pred_num_of(dest) > 1 + && dest != self.ctx.entry + && !self.planned_blocks.contains(dest) + { + debug_assert!( + self.ctx.scc.is_reachable(dest), + "pending edge target must be reachable" + ); + self.pending_edges + .entry(dest) + .or_default() + .push(PendingEdge { + pred: state.block, + inst, + stack: state.stack.clone(), + free_slots: state.free_slots.clone(), + action_start, + trace_token: self.observer.on_deferred_inst_jump(inst, dest), + }); + return true; + } else if self.ctx.cfg.pred_num_of(dest) == 1 + && dest != self.ctx.entry + && !self.planned_blocks.contains(dest) + { + self.inherited_stack + .entry(dest) + .or_insert_with(|| (state.block, state.stack.clone())); + } else { + let tmpl = self.templates[dest].clone(); + let src = state.block; + self.with_planner( + &mut state.stack, + &mut state.free_slots, + PlannerActionSink::Pre(inst), + |planner| planner.plan_edge_fixup_to_template(&tmpl, src, dest), + ); + } + false + } + + pub(super) fn record_branch_edge( + &mut self, + state: &mut BlockSimState, + succ: BlockId, + stack: SymStack, + ) { + assert!( + !self.planned_blocks.contains(succ), + "multiway branch edge to already-planned block {succ:?}: run StackifyEdgeSplitter \ + before stackify to split in-cycle multiway edges" + ); + debug_assert_eq!( + self.ctx.cfg.pred_num_of(succ), + 1, + "no critical edges: branch target must be single-pred" + ); + self.inherited_stack + .entry(succ) + .or_insert_with(|| (state.block, stack)); + } + + pub(super) fn record_br_table_edge( + &mut self, + state: &mut BlockSimState, + succ: BlockId, + stack: SymStack, + ) { + assert!( + !self.planned_blocks.contains(succ), + "multiway br_table edge to already-planned block {succ:?}: run StackifyEdgeSplitter \ + before stackify to split in-cycle multiway edges" + ); + debug_assert_eq!( + self.ctx.cfg.pred_num_of(succ), + 1, + "no critical edges: br_table target must be single-pred" + ); + self.inherited_stack + .entry(succ) + .or_insert_with(|| (state.block, stack)); + } +} diff --git a/crates/codegen/src/stackalloc/stackify/iteration.rs b/crates/codegen/src/stackalloc/stackify/iteration.rs deleted file mode 100644 index b57cfe508..000000000 --- a/crates/codegen/src/stackalloc/stackify/iteration.rs +++ /dev/null @@ -1,905 +0,0 @@ -use cranelift_entity::SecondaryMap; -use smallvec::SmallVec; -use sonatina_ir::{BlockId, Function, I256, InstId, ValueId}; -use std::collections::BTreeMap; - -use crate::{bitset::BitSet, isa::evm::immediate_materialization_code_len, stackalloc::Actions}; - -use super::{ - alloc::StackifyAlloc, - block_sim::{BlockSimState, PlannerActionSink, run_block_sim}, - builder::{StackifyContext, StackifyReachability}, - planner::{self, NormalizeSearchScratch, Planner}, - slots::{FreeSlotPools, FreeSlots, SlotPool, SpillSlotPools}, - spill::SpillSet, - sym_stack::{StackItem, SymStack}, - templates::{ - BlockTemplate, TransferOrder, canonical_transfer_order, choose_transfer, project_transfer, - }, - trace::StackifyObserver, -}; - -pub(super) struct IterationPlanner<'a, 'ctx, O: StackifyObserver> { - ctx: &'a StackifyContext<'ctx>, - spill: SpillSet<'a>, - slots: &'a mut SpillSlotPools, - templates: &'a mut SecondaryMap, - terminal_chain_blocks: &'a BitSet, - carry_in: &'a SecondaryMap>, - alloc: &'a mut StackifyAlloc, - /// Provisional per-iteration object-id assignment (`assign_spill_obj_ids`), read by - /// `MemPlan` during planning before storage is finalized. - spill_obj: &'a SecondaryMap>, - spill_requests: &'a mut BitSet, - object_spill_requests: &'a mut BitSet, - forced_object_spills: &'a BitSet, - inherited_stack: BTreeMap, - pending_edges: BTreeMap>, - planned_blocks: BitSet, - search_scratch: &'a mut NormalizeSearchScratch, - observer: &'a mut O, -} - -struct PendingEdge { - pred: BlockId, - inst: InstId, - stack: SymStack, - free_slots: FreeSlotPools, - action_start: usize, - /// Index of this edge's `DeferredExit` event in the observer's trace (0 for `NullObserver`), - /// used to backfill the exit fixup actions once the merge template is resolved. - trace_token: usize, -} - -#[derive(Clone, Copy)] -pub(super) struct ReachabilityValues<'a> { - pub(super) func: &'a Function, -} - -impl ReachabilityValues<'_> { - fn retains(self, value: ValueId) -> bool { - !self.func.dfg.value_is_imm(value) - } -} - -impl<'a, 'ctx, O: StackifyObserver> IterationPlanner<'a, 'ctx, O> { - #[allow(clippy::too_many_arguments)] - pub(super) fn new( - ctx: &'a StackifyContext<'ctx>, - spill: SpillSet<'a>, - slots: &'a mut SpillSlotPools, - templates: &'a mut SecondaryMap, - terminal_chain_blocks: &'a BitSet, - carry_in: &'a SecondaryMap>, - alloc: &'a mut StackifyAlloc, - spill_obj: &'a SecondaryMap< - ValueId, - Option, - >, - spill_requests: &'a mut BitSet, - object_spill_requests: &'a mut BitSet, - forced_object_spills: &'a BitSet, - inherited_stack: BTreeMap, - search_scratch: &'a mut NormalizeSearchScratch, - observer: &'a mut O, - ) -> Self { - Self { - ctx, - spill, - slots, - templates, - terminal_chain_blocks, - carry_in, - alloc, - spill_obj, - spill_requests, - object_spill_requests, - forced_object_spills, - inherited_stack, - pending_edges: BTreeMap::new(), - planned_blocks: BitSet::default(), - search_scratch, - observer, - } - } - - fn with_actions_planner( - &mut self, - stack: &mut SymStack, - actions: &mut Actions, - free_slots: &mut FreeSlotPools, - f: impl FnOnce(&mut Planner) -> R, - ) -> R { - let mem = planner::MemPlan::new( - self.spill, - &mut *self.spill_requests, - self.ctx, - self.spill_obj, - &self.alloc.exact_local_addr, - self.object_spill_requests, - self.forced_object_spills, - free_slots, - &mut *self.slots, - ); - let mut planner = Planner::new(self.ctx, stack, actions, mem, &mut *self.search_scratch); - f(&mut planner) - } - - pub(super) fn plan_blocks(&mut self) { - for &block in self.ctx.dom.rpo() { - if block != self.ctx.entry && !self.ctx.dom.is_reachable(block) { - continue; - } - - self.plan_block(block); - } - - debug_assert!( - self.pending_edges.is_empty(), - "unresolved stackify edges remain" - ); - } - - fn plan_block(&mut self, block: BlockId) { - let mut free_slots: FreeSlotPools = FreeSlotPools::default(); - let mut prologue: Actions = Actions::new(); - - let live_sets = BlockSimState::block_live_sets(self.ctx, block); - self.resolve_pending_edges(block); - - let inherited = self.inherited_stack.remove(&block); - if self.terminal_chain_blocks.contains(block) { - self.freeze_template(block, TransferOrder::new()); - } else if let Some((_pred, stack)) = inherited.as_ref() { - self.freeze_template_from_stack(block, stack); - } else if block != self.ctx.entry { - self.freeze_template_canonical(block); - } - - self.observer - .on_block_header(self.ctx.func, block, &self.templates[block]); - self.planned_blocks.insert(block); - - let stack = if self.terminal_chain_blocks.contains(block) { - SymStack::opaque_prefix_empty(self.ctx.has_internal_return) - } else if let Some((pred, mut inh)) = inherited { - // Dynamic entry stack (single predecessor). - if block != self.ctx.entry { - debug_assert_eq!( - self.ctx.cfg.pred_num_of(block), - 1, - "inherited stack implies single-predecessor block" - ); - self.observer.on_block_inherited( - self.ctx.func, - block, - pred, - &inh, - &live_sets.live_future, - &live_sets.live_out, - ); - let has_phi_params = !self.ctx.phi_results[block].is_empty(); - - // Single-predecessor blocks without phis do not need exact entry - // template normalization. Keeping the inherited stack avoids pointless bottom - // reshuffling that can cascade into SWAP/POP churn. - if has_phi_params { - let tmpl = self.templates[block].clone(); - self.with_actions_planner( - &mut inh, - &mut prologue, - &mut free_slots, - |planner| { - planner.plan_edge_fixup_to_template(&tmpl, pred, block); - }, - ); - } - self.observer.on_block_prologue(&prologue); - } - inh - } else { - SymStack::from_template(&self.templates[block], self.ctx.has_internal_return) - }; - - let state = BlockSimState::with_live_sets(block, stack, free_slots, prologue, live_sets); - let state = run_block_sim(self, state); - - // `on_inst_start` runs for every non-phi instruction (including the terminator that every - // block has), so a non-empty prologue is always injected during `run_block_sim`. - debug_assert!( - state.prologue.is_empty() || state.injected_prologue, - "prologue was not injected during the block walk" - ); - } - - fn resolve_pending_edges(&mut self, block: BlockId) { - let Some(mut edges) = self.pending_edges.remove(&block) else { - return; - }; - debug_assert!( - !self.inherited_stack.contains_key(&block), - "pending merge edges cannot also inherit one stack" - ); - debug_assert!( - !self.planned_blocks.contains(block), - "pending edge target already planned" - ); - - let projected: Vec<(BlockId, TransferOrder)> = edges - .iter() - .map(|edge| { - ( - edge.pred, - project_transfer(&edge.stack, &self.carry_in[block]), - ) - }) - .collect(); - if !projected.is_empty() { - self.freeze_template(block, choose_transfer(self.ctx, block, &projected)); - } - - let tmpl = self.templates[block].clone(); - for edge in edges.iter_mut() { - debug_assert_eq!( - self.alloc.pre_actions[edge.inst].len(), - edge.action_start, - "deferred edge action list changed before resolution" - ); - self.with_planner( - &mut edge.stack, - &mut edge.free_slots, - PlannerActionSink::Pre(edge.inst), - |planner| planner.plan_edge_fixup_to_template(&tmpl, edge.pred, block), - ); - self.observer.on_deferred_exit_actions( - edge.trace_token, - &self.alloc.pre_actions[edge.inst][edge.action_start..], - ); - } - } - - fn freeze_template_from_stack(&mut self, block: BlockId, stack: &SymStack) { - self.freeze_template(block, project_transfer(stack, &self.carry_in[block])); - } - - fn freeze_template_canonical(&mut self, block: BlockId) { - let transfer = canonical_transfer_order( - &self.carry_in[block], - &self.ctx.dom_depth, - &self.ctx.def_info, - ); - self.freeze_template(block, transfer); - } - - fn freeze_template(&mut self, block: BlockId, transfer: TransferOrder) { - self.templates[block].freeze_transfer(transfer); - } - - pub(super) fn ctx(&self) -> &StackifyContext<'ctx> { - self.ctx - } - - pub(super) fn call_uses_stack_continuation(&self, inst: InstId) -> bool { - call_has_local_return(self.ctx.func, inst) - } - - pub(super) fn scratch_slots(&self) -> &SlotPool { - &self.slots.scratch - } - - pub(super) fn debug_assert_inst_actions_empty(&self, inst: InstId) { - // Within one fixed-point iteration each inst is visited once on a fresh `StackifyAlloc`, - // so its action buffers are always still empty when the walk reaches it. - debug_assert!( - self.alloc.pre_actions[inst].is_empty() - && self.alloc.post_actions[inst].is_empty() - && self.alloc.brtable_actions[inst].is_empty(), - "inst action buffers are not empty at the start of the block walk" - ); - } - - pub(super) fn pre_actions_len(&self, inst: InstId) -> usize { - self.alloc.pre_actions[inst].len() - } - - pub(super) fn with_pre_actions( - &mut self, - inst: InstId, - f: impl FnOnce(&mut Actions) -> R, - ) -> R { - f(&mut self.alloc.pre_actions[inst]) - } - - pub(super) fn with_planner( - &mut self, - stack: &mut SymStack, - free_slots: &mut FreeSlotPools, - sink: PlannerActionSink, - f: impl FnOnce(&mut Planner<'_, '_>) -> R, - ) -> R { - match sink { - PlannerActionSink::Pre(inst) => { - let mem = planner::MemPlan::new( - self.spill, - &mut *self.spill_requests, - self.ctx, - self.spill_obj, - &self.alloc.exact_local_addr, - self.object_spill_requests, - self.forced_object_spills, - free_slots, - &mut *self.slots, - ); - let mut planner = Planner::new( - self.ctx, - stack, - &mut self.alloc.pre_actions[inst], - mem, - &mut *self.search_scratch, - ); - f(&mut planner) - } - PlannerActionSink::Post(inst) => { - let mem = planner::MemPlan::new( - self.spill, - &mut *self.spill_requests, - self.ctx, - self.spill_obj, - &self.alloc.exact_local_addr, - self.object_spill_requests, - self.forced_object_spills, - free_slots, - &mut *self.slots, - ); - let mut planner = Planner::new( - self.ctx, - stack, - &mut self.alloc.post_actions[inst], - mem, - &mut *self.search_scratch, - ); - f(&mut planner) - } - PlannerActionSink::BrTableCase { inst, case_idx } => { - let mut actions = Actions::new(); - let mem = planner::MemPlan::new( - self.spill, - &mut *self.spill_requests, - self.ctx, - self.spill_obj, - &self.alloc.exact_local_addr, - self.object_spill_requests, - self.forced_object_spills, - free_slots, - &mut *self.slots, - ); - let result = { - let mut planner = Planner::new( - self.ctx, - stack, - &mut actions, - mem, - &mut *self.search_scratch, - ); - f(&mut planner) - }; - debug_assert_eq!(self.alloc.brtable_actions[inst].len(), case_idx); - self.alloc.brtable_actions[inst].push(actions); - result - } - } - } - - pub(super) fn on_inst_start( - &mut self, - state: &mut BlockSimState, - inst: InstId, - last_use: &BitSet, - ) { - if !state.prologue.is_empty() && !state.injected_prologue { - self.alloc.pre_actions[inst].extend_from_slice(&state.prologue); - state.injected_prologue = true; - } - self.observer.on_inst_start( - self.ctx.func, - inst, - &state.stack, - state.live_future(), - state.live_out(), - last_use, - ); - } - - pub(super) fn on_cleanup_actions(&mut self, inst: InstId, start: usize, end: usize) { - self.observer - .on_inst_actions("cleanup", &self.alloc.pre_actions[inst][start..end], None); - } - - pub(super) fn on_pre_actions(&mut self, inst: InstId, start: usize) { - self.observer - .on_inst_actions("pre", &self.alloc.pre_actions[inst][start..], None); - } - - pub(super) fn on_post_actions(&mut self, inst: InstId) { - self.observer - .on_inst_actions("post", &self.alloc.post_actions[inst], None); - } - - pub(super) fn on_normal_inst(&mut self, inst: InstId, args: &[ValueId], results: &[ValueId]) { - self.observer - .on_inst_normal(self.ctx.func, inst, args, results); - } - - pub(super) fn on_return(&mut self, inst: InstId, start: usize) { - self.observer - .on_inst_actions("return", &self.alloc.pre_actions[inst][start..], None); - let ret_vals: SmallVec<[ValueId; 16]> = self - .ctx - .func - .dfg - .return_args(inst) - .map(|args| args.iter().copied().collect()) - .unwrap_or_default(); - self.observer - .on_inst_return(self.ctx.func, inst, ret_vals.as_slice()); - } - - pub(super) fn on_jump( - &mut self, - state: &mut BlockSimState, - inst: InstId, - dest: BlockId, - action_start: usize, - ) { - if self.terminal_chain_blocks.contains(dest) { - } else if self.ctx.cfg.pred_num_of(dest) > 1 - && dest != self.ctx.entry - && !self.planned_blocks.contains(dest) - { - debug_assert!( - self.ctx.scc.is_reachable(dest), - "pending edge target must be reachable" - ); - self.pending_edges - .entry(dest) - .or_default() - .push(PendingEdge { - pred: state.block, - inst, - stack: state.stack.clone(), - free_slots: state.free_slots.clone(), - action_start, - trace_token: self.observer.on_deferred_inst_jump(inst, dest), - }); - return; - } else if self.ctx.cfg.pred_num_of(dest) == 1 - && dest != self.ctx.entry - && !self.planned_blocks.contains(dest) - { - self.inherited_stack - .entry(dest) - .or_insert_with(|| (state.block, state.stack.clone())); - } else { - let tmpl = self.templates[dest].clone(); - let src = state.block; - self.with_planner( - &mut state.stack, - &mut state.free_slots, - PlannerActionSink::Pre(inst), - |planner| planner.plan_edge_fixup_to_template(&tmpl, src, dest), - ); - } - - self.observer.on_inst_actions( - "exit", - &self.alloc.pre_actions[inst][action_start..], - Some(dest), - ); - self.observer.on_inst_jump(inst, dest); - } - - pub(super) fn on_branch(&mut self, inst: InstId, cond: ValueId, dests: &[BlockId]) { - self.observer.on_inst_br(self.ctx.func, inst, cond, dests); - } - - pub(super) fn on_branch_edge( - &mut self, - state: &mut BlockSimState, - succ: BlockId, - stack: SymStack, - ) { - assert!( - !self.planned_blocks.contains(succ), - "multiway branch edge to already-planned block {succ:?}: run StackifyEdgeSplitter \ - before stackify to split in-cycle multiway edges" - ); - debug_assert_eq!( - self.ctx.cfg.pred_num_of(succ), - 1, - "no critical edges: branch target must be single-pred" - ); - self.inherited_stack - .entry(succ) - .or_insert_with(|| (state.block, stack)); - } - - pub(super) fn on_br_table_edge( - &mut self, - state: &mut BlockSimState, - succ: BlockId, - stack: SymStack, - ) { - assert!( - !self.planned_blocks.contains(succ), - "multiway br_table edge to already-planned block {succ:?}: run StackifyEdgeSplitter \ - before stackify to split in-cycle multiway edges" - ); - debug_assert_eq!( - self.ctx.cfg.pred_num_of(succ), - 1, - "no critical edges: br_table target must be single-pred" - ); - self.inherited_stack - .entry(succ) - .or_insert_with(|| (state.block, stack)); - } - - pub(super) fn on_br_table(&mut self, inst: InstId) { - self.observer.on_inst_br_table(inst); - } -} - -pub(super) fn skip_pre_exit_cleanup(func: &Function, inst: InstId) -> bool { - func.dfg.is_exit(inst) && !func.dfg.is_return(inst) -} - -pub(super) fn count_block_uses( - ctx: &StackifyContext<'_>, - block: BlockId, -) -> (BTreeMap, BitSet, BTreeMap) { - let mut counts: BTreeMap = BTreeMap::new(); - let mut cached_counts: BTreeMap = BTreeMap::new(); - for inst in ctx.func.layout.iter_inst(block) { - if ctx.func.dfg.is_phi(inst) { - continue; - } - for v in operand_order_for_stackify(ctx.func, inst, &ctx.value_aliases) { - if ctx.retains_value(v) { - *counts.entry(v).or_insert(0) += 1; - } else if ctx.stack_caches_immediate(v) - && let Some(imm) = ctx.func.dfg.value_imm(v) - { - *cached_counts.entry(imm.as_i256()).or_insert(0) += 1; - } - } - } - let live_future: BitSet = counts.keys().copied().collect(); - (counts, live_future, cached_counts) -} - -pub(super) fn cached_immediate_preserve_values_in_inst( - ctx: &StackifyContext<'_>, - args: &[ValueId], - cached_remaining_uses: &BTreeMap, -) -> BitSet { - let mut inst_counts: BTreeMap = BTreeMap::new(); - for &value in args { - if ctx.stack_caches_immediate(value) - && let Some(imm) = ctx.func.dfg.value_imm(value) - { - *inst_counts.entry(imm.as_i256()).or_insert(0) += 1; - } - } - - let mut preserve: BitSet = BitSet::default(); - for &value in args { - if ctx.stack_caches_immediate(value) - && let Some(imm) = ctx.func.dfg.value_imm(value) - && cached_remaining_uses - .get(&imm.as_i256()) - .copied() - .unwrap_or(0) - > inst_counts.get(&imm.as_i256()).copied().unwrap_or(0) - { - preserve.insert(value); - } - } - preserve -} - -pub(super) fn consume_cached_immediate_uses( - ctx: &StackifyContext<'_>, - args: &[ValueId], - cached_remaining_uses: &mut BTreeMap, -) { - for &value in args { - if ctx.stack_caches_immediate(value) - && let Some(imm) = ctx.func.dfg.value_imm(value) - && let Some(count) = cached_remaining_uses.get_mut(&imm.as_i256()) - { - *count = count.saturating_sub(1); - } - } - cached_remaining_uses.retain(|_, count| *count != 0); -} - -fn pop_dead_tops( - stack: &mut SymStack, - live_future: &BitSet, - live_out: &BitSet, - actions: &mut Actions, -) { - while let Some(StackItem::Value(v)) = stack.top() { - if live_future.contains(*v) || live_out.contains(*v) { - break; - } - stack.pop(actions); - } -} - -pub(super) fn clean_dead_stack_prefix( - reach: StackifyReachability, - stack: &mut SymStack, - live_future: &BitSet, - live_out: &BitSet, - actions: &mut Actions, -) { - // Two local cleanups: - // - pop any dead values that reach the top - // - if a live value is on top and there is a contiguous dead chain directly beneath it, - // swap the live value down and pop the dead chain off the top. - loop { - let before_len = stack.len(); - pop_dead_tops(stack, live_future, live_out, actions); - - // If the top is not a normal value, don't try to reorder under it. - let Some(StackItem::Value(top)) = stack.top() else { - break; - }; - // `pop_dead_tops` only stops on a live value (or a non-`Value`, handled above). - debug_assert!( - live_future.contains(*top) || live_out.contains(*top), - "pop_dead_tops left a dead value on top" - ); - - let is_dead = |v: ValueId| !live_future.contains(v) && !live_out.contains(v); - let dead_run = stack - .iter() - .skip(1) - .take_while(|&v| matches!(v, StackItem::Value(v) if is_dead(*v))) - .count(); - if dead_run == 0 { - break; - } - - // Swap the top live value with the deepest dead value in the contiguous chain (within - // SWAP16 reach), then pop that chunk off. Repeat until the chain is gone. - let mut remaining = dead_run; - while remaining > 0 { - let swap_depth = remaining.min(reach.swap_max.saturating_sub(1)); - stack.swap(swap_depth, actions); - stack.pop_n(swap_depth, actions); - remaining -= swap_depth; - } - - if stack.len() == before_len { - break; - } - } -} - -pub(super) fn operand_order_for_stackify( - func: &Function, - inst: InstId, - value_aliases: &SecondaryMap>, -) -> SmallVec<[ValueId; 8]> { - let mut args: SmallVec<[ValueId; 8]> = func - .dfg - .inst(inst) - .collect_values() - .into_iter() - .map(|v| value_aliases[v].unwrap_or(v)) - .collect(); - - if call_has_local_return(func, inst) && !args.is_empty() { - // Rotation pairs with `Planner::prepare_internal_call`; see the ABI note there. - args.as_mut_slice().rotate_left(1); - } - - args -} - -fn call_has_local_return(func: &Function, inst: InstId) -> bool { - func.dfg.call_info(inst).is_some_and(|call| { - func.ctx() - .func_effects(call.callee()) - .may_return_to_caller() - }) -} - -pub(super) fn consume_operand_uses( - ctx: &StackifyContext<'_>, - args: &[ValueId], - remaining_uses: &mut BTreeMap, - live_future: &mut BitSet, - live_out: &BitSet, - scratch_slots: &SlotPool, - free_scratch_slots: &mut FreeSlots, -) { - for &v in args { - if ctx.retains_value(v) - && let Some(n) = remaining_uses.get_mut(&v) - { - let before = *n; - *n = n.saturating_sub(1); - if before != 0 && *n == 0 { - live_future.remove(v); - if !ctx.func.dfg.value_is_imm(v) && !live_out.contains(v) { - scratch_slots.release_if_assigned(v, free_scratch_slots); - } - } - } - } -} - -pub(super) fn last_use_values_in_inst( - ctx: &StackifyContext<'_>, - args: &[ValueId], - remaining_uses: &BTreeMap, - live_out: &BitSet, -) -> BitSet { - let mut inst_counts: BTreeMap = BTreeMap::new(); - for &v in args.iter() { - if !ctx.retains_value(v) { - continue; - } - *inst_counts.entry(v).or_insert(0) += 1; - } - - let mut last_use: BitSet = BitSet::default(); - for (v, count) in inst_counts { - let rem = remaining_uses.get(&v).copied().unwrap_or(0); - if rem == count && !live_out.contains(v) { - last_use.insert(v); - } - } - last_use -} - -pub(super) fn improve_reachability_before_operands( - values: ReachabilityValues<'_>, - args: &[ValueId], - reach: StackifyReachability, - stack: &mut SymStack, - live_future: &BitSet, - live_out: &BitSet, - actions: &mut Actions, -) { - const AGGRESSIVE_REACHABILITY_DEPTH: usize = 20; - - let mut protected_args: BitSet = BitSet::default(); - for &arg in args.iter() { - if values.retains(arg) { - protected_args.insert(arg); - } - } - - // If at least one operand is on the stack but unreachable by `DUP16`, attempt a more - // aggressive cleanup to bring it back into reach. This extends slightly past `SWAP16` reach - // by allowing deletions that shift the stack (e.g. popping dead/cheap values above an - // operand at depth 18-20). - let mut needs_aggressive = false; - for &arg in args.iter() { - if values.retains(arg) - && stack.find_reachable_value(arg, reach.dup_max).is_none() - && stack - .find_reachable_value(arg, AGGRESSIVE_REACHABILITY_DEPTH) - .is_some() - { - needs_aggressive = true; - break; - } - } - if !needs_aggressive { - return; - } - - // Bound the amount of cleanup we do per instruction to avoid pathological swap chains. - const MAX_DELETIONS: usize = 8; - let mut deletions: usize = 0; - - while deletions < MAX_DELETIONS { - let mut progressed = false; - - for &arg in args.iter() { - if values.retains(arg) - && stack.find_reachable_value(arg, reach.dup_max).is_none() - && let Some(pos) = stack.find_reachable_value(arg, AGGRESSIVE_REACHABILITY_DEPTH) - && let Some(victim) = choose_reachability_victim( - values, - stack, - pos, - &protected_args, - reach, - live_future, - live_out, - ) - { - stack.stable_delete_at_depth(victim + 1, actions); - deletions += 1; - progressed = true; - break; - } - } - - if !progressed { - break; - } - } -} - -fn choose_reachability_victim( - values: ReachabilityValues<'_>, - stack: &SymStack, - above: usize, - protected_args: &BitSet, - reach: StackifyReachability, - live_future: &BitSet, - live_out: &BitSet, -) -> Option { - let limit = stack.len_above_func_ret().min(reach.swap_max); - let above = above.min(limit); - - // 1) Prefer deleting dead values, starting from the shallowest depth to minimize `SWAP*` - // chains. This includes immediates: if they're dead, removing them cannot introduce new - // rematerialization cost. - for (i, item) in stack.iter().take(above).enumerate() { - if let StackItem::Value(v) = item - && !protected_args.contains(*v) - && !live_future.contains(*v) - && !live_out.contains(*v) - { - return Some(i); - } - } - - // 2) Then evict "cheap" immediates (they are always rematerializable). - for (i, item) in stack.iter().take(above).enumerate() { - if let StackItem::Value(v) = item - && !protected_args.contains(*v) - && values.func.dfg.value_is_imm(*v) - && is_evictable_imm(values.func, *v) - { - return Some(i); - } - } - - // 3) Then delete redundant duplicates of non-operands (keeping the shallowest copy). - let mut first_index: BTreeMap = BTreeMap::new(); - for (i, item) in stack.iter().take(above).enumerate() { - let StackItem::Value(v) = item else { - continue; - }; - first_index.entry(*v).or_insert(i); - } - - for (i, item) in stack.iter().take(above).enumerate() { - if let StackItem::Value(v) = item - && !protected_args.contains(*v) - && let Some(&first) = first_index.get(v) - && first != i - { - return Some(i); - } - } - - None -} - -fn is_evictable_imm(func: &Function, v: ValueId) -> bool { - const MAX_MATERIALIZATION_BYTES: usize = 3; // PUSH2 or smaller. - let Some(imm) = func.dfg.value_imm(v) else { - return false; - }; - immediate_materialization_code_len(imm) <= MAX_MATERIALIZATION_BYTES -} diff --git a/crates/codegen/src/stackalloc/stackify/mod.rs b/crates/codegen/src/stackalloc/stackify/mod.rs index 30d200f76..d819f019b 100644 --- a/crates/codegen/src/stackalloc/stackify/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/mod.rs @@ -29,17 +29,19 @@ //! into the caller's preserved stack segment. mod alloc; -mod block_sim; +mod block; mod br_table; mod builder; -mod iteration; +mod driver; mod planner; +mod rescue; mod slots; mod spill; mod sym_stack; mod templates; mod terminal_chain; mod trace; +mod uses; pub use alloc::StackifyAlloc; pub(crate) use builder::{ diff --git a/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs b/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs index 811eb76f7..b9969deae 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs @@ -163,7 +163,7 @@ mod tests { stackify::{ builder::StackifyReachability, planner::{ - MemPlan, NormalizeSearchScratch, Planner, + MemPlan, MemState, NormalizeSearchScratch, Planner, test_utils::build_stackify_test_context, }, slots::{FreeSlotPools, SpillSlotPools}, @@ -263,17 +263,15 @@ block1: ) .expect("source scratch slot"); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::opaque_prefix_empty(false); let mut actions = Actions::new(); let mut search_scratch = NormalizeSearchScratch::default(); @@ -383,17 +381,15 @@ block1: ) .expect("source scratch slot"); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::opaque_prefix_empty(false); let mut actions = Actions::new(); let mut search_scratch = NormalizeSearchScratch::default(); diff --git a/crates/codegen/src/stackalloc/stackify/planner/mod.rs b/crates/codegen/src/stackalloc/stackify/planner/mod.rs index 0fa8087b3..a6aac3444 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/mod.rs @@ -38,6 +38,21 @@ pub(super) fn must_use_object_storage( scratch_spill_slots == 0 || scratch_live_values.contains(v) || forced_object_spills.contains(v) } +/// Driver-owned memory-planning state, grouped so `MemPlan::new` takes it in one reborrow rather +/// than six separate arguments. Lives in `FunctionPlanner`; a `MemPlan` borrows its fields for the +/// span of a single instruction/edge plan. +pub(super) struct MemState<'a> { + pub(super) spill: SpillSet<'a>, + /// Provisional per-iteration object-id assignment (`assign_spill_obj_ids`), read by `MemPlan` + /// during planning before storage is finalized. + pub(super) spill_obj: + &'a SecondaryMap>, + pub(super) spill_requests: &'a mut BitSet, + pub(super) object_spill_requests: &'a mut BitSet, + pub(super) forced_object_spills: &'a BitSet, + pub(super) slots: &'a mut SpillSlotPools, +} + #[derive(Clone)] pub(super) struct MemPlanSnapshot { free_slots: FreeSlotPools, @@ -60,32 +75,23 @@ pub(super) struct MemPlan<'a> { } impl<'a> MemPlan<'a> { - #[allow(clippy::too_many_arguments)] pub(super) fn new( - spill: SpillSet<'a>, - spill_requests: &'a mut BitSet, + mem: &'a mut MemState<'_>, ctx: &'a StackifyContext<'_>, - spill_obj: &'a SecondaryMap< - ValueId, - Option, - >, exact_local_addr: &'a SecondaryMap>, - object_spill_requests: &'a mut BitSet, - forced_object_spills: &'a BitSet, free_slots: &'a mut FreeSlotPools, - slots: &'a mut SpillSlotPools, ) -> Self { Self { scratch_live_values: &ctx.scratch_live_values, scratch_spill_slots: ctx.scratch_spill_slots, - spill_obj, + spill_obj: mem.spill_obj, exact_local_addr, - object_spill_requests, - forced_object_spills, + object_spill_requests: &mut *mem.object_spill_requests, + forced_object_spills: mem.forced_object_spills, free_slots, - slots, + slots: &mut *mem.slots, spill_slot_interference: &ctx.spill_slot_interference, - spill: SpillDiscovery::new(spill, spill_requests), + spill: SpillDiscovery::new(mem.spill, &mut *mem.spill_requests), } } diff --git a/crates/codegen/src/stackalloc/stackify/planner/normalize_search.rs b/crates/codegen/src/stackalloc/stackify/planner/normalize_search.rs index c596b45f5..2ffffd1c9 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/normalize_search.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/normalize_search.rs @@ -4074,7 +4074,7 @@ mod tests { liveness::Liveness, stackalloc::stackify::{ builder::StackifyReachability, - planner::{MemPlan, Planner, test_utils::build_stackify_test_context}, + planner::{MemPlan, MemState, Planner, test_utils::build_stackify_test_context}, slots::{FreeSlotPools, SpillSlotPools}, spill::SpillSet, sym_stack::SymStack, @@ -5541,17 +5541,15 @@ func public %f() { let spill_obj = SecondaryMap::new(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::entry_stack(func, false); let mut actions = crate::stackalloc::Actions::new(); @@ -5594,17 +5592,15 @@ func public %f() { let spill_obj = SecondaryMap::new(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::entry_stack(func, false); let mut actions = crate::stackalloc::Actions::new(); diff --git a/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs b/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs index cfc13d876..e84598714 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs @@ -996,7 +996,7 @@ mod tests { stackify::{ builder::StackifyReachability, planner::{ - MemPlan, NormalizeSearchScratch, Planner, + MemPlan, MemState, NormalizeSearchScratch, Planner, normalize_search::{Cost, EstimatedCostModel, KeyInfo, SearchCfg, Step}, test_utils::build_stackify_test_context, }, @@ -1131,17 +1131,15 @@ block0: let spill_obj = SecondaryMap::new(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let old_args = [old_imm, imm2, imm3]; let current_args = [current_imm, imm2, imm3]; @@ -1242,17 +1240,15 @@ block0: let spill_obj = SecondaryMap::new(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::entry_stack(func, false); stack.push_value(arg); @@ -1318,17 +1314,15 @@ block0: let spill_obj = SecondaryMap::new(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let args: Vec<_> = func.arg_values.iter().copied().collect(); let mut stack = SymStack::entry_stack(func, false); @@ -1402,17 +1396,16 @@ block0: let forced_object_spills = BitSet::default(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = + MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::entry_stack(func, false); for &value in values.iter().rev() { stack.push_value(value); @@ -1480,17 +1473,16 @@ block0: let forced_object_spills = BitSet::default(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = + MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::entry_stack(func, false); stack.push_value(source); stack.push_value(top); @@ -1561,17 +1553,15 @@ block0: let forced_object_spills = BitSet::default(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::entry_stack(func, false); stack.push_value(stack_source); stack.push_value(ignored); @@ -1640,17 +1630,16 @@ block0: let forced_object_spills = BitSet::default(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = + MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::entry_stack(func, false); for &value in values.iter().rev() { stack.push_value(value); @@ -1721,17 +1710,16 @@ block0: let forced_object_spills = BitSet::default(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = + MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = SymStack::entry_stack(func, false); stack.push_value(ignored); stack.push_value(source); @@ -1798,17 +1786,16 @@ block0: let forced_object_spills = BitSet::default(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = + MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut actions = crate::stackalloc::Actions::new(); let mut search_scratch = NormalizeSearchScratch::default(); @@ -1903,17 +1890,16 @@ block0: let forced_object_spills = BitSet::default(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = + MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let mut stack = stack_from_top(stack_values); let mut args: SmallVec<[ValueId; 8]> = args_values.iter().copied().collect(); let mut last_use = BitSet::default(); @@ -2007,17 +1993,15 @@ block0: let spill_obj = SecondaryMap::new(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let args = [func.arg_values[1], func.arg_values[0]]; let mut stack = SymStack::entry_stack(func, false); @@ -2087,17 +2071,15 @@ block0: let spill_obj = SecondaryMap::new(); let mut free_slots = FreeSlotPools::default(); let mut slots = SpillSlotPools::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); let args = [func.arg_values[2]]; let mut stack = SymStack::entry_stack(func, false); @@ -2126,17 +2108,15 @@ block0: let mut actions = crate::stackalloc::Actions::new(); let mut object_spill_requests = BitSet::default(); let forced_object_spills = BitSet::default(); - let mem = MemPlan::new( - SpillSet::new(&spill_set), - &mut spill_requests, - &ctx, - &spill_obj, - &ctx.exact_local_addr, - &mut object_spill_requests, - &forced_object_spills, - &mut free_slots, - &mut slots, - ); + let mut mem_state = MemState { + spill: SpillSet::new(&spill_set), + spill_obj: &spill_obj, + spill_requests: &mut spill_requests, + object_spill_requests: &mut object_spill_requests, + forced_object_spills: &forced_object_spills, + slots: &mut slots, + }; + let mem = MemPlan::new(&mut mem_state, &ctx, &ctx.exact_local_addr, &mut free_slots); { let mut search_scratch = NormalizeSearchScratch::default(); let mut planner = diff --git a/crates/codegen/src/stackalloc/stackify/rescue.rs b/crates/codegen/src/stackalloc/stackify/rescue.rs new file mode 100644 index 000000000..eaae7d76d --- /dev/null +++ b/crates/codegen/src/stackalloc/stackify/rescue.rs @@ -0,0 +1,217 @@ +//! Purely local, bounded-effort stack cleanups run before operand preparation: +//! +//! - dead-prefix cleanup (`clean_dead_stack_prefix`): pop dead values off the top and swap a live +//! top past a contiguous dead chain beneath it; +//! - reachability rescue (`improve_reachability_before_operands`): when an operand sits just past +//! `DUP16` reach, evict a bounded number of dead / cheap / duplicate stack items to bring it +//! back into range. +//! +//! Both are opportunistic heuristics; the bounds below keep them from generating pathological swap +//! chains. + +use sonatina_ir::{Function, ValueId}; +use std::collections::BTreeMap; + +use crate::{bitset::BitSet, isa::evm::immediate_materialization_code_len, stackalloc::Actions}; + +use super::{ + builder::StackifyReachability, + sym_stack::{StackItem, SymStack}, + uses::UseTracker, +}; + +/// Reachability rescue looks slightly past `SWAP16` reach (e.g. an operand at depth 18-20) when +/// deciding whether an aggressive cleanup can bring it back into `DUP16` range. +const AGGRESSIVE_REACHABILITY_DEPTH: usize = 20; + +/// Upper bound on evictions per instruction, to avoid pathological swap chains. +const MAX_DELETIONS: usize = 8; + +fn pop_dead_tops( + stack: &mut SymStack, + live_future: &BitSet, + live_out: &BitSet, + actions: &mut Actions, +) { + while let Some(StackItem::Value(v)) = stack.top() { + if live_future.contains(*v) || live_out.contains(*v) { + break; + } + stack.pop(actions); + } +} + +pub(super) fn clean_dead_stack_prefix( + reach: StackifyReachability, + stack: &mut SymStack, + live_future: &BitSet, + live_out: &BitSet, + actions: &mut Actions, +) { + // Two local cleanups: + // - pop any dead values that reach the top + // - if a live value is on top and there is a contiguous dead chain directly beneath it, + // swap the live value down and pop the dead chain off the top. + loop { + let before_len = stack.len(); + pop_dead_tops(stack, live_future, live_out, actions); + + // If the top is not a normal value, don't try to reorder under it. + let Some(StackItem::Value(top)) = stack.top() else { + break; + }; + // `pop_dead_tops` only stops on a live value (or a non-`Value`, handled above). + debug_assert!( + live_future.contains(*top) || live_out.contains(*top), + "pop_dead_tops left a dead value on top" + ); + + let is_dead = |v: ValueId| !live_future.contains(v) && !live_out.contains(v); + let dead_run = stack + .iter() + .skip(1) + .take_while(|&v| matches!(v, StackItem::Value(v) if is_dead(*v))) + .count(); + if dead_run == 0 { + break; + } + + // Swap the top live value with the deepest dead value in the contiguous chain (within + // SWAP16 reach), then pop that chunk off. Repeat until the chain is gone. + let mut remaining = dead_run; + while remaining > 0 { + let swap_depth = remaining.min(reach.swap_max.saturating_sub(1)); + stack.swap(swap_depth, actions); + stack.pop_n(swap_depth, actions); + remaining -= swap_depth; + } + + if stack.len() == before_len { + break; + } + } +} + +pub(super) fn improve_reachability_before_operands( + func: &Function, + args: &[ValueId], + reach: StackifyReachability, + stack: &mut SymStack, + uses: &UseTracker, + actions: &mut Actions, +) { + let mut protected_args: BitSet = BitSet::default(); + for &arg in args.iter() { + if !func.dfg.value_is_imm(arg) { + protected_args.insert(arg); + } + } + + // If at least one operand is on the stack but unreachable by `DUP16`, attempt a more + // aggressive cleanup to bring it back into reach. This extends slightly past `SWAP16` reach + // by allowing deletions that shift the stack (e.g. popping dead/cheap values above an + // operand at depth 18-20). + let mut needs_aggressive = false; + for &arg in args.iter() { + if !func.dfg.value_is_imm(arg) + && stack.find_reachable_value(arg, reach.dup_max).is_none() + && stack + .find_reachable_value(arg, AGGRESSIVE_REACHABILITY_DEPTH) + .is_some() + { + needs_aggressive = true; + break; + } + } + if !needs_aggressive { + return; + } + + let mut deletions: usize = 0; + + while deletions < MAX_DELETIONS { + let mut progressed = false; + + for &arg in args.iter() { + if !func.dfg.value_is_imm(arg) + && stack.find_reachable_value(arg, reach.dup_max).is_none() + && let Some(pos) = stack.find_reachable_value(arg, AGGRESSIVE_REACHABILITY_DEPTH) + && let Some(victim) = + choose_reachability_victim(func, stack, pos, &protected_args, reach, uses) + { + stack.stable_delete_at_depth(victim + 1, actions); + deletions += 1; + progressed = true; + break; + } + } + + if !progressed { + break; + } + } +} + +fn choose_reachability_victim( + func: &Function, + stack: &SymStack, + above: usize, + protected_args: &BitSet, + reach: StackifyReachability, + uses: &UseTracker, +) -> Option { + let limit = stack.len_above_func_ret().min(reach.swap_max); + let above = above.min(limit); + + // 1) Prefer deleting dead values, starting from the shallowest depth to minimize `SWAP*` + // chains. This includes immediates: if they're dead, removing them cannot introduce new + // rematerialization cost. + for (i, item) in stack.iter().take(above).enumerate() { + if let StackItem::Value(v) = item + && !protected_args.contains(*v) + && uses.is_dead(*v) + { + return Some(i); + } + } + + // 2) Then evict "cheap" immediates (they are always rematerializable). + for (i, item) in stack.iter().take(above).enumerate() { + if let StackItem::Value(v) = item + && !protected_args.contains(*v) + && func.dfg.value_is_imm(*v) + && is_evictable_imm(func, *v) + { + return Some(i); + } + } + + // 3) Then delete redundant duplicates of non-operands (keeping the shallowest copy). + let mut first_index: BTreeMap = BTreeMap::new(); + for (i, item) in stack.iter().take(above).enumerate() { + let StackItem::Value(v) = item else { + continue; + }; + first_index.entry(*v).or_insert(i); + } + + for (i, item) in stack.iter().take(above).enumerate() { + if let StackItem::Value(v) = item + && !protected_args.contains(*v) + && let Some(&first) = first_index.get(v) + && first != i + { + return Some(i); + } + } + + None +} + +fn is_evictable_imm(func: &Function, v: ValueId) -> bool { + const MAX_MATERIALIZATION_BYTES: usize = 3; // PUSH2 or smaller. + let Some(imm) = func.dfg.value_imm(v) else { + return false; + }; + immediate_materialization_code_len(imm) <= MAX_MATERIALIZATION_BYTES +} diff --git a/crates/codegen/src/stackalloc/stackify/uses.rs b/crates/codegen/src/stackalloc/stackify/uses.rs new file mode 100644 index 000000000..03a595e92 --- /dev/null +++ b/crates/codegen/src/stackalloc/stackify/uses.rs @@ -0,0 +1,185 @@ +use sonatina_ir::{BlockId, I256, ValueId}; +use std::collections::BTreeMap; + +use crate::bitset::BitSet; + +use super::{ + block::operand_order_for_stackify, + builder::StackifyContext, + slots::{FreeSlots, SlotPool}, + sym_stack::{StackItem, SymStack}, +}; + +/// Per-block value-use bookkeeping threaded through the instruction walk. +/// +/// Tracks how many in-block uses remain for each retained value (`remaining`) and each cached +/// immediate (`cached_imm_remaining`), the set of values still used later in the block +/// (`live_future`), and the block's live-out set unioned with phi-out sources (`live_out`). +pub(super) struct UseTracker { + remaining: BTreeMap, + cached_imm_remaining: BTreeMap, + live_future: BitSet, + live_out: BitSet, +} + +impl UseTracker { + pub(super) fn for_block(ctx: &StackifyContext<'_>, block: BlockId) -> Self { + let mut remaining: BTreeMap = BTreeMap::new(); + let mut cached_imm_remaining: BTreeMap = BTreeMap::new(); + for inst in ctx.func.layout.iter_inst(block) { + if ctx.func.dfg.is_phi(inst) { + continue; + } + for v in operand_order_for_stackify(ctx.func, inst, &ctx.value_aliases) { + if ctx.retains_value(v) { + *remaining.entry(v).or_insert(0) += 1; + } else if ctx.stack_caches_immediate(v) + && let Some(imm) = ctx.func.dfg.value_imm(v) + { + *cached_imm_remaining.entry(imm.as_i256()).or_insert(0) += 1; + } + } + } + let live_future: BitSet = remaining.keys().copied().collect(); + let mut live_out = ctx.liveness.block_live_outs(block).clone(); + live_out.union_with(&ctx.phi_out_sources[block]); + Self { + remaining, + cached_imm_remaining, + live_future, + live_out, + } + } + + pub(super) fn live_future(&self) -> &BitSet { + &self.live_future + } + + pub(super) fn live_out(&self) -> &BitSet { + &self.live_out + } + + pub(super) fn is_dead(&self, v: ValueId) -> bool { + !self.live_future.contains(v) && !self.live_out.contains(v) + } + + /// Values in `args` whose in-block uses all end at this instruction (and which are not + /// live-out), so they may be consumed directly from the stack. + pub(super) fn last_uses_in( + &self, + ctx: &StackifyContext<'_>, + args: &[ValueId], + ) -> BitSet { + let mut inst_counts: BTreeMap = BTreeMap::new(); + for &v in args.iter() { + if !ctx.retains_value(v) { + continue; + } + *inst_counts.entry(v).or_insert(0) += 1; + } + + let mut last_use: BitSet = BitSet::default(); + for (v, count) in inst_counts { + let rem = self.remaining.get(&v).copied().unwrap_or(0); + if rem == count && !self.live_out.contains(v) { + last_use.insert(v); + } + } + last_use + } + + /// Cached immediates in `args` that have further in-block uses beyond this instruction and so + /// should be kept on the stack rather than consumed. + pub(super) fn cache_preserve_in( + &self, + ctx: &StackifyContext<'_>, + args: &[ValueId], + ) -> BitSet { + let mut inst_counts: BTreeMap = BTreeMap::new(); + for &value in args { + if ctx.stack_caches_immediate(value) + && let Some(imm) = ctx.func.dfg.value_imm(value) + { + *inst_counts.entry(imm.as_i256()).or_insert(0) += 1; + } + } + + let mut preserve: BitSet = BitSet::default(); + for &value in args { + if ctx.stack_caches_immediate(value) + && let Some(imm) = ctx.func.dfg.value_imm(value) + && self + .cached_imm_remaining + .get(&imm.as_i256()) + .copied() + .unwrap_or(0) + > inst_counts.get(&imm.as_i256()).copied().unwrap_or(0) + { + preserve.insert(value); + } + } + preserve + } + + /// Consume one use of each operand in `args`, releasing scratch slots for values that die here + /// and decaying cached-immediate counts. + pub(super) fn consume( + &mut self, + ctx: &StackifyContext<'_>, + args: &[ValueId], + scratch_slots: &SlotPool, + free_scratch: &mut FreeSlots, + ) { + for &v in args { + if ctx.retains_value(v) + && let Some(n) = self.remaining.get_mut(&v) + { + let before = *n; + *n = n.saturating_sub(1); + if before != 0 && *n == 0 { + self.live_future.remove(v); + if !ctx.func.dfg.value_is_imm(v) && !self.live_out.contains(v) { + scratch_slots.release_if_assigned(v, free_scratch); + } + } + } + } + + for &value in args { + if ctx.stack_caches_immediate(value) + && let Some(imm) = ctx.func.dfg.value_imm(value) + && let Some(count) = self.cached_imm_remaining.get_mut(&imm.as_i256()) + { + *count = count.saturating_sub(1); + } + } + self.cached_imm_remaining.retain(|_, count| *count != 0); + } + + /// `live_future`, extended with any cached immediates still on `stack` that have remaining + /// in-block uses, so dead-prefix cleanup does not drop a reusable cached immediate. + pub(super) fn cleanup_live( + &self, + ctx: &StackifyContext<'_>, + stack: &SymStack, + ) -> BitSet { + let mut live = self.live_future.clone(); + if self.cached_imm_remaining.is_empty() { + return live; + } + + for item in stack.iter() { + if let StackItem::Value(value) = item + && ctx.stack_caches_immediate(*value) + && let Some(imm) = ctx.func.dfg.value_imm(*value) + && self + .cached_imm_remaining + .get(&imm.as_i256()) + .is_some_and(|count| *count != 0) + { + live.insert(*value); + } + } + live + } +} From ce3ae52306b0b89e9972f77b46b03c80bbdf2b5d Mon Sep 17 00:00:00 2001 From: sbillig Date: Tue, 7 Jul 2026 01:26:35 -0700 Subject: [PATCH 09/14] Reify the stackify block-entry state machine New entry.rs owns the per-block EntryKind classification (function entry / terminal-chain opaque / single-pred inherited / merge) and the Pending -> Frozen entry lifecycle. EntryTable::record_edge is the one edge-dispatch point and EntryTable::freeze_entry the one transfer- selection point; frozen templates live in the entry state. This replaces the templates / inherited_stack / pending_edges / planned_blocks / terminal_chain side tables, and double-freezing a template is now structurally unreachable instead of silently ignored. plan_block becomes a short dispatch over ResolvedEntry. --- .../codegen/src/stackalloc/stackify/block.rs | 6 +- .../src/stackalloc/stackify/builder.rs | 27 +- .../codegen/src/stackalloc/stackify/driver.rs | 307 ++++++--------- .../codegen/src/stackalloc/stackify/entry.rs | 365 ++++++++++++++++++ crates/codegen/src/stackalloc/stackify/mod.rs | 4 + .../stackify/planner/control_flow.rs | 6 +- .../src/stackalloc/stackify/templates.rs | 18 +- 7 files changed, 507 insertions(+), 226 deletions(-) create mode 100644 crates/codegen/src/stackalloc/stackify/entry.rs diff --git a/crates/codegen/src/stackalloc/stackify/block.rs b/crates/codegen/src/stackalloc/stackify/block.rs index 39a38f245..cb206feaf 100644 --- a/crates/codegen/src/stackalloc/stackify/block.rs +++ b/crates/codegen/src/stackalloc/stackify/block.rs @@ -302,7 +302,7 @@ impl<'d, 'a, 'ctx, O: StackifyObserver> BlockPlanner<'d, 'a, 'ctx, O> { for succ in dests.iter().copied() { self.driver - .record_branch_edge(&mut self.state, succ, post_branch_stack.clone()); + .record_branch_edge(&mut self.state, succ, &post_branch_stack); } } @@ -338,11 +338,11 @@ impl<'d, 'a, 'ctx, O: StackifyObserver> BlockPlanner<'d, 'a, 'ctx, O> { for case in case_stacks { self.driver - .record_br_table_edge(&mut self.state, case.dest, case.post_compare_stack); + .record_br_table_edge(&mut self.state, case.dest, &case.post_compare_stack); } if let Some(default) = default { self.driver - .record_br_table_edge(&mut self.state, default, default_stack); + .record_br_table_edge(&mut self.state, default, &default_stack); } self.driver.observer().on_inst_br_table(inst); diff --git a/crates/codegen/src/stackalloc/stackify/builder.rs b/crates/codegen/src/stackalloc/stackify/builder.rs index 9adaee3ca..afbd88741 100644 --- a/crates/codegen/src/stackalloc/stackify/builder.rs +++ b/crates/codegen/src/stackalloc/stackify/builder.rs @@ -16,12 +16,13 @@ use super::{ alloc::{SpillStorage, StackifyAlloc}, block::operand_order_for_stackify, driver::FunctionPlanner, + entry::EntryTable, planner::{MemState, NormalizeSearchScratch, must_use_object_storage}, slots::{FreeSlotPools, SpillSlotInterference, SpillSlotPools}, spill::SpillSet, sym_stack::SymStack, templates::{ - BlockTemplate, DefInfo, compute_block_interfaces, compute_def_info, compute_dom_depth, + DefInfo, compute_block_interfaces, compute_def_info, compute_dom_depth, compute_phi_out_sources, compute_phi_results, function_has_internal_return, }, terminal_chain::compute_terminal_chain_blocks, @@ -444,8 +445,6 @@ impl<'a> StackifyBuilder<'a> { let spill_obj = assign_spill_obj_ids(ctx.func, spill, &ctx.exact_local_addr); let interfaces = compute_block_interfaces(ctx, spill); - let mut templates = initial_templates(ctx, &interfaces.params); - let mut alloc = StackifyAlloc { pre_actions: SecondaryMap::new(), post_actions: SecondaryMap::new(), @@ -457,14 +456,13 @@ impl<'a> StackifyBuilder<'a> { let mut spill_requests: BitSet = BitSet::default(); let terminal_chain_blocks = compute_terminal_chain_blocks(ctx, &interfaces); - // Blocks that are reached from multi-way branches inherit a dynamic stack and - // run an entry normalization prologue (single-pred only; critical edges split). - let mut inherited_stack: BTreeMap = BTreeMap::new(); + // The entry block enters with its ABI stack (function args ++ optional return address); the + // entry-state machine seeds this as the entry's inherited predecessor stack. let mut entry_stack = SymStack::entry_stack(ctx.func, ctx.has_internal_return); for (idx, &arg) in ctx.func.arg_values.iter().enumerate() { entry_stack.rename_value_at_depth(idx, ctx.canonicalize_value(arg)); } - inherited_stack.insert(ctx.entry, (ctx.entry, entry_stack)); + let entries = EntryTable::classify(ctx, &interfaces, &terminal_chain_blocks, entry_stack); let mem = MemState { spill, @@ -477,11 +475,9 @@ impl<'a> StackifyBuilder<'a> { let mut planner = FunctionPlanner::new( ctx, mem, - &mut templates, - &terminal_chain_blocks, &interfaces.carry_in, &mut alloc, - inherited_stack, + entries, search_scratch, observer, ); @@ -571,17 +567,6 @@ fn compute_hot_stack_cached_immediates( hot } -fn initial_templates( - ctx: &StackifyContext<'_>, - params: &SecondaryMap>, -) -> SecondaryMap { - let mut templates = SecondaryMap::new(); - for block in ctx.func.layout.iter_block() { - templates[block] = BlockTemplate::new(params[block].clone()); - } - templates -} - fn assign_spill_obj_ids( func: &Function, spill: SpillSet<'_>, diff --git a/crates/codegen/src/stackalloc/stackify/driver.rs b/crates/codegen/src/stackalloc/stackify/driver.rs index 52995f576..8bbd99783 100644 --- a/crates/codegen/src/stackalloc/stackify/driver.rs +++ b/crates/codegen/src/stackalloc/stackify/driver.rs @@ -1,6 +1,5 @@ use cranelift_entity::SecondaryMap; use sonatina_ir::{BlockId, InstId, ValueId}; -use std::collections::BTreeMap; use crate::{bitset::BitSet, stackalloc::Actions}; @@ -8,12 +7,10 @@ use super::{ alloc::StackifyAlloc, block::{BlockPlanner, BlockSimState, PlannerActionSink}, builder::StackifyContext, + entry::{DeferredEdge, EdgeDisposition, EntryTable, RecordEdge, ResolvedEntry}, planner::{self, MemState, NormalizeSearchScratch, Planner}, slots::{FreeSlotPools, SlotPool}, sym_stack::SymStack, - templates::{ - BlockTemplate, TransferOrder, canonical_transfer_order, choose_transfer, project_transfer, - }, trace::StackifyObserver, uses::UseTracker, }; @@ -23,51 +20,33 @@ pub(super) struct FunctionPlanner<'a, 'ctx, O: StackifyObserver> { /// Memory-planning state (spill set, provisional object ids, spill/object requests, slot /// pools) handed to each `MemPlan` in one reborrow. mem: MemState<'a>, - templates: &'a mut SecondaryMap, - terminal_chain_blocks: &'a BitSet, carry_in: &'a SecondaryMap>, alloc: &'a mut StackifyAlloc, - inherited_stack: BTreeMap, - pending_edges: BTreeMap>, - planned_blocks: BitSet, + /// The block-entry state machine: classification + `Pending -> Frozen` lifecycle, owning the + /// frozen templates. Replaces the old `templates` / `inherited_stack` / `pending_edges` / + /// `planned_blocks` / `terminal_chain_blocks` tables. + entries: EntryTable, search_scratch: &'a mut NormalizeSearchScratch, observer: &'a mut O, } -struct PendingEdge { - pred: BlockId, - inst: InstId, - stack: SymStack, - free_slots: FreeSlotPools, - action_start: usize, - /// Index of this edge's `DeferredExit` event in the observer's trace (0 for `NullObserver`), - /// used to backfill the exit fixup actions once the merge template is resolved. - trace_token: usize, -} - impl<'a, 'ctx, O: StackifyObserver> FunctionPlanner<'a, 'ctx, O> { #[allow(clippy::too_many_arguments)] pub(super) fn new( ctx: &'a StackifyContext<'ctx>, mem: MemState<'a>, - templates: &'a mut SecondaryMap, - terminal_chain_blocks: &'a BitSet, carry_in: &'a SecondaryMap>, alloc: &'a mut StackifyAlloc, - inherited_stack: BTreeMap, + entries: EntryTable, search_scratch: &'a mut NormalizeSearchScratch, observer: &'a mut O, ) -> Self { Self { ctx, mem, - templates, - terminal_chain_blocks, carry_in, alloc, - inherited_stack, - pending_edges: BTreeMap::new(), - planned_blocks: BitSet::default(), + entries, search_scratch, observer, } @@ -99,37 +78,65 @@ impl<'a, 'ctx, O: StackifyObserver> FunctionPlanner<'a, 'ctx, O> { self.plan_block(block); } - debug_assert!( - self.pending_edges.is_empty(), - "unresolved stackify edges remain" - ); + self.entries.debug_assert_no_pending_edges(); } fn plan_block(&mut self, block: BlockId) { let mut free_slots: FreeSlotPools = FreeSlotPools::default(); - let mut prologue: Actions = Actions::new(); let uses = UseTracker::for_block(self.ctx, block); - self.resolve_pending_edges(block); + let (stack, prologue) = self.resolve_entry(block, &uses, &mut free_slots); + + let state = BlockSimState::new(block, stack, free_slots, prologue, uses); + let state = BlockPlanner::new(self, state).run(); - let inherited = self.inherited_stack.remove(&block); - if self.terminal_chain_blocks.contains(block) { - self.freeze_template(block, TransferOrder::new()); - } else if let Some((_pred, stack)) = inherited.as_ref() { - self.freeze_template_from_stack(block, stack); - } else if block != self.ctx.entry { - self.freeze_template_canonical(block); + // The block walk injects the prologue on every non-phi instruction (including the + // terminator that every block has), so a non-empty prologue is always injected. + debug_assert!( + state.prologue.is_empty() || state.injected_prologue, + "prologue was not injected during the block walk" + ); + } + + /// Resolve `block`'s entry: freeze its template (via the entry table's single transfer-selection + /// point), plan the entry-specific fixups, and return the initial stack plus any prologue + /// actions. This subsumes the old `resolve_pending_edges` + the template/stack selection that + /// used to be inlined in `plan_block`. + fn resolve_entry( + &mut self, + block: BlockId, + uses: &UseTracker, + free_slots: &mut FreeSlotPools, + ) -> (SymStack, Actions) { + let mut prologue = Actions::new(); + let resolved = self + .entries + .freeze_entry(self.ctx, &self.carry_in[block], block); + + // Merge templates fix up their deferred edges *before* the block header fires, reproducing + // the old `resolve_pending_edges` ordering (deferred exit actions are traced ahead of + // `on_block_header`). + if let ResolvedEntry::Merge { deferred } = resolved { + self.plan_deferred_edges(block, deferred); + self.emit_block_header(block); + let stack = + SymStack::from_template(self.entries.template(block), self.ctx.has_internal_return); + return (stack, prologue); } - self.observer - .on_block_header(self.ctx.func, block, &self.templates[block]); - self.planned_blocks.insert(block); + self.emit_block_header(block); - let stack = if self.terminal_chain_blocks.contains(block) { - SymStack::opaque_prefix_empty(self.ctx.has_internal_return) - } else if let Some((pred, mut inh)) = inherited { - // Dynamic entry stack (single predecessor). - if block != self.ctx.entry { + let stack = match resolved { + ResolvedEntry::Merge { .. } => unreachable!("merge resolved above"), + ResolvedEntry::Opaque => SymStack::opaque_prefix_empty(self.ctx.has_internal_return), + ResolvedEntry::Fallback => { + SymStack::from_template(self.entries.template(block), self.ctx.has_internal_return) + } + ResolvedEntry::Entry { stack } => stack, + ResolvedEntry::Inherited { + pred, + stack: mut inh, + } => { debug_assert_eq!( self.ctx.cfg.pred_num_of(block), 1, @@ -143,68 +150,29 @@ impl<'a, 'ctx, O: StackifyObserver> FunctionPlanner<'a, 'ctx, O> { uses.live_future(), uses.live_out(), ); - let has_phi_params = !self.ctx.phi_results[block].is_empty(); - // Single-predecessor blocks without phis do not need exact entry - // template normalization. Keeping the inherited stack avoids pointless bottom - // reshuffling that can cascade into SWAP/POP churn. - if has_phi_params { - let tmpl = self.templates[block].clone(); - self.with_actions_planner( - &mut inh, - &mut prologue, - &mut free_slots, - |planner| { - planner.plan_edge_fixup_to_template(&tmpl, pred, block); - }, - ); + // Single-predecessor blocks without phis do not need exact entry template + // normalization. Keeping the inherited stack avoids pointless bottom reshuffling + // that can cascade into SWAP/POP churn. + if !self.ctx.phi_results[block].is_empty() { + let tmpl = self.entries.template(block).clone(); + self.with_actions_planner(&mut inh, &mut prologue, free_slots, |planner| { + planner.plan_edge_fixup_to_template(&tmpl, pred, block); + }); } self.observer.on_block_prologue(&prologue); + inh } - inh - } else { - SymStack::from_template(&self.templates[block], self.ctx.has_internal_return) }; - let state = BlockSimState::new(block, stack, free_slots, prologue, uses); - let state = BlockPlanner::new(self, state).run(); - - // The block walk injects the prologue on every non-phi instruction (including the - // terminator that every block has), so a non-empty prologue is always injected. - debug_assert!( - state.prologue.is_empty() || state.injected_prologue, - "prologue was not injected during the block walk" - ); + (stack, prologue) } - fn resolve_pending_edges(&mut self, block: BlockId) { - let Some(mut edges) = self.pending_edges.remove(&block) else { - return; - }; - debug_assert!( - !self.inherited_stack.contains_key(&block), - "pending merge edges cannot also inherit one stack" - ); - debug_assert!( - !self.planned_blocks.contains(block), - "pending edge target already planned" - ); - - let projected: Vec<(BlockId, TransferOrder)> = edges - .iter() - .map(|edge| { - ( - edge.pred, - project_transfer(&edge.stack, &self.carry_in[block]), - ) - }) - .collect(); - if !projected.is_empty() { - self.freeze_template(block, choose_transfer(self.ctx, block, &projected)); - } - - let tmpl = self.templates[block].clone(); - for edge in edges.iter_mut() { + /// Fix each deferred merge edge up to the now-frozen template, in arrival order, backfilling the + /// observer's deferred-exit trace. Mirrors the old `resolve_pending_edges` fixup loop. + fn plan_deferred_edges(&mut self, block: BlockId, mut deferred: Vec) { + let tmpl = self.entries.template(block).clone(); + for edge in deferred.iter_mut() { debug_assert_eq!( self.alloc.pre_actions[edge.inst].len(), edge.action_start, @@ -223,21 +191,9 @@ impl<'a, 'ctx, O: StackifyObserver> FunctionPlanner<'a, 'ctx, O> { } } - fn freeze_template_from_stack(&mut self, block: BlockId, stack: &SymStack) { - self.freeze_template(block, project_transfer(stack, &self.carry_in[block])); - } - - fn freeze_template_canonical(&mut self, block: BlockId) { - let transfer = canonical_transfer_order( - &self.carry_in[block], - &self.ctx.dom_depth, - &self.ctx.def_info, - ); - self.freeze_template(block, transfer); - } - - fn freeze_template(&mut self, block: BlockId, transfer: TransferOrder) { - self.templates[block].freeze_transfer(transfer); + fn emit_block_header(&mut self, block: BlockId) { + self.observer + .on_block_header(self.ctx.func, block, self.entries.template(block)); } pub(super) fn ctx(&self) -> &StackifyContext<'ctx> { @@ -316,11 +272,11 @@ impl<'a, 'ctx, O: StackifyObserver> FunctionPlanner<'a, 'ctx, O> { (&mut *self.observer, &*self.alloc) } - /// Record a `jump` edge to `dest`, applying immediate normalization / inheritance / pending - /// deferral as appropriate. Returns `true` when the edge was deferred (a pending merge edge): - /// its exit + jump trace has already been emitted here (synchronously with capturing the - /// deferred trace token), so the caller emits nothing further. Returns `false` for a resolved - /// edge, so the caller emits the exit-normalization actions and the jump trace. + /// Record a `jump` edge to `dest`. Returns `true` when the edge was deferred (a pending merge + /// edge): its exit + jump trace has already been emitted synchronously with capturing the + /// deferred trace token, so the caller emits nothing further. Returns `false` for an opaque, + /// inherited, or fixed-up-now edge, so the caller emits the exit-normalization actions and the + /// jump trace. pub(super) fn record_jump_edge( &mut self, state: &mut BlockSimState, @@ -328,86 +284,63 @@ impl<'a, 'ctx, O: StackifyObserver> FunctionPlanner<'a, 'ctx, O> { dest: BlockId, action_start: usize, ) -> bool { - if self.terminal_chain_blocks.contains(dest) { - } else if self.ctx.cfg.pred_num_of(dest) > 1 - && dest != self.ctx.entry - && !self.planned_blocks.contains(dest) - { - debug_assert!( - self.ctx.scc.is_reachable(dest), - "pending edge target must be reachable" - ); - self.pending_edges - .entry(dest) - .or_default() - .push(PendingEdge { - pred: state.block, - inst, - stack: state.stack.clone(), - free_slots: state.free_slots.clone(), - action_start, - trace_token: self.observer.on_deferred_inst_jump(inst, dest), - }); - return true; - } else if self.ctx.cfg.pred_num_of(dest) == 1 - && dest != self.ctx.entry - && !self.planned_blocks.contains(dest) - { - self.inherited_stack - .entry(dest) - .or_insert_with(|| (state.block, state.stack.clone())); - } else { - let tmpl = self.templates[dest].clone(); - let src = state.block; - self.with_planner( - &mut state.stack, - &mut state.free_slots, - PlannerActionSink::Pre(inst), - |planner| planner.plan_edge_fixup_to_template(&tmpl, src, dest), - ); + let disposition = self.entries.record_edge( + self.ctx, + self.observer, + RecordEdge::Jump { + inst, + free_slots: &state.free_slots, + action_start, + }, + state.block, + dest, + &state.stack, + ); + match disposition { + EdgeDisposition::Deferred => true, + EdgeDisposition::Recorded => false, + EdgeDisposition::FixupNow(tmpl) => { + let src = state.block; + self.with_planner( + &mut state.stack, + &mut state.free_slots, + PlannerActionSink::Pre(inst), + |planner| planner.plan_edge_fixup_to_template(&tmpl, src, dest), + ); + false + } } - false } pub(super) fn record_branch_edge( &mut self, state: &mut BlockSimState, succ: BlockId, - stack: SymStack, + stack: &SymStack, ) { - assert!( - !self.planned_blocks.contains(succ), - "multiway branch edge to already-planned block {succ:?}: run StackifyEdgeSplitter \ - before stackify to split in-cycle multiway edges" - ); - debug_assert_eq!( - self.ctx.cfg.pred_num_of(succ), - 1, - "no critical edges: branch target must be single-pred" + self.entries.record_edge( + self.ctx, + self.observer, + RecordEdge::Multiway { label: "branch" }, + state.block, + succ, + stack, ); - self.inherited_stack - .entry(succ) - .or_insert_with(|| (state.block, stack)); } pub(super) fn record_br_table_edge( &mut self, state: &mut BlockSimState, succ: BlockId, - stack: SymStack, + stack: &SymStack, ) { - assert!( - !self.planned_blocks.contains(succ), - "multiway br_table edge to already-planned block {succ:?}: run StackifyEdgeSplitter \ - before stackify to split in-cycle multiway edges" - ); - debug_assert_eq!( - self.ctx.cfg.pred_num_of(succ), - 1, - "no critical edges: br_table target must be single-pred" + self.entries.record_edge( + self.ctx, + self.observer, + RecordEdge::Multiway { label: "br_table" }, + state.block, + succ, + stack, ); - self.inherited_stack - .entry(succ) - .or_insert_with(|| (state.block, stack)); } } diff --git a/crates/codegen/src/stackalloc/stackify/entry.rs b/crates/codegen/src/stackalloc/stackify/entry.rs new file mode 100644 index 000000000..4595d7163 --- /dev/null +++ b/crates/codegen/src/stackalloc/stackify/entry.rs @@ -0,0 +1,365 @@ +//! The block-entry state machine. +//! +//! Every block obtains its entry stack `StackIn(B) = P(B) ++ T(B)` from exactly one of a few +//! sources. This module reifies that decision as a per-block [`EntryKind`] (a static +//! classification) plus an [`EntryState`] lifecycle (`Pending -> Frozen`), replacing the parallel +//! side tables (`terminal_chain_blocks`, `inherited_stack`, `pending_edges`, `templates`) and the +//! `planned_blocks` set that the driver used to consult. +//! +//! Two operations drive the lifecycle: +//! - [`EntryTable::record_edge`] is called at every terminator for every successor edge and decides +//! whether the edge is recorded for later (opaque no-op, single-pred inherit, deferred merge) or +//! must be fixed up now against an already-frozen template (backedge / planned-early block). +//! - [`EntryTable::freeze_entry`] is called once when a block is planned: it consumes the pending +//! state, selects the transfer region (the single place `choose_transfer` / `project_transfer` / +//! `canonical_transfer_order` are consulted), freezes the template, and hands the driver the +//! payload it needs to finish (deferred-edge fixups, the inherited prologue, the initial stack). +//! +//! Freezing is single-assignment by construction: the only `Pending -> Frozen` transition is in +//! `freeze_entry`, so the old `freeze_transfer` "second freeze is silently ignored" convention is +//! gone and a double freeze is unreachable rather than tolerated. + +use cranelift_entity::SecondaryMap; +use smallvec::SmallVec; +use sonatina_ir::{BlockId, InstId, ValueId}; + +use crate::bitset::BitSet; + +use super::{ + builder::StackifyContext, + slots::FreeSlotPools, + sym_stack::SymStack, + templates::{ + BlockInterfaces, BlockTemplate, TransferOrder, canonical_transfer_order, choose_transfer, + project_transfer, + }, + trace::StackifyObserver, +}; + +/// Static per-block classification, derived once per fixed-point iteration from the CFG plus the +/// terminal-chain analysis. Determines *which* rule supplies the block's entry stack. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)] +pub(super) enum EntryKind { + /// Function entry: the entry stack is the ABI stack (args ++ optional `FuncRetAddr`). + FunctionEntry, + /// Terminal chain: entered with an opaque stack; no normalization anywhere. + Opaque, + /// Exactly one predecessor: the entry stack is whatever that predecessor leaves. + #[default] + Inherited, + /// Two or more predecessors (jump-only, by the split-edge precondition): the entry stack is the + /// frozen template; each incoming jump normalizes to it. + Merge, +} + +/// A merge-block incoming jump whose exit-normalization fixup is deferred until the merge block is +/// planned (its template is frozen from the arrived predecessor stacks). Mirrors what the old +/// `PendingEdge` carried. +#[derive(Clone)] +pub(super) struct DeferredEdge { + pub(super) pred: BlockId, + pub(super) inst: InstId, + pub(super) stack: SymStack, + pub(super) free_slots: FreeSlotPools, + pub(super) action_start: usize, + /// Index of this edge's `DeferredExit` event in the observer's trace (0 for `NullObserver`), + /// used to backfill the exit fixup actions once the merge template is resolved. + pub(super) trace_token: usize, +} + +/// Dynamic entry-state lifecycle for one block within one fixed-point iteration. +#[derive(Clone)] +pub(super) enum EntryState { + /// Not planned yet; collects whatever has arrived. `inherited` is used by `FunctionEntry` + /// (seeded) and `Inherited` kinds; `deferred` by `Merge`. `params` is the stack-resident + /// parameter prefix, moved into the frozen template. + Pending { + params: SmallVec<[ValueId; 4]>, + inherited: Option<(BlockId, SymStack)>, + deferred: Vec, + }, + /// Block planned; template frozen. Later (back)edges fix up against this. + Frozen(BlockTemplate), +} + +impl Default for EntryState { + fn default() -> Self { + EntryState::Pending { + params: SmallVec::new(), + inherited: None, + deferred: Vec::new(), + } + } +} + +#[derive(Clone, Default)] +struct Entry { + kind: EntryKind, + state: EntryState, +} + +/// Which record path a terminator edge takes, carrying the path-specific data. +pub(super) enum RecordEdge<'a> { + /// Unconditional jump: may defer (forward merge), inherit (forward single-pred), or fix up now + /// (frozen dest / planned-early block). + Jump { + inst: InstId, + free_slots: &'a FreeSlotPools, + action_start: usize, + }, + /// Conditional branch / `br_table` edge. Critical edges are pre-split, so the target is + /// single-pred and always inherits; a frozen target means an unsplit in-cycle multiway edge. + /// `label` ("branch" / "br_table") only feeds the assert messages. + Multiway { label: &'static str }, +} + +/// What [`EntryTable::record_edge`] tells the caller to do at the terminator. +pub(super) enum EdgeDisposition { + /// Forward edge recorded for later (opaque no-op, single-pred inherit): caller emits the exit + + /// jump trace itself. + Recorded, + /// Forward merge edge deferred; its exit + jump trace was captured synchronously here, so the + /// caller emits nothing further. + Deferred, + /// Dest already frozen (backedge / planned-early): normalize to this template now. + FixupNow(BlockTemplate), +} + +/// What [`EntryTable::freeze_entry`] hands back to the driver so it can finish planning the entry +/// (run deferred-edge fixups, the inherited prologue, and build the initial stack). Transfer +/// selection and the `Pending -> Frozen` transition have already happened. +pub(super) enum ResolvedEntry { + /// Terminal chain: opaque entry stack, no normalization. + Opaque, + /// Function entry: the seeded ABI stack, used unchanged. + Entry { stack: SymStack }, + /// Single predecessor arrived: inherit its exit stack (with a prologue fixup if the block has + /// phi params). + Inherited { pred: BlockId, stack: SymStack }, + /// Nothing arrived before the block's RPO visit (irreducible / planned-early merge or + /// single-pred): the template was frozen from the canonical fallback; build the stack from it. + Fallback, + /// Merge with `>= 1` arrived edges: fix each deferred edge up to the frozen template, then build + /// the stack from it. + Merge { deferred: Vec }, +} + +pub(super) struct EntryTable { + entries: SecondaryMap, +} + +impl EntryTable { + /// Classify every block and seed the function entry's ABI stack. Runs once per fixed-point + /// iteration, after the terminal-chain analysis it consumes. + pub(super) fn classify( + ctx: &StackifyContext<'_>, + interfaces: &BlockInterfaces, + terminal_chain: &BitSet, + entry_stack: SymStack, + ) -> Self { + let mut entries: SecondaryMap = SecondaryMap::new(); + for block in ctx.func.layout.iter_block() { + let kind = if block == ctx.entry { + EntryKind::FunctionEntry + } else if terminal_chain.contains(block) { + EntryKind::Opaque + } else if ctx.cfg.pred_num_of(block) > 1 { + EntryKind::Merge + } else { + EntryKind::Inherited + }; + entries[block] = Entry { + kind, + state: EntryState::Pending { + params: interfaces.params[block].clone(), + inherited: None, + deferred: Vec::new(), + }, + }; + } + + // Seed the entry block's ABI stack as its (single) inherited predecessor stack. + if let EntryState::Pending { inherited, .. } = &mut entries[ctx.entry].state { + *inherited = Some((ctx.entry, entry_stack)); + } + + Self { entries } + } + + pub(super) fn kind(&self, block: BlockId) -> EntryKind { + self.entries[block].kind + } + + pub(super) fn is_frozen(&self, block: BlockId) -> bool { + matches!(self.entries[block].state, EntryState::Frozen(_)) + } + + /// The frozen template for `block`. Panics if the block has not been planned yet. + pub(super) fn template(&self, block: BlockId) -> &BlockTemplate { + match &self.entries[block].state { + EntryState::Frozen(tmpl) => tmpl, + EntryState::Pending { .. } => panic!("block template is not frozen"), + } + } + + /// Record one successor edge at a terminator; see [`EdgeDisposition`]. + pub(super) fn record_edge( + &mut self, + ctx: &StackifyContext<'_>, + observer: &mut O, + edge: RecordEdge<'_>, + pred: BlockId, + dest: BlockId, + stack: &SymStack, + ) -> EdgeDisposition { + match edge { + RecordEdge::Multiway { label } => { + // Critical edges are pre-split, so every multiway successor is single-pred and only + // ever inherits. An already-frozen (planned) target is an in-cycle multiway edge + // that `StackifyEdgeSplitter` should have split. + assert!( + !self.is_frozen(dest), + "multiway {label} edge to already-planned block {dest:?}: run \ + StackifyEdgeSplitter before stackify to split in-cycle multiway edges" + ); + debug_assert_eq!( + ctx.cfg.pred_num_of(dest), + 1, + "no critical edges: {label} target must be single-pred" + ); + self.set_inherited(dest, pred, stack); + EdgeDisposition::Recorded + } + RecordEdge::Jump { + inst, + free_slots, + action_start, + } => match self.kind(dest) { + // Terminal-chain dest: no normalization on any edge. + EntryKind::Opaque => EdgeDisposition::Recorded, + // Forward merge edge: defer the fixup until the merge template is frozen. + EntryKind::Merge if !self.is_frozen(dest) => { + debug_assert!( + ctx.scc.is_reachable(dest), + "pending edge target must be reachable" + ); + let trace_token = observer.on_deferred_inst_jump(inst, dest); + self.push_deferred( + dest, + DeferredEdge { + pred, + inst, + stack: stack.clone(), + free_slots: free_slots.clone(), + action_start, + trace_token, + }, + ); + EdgeDisposition::Deferred + } + // Forward single-pred edge: the successor inherits this exit stack. + EntryKind::Inherited if !self.is_frozen(dest) => { + self.set_inherited(dest, pred, stack); + EdgeDisposition::Recorded + } + // Frozen `Merge`/`Inherited` (backedge / planned-early), or the function entry + // (always frozen by the time an edge reaches it): fix up against the frozen template. + EntryKind::FunctionEntry | EntryKind::Merge | EntryKind::Inherited => { + EdgeDisposition::FixupNow(self.template(dest).clone()) + } + }, + } + } + + /// Store a single-pred block's inherited predecessor stack. First writer wins: with split + /// critical edges a single-pred block receives exactly one edge, except the degenerate + /// `br c b b` double edge whose two stacks are identical, so keeping the first is correct. + fn set_inherited(&mut self, block: BlockId, pred: BlockId, stack: &SymStack) { + match &mut self.entries[block].state { + EntryState::Pending { inherited, .. } => { + inherited.get_or_insert_with(|| (pred, stack.clone())); + } + EntryState::Frozen(_) => unreachable!("inherit into already-frozen block {block:?}"), + } + } + + fn push_deferred(&mut self, block: BlockId, edge: DeferredEdge) { + match &mut self.entries[block].state { + EntryState::Pending { deferred, .. } => deferred.push(edge), + EntryState::Frozen(_) => unreachable!("defer edge to already-frozen block {block:?}"), + } + } + + /// Consume `block`'s pending state, select its transfer region, freeze its template, and return + /// the payload the driver needs to finish planning the entry. This is the single place transfer + /// selection happens (`project_transfer` / `choose_transfer` / `canonical_transfer_order`). + pub(super) fn freeze_entry( + &mut self, + ctx: &StackifyContext<'_>, + carry_in: &BitSet, + block: BlockId, + ) -> ResolvedEntry { + let kind = self.entries[block].kind; + let EntryState::Pending { + params, + inherited, + deferred, + } = std::mem::take(&mut self.entries[block].state) + else { + panic!("block {block:?} entry already frozen"); + }; + + let (transfer, resolved) = match kind { + EntryKind::Opaque => (TransferOrder::new(), ResolvedEntry::Opaque), + EntryKind::FunctionEntry => { + let (_, stack) = inherited.expect("function entry stack seeded during classify"); + let transfer = project_transfer(&stack, carry_in); + (transfer, ResolvedEntry::Entry { stack }) + } + EntryKind::Inherited => match inherited { + Some((pred, stack)) => { + let transfer = project_transfer(&stack, carry_in); + (transfer, ResolvedEntry::Inherited { pred, stack }) + } + None => ( + canonical_transfer_order(carry_in, &ctx.dom_depth, &ctx.def_info), + ResolvedEntry::Fallback, + ), + }, + EntryKind::Merge => { + debug_assert!( + inherited.is_none(), + "pending merge edges cannot also inherit one stack" + ); + if deferred.is_empty() { + ( + canonical_transfer_order(carry_in, &ctx.dom_depth, &ctx.def_info), + ResolvedEntry::Fallback, + ) + } else { + let projected: Vec<(BlockId, TransferOrder)> = deferred + .iter() + .map(|edge| (edge.pred, project_transfer(&edge.stack, carry_in))) + .collect(); + let transfer = choose_transfer(ctx, block, &projected); + (transfer, ResolvedEntry::Merge { deferred }) + } + } + }; + + self.entries[block].state = EntryState::Frozen(BlockTemplate::new(params, transfer)); + resolved + } + + /// After `plan_blocks`, every planned block is frozen and no pending block still holds deferred + /// merge edges (the old "unresolved stackify edges remain" invariant). + pub(super) fn debug_assert_no_pending_edges(&self) { + debug_assert!( + self.entries.values().all(|entry| match &entry.state { + EntryState::Frozen(_) => true, + EntryState::Pending { deferred, .. } => deferred.is_empty(), + }), + "unresolved stackify edges remain" + ); + } +} diff --git a/crates/codegen/src/stackalloc/stackify/mod.rs b/crates/codegen/src/stackalloc/stackify/mod.rs index d819f019b..34eed71e3 100644 --- a/crates/codegen/src/stackalloc/stackify/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/mod.rs @@ -7,6 +7,9 @@ //! - `T(B)` is a transfer region: live-in, non-phi values in a chosen order. //! - `T(B)` is frozen from the first real predecessor stack, a chosen pending merge stack, or //! a deterministic fallback when the block must be emitted before any predecessor is planned. +//! - How each block obtains its entry stack (opaque terminal chain / inherited single-pred / +//! frozen merge / function entry) and the `Pending -> Frozen` freeze lifecycle live in one +//! place: the block-entry state machine in `entry.rs` (`record_edge` + `resolve_entry`). //! - For merge blocks, all incoming edges are normalized to the same `StackIn(B)` (often a no-op); //! spilled phi results are stored directly on the incoming edge instead of being carried in //! `P(B)`. @@ -33,6 +36,7 @@ mod block; mod br_table; mod builder; mod driver; +mod entry; mod planner; mod rescue; mod slots; diff --git a/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs b/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs index b9969deae..086972835 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/control_flow.rs @@ -278,8 +278,7 @@ block1: let mut planner = Planner::new(&ctx, &mut stack, &mut actions, mem, &mut search_scratch); - let mut template = BlockTemplate::new(smallvec![]); - template.freeze_transfer(smallvec![]); + let template = BlockTemplate::new(smallvec![], smallvec![]); planner.plan_edge_fixup_to_template(&template, BlockId(0), BlockId(1)); assert_eq!( @@ -396,8 +395,7 @@ block1: let mut planner = Planner::new(&ctx, &mut stack, &mut actions, mem, &mut search_scratch); - let mut template = BlockTemplate::new(smallvec![stack_phi]); - template.freeze_transfer(smallvec![]); + let template = BlockTemplate::new(smallvec![stack_phi], smallvec![]); planner.plan_edge_fixup_to_template(&template, BlockId(0), BlockId(1)); assert_eq!( diff --git a/crates/codegen/src/stackalloc/stackify/templates.rs b/crates/codegen/src/stackalloc/stackify/templates.rs index 77e57c128..9a57e555f 100644 --- a/crates/codegen/src/stackalloc/stackify/templates.rs +++ b/crates/codegen/src/stackalloc/stackify/templates.rs @@ -19,28 +19,24 @@ pub(super) struct DefInfo { def_index: u32, } +/// A block's frozen entry template `StackIn(B) = P(B) ++ T(B)`. Constructed frozen (transfer +/// already selected); the freeze lifecycle lives in `entry.rs`, so there is no lazy "freeze +/// later" step here. #[derive(Clone, Debug, Default)] pub(super) struct BlockTemplate { /// Stack-resident parameter prefix (entry args; non-spilled phi results elsewhere). pub(super) params: SmallVec<[ValueId; 4]>, /// Transfer region (top-first). - transfer: Option, + transfer: TransferOrder, } impl BlockTemplate { - pub(super) fn new(params: SmallVec<[ValueId; 4]>) -> Self { - Self { - params, - transfer: None, - } + pub(super) fn new(params: SmallVec<[ValueId; 4]>, transfer: TransferOrder) -> Self { + Self { params, transfer } } pub(super) fn transfer(&self) -> &TransferOrder { - self.transfer.as_ref().expect("block template is frozen") - } - - pub(super) fn freeze_transfer(&mut self, transfer: TransferOrder) { - self.transfer.get_or_insert(transfer); + &self.transfer } } From a70c0805d552b4064d7c94d7bd8d6fe108f00263 Mon Sep 17 00:00:00 2001 From: sbillig Date: Tue, 7 Jul 2026 01:40:27 -0700 Subject: [PATCH 10/14] Canonicalize numerically-equal immediates to one ValueId Stackify already treats equal as_i256 words as interchangeable stack items (rename_immediate_slots_to_match, semantic suffix matching). Give each distinct word one representative ValueId via the existing value-alias mechanism, choosing the class member with the cheapest materialization (then lowest id) so a class never pushes worse bytes than its best member. This is stage 1 of folding the parallel immediate-identity machinery into ordinary ValueId equality. --- .../src/stackalloc/stackify/builder.rs | 125 ++++++++++++++++++ 1 file changed, 125 insertions(+) diff --git a/crates/codegen/src/stackalloc/stackify/builder.rs b/crates/codegen/src/stackalloc/stackify/builder.rs index afbd88741..3ee4b1a71 100644 --- a/crates/codegen/src/stackalloc/stackify/builder.rs +++ b/crates/codegen/src/stackalloc/stackify/builder.rs @@ -312,6 +312,11 @@ impl<'a> StackifyBuilder<'a> { aliases }; normalize_value_alias_map(self.func, &mut value_aliases); + // Collapse numerically-equal immediates onto one representative ValueId, then re-normalize. + // The imm pass keeps the map one-hop canonical, so re-running is idempotent; it re-validates + // the invariant in debug builds over the post-canonicalization map. + canonicalize_immediate_aliases(self.func, &mut value_aliases); + normalize_value_alias_map(self.func, &mut value_aliases); let mut stack_cached_immediates = self.stack_cached_immediates; if self.cache_hot_immediates { @@ -567,6 +572,49 @@ fn compute_hot_stack_cached_immediates( hot } +/// Aliases every immediate value to one canonical representative `ValueId` per distinct materialized +/// 256-bit word. +/// +/// Numerically equal immediates are interchangeable on the EVM stack: equal `as_i256()` means an +/// identical stack word (types like `I8(-1)` and `I256(-1)` sign-extend to the same word), which the +/// materialization/rename machinery already relies on. Giving them one identity lets ordinary +/// `ValueId` equality subsume the semantic-equality checks scattered across stackify. +/// +/// The representative is the class member with the cheapest materialization, then the lowest id — +/// deterministic, and the whole class is pushed via `Action::Push(rep's Immediate)`, so it never +/// materializes worse than its best member (equal-word plans can differ across immediate types, +/// e.g. `I256(-1)` gets a compact `NOT`-based plan that `I8(-1)` does not). The canonical `c` is +/// read first, and imm-ness is keyed off `c`, so pre-existing (e.g. GVN-provided) aliases are +/// respected and folded through. +/// +/// Assumes `value_aliases` is already one-hop canonical (post `normalize_value_alias_map`). The pass +/// preserves that: each re-aliased `v` points to `rep`, which is the canonical of some value and is +/// itself never re-aliased (it is its own word's representative), hence self-canonical. +fn canonicalize_immediate_aliases( + func: &Function, + value_aliases: &mut SecondaryMap>, +) { + let mut rep_of: BTreeMap = BTreeMap::new(); + for v in func.dfg.value_ids() { + let c = value_aliases[v].unwrap_or(v); + if let Some(imm) = func.dfg.value_imm(c) { + let key = (immediate_materialization_code_len(imm), c); + let entry = rep_of.entry(imm.as_i256()).or_insert(key); + *entry = key.min(*entry); + } + } + + for v in func.dfg.value_ids() { + let c = value_aliases[v].unwrap_or(v); + if let Some(imm) = func.dfg.value_imm(c) { + let (_, rep) = rep_of[&imm.as_i256()]; + if rep != c { + value_aliases[v] = Some(rep); + } + } + } +} + fn assign_spill_obj_ids( func: &Function, spill: SpillSet<'_>, @@ -770,6 +818,83 @@ block0: }); } + #[test] + fn canonicalize_immediates_collapses_numerically_equal_differently_typed_immediates() { + use sonatina_ir::I256; + + // The DFG interns immediates per `Immediate` (type-inclusive), so `1.i8` and `1.i256` become + // distinct ValueIds that share `as_i256() == 1`. Stage-1 canonicalization must collapse them + // onto one representative, and the resulting map must stay one-hop canonical. + const SRC: &str = r#" +target = "evm-ethereum-osaka" + +func public %f(v0.i256, v1.i8) -> i256 { +block0: + v2.i8 = add v1 1.i8; + v3.i256 = zext v2 i256; + v4.i256 = add v0 1.i256; + v5.i256 = add v3 v4; + return v5; +} +"#; + + let parsed = parse_module(SRC).expect("module parses"); + let func_ref = parsed.debug.func_order[0]; + + parsed.module.func_store.view(func_ref, |func| { + // Collect the immediate value ids whose 256-bit word is 1, lowest id first. + let mut imm_ones: Vec<_> = func + .dfg + .value_ids() + .filter(|&v| { + func.dfg + .value_imm(v) + .is_some_and(|imm| imm.as_i256() == I256::one()) + }) + .collect(); + imm_ones.sort_unstable_by_key(|v| v.as_u32()); + + // The scenario is real: two distinct ValueIds, equal word, different types. + assert_eq!( + imm_ones.len(), + 2, + "parser should intern 1.i8 and 1.i256 as distinct value ids" + ); + let ty0 = func.dfg.value_imm(imm_ones[0]).unwrap().ty(); + let ty1 = func.dfg.value_imm(imm_ones[1]).unwrap().ty(); + assert_ne!(ty0, ty1, "the two immediates must have different types"); + + let mut aliases: SecondaryMap<_, Option<_>> = SecondaryMap::new(); + for value in func.dfg.value_ids() { + aliases[value] = Some(value); + } + normalize_value_alias_map(func, &mut aliases); + super::canonicalize_immediate_aliases(func, &mut aliases); + + // Both immediates canonicalize to the lowest-id representative. + let rep = imm_ones[0]; + for &v in &imm_ones { + assert_eq!(aliases[v].unwrap_or(v), rep); + } + + // The map remains one-hop canonical: re-running normalization is a no-op. + let before = aliases.clone(); + normalize_value_alias_map(func, &mut aliases); + for value in func.dfg.value_ids() { + assert_eq!(aliases[value], before[value]); + } + + // compute() integrates the canonicalization end-to-end without panicking. + let mut cfg = ControlFlowGraph::new(); + cfg.compute(func); + let mut liveness = Liveness::new(); + liveness.compute(func, &cfg); + let mut dom = DomTree::new(); + dom.compute(&cfg); + let _ = StackifyBuilder::new(func, &cfg, &dom, &liveness, 16).compute(); + }); + } + #[test] fn hot_immediate_caching_can_use_size_mode_use_threshold() { const BIG: &str = "340282366920938463463374607431768211454"; From f81e149e1b07998725470d2032dbe87d3da326ea Mon Sep 17 00:00:00 2001 From: sbillig Date: Tue, 7 Jul 2026 01:53:31 -0700 Subject: [PATCH 11/14] Replace immediate semantic-equality shims with ValueId equality With one canonical ValueId per immediate word, plain equality subsumes the word-level matching: delete rename_immediate_slots_to_match, common_suffix_len_semantic, operand_prep_item_matches_arg, and unary_operand_prep_find_arg; plan-replay mismatches still roll back to the greedy fallback, and flush_rebuild asserts exact equality. --- .../stackalloc/stackify/planner/normalize.rs | 89 ++----------- .../stackify/planner/operand_prep.rs | 120 +++++------------- 2 files changed, 37 insertions(+), 172 deletions(-) diff --git a/crates/codegen/src/stackalloc/stackify/planner/normalize.rs b/crates/codegen/src/stackalloc/stackify/planner/normalize.rs index 3d406a12e..9a9805187 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/normalize.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/normalize.rs @@ -15,36 +15,6 @@ use std::sync::{ }; impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { - pub(super) fn rename_immediate_slots_to_match(&mut self, desired: &[ValueId]) -> bool { - for (depth, &want) in desired.iter().enumerate() { - if self.ctx.func.dfg.value_is_imm(want) { - let want_imm = self - .ctx - .func - .dfg - .value_imm(want) - .expect("imm value missing payload") - .as_i256(); - let Some(StackItem::Value(cur)) = self.stack.item_at(depth) else { - return false; - }; - let cur_imm = self - .ctx - .func - .dfg - .value_imm(*cur) - .expect("expected immediate value on stack") - .as_i256(); - if cur_imm != want_imm { - return false; - } - self.stack.rename_value_at_depth(depth, want); - } - } - - true - } - pub(super) fn normalize_to_exact(&mut self, desired: &[ValueId]) { // Contract: rewrite the current symbolic stack (above the function return barrier, if any) // so that it matches `desired` exactly (top-first). @@ -75,7 +45,9 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { } } - if !self.rename_immediate_slots_to_match(desired) || !matches_exact(self.stack, desired) { + // Immediates are canonicalized to one ValueId per word, so a successful replay leaves the + // stack equal to `desired` by plain ValueId equality; any mismatch falls back. + if !matches_exact(self.stack, desired) { *self.stack = stack_before; self.actions.truncate(actions_before); self.mem.restore(mem_before); @@ -125,7 +97,7 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { }; if plan.is_none() { - let base_len = common_suffix_len_semantic(self.ctx.func, self.stack, desired); + let base_len = self.stack.common_suffix_len(desired); let start_repair = limit.saturating_sub(base_len); let goal_repair = desired.len().saturating_sub(base_len); @@ -190,7 +162,7 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { } fn flush_rebuild(&mut self, desired: &[ValueId]) { - let base_len = common_suffix_len_semantic(self.ctx.func, self.stack, desired); + let base_len = self.stack.common_suffix_len(desired); // Pop everything above the common base suffix. while self.stack.len_above_func_ret() > base_len { @@ -215,11 +187,10 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { } } - assert!( - self.rename_immediate_slots_to_match(desired), - "flush_rebuild failed to restore immediate value ids" + debug_assert!( + matches_exact(self.stack, desired), + "flush_rebuild must leave the stack equal to `desired`" ); - debug_assert!(matches_exact(self.stack, desired)); } } @@ -303,47 +274,3 @@ fn dump_failed_normalization(stack: &SymStack, desired: &[ValueId]) { eprintln!("normalize_to_exact: start={start:?}"); eprintln!("normalize_to_exact: desired={desired:?}"); } - -fn common_suffix_len_semantic( - func: &sonatina_ir::Function, - stack: &SymStack, - desired: &[ValueId], -) -> usize { - let limit = stack.len_above_func_ret(); - let max = limit.min(desired.len()); - let mut k = 0usize; - - for off in 0..max { - let depth = limit - 1 - off; - let want = desired[desired.len() - 1 - off]; - - let Some(item) = stack.item_at(depth) else { - break; - }; - let StackItem::Value(cur) = item else { - break; - }; - let cur = *cur; - - if func.dfg.value_is_imm(want) { - if !func.dfg.value_is_imm(cur) { - break; - } - let Some(want_imm) = func.dfg.value_imm(want) else { - break; - }; - let Some(cur_imm) = func.dfg.value_imm(cur) else { - break; - }; - if want_imm.as_i256() != cur_imm.as_i256() { - break; - } - } else if cur != want { - break; - } - - k += 1; - } - - k -} diff --git a/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs b/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs index e84598714..70958e70c 100644 --- a/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs +++ b/crates/codegen/src/stackalloc/stackify/planner/operand_prep.rs @@ -214,12 +214,8 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { continue; } - let arg_imm = self.ctx.func.dfg.value_imm(arg); - let found_in_tail = (window_len..start_limit).any(|depth| { - self.stack - .item_at(depth) - .is_some_and(|item| self.operand_prep_item_matches_arg(item, arg, arg_imm)) - }); + let found_in_tail = (window_len..start_limit) + .any(|depth| self.stack.item_at(depth) == Some(&StackItem::Value(arg))); if found_in_tail { mask |= 1u64 << idx; } @@ -411,11 +407,8 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { let last_use = consume_last_use.contains(arg); let preserve_needed = cache_preserve.contains(arg) || (self.ctx.retains_value(arg) && !last_use); - let copy_count = self.unary_operand_prep_copy_count(arg, arg_imm); - let top_matches = self - .stack - .top() - .is_some_and(|item| self.operand_prep_item_matches_arg(item, arg, arg_imm)); + let copy_count = self.unary_operand_prep_copy_count(arg); + let top_matches = self.stack.top() == Some(&StackItem::Value(arg)); let preserve_satisfied = !preserve_needed || copy_count >= 2; let surplus_last_use_penalty = cost.cost_pop().saturating_add(cost.cost_swap(1)); @@ -453,7 +446,7 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { }); } - if let Some(pos) = self.unary_operand_prep_find_arg(arg, arg_imm, 0, search_cfg.dup_max) { + if let Some(pos) = self.stack.find_reachable_value(arg, search_cfg.dup_max) { let emitted_cost = cost.cost_dup(pos as u8); candidates.push(UnaryOperandPrepCandidate { modeled_cost: copy_cost(emitted_cost), @@ -463,7 +456,7 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { }); } - if let Some(pos) = self.unary_operand_prep_find_arg(arg, arg_imm, 0, search_cfg.swap_max) + if let Some(pos) = self.stack.find_reachable_value(arg, search_cfg.swap_max) && pos != 0 && preserve_satisfied { @@ -479,12 +472,9 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { if !arg_is_imm && !last_use && copy_count < 2 - && let Some(pos) = self.unary_operand_prep_find_arg( - arg, - arg_imm, - search_cfg.dup_max, - search_cfg.swap_max, - ) + && let Some(pos) = + self.stack + .find_reachable_value_from(arg, search_cfg.dup_max, search_cfg.swap_max) { let emitted_cost = cost.cost_swap(pos as u8).saturating_add(cost.cost_dup(0)); candidates.push(UnaryOperandPrepCandidate { @@ -529,56 +519,14 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { } } - fn operand_prep_item_matches_arg( - &self, - item: &StackItem, - arg: ValueId, - arg_imm: Option, - ) -> bool { - let StackItem::Value(value) = *item else { - return false; - }; - - if let Some(arg_imm) = arg_imm { - return self - .ctx - .func - .dfg - .value_imm(value) - .is_some_and(|imm| imm.as_i256() == arg_imm.as_i256()); - } - - value == arg - } - - fn unary_operand_prep_copy_count(&self, arg: ValueId, arg_imm: Option) -> usize { + fn unary_operand_prep_copy_count(&self, arg: ValueId) -> usize { self.stack .iter() .take(self.stack.len_above_func_ret()) - .filter(|item| self.operand_prep_item_matches_arg(item, arg, arg_imm)) + .filter(|&item| item == &StackItem::Value(arg)) .count() } - fn unary_operand_prep_find_arg( - &self, - arg: ValueId, - arg_imm: Option, - start: usize, - max_depth: usize, - ) -> Option { - let limit = self.stack.len_above_func_ret().min(max_depth); - if start >= limit { - return None; - } - - self.stack - .iter() - .skip(start) - .take(limit - start) - .position(|item| self.operand_prep_item_matches_arg(item, arg, arg_imm)) - .map(|off| start + off) - } - fn inst_is_commutative(&self, inst: InstId) -> bool { use sonatina_ir::{ InstDowncast, @@ -708,9 +656,7 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { } fn stack_item_matches_arg(&self, depth: usize, arg: ValueId) -> bool { - self.stack.item_at(depth).is_some_and(|item| { - self.operand_prep_item_matches_arg(item, arg, self.ctx.func.dfg.value_imm(arg)) - }) + self.stack.item_at(depth) == Some(&StackItem::Value(arg)) } fn stack_prefix_matches_and_preserved( @@ -719,16 +665,10 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { consume_last_use: &BitSet, cache_preserve: &BitSet, ) -> bool { - if args.is_empty() || self.stack.len_above_func_ret() < args.len() { + if args.is_empty() || !self.stack_prefix_matches(args) { return false; } - for (depth, &arg) in args.iter().enumerate() { - if !self.stack_item_matches_arg(depth, arg) { - return false; - } - } - let mut checked = BitSet::default(); let stack_len = self.stack.len_above_func_ret(); for &arg in args { @@ -739,19 +679,20 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { continue; } - let arg_imm = self.ctx.func.dfg.value_imm(arg); let preserved_on_stack = self .stack .iter() .take(stack_len) .skip(args.len()) - .any(|item| self.operand_prep_item_matches_arg(item, arg, arg_imm)); - if !preserved_on_stack && (arg_imm.is_some() || !self.mem.spill_set().contains(arg)) { + .any(|item| item == &StackItem::Value(arg)); + if !preserved_on_stack + && (self.ctx.func.dfg.value_is_imm(arg) || !self.mem.spill_set().contains(arg)) + { return false; } } - self.rename_immediate_slots_to_match(args) + true } fn prepare_trivial_binary_operands( @@ -831,7 +772,9 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { } } - if !self.rename_immediate_slots_to_match(args) || !self.stack_prefix_matches(args) { + // Immediates are canonicalized to one ValueId per word, so a successful replay leaves the + // operand prefix equal to `args` by plain ValueId equality; any mismatch falls back. + if !self.stack_prefix_matches(args) { *self.stack = stack_before; self.actions.truncate(actions_before); self.mem.restore(mem_before); @@ -917,12 +860,9 @@ impl<'a, 'ctx: 'a> Planner<'a, 'ctx> { .value_imm(v) .expect("imm value missing payload"); if self.ctx.stack_caches_immediate(v) - && let Some(pos) = self.unary_operand_prep_find_arg( - v, - Some(imm), - prepared, - self.ctx.reach.dup_max, - ) + && let Some(pos) = + self.stack + .find_reachable_value_from(v, prepared, self.ctx.reach.dup_max) { self.stack.dup(pos, self.actions); prepared += 1; @@ -1857,11 +1797,7 @@ block0: let func_ref = parsed.debug.func_order[0]; parsed.module.func_store.modify(func_ref, |func| { - let old_imm = func.dfg.make_imm_value(Immediate::I8(7)); - let current_imm = func.dfg.make_value(Value::Immediate { - imm: Immediate::I8(7), - ty: Type::I8, - }); + let imm = func.dfg.make_imm_value(Immediate::I8(7)); let mut cfg = ControlFlowGraph::new(); cfg.compute(func); @@ -1949,8 +1885,10 @@ block0: let (_, _, actions, _) = run(&[*y, *x, *x, *y], &[*x, *y], &[], &[], false); assert_eq!(actions.as_slice(), &[Action::StackSwap(1)]); - let (stack, _, actions, _) = run(&[old_imm], &[current_imm], &[], &[], false); - assert_eq!(stack.item_at(0), Some(&StackItem::Value(current_imm))); + // Immediates are canonicalized to one ValueId per word before planning, so a stack + // slot already holding that ValueId satisfies the prefix check with no actions. + let (stack, _, actions, _) = run(&[imm], &[imm], &[], &[], false); + assert_eq!(stack.item_at(0), Some(&StackItem::Value(imm))); assert!(actions.is_empty()); }); } From 1e9fcdcd7dbfee672de74a074b85ba1a464e9248 Mon Sep 17 00:00:00 2001 From: sbillig Date: Tue, 7 Jul 2026 02:04:57 -0700 Subject: [PATCH 12/14] Fold hot-immediate tracking into ordinary use counting With one canonical ValueId per immediate word, UseTracker's parallel I256-keyed cached_imm_remaining map merges into the ValueId-keyed remaining map: cached immediates enter live_future until their last in-block use, cleanup_live reduces to live_future, and cache_preserve becomes the plain not-a-last-use rule. Rescue keeps immediates evictable regardless of liveness, preserving its prior behavior. --- .../codegen/src/stackalloc/stackify/block.rs | 9 +- .../src/stackalloc/stackify/builder.rs | 10 ++ .../codegen/src/stackalloc/stackify/rescue.rs | 7 +- .../codegen/src/stackalloc/stackify/uses.rs | 92 +++++-------------- 4 files changed, 39 insertions(+), 79 deletions(-) diff --git a/crates/codegen/src/stackalloc/stackify/block.rs b/crates/codegen/src/stackalloc/stackify/block.rs index cb206feaf..9f4bba40a 100644 --- a/crates/codegen/src/stackalloc/stackify/block.rs +++ b/crates/codegen/src/stackalloc/stackify/block.rs @@ -163,17 +163,16 @@ impl<'d, 'a, 'ctx, O: StackifyObserver> BlockPlanner<'d, 'a, 'ctx, O> { let before_cleanup_len = self.driver.pre_actions_len(inst); if !skip_cleanup { - let cleanup_live_future = self - .state - .uses - .cleanup_live(self.driver.ctx(), &self.state.stack); + // `live_future` includes cached immediates until their last in-block use, so + // dead-prefix cleanup does not drop a reusable cached-immediate stack copy. + let live_future = self.state.uses.live_future(); let live_out = self.state.uses.live_out(); let reach = self.driver.ctx().reach; self.driver.with_pre_actions(inst, |actions| { clean_dead_stack_prefix( reach, &mut self.state.stack, - &cleanup_live_future, + live_future, live_out, actions, ); diff --git a/crates/codegen/src/stackalloc/stackify/builder.rs b/crates/codegen/src/stackalloc/stackify/builder.rs index 3ee4b1a71..da93733f0 100644 --- a/crates/codegen/src/stackalloc/stackify/builder.rs +++ b/crates/codegen/src/stackalloc/stackify/builder.rs @@ -184,16 +184,26 @@ impl StackifyContext<'_> { self.value_aliases[value].unwrap_or(value) } + /// The value is pinned to the stack: losing its last copy costs a spill/reload, so operand + /// preparation must preserve it. Immediates are never pinned — they can always be re-pushed. pub(super) fn retains_value(&self, value: ValueId) -> bool { !self.func.dfg.value_is_imm(value) } + /// The value is a hot cached immediate: rematerializable, but expensive enough to push that + /// planning prefers `DUP`ing an existing stack copy over re-pushing it. pub(super) fn stack_caches_immediate(&self, value: ValueId) -> bool { self.func .dfg .value_imm(value) .is_some_and(|imm| self.stack_cached_immediates.contains(&imm.as_i256())) } + + /// The value participates in per-block use counting (`UseTracker`): pinned values and cached + /// immediates alike are counted so they stay in `live_future` until their last in-block use. + pub(super) fn is_use_tracked(&self, value: ValueId) -> bool { + self.retains_value(value) || self.stack_caches_immediate(value) + } } impl<'a> StackifyBuilder<'a> { diff --git a/crates/codegen/src/stackalloc/stackify/rescue.rs b/crates/codegen/src/stackalloc/stackify/rescue.rs index eaae7d76d..25aa9ed2b 100644 --- a/crates/codegen/src/stackalloc/stackify/rescue.rs +++ b/crates/codegen/src/stackalloc/stackify/rescue.rs @@ -164,12 +164,13 @@ fn choose_reachability_victim( let above = above.min(limit); // 1) Prefer deleting dead values, starting from the shallowest depth to minimize `SWAP*` - // chains. This includes immediates: if they're dead, removing them cannot introduce new - // rematerialization cost. + // chains. This includes immediates — even cached ones still in `live_future`: they are + // rematerializable, so deleting a copy never loses the value, and rescue deliberately trades + // a possible re-push for operand reachability. for (i, item) in stack.iter().take(above).enumerate() { if let StackItem::Value(v) = item && !protected_args.contains(*v) - && uses.is_dead(*v) + && (func.dfg.value_is_imm(*v) || uses.is_dead(*v)) { return Some(i); } diff --git a/crates/codegen/src/stackalloc/stackify/uses.rs b/crates/codegen/src/stackalloc/stackify/uses.rs index 03a595e92..2baff65bf 100644 --- a/crates/codegen/src/stackalloc/stackify/uses.rs +++ b/crates/codegen/src/stackalloc/stackify/uses.rs @@ -1,4 +1,4 @@ -use sonatina_ir::{BlockId, I256, ValueId}; +use sonatina_ir::{BlockId, ValueId}; use std::collections::BTreeMap; use crate::bitset::BitSet; @@ -7,17 +7,18 @@ use super::{ block::operand_order_for_stackify, builder::StackifyContext, slots::{FreeSlots, SlotPool}, - sym_stack::{StackItem, SymStack}, }; /// Per-block value-use bookkeeping threaded through the instruction walk. /// -/// Tracks how many in-block uses remain for each retained value (`remaining`) and each cached -/// immediate (`cached_imm_remaining`), the set of values still used later in the block -/// (`live_future`), and the block's live-out set unioned with phi-out sources (`live_out`). +/// Tracks how many in-block uses remain for each use-tracked value (`remaining`) — pinned values +/// and cached immediates alike, since immediates are canonicalized to one `ValueId` per word — the +/// set of values still used later in the block (`live_future`), and the block's live-out set +/// unioned with phi-out sources (`live_out`). `live_out` never contains immediates (liveness only +/// `mark_use`s them and `phi_out_sources` filters them), so a cached immediate is live exactly +/// until its last in-block use. pub(super) struct UseTracker { remaining: BTreeMap, - cached_imm_remaining: BTreeMap, live_future: BitSet, live_out: BitSet, } @@ -25,18 +26,13 @@ pub(super) struct UseTracker { impl UseTracker { pub(super) fn for_block(ctx: &StackifyContext<'_>, block: BlockId) -> Self { let mut remaining: BTreeMap = BTreeMap::new(); - let mut cached_imm_remaining: BTreeMap = BTreeMap::new(); for inst in ctx.func.layout.iter_inst(block) { if ctx.func.dfg.is_phi(inst) { continue; } for v in operand_order_for_stackify(ctx.func, inst, &ctx.value_aliases) { - if ctx.retains_value(v) { + if ctx.is_use_tracked(v) { *remaining.entry(v).or_insert(0) += 1; - } else if ctx.stack_caches_immediate(v) - && let Some(imm) = ctx.func.dfg.value_imm(v) - { - *cached_imm_remaining.entry(imm.as_i256()).or_insert(0) += 1; } } } @@ -45,7 +41,6 @@ impl UseTracker { live_out.union_with(&ctx.phi_out_sources[block]); Self { remaining, - cached_imm_remaining, live_future, live_out, } @@ -64,7 +59,9 @@ impl UseTracker { } /// Values in `args` whose in-block uses all end at this instruction (and which are not - /// live-out), so they may be consumed directly from the stack. + /// live-out), so they may be consumed directly from the stack. Restricted to pinned values: + /// cached immediates stay out of the consume-last-use search mask (they are rematerializable, + /// and their keep-vs-consume choice is expressed through `cache_preserve_in` instead). pub(super) fn last_uses_in( &self, ctx: &StackifyContext<'_>, @@ -88,41 +85,31 @@ impl UseTracker { last_use } - /// Cached immediates in `args` that have further in-block uses beyond this instruction and so - /// should be kept on the stack rather than consumed. + /// Cached immediates in `args` that have further in-block uses beyond this instruction — the + /// inverse of a last use — and so should be kept on the stack rather than consumed. pub(super) fn cache_preserve_in( &self, ctx: &StackifyContext<'_>, args: &[ValueId], ) -> BitSet { - let mut inst_counts: BTreeMap = BTreeMap::new(); + let mut inst_counts: BTreeMap = BTreeMap::new(); for &value in args { - if ctx.stack_caches_immediate(value) - && let Some(imm) = ctx.func.dfg.value_imm(value) - { - *inst_counts.entry(imm.as_i256()).or_insert(0) += 1; + if ctx.stack_caches_immediate(value) { + *inst_counts.entry(value).or_insert(0) += 1; } } let mut preserve: BitSet = BitSet::default(); - for &value in args { - if ctx.stack_caches_immediate(value) - && let Some(imm) = ctx.func.dfg.value_imm(value) - && self - .cached_imm_remaining - .get(&imm.as_i256()) - .copied() - .unwrap_or(0) - > inst_counts.get(&imm.as_i256()).copied().unwrap_or(0) - { + for (value, count) in inst_counts { + if self.remaining.get(&value).copied().unwrap_or(0) > count { preserve.insert(value); } } preserve } - /// Consume one use of each operand in `args`, releasing scratch slots for values that die here - /// and decaying cached-immediate counts. + /// Consume one use of each operand in `args`, dropping values from `live_future` at their last + /// in-block use and releasing scratch slots for non-immediate values that die here. pub(super) fn consume( &mut self, ctx: &StackifyContext<'_>, @@ -131,7 +118,7 @@ impl UseTracker { free_scratch: &mut FreeSlots, ) { for &v in args { - if ctx.retains_value(v) + if ctx.is_use_tracked(v) && let Some(n) = self.remaining.get_mut(&v) { let before = *n; @@ -144,42 +131,5 @@ impl UseTracker { } } } - - for &value in args { - if ctx.stack_caches_immediate(value) - && let Some(imm) = ctx.func.dfg.value_imm(value) - && let Some(count) = self.cached_imm_remaining.get_mut(&imm.as_i256()) - { - *count = count.saturating_sub(1); - } - } - self.cached_imm_remaining.retain(|_, count| *count != 0); - } - - /// `live_future`, extended with any cached immediates still on `stack` that have remaining - /// in-block uses, so dead-prefix cleanup does not drop a reusable cached immediate. - pub(super) fn cleanup_live( - &self, - ctx: &StackifyContext<'_>, - stack: &SymStack, - ) -> BitSet { - let mut live = self.live_future.clone(); - if self.cached_imm_remaining.is_empty() { - return live; - } - - for item in stack.iter() { - if let StackItem::Value(value) = item - && ctx.stack_caches_immediate(*value) - && let Some(imm) = ctx.func.dfg.value_imm(*value) - && self - .cached_imm_remaining - .get(&imm.as_i256()) - .is_some_and(|count| *count != 0) - { - live.insert(*value); - } - } - live } } From b51c46f83e29cbaa6a447fff9e005cc50ab7a2e9 Mon Sep 17 00:00:00 2001 From: sbillig Date: Sat, 11 Jul 2026 14:17:54 -0700 Subject: [PATCH 13/14] Handle duplicate stackify branch targets --- crates/codegen/src/cfg_edit.rs | 141 ++++++++++++++-- crates/codegen/src/critical_edge.rs | 100 +++++++++++- crates/codegen/src/isa/evm/tests.rs | 49 ++++-- crates/codegen/src/stackalloc/edge_split.rs | 152 ++++++++++++++++-- .../codegen/src/stackalloc/stackify/entry.rs | 12 +- crates/codegen/src/stackalloc/stackify/mod.rs | 9 +- 6 files changed, 410 insertions(+), 53 deletions(-) diff --git a/crates/codegen/src/cfg_edit.rs b/crates/codegen/src/cfg_edit.rs index 30350057e..0ca2f93aa 100644 --- a/crates/codegen/src/cfg_edit.rs +++ b/crates/codegen/src/cfg_edit.rs @@ -335,17 +335,7 @@ impl<'f> CfgEditor<'f> { (from, new_block) } - pub fn split_edge(&mut self, from: BlockId, to: BlockId) -> BlockId { - assert!(self.func.layout.is_block_inserted(from)); - assert!(self.func.layout.is_block_inserted(to)); - - let term = self.branch_terminator(from); - let branch_info = self.func.dfg.branch_info(term).unwrap(); - assert!( - branch_info.dests().into_iter().any(|dest| dest == to), - "edge {from:?} -> {to:?} does not exist" - ); - + fn insert_edge_block(&mut self, to: BlockId) -> BlockId { let mid = self.func.dfg.make_block(); if self.func.layout.entry_block() == Some(to) { // Splitting an edge whose destination is the entry block (e.g. a multiway self-loop @@ -361,6 +351,26 @@ impl<'f> CfgEditor<'f> { let mut cursor = InstInserter::at_location(CursorLocation::BlockTop(mid)); cursor.append_inst_data(self.func, Jump::new(self.func.dfg.inst_set().jump(), to)); + mid + } + + /// Split every parallel edge from `from` to `to` through one new jump block. + /// + /// Use [`Self::split_edge_at`] when parallel edge slots must remain distinct (for example, + /// stackify `br_table` cases whose outgoing symbolic stacks differ). + pub fn split_edge(&mut self, from: BlockId, to: BlockId) -> BlockId { + assert!(self.func.layout.is_block_inserted(from)); + assert!(self.func.layout.is_block_inserted(to)); + + let term = self.branch_terminator(from); + let branch_info = self.func.dfg.branch_info(term).unwrap(); + assert!( + branch_info.dests().into_iter().any(|dest| dest == to), + "edge {from:?} -> {to:?} does not exist" + ); + + let mid = self.insert_edge_block(to); + self.func.dfg.rewrite_branch_edges_to_block(term, to, mid); replace_phi_incoming_block(self.func, to, from, mid); @@ -368,6 +378,41 @@ impl<'f> CfgEditor<'f> { mid } + /// Split one outgoing branch edge slot through its own new jump block. + /// + /// Phi nodes are keyed by predecessor block rather than edge slot. If another parallel edge + /// from `from` to the same target remains, its incoming value is copied for the new block; + /// otherwise the predecessor label is moved from `from` to the new block. + pub fn split_edge_at(&mut self, from: BlockId, branch_slot: usize) -> BlockId { + assert!(self.func.layout.is_block_inserted(from)); + + let term = self.branch_terminator(from); + let branch_info = self.func.dfg.branch_info(term).unwrap(); + let dests = branch_info.dests(); + let to = *dests + .get(branch_slot) + .unwrap_or_else(|| panic!("outgoing edge slot out of bounds: {branch_slot}")); + assert!(self.func.layout.is_block_inserted(to)); + + let has_parallel_edge = dests + .iter() + .enumerate() + .any(|(slot, &dest)| slot != branch_slot && dest == to); + let mid = self.insert_edge_block(to); + self.func + .dfg + .rewrite_branch_edge_dest(term, branch_slot, mid); + + if has_parallel_edge { + copy_phi_incoming_block(self.func, to, from, mid); + } else { + replace_phi_incoming_block(self.func, to, from, mid); + } + + self.recompute_cfg(); + mid + } + /// Create or reuse a loop preheader for `lp_header` using predecessors outside the loop. /// /// Returns: @@ -1010,6 +1055,33 @@ fn replace_phi_incoming_block( } } +fn copy_phi_incoming_block( + func: &mut Function, + block: BlockId, + old_pred: BlockId, + new_pred: BlockId, +) { + let phi_inputs = iter_phis_in_block(func, block) + .map(|phi_inst| { + let phi = func.dfg.cast_phi(phi_inst).unwrap(); + let mut incoming = phi + .args() + .iter() + .filter(|(_, pred)| *pred == old_pred) + .map(|(value, _)| *value); + let value = incoming.next().unwrap_or_else(|| { + panic!("phi {phi_inst:?} in {block:?} missing incoming from {old_pred:?}") + }); + assert!( + incoming.next().is_none(), + "phi {phi_inst:?} in {block:?} has duplicate incoming from {old_pred:?}" + ); + (phi_inst, value) + }) + .collect::>(); + append_phi_inputs_for_new_pred(func, block, new_pred, &phi_inputs); +} + pub(crate) fn simplify_trivial_phis_in_block(func: &mut Function, block: BlockId) -> bool { let mut changed = false; let mut next_inst = func.layout.first_inst_of(block); @@ -1341,6 +1413,53 @@ block2: }); } + #[test] + fn split_edge_at_preserves_parallel_phi_inputs() { + let module = parse_test_module( + r#" +target = "evm-ethereum-osaka" + +func private %f() -> i32 { +block0: + br_table 0.i8 block1 (0.i8 block1) (1.i8 block2); + +block1: + v0.i32 = phi (7.i32 block0); + return v0; + +block2: + return 9.i32; +} +"#, + ); + let func_ref = module.funcs()[0]; + module.func_store.modify(func_ref, |func| { + let blocks: Vec<_> = func.layout.iter_block().collect(); + let [b0, b1, b2] = blocks.as_slice() else { + panic!("expected three blocks"); + }; + + let phi_inst = func.layout.first_inst_of(*b1).unwrap(); + let incoming = func.dfg.cast_phi(phi_inst).unwrap().args()[0].0; + let mut editor = CfgEditor::new(func, CleanupMode::Strict); + let first_mid = editor.split_edge_at(*b0, 0); + let second_mid = editor.split_edge_at(*b0, 1); + + let term = editor.func().layout.last_inst_of(*b0).unwrap(); + let dests = editor.func().dfg.branch_info(term).unwrap().dests(); + assert_eq!(dests.as_slice(), &[first_mid, second_mid, *b2]); + + let phi = editor.func().dfg.cast_phi(phi_inst).unwrap(); + assert_eq!(phi.args().len(), 2); + assert!(phi.args().iter().all( + |&(value, pred)| value == incoming && (pred == first_mid || pred == second_mid) + )); + assert_eq!(editor.cfg().preds_as_slice(*b1).len(), 2); + assert!(editor.cfg().preds_of(*b1).any(|&pred| pred == first_mid)); + assert!(editor.cfg().preds_of(*b1).any(|&pred| pred == second_mid)); + }); + } + #[test] fn truncate_block_from_inst_and_append_inst_with_result_rebuild_tail() { let mb = test_module_builder(); diff --git a/crates/codegen/src/critical_edge.rs b/crates/codegen/src/critical_edge.rs index 00380f128..cd5fc1d15 100644 --- a/crates/codegen/src/critical_edge.rs +++ b/crates/codegen/src/critical_edge.rs @@ -1,4 +1,4 @@ -use sonatina_ir::{BlockId, ControlFlowGraph, Function, InstId}; +use sonatina_ir::{ControlFlowGraph, Function, InstId}; use crate::cfg_edit::{CfgEditor, CleanupMode}; @@ -33,7 +33,7 @@ impl CriticalEdgeSplitter { let mut editor = CfgEditor::new(func, CleanupMode::Strict); for edge in edges { let from = editor.func().layout.inst_block(edge.inst); - editor.split_edge(from, edge.to); + editor.split_edge_at(from, edge.branch_slot); } cfg.compute(editor.func()); @@ -49,10 +49,19 @@ impl CriticalEdgeSplitter { return; } - for &succ in cfg.succs_of(block) { - if cfg.pred_num_of(succ) > 1 { - self.critical_edges.push(CriticalEdge::new(inst_id, succ)); + // Preserve the historical destination ordering for ordinary distinct edges so bridge + // creation does not perturb later fallthrough placement. Parallel slots for one target + // are still kept distinct and ordered by their branch slot. + for &to in cfg.succs_of(block) { + if cfg.pred_num_of(to) < 2 { + continue; } + self.critical_edges.extend( + cfg.succ_edges_of(block) + .map(|&edge| cfg.edge_data(edge)) + .filter(|edge| edge.to == to) + .map(|edge| CriticalEdge::new(inst_id, edge.branch_slot)), + ); } } } @@ -60,12 +69,12 @@ impl CriticalEdgeSplitter { #[derive(Debug)] struct CriticalEdge { inst: InstId, - to: BlockId, + branch_slot: usize, } impl CriticalEdge { - fn new(inst: InstId, to: BlockId) -> Self { - Self { inst, to } + fn new(inst: InstId, branch_slot: usize) -> Self { + Self { inst, branch_slot } } } @@ -80,6 +89,7 @@ mod tests { }, isa::Isa, }; + use sonatina_parser::parse_module; use super::*; @@ -414,4 +424,78 @@ mod tests { .view(func_ref, |func| cfg_split.compute(func)); assert_eq!(cfg, cfg_split); } + + #[test] + fn critical_edge_br_table_splits_duplicate_target_slots_individually() { + let parsed = parse_module( + r#" +target = "evm-ethereum-osaka" + +func private %f(v0.i1, v1.i8) -> i8 { +block0: + br v0 block1 block2; + +block1: + br_table v1 block3 (0.i8 block3) (1.i8 block4); + +block2: + jump block3; + +block3: + v2.i8 = phi (7.i8 block1) (9.i8 block2); + return v2; + +block4: + return 11.i8; +} +"#, + ) + .expect("module parses"); + let func_ref = parsed.module.funcs()[0]; + parsed.module.func_store.modify(func_ref, |func| { + let blocks: Vec<_> = func.layout.iter_block().collect(); + let [_, duplicate_pred, other_pred, target, distinct_target] = blocks.as_slice() else { + panic!("expected five blocks"); + }; + let phi_inst = func.layout.first_inst_of(*target).unwrap(); + let duplicate_incoming = func + .dfg + .cast_phi(phi_inst) + .unwrap() + .args() + .iter() + .find(|(_, pred)| pred == duplicate_pred) + .unwrap() + .0; + + let mut cfg = ControlFlowGraph::default(); + cfg.compute(func); + CriticalEdgeSplitter::new().run(func, &mut cfg); + + let term = func.layout.last_inst_of(*duplicate_pred).unwrap(); + let dests = func.dfg.branch_info(term).unwrap().dests(); + assert_eq!(dests.len(), 3); + assert_ne!(dests[0], dests[1]); + assert_eq!(dests[2], *distinct_target); + assert!( + dests[..2] + .iter() + .all(|&mid| { cfg.succs_of(mid).eq(std::iter::once(target)) }) + ); + + let phi = func.dfg.cast_phi(phi_inst).unwrap(); + assert_eq!(phi.args().len(), 3); + assert!(!phi.args().iter().any(|(_, pred)| pred == duplicate_pred)); + assert!(phi.args().iter().any(|(_, pred)| pred == other_pred)); + assert_eq!( + phi.args() + .iter() + .filter(|(value, pred)| { + *value == duplicate_incoming && *pred != *other_pred + }) + .count(), + 2 + ); + }); + } } diff --git a/crates/codegen/src/isa/evm/tests.rs b/crates/codegen/src/isa/evm/tests.rs index 67a7be082..afeb29b1d 100644 --- a/crates/codegen/src/isa/evm/tests.rs +++ b/crates/codegen/src/isa/evm/tests.rs @@ -2769,15 +2769,18 @@ object @Contract { } #[test] -fn machine_pipeline_compiles_entry_self_loop() { +fn machine_pipeline_compiles_multiway_entry_self_loops() { // Regression for the `StackifyEdgeSplitter` pipeline gap: a multiway self-loop on the entry // block (`br v0 block0 block1`) is an in-cycle multiway edge that is *not* critical (block0's // only predecessor is itself), so plain critical-edge splitting leaves it intact. Stackify // then plans a branch edge back to the already-planned entry block, which trips the guard in // `on_branch_edge` (previously a debug assert / silent miscompile). The machine pipeline runs - // `StackifyEdgeSplitter`, which splits the edge and makes the shape compile. - let parsed = parse_module( - r#" + // `StackifyEdgeSplitter`, which splits the edge and makes the shape compile. The + // all-duplicate form is canonicalized to a jump before edge classification. + let sources = [ + ( + "mixed targets", + r#" target = "evm-ethereum-osaka" func public %f(v0.i1, v1.i256) { @@ -2795,18 +2798,38 @@ object @Contract { } } "#, - ) - .unwrap(); + ), + ( + "duplicate targets", + r#" +target = "evm-ethereum-osaka" - let func = parsed.module.funcs()[0]; +func public %f(v0.i1, v1.i256) { +block0: + mstore v1 v1 i256; + br v0 block0 block0; +} + +object @Contract { + section runtime { + entry %f; + } +} +"#, + ), + ]; let backend = test_backend(); - let prepared = backend - .prepare_section(work_module(&parsed.module, &[func])) - .expect("prepare should succeed for an entry self-loop"); + for (label, source) in sources { + let parsed = parse_module(source).unwrap(); + let func = parsed.module.funcs()[0]; + let prepared = backend + .prepare_section(work_module(&parsed.module, &[func])) + .unwrap_or_else(|err| panic!("prepare should succeed for {label}: {err}")); - backend - .lower_function(&prepared, func) - .expect("lowering an entry self-loop should succeed"); + backend + .lower_function(&prepared, func) + .unwrap_or_else(|err| panic!("lowering should succeed for {label}: {err}")); + } } #[test] diff --git a/crates/codegen/src/stackalloc/edge_split.rs b/crates/codegen/src/stackalloc/edge_split.rs index 309950742..e184d1501 100644 --- a/crates/codegen/src/stackalloc/edge_split.rs +++ b/crates/codegen/src/stackalloc/edge_split.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeSet; +use std::collections::BTreeMap; use cranelift_entity::SecondaryMap; use sonatina_ir::{BlockId, ControlFlowGraph, Function}; @@ -11,8 +11,9 @@ use crate::{ /// Establishes the edge preconditions of the stackify allocator (`stackalloc::stackify`). /// -/// It runs [`CriticalEdgeSplitter`] and, in addition, splits every *multiway* edge whose target -/// is planned no later than the branching block itself. +/// It canonicalizes all-identical multiway terminators to jumps, runs +/// [`CriticalEdgeSplitter`], and then splits each remaining duplicate-target edge slot and every +/// *multiway* edge whose target is planned no later than the branching block itself. /// /// Stackify plans blocks in dominator-tree RPO, threading a symbolic stack forward, and marks a /// block planned before simulating its terminator. A multiway terminator whose target is a block @@ -22,13 +23,42 @@ use crate::{ /// or silently drops the fixup. Splitting the edge inserts a single-jump block that carries the /// fixup instead. /// -/// Forward multiway edges — the common in-loop conditional branch, whose target is planned after -/// the branch — are handled by the planner directly and are deliberately left intact, so block -/// layout (and emitted bytecode) is unchanged for functions without a retreating multiway edge. +/// Forward multiway edges with distinct targets — the common in-loop conditional branch, whose +/// target is planned after the branch — are handled by the planner directly and are deliberately +/// left intact, so block layout (and emitted bytecode) is unchanged for functions without a +/// retreating or duplicate-target multiway edge. pub struct StackifyEdgeSplitter; impl StackifyEdgeSplitter { pub fn run(func: &mut Function, cfg: &mut ControlFlowGraph) { + // An all-identical multiway terminator has one semantic destination and needs no per-edge + // stack state. Canonicalize it before classifying critical or retreating edges; duplicate + // destinations that remain (e.g. a subset of br_table cases) do require distinct bridges. + let terms: Vec<_> = func + .layout + .iter_block() + .filter_map(|block| func.layout.last_inst_of(block)) + .collect(); + let mut canonicalized = false; + for term in terms { + let Some(branch) = func.dfg.branch_info(term) else { + continue; + }; + let dests = branch.dests(); + if dests.len() < 2 || dests.iter().any(|&dest| dest != dests[0]) { + continue; + } + + // Retaining every edge invokes the branch instruction's canonical representation: + // `Br` and `BrTable` both collapse an all-identical destination set to `Jump`. + let keep_mask = vec![true; dests.len()]; + func.dfg.retain_branch_edges(term, &keep_mask); + canonicalized = true; + } + if canonicalized { + cfg.compute(func); + } + CriticalEdgeSplitter::new().run(func, cfg); // Rank reachable blocks by stackify's planning order (dominator-tree RPO). @@ -39,19 +69,36 @@ impl StackifyEdgeSplitter { plan_rank[block] = Some(rank as u32); } - let mut edges = BTreeSet::<(BlockId, BlockId)>::new(); + let mut edges = Vec::<(BlockId, BlockId, usize)>::new(); for from in func.layout.iter_block() { - if cfg.succ_num_of(from) < 2 { + let Some(term) = func.layout.last_inst_of(from) else { + continue; + }; + let Some(branch) = func.dfg.branch_info(term) else { + continue; + }; + let dests = branch.dests(); + if dests.len() < 2 { continue; } let Some(from_rank) = plan_rank[from] else { continue; }; - for &to in cfg.succs_of(from) { + + let mut dest_counts = BTreeMap::::new(); + for &to in &dests { + *dest_counts.entry(to).or_default() += 1; + } + + for (branch_slot, to) in dests.into_iter().enumerate() { // Retreating edge: `to` is planned before (backedge) or together with (self-loop) // `from`, so `to` is already planned when `from`'s terminator is simulated. - if plan_rank[to].is_some_and(|to_rank| to_rank <= from_rank) { - edges.insert((from, to)); + // Duplicate-target edge slots also need separate bridges: br_table cases can + // reach the same block with different post-comparison symbolic stacks. + let duplicate = dest_counts[&to] > 1; + let retreating = plan_rank[to].is_some_and(|to_rank| to_rank <= from_rank); + if duplicate || retreating { + edges.push((from, to, branch_slot)); } } } @@ -60,9 +107,12 @@ impl StackifyEdgeSplitter { return; } + // Preserve the old `(from, to)` bridge-creation order for distinct edges; use the branch + // slot only to order parallel edges that previously collapsed into one operation. + edges.sort_unstable(); let mut editor = CfgEditor::new(func, CleanupMode::Strict); - for (from, to) in edges { - editor.split_edge(from, to); + for (from, _, branch_slot) in edges { + editor.split_edge_at(from, branch_slot); } cfg.compute(editor.func()); } @@ -70,6 +120,8 @@ impl StackifyEdgeSplitter { #[cfg(test)] mod tests { + use std::collections::BTreeSet; + use sonatina_ir::cfg::ControlFlowGraph; use sonatina_parser::parse_module; @@ -108,4 +160,78 @@ block1: assert!(cfg.preds_of(entry).any(|&pred| cfg.succ_num_of(pred) == 1)); }); } + + #[test] + fn canonicalizes_all_duplicate_self_loop_targets() { + const SRC: &str = r#" +target = "evm-ethereum-osaka" + +func public %f(v0.i1) { +block0: + br v0 block0 block0; +} +"#; + + let parsed = parse_module(SRC).expect("module parses"); + let func = parsed.module.funcs()[0]; + + parsed.module.func_store.modify(func, |function| { + let mut cfg = ControlFlowGraph::new(); + cfg.compute(function); + let entry = cfg.entry().expect("missing entry"); + assert_eq!(cfg.succ_edges_as_slice(entry).len(), 2); + + StackifyEdgeSplitter::run(function, &mut cfg); + + assert_eq!(cfg.entry(), Some(entry)); + assert_eq!(cfg.succ_edges_as_slice(entry).len(), 1); + let term = function.layout.last_inst_of(entry).unwrap(); + let jump = function.dfg.cast_jump(term).expect("branch becomes jump"); + assert_eq!(*jump.dest(), entry); + }); + } + + #[test] + fn splits_duplicate_br_table_targets_per_edge_slot() { + const SRC: &str = r#" +target = "evm-ethereum-osaka" + +func public %f(v0.i8) { +block0: + br_table v0 block1 (0.i8 block1) (1.i8 block2); + +block1: + return; + +block2: + return; +} +"#; + + let parsed = parse_module(SRC).expect("module parses"); + let func = parsed.module.funcs()[0]; + + parsed.module.func_store.modify(func, |function| { + let blocks: Vec<_> = function.layout.iter_block().collect(); + let [entry, duplicate_target, other_target] = blocks.as_slice() else { + panic!("expected three blocks"); + }; + let mut cfg = ControlFlowGraph::new(); + cfg.compute(function); + + StackifyEdgeSplitter::run(function, &mut cfg); + + let term = function.layout.last_inst_of(*entry).unwrap(); + let dests = function.dfg.branch_info(term).unwrap().dests(); + assert_eq!(dests.len(), 3); + assert_eq!(dests.iter().copied().collect::>().len(), 3); + assert_eq!(dests[2], *other_target); + assert!( + dests[..2] + .iter() + .all(|&mid| { cfg.succs_of(mid).eq(std::iter::once(duplicate_target)) }) + ); + assert_eq!(cfg.pred_num_of(*duplicate_target), 2); + }); + } } diff --git a/crates/codegen/src/stackalloc/stackify/entry.rs b/crates/codegen/src/stackalloc/stackify/entry.rs index 4595d7163..f32f98e07 100644 --- a/crates/codegen/src/stackalloc/stackify/entry.rs +++ b/crates/codegen/src/stackalloc/stackify/entry.rs @@ -271,13 +271,17 @@ impl EntryTable { } } - /// Store a single-pred block's inherited predecessor stack. First writer wins: with split - /// critical edges a single-pred block receives exactly one edge, except the degenerate - /// `br c b b` double edge whose two stacks are identical, so keeping the first is correct. + /// Store a single-pred block's inherited predecessor stack. Critical, retreating, and + /// duplicate-target multiway edges are split beforehand, so exactly one edge may write it. fn set_inherited(&mut self, block: BlockId, pred: BlockId, stack: &SymStack) { match &mut self.entries[block].state { EntryState::Pending { inherited, .. } => { - inherited.get_or_insert_with(|| (pred, stack.clone())); + assert!( + inherited.is_none(), + "multiple edges attempted to inherit block {block:?}: run \ + StackifyEdgeSplitter before stackify" + ); + *inherited = Some((pred, stack.clone())); } EntryState::Frozen(_) => unreachable!("inherit into already-frozen block {block:?}"), } diff --git a/crates/codegen/src/stackalloc/stackify/mod.rs b/crates/codegen/src/stackalloc/stackify/mod.rs index 34eed71e3..d2f329c00 100644 --- a/crates/codegen/src/stackalloc/stackify/mod.rs +++ b/crates/codegen/src/stackalloc/stackify/mod.rs @@ -23,10 +23,11 @@ //! edge stores can be emitted directly without first staging every phi source on the stack. //! //! Notes specific to this codebase: -//! - Run `StackifyEdgeSplitter` before this allocator: it establishes both split -//! preconditions, splitting critical edges *and* every multiway edge whose target is already -//! planned when the branch is reached (a self-loop or backedge, e.g. a multiway self-loop on -//! the entry block). Splitting only critical edges is not enough. +//! - Run `StackifyEdgeSplitter` before this allocator: it canonicalizes all-identical branches +//! and establishes all split preconditions, splitting critical edges, duplicate-target edge +//! slots, and every multiway edge whose target is already planned when the branch is reached +//! (a self-loop or backedge, e.g. a multiway self-loop on the entry block). Splitting only +//! critical edges is not enough. //! - Internal calls rely on an implicit return address value on the EVM stack. //! The allocator models this as a special stack item barrier to avoid popping //! into the caller's preserved stack segment. From 1e7a242a5736e956154ba0b66c40385e77fc4070 Mon Sep 17 00:00:00 2001 From: sbillig Date: Sat, 11 Jul 2026 17:57:42 -0700 Subject: [PATCH 14/14] Simplify slot-aware CFG edge splitting --- crates/codegen/src/cfg_edit.rs | 42 ++-------- crates/codegen/src/critical_edge.rs | 85 +++++++-------------- crates/codegen/src/stackalloc/edge_split.rs | 48 +++++------- 3 files changed, 53 insertions(+), 122 deletions(-) diff --git a/crates/codegen/src/cfg_edit.rs b/crates/codegen/src/cfg_edit.rs index 0ca2f93aa..ab4828a35 100644 --- a/crates/codegen/src/cfg_edit.rs +++ b/crates/codegen/src/cfg_edit.rs @@ -354,54 +354,28 @@ impl<'f> CfgEditor<'f> { mid } - /// Split every parallel edge from `from` to `to` through one new jump block. - /// - /// Use [`Self::split_edge_at`] when parallel edge slots must remain distinct (for example, - /// stackify `br_table` cases whose outgoing symbolic stacks differ). - pub fn split_edge(&mut self, from: BlockId, to: BlockId) -> BlockId { - assert!(self.func.layout.is_block_inserted(from)); - assert!(self.func.layout.is_block_inserted(to)); - - let term = self.branch_terminator(from); - let branch_info = self.func.dfg.branch_info(term).unwrap(); - assert!( - branch_info.dests().into_iter().any(|dest| dest == to), - "edge {from:?} -> {to:?} does not exist" - ); - - let mid = self.insert_edge_block(to); - - self.func.dfg.rewrite_branch_edges_to_block(term, to, mid); - replace_phi_incoming_block(self.func, to, from, mid); - - self.recompute_cfg(); - mid - } - /// Split one outgoing branch edge slot through its own new jump block. /// /// Phi nodes are keyed by predecessor block rather than edge slot. If another parallel edge /// from `from` to the same target remains, its incoming value is copied for the new block; /// otherwise the predecessor label is moved from `from` to the new block. - pub fn split_edge_at(&mut self, from: BlockId, branch_slot: usize) -> BlockId { + pub fn split_out_edge(&mut self, from: BlockId, edge_idx: usize) -> BlockId { assert!(self.func.layout.is_block_inserted(from)); let term = self.branch_terminator(from); let branch_info = self.func.dfg.branch_info(term).unwrap(); let dests = branch_info.dests(); let to = *dests - .get(branch_slot) - .unwrap_or_else(|| panic!("outgoing edge slot out of bounds: {branch_slot}")); + .get(edge_idx) + .unwrap_or_else(|| panic!("outgoing edge index out of bounds: {edge_idx}")); assert!(self.func.layout.is_block_inserted(to)); let has_parallel_edge = dests .iter() .enumerate() - .any(|(slot, &dest)| slot != branch_slot && dest == to); + .any(|(other_idx, &dest)| other_idx != edge_idx && dest == to); let mid = self.insert_edge_block(to); - self.func - .dfg - .rewrite_branch_edge_dest(term, branch_slot, mid); + self.func.dfg.rewrite_branch_edge_dest(term, edge_idx, mid); if has_parallel_edge { copy_phi_incoming_block(self.func, to, from, mid); @@ -1414,7 +1388,7 @@ block2: } #[test] - fn split_edge_at_preserves_parallel_phi_inputs() { + fn split_out_edge_preserves_parallel_phi_inputs() { let module = parse_test_module( r#" target = "evm-ethereum-osaka" @@ -1442,8 +1416,8 @@ block2: let phi_inst = func.layout.first_inst_of(*b1).unwrap(); let incoming = func.dfg.cast_phi(phi_inst).unwrap().args()[0].0; let mut editor = CfgEditor::new(func, CleanupMode::Strict); - let first_mid = editor.split_edge_at(*b0, 0); - let second_mid = editor.split_edge_at(*b0, 1); + let first_mid = editor.split_out_edge(*b0, 0); + let second_mid = editor.split_out_edge(*b0, 1); let term = editor.func().layout.last_inst_of(*b0).unwrap(); let dests = editor.func().dfg.branch_info(term).unwrap().dests(); diff --git a/crates/codegen/src/critical_edge.rs b/crates/codegen/src/critical_edge.rs index cd5fc1d15..93f52046e 100644 --- a/crates/codegen/src/critical_edge.rs +++ b/crates/codegen/src/critical_edge.rs @@ -1,80 +1,47 @@ -use sonatina_ir::{ControlFlowGraph, Function, InstId}; +use sonatina_ir::{BlockId, ControlFlowGraph, Function}; use crate::cfg_edit::{CfgEditor, CleanupMode}; -#[derive(Debug)] -pub struct CriticalEdgeSplitter { - critical_edges: Vec, -} - -impl Default for CriticalEdgeSplitter { - fn default() -> Self { - Self::new() - } -} +#[derive(Debug, Default)] +pub struct CriticalEdgeSplitter; impl CriticalEdgeSplitter { pub fn new() -> Self { - Self { - critical_edges: Vec::default(), - } + Self } pub fn run(&mut self, func: &mut Function, cfg: &mut ControlFlowGraph) { - self.clear(); - + let mut critical_edges = Vec::<(BlockId, usize)>::new(); for block in func.layout.iter_block() { - if let Some(last_inst) = func.layout.last_inst_of(block) { - self.add_critical_edges(last_inst, func, cfg); + if cfg.succ_num_of(block) < 2 { + continue; } - } - let edges = std::mem::take(&mut self.critical_edges); - let mut editor = CfgEditor::new(func, CleanupMode::Strict); - for edge in edges { - let from = editor.func().layout.inst_block(edge.inst); - editor.split_edge_at(from, edge.branch_slot); + // Preserve the historical destination ordering for ordinary distinct edges so bridge + // creation does not perturb later fallthrough placement. Parallel slots for one + // target are still kept distinct and ordered by their branch slot. + for &to in cfg.succs_of(block) { + if cfg.pred_num_of(to) < 2 { + continue; + } + critical_edges.extend( + cfg.succ_edges_of(block) + .map(|&edge| cfg.edge_data(edge)) + .filter(|edge| edge.to == to) + .map(|edge| (block, edge.branch_slot)), + ); + } } - cfg.compute(editor.func()); - } - - pub fn clear(&mut self) { - self.critical_edges.clear(); - } - - fn add_critical_edges(&mut self, inst_id: InstId, func: &Function, cfg: &ControlFlowGraph) { - let block = func.layout.inst_block(inst_id); - if cfg.succ_num_of(block) < 2 { + if critical_edges.is_empty() { return; } - // Preserve the historical destination ordering for ordinary distinct edges so bridge - // creation does not perturb later fallthrough placement. Parallel slots for one target - // are still kept distinct and ordered by their branch slot. - for &to in cfg.succs_of(block) { - if cfg.pred_num_of(to) < 2 { - continue; - } - self.critical_edges.extend( - cfg.succ_edges_of(block) - .map(|&edge| cfg.edge_data(edge)) - .filter(|edge| edge.to == to) - .map(|edge| CriticalEdge::new(inst_id, edge.branch_slot)), - ); + let mut editor = CfgEditor::new(func, CleanupMode::Strict); + for (from, branch_slot) in critical_edges { + editor.split_out_edge(from, branch_slot); } - } -} - -#[derive(Debug)] -struct CriticalEdge { - inst: InstId, - branch_slot: usize, -} - -impl CriticalEdge { - fn new(inst: InstId, branch_slot: usize) -> Self { - Self { inst, branch_slot } + cfg.clone_from(editor.cfg()); } } diff --git a/crates/codegen/src/stackalloc/edge_split.rs b/crates/codegen/src/stackalloc/edge_split.rs index e184d1501..39a754be9 100644 --- a/crates/codegen/src/stackalloc/edge_split.rs +++ b/crates/codegen/src/stackalloc/edge_split.rs @@ -34,28 +34,22 @@ impl StackifyEdgeSplitter { // An all-identical multiway terminator has one semantic destination and needs no per-edge // stack state. Canonicalize it before classifying critical or retreating edges; duplicate // destinations that remain (e.g. a subset of br_table cases) do require distinct bridges. - let terms: Vec<_> = func + let redundant_terms: Vec<_> = func .layout .iter_block() - .filter_map(|block| func.layout.last_inst_of(block)) + .filter_map(|block| { + let edge_count = cfg.succ_edges_as_slice(block).len(); + (edge_count > 1 && cfg.succ_num_of(block) == 1) + .then(|| (func.layout.last_inst_of(block).unwrap(), edge_count)) + }) .collect(); - let mut canonicalized = false; - for term in terms { - let Some(branch) = func.dfg.branch_info(term) else { - continue; - }; - let dests = branch.dests(); - if dests.len() < 2 || dests.iter().any(|&dest| dest != dests[0]) { - continue; - } - + for &(term, edge_count) in &redundant_terms { // Retaining every edge invokes the branch instruction's canonical representation: // `Br` and `BrTable` both collapse an all-identical destination set to `Jump`. - let keep_mask = vec![true; dests.len()]; + let keep_mask = vec![true; edge_count]; func.dfg.retain_branch_edges(term, &keep_mask); - canonicalized = true; } - if canonicalized { + if !redundant_terms.is_empty() { cfg.compute(func); } @@ -71,14 +65,8 @@ impl StackifyEdgeSplitter { let mut edges = Vec::<(BlockId, BlockId, usize)>::new(); for from in func.layout.iter_block() { - let Some(term) = func.layout.last_inst_of(from) else { - continue; - }; - let Some(branch) = func.dfg.branch_info(term) else { - continue; - }; - let dests = branch.dests(); - if dests.len() < 2 { + let outgoing = cfg.succ_edges_as_slice(from); + if outgoing.len() < 2 { continue; } let Some(from_rank) = plan_rank[from] else { @@ -86,11 +74,13 @@ impl StackifyEdgeSplitter { }; let mut dest_counts = BTreeMap::::new(); - for &to in &dests { - *dest_counts.entry(to).or_default() += 1; + for &edge in outgoing { + *dest_counts.entry(cfg.edge_data(edge).to).or_default() += 1; } - for (branch_slot, to) in dests.into_iter().enumerate() { + for &edge in outgoing { + let edge = cfg.edge_data(edge); + let to = edge.to; // Retreating edge: `to` is planned before (backedge) or together with (self-loop) // `from`, so `to` is already planned when `from`'s terminator is simulated. // Duplicate-target edge slots also need separate bridges: br_table cases can @@ -98,7 +88,7 @@ impl StackifyEdgeSplitter { let duplicate = dest_counts[&to] > 1; let retreating = plan_rank[to].is_some_and(|to_rank| to_rank <= from_rank); if duplicate || retreating { - edges.push((from, to, branch_slot)); + edges.push((from, to, edge.branch_slot)); } } } @@ -112,9 +102,9 @@ impl StackifyEdgeSplitter { edges.sort_unstable(); let mut editor = CfgEditor::new(func, CleanupMode::Strict); for (from, _, branch_slot) in edges { - editor.split_edge_at(from, branch_slot); + editor.split_out_edge(from, branch_slot); } - cfg.compute(editor.func()); + cfg.clone_from(editor.cfg()); } }