diff --git a/Cargo.toml b/Cargo.toml index dfff957..ae5a395 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ repository = "https://github.com/bytecodealliance/regalloc2" [dependencies] log = { version = "0.4.8", default-features = false } -smallvec = { version = "1.6.1", features = ["union"] } +smallvec = { version = "1.6.1", features = ["union", "const_generics"] } rustc-hash = { version = "1.1.0", default-features = false } slice-group-by = { version = "0.3.0", default-features = false } hashbrown = "0.13.2" diff --git a/src/ion/fast_alloc2.rs b/src/ion/fast_alloc2.rs index 6d361ea..eee3735 100644 --- a/src/ion/fast_alloc2.rs +++ b/src/ion/fast_alloc2.rs @@ -626,6 +626,19 @@ impl<'a, F: Function> FastAlloc<'a, F> { self.livein_locs .push(LiveinLoc::init(data.preg(), data.stack_slot())); } + + // also save block params + for &vreg in self.func.block_params(block) { + let data = &self.vregs[vreg.vreg()]; + trace!( + " -> {} at {:?} and slot {:?} (param)", + vreg, + data.preg(), + data.stack_slot() + ); + self.livein_locs + .push(LiveinLoc::init(data.preg(), data.stack_slot())); + } } // load livein locations if our predecessor has multiple successors @@ -670,6 +683,51 @@ impl<'a, F: Function> FastAlloc<'a, F> { loc_idx += 1; } self.liveins = liveins; + + // also restore block params + for &vreg in self.func.block_params(block) { + let data = &mut self.vregs[vreg.vreg()]; + let loc = self.livein_locs[loc_idx]; + trace!( + " -> {} at {:?} and slot {:?} (param)", + vreg, + loc.preg(), + loc.slot() + ); + match (data.preg(), loc.preg()) { + (Some(cur_preg), Some(loc_preg)) => { + if cur_preg != loc_preg { + self.clear_preg(cur_preg.index()); + self.clear_preg(loc_preg.index()); + self.assign_preg( + loc_preg, + VReg::new(vreg.vreg(), loc_preg.class()), + ); + } + } + (Some(cur_preg), None) => { + self.clear_preg(cur_preg.index()); + } + (None, Some(loc_preg)) => { + self.clear_preg(loc_preg.index()); + self.assign_preg(loc_preg, VReg::new(vreg.vreg(), loc_preg.class())); + } + (None, None) => {} + } + let data = &mut self.vregs[vreg.vreg()]; + match loc.slot() { + Some(slot) => { + debug_assert!(data.allocated_stack_slot().is_some()); + debug_assert_eq!(data.allocated_stack_slot().unwrap(), slot); + data.set_stack_slot_valid(); + } + None => { + data.clear_stack_slot(true); + } + } + + loc_idx += 1; + } } self.alloc_block_insts(block)?; @@ -1208,6 +1266,7 @@ impl<'a, F: Function> FastAlloc<'a, F> { return res; } + // TODO: check if the reftype count is correctly managed fn alloc_block_edge_single_unallocated( &mut self, block: Block, @@ -1291,6 +1350,7 @@ impl<'a, F: Function> FastAlloc<'a, F> { to: Allocation::stack(SpillSlot::new(slot as usize)), }, )); + self.vregs[succ_param.vreg()].set_stack_slot_valid(); self.vregs[out_vreg.vreg()].cur_use_idx -= 1; continue; @@ -1349,12 +1409,233 @@ impl<'a, F: Function> FastAlloc<'a, F> { todo!() } + // TODO: check if the reftype count is correctly managed fn alloc_block_edge_multiple( &mut self, block: Block, succs: &[Block], ) -> Result<(), RegAllocError> { - todo!() + trace!(" -> Block has multiple unallocated edges. Copying allocations"); + let mut slots_made_valid = SmallVec::<[VReg; 4]>::new(); + let mut slots_made_invalid: SmallVec<[VReg; 4]> = SmallVec::<[VReg; 4]>::new(); + let mut preg_allocs: SmallVec<[(u16, VReg); PReg::NUM_INDEX + 1]> = SmallVec::new(); + + for idx in 0..self.pregs.len() { + debug_assert!(idx < PReg::NUM_INDEX); + let data = &self.pregs[idx]; + match data.vreg() { + Some(vreg) => preg_allocs.push((idx as u16, vreg)), + None => {} + } + } + preg_allocs.push((u16::MAX, VReg::invalid())); + + let last_inst = self.func.block_insns(block).last(); + debug_assert!(self.func.is_branch(last_inst)); + let out_vregs = self.func.branch_blockparams(block, last_inst, 0); + + // we basically run the block_edge_single_unallocated algo but inside the target block and rewind afterwards + for &succ in succs { + assert_ne!(block, succ); + trace!(" -> Allocating edge for block {:?}", succ); + + let succ_params = self.func.block_params(succ); + // prevent these from being allocated as tmp regs + let mut pregs_allocated = PRegSet::empty(); + let mut tmp_reg = None; + for i in 0..succ_params.len() { + let succ_param = succ_params[i]; + let out_vreg = out_vregs[i]; + + if self.vreg_killed(out_vreg.vreg(), block, self.cur_inst_pos + 1, false) { + trace!(" -> {} mapping to {} is killed", out_vreg, succ_param); + if let Some(slot) = self.vregs[out_vreg.vreg()].stack_slot() { + trace!(" -> Transferring slot {} to {}", slot, succ_param); + self.vregs[succ_param.vreg()].set_stack_slot_allocated(slot); + self.vregs[succ_param.vreg()].set_stack_slot_valid(); + slots_made_valid.push(succ_param); + if cfg!(debug_assertions) { + self.vregs[out_vreg.vreg()].clear_stack_slot(false); + slots_made_invalid.push(out_vreg); + } + } + if let Some(preg) = self.vregs[out_vreg.vreg()].preg() { + trace!(" -> Transferring {} to {}", preg, succ_param); + self.vregs[succ_param.vreg()].set_preg(preg); + if cfg!(debug_assertions) { + self.vregs[out_vreg.vreg()].clear_preg(); + } + self.pregs[preg.index()].set_vreg(succ_param); + pregs_allocated.add(preg); + } + self.vregs[out_vreg.vreg()].cur_use_idx -= 1; + continue; + } + + trace!(" -> {} mapping to {} is not killed", out_vreg, succ_param); + // the out_vreg might only be alive until the end of the block but used by another param or still be live after the block + if let Some(preg) = self.vregs[out_vreg.vreg()].preg() { + // TODO: we could do some fancy calculation here which of the block params should get to stay in the preg but rn we dont + if self.next_use(out_vreg.vreg()).unwrap() + > self.next_use(succ_param.vreg()).unwrap() + || self.vregs[out_vreg.vreg()].stack_slot().is_some() + { + trace!(" -> Transferring preg from {} to {}", out_vreg, succ_param); + + if self.vregs[out_vreg.vreg()].stack_slot().is_none() { + // save to stack if necessary + self.alloc_and_move_to_stack(out_vreg, ProgPoint::before(last_inst)); + slots_made_valid.push(out_vreg); + } + + self.vregs[out_vreg.vreg()].clear_preg(); + self.vregs[succ_param.vreg()].set_preg(preg); + self.pregs[preg.index()].set_vreg(succ_param); + pregs_allocated.add(preg); + + // don't need to care about stack slot + self.vregs[out_vreg.vreg()].cur_use_idx -= 1; + continue; + } + + let slot = self.alloc_stack_slot(succ_param); + trace!(" -> Transferring preg from {} to slot {}", preg, slot); + self.edits.push(( + ProgPoint::before(last_inst), + Edit::Move { + from: Allocation::reg(preg), + to: Allocation::stack(SpillSlot::new(slot as usize)), + }, + )); + self.vregs[succ_param.vreg()].set_stack_slot_valid(); + slots_made_valid.push(succ_param); + + self.vregs[out_vreg.vreg()].cur_use_idx -= 1; + continue; + } + + let dst_slot = self.alloc_stack_slot(succ_param) as usize; + let src_slot = self.vregs[out_vreg.vreg()].stack_slot().unwrap() as usize; + trace!( + " -> Copying {} from slot {} to slot {} for {}", + out_vreg, + src_slot, + dst_slot, + succ_param + ); + // out_vreg is on the stack so create a stack slot for succ_param + let preg = match tmp_reg { + Some(preg) => preg, + None => { + let preg = self.find_free_reg( + succ_param.class(), + pregs_allocated, + Some(ProgPoint::before(last_inst)), + block, + self.cur_inst_pos + 1, + )?; + tmp_reg = Some(preg); + preg + } + }; + + self.edits.push(( + ProgPoint::before(last_inst), + Edit::Move { + from: Allocation::stack(SpillSlot::new(src_slot)), + to: Allocation::reg(preg), + }, + )); + self.edits.push(( + ProgPoint::before(last_inst), + Edit::Move { + to: Allocation::stack(SpillSlot::new(dst_slot)), + from: Allocation::reg(preg), + }, + )); + slots_made_valid.push(succ_param); + self.vregs[succ_param.vreg()].set_stack_slot_valid(); + + self.vregs[out_vreg.vreg()].cur_use_idx -= 1; + } + + // save liveins + trace!(" -> Saving livein locations"); + debug_assert_eq!(self.livein_loc_lookup[succ.index()], 0); + self.livein_loc_lookup[succ.index()] = self.livein_locs.len() as u32; + for vreg in self.liveins[succ.index()].iter() { + let data = &self.vregs[vreg]; + trace!( + " -> {} at {:?} and slot {:?}", + vreg, + data.preg(), + data.stack_slot() + ); + self.livein_locs + .push(LiveinLoc::init(data.preg(), data.stack_slot())); + } + + // also save block params + for &vreg in self.func.block_params(succ) { + let data = &self.vregs[vreg.vreg()]; + trace!( + " -> {} at {:?} and slot {:?} (param)", + vreg, + data.preg(), + data.stack_slot() + ); + self.livein_locs + .push(LiveinLoc::init(data.preg(), data.stack_slot())); + } + + // reset pregs & slots + trace!(" -> Resetting pregs and slots"); + let mut vec_idx = 0; + for idx in 0..self.pregs.len() { + debug_assert!(idx < PReg::NUM_INDEX); + let data = &self.pregs[idx]; + let next_preg = &preg_allocs[vec_idx]; + if next_preg.0 as usize == idx { + if data.vreg().is_none() || data.vreg().unwrap() != next_preg.1 { + trace!(" -> Restoring {} to {}", idx, next_preg.1); + self.clear_preg(idx); + if let Some(preg) = self.vregs[next_preg.1.vreg()].preg() { + self.clear_preg(preg.index()); + } + let preg = PReg::new(idx, next_preg.1.class()); + self.assign_preg(preg, next_preg.1); + } + vec_idx += 1; + continue; + } + + self.clear_preg(idx); + } + + for &vreg in &slots_made_valid { + trace!( + " -> Setting slot {} for {} valid", + self.vregs[vreg.vreg()].allocated_stack_slot().unwrap(), + vreg + ); + self.vregs[vreg.vreg()].clear_stack_slot(true); + } + slots_made_valid.clear(); + + if cfg!(debug_assertions) { + for &vreg in &slots_made_invalid { + trace!( + " -> Setting slot {} for {} invalid", + self.vregs[vreg.vreg()].allocated_stack_slot().unwrap(), + vreg + ); + self.vregs[vreg.vreg()].set_stack_slot_valid(); + } + slots_made_invalid.clear(); + } + } + + Ok(()) } fn run(&mut self) -> Result<(), RegAllocError> { diff --git a/src/lib.rs b/src/lib.rs index 252f0a8..6944383 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1491,7 +1491,13 @@ pub fn run( env: &MachineEnv, options: &RegallocOptions, ) -> Result { - ion::run(func, env, options.verbose_log, options.validate_ssa, options.fast_alloc) + ion::run( + func, + env, + options.verbose_log, + options.validate_ssa, + options.fast_alloc, + ) } /// Options for allocation.