From 8fb8aa15b0e171573c29699c3e9a5d4e9fb438c5 Mon Sep 17 00:00:00 2001 From: T0b1 Date: Fri, 5 May 2023 02:10:23 +0200 Subject: [PATCH] fixed first reg version --- src/ion/fast_alloc.rs | 691 +++++++++++++++++++++++++++++++++++------- 1 file changed, 577 insertions(+), 114 deletions(-) diff --git a/src/ion/fast_alloc.rs b/src/ion/fast_alloc.rs index 1b064ad..d4b3e4c 100644 --- a/src/ion/fast_alloc.rs +++ b/src/ion/fast_alloc.rs @@ -140,6 +140,38 @@ impl ReadOnlyData { } } +// https://burtleburtle.net/bob/rand/smallprng.html +struct PRNG { + a: u64, + b: u64, + c: u64, + d: u64, +} + +impl PRNG { + fn new(seed: u64) -> Self { + Self { + a: 0xf1ea5eed, + b: seed, + c: seed, + d: seed, + } + } + + fn val(&mut self) -> u64 { + let e = self.a - PRNG::rot(self.b, 27); + self.a = self.b ^ PRNG::rot(self.c, 17); + self.b = self.c + self.d; + self.c = self.d + e; + self.d = e + self.a; + self.d + } + + fn rot(x: u64, k: u64) -> u64 { + (x << k) | (x >> (32 - k)) + } +} + struct FastAllocState<'a, F: Function> { pub vregs: Vec, pub pregs: Vec, @@ -164,6 +196,7 @@ struct FastAllocState<'a, F: Function> { pub func: &'a F, pub mach_env: &'a MachineEnv, + pub prng: PRNG, } impl<'a, F: Function> FastAllocState<'a, F> { @@ -198,6 +231,11 @@ impl<'a, F: Function> FastAllocState<'a, F> { blocks }; + trace!( + "Num Insts: {} Num Blocks: {}", + func.num_insts(), + func.num_blocks() + ); let mut inst_alloc_offsets = Vec::with_capacity(func.num_insts()); inst_alloc_offsets.resize(func.num_insts(), 0); @@ -220,6 +258,12 @@ impl<'a, F: Function> FastAllocState<'a, F> { allocs }; + trace!("InstAllocOffsets: {:?}", inst_alloc_offsets); + trace!("Allocs Len: {}", allocs.len()); + + let prng = PRNG::new( + (blocks.len() as u64) << 48 | (vregs.len() as u64) << 32 | (allocs.len() as u64), + ); Self { vregs, pregs, @@ -245,6 +289,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { func, mach_env, + prng, } } @@ -259,7 +304,10 @@ impl<'a, F: Function> FastAllocState<'a, F> { pub fn alloc_stack_slot(&mut self, vreg: VReg) -> u32 { let data = &mut self.vregs[vreg.vreg()]; if data.slot_idx.is_some() { - panic!("Trying to allocate already allocated stack slot"); + panic!( + "Trying to allocate already allocated stack slot for {}", + vreg + ); } let size = if vreg.class() == RegClass::Int { @@ -286,14 +334,21 @@ impl<'a, F: Function> FastAllocState<'a, F> { } pub fn move_to_preg(&mut self, vreg: VReg, preg: PReg, pos: ProgPoint) { - if let Some(vreg) = &self.pregs[preg.index()].vreg { + trace!("Move {} to {} at {:?}", vreg, preg, pos); + /*if let Some(vreg) = &self.pregs[preg.index()].vreg { let vdata = &mut self.vregs[vreg.vreg() as usize]; debug_assert!(vdata.preg.is_some()); debug_assert_eq!(vdata.preg.unwrap(), preg); vdata.preg = None; - } + }*/ if let Some(cur_preg) = &self.vregs[vreg.vreg()].preg { + if *cur_preg == preg { + trace!("{} already in target reg", vreg); + return; + } + + trace!("Moving directly from {} to {}", cur_preg, preg); // Do a reg->reg move self.edits.push(( pos, @@ -310,9 +365,11 @@ impl<'a, F: Function> FastAllocState<'a, F> { pdata.vreg = None; self.pregs[preg.index()].vreg = Some(vreg); + self.vregs[vreg.vreg()].preg = Some(preg); return; } + self.clear_preg(preg); let vdata = &mut self.vregs[vreg.vreg()]; let pdata = &mut self.pregs[preg.index()]; @@ -320,10 +377,12 @@ impl<'a, F: Function> FastAllocState<'a, F> { panic!("Trying to move from vreg that has no stack slot to preg"); } + let slot = vdata.slot_idx.unwrap() as usize; + trace!("Moving from slot {}", slot); self.edits.push(( pos, Edit::Move { - from: Allocation::stack(SpillSlot::new(vdata.slot_idx.unwrap() as usize)), + from: Allocation::stack(SpillSlot::new(slot)), to: Allocation::reg(preg), }, )); @@ -336,6 +395,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { } pub fn move_to_stack(&mut self, preg: PReg, vreg: VReg, pos: ProgPoint) { + trace!("Move {} of {} to stack at {:?}", preg, vreg, pos); let vdata = &mut self.vregs[vreg.vreg()]; let pdata = &mut self.pregs[preg.index()]; if pdata.vreg.is_none() || vdata.preg.is_none() { @@ -358,6 +418,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { } pub fn assign_preg(&mut self, preg: PReg, vreg: VReg) { + trace!("Assigning {} to {}", vreg, preg); // TODO: somewhere assign_preg is called without making sure the vreg is clear (or inspite of it) // need to make sure this is intended behavior self.clear_preg(preg); @@ -374,6 +435,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { } fn clear_preg_idx(&mut self, preg: usize) { + trace!("Clearing preg {}", preg); let pdata = &mut self.pregs[preg]; if let Some(vreg) = pdata.vreg { let vdata = &mut self.vregs[vreg.vreg()]; @@ -387,8 +449,10 @@ impl<'a, F: Function> FastAllocState<'a, F> { } pub fn clear_vreg_from_reg(&mut self, vreg: VReg) { + trace!("Clearing vreg {} from reg", vreg); let vdata = &mut self.vregs[vreg.vreg()]; if let Some(preg) = vdata.preg { + debug_assert!(self.pregs[preg.index()].vreg.is_some()); debug_assert_eq!(self.pregs[preg.index()].vreg.unwrap().vreg(), vreg.vreg()); self.pregs[preg.index()].vreg = None; vdata.preg = None; @@ -413,6 +477,23 @@ impl<'a, F: Function> FastAllocState<'a, F> { } } } + + pub fn vreg_used_at_cur_inst(&mut self, vreg: VReg) -> bool { + let vdata = &self.vregs[vreg.vreg()]; + if vdata.cur_use_idx as usize >= vdata.uses.len() { + return false; + } + vdata.uses[vdata.cur_use_idx as usize] == self.cur_inst_pos as u32 + } + + pub fn vreg_next_use(&self, vreg: VReg) -> Option { + let vdata = &self.vregs[vreg.vreg()]; + if vdata.cur_use_idx as usize >= vdata.uses.len() { + return None; + } + + Some(vdata.uses[vdata.cur_use_idx as usize]) + } } pub fn run(func: &F, mach_env: &MachineEnv) -> Result { @@ -437,6 +518,12 @@ pub fn run(func: &F, mach_env: &MachineEnv) -> Result( block: Block, block_last_inst: usize, vreg: usize, + // if the vreg is used at the current instruction, count it as not killed + // TODO: this is currently always true but can be used for optimization later on + save_on_current_use: bool, ) -> bool { let info = &state.vregs[vreg]; let block_after_pos = state.cur_inst_pos + (block_last_inst - inst.index()) + 1; let cur_use_idx = info.cur_use_idx as usize; + let cur_pos = if save_on_current_use { + state.cur_inst_pos + } else { + state.cur_inst_pos + 1 + }; + + trace!( + "Checking live-status of v{} in {:?} at inst {:?} (CurPos: {} SaveCurPos? {}): Liveout {}, block_after: {}", + vreg, + block, + inst, + state.cur_inst_pos, + save_on_current_use, + state.liveouts[block.index()].get(vreg), + block_after_pos + ); + trace!( + "Uses of v{}: {:?}. Currently at {}", + vreg, + info.uses, + info.cur_use_idx + ); if !state.liveouts[block.index()].get(vreg) { - if info.uses.len() >= cur_use_idx { + if info.uses.len() <= cur_use_idx { + trace!("Uses exhausted, vreg must be dead"); return true; } - if info.uses[cur_use_idx] == state.cur_inst_pos as u32 { - if info.uses.len() >= cur_use_idx + 1 + + if info.uses.len() <= cur_use_idx + 1 { + trace!("next use: {}, no use after that", info.uses[cur_use_idx]); + } else { + trace!( + "next use: {}, {}", + info.uses[cur_use_idx], + info.uses[cur_use_idx + 1] + ); + } + + if !save_on_current_use && info.uses[cur_use_idx] == state.cur_inst_pos as u32 { + if info.uses.len() <= cur_use_idx + 1 || info.uses[cur_use_idx + 1] >= block_after_pos as u32 { + trace!("v{} is killed", vreg); return true; } } @@ -574,25 +699,61 @@ fn allocate_block_insts<'a, F: Function>( } } + for preg in clobbers { + // TODO: this might save a use that is killed at this inst i think + let vreg = if let Some(vreg) = &state.pregs[preg.index()].vreg { + *vreg + } else { + continue; + }; + + if state.vregs[vreg.vreg()].slot_idx.is_some() { + trace!("{} with {} clobbered but saved on stack", preg, vreg); + state.clear_preg(preg); + continue; + } + + if !state.vreg_used_at_cur_inst(vreg) + && vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg(), true) + { + trace!("{} with {} clobbered but vreg killed", preg, vreg); + state.clear_preg(preg); + continue; + } + + state.alloc_stack_slot(vreg); + state.move_to_stack(preg, vreg, ProgPoint::before(inst)); + state.clear_preg(preg); + } + // we allocate fixed defs/uses and stack allocations first + // TODO: if a fixed def is before a fixed use with the same preg here, it will incorrectly update state!!! trace!("First alloc pass"); for (i, op) in operands.iter().enumerate() { let vreg = op.vreg(); trace!("Operand {}: {}", i, op); - if vreg == VReg::invalid() { + if op.as_fixed_nonallocatable().is_some() { // it seems cranelift emits fixed reg uses with invalid vregs, handle them here // TODO: treat them like normal vregs by just using last_vreg_index+1 for them? match op.constraint() { OperandConstraint::FixedReg(reg) => { // Save vreg if needed if let Some(vreg) = state.pregs[reg.index()].vreg { - let vreg = vreg.vreg(); - if state.vregs[vreg].slot_idx.is_none() - && !vreg_killed(state, inst, block, block_last_inst_idx, vreg) + let vreg_idx = vreg.vreg(); + if state.vregs[vreg_idx].slot_idx.is_none() + && (state.vreg_used_at_cur_inst(vreg) + || !vreg_killed( + state, + inst, + block, + block_last_inst_idx, + vreg_idx, + true, + )) { let slot = state.create_stack_slot(reg.class()); - state.vregs[vreg].slot_idx = Some(slot); + state.vregs[vreg_idx].slot_idx = Some(slot); state.edits.push(( ProgPoint::before(inst), Edit::Move { @@ -642,18 +803,27 @@ fn allocate_block_insts<'a, F: Function>( if let Some(cur_preg) = state.vregs[vreg.vreg()].preg { if cur_preg == reg { trace!("{} already in target {}", vreg, cur_preg); + state.allocs[alloc_idx + i] = Allocation::reg(cur_preg); continue; } } // Save vreg if needed if let Some(vreg) = state.pregs[reg.index()].vreg { - let vreg = vreg.vreg(); - if state.vregs[vreg].slot_idx.is_none() - && !vreg_killed(state, inst, block, block_last_inst_idx, vreg) + let vreg_idx = vreg.vreg(); + if state.vregs[vreg_idx].slot_idx.is_none() + && (state.vreg_used_at_cur_inst(vreg) + || !vreg_killed( + state, + inst, + block, + block_last_inst_idx, + vreg_idx, + true, + )) { let slot = state.create_stack_slot(reg.class()); - state.vregs[vreg].slot_idx = Some(slot); + state.vregs[vreg_idx].slot_idx = Some(slot); state.edits.push(( ProgPoint::before(inst), Edit::Move { @@ -665,6 +835,7 @@ fn allocate_block_insts<'a, F: Function>( } if let Some(cur_preg) = state.vregs[vreg.vreg()].preg { + trace!("Move {} directly from {} to {}", vreg, cur_preg, reg); // Move from preg to preg state.edits.push(( ProgPoint::before(inst), @@ -673,9 +844,12 @@ fn allocate_block_insts<'a, F: Function>( to: Allocation::reg(reg), }, )); + debug_assert_eq!( + state.pregs[cur_preg.index()].vreg.unwrap().vreg(), + vreg.vreg() + ); state.pregs[cur_preg.index()].vreg = None; - state.vregs[vreg.vreg()].preg = Some(reg); - state.pregs[reg.index()].vreg = Some(vreg); + state.assign_preg(reg, vreg); } else { state.move_to_preg(vreg, reg, ProgPoint::before(inst)); } @@ -715,12 +889,21 @@ fn allocate_block_insts<'a, F: Function>( // Save vreg if needed if let Some(vreg) = state.pregs[reg.index()].vreg { - let vreg = vreg.vreg(); - if state.vregs[vreg].slot_idx.is_none() - && !vreg_killed(state, inst, block, block_last_inst_idx, vreg) + let vreg_idx = vreg.vreg(); + if state.vregs[vreg_idx].slot_idx.is_none() + && (op.pos() != OperandPos::Late + && state.vreg_used_at_cur_inst(vreg) + || !vreg_killed( + state, + inst, + block, + block_last_inst_idx, + vreg_idx, + true, + )) { let slot = state.create_stack_slot(reg.class()); - state.vregs[vreg].slot_idx = Some(slot); + state.vregs[vreg_idx].slot_idx = Some(slot); state.edits.push(( ProgPoint::before(inst), Edit::Move { @@ -769,18 +952,21 @@ fn allocate_block_insts<'a, F: Function>( let tmp_reg = pregs.into_iter().next().unwrap(); if let Some(vreg) = state.pregs[tmp_reg.index()].vreg { // Save vreg if needed - let vreg = vreg.vreg(); - if state.vregs[vreg].slot_idx.is_none() - && !vreg_killed( - state, - inst, - block, - block_last_inst_idx, - vreg, - ) + let vreg_idx = vreg.vreg(); + if state.vregs[vreg_idx].slot_idx.is_none() + && (op.pos() != OperandPos::Late + && state.vreg_used_at_cur_inst(vreg) + || !vreg_killed( + state, + inst, + block, + block_last_inst_idx, + vreg_idx, + true, + )) { let slot = state.create_stack_slot(reg.class()); - state.vregs[vreg].slot_idx = Some(slot); + state.vregs[vreg_idx].slot_idx = Some(slot); state.edits.push(( ProgPoint::before(inst), Edit::Move { @@ -851,12 +1037,24 @@ fn allocate_block_insts<'a, F: Function>( trace!("Operand {}: {}", i, op); let vreg = op.vreg(); - if vreg == VReg::invalid() { + if op.as_fixed_nonallocatable().is_some() { continue; } match op.constraint() { OperandConstraint::Reg => { + // Are we alredy in a reg? + if let Some(cur_preg) = &state.vregs[vreg.vreg()].preg { + assert_eq!(op.kind(), OperandKind::Use); + // Late uses need to survive the instruction + if op.pos() == OperandPos::Early || !clobbers.contains(*cur_preg) { + trace!("{} already in reg {}. Using that", vreg, cur_preg); + state.allocs[alloc_idx + i] = Allocation::reg(*cur_preg); + regs_allocated.add(*cur_preg); + continue; + } + } + // find first non-allocated register let reg_order = const_state.reg_order(op.class()); let mut allocated = false; @@ -865,15 +1063,34 @@ fn allocate_block_insts<'a, F: Function>( continue; } - if state.pregs[reg.index()].vreg.is_some() { - continue; + if let Some(cur_vreg) = &state.pregs[reg.index()].vreg { + // we can override the reg if the vreg was killed already + if !vreg_killed( + state, + inst, + block, + block_last_inst_idx, + cur_vreg.vreg(), + true, + ) { + continue; + } + state.clear_preg(reg); } // reg should not contain anything debug_assert!(state.pregs[reg.index()].vreg.is_none()); + if op.kind() == OperandKind::Use + && op.pos() == OperandPos::Late + && clobbers.contains(reg) + { + continue; + } + state.allocs[alloc_idx + i] = Allocation::reg(reg); regs_allocated.add(reg); + trace!("Chose {} for operand {}", reg, i); if op.kind() == OperandKind::Use { if req_refs_on_stack && state.vregs[vreg.vreg()].reftype { panic!("reftype required to be in reg at safepoint"); @@ -890,7 +1107,6 @@ fn allocate_block_insts<'a, F: Function>( state.move_to_stack(reg, vreg, ProgPoint::after(inst)); } - trace!("Chose {} for operand {}", reg, i); allocated = true; break; } @@ -903,44 +1119,69 @@ fn allocate_block_insts<'a, F: Function>( // No register available // TODO: first evict pregs that already have a stack slot even if they are used earlier? let mut evict_candidate = None; + let mut ffa_reg_pool = PRegSet::empty(); for ® in reg_order { if regs_allocated.contains(reg) { continue; } - debug_assert!(state.pregs[reg.index()].vreg.is_some()); - let vreg = state.pregs[reg.index()].vreg.unwrap().vreg(); - let next_use = - state.vregs[vreg].uses[state.vregs[vreg].cur_use_idx as usize]; - if next_use == state.cur_inst_pos as u32 { + if op.kind() == OperandKind::Use + && op.pos() == OperandPos::Late + && clobbers.contains(reg) + { continue; } - if let Some((_, pos)) = &evict_candidate { - if *pos < next_use { + debug_assert!(state.pregs[reg.index()].vreg.is_some()); + let vreg = state.pregs[reg.index()].vreg.unwrap(); + if let Some(next_use) = state.vreg_next_use(vreg) { + if next_use == state.cur_inst_pos as u32 { + continue; + } + + if let Some((_, pos)) = &evict_candidate { + if *pos < next_use { + evict_candidate = Some((reg, next_use)); + } + } else { evict_candidate = Some((reg, next_use)); } } else { - evict_candidate = Some((reg, next_use)); + // see further below + ffa_reg_pool.add(reg); + + //panic!("preg should have already been chosen") } } + // TODO: we need some logic to shuffle assignments around if there is a late use that needs to survive a clobber + // and another reg is available but taken by an early use so it would not be an eviction_candidate + if let Some((reg, next_use)) = evict_candidate { // Save vreg if needed - let vreg = state.pregs[reg.index()].vreg.unwrap(); - log::trace!("Evicting {} with v{}", reg, vreg); - if state.vregs[vreg.vreg()].slot_idx.is_none() - && !vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg()) { - let slot = state.create_stack_slot(reg.class()); - state.vregs[vreg.vreg()].slot_idx = Some(slot); - state.edits.push(( - ProgPoint::before(inst), - Edit::Move { - from: Allocation::reg(reg), - to: Allocation::stack(SpillSlot::new(slot as usize)), - }, - )); + let vreg = state.pregs[reg.index()].vreg.unwrap(); + trace!("Evicting {} with v{}", reg, vreg); + if state.vregs[vreg.vreg()].slot_idx.is_none() + && !vreg_killed( + state, + inst, + block, + block_last_inst_idx, + vreg.vreg(), + true, + ) + { + let slot = state.create_stack_slot(reg.class()); + state.vregs[vreg.vreg()].slot_idx = Some(slot); + state.edits.push(( + ProgPoint::before(inst), + Edit::Move { + from: Allocation::reg(reg), + to: Allocation::stack(SpillSlot::new(slot as usize)), + }, + )); + } } state.clear_preg(reg); @@ -964,8 +1205,61 @@ fn allocate_block_insts<'a, F: Function>( trace!("Chose {} for operand {}", reg, i); } else { - panic!("Out of registers: {:?}", regs_allocated); - return Err(RegAllocError::TooManyLiveRegs); + if ffa_reg_pool == PRegSet::empty() { + panic!("Out of registers: {:?}", regs_allocated); + return Err(RegAllocError::TooManyLiveRegs); + } + + let preg = 'block: { + let len = ffa_reg_pool.bits.count_ones() as usize; + let mut idx = (state.prng.val() as usize % 128) % len; + for preg in ffa_reg_pool.into_iter() { + if idx == 0 { + break 'block preg; + } + + idx -= 1; + } + + panic!("I can't do math"); + }; + + trace!("Chose {} from ffa_reg_pool", preg); + { + let vreg = state.pregs[preg.index()].vreg.unwrap(); + // need to save vreg if it does not have a slot + if state.vregs[vreg.vreg()].slot_idx.is_none() { + let slot = state.create_stack_slot(preg.class()); + state.vregs[vreg.vreg()].slot_idx = Some(slot); + state.edits.push(( + ProgPoint::before(inst), + Edit::Move { + from: Allocation::reg(preg), + to: Allocation::stack(SpillSlot::new(slot as usize)), + }, + )); + } + } + state.clear_preg(preg); + state.allocs[alloc_idx + i] = Allocation::reg(preg); + regs_allocated.add(preg); + if op.kind() == OperandKind::Use { + if req_refs_on_stack && state.vregs[vreg.vreg()].reftype { + panic!("reftype required to be in reg at safepoint"); + return Err(RegAllocError::TooManyLiveRegs); + } + + // need to move from stack to reg + state.move_to_preg(vreg, preg, ProgPoint::before(inst)); + } else { + // early def + state.vregs[vreg.vreg()].def_block = Some(block); + state.assign_preg(preg, vreg); + state.alloc_stack_slot(vreg); + state.move_to_stack(preg, vreg, ProgPoint::after(inst)); + } + + trace!("Chose {} for operand {}", preg, i); } } OperandConstraint::Reuse(_) => { @@ -975,6 +1269,62 @@ fn allocate_block_insts<'a, F: Function>( } } + // advance use_idx + for op in operands { + if op.kind() != OperandKind::Use || op.as_fixed_nonallocatable().is_some() { + continue; + } + + let vreg_idx = op.vreg().vreg(); + let info = &mut state.vregs[vreg_idx]; + info.cur_use_idx += 1; + + if vreg_killed(state, inst, block, block_last_inst_idx, vreg_idx, true) { + // TODO: clear stack slot + state.clear_vreg_from_reg(op.vreg()); + } + /*let block_after_pos = state.cur_inst_pos + (block_last_inst_idx - inst.index()) + 1; + // check if vreg dies here + if !state.liveouts[block.index()].get(vreg_idx) + && (info.uses.len() <= info.cur_use_idx as usize + || info.uses[info.cur_use_idx as usize] > block_after_pos as u32) + { + // TODO: clear stack slot + state.clear_vreg_from_reg(op.vreg()); + }*/ + } + + // TODO: this is currently a fix for the register state for uses that are clobbered as it is incorrectly set + // but this is inefficient as we could check for this when handling uses + trace!("Late clobber handling"); + for preg in clobbers { + // TODO: this might save a use that is killed at this inst i think + let vreg = if let Some(vreg) = &state.pregs[preg.index()].vreg { + *vreg + } else { + continue; + }; + + if state.vregs[vreg.vreg()].slot_idx.is_some() { + trace!("{} with {} clobbered but saved on stack", preg, vreg); + state.clear_preg(preg); + continue; + } + + // we don't care if the reg is used at the current inst + if vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg(), true) { + trace!("{} with {} clobbered but vreg killed", preg, vreg); + state.clear_preg(preg); + continue; + } + + // TODO: this should not be hit i think as all we should be clearing here are use assignments and the vregs + // that need to be saved should have been saved at the check before + state.alloc_stack_slot(vreg); + state.move_to_stack(preg, vreg, ProgPoint::before(inst)); + state.clear_preg(preg); + } + // alloc non-fixed late defs and reuse trace!("Third alloc pass"); for (i, op) in operands.iter().enumerate() { @@ -984,7 +1334,7 @@ fn allocate_block_insts<'a, F: Function>( trace!("Operand {}: {}", i, op); let vreg = op.vreg(); - if vreg == VReg::invalid() { + if op.as_fixed_nonallocatable().is_some() { continue; } @@ -999,8 +1349,19 @@ fn allocate_block_insts<'a, F: Function>( if regs_allocated.contains(reg) || late_write_disallow_regs.contains(reg) { continue; } - if state.pregs[reg.index()].vreg.is_some() { - continue; + if let Some(cur_vreg) = &state.pregs[reg.index()].vreg { + // we can override the reg if the vreg was killed already + if !vreg_killed( + state, + inst, + block, + block_last_inst_idx, + cur_vreg.vreg(), + true, + ) { + continue; + } + state.clear_preg(reg); } // reg should not contain anything @@ -1024,44 +1385,60 @@ fn allocate_block_insts<'a, F: Function>( // TODO: first evict pregs that already have a stack slot even if they are used earlier? let mut evict_candidate = None; + let mut ffa_reg_pool = PRegSet::empty(); for ® in reg_order { if regs_allocated.contains(reg) || late_write_disallow_regs.contains(reg) { continue; } debug_assert!(state.pregs[reg.index()].vreg.is_some()); - let vreg = state.pregs[reg.index()].vreg.unwrap().vreg(); - let next_use = - state.vregs[vreg].uses[state.vregs[vreg].cur_use_idx as usize]; - if next_use == state.cur_inst_pos as u32 { - continue; - } + let vreg = state.pregs[reg.index()].vreg.unwrap(); + if let Some(next_use) = state.vreg_next_use(vreg) { + if next_use == state.cur_inst_pos as u32 { + continue; + } - if let Some((_, pos)) = &evict_candidate { - if *pos < next_use { + if let Some((_, pos)) = &evict_candidate { + if *pos < next_use { + evict_candidate = Some((reg, next_use)); + } + } else { evict_candidate = Some((reg, next_use)); } } else { - evict_candidate = Some((reg, next_use)); + // if we hit this it means that all uses are "before" this one in lowering-order + // we should probably find a nice heuristic for chosing which register to choose + // here. tbf we should probably find an overall better heuristic for chosing which register to evict + // rn just add the reg to a set and we pick a random one later + ffa_reg_pool.add(reg); } } if let Some((reg, next_use)) = evict_candidate { // Save vreg if needed - let vreg = state.pregs[reg.index()].vreg.unwrap(); - log::trace!("Evicting {} with v{}", reg, vreg); - if state.vregs[vreg.vreg()].slot_idx.is_none() - && !vreg_killed(state, inst, block, block_last_inst_idx, vreg.vreg()) { - let slot = state.create_stack_slot(reg.class()); - state.vregs[vreg.vreg()].slot_idx = Some(slot); - state.edits.push(( - ProgPoint::before(inst), - Edit::Move { - from: Allocation::reg(reg), - to: Allocation::stack(SpillSlot::new(slot as usize)), - }, - )); + let vreg = state.pregs[reg.index()].vreg.unwrap(); + trace!("Evicting {} with {}", reg, vreg); + if state.vregs[vreg.vreg()].slot_idx.is_none() + && !vreg_killed( + state, + inst, + block, + block_last_inst_idx, + vreg.vreg(), + true, + ) + { + let slot = state.create_stack_slot(reg.class()); + state.vregs[vreg.vreg()].slot_idx = Some(slot); + state.edits.push(( + ProgPoint::before(inst), + Edit::Move { + from: Allocation::reg(reg), + to: Allocation::stack(SpillSlot::new(slot as usize)), + }, + )); + } } state.clear_preg(reg); @@ -1073,8 +1450,61 @@ fn allocate_block_insts<'a, F: Function>( state.move_to_stack(reg, vreg, ProgPoint::after(inst)); trace!("Chose {} for operand {}", reg, i); } else { - panic!("out of registers"); - return Err(RegAllocError::TooManyLiveRegs); + if ffa_reg_pool == PRegSet::empty() { + panic!("Out of registers: {:?}", regs_allocated); + return Err(RegAllocError::TooManyLiveRegs); + } + + let preg = 'block: { + let len = ffa_reg_pool.bits.count_ones() as usize; + let mut idx = (state.prng.val() as usize % 128) % len; + for preg in ffa_reg_pool.into_iter() { + if idx == 0 { + break 'block preg; + } + + idx -= 1; + } + + panic!("I can't do math"); + }; + + trace!("Chose {} from ffa_reg_pool", preg); + { + let vreg = state.pregs[preg.index()].vreg.unwrap(); + // need to save vreg if it does not have a slot + if state.vregs[vreg.vreg()].slot_idx.is_none() { + let slot = state.create_stack_slot(preg.class()); + state.vregs[vreg.vreg()].slot_idx = Some(slot); + state.edits.push(( + ProgPoint::before(inst), + Edit::Move { + from: Allocation::reg(preg), + to: Allocation::stack(SpillSlot::new(slot as usize)), + }, + )); + } + } + state.clear_preg(preg); + state.allocs[alloc_idx + i] = Allocation::reg(preg); + regs_allocated.add(preg); + if op.kind() == OperandKind::Use { + if req_refs_on_stack && state.vregs[vreg.vreg()].reftype { + panic!("reftype required to be in reg at safepoint"); + return Err(RegAllocError::TooManyLiveRegs); + } + + // need to move from stack to reg + state.move_to_preg(vreg, preg, ProgPoint::before(inst)); + } else { + // early def + state.vregs[vreg.vreg()].def_block = Some(block); + state.assign_preg(preg, vreg); + state.alloc_stack_slot(vreg); + state.move_to_stack(preg, vreg, ProgPoint::after(inst)); + } + + trace!("Chose {} for operand {}", preg, i); } } OperandConstraint::Reuse(idx) => { @@ -1085,12 +1515,13 @@ fn allocate_block_insts<'a, F: Function>( // Save vreg on stack if it is not killed if let Some(vreg) = state.pregs[preg.index()].vreg { - let vreg = vreg.vreg(); - if state.vregs[vreg].slot_idx.is_none() - && !vreg_killed(state, inst, block, block_last_inst_idx, vreg) + let vreg_idx = vreg.vreg(); + if state.vregs[vreg_idx].slot_idx.is_none() + && !vreg_killed(state, inst, block, block_last_inst_idx, vreg_idx, true) { + trace!("Saving {}", vreg); let slot = state.create_stack_slot(preg.class()); - state.vregs[vreg].slot_idx = Some(slot); + state.vregs[vreg_idx].slot_idx = Some(slot); state.edits.push(( ProgPoint::before(inst), Edit::Move { @@ -1112,26 +1543,6 @@ fn allocate_block_insts<'a, F: Function>( } } - for op in operands { - if op.kind() != OperandKind::Use || op.as_fixed_nonallocatable().is_some() { - continue; - } - - let vreg_idx = op.vreg().vreg(); - let info = &mut state.vregs[vreg_idx]; - info.cur_use_idx += 1; - - let block_after_pos = state.cur_inst_pos + (block_last_inst_idx - inst.index()) + 1; - // check if vreg dies here - if !state.liveouts[block.index()].get(vreg_idx) - && (info.uses.len() >= info.cur_use_idx as usize - || info.uses[info.cur_use_idx as usize] > block_after_pos as u32) - { - // TODO: clear stack slot - state.clear_vreg_from_reg(op.vreg()); - } - } - // fixup edit order let mut first_post_pos = None; for i in edit_start_idx..state.edits.len() { @@ -1151,16 +1562,45 @@ fn allocate_block_insts<'a, F: Function>( } } + assert!(!state.allocs[alloc_idx..alloc_idx + operands.len()] + .iter() + .any(|a| a.is_none())); + + trace!( + "Instruction Allocs: {:?}", + &state.allocs[alloc_idx..alloc_idx + operands.len()] + ); + state.cur_inst_pos += 1; } - // Move all liveout vregs to a stack slot if they dont have one and clear pregs + // Move all liveout/block param vregs to a stack slot if they dont have one and clear pregs for i in 0..state.pregs.len() { match state.pregs[i].vreg { None => {} Some(vreg) => { + trace!("Clearing {} from p{}", vreg, i); let idx = vreg.vreg(); - if state.liveouts[block.index()].get(idx) && state.vregs[idx].slot_idx.is_none() { + // TODO: obv dont need that if the block param handle funcs can handle reg locations + let is_out_param = 'block: { + let last_inst = state.func.block_insns(block).last(); + if !state.func.is_branch(last_inst) { + break 'block false; + } + + for succ_idx in 0..state.func.block_succs(block).len() { + for out_vreg in state.func.branch_blockparams(block, last_inst, succ_idx) { + if *out_vreg == vreg { + break 'block true; + } + } + } + + false + }; + if (is_out_param || state.liveouts[block.index()].get(idx)) + && state.vregs[idx].slot_idx.is_none() + { let preg = state.vregs[idx].preg.unwrap(); let slot = state.create_stack_slot(preg.class()); state.edits.push(( @@ -1187,8 +1627,12 @@ fn handle_out_block_params<'a, F: Function>( const_state: &ReadOnlyData, block: Block, ) -> Result<(), RegAllocError> { - trace!("Allocating outgoing blockparams for {}", block.index()); let last_inst = state.func.block_insns(block).last(); + trace!( + "Allocating outgoing blockparams for {}, last_inst: {}", + block.index(), + last_inst.index() + ); if !state.func.is_branch(last_inst) { trace!("Last inst {} is not a branch", last_inst.index()); return Ok(()); @@ -1198,10 +1642,11 @@ fn handle_out_block_params<'a, F: Function>( { let alloc_start = state.inst_alloc_offsets[last_inst.index()] as usize; let alloc_end = if last_inst.index() + 1 == state.inst_alloc_offsets.len() { - state.inst_alloc_offsets.len() + state.allocs.len() } else { state.inst_alloc_offsets[last_inst.index() + 1] as usize }; + trace!("alloc_start: {}, alloc_end: {}", alloc_start, alloc_end); for i in alloc_start..alloc_end { if let Some(reg) = state.allocs[i].clone().as_reg() { pregs_used_by_br.add(reg); @@ -1686,6 +2131,7 @@ fn handle_out_block_params<'a, F: Function>( let out_vreg = out_params[i]; let in_vreg = in_params[i]; debug_assert!(state.vregs[out_vreg.vreg()].slot_idx.is_some()); + debug_assert!(state.vregs[in_vreg.vreg()].slot_idx.is_none()); let out_slot_idx = state.vregs[out_vreg.vreg()].slot_idx.unwrap(); if out_vreg == VReg::invalid() { @@ -1807,6 +2253,12 @@ fn calc_use_positions_and_live_bitmaps<'a, F: Function>( continue; } + trace!( + "Use of {} at {} (inst {})", + op.vreg(), + cur_pos, + inst.index() + ); state.vregs[op.vreg().vreg()].uses.push(cur_pos); } @@ -1820,6 +2272,12 @@ fn calc_use_positions_and_live_bitmaps<'a, F: Function>( for i in 0..state.func.block_succs(block).len() { for vreg in state.func.branch_blockparams(block, last_inst, i) { + trace!( + "Use of {} in blockparam at {} (inst {})", + vreg, + cur_pos, + last_inst.index() + ); state.vregs[vreg.vreg()].uses.push(cur_pos); } } @@ -1955,5 +2413,10 @@ fn calc_live_bitmaps<'a, F: Function>( return Err(RegAllocError::EntryLivein); } + for idx in 0..state.blocks.len() { + trace!("Livein for block {}: {:?}", idx, state.liveins[idx]); + trace!("Liveouts for block {}: {:?}", idx, state.liveouts[idx]); + } + Ok(()) }