use crate::cursor::{Cursor, EncCursor};
use crate::dominator_tree::DominatorTree;
use crate::ir::{ArgumentLoc, Block, Function, Inst, InstBuilder, SigRef, Value, ValueLoc};
use crate::isa::registers::{RegClass, RegClassIndex, RegClassMask, RegUnit};
use crate::isa::{ConstraintKind, EncInfo, RecipeConstraints, RegInfo, TargetIsa};
use crate::regalloc::affinity::Affinity;
use crate::regalloc::live_value_tracker::{LiveValue, LiveValueTracker};
use crate::regalloc::liveness::Liveness;
use crate::regalloc::pressure::Pressure;
use crate::regalloc::virtregs::VirtRegs;
use crate::timing;
use crate::topo_order::TopoOrder;
use alloc::vec::Vec;
use core::fmt;
use log::debug;
fn toprc_containing_regunit(unit: RegUnit, reginfo: &RegInfo) -> RegClass {
let bank = reginfo.bank_containing_regunit(unit).unwrap();
reginfo.classes[bank.first_toprc..(bank.first_toprc + bank.num_toprcs)]
.iter()
.find(|&rc| rc.contains(unit))
.expect("reg unit should be in a toprc")
}
pub struct Spilling {
spills: Vec<Value>,
reg_uses: Vec<RegUse>,
}
struct Context<'a> {
cur: EncCursor<'a>,
reginfo: RegInfo,
encinfo: EncInfo,
domtree: &'a DominatorTree,
liveness: &'a mut Liveness,
virtregs: &'a VirtRegs,
topo: &'a mut TopoOrder,
pressure: Pressure,
spills: &'a mut Vec<Value>,
reg_uses: &'a mut Vec<RegUse>,
}
impl Spilling {
pub fn new() -> Self {
Self {
spills: Vec::new(),
reg_uses: Vec::new(),
}
}
pub fn clear(&mut self) {
self.spills.clear();
self.reg_uses.clear();
}
pub fn run(
&mut self,
isa: &dyn TargetIsa,
func: &mut Function,
domtree: &DominatorTree,
liveness: &mut Liveness,
virtregs: &VirtRegs,
topo: &mut TopoOrder,
tracker: &mut LiveValueTracker,
) {
let _tt = timing::ra_spilling();
debug!("Spilling for:\n{}", func.display(isa));
let reginfo = isa.register_info();
let usable_regs = isa.allocatable_registers(func);
let mut ctx = Context {
cur: EncCursor::new(func, isa),
reginfo: isa.register_info(),
encinfo: isa.encoding_info(),
domtree,
liveness,
virtregs,
topo,
pressure: Pressure::new(®info, &usable_regs),
spills: &mut self.spills,
reg_uses: &mut self.reg_uses,
};
ctx.run(tracker)
}
}
impl<'a> Context<'a> {
fn run(&mut self, tracker: &mut LiveValueTracker) {
self.topo.reset(self.cur.func.layout.blocks());
while let Some(block) = self.topo.next(&self.cur.func.layout, self.domtree) {
self.visit_block(block, tracker);
}
}
fn visit_block(&mut self, block: Block, tracker: &mut LiveValueTracker) {
debug!("Spilling {}:", block);
self.cur.goto_top(block);
self.visit_block_header(block, tracker);
tracker.drop_dead_params();
self.process_spills(tracker);
while let Some(inst) = self.cur.next_inst() {
if !self.cur.func.dfg[inst].opcode().is_ghost() {
self.visit_inst(inst, block, tracker);
} else {
let (_throughs, kills) = tracker.process_ghost(inst);
self.free_regs(kills);
}
tracker.drop_dead(inst);
self.process_spills(tracker);
}
}
fn take_live_regs(&mut self, regs: &[LiveValue]) {
for lv in regs {
if !lv.is_dead {
if let Affinity::Reg(rci) = lv.affinity {
let rc = self.reginfo.rc(rci);
self.pressure.take(rc);
}
}
}
}
fn free_regs(&mut self, kills: &[LiveValue]) {
for lv in kills {
if let Affinity::Reg(rci) = lv.affinity {
if !self.spills.contains(&lv.value) {
let rc = self.reginfo.rc(rci);
self.pressure.free(rc);
}
}
}
}
fn free_dead_regs(&mut self, regs: &[LiveValue]) {
for lv in regs {
if lv.is_dead {
if let Affinity::Reg(rci) = lv.affinity {
if !self.spills.contains(&lv.value) {
let rc = self.reginfo.rc(rci);
self.pressure.free(rc);
}
}
}
}
}
fn visit_block_header(&mut self, block: Block, tracker: &mut LiveValueTracker) {
let (liveins, params) = tracker.block_top(
block,
&self.cur.func.dfg,
self.liveness,
&self.cur.func.layout,
self.domtree,
);
self.pressure.reset();
self.take_live_regs(liveins);
for lv in params {
if let Affinity::Reg(rci) = lv.affinity {
let rc = self.reginfo.rc(rci);
'try_take: while let Err(mask) = self.pressure.take_transient(rc) {
debug!("Need {} reg for block param {}", rc, lv.value);
match self.spill_candidate(mask, liveins) {
Some(cand) => {
debug!(
"Spilling live-in {} to make room for {} block param {}",
cand, rc, lv.value
);
self.spill_reg(cand);
}
None => {
debug!("Spilling {} block argument {}", rc, lv.value);
self.pressure.take(rc);
self.spill_reg(lv.value);
break 'try_take;
}
}
}
}
}
self.pressure.preserve_transient();
self.free_dead_regs(params);
}
fn visit_inst(&mut self, inst: Inst, block: Block, tracker: &mut LiveValueTracker) {
debug!("Inst {}, {}", self.cur.display_inst(inst), self.pressure);
debug_assert_eq!(self.cur.current_inst(), Some(inst));
debug_assert_eq!(self.cur.current_block(), Some(block));
let constraints = self
.encinfo
.operand_constraints(self.cur.func.encodings[inst]);
debug_assert!(self.reg_uses.is_empty());
self.collect_reg_uses(inst, block, constraints);
let call_sig = self.cur.func.dfg.call_signature(inst);
if let Some(sig) = call_sig {
self.collect_abi_reg_uses(inst, sig);
}
if !self.reg_uses.is_empty() {
self.process_reg_uses(inst, tracker);
}
let (throughs, kills, defs) = tracker.process_inst(inst, &self.cur.func.dfg, self.liveness);
self.free_regs(kills);
let opcode = self.cur.func.dfg[inst].opcode();
if call_sig.is_some() || opcode.clobbers_all_regs() {
for lv in throughs {
if lv.affinity.is_reg() && !self.spills.contains(&lv.value) {
self.spill_reg(lv.value);
}
}
}
if let Some(constraints) = constraints {
for op in constraints.outs {
if op.kind != ConstraintKind::Stack {
while let Err(mask) = self.pressure.take_transient(op.regclass) {
debug!("Need {} reg from {} throughs", op.regclass, throughs.len());
match self.spill_candidate(mask, throughs) {
Some(cand) => self.spill_reg(cand),
None => panic!(
"Ran out of {} registers for {}",
op.regclass,
self.cur.display_inst(inst)
),
}
}
}
}
self.pressure.reset_transient();
}
self.take_live_regs(defs);
}
fn collect_reg_uses(
&mut self,
inst: Inst,
block: Block,
constraints: Option<&RecipeConstraints>,
) {
let args = self.cur.func.dfg.inst_args(inst);
let num_fixed_ins = if let Some(constraints) = constraints {
for (idx, (op, &arg)) in constraints.ins.iter().zip(args).enumerate() {
let mut reguse = RegUse::new(arg, idx, op.regclass.into());
let lr = &self.liveness[arg];
match op.kind {
ConstraintKind::Stack => continue,
ConstraintKind::FixedReg(_) => reguse.fixed = true,
ConstraintKind::Tied(_) => {
reguse.tied = !lr.killed_at(inst, block, &self.cur.func.layout);
}
ConstraintKind::FixedTied(_) => {
reguse.fixed = true;
reguse.tied = !lr.killed_at(inst, block, &self.cur.func.layout);
}
ConstraintKind::Reg => {}
}
if lr.affinity.is_stack() {
reguse.spilled = true;
}
if reguse.fixed || reguse.tied || reguse.spilled {
debug!(" reguse: {}", reguse);
self.reg_uses.push(reguse);
}
}
constraints.ins.len()
} else {
0
};
if self.cur.func.dfg[inst].opcode().is_return() {
debug_assert_eq!(
self.cur.func.dfg.inst_variable_args(inst).len(),
self.cur.func.signature.returns.len(),
"The non-fixed arguments in a return should follow the function's signature."
);
for (ret_idx, (ret, &arg)) in
self.cur.func.signature.returns.iter().zip(args).enumerate()
{
let idx = num_fixed_ins + ret_idx;
let unit = match ret.location {
ArgumentLoc::Unassigned => {
panic!("function return signature should be legalized")
}
ArgumentLoc::Reg(unit) => unit,
ArgumentLoc::Stack(_) => continue,
};
let toprc = toprc_containing_regunit(unit, &self.reginfo);
let mut reguse = RegUse::new(arg, idx, toprc.into());
reguse.fixed = true;
debug!(" reguse: {}", reguse);
self.reg_uses.push(reguse);
}
}
}
fn collect_abi_reg_uses(&mut self, inst: Inst, sig: SigRef) {
let num_fixed_args = self.cur.func.dfg[inst]
.opcode()
.constraints()
.num_fixed_value_arguments();
let args = self.cur.func.dfg.inst_variable_args(inst);
for (idx, (abi, &arg)) in self.cur.func.dfg.signatures[sig]
.params
.iter()
.zip(args)
.enumerate()
{
if abi.location.is_reg() {
let (rci, spilled) = match self.liveness[arg].affinity {
Affinity::Reg(rci) => (rci, false),
Affinity::Stack => (
self.cur.isa.regclass_for_abi_type(abi.value_type).into(),
true,
),
Affinity::Unassigned => panic!("Missing affinity for {}", arg),
};
let mut reguse = RegUse::new(arg, num_fixed_args + idx, rci);
reguse.fixed = true;
reguse.spilled = spilled;
self.reg_uses.push(reguse);
}
}
}
fn process_reg_uses(&mut self, inst: Inst, tracker: &LiveValueTracker) {
self.reg_uses.sort_unstable_by_key(|u| (u.value, u.opidx));
self.cur.use_srcloc(inst);
for i in 0..self.reg_uses.len() {
let ru = self.reg_uses[i];
let need_copy = if ru.tied {
true
} else if ru.fixed {
self.reg_uses
.get(i.wrapping_sub(1))
.map_or(false, |ru2| ru2.value == ru.value)
} else {
false
};
if need_copy {
let copy = self.insert_copy(ru.value, ru.rci);
self.cur.func.dfg.inst_args_mut(inst)[ru.opidx as usize] = copy;
}
if need_copy || ru.spilled {
let rc = self.reginfo.rc(ru.rci);
while let Err(mask) = self.pressure.take_transient(rc) {
debug!("Copy of {} reg causes spill", rc);
match {
let args = if self.cur.func.dfg[inst].opcode().is_branch() {
self.cur.func.dfg.inst_fixed_args(inst)
} else {
self.cur.func.dfg.inst_args(inst)
};
self.spill_candidate(
mask,
tracker.live().iter().filter(|lv| !args.contains(&lv.value)),
)
} {
Some(cand) => self.spill_reg(cand),
None => panic!(
"Ran out of {} registers when inserting copy before {}",
rc,
self.cur.display_inst(inst)
),
}
}
}
}
self.pressure.reset_transient();
self.reg_uses.clear()
}
fn spill_candidate<'ii, II>(&self, mask: RegClassMask, candidates: II) -> Option<Value>
where
II: IntoIterator<Item = &'ii LiveValue>,
{
candidates
.into_iter()
.filter_map(|lv| {
if let Affinity::Reg(rci) = lv.affinity {
let rc = self.reginfo.rc(rci);
if (mask & (1 << rc.toprc)) != 0 && !self.spills.contains(&lv.value) {
return Some(lv.value);
}
}
None
})
.min_by(|&a, &b| {
self.domtree.rpo_cmp(
self.cur.func.dfg.value_def(a),
self.cur.func.dfg.value_def(b),
&self.cur.func.layout,
)
})
}
fn spill_reg(&mut self, value: Value) {
if let Affinity::Reg(rci) = self.liveness.spill(value) {
let rc = self.reginfo.rc(rci);
self.pressure.free(rc);
self.spills.push(value);
debug!("Spilled {}:{} -> {}", value, rc, self.pressure);
} else {
panic!("Cannot spill {} that was already on the stack", value);
}
let ss = self
.cur
.func
.stack_slots
.make_spill_slot(self.cur.func.dfg.value_type(value));
for &v in self.virtregs.congruence_class(&value) {
self.liveness.spill(v);
self.cur.func.locations[v] = ValueLoc::Stack(ss);
}
}
fn process_spills(&mut self, tracker: &mut LiveValueTracker) {
if !self.spills.is_empty() {
tracker.process_spills(|v| self.spills.contains(&v));
self.spills.clear()
}
}
fn insert_copy(&mut self, value: Value, rci: RegClassIndex) -> Value {
let copy = self.cur.ins().copy(value);
let inst = self.cur.built_inst();
self.liveness.create_dead(copy, inst, Affinity::Reg(rci));
self.liveness.extend_locally(
copy,
self.cur.func.layout.pp_block(inst),
self.cur.current_inst().expect("must be at an instruction"),
&self.cur.func.layout,
);
copy
}
}
#[derive(Clone, Copy)]
struct RegUse {
value: Value,
opidx: u16,
rci: RegClassIndex,
fixed: bool,
spilled: bool,
tied: bool,
}
impl RegUse {
fn new(value: Value, idx: usize, rci: RegClassIndex) -> Self {
Self {
value,
opidx: idx as u16,
rci,
fixed: false,
spilled: false,
tied: false,
}
}
}
impl fmt::Display for RegUse {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}@op{}", self.value, self.opidx)?;
if self.fixed {
write!(f, "/fixed")?;
}
if self.spilled {
write!(f, "/spilled")?;
}
if self.tied {
write!(f, "/tied")?;
}
Ok(())
}
}