352 lines
12 KiB
Rust
352 lines
12 KiB
Rust
use crate::ir::ir_statement::IrStatement;
|
|
use crate::ir::ir_variable::IrVirtualRegisterVariable;
|
|
use std::cell::RefCell;
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::fmt::{Display, Formatter};
|
|
use std::rc::Rc;
|
|
|
|
pub struct IrBlock {
|
|
id: usize,
|
|
debug_label: String,
|
|
predecessors: Vec<Rc<RefCell<IrBlock>>>,
|
|
successors: Vec<Rc<RefCell<IrBlock>>>,
|
|
statements: Vec<IrStatement>,
|
|
}
|
|
|
|
type LivenessMapByStatement = HashMap<usize, HashSet<Rc<IrVirtualRegisterVariable>>>;
|
|
|
|
impl IrBlock {
|
|
pub fn new(id: usize, debug_label: &str, statements: Vec<IrStatement>) -> Self {
|
|
Self {
|
|
id,
|
|
debug_label: debug_label.into(),
|
|
predecessors: vec![],
|
|
successors: vec![],
|
|
statements,
|
|
}
|
|
}
|
|
|
|
pub fn id(&self) -> usize {
|
|
self.id
|
|
}
|
|
|
|
pub fn statements(&self) -> &[IrStatement] {
|
|
&self.statements
|
|
}
|
|
|
|
fn vr_definitions(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
|
|
let mut set = HashSet::new();
|
|
for statement in &self.statements {
|
|
set.extend(statement.vr_definitions());
|
|
}
|
|
set
|
|
}
|
|
|
|
fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
|
|
let mut set = HashSet::new();
|
|
for statement in &self.statements {
|
|
set.extend(statement.vr_uses());
|
|
}
|
|
set
|
|
}
|
|
|
|
fn live_in_live_out(&self) -> (LivenessMapByStatement, LivenessMapByStatement) {
|
|
let mut live_in: LivenessMapByStatement = HashMap::new();
|
|
let mut live_out: LivenessMapByStatement = HashMap::new();
|
|
|
|
loop {
|
|
let mut did_work = false;
|
|
for (statement_index, statement) in self.statements.iter().enumerate().rev() {
|
|
// init if necessary
|
|
if !live_in.contains_key(&statement_index) {
|
|
live_in.insert(statement_index, HashSet::new());
|
|
}
|
|
if !live_out.contains_key(&statement_index) {
|
|
live_out.insert(statement_index, HashSet::new());
|
|
}
|
|
|
|
// out (union of successors ins)
|
|
// for now, a statement can only have one successor
|
|
// this will need to be updated when we add jumps
|
|
if let Some(successor_live_in) = live_in.get(&(statement_index + 1)) {
|
|
let statement_live_out = live_out.get_mut(&statement_index).unwrap();
|
|
for vr_variable in successor_live_in {
|
|
if statement_live_out.insert(vr_variable.clone()) {
|
|
did_work = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
// in: use(s) U ( out(s) - def(s) )
|
|
let mut new_ins = statement.vr_uses();
|
|
let statement_live_out = live_out.get(&statement_index).unwrap();
|
|
let defs = statement.vr_definitions();
|
|
let rhs = statement_live_out - &defs;
|
|
new_ins.extend(rhs);
|
|
|
|
let statement_live_in = live_in.get_mut(&statement_index).unwrap();
|
|
for new_in in new_ins {
|
|
if statement_live_in.insert(new_in) {
|
|
did_work = true;
|
|
}
|
|
}
|
|
}
|
|
if !did_work {
|
|
break;
|
|
}
|
|
}
|
|
(live_in, live_out)
|
|
}
|
|
|
|
fn interference_graph(
|
|
&self,
|
|
) -> HashMap<Rc<IrVirtualRegisterVariable>, HashSet<Rc<IrVirtualRegisterVariable>>> {
|
|
let mut all_vr_variables: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
|
|
for statement in &self.statements {
|
|
all_vr_variables.extend(statement.vr_definitions());
|
|
all_vr_variables.extend(statement.vr_uses());
|
|
}
|
|
|
|
let mut graph: HashMap<
|
|
Rc<IrVirtualRegisterVariable>,
|
|
HashSet<Rc<IrVirtualRegisterVariable>>,
|
|
> = HashMap::new();
|
|
for variable in all_vr_variables {
|
|
graph.insert(variable, HashSet::new());
|
|
}
|
|
|
|
let (_, live_out) = self.live_in_live_out();
|
|
|
|
for (statement_index, statement) in self.statements.iter().enumerate() {
|
|
let statement_live_out = live_out.get(&statement_index).unwrap();
|
|
for definition_vr_variable in statement.vr_definitions() {
|
|
for live_out_variable in statement_live_out {
|
|
// we do the following check to avoid adding an edge to itself
|
|
if definition_vr_variable != *live_out_variable {
|
|
graph
|
|
.get_mut(&definition_vr_variable)
|
|
.unwrap()
|
|
.insert(live_out_variable.clone());
|
|
graph
|
|
.get_mut(live_out_variable)
|
|
.unwrap()
|
|
.insert(definition_vr_variable.clone());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
graph
|
|
}
|
|
|
|
pub fn assign_registers(&mut self) {
|
|
let mut spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
|
|
loop {
|
|
let (registers, new_spills) = register_assignment::assign_registers(self);
|
|
if spills != new_spills {
|
|
spills = new_spills;
|
|
// mutate all IrVirtualRegisters to constituent statements
|
|
for statement in &mut self.statements {
|
|
statement.propagate_spills(&spills);
|
|
}
|
|
} else {
|
|
println!("{:?}", registers);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Display for IrBlock {
|
|
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
|
writeln!(f, " {}:", self.debug_label)?;
|
|
|
|
let (live_in, live_out) = self.live_in_live_out();
|
|
|
|
for (statement_index, statement) in self.statements.iter().enumerate() {
|
|
let statement_live_in = live_in.get(&statement_index).unwrap();
|
|
let statement_live_out = live_out.get(&statement_index).unwrap();
|
|
|
|
writeln!(
|
|
f,
|
|
" {} // live_in: {:?}, live_out: {:?}",
|
|
statement, statement_live_in, statement_live_out
|
|
)?;
|
|
}
|
|
writeln!(f, " // ---- {} meta ----", self.debug_label)?;
|
|
writeln!(f, " // definitions: {:?}", self.vr_definitions())?;
|
|
writeln!(f, " // uses: {:?}", self.vr_uses())?;
|
|
writeln!(
|
|
f,
|
|
" // interference graph: {:?}",
|
|
self.interference_graph()
|
|
)?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
mod register_assignment {
|
|
use super::*;
|
|
|
|
#[derive(Debug)]
|
|
struct WorkItem {
|
|
vr: Rc<IrVirtualRegisterVariable>,
|
|
edges: HashSet<Rc<IrVirtualRegisterVariable>>,
|
|
color: bool,
|
|
}
|
|
|
|
pub fn assign_registers(
|
|
block: &IrBlock,
|
|
) -> (
|
|
HashMap<Rc<IrVirtualRegisterVariable>, usize>,
|
|
HashSet<Rc<IrVirtualRegisterVariable>>,
|
|
) {
|
|
let k = 8;
|
|
let mut spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
|
|
|
|
loop {
|
|
// 1. get interference graph
|
|
let mut interference_graph = block.interference_graph();
|
|
let mut work_stack: Vec<WorkItem> = vec![];
|
|
|
|
loop {
|
|
// 2. coloring by simplification
|
|
// try to find a node (virtual register) with less than k outgoing edges,
|
|
// and mark as color
|
|
// if not, pick any, and mark as spill for step 3
|
|
let register_lt_k = interference_graph.iter().find_map(|(vr, neighbors)| {
|
|
if neighbors.len() < k { Some(vr) } else { None }
|
|
});
|
|
if let Some(vr) = register_lt_k {
|
|
let vr = vr.clone();
|
|
// remove both outgoing and incoming edges; save either set for WorkItem
|
|
// first, outgoing:
|
|
let outgoing_edges = interference_graph.remove(&vr).unwrap();
|
|
|
|
// second, incoming
|
|
interference_graph.iter_mut().for_each(|(_, neighbors)| {
|
|
neighbors.remove(&vr);
|
|
});
|
|
|
|
// push to work stack
|
|
work_stack.push(WorkItem {
|
|
vr,
|
|
edges: outgoing_edges,
|
|
color: true,
|
|
})
|
|
} else {
|
|
// pick any
|
|
let vr = interference_graph.iter().last().unwrap().0.clone();
|
|
|
|
// first, outgoing
|
|
let outgoing_edges = interference_graph.remove(&vr).unwrap();
|
|
|
|
// second, incoming
|
|
interference_graph.iter_mut().for_each(|(_, neighbors)| {
|
|
neighbors.remove(&vr);
|
|
});
|
|
|
|
work_stack.push(WorkItem {
|
|
vr,
|
|
edges: outgoing_edges,
|
|
color: false, // spill
|
|
});
|
|
}
|
|
|
|
if interference_graph.is_empty() {
|
|
break;
|
|
}
|
|
}
|
|
|
|
// 3. assign colors to registers
|
|
let mut rebuilt_graph: HashMap<
|
|
Rc<IrVirtualRegisterVariable>,
|
|
HashSet<Rc<IrVirtualRegisterVariable>>,
|
|
> = HashMap::new();
|
|
let mut register_assignments: HashMap<Rc<IrVirtualRegisterVariable>, usize> =
|
|
HashMap::new();
|
|
let mut new_spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
|
|
|
|
while let Some(work_item) = work_stack.pop() {
|
|
if work_item.color {
|
|
assign_register(&work_item, &mut rebuilt_graph, k, &mut register_assignments);
|
|
} else {
|
|
// first, see if we can optimistically color
|
|
// find how many assignments have been made for the outgoing edges
|
|
// if it's less than k, we can do it
|
|
let mut number_of_assigned_edges = 0;
|
|
for edge in &work_item.edges {
|
|
if register_assignments.contains_key(edge) {
|
|
number_of_assigned_edges += 1;
|
|
}
|
|
}
|
|
|
|
if number_of_assigned_edges < k {
|
|
// optimistically color
|
|
assign_register(
|
|
&work_item,
|
|
&mut rebuilt_graph,
|
|
k,
|
|
&mut register_assignments,
|
|
);
|
|
} else {
|
|
// spill
|
|
new_spills.insert(work_item.vr.clone());
|
|
}
|
|
}
|
|
}
|
|
|
|
if spills.eq(&new_spills) {
|
|
return (register_assignments, spills);
|
|
} else {
|
|
spills = new_spills;
|
|
}
|
|
}
|
|
}
|
|
|
|
fn assign_register(
|
|
work_item: &WorkItem,
|
|
rebuilt_graph: &mut HashMap<
|
|
Rc<IrVirtualRegisterVariable>,
|
|
HashSet<Rc<IrVirtualRegisterVariable>>,
|
|
>,
|
|
k: usize,
|
|
register_assignments: &mut HashMap<Rc<IrVirtualRegisterVariable>, usize>,
|
|
) {
|
|
let this_vertex_vr = &work_item.vr;
|
|
// init the vertex
|
|
rebuilt_graph.insert(this_vertex_vr.clone(), HashSet::new());
|
|
|
|
// add edges, both outgoing and incoming
|
|
let neighbors = rebuilt_graph.get_mut(this_vertex_vr).unwrap();
|
|
// outgoing
|
|
for edge in &work_item.edges {
|
|
neighbors.insert(edge.clone());
|
|
}
|
|
// incoming
|
|
for neighbor in neighbors.clone() {
|
|
if rebuilt_graph.contains_key(&neighbor) {
|
|
rebuilt_graph
|
|
.get_mut(&neighbor)
|
|
.unwrap()
|
|
.insert(this_vertex_vr.clone());
|
|
}
|
|
}
|
|
|
|
// find a register which is not yet shared by all outgoing edges' vertices
|
|
// I think the bug is somewhere here
|
|
'outer: for i in 0..k {
|
|
for edge in rebuilt_graph.get_mut(this_vertex_vr).unwrap().iter() {
|
|
if register_assignments.contains_key(edge) {
|
|
let assignment = register_assignments.get(edge).unwrap();
|
|
if *assignment == i {
|
|
continue 'outer;
|
|
}
|
|
}
|
|
}
|
|
register_assignments.insert(this_vertex_vr.clone(), i);
|
|
break;
|
|
}
|
|
}
|
|
}
|