From b6cc306d7a16e92238946f6ef0f8beb4234676ef Mon Sep 17 00:00:00 2001 From: T0b1 Date: Sun, 30 Apr 2023 23:10:59 +0200 Subject: [PATCH] first steps to keeping stuff in register --- src/ion/fast_alloc.rs | 386 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 358 insertions(+), 28 deletions(-) diff --git a/src/ion/fast_alloc.rs b/src/ion/fast_alloc.rs index bd76b1c..1b064ad 100644 --- a/src/ion/fast_alloc.rs +++ b/src/ion/fast_alloc.rs @@ -30,7 +30,7 @@ struct VRegData { #[derive(Default, Clone, Copy)] struct PRegData { - pub vreg: Option, + pub vreg: Option, pub stack_pseudo: bool, } @@ -153,6 +153,8 @@ struct FastAllocState<'a, F: Function> { pub stack_slot_count_int: u8, pub stack_slot_count_float: u8, + pub cur_inst_pos: usize, + pub allocs: Vec, pub inst_alloc_offsets: Vec, pub edits: Vec<(ProgPoint, Edit)>, @@ -227,6 +229,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { cur_stack_slot_idx: 0, reftype_vregs_in_pregs_count: 0, + cur_inst_pos: 0, stack_slot_count_int: u8::try_from(func.spillslot_size(RegClass::Int)) .expect("that's a big integer"), @@ -284,18 +287,30 @@ 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 { - let vdata = &mut self.vregs[*vreg as usize]; + 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(preg) = &self.vregs[vreg.vreg()].preg { + if let Some(cur_preg) = &self.vregs[vreg.vreg()].preg { + // Do a reg->reg move + self.edits.push(( + pos, + Edit::Move { + from: Allocation::reg(*cur_preg), + to: Allocation::reg(preg), + }, + )); + // TODO: allow multiple pregs for a single vreg? - let pdata = &mut self.pregs[preg.index()]; + let pdata = &mut self.pregs[cur_preg.index()]; debug_assert!(pdata.vreg.is_some()); - debug_assert_eq!(pdata.vreg.unwrap(), vreg.vreg() as u32); + debug_assert_eq!(pdata.vreg.unwrap().vreg(), vreg.vreg()); pdata.vreg = None; + + self.pregs[preg.index()].vreg = Some(vreg); + return; } let vdata = &mut self.vregs[vreg.vreg()]; @@ -313,7 +328,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { }, )); vdata.preg = Some(preg); - pdata.vreg = Some(vreg.vreg() as u32); + pdata.vreg = Some(vreg); if vdata.reftype { self.reftype_vregs_in_pregs_count += 1; @@ -327,7 +342,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { panic!("Trying to move from unallocated preg/vreg to stack"); } debug_assert_eq!(vdata.preg.unwrap(), preg); - debug_assert_eq!(pdata.vreg.unwrap(), vreg.vreg() as u32); + debug_assert_eq!(pdata.vreg.unwrap().vreg(), vreg.vreg()); if vdata.slot_idx.is_none() { panic!("Trying to move to vreg without stack slot"); @@ -347,7 +362,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { // need to make sure this is intended behavior self.clear_preg(preg); - self.pregs[preg.index()].vreg = Some(vreg.vreg() as u32); + self.pregs[preg.index()].vreg = Some(vreg); self.vregs[vreg.vreg()].preg = Some(preg); if self.vregs[vreg.vreg()].reftype { self.reftype_vregs_in_pregs_count += 1; @@ -361,7 +376,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { fn clear_preg_idx(&mut self, preg: usize) { let pdata = &mut self.pregs[preg]; if let Some(vreg) = pdata.vreg { - let vdata = &mut self.vregs[vreg as usize]; + let vdata = &mut self.vregs[vreg.vreg()]; debug_assert_eq!(vdata.preg.unwrap().index(), preg); vdata.preg = None; pdata.vreg = None; @@ -374,7 +389,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { pub fn clear_vreg_from_reg(&mut self, vreg: VReg) { let vdata = &mut self.vregs[vreg.vreg()]; if let Some(preg) = vdata.preg { - debug_assert_eq!(self.pregs[preg.index()].vreg.unwrap(), vreg.vreg() as u32); + debug_assert_eq!(self.pregs[preg.index()].vreg.unwrap().vreg(), vreg.vreg()); self.pregs[preg.index()].vreg = None; vdata.preg = None; @@ -391,7 +406,7 @@ impl<'a, F: Function> FastAllocState<'a, F> { for i in 0..self.pregs.len() { if let Some(vreg) = self.pregs[i].vreg.clone() { - let vreg = vreg as usize; + let vreg = vreg.vreg(); if self.vregs[vreg].reftype { self.clear_preg_idx(i); } @@ -425,6 +440,11 @@ pub fn run(func: &F, mach_env: &MachineEnv) -> Result(func: &F, mach_env: &MachineEnv) -> Result( + state: &FastAllocState<'a, F>, + inst: Inst, + block: Block, + block_last_inst: usize, + vreg: usize, +) -> 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; + + if !state.liveouts[block.index()].get(vreg) { + if info.uses.len() >= cur_use_idx { + return true; + } + if 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 + { + return true; + } + } + } + + return false; +} + fn allocate_block_insts<'a, F: Function>( state: &mut FastAllocState<'a, F>, const_state: &ReadOnlyData, block: Block, ) -> Result<(), RegAllocError> { + let block_last_inst_idx = state.func.block_insns(block).last().index(); for inst in state.func.block_insns(block).iter() { let edit_start_idx = state.edits.len(); let clobbers = state.func.inst_clobbers(inst); @@ -537,6 +585,24 @@ fn allocate_block_insts<'a, F: Function>( // 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 slot = state.create_stack_slot(reg.class()); + state.vregs[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); regs_allocated.add(reg); state.allocs[alloc_idx + i] = Allocation::reg(reg); @@ -562,7 +628,8 @@ fn allocate_block_insts<'a, F: Function>( return Err(RegAllocError::TooManyLiveRegs); } - if state.pregs[reg.index()].vreg.is_some() { + // TODO: make this proper + if regs_allocated.contains(reg) { // if the reg was allocated by another early use/write or late use // OR it is allocated and we have a late use we cannot do a correct allocation if op.pos() == OperandPos::Late || !late_write_regs.contains(reg) { @@ -571,7 +638,48 @@ fn allocate_block_insts<'a, F: Function>( } } - state.move_to_preg(vreg, reg, ProgPoint::before(inst)); + // are we already in the correct reg? + if let Some(cur_preg) = state.vregs[vreg.vreg()].preg { + if cur_preg == reg { + trace!("{} already in target {}", vreg, 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 slot = state.create_stack_slot(reg.class()); + state.vregs[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)), + }, + )); + } + } + + if let Some(cur_preg) = state.vregs[vreg.vreg()].preg { + // Move from preg to preg + state.edits.push(( + ProgPoint::before(inst), + Edit::Move { + from: Allocation::reg(cur_preg), + to: Allocation::reg(reg), + }, + )); + state.pregs[cur_preg.index()].vreg = None; + state.vregs[vreg.vreg()].preg = Some(reg); + state.pregs[reg.index()].vreg = Some(vreg); + } else { + state.move_to_preg(vreg, reg, ProgPoint::before(inst)); + } + state.allocs[alloc_idx + i] = Allocation::reg(reg); if op.pos() == OperandPos::Late { if clobbers.contains(reg) { @@ -605,6 +713,24 @@ fn allocate_block_insts<'a, F: Function>( late_write_disallow_regs.add(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 slot = state.create_stack_slot(reg.class()); + state.vregs[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.vregs[vreg.vreg()].def_block = Some(block); state.allocs[alloc_idx + i] = Allocation::reg(reg); state.assign_preg(reg, vreg); @@ -641,7 +767,30 @@ fn allocate_block_insts<'a, F: Function>( // Move from pseudoreg to tmp_reg and then to stack let tmp_reg = pregs.into_iter().next().unwrap(); - if state.pregs[tmp_reg.index()].vreg.is_some() { + 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 slot = state.create_stack_slot(reg.class()); + state.vregs[vreg].slot_idx = Some(slot); + state.edits.push(( + ProgPoint::before(inst), + Edit::Move { + from: Allocation::reg(tmp_reg), + to: Allocation::stack(SpillSlot::new( + slot as usize, + )), + }, + )); + } state.clear_preg(tmp_reg); } @@ -716,6 +865,10 @@ fn allocate_block_insts<'a, F: Function>( continue; } + if state.pregs[reg.index()].vreg.is_some() { + continue; + } + // reg should not contain anything debug_assert!(state.pregs[reg.index()].vreg.is_none()); @@ -746,11 +899,74 @@ fn allocate_block_insts<'a, F: Function>( continue; } - trace!("Ran out of registers for operand {}", i); + trace!("No free register found for operand {}", i); // No register available - // TODO: try to evict vreg that does not need to be in a preg - panic!("Out of registers: {:?}", regs_allocated); - return Err(RegAllocError::TooManyLiveRegs); + // TODO: first evict pregs that already have a stack slot even if they are used earlier? + let mut evict_candidate = None; + 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 { + continue; + } + + if let Some((_, pos)) = &evict_candidate { + if *pos < next_use { + evict_candidate = Some((reg, next_use)); + } + } else { + evict_candidate = Some((reg, next_use)); + } + } + + 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)), + }, + )); + } + state.clear_preg(reg); + + state.allocs[alloc_idx + i] = Allocation::reg(reg); + regs_allocated.add(reg); + 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, reg, ProgPoint::before(inst)); + } else { + // early def + state.vregs[vreg.vreg()].def_block = Some(block); + state.assign_preg(reg, vreg); + state.alloc_stack_slot(vreg); + state.move_to_stack(reg, vreg, ProgPoint::after(inst)); + } + + trace!("Chose {} for operand {}", reg, i); + } else { + panic!("Out of registers: {:?}", regs_allocated); + return Err(RegAllocError::TooManyLiveRegs); + } } OperandConstraint::Reuse(_) => { panic!("Illegal register constraint reuse for early def or use"); @@ -783,6 +999,9 @@ 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; + } // reg should not contain anything regs_allocated.add(reg); @@ -801,10 +1020,62 @@ fn allocate_block_insts<'a, F: Function>( continue; } - // No register available - // TODO: try to evict vreg that does not need to be in a preg - panic!("out of registers"); - return Err(RegAllocError::TooManyLiveRegs); + trace!("No free register found for {}", vreg); + + // TODO: first evict pregs that already have a stack slot even if they are used earlier? + let mut evict_candidate = None; + 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; + } + + if let Some((_, pos)) = &evict_candidate { + if *pos < next_use { + evict_candidate = Some((reg, next_use)); + } + } else { + evict_candidate = Some((reg, next_use)); + } + } + + 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)), + }, + )); + } + state.clear_preg(reg); + + regs_allocated.add(reg); + state.allocs[alloc_idx + i] = Allocation::reg(reg); + + state.assign_preg(reg, vreg); + state.alloc_stack_slot(vreg); + state.move_to_stack(reg, vreg, ProgPoint::after(inst)); + trace!("Chose {} for operand {}", reg, i); + } else { + panic!("out of registers"); + return Err(RegAllocError::TooManyLiveRegs); + } } OperandConstraint::Reuse(idx) => { debug_assert!(state.allocs[alloc_idx + idx].is_reg()); @@ -812,10 +1083,28 @@ fn allocate_block_insts<'a, F: Function>( debug_assert!(regs_allocated.contains(preg)); state.allocs[alloc_idx + i] = Allocation::reg(preg); + // 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 slot = state.create_stack_slot(preg.class()); + state.vregs[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.assign_preg(preg, vreg); - state.alloc_stack_slot(vreg); - state.move_to_stack(preg, vreg, ProgPoint::after(inst)); + //state.alloc_stack_slot(vreg); + //state.move_to_stack(preg, vreg, ProgPoint::after(inst)); } _ => { debug_assert!(!state.allocs[alloc_idx + i].is_none()); @@ -823,10 +1112,24 @@ fn allocate_block_insts<'a, F: Function>( } } - // clear out all allocated regs - for reg in regs_allocated { - trace!("Clearing {}", reg); - state.clear_preg(reg); + 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 @@ -847,6 +1150,33 @@ fn allocate_block_insts<'a, F: Function>( } } } + + state.cur_inst_pos += 1; + } + + // Move all liveout 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) => { + let idx = vreg.vreg(); + if 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(( + ProgPoint::before(Inst::new(block_last_inst_idx)), + Edit::Move { + from: Allocation::reg(preg), + to: Allocation::stack(SpillSlot::new(slot as usize)), + }, + )); + state.vregs[idx].slot_idx = Some(slot); + } + + state.vregs[idx].preg = None; + state.pregs[i].vreg = None; + } + } } Ok(())