From db8e06d3a86cb24510ade2d4dfe591bb20449d2f Mon Sep 17 00:00:00 2001 From: Tobias Schwarz Date: Wed, 21 Jun 2023 20:43:36 +0200 Subject: [PATCH] WIP --- src/ion/fast_alloc2.rs | 204 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 197 insertions(+), 7 deletions(-) diff --git a/src/ion/fast_alloc2.rs b/src/ion/fast_alloc2.rs index eee3735..48eac9c 100644 --- a/src/ion/fast_alloc2.rs +++ b/src/ion/fast_alloc2.rs @@ -940,7 +940,7 @@ impl<'a, F: Function> FastAlloc<'a, F> { Some(slot) => slot, None => self.alloc_stack_slot(op.vreg()), }; - self.move_to_stack(op.vreg(), ProgPoint::before(inst)); + self.move_to_stack(op.vreg().vreg(), ProgPoint::before(inst)); trace!(" -> Allocated op {} to slot {}", op_idx, slot); } else { self.allocs[alloc_idx + op_idx] = Allocation::reg(preg); @@ -1406,7 +1406,197 @@ impl<'a, F: Function> FastAlloc<'a, F> { block: Block, succ: Block, ) -> Result<(), RegAllocError> { - todo!() + trace!(" -> Block has a single allocated edge. Moving"); + let succ_params = self.func.block_params(succ); + 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); + debug_assert_eq!(succ_params.len(), out_vregs.len()); + + // first handle all the stack moves + // then registers + let succ_liveins = &self.liveins[succ.index()]; + let succ_livein_loc_idx = self.livein_loc_lookup[succ.index()] as usize; + + // live vregs that need to be in certain registers + let mut live_vregs_in_pregs = Vec::new(); + + let mut loc_idx = succ_livein_loc_idx; + for live_vreg in succ_liveins.iter() { + let loc = self.livein_locs[loc_idx]; + match loc.slot() { + Some(idx) => { + if self.vregs[live_vreg].stack_slot().is_none() { + // inlined because i dont have the nerve to deal with the borrow checker + // self.move_to_stack(live_vreg, ProgPoint::before(last_inst)); + + // must exist + let slot = self.vregs[live_vreg].allocated_stack_slot().unwrap(); + + let preg = if let Some(preg) = self.vregs[live_vreg].preg() { + preg + } else { + panic!( + "Trying to move {} to stack but it is not in a preg", + live_vreg + ); + }; + + trace!("Moving {} from {} to slot {}", live_vreg, 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[live_vreg].set_stack_slot_valid(); + } else { + debug_assert_eq!(self.vregs[live_vreg].stack_slot().unwrap(), idx); + } + } + None => { + // don't need to reset stack alloc since that is handled when allocating the block + } + } + if let Some(preg) = loc.preg() { + live_vregs_in_pregs.push((live_vreg, preg)); + } else { + self.vregs[live_vreg].cur_use_idx -= 1; + } + loc_idx += 1; + } + + for i in 0..out_vregs.len() { + let out_vreg = out_vregs[i]; + let succ_vreg = succ_params[i]; + let loc = self.livein_locs[loc_idx]; + match loc.slot() { + Some(slot) => { + if self.vregs[out_vreg.vreg()].preg().is_none() { + // TODO: we could also use a single temp register + let preg = self.find_free_reg( + out_vreg.class(), + PRegSet::empty(), + Some(ProgPoint::before(last_inst)), + block, + self.cur_inst_pos + 1, + )?; + self.move_vreg_to_preg( + out_vreg, + preg, + ProgPoint::before(last_inst), + block, + self.cur_inst_pos + 1, + ); + } + + let preg = self.vregs[out_vreg.vreg()].preg().unwrap(); + // since this is the stack slot of + trace!("Moving {} from {} to slot {}", out_vreg, 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_vreg.vreg()].set_stack_slot_valid(); + } + None => {} + } + if loc.preg().is_none() { + self.vregs[out_vreg.vreg()].cur_use_idx -= 1; + } + loc_idx += 1; + } + + // move to pregs, note: this can unnecessarily spill + for &(live_vreg, target_preg) in &live_vregs_in_pregs { + self.vregs[live_vreg].cur_use_idx -= 1; + if let Some(cur_preg) = self.vregs[live_vreg].preg() { + if cur_preg == target_preg { + continue; + } + } + + if let Some(cur_vreg) = self.pregs[target_preg.index()].vreg() { + if !self.vreg_killed(cur_vreg.vreg(), block, self.cur_inst_pos + 1, true) { + if self.vregs[cur_vreg.vreg()].stack_slot().is_none() { + self.alloc_and_move_to_stack(cur_vreg, ProgPoint::before(last_inst)); + } + self.clear_preg(target_preg.index()); + } + } + + if let Some(cur_preg) = self.vregs[live_vreg].preg() { + self.vregs[live_vreg].set_preg(target_preg); + self.pregs[cur_preg.index()].clear_vreg(); + self.pregs[target_preg.index()].set_vreg(VReg::new(live_vreg, target_preg.class())); + self.edits.push(( + ProgPoint::before(last_inst), + Edit::Move { + from: Allocation::reg(cur_preg), + to: Allocation::reg(target_preg), + }, + )); + } else { + self.move_vreg_to_preg( + VReg::new(live_vreg, target_preg.class()), + target_preg, + ProgPoint::before(last_inst), + block, + self.cur_inst_pos + 1, + ); + } + } + + loc_idx -= out_vregs.len(); + for &out_vreg in out_vregs { + let loc = self.livein_locs[loc_idx]; + if let Some(target_preg) = loc.preg() { + self.vregs[out_vreg.vreg()].cur_use_idx -= 1; + if let Some(cur_preg) = self.vregs[out_vreg.vreg()].preg() { + if cur_preg == target_preg { + loc_idx += 1; + continue; + } + } + + if let Some(cur_vreg) = self.pregs[target_preg.index()].vreg() { + if !self.vreg_killed(cur_vreg.vreg(), block, self.cur_inst_pos + 1, true) { + if self.vregs[cur_vreg.vreg()].stack_slot().is_none() { + self.alloc_and_move_to_stack(cur_vreg, ProgPoint::before(last_inst)); + } + self.clear_preg(target_preg.index()); + } + } + + if let Some(cur_preg) = self.vregs[out_vreg.vreg()].preg() { + self.vregs[out_vreg.vreg()].set_preg(target_preg); + self.pregs[cur_preg.index()].clear_vreg(); + self.pregs[target_preg.index()].set_vreg(out_vreg); + self.edits.push(( + ProgPoint::before(last_inst), + Edit::Move { + from: Allocation::reg(cur_preg), + to: Allocation::reg(target_preg), + }, + )); + } else { + self.move_vreg_to_preg( + out_vreg, + target_preg, + ProgPoint::before(last_inst), + block, + self.cur_inst_pos + 1, + ); + } + } + loc_idx += 1; + } + + return Ok(()); } // TODO: check if the reftype count is correctly managed @@ -1797,19 +1987,19 @@ impl<'a, F: Function> FastAlloc<'a, F> { // panics if there is a stack slot already or the vreg is not in a preg fn alloc_and_move_to_stack(&mut self, vreg: VReg, pos: ProgPoint) { self.alloc_stack_slot(vreg); - self.move_to_stack(vreg, pos); + self.move_to_stack(vreg.vreg(), pos); } // Moves a vreg to stack // panics if there is no stack slot or the vreg is not in a preg - fn move_to_stack(&mut self, vreg: VReg, pos: ProgPoint) { - let slot = if let Some(slot) = self.vregs[vreg.vreg()].allocated_stack_slot() { + fn move_to_stack(&mut self, vreg: usize, pos: ProgPoint) { + let slot = if let Some(slot) = self.vregs[vreg].allocated_stack_slot() { slot } else { panic!("Trying to move {} to stack but it has no slot", vreg); }; - let preg = if let Some(preg) = self.vregs[vreg.vreg()].preg() { + let preg = if let Some(preg) = self.vregs[vreg].preg() { preg } else { panic!("Trying to move {} to stack but it is not in a preg", vreg); @@ -1823,7 +2013,7 @@ impl<'a, F: Function> FastAlloc<'a, F> { to: Allocation::stack(SpillSlot::new(slot as usize)), }, )); - self.vregs[vreg.vreg()].set_stack_slot_valid(); + self.vregs[vreg].set_stack_slot_valid(); } // Allocates a stack slot for a vreg