use crate::moves::{MoveAndScratchResolver, ParallelMoves};
use crate::{cfg::CFGInfo, ion::Stats, Allocation, RegAllocError};
use crate::{ssa::validate_ssa, Edit, Function, MachineEnv, Output, ProgPoint};
use crate::{
AllocationKind, Block, FxHashMap, Inst, InstPosition, Operand, OperandConstraint, OperandKind,
OperandPos, PReg, PRegSet, RegClass, SpillSlot, VReg,
};
use alloc::vec::Vec;
use core::convert::TryInto;
use core::iter::FromIterator;
use core::ops::{Index, IndexMut};
mod iter;
mod lru;
mod vregset;
use iter::*;
use lru::*;
use vregset::VRegSet;
#[cfg(test)]
mod tests;
#[derive(Debug)]
struct Allocs {
allocs: Vec<Allocation>,
inst_alloc_offsets: Vec<u32>,
}
impl Allocs {
fn new<F: Function>(func: &F) -> (Self, u32) {
let mut allocs = Vec::new();
let mut inst_alloc_offsets = Vec::with_capacity(func.num_insts());
let mut max_operand_len = 0;
let mut no_of_operands = 0;
for inst in 0..func.num_insts() {
let operands_len = func.inst_operands(Inst::new(inst)).len() as u32;
max_operand_len = max_operand_len.max(operands_len);
inst_alloc_offsets.push(no_of_operands as u32);
no_of_operands += operands_len;
}
allocs.resize(no_of_operands as usize, Allocation::none());
(
Self {
allocs,
inst_alloc_offsets,
},
max_operand_len,
)
}
}
impl Index<(usize, usize)> for Allocs {
type Output = Allocation;
fn index(&self, idx: (usize, usize)) -> &Allocation {
&self.allocs[self.inst_alloc_offsets[idx.0] as usize + idx.1]
}
}
impl IndexMut<(usize, usize)> for Allocs {
fn index_mut(&mut self, idx: (usize, usize)) -> &mut Allocation {
&mut self.allocs[self.inst_alloc_offsets[idx.0] as usize + idx.1]
}
}
#[derive(Debug)]
struct Stack<'a, F: Function> {
num_spillslots: u32,
func: &'a F,
}
impl<'a, F: Function> Stack<'a, F> {
fn new(func: &'a F) -> Self {
Self {
num_spillslots: 0,
func,
}
}
fn allocstack(&mut self, class: RegClass) -> SpillSlot {
trace!("Allocating a spillslot for class {class:?}");
let size: u32 = self.func.spillslot_size(class).try_into().unwrap();
let mut offset = self.num_spillslots;
debug_assert!(size.is_power_of_two());
offset = (offset + size - 1) & !(size - 1);
let slot = if self.func.multi_spillslot_named_by_last_slot() {
offset + size - 1
} else {
offset
};
offset += size;
self.num_spillslots = offset;
trace!("Allocated slot: {slot}");
SpillSlot::new(slot as usize)
}
}
#[derive(Debug)]
struct Edits {
edits: Vec<(ProgPoint, Edit)>,
fixed_stack_slots: PRegSet,
scratch_regs: PartedByRegClass<Option<PReg>>,
dedicated_scratch_regs: PartedByRegClass<Option<PReg>>,
}
impl Edits {
fn new(
fixed_stack_slots: PRegSet,
num_insts: usize,
dedicated_scratch_regs: PartedByRegClass<Option<PReg>>,
) -> Self {
let edits_len_guess = num_insts;
Self {
edits: Vec::with_capacity(edits_len_guess),
fixed_stack_slots,
scratch_regs: dedicated_scratch_regs.clone(),
dedicated_scratch_regs,
}
}
}
impl Edits {
fn is_stack(&self, alloc: Allocation) -> bool {
alloc.is_stack()
|| (alloc.is_reg() && self.fixed_stack_slots.contains(alloc.as_reg().unwrap()))
}
fn add_move(
&mut self,
inst: Inst,
from: Allocation,
to: Allocation,
class: RegClass,
pos: InstPosition,
) {
trace!(
"Recording edit: {:?}",
(ProgPoint::new(inst, pos), Edit::Move { from, to }, class)
);
if self.is_stack(from) && self.is_stack(to) {
trace!("Edit is stack-to-stack. Generating two edits with a scratch register");
let scratch_reg = self.scratch_regs[class].unwrap();
let scratch_alloc = Allocation::reg(scratch_reg);
trace!("Move 1: {scratch_alloc:?} to {to:?}");
self.edits.push((
ProgPoint::new(inst, pos),
Edit::Move {
from: scratch_alloc,
to,
},
));
trace!("Move 2: {from:?} to {scratch_alloc:?}");
self.edits.push((
ProgPoint::new(inst, pos),
Edit::Move {
from,
to: scratch_alloc,
},
));
} else {
self.edits
.push((ProgPoint::new(inst, pos), Edit::Move { from, to }));
}
}
}
#[derive(Debug, Clone)]
struct PartedByOperandPos<T> {
items: [T; 2],
}
impl<T> Index<OperandPos> for PartedByOperandPos<T> {
type Output = T;
fn index(&self, index: OperandPos) -> &Self::Output {
&self.items[index as usize]
}
}
impl<T> IndexMut<OperandPos> for PartedByOperandPos<T> {
fn index_mut(&mut self, index: OperandPos) -> &mut Self::Output {
&mut self.items[index as usize]
}
}
use core::fmt;
impl<T: fmt::Display> fmt::Display for PartedByOperandPos<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{{ early: {}, late: {} }}", self.items[0], self.items[1])
}
}
#[derive(Debug)]
pub struct Env<'a, F: Function> {
func: &'a F,
vreg_allocs: Vec<Allocation>,
vreg_spillslots: Vec<SpillSlot>,
live_vregs: VRegSet,
lrus: Lrus,
vreg_in_preg: Vec<VReg>,
reused_input_to_reuse_op: Vec<usize>,
available_pregs: PartedByOperandPos<PRegSet>,
init_available_pregs: PRegSet,
allocatable_regs: PRegSet,
stack: Stack<'a, F>,
preferred_victim: PartedByRegClass<PReg>,
vreg_to_live_inst_range: Vec<(ProgPoint, ProgPoint, Allocation)>,
fixed_stack_slots: PRegSet,
allocs: Allocs,
edits: Edits,
stats: Stats,
debug_locations: Vec<(u32, ProgPoint, ProgPoint, Allocation)>,
}
impl<'a, F: Function> Env<'a, F> {
fn new(func: &'a F, env: &'a MachineEnv) -> Self {
use alloc::vec;
let mut regs = [
env.preferred_regs_by_class[RegClass::Int as usize].clone(),
env.preferred_regs_by_class[RegClass::Float as usize].clone(),
env.preferred_regs_by_class[RegClass::Vector as usize].clone(),
];
regs[0].extend(
env.non_preferred_regs_by_class[RegClass::Int as usize]
.iter()
.cloned(),
);
regs[1].extend(
env.non_preferred_regs_by_class[RegClass::Float as usize]
.iter()
.cloned(),
);
regs[2].extend(
env.non_preferred_regs_by_class[RegClass::Vector as usize]
.iter()
.cloned(),
);
let allocatable_regs = PRegSet::from(env);
let init_available_pregs = {
let mut regs = allocatable_regs;
for preg in env.fixed_stack_slots.iter() {
regs.add(*preg);
}
regs
};
let dedicated_scratch_regs = PartedByRegClass {
items: [
env.scratch_by_class[0],
env.scratch_by_class[1],
env.scratch_by_class[2],
],
};
trace!("{:?}", env);
let (allocs, max_operand_len) = Allocs::new(func);
let fixed_stack_slots = PRegSet::from_iter(env.fixed_stack_slots.iter().cloned());
Self {
func,
allocatable_regs,
vreg_allocs: vec![Allocation::none(); func.num_vregs()],
vreg_spillslots: vec![SpillSlot::invalid(); func.num_vregs()],
live_vregs: VRegSet::with_capacity(func.num_vregs()),
lrus: Lrus::new(®s[0], ®s[1], ®s[2]),
vreg_in_preg: vec![VReg::invalid(); PReg::NUM_INDEX],
stack: Stack::new(func),
fixed_stack_slots,
vreg_to_live_inst_range: vec![
(
ProgPoint::invalid(),
ProgPoint::invalid(),
Allocation::none()
);
func.num_vregs()
],
preferred_victim: PartedByRegClass {
items: [
regs[0].last().cloned().unwrap_or(PReg::invalid()),
regs[1].last().cloned().unwrap_or(PReg::invalid()),
regs[2].last().cloned().unwrap_or(PReg::invalid()),
],
},
reused_input_to_reuse_op: vec![usize::MAX; max_operand_len as usize],
init_available_pregs,
available_pregs: PartedByOperandPos {
items: [init_available_pregs, init_available_pregs],
},
allocs,
edits: Edits::new(fixed_stack_slots, func.num_insts(), dedicated_scratch_regs),
stats: Stats::default(),
debug_locations: Vec::with_capacity(func.debug_value_labels().len()),
}
}
fn reset_available_pregs_and_scratch_regs(&mut self) {
trace!("Resetting the available pregs");
self.available_pregs = PartedByOperandPos {
items: [self.init_available_pregs, self.init_available_pregs],
};
self.edits.scratch_regs = self.edits.dedicated_scratch_regs.clone();
}
fn alloc_scratch_reg(&mut self, inst: Inst, class: RegClass) -> Result<(), RegAllocError> {
use OperandPos::{Early, Late};
let reg = self.get_scratch_reg(
inst,
class,
self.available_pregs[Late] & self.available_pregs[Early],
)?;
self.edits.scratch_regs[class] = Some(reg);
self.available_pregs[OperandPos::Early].remove(reg);
self.available_pregs[OperandPos::Late].remove(reg);
Ok(())
}
fn get_scratch_reg_for_reload(
&mut self,
inst: Inst,
class: RegClass,
avail_regs: PRegSet,
) -> Result<PReg, RegAllocError> {
let Some(preg) = self.lrus[class].last(avail_regs) else {
return Err(RegAllocError::TooManyLiveRegs);
};
if self.vreg_in_preg[preg.index()] != VReg::invalid() {
self.evict_vreg_in_preg_before_inst(inst, preg);
}
Ok(preg)
}
fn get_scratch_reg(
&mut self,
inst: Inst,
class: RegClass,
avail_regs: PRegSet,
) -> Result<PReg, RegAllocError> {
let Some(preg) = self.lrus[class].last(avail_regs) else {
return Err(RegAllocError::TooManyLiveRegs);
};
if self.vreg_in_preg[preg.index()] != VReg::invalid() {
self.evict_vreg_in_preg(inst, preg);
}
Ok(preg)
}
fn reserve_reg_for_fixed_operand(
&mut self,
op: Operand,
op_idx: usize,
preg: PReg,
) -> Result<(), RegAllocError> {
trace!("Reserving register {preg} for fixed operand {op}");
let early_avail_pregs = self.available_pregs[OperandPos::Early];
let late_avail_pregs = self.available_pregs[OperandPos::Late];
match (op.pos(), op.kind()) {
(OperandPos::Early, OperandKind::Use) => {
if op.as_fixed_nonallocatable().is_none() && !early_avail_pregs.contains(preg) {
return Err(RegAllocError::TooManyLiveRegs);
}
self.available_pregs[OperandPos::Early].remove(preg);
if self.reused_input_to_reuse_op[op_idx] != usize::MAX {
if op.as_fixed_nonallocatable().is_none() && !late_avail_pregs.contains(preg) {
return Err(RegAllocError::TooManyLiveRegs);
}
self.available_pregs[OperandPos::Late].remove(preg);
}
}
(OperandPos::Late, OperandKind::Def) => {
if op.as_fixed_nonallocatable().is_none() && !late_avail_pregs.contains(preg) {
return Err(RegAllocError::TooManyLiveRegs);
}
self.available_pregs[OperandPos::Late].remove(preg);
}
_ => {
if op.as_fixed_nonallocatable().is_none()
&& (!early_avail_pregs.contains(preg) || !late_avail_pregs.contains(preg))
{
return Err(RegAllocError::TooManyLiveRegs);
}
self.available_pregs[OperandPos::Early].remove(preg);
self.available_pregs[OperandPos::Late].remove(preg);
}
}
Ok(())
}
fn allocd_within_constraint(&self, op: Operand) -> bool {
let alloc = self.vreg_allocs[op.vreg().vreg()];
match op.constraint() {
OperandConstraint::Any => {
if let Some(preg) = alloc.as_reg() {
if !self.available_pregs[op.pos()].contains(preg) {
trace!("The vreg in {preg}: {}", self.vreg_in_preg[preg.index()]);
self.vreg_in_preg[preg.index()] == op.vreg()
} else {
true
}
} else {
!alloc.is_none()
}
}
OperandConstraint::Reg => {
if self.edits.is_stack(alloc) {
return false;
}
if let Some(preg) = alloc.as_reg() {
if !self.available_pregs[op.pos()].contains(preg) {
trace!("The vreg in {preg}: {}", self.vreg_in_preg[preg.index()]);
self.vreg_in_preg[preg.index()] == op.vreg()
} else {
true
}
} else {
false
}
}
OperandConstraint::FixedReg(preg) => alloc.is_reg() && alloc.as_reg().unwrap() == preg,
OperandConstraint::Reuse(_) => {
unreachable!()
}
}
}
fn base_evict_vreg_in_preg(&mut self, inst: Inst, preg: PReg, pos: InstPosition) {
trace!("Removing the vreg in preg {} for eviction", preg);
let evicted_vreg = self.vreg_in_preg[preg.index()];
trace!("The removed vreg: {}", evicted_vreg);
debug_assert_ne!(evicted_vreg, VReg::invalid());
if self.vreg_spillslots[evicted_vreg.vreg()].is_invalid() {
self.vreg_spillslots[evicted_vreg.vreg()] = self.stack.allocstack(evicted_vreg.class());
}
let slot = self.vreg_spillslots[evicted_vreg.vreg()];
self.vreg_allocs[evicted_vreg.vreg()] = Allocation::stack(slot);
trace!("Move reason: eviction");
self.edits.add_move(
inst,
self.vreg_allocs[evicted_vreg.vreg()],
Allocation::reg(preg),
evicted_vreg.class(),
pos,
);
}
fn evict_vreg_in_preg_before_inst(&mut self, inst: Inst, preg: PReg) {
self.base_evict_vreg_in_preg(inst, preg, InstPosition::Before)
}
fn evict_vreg_in_preg(&mut self, inst: Inst, preg: PReg) {
self.base_evict_vreg_in_preg(inst, preg, InstPosition::After)
}
fn freealloc(&mut self, vreg: VReg) {
trace!("Freeing vreg {}", vreg);
let alloc = self.vreg_allocs[vreg.vreg()];
match alloc.kind() {
AllocationKind::Reg => {
let preg = alloc.as_reg().unwrap();
self.vreg_in_preg[preg.index()] = VReg::invalid();
}
AllocationKind::Stack => (),
AllocationKind::None => unreachable!("Attempting to free an unallocated operand!"),
}
self.vreg_allocs[vreg.vreg()] = Allocation::none();
self.live_vregs.remove(vreg.vreg());
trace!(
"{} curr alloc is now {}",
vreg,
self.vreg_allocs[vreg.vreg()]
);
}
fn alloc_reg_for_operand(
&mut self,
inst: Inst,
op: Operand,
) -> Result<Allocation, RegAllocError> {
trace!("available regs: {}", self.available_pregs);
trace!("Int LRU: {:?}", self.lrus[RegClass::Int]);
trace!("Float LRU: {:?}", self.lrus[RegClass::Float]);
trace!("Vector LRU: {:?}", self.lrus[RegClass::Vector]);
trace!("");
let draw_from = match (op.pos(), op.kind()) {
(OperandPos::Late, OperandKind::Use)
| (OperandPos::Early, OperandKind::Def)
| (OperandPos::Late, OperandKind::Def)
if matches!(op.constraint(), OperandConstraint::Reuse(_)) =>
{
self.available_pregs[OperandPos::Late] & self.available_pregs[OperandPos::Early]
}
_ => self.available_pregs[op.pos()],
};
if draw_from.is_empty(op.class()) {
trace!("No registers available for {op}");
return Err(RegAllocError::TooManyLiveRegs);
}
let Some(preg) = self.lrus[op.class()].last(draw_from) else {
trace!(
"Failed to find an available {:?} register in the LRU for operand {op}",
op.class()
);
return Err(RegAllocError::TooManyLiveRegs);
};
if self.vreg_in_preg[preg.index()] != VReg::invalid() {
self.evict_vreg_in_preg(inst, preg);
}
trace!("The allocated register for vreg {}: {}", op.vreg(), preg);
self.lrus[op.class()].poke(preg);
self.available_pregs[op.pos()].remove(preg);
match (op.pos(), op.kind()) {
(OperandPos::Late, OperandKind::Use) => {
self.available_pregs[OperandPos::Early].remove(preg);
}
(OperandPos::Early, OperandKind::Def) => {
self.available_pregs[OperandPos::Late].remove(preg);
}
(OperandPos::Late, OperandKind::Def)
if matches!(op.constraint(), OperandConstraint::Reuse(_)) =>
{
self.available_pregs[OperandPos::Early].remove(preg);
}
_ => (),
};
Ok(Allocation::reg(preg))
}
fn alloc_operand(
&mut self,
inst: Inst,
op: Operand,
op_idx: usize,
) -> Result<Allocation, RegAllocError> {
let new_alloc = match op.constraint() {
OperandConstraint::Any => self.alloc_reg_for_operand(inst, op)?,
OperandConstraint::Reg => self.alloc_reg_for_operand(inst, op)?,
OperandConstraint::FixedReg(preg) => {
trace!("The fixed preg: {} for operand {}", preg, op);
Allocation::reg(preg)
}
OperandConstraint::Reuse(_) => {
unreachable!();
}
};
self.allocs[(inst.index(), op_idx)] = new_alloc;
Ok(new_alloc)
}
fn process_operand_allocation(
&mut self,
inst: Inst,
op: Operand,
op_idx: usize,
) -> Result<(), RegAllocError> {
if let Some(preg) = op.as_fixed_nonallocatable() {
self.allocs[(inst.index(), op_idx)] = Allocation::reg(preg);
trace!(
"Allocation for instruction {:?} and operand {}: {}",
inst,
op,
self.allocs[(inst.index(), op_idx)]
);
return Ok(());
}
if !self.allocd_within_constraint(op) {
trace!("{op} isn't allocated within constraints.");
let curr_alloc = self.vreg_allocs[op.vreg().vreg()];
let new_alloc = self.alloc_operand(inst, op, op_idx)?;
if curr_alloc.is_none() {
self.live_vregs.insert(op.vreg());
self.vreg_to_live_inst_range[op.vreg().vreg()].1 = match (op.pos(), op.kind()) {
(OperandPos::Late, OperandKind::Use) | (_, OperandKind::Def) => {
ProgPoint::before(Inst::new(inst.index() + 1))
}
(OperandPos::Early, OperandKind::Use) => {
ProgPoint::after(inst)
}
};
self.vreg_to_live_inst_range[op.vreg().vreg()].2 = new_alloc;
trace!("Setting vreg_allocs[{op}] to {new_alloc:?}");
self.vreg_allocs[op.vreg().vreg()] = new_alloc;
if let Some(preg) = new_alloc.as_reg() {
self.vreg_in_preg[preg.index()] = op.vreg();
}
}
else {
trace!("Move reason: Prev allocation doesn't meet constraints");
if self.edits.is_stack(new_alloc)
&& self.edits.is_stack(curr_alloc)
&& self.edits.scratch_regs[op.class()].is_none()
{
self.alloc_scratch_reg(inst, op.class())?;
}
if op.kind() == OperandKind::Def {
trace!("Adding edit from {new_alloc:?} to {curr_alloc:?} after inst {inst:?} for {op}");
self.edits.add_move(
inst,
new_alloc,
curr_alloc,
op.class(),
InstPosition::After,
);
}
if let Some(preg) = new_alloc.as_reg() {
self.vreg_in_preg[preg.index()] = VReg::invalid();
}
}
trace!(
"Allocation for instruction {:?} and operand {}: {}",
inst,
op,
self.allocs[(inst.index(), op_idx)]
);
} else {
trace!("{op} is already allocated within constraints");
self.allocs[(inst.index(), op_idx)] = self.vreg_allocs[op.vreg().vreg()];
if let Some(preg) = self.allocs[(inst.index(), op_idx)].as_reg() {
if self.allocatable_regs.contains(preg) {
self.lrus[preg.class()].poke(preg);
}
self.available_pregs[op.pos()].remove(preg);
match (op.pos(), op.kind()) {
(OperandPos::Late, OperandKind::Use) => {
self.available_pregs[OperandPos::Early].remove(preg);
}
(OperandPos::Early, OperandKind::Def) => {
self.available_pregs[OperandPos::Late].remove(preg);
}
_ => (),
};
}
trace!(
"Allocation for instruction {:?} and operand {}: {}",
inst,
op,
self.allocs[(inst.index(), op_idx)]
);
}
trace!(
"Late available regs: {}",
self.available_pregs[OperandPos::Late]
);
trace!(
"Early available regs: {}",
self.available_pregs[OperandPos::Early]
);
Ok(())
}
fn remove_clobbers_from_available_pregs(&mut self, clobbers: PRegSet) {
trace!("Removing clobbers {clobbers} from available reg sets");
let all_but_clobbers = clobbers.invert();
self.available_pregs[OperandPos::Late].intersect_from(all_but_clobbers);
self.available_pregs[OperandPos::Early].intersect_from(all_but_clobbers);
}
fn process_branch(&mut self, block: Block, inst: Inst) -> Result<(), RegAllocError> {
use OperandPos::*;
trace!("Processing branch instruction {inst:?} in block {block:?}");
let mut int_parallel_moves = ParallelMoves::new();
let mut float_parallel_moves = ParallelMoves::new();
let mut vec_parallel_moves = ParallelMoves::new();
for (succ_idx, succ) in self.func.block_succs(block).iter().enumerate() {
for (pos, vreg) in self
.func
.branch_blockparams(block, inst, succ_idx)
.iter()
.enumerate()
{
let succ_params = self.func.block_params(*succ);
let succ_param_vreg = succ_params[pos];
if self.vreg_spillslots[succ_param_vreg.vreg()].is_invalid() {
self.vreg_spillslots[succ_param_vreg.vreg()] =
self.stack.allocstack(succ_param_vreg.class());
}
if self.vreg_spillslots[vreg.vreg()].is_invalid() {
self.vreg_spillslots[vreg.vreg()] = self.stack.allocstack(vreg.class());
}
let vreg_spill = Allocation::stack(self.vreg_spillslots[vreg.vreg()]);
let curr_alloc = self.vreg_allocs[vreg.vreg()];
if curr_alloc.is_none() {
self.live_vregs.insert(*vreg);
self.vreg_to_live_inst_range[vreg.vreg()].1 = ProgPoint::before(inst);
} else if curr_alloc != vreg_spill {
if self.edits.is_stack(curr_alloc)
&& self.edits.scratch_regs[vreg.class()].is_none()
{
let reg = self.get_scratch_reg_for_reload(
inst,
vreg.class(),
self.available_pregs[Early] & self.available_pregs[Late],
)?;
self.edits.scratch_regs[vreg.class()] = Some(reg);
self.available_pregs[OperandPos::Early].remove(reg);
self.available_pregs[OperandPos::Late].remove(reg);
}
self.edits.add_move(
inst,
vreg_spill,
curr_alloc,
vreg.class(),
InstPosition::Before,
);
}
self.vreg_allocs[vreg.vreg()] = vreg_spill;
let parallel_moves = match vreg.class() {
RegClass::Int => &mut int_parallel_moves,
RegClass::Float => &mut float_parallel_moves,
RegClass::Vector => &mut vec_parallel_moves,
};
let from = Allocation::stack(self.vreg_spillslots[vreg.vreg()]);
let to = Allocation::stack(self.vreg_spillslots[succ_param_vreg.vreg()]);
trace!("Recording parallel move from {from} to {to}");
parallel_moves.add(from, to, Some(*vreg));
}
}
let resolved_int = int_parallel_moves.resolve();
let resolved_float = float_parallel_moves.resolve();
let resolved_vec = vec_parallel_moves.resolve();
let mut scratch_regs = self.edits.scratch_regs.clone();
let mut num_spillslots = self.stack.num_spillslots;
let mut avail_regs = self.available_pregs[Early] & self.available_pregs[Late];
trace!("Resolving parallel moves");
for (resolved, class) in [
(resolved_int, RegClass::Int),
(resolved_float, RegClass::Float),
(resolved_vec, RegClass::Vector),
] {
let scratch_resolver = MoveAndScratchResolver {
find_free_reg: || {
if let Some(reg) = scratch_regs[class] {
trace!("Retrieved reg {reg} for scratch resolver");
scratch_regs[class] = None;
Some(Allocation::reg(reg))
} else {
let Some(preg) = self.lrus[class].last(avail_regs) else {
trace!("Couldn't find any reg for scratch resolver");
return None;
};
avail_regs.remove(preg);
trace!("Retrieved reg {preg} for scratch resolver");
Some(Allocation::reg(preg))
}
},
get_stackslot: || {
let size: u32 = self.func.spillslot_size(class).try_into().unwrap();
let mut offset = num_spillslots;
debug_assert!(size.is_power_of_two());
offset = (offset + size - 1) & !(size - 1);
let slot = if self.func.multi_spillslot_named_by_last_slot() {
offset + size - 1
} else {
offset
};
offset += size;
num_spillslots = offset;
trace!("Retrieved slot {slot} for scratch resolver");
Allocation::stack(SpillSlot::new(slot as usize))
},
is_stack_alloc: |alloc| self.edits.is_stack(alloc),
borrowed_scratch_reg: self.preferred_victim[class],
};
let moves = scratch_resolver.compute(resolved);
trace!("Resolved {class:?} parallel moves");
for (from, to, _) in moves.into_iter().rev() {
self.edits
.edits
.push((ProgPoint::before(inst), Edit::Move { from, to }))
}
self.stack.num_spillslots = num_spillslots;
}
trace!("Completed processing branch");
Ok(())
}
fn alloc_inst(&mut self, block: Block, inst: Inst) -> Result<(), RegAllocError> {
trace!("Allocating instruction {:?}", inst);
self.reset_available_pregs_and_scratch_regs();
let operands = Operands::new(self.func.inst_operands(inst));
let clobbers = self.func.inst_clobbers(inst);
for (op_idx, op) in operands.reuse() {
trace!("Initializing reused_input_to_reuse_op for {op}");
let OperandConstraint::Reuse(reused_idx) = op.constraint() else {
unreachable!()
};
self.reused_input_to_reuse_op[reused_idx] = op_idx;
}
for (op_idx, op) in operands.fixed() {
let OperandConstraint::FixedReg(preg) = op.constraint() else {
unreachable!();
};
self.reserve_reg_for_fixed_operand(op, op_idx, preg)?;
if self.allocatable_regs.contains(preg) {
self.lrus[preg.class()].poke(preg);
}
}
for (_, op) in operands.fixed() {
let OperandConstraint::FixedReg(preg) = op.constraint() else {
unreachable!();
};
if self.vreg_in_preg[preg.index()] != VReg::invalid()
&& self.vreg_in_preg[preg.index()] != op.vreg()
{
trace!(
"Evicting {} from fixed register {preg}",
self.vreg_in_preg[preg.index()]
);
if self.fixed_stack_slots.contains(preg)
&& self.edits.scratch_regs[preg.class()].is_none()
{
self.alloc_scratch_reg(inst, preg.class())?;
}
self.evict_vreg_in_preg(inst, preg);
self.vreg_in_preg[preg.index()] = VReg::invalid();
}
}
self.remove_clobbers_from_available_pregs(clobbers);
for preg in clobbers {
if self.vreg_in_preg[preg.index()] != VReg::invalid() {
trace!(
"Evicting {} from clobber {preg}",
self.vreg_in_preg[preg.index()]
);
if self.fixed_stack_slots.contains(preg)
&& self.edits.scratch_regs[preg.class()].is_none()
{
self.alloc_scratch_reg(inst, preg.class())?;
}
self.evict_vreg_in_preg(inst, preg);
self.vreg_in_preg[preg.index()] = VReg::invalid();
}
}
for (op_idx, op) in operands.def_ops() {
trace!("Allocating def operands {op}");
if let OperandConstraint::Reuse(reused_idx) = op.constraint() {
let reused_op = operands[reused_idx];
let new_reuse_op =
Operand::new(op.vreg(), reused_op.constraint(), op.kind(), op.pos());
trace!("allocating reuse op {op} as {new_reuse_op}");
self.process_operand_allocation(inst, new_reuse_op, op_idx)?;
} else {
self.process_operand_allocation(inst, op, op_idx)?;
}
let slot = self.vreg_spillslots[op.vreg().vreg()];
if slot.is_valid() {
self.vreg_to_live_inst_range[op.vreg().vreg()].2 = Allocation::stack(slot);
let curr_alloc = self.vreg_allocs[op.vreg().vreg()];
let vreg_slot = self.vreg_spillslots[op.vreg().vreg()];
let (is_stack_to_stack, src_and_dest_are_same) =
if let Some(curr_alloc) = curr_alloc.as_stack() {
(true, curr_alloc == vreg_slot)
} else {
(self.edits.is_stack(curr_alloc), false)
};
if !src_and_dest_are_same {
if is_stack_to_stack && self.edits.scratch_regs[op.class()].is_none() {
self.alloc_scratch_reg(inst, op.class())?;
};
self.edits.add_move(
inst,
self.vreg_allocs[op.vreg().vreg()],
Allocation::stack(self.vreg_spillslots[op.vreg().vreg()]),
op.class(),
InstPosition::After,
);
}
}
self.vreg_to_live_inst_range[op.vreg().vreg()].0 = ProgPoint::after(inst);
self.freealloc(op.vreg());
}
for (op_idx, op) in operands.use_ops() {
trace!("Allocating use operand {op}");
if self.reused_input_to_reuse_op[op_idx] != usize::MAX {
let reuse_op_idx = self.reused_input_to_reuse_op[op_idx];
let reuse_op_alloc = self.allocs[(inst.index(), reuse_op_idx)];
let Some(preg) = reuse_op_alloc.as_reg() else {
unreachable!();
};
let new_reused_input_constraint = OperandConstraint::FixedReg(preg);
let new_reused_input =
Operand::new(op.vreg(), new_reused_input_constraint, op.kind(), op.pos());
trace!("Allocating reused input {op} as {new_reused_input}");
self.process_operand_allocation(inst, new_reused_input, op_idx)?;
} else {
self.process_operand_allocation(inst, op, op_idx)?;
}
}
for (op_idx, op) in operands.use_ops() {
if op.as_fixed_nonallocatable().is_some() {
continue;
}
if self.vreg_allocs[op.vreg().vreg()] != self.allocs[(inst.index(), op_idx)] {
let curr_alloc = self.vreg_allocs[op.vreg().vreg()];
let new_alloc = self.allocs[(inst.index(), op_idx)];
trace!("Adding edit from {curr_alloc:?} to {new_alloc:?} before inst {inst:?} for {op}");
self.edits.add_move(
inst,
curr_alloc,
new_alloc,
op.class(),
InstPosition::Before,
);
}
}
if self.func.is_branch(inst) {
self.process_branch(block, inst)?;
}
for entry in self.reused_input_to_reuse_op.iter_mut() {
*entry = usize::MAX;
}
if trace_enabled!() {
self.log_post_inst_processing_state(inst);
}
Ok(())
}
fn reload_at_begin(&mut self, block: Block) -> Result<(), RegAllocError> {
trace!(
"Reloading live registers at the beginning of block {:?}",
block
);
trace!(
"Live registers at the beginning of block {:?}: {:?}",
block,
self.live_vregs
);
trace!(
"Block params at block {:?} beginning: {:?}",
block,
self.func.block_params(block)
);
trace!(
"Available pregs: {}",
self.available_pregs[OperandPos::Early]
);
self.reset_available_pregs_and_scratch_regs();
let avail_regs_for_scratch = self.available_pregs[OperandPos::Early];
let first_inst = self.func.block_insns(block).first();
for vreg in self.func.block_params(block).iter().cloned() {
trace!("Processing {}", vreg);
if self.vreg_allocs[vreg.vreg()] == Allocation::none() {
continue;
}
if self.vreg_spillslots[vreg.vreg()].is_invalid() {
self.vreg_spillslots[vreg.vreg()] = self.stack.allocstack(vreg.class());
}
let prev_alloc = self.vreg_allocs[vreg.vreg()];
let slot = Allocation::stack(self.vreg_spillslots[vreg.vreg()]);
self.vreg_to_live_inst_range[vreg.vreg()].2 = slot;
self.vreg_to_live_inst_range[vreg.vreg()].0 = ProgPoint::before(first_inst);
trace!("{} is a block param. Freeing it", vreg);
self.freealloc(vreg);
if slot == prev_alloc {
trace!(
"No need to reload {} because it's already in its expected allocation",
vreg
);
continue;
}
trace!(
"Move reason: reload {} at begin - move from its spillslot",
vreg
);
if self.edits.is_stack(prev_alloc) && self.edits.scratch_regs[vreg.class()].is_none() {
let reg = self.get_scratch_reg_for_reload(
first_inst,
vreg.class(),
avail_regs_for_scratch,
)?;
self.edits.scratch_regs[vreg.class()] = Some(reg);
}
self.edits.add_move(
self.func.block_insns(block).first(),
slot,
prev_alloc,
vreg.class(),
InstPosition::Before,
);
}
for vreg in self.live_vregs.iter() {
trace!("Processing {}", vreg);
trace!(
"{} is not a block param. It's a liveout vreg from some predecessor",
vreg
);
if self.vreg_spillslots[vreg.vreg()].is_invalid() {
self.vreg_spillslots[vreg.vreg()] = self.stack.allocstack(vreg.class());
}
let prev_alloc = self.vreg_allocs[vreg.vreg()];
let slot = Allocation::stack(self.vreg_spillslots[vreg.vreg()]);
trace!("Setting {}'s current allocation to its spillslot", vreg);
self.vreg_allocs[vreg.vreg()] = slot;
if let Some(preg) = prev_alloc.as_reg() {
trace!("{} was in {}. Removing it", preg, vreg);
self.vreg_in_preg[preg.index()] = VReg::invalid();
}
if slot == prev_alloc {
trace!(
"No need to reload {} because it's already in its expected allocation",
vreg
);
continue;
}
trace!(
"Move reason: reload {} at begin - move from its spillslot",
vreg
);
if self.edits.is_stack(prev_alloc) && self.edits.scratch_regs[vreg.class()].is_none() {
let Some(preg) = self.lrus[vreg.class()].last(avail_regs_for_scratch) else {
return Err(RegAllocError::TooManyLiveRegs);
};
if self.vreg_in_preg[preg.index()] != VReg::invalid() {
trace!("Removing the vreg in preg {} for eviction", preg);
let evicted_vreg = self.vreg_in_preg[preg.index()];
trace!("The removed vreg: {}", evicted_vreg);
debug_assert_ne!(evicted_vreg, VReg::invalid());
if self.vreg_spillslots[evicted_vreg.vreg()].is_invalid() {
self.vreg_spillslots[evicted_vreg.vreg()] =
self.stack.allocstack(evicted_vreg.class());
}
let slot = self.vreg_spillslots[evicted_vreg.vreg()];
self.vreg_allocs[evicted_vreg.vreg()] = Allocation::stack(slot);
trace!("Move reason: eviction");
self.edits.add_move(
first_inst,
self.vreg_allocs[evicted_vreg.vreg()],
Allocation::reg(preg),
evicted_vreg.class(),
InstPosition::Before,
);
}
self.edits.scratch_regs[vreg.class()] = Some(preg);
}
self.edits.add_move(
first_inst,
slot,
prev_alloc,
vreg.class(),
InstPosition::Before,
);
}
if trace_enabled!() {
self.log_post_reload_at_begin_state(block);
}
Ok(())
}
fn log_post_reload_at_begin_state(&self, block: Block) {
use alloc::format;
trace!("");
trace!("State after instruction reload_at_begin of {:?}", block);
let mut map = FxHashMap::default();
for (vreg_idx, alloc) in self.vreg_allocs.iter().enumerate() {
if *alloc != Allocation::none() {
map.insert(format!("vreg{vreg_idx}"), alloc);
}
}
trace!("vreg_allocs: {:?}", map);
let mut map = FxHashMap::default();
for i in 0..self.vreg_in_preg.len() {
if self.vreg_in_preg[i] != VReg::invalid() {
map.insert(PReg::from_index(i), self.vreg_in_preg[i]);
}
}
trace!("vreg_in_preg: {:?}", map);
trace!("Int LRU: {:?}", self.lrus[RegClass::Int]);
trace!("Float LRU: {:?}", self.lrus[RegClass::Float]);
trace!("Vector LRU: {:?}", self.lrus[RegClass::Vector]);
}
fn log_post_inst_processing_state(&self, inst: Inst) {
use alloc::format;
trace!("");
trace!("State after instruction {:?}", inst);
let mut map = FxHashMap::default();
for (vreg_idx, alloc) in self.vreg_allocs.iter().enumerate() {
if *alloc != Allocation::none() {
map.insert(format!("vreg{vreg_idx}"), alloc);
}
}
trace!("vreg_allocs: {:?}", map);
let mut v = Vec::new();
for i in 0..self.vreg_in_preg.len() {
if self.vreg_in_preg[i] != VReg::invalid() {
v.push(format!(
"{}: {}, ",
PReg::from_index(i),
self.vreg_in_preg[i]
));
}
}
trace!("vreg_in_preg: {:?}", v);
trace!("Int LRU: {:?}", self.lrus[RegClass::Int]);
trace!("Float LRU: {:?}", self.lrus[RegClass::Float]);
trace!("Vector LRU: {:?}", self.lrus[RegClass::Vector]);
trace!("");
}
fn alloc_block(&mut self, block: Block) -> Result<(), RegAllocError> {
trace!("{:?} start", block);
for inst in self.func.block_insns(block).iter().rev() {
self.alloc_inst(block, inst)?;
}
self.reload_at_begin(block)?;
trace!("{:?} end\n", block);
Ok(())
}
fn build_debug_info(&mut self) {
trace!("Building debug location info");
for &(vreg, start, end, label) in self.func.debug_value_labels() {
let (point_start, point_end, alloc) = self.vreg_to_live_inst_range[vreg.vreg()];
if point_start.inst() <= start && end <= point_end.inst().next() {
self.debug_locations
.push((label, point_start, point_end, alloc));
}
}
self.debug_locations.sort_by_key(|loc| loc.0);
}
fn run(&mut self) -> Result<(), RegAllocError> {
debug_assert_eq!(self.func.entry_block().index(), 0);
for block in (0..self.func.num_blocks()).rev() {
self.alloc_block(Block::new(block))?;
}
self.edits.edits.reverse();
self.build_debug_info();
Ok(())
}
}
fn log_function<F: Function>(func: &F) {
trace!("Processing a new function");
for block in 0..func.num_blocks() {
let block = Block::new(block);
trace!(
"Block {:?}. preds: {:?}. succs: {:?}, params: {:?}",
block,
func.block_preds(block),
func.block_succs(block),
func.block_params(block)
);
for inst in func.block_insns(block).iter() {
let clobbers = func.inst_clobbers(inst);
trace!(
"inst{:?}: {:?}. Clobbers: {}",
inst.index(),
func.inst_operands(inst),
clobbers
);
if func.is_branch(inst) {
trace!("Block args: ");
for (succ_idx, _succ) in func.block_succs(block).iter().enumerate() {
trace!(" {:?}", func.branch_blockparams(block, inst, succ_idx));
}
}
}
trace!("");
}
}
fn log_output<'a, F: Function>(env: &Env<'a, F>) {
trace!("Done!");
use alloc::format;
let mut v = Vec::new();
for i in 0..env.func.num_vregs() {
if env.vreg_spillslots[i].is_valid() {
v.push((
format!("{}", VReg::new(i, RegClass::Int)),
format!("{}", Allocation::stack(env.vreg_spillslots[i])),
));
}
}
trace!("VReg spillslots: {:?}", v);
trace!("Final edits: {:?}", env.edits.edits);
}
pub fn run<F: Function>(
func: &F,
mach_env: &MachineEnv,
verbose_log: bool,
enable_ssa_checker: bool,
) -> Result<Output, RegAllocError> {
if enable_ssa_checker {
validate_ssa(func, &CFGInfo::new(func)?)?;
}
if trace_enabled!() || verbose_log {
log_function(func);
}
let mut env = Env::new(func, mach_env);
env.run()?;
if trace_enabled!() || verbose_log {
log_output(&env);
}
Ok(Output {
edits: env.edits.edits,
allocs: env.allocs.allocs,
inst_alloc_offsets: env.allocs.inst_alloc_offsets,
num_spillslots: env.stack.num_spillslots as usize,
debug_locations: env.debug_locations,
stats: env.stats,
})
}