From 993074a974369b9830bce132461e829822009a1e Mon Sep 17 00:00:00 2001 From: T0b1 Date: Thu, 13 Apr 2023 03:38:59 +0200 Subject: [PATCH] something --- src/ion/fast_alloc.rs | 620 ++++++++++++++++++++++++++++++++++++------ 1 file changed, 542 insertions(+), 78 deletions(-) diff --git a/src/ion/fast_alloc.rs b/src/ion/fast_alloc.rs index 2b8d95a..e3c2dc0 100644 --- a/src/ion/fast_alloc.rs +++ b/src/ion/fast_alloc.rs @@ -1,13 +1,16 @@ use alloc::vec; use alloc::vec::Vec; use smallvec::{smallvec, SmallVec}; -use std::convert::TryFrom; +use std::{convert::TryFrom, println}; use crate::{ cfg::CFGInfo, Allocation, Block, Edit, Function, Inst, MachineEnv, Operand, OperandConstraint, - OperandKind, OperandPos, Output, PReg, ProgPoint, RegAllocError, RegClass, SpillSlot, VReg, + OperandKind, OperandPos, Output, PReg, PRegSet, ProgPoint, RegAllocError, RegClass, SpillSlot, + VReg, }; +use super::Stats; + #[derive(Default, Clone, Copy)] struct VRegData { pub preg: Option, @@ -26,15 +29,98 @@ struct BlockData { pub allocated: bool, } +struct ReadOnlyData { + pub preorder: Vec, + pub reg_order_int: Vec, + pub reg_order_float: Vec, +} + +impl ReadOnlyData { + pub fn init(func: &F, mach_env: &MachineEnv) -> Self { + let reg_order_int = { + let class = RegClass::Int as usize; + let amount = mach_env.preferred_regs_by_class[class].len() + + mach_env.non_preferred_regs_by_class[class].len(); + let mut reg_order = Vec::with_capacity(amount); + reg_order.extend_from_slice(&mach_env.preferred_regs_by_class[class]); + reg_order.extend_from_slice(&mach_env.non_preferred_regs_by_class[class]); + reg_order + }; + + let reg_order_float = { + let class = RegClass::Float as usize; + let amount = mach_env.preferred_regs_by_class[class].len() + + mach_env.non_preferred_regs_by_class[class].len(); + let mut reg_order = Vec::with_capacity(amount); + reg_order.extend_from_slice(&mach_env.preferred_regs_by_class[class]); + reg_order.extend_from_slice(&mach_env.non_preferred_regs_by_class[class]); + reg_order + }; + + Self { + reg_order_int, + reg_order_float, + preorder: Self::calc_preorder(func), + } + } + + pub fn reg_order(&self, class: RegClass) -> &[PReg] { + match class { + RegClass::Int => &self.reg_order_int, + RegClass::Float => &self.reg_order_int, + } + } + + fn calc_preorder(func: &F) -> Vec { + let entry = func.entry_block(); + let mut ret = vec![entry]; + + struct State<'a> { + block: Block, + succs: &'a [Block], + next_succ: usize, + } + let mut stack: SmallVec<[State; 64]> = smallvec![]; + + stack.push(State { + block: entry, + succs: func.block_succs(entry), + next_succ: 0, + }); + + while let Some(ref mut state) = stack.last_mut() { + if state.next_succ >= state.succs.len() { + stack.pop(); + continue; + } + + let block = state.succs[state.next_succ]; + let succs = func.block_succs(block); + ret.push(block); + + state.next_succ += 1; + if state.next_succ >= state.succs.len() { + stack.pop(); + } + + if !succs.is_empty() { + stack.push(State { + block, + succs: func.block_succs(block), + next_succ: 0, + }); + } + } + + ret + } +} + struct FastAllocState<'a, F: Function> { pub vregs: Vec, pub pregs: Vec, pub blocks: Vec, - pub preorder: Vec, - pub reg_order_int: Vec, - pub reg_order_float: Vec, - pub cur_stack_slot_idx: u32, pub reftype_vregs_in_pregs_count: u32, @@ -46,6 +132,8 @@ struct FastAllocState<'a, F: Function> { pub edits: Vec<(ProgPoint, Edit)>, pub safepoint_slots: Vec<(ProgPoint, Allocation)>, + pub reftype_vregs: &'a [VReg], + pub func: &'a F, pub mach_env: &'a MachineEnv, pub cfg: &'a CFGInfo, @@ -53,11 +141,13 @@ struct FastAllocState<'a, F: Function> { impl<'a, F: Function> FastAllocState<'a, F> { pub fn init(func: &'a F, mach_env: &'a MachineEnv, cfg: &'a CFGInfo) -> Self { + let reftype_vregs = func.reftype_vregs(); + let vregs = { let mut vregs = Vec::with_capacity(func.num_vregs()); vregs.resize(func.num_vregs(), VRegData::default()); - for vreg in func.reftype_vregs() { + for vreg in reftype_vregs { vregs[vreg.vreg()].reftype = true; } @@ -80,26 +170,6 @@ impl<'a, F: Function> FastAllocState<'a, F> { blocks }; - let reg_order_int = { - let class = RegClass::Int as usize; - let amount = mach_env.preferred_regs_by_class[class].len() - + mach_env.non_preferred_regs_by_class[class].len(); - let mut reg_order = Vec::with_capacity(amount); - reg_order.extend_from_slice(&mach_env.preferred_regs_by_class[class]); - reg_order.extend_from_slice(&mach_env.non_preferred_regs_by_class[class]); - reg_order - }; - - let reg_order_float = { - let class = RegClass::Float as usize; - let amount = mach_env.preferred_regs_by_class[class].len() - + mach_env.non_preferred_regs_by_class[class].len(); - let mut reg_order = Vec::with_capacity(amount); - reg_order.extend_from_slice(&mach_env.preferred_regs_by_class[class]); - reg_order.extend_from_slice(&mach_env.non_preferred_regs_by_class[class]); - reg_order - }; - let mut inst_alloc_offsets = Vec::with_capacity(func.num_insts()); inst_alloc_offsets.resize(func.num_insts(), 0); @@ -108,10 +178,6 @@ impl<'a, F: Function> FastAllocState<'a, F> { pregs, blocks, - preorder: Self::calc_preorder(func), - reg_order_int, - reg_order_float, - cur_stack_slot_idx: 0, reftype_vregs_in_pregs_count: 0, @@ -125,6 +191,8 @@ impl<'a, F: Function> FastAllocState<'a, F> { edits: Vec::new(), safepoint_slots: Vec::new(), + reftype_vregs, + func, mach_env, cfg, @@ -156,6 +224,17 @@ impl<'a, F: Function> FastAllocState<'a, F> { idx } + pub fn create_stack_slot(&mut self, class: RegClass) -> u32 { + let size = if class == RegClass::Int { + self.stack_slot_count_int + } else { + self.stack_slot_count_float + }; + let idx = self.cur_stack_slot_idx; + self.cur_stack_slot_idx += size as u32; + idx + } + 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]; @@ -216,11 +295,20 @@ impl<'a, F: Function> FastAllocState<'a, F> { )); } + pub fn assign_preg(&mut self, preg: PReg, vreg: VReg) { + self.pregs[preg.index()].vreg = Some(vreg.vreg() as u32); + self.vregs[vreg.vreg()].preg = Some(preg); + } + pub fn clear_preg(&mut self, preg: PReg) { - let pdata = &mut self.pregs[preg.index()]; + self.clear_preg_idx(preg.index()); + } + + 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]; - debug_assert_eq!(vdata.preg.unwrap(), preg); + debug_assert_eq!(vdata.preg.unwrap().index(), preg); vdata.preg = None; pdata.vreg = None; if vdata.reftype { @@ -242,48 +330,19 @@ impl<'a, F: Function> FastAllocState<'a, F> { } } - fn calc_preorder(func: &F) -> Vec { - let entry = func.entry_block(); - let mut ret = vec![entry]; - - struct State<'a> { - block: Block, - succs: &'a [Block], - next_succ: usize, - } - let mut stack: SmallVec<[State; 64]> = smallvec![]; - - stack.push(State { - block: entry, - succs: func.block_succs(entry), - next_succ: 0, - }); - - while let Some(ref mut state) = stack.last_mut() { - if state.next_succ >= state.succs.len() { - stack.pop(); - continue; - } - - let block = state.succs[state.next_succ]; - let succs = func.block_succs(block); - ret.push(block); - - state.next_succ += 1; - if state.next_succ >= state.succs.len() { - stack.pop(); - } - - if !succs.is_empty() { - stack.push(State { - block, - succs: func.block_succs(block), - next_succ: 0, - }); - } + pub fn clear_reftype_vregs(&mut self) { + if self.reftype_vregs_in_pregs_count == 0 { + return; } - ret + for i in 0..self.pregs.len() { + if let Some(vreg) = self.pregs[i].vreg.clone() { + let vreg = vreg as usize; + if self.vregs[vreg].reftype { + self.clear_preg_idx(i); + } + } + } } } @@ -297,15 +356,420 @@ pub fn run( } let mut state = FastAllocState::init(func, mach_env, &cfg); + let const_state = ReadOnlyData::init(func, mach_env); - todo!("") + let len = const_state.preorder.len(); + for i in 0..len { + let block = const_state.preorder[i]; + // when handling branches later, we already have an input mapping for the block params + state.blocks[block.index()].allocated = true; + + allocate_block_insts(&mut state, &const_state, block)?; + handle_out_block_params(&mut state, &const_state, block)?; + } + + Ok(Output { + num_spillslots: state.cur_stack_slot_idx as usize, + edits: state.edits, + allocs: state.allocs, + inst_alloc_offsets: state.inst_alloc_offsets, + safepoint_slots: state.safepoint_slots, + debug_locations: Vec::new(), + stats: Stats::default(), + }) } -fn setup_entry_params<'a, F: Function>( +fn allocate_block_insts<'a, F: Function>( state: &mut FastAllocState<'a, F>, + const_state: &ReadOnlyData, + block: Block, ) -> Result<(), RegAllocError> { - // we need to set the vreg location for the initial block parameters and copy them to the stack - let entry = state.func.entry_block(); + for inst in state.func.block_insns(block).iter() { + let clobbers = state.func.inst_clobbers(inst); + let operands = state.func.inst_operands(inst); + let req_refs_on_stack = state.func.requires_refs_on_stack(inst); - todo!("") + let alloc_idx = state.allocs.len(); + state.inst_alloc_offsets[inst.index()] = alloc_idx as u32; + state + .allocs + .resize(alloc_idx + operands.len(), Allocation::none()); + + // keep track of which pregs where allocated so we can clear them later on + // TODO: wouldnt need this if we look up the inst a vreg was allocated at + let mut regs_allocated: SmallVec<[PReg; 8]> = smallvec![]; + + // keep track of which pregs hold late uses/early writes and so are unelligible + // as destinations for late writes + let mut late_write_disallow_regs = PRegSet::empty(); + // we need to keep track of late defs allocated during the fixed reg stage + // as they may not overlap with late uses and there is no order guarantee for inst_operands + let mut late_write_regs = PRegSet::empty(); + + if req_refs_on_stack { + state.clear_reftype_vregs(); + let pos = ProgPoint::before(inst); + + for vreg in state.reftype_vregs { + let data = &state.vregs[vreg.vreg()]; + if let Some(slot) = data.slot_idx { + state + .safepoint_slots + .push((pos, Allocation::stack(SpillSlot::new(slot as usize)))); + } + } + } + + // we allocate fixed defs/uses and stack allocations first + for (i, op) in operands.iter().enumerate() { + let vreg = op.vreg(); + + match op.constraint() { + OperandConstraint::FixedReg(reg) => { + match op.kind() { + OperandKind::Use => { + if req_refs_on_stack && state.vregs[vreg.vreg()].reftype { + panic!("reftype has fixed use when its required to be on stack"); + return Err(RegAllocError::TooManyLiveRegs); + } + + if state.pregs[reg.index()].vreg.is_some() { + // 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) { + panic!("fixed reg late use would overwrite another fixed reg use/early write"); + return Err(RegAllocError::TooManyLiveRegs); + } + } + + 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) { + panic!("fixed late use would be clobbered"); + return Err(RegAllocError::TooManyLiveRegs); + } + + // late uses cannot share a register with late defs + late_write_disallow_regs.add(reg); + } + regs_allocated.push(reg); + } + OperandKind::Def => { + if op.pos() == OperandPos::Late { + if late_write_disallow_regs.contains(reg) { + panic!("fixed late def would overwrite late use/early def"); + return Err(RegAllocError::TooManyLiveRegs); + } + late_write_regs.add(reg); + } else { + if state.pregs[reg.index()].vreg.is_some() || clobbers.contains(reg) + { + // early defs cannot share a register with anything and cannot be clobbered + panic!("early def shares reg or is clobbered"); + return Err(RegAllocError::TooManyLiveRegs); + } + // early defs cannot share a register with late defs + late_write_disallow_regs.add(reg); + } + + state.allocs[alloc_idx + i] = Allocation::reg(reg); + state.assign_preg(reg, vreg); + + // some pseudoops use the pseudo stack pregs as defs + if state.pregs[reg.index()].stack_pseudo { + // find preg to use as a tmp register + let mut pregs = PRegSet::empty(); + + for reg in const_state.reg_order(vreg.class()) { + if state.pregs[reg.index()].vreg.is_some() { + continue; + } + pregs.add(*reg); + } + + for op in operands { + match op.constraint() { + OperandConstraint::FixedReg(reg) => { + if op.kind() == OperandKind::Use + && op.pos() == OperandPos::Early + { + continue; + } + pregs.remove(reg); + } + _ => {} + } + } + + if pregs == PRegSet::empty() { + panic!("No way to solve pseudo-stack preg"); + } + + // 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() { + state.clear_preg(tmp_reg); + } + + state.edits.push(( + ProgPoint::after(inst), + Edit::Move { + from: Allocation::reg(reg), + to: Allocation::reg(tmp_reg), + }, + )); + if state.pregs[reg.index()].vreg.is_some() { + state.clear_preg(reg); + } + state.assign_preg(tmp_reg, vreg); + println!("1"); + state.move_to_stack(tmp_reg, vreg, ProgPoint::after(inst)); + regs_allocated.push(tmp_reg); + } else { + println!("2"); + state.alloc_stack_slot(vreg); + state.move_to_stack(reg, vreg, ProgPoint::after(inst)); + regs_allocated.push(reg); + } + } + } + } + OperandConstraint::Stack | OperandConstraint::Any => { + // we allocate Any on the stack for now + match op.kind() { + OperandKind::Use => { + if let Some(slot) = &state.vregs[vreg.vreg()].slot_idx { + state.allocs[alloc_idx + i] = + Allocation::stack(SpillSlot::new(*slot as usize)); + } else { + return Err(RegAllocError::SSA(vreg, inst)); + } + } + OperandKind::Def => { + state.allocs[alloc_idx + i] = Allocation::stack(SpillSlot::new( + state.alloc_stack_slot(vreg) as usize, + )); + } + } + } + _ => continue, + } + } + + // alloc non-fixed uses and early defs in registers + for (i, op) in operands.iter().enumerate() { + if op.kind() == OperandKind::Def && op.pos() == OperandPos::Late { + continue; + } + + let vreg = op.vreg(); + + match op.constraint() { + OperandConstraint::Reg => { + // find first non-allocated register + let reg_order = const_state.reg_order(op.class()); + let mut allocated = false; + for ® in reg_order { + if regs_allocated.contains(®) { + continue; + } + + // reg should not contain anything + debug_assert!(state.pregs[reg.index()].vreg.is_none()); + + state.allocs[alloc_idx + i] = Allocation::reg(reg); + regs_allocated.push(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.assign_preg(reg, vreg); + state.alloc_stack_slot(vreg); + println!("3"); + state.move_to_stack(reg, vreg, ProgPoint::after(inst)); + } + allocated = true; + break; + } + + if allocated { + continue; + } + + // 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); + } + OperandConstraint::Reuse(_) => { + panic!("Illegal register constraint reuse for early def or use"); + } + _ => {} + } + } + + // alloc non-fixed late defs and reuse + for (i, op) in operands.iter().enumerate() { + if op.kind() != OperandKind::Def || op.pos() != OperandPos::Late { + continue; + } + + let vreg = op.vreg(); + match op.constraint() { + OperandConstraint::Reg => { + // find first non-allocated register + let reg_order = const_state.reg_order(op.class()); + let mut allocated = false; + for ® in reg_order { + if regs_allocated.contains(®) || late_write_disallow_regs.contains(reg) { + continue; + } + + // reg should not contain anything + regs_allocated.push(reg); + state.allocs[alloc_idx + i] = Allocation::reg(reg); + + state.clear_preg(reg); + state.assign_preg(reg, vreg); + state.alloc_stack_slot(vreg); + println!("4"); + state.move_to_stack(reg, vreg, ProgPoint::after(inst)); + allocated = true; + break; + } + + if allocated { + 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); + } + OperandConstraint::Reuse(idx) => { + debug_assert!(state.allocs[alloc_idx + idx].is_reg()); + let preg = state.allocs[alloc_idx + idx].as_reg().unwrap(); + debug_assert!(regs_allocated.contains(&preg)); + state.allocs[alloc_idx + i] = Allocation::reg(preg); + + state.clear_preg(preg); + state.assign_preg(preg, vreg); + state.alloc_stack_slot(vreg); + println!("5"); + state.move_to_stack(preg, vreg, ProgPoint::after(inst)); + } + _ => { + debug_assert!(!state.allocs[alloc_idx + i].is_none()); + } + } + } + + // clear out all allocated regs + for reg in regs_allocated { + state.clear_preg(reg); + } + } + + Ok(()) +} + +fn handle_out_block_params<'a, F: Function>( + state: &mut FastAllocState<'a, F>, + const_state: &ReadOnlyData, + block: Block, +) -> Result<(), RegAllocError> { + let last_inst = state.func.block_insns(block).last(); + if !state.func.is_branch(last_inst) { + return Ok(()); + } + + let succs = state.func.block_succs(block); + if succs.len() == 1 && state.blocks[succs[0].index()].allocated { + let succ = succs[0]; + // move values to the already allocated places + let in_params = state.func.block_params(succ); + let out_params = state.func.branch_blockparams(block, last_inst, 0); + debug_assert_eq!(in_params.len(), out_params.len()); + + for i in 0..in_params.len() { + let in_vreg = in_params[i]; + let out_vreg = out_params[i]; + + debug_assert!(state.vregs[in_vreg.vreg()].slot_idx.is_some()); + debug_assert!(state.vregs[out_vreg.vreg()].slot_idx.is_some()); + + let tmp_reg = const_state.reg_order(out_vreg.class())[0]; + let out_slot = state.vregs[out_vreg.vreg()].slot_idx.unwrap(); + let in_slot = state.vregs[in_vreg.vreg()].slot_idx.unwrap(); + + state.edits.push(( + ProgPoint::before(last_inst), + Edit::Move { + from: Allocation::stack(SpillSlot::new(out_slot as usize)), + to: Allocation::reg(tmp_reg), + }, + )); + state.edits.push(( + ProgPoint::before(last_inst), + Edit::Move { + from: Allocation::reg(tmp_reg), + to: Allocation::stack(SpillSlot::new(in_slot as usize)), + }, + )); + } + } else { + // set incoming block params of successor to the current stack slot + for (i, &succ) in state.func.block_succs(block).iter().enumerate() { + if state.blocks[succ.index()].allocated { + return Err(RegAllocError::CritEdge(block, succ)); + } + + let in_params = state.func.block_params(succ); + let out_params = state.func.branch_blockparams(block, last_inst, i); + debug_assert_eq!(in_params.len(), out_params.len()); + + let mut vregs_passed = SmallVec::<[VReg; 4]>::new(); + for i in 0..in_params.len() { + let out_vreg = out_params[i]; + let in_vreg = in_params[i]; + debug_assert!(state.vregs[out_vreg.vreg()].slot_idx.is_some()); + let out_slot_idx = state.vregs[out_vreg.vreg()].slot_idx.unwrap(); + + if !vregs_passed.contains(&out_vreg) { + state.vregs[in_vreg.vreg()].slot_idx = Some(out_slot_idx); + vregs_passed.push(out_vreg); + } else { + // need to duplicate to avoid aliasing + // TODO: this creates multiple duplications for multiple blocks, can be avoided + let tmp_reg = const_state.reg_order(out_vreg.class())[0]; + let slot = state.create_stack_slot(out_vreg.class()); + state.edits.push(( + ProgPoint::before(last_inst), + Edit::Move { + from: Allocation::stack(SpillSlot::new(out_slot_idx as usize)), + to: Allocation::reg(tmp_reg), + }, + )); + state.edits.push(( + ProgPoint::before(last_inst), + Edit::Move { + from: Allocation::reg(tmp_reg), + to: Allocation::stack(SpillSlot::new(slot as usize)), + }, + )); + state.vregs[in_vreg.vreg()].slot_idx = Some(slot); + } + } + } + } + + Ok(()) }