463 lines
15 KiB
Rust
463 lines
15 KiB
Rust
use crate::ir::ir_variable::IrVrVariableDescriptor;
|
|
use std::collections::{HashMap, HashSet};
|
|
|
|
pub type InterferenceGraph = HashMap<IrVrVariableDescriptor, HashSet<IrVrVariableDescriptor>>;
|
|
pub type LivenessMap = HashMap<usize, HashSet<IrVrVariableDescriptor>>;
|
|
|
|
pub trait HasVrUsers {
|
|
fn vr_users(&self) -> Vec<&dyn VrUser>;
|
|
fn vr_users_mut(&mut self) -> Vec<&mut dyn VrUser>;
|
|
|
|
fn live_in_live_out(&self) -> (LivenessMap, LivenessMap) {
|
|
let mut live_in: LivenessMap = HashMap::new();
|
|
let mut live_out: LivenessMap = HashMap::new();
|
|
|
|
loop {
|
|
let mut did_work = false;
|
|
for (block_index, statement) in self.vr_users().iter().enumerate().rev() {
|
|
// init if necessary
|
|
if !live_in.contains_key(&block_index) {
|
|
live_in.insert(block_index, HashSet::new());
|
|
}
|
|
if !live_out.contains_key(&block_index) {
|
|
live_out.insert(block_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(&(block_index + 1)) {
|
|
let statement_live_out = live_out.get_mut(&block_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()
|
|
.iter()
|
|
.map(|u| (*u).clone())
|
|
.collect::<HashSet<_>>();
|
|
let statement_live_out = live_out.get(&block_index).unwrap();
|
|
let defs = statement
|
|
.vr_definitions()
|
|
.iter()
|
|
.map(|d| (*d).clone())
|
|
.collect::<HashSet<_>>();
|
|
let rhs = statement_live_out - &defs;
|
|
new_ins.extend(rhs);
|
|
|
|
let statement_live_in = live_in.get_mut(&block_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) -> InterferenceGraph {
|
|
let mut all_vr_variables: HashSet<IrVrVariableDescriptor> = HashSet::new();
|
|
|
|
for vr_user in self.vr_users() {
|
|
all_vr_variables.extend(vr_user.vr_definitions());
|
|
all_vr_variables.extend(vr_user.vr_uses());
|
|
}
|
|
|
|
let mut graph: InterferenceGraph = HashMap::new();
|
|
for variable in all_vr_variables {
|
|
graph.insert(variable, HashSet::new());
|
|
}
|
|
|
|
let (_, live_out) = self.live_in_live_out();
|
|
|
|
for (index, vr_user) in self.vr_users().iter().enumerate() {
|
|
let user_live_in = live_out.get(&index).unwrap();
|
|
for definition_vr_variable in vr_user.vr_definitions() {
|
|
for live_out_variable in user_live_in {
|
|
// we do the following check to avoid adding an edge to itself
|
|
if definition_vr_variable.ne(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
|
|
}
|
|
|
|
fn assign_registers(
|
|
&mut self,
|
|
register_count: usize,
|
|
) -> (HashMap<IrVrVariableDescriptor, usize>, isize) {
|
|
let mut spills: HashSet<IrVrVariableDescriptor> = HashSet::new();
|
|
loop {
|
|
let mut interference_graph = self.interference_graph();
|
|
let (registers, new_spills) =
|
|
registers_and_spills(&mut interference_graph, register_count);
|
|
|
|
if spills != new_spills {
|
|
spills = new_spills;
|
|
// propagate all spills, since those won't be used for the next interference graph
|
|
for vr_user in &mut self.vr_users_mut() {
|
|
vr_user.propagate_spills(&spills);
|
|
}
|
|
} else {
|
|
// we've calculated final assignments, so propagate them
|
|
for vr_user in self.vr_users_mut() {
|
|
vr_user.propagate_register_assignments(®isters);
|
|
}
|
|
// also set offsets
|
|
let mut offset_counter = OffsetCounter::new();
|
|
for vr_user in self.vr_users_mut() {
|
|
vr_user.propagate_stack_offsets(&mut offset_counter);
|
|
}
|
|
|
|
return (registers, offset_counter.get_count());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub trait VrUser {
|
|
fn vr_definitions(&self) -> HashSet<IrVrVariableDescriptor>;
|
|
fn vr_uses(&self) -> HashSet<IrVrVariableDescriptor>;
|
|
fn propagate_spills(&mut self, spills: &HashSet<IrVrVariableDescriptor>);
|
|
fn propagate_register_assignments(
|
|
&mut self,
|
|
assignments: &HashMap<IrVrVariableDescriptor, usize>,
|
|
);
|
|
fn propagate_stack_offsets(&mut self, counter: &mut OffsetCounter);
|
|
}
|
|
|
|
pub struct OffsetCounter {
|
|
counter: isize,
|
|
}
|
|
|
|
impl OffsetCounter {
|
|
pub fn new() -> Self {
|
|
Self { counter: 0 }
|
|
}
|
|
|
|
pub fn next(&mut self) -> isize {
|
|
let offset = self.counter;
|
|
self.counter += 1;
|
|
offset
|
|
}
|
|
|
|
pub fn get_count(&self) -> isize {
|
|
self.counter
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct WorkItem {
|
|
vr: IrVrVariableDescriptor,
|
|
edges: HashSet<IrVrVariableDescriptor>,
|
|
color: bool,
|
|
}
|
|
|
|
pub fn registers_and_spills(
|
|
interference_graph: &mut InterferenceGraph,
|
|
k: usize,
|
|
) -> (
|
|
HashMap<IrVrVariableDescriptor, usize>,
|
|
HashSet<IrVrVariableDescriptor>,
|
|
) {
|
|
let mut work_stack: Vec<WorkItem> = vec![];
|
|
|
|
while !interference_graph.is_empty() {
|
|
work_stack.push(next_work_item(interference_graph, k));
|
|
}
|
|
|
|
// 3. assign colors to registers
|
|
let mut rebuilt_graph: InterferenceGraph = HashMap::new();
|
|
let mut register_assignments: HashMap<IrVrVariableDescriptor, usize> = HashMap::new();
|
|
let mut spills: HashSet<IrVrVariableDescriptor> = 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 if can_optimistically_color(&work_item, &mut register_assignments, k) {
|
|
assign_register(&work_item, &mut rebuilt_graph, k, &mut register_assignments);
|
|
} else {
|
|
// spill
|
|
spills.insert(work_item.vr.clone());
|
|
}
|
|
}
|
|
|
|
(register_assignments, spills)
|
|
}
|
|
|
|
fn assign_register(
|
|
work_item: &WorkItem,
|
|
graph: &mut InterferenceGraph,
|
|
k: usize,
|
|
register_assignments: &mut HashMap<IrVrVariableDescriptor, usize>,
|
|
) {
|
|
rebuild_vr_and_edges(graph, work_item);
|
|
|
|
// find a register which is not yet shared by all outgoing edges' vertices
|
|
'outer: for i in 0..k {
|
|
for edge in graph.get_mut(&work_item.vr).unwrap().iter() {
|
|
if register_assignments.contains_key(edge) {
|
|
let assignment = register_assignments.get(edge).unwrap();
|
|
if *assignment == i {
|
|
continue 'outer;
|
|
}
|
|
}
|
|
}
|
|
register_assignments.insert(work_item.vr.clone(), i);
|
|
break;
|
|
}
|
|
}
|
|
|
|
fn find_vr_lt_k(
|
|
interference_graph: &InterferenceGraph,
|
|
k: usize,
|
|
) -> Option<&IrVrVariableDescriptor> {
|
|
interference_graph.iter().find_map(
|
|
|(vr, neighbors)| {
|
|
if neighbors.len() < k { Some(vr) } else { None }
|
|
},
|
|
)
|
|
}
|
|
|
|
fn remove_vr_and_edges(
|
|
interference_graph: &mut InterferenceGraph,
|
|
vr: &IrVrVariableDescriptor,
|
|
) -> HashSet<IrVrVariableDescriptor> {
|
|
// first, outgoing
|
|
let outgoing_edges = interference_graph.remove(vr).unwrap();
|
|
|
|
// second, incoming
|
|
for neighbor in &outgoing_edges {
|
|
let neighbor_edges = interference_graph.get_mut(neighbor).unwrap();
|
|
neighbor_edges.remove(vr);
|
|
}
|
|
|
|
outgoing_edges
|
|
}
|
|
|
|
fn next_work_item(interference_graph: &mut InterferenceGraph, k: usize) -> WorkItem {
|
|
// try to find a node (virtual register) with less than k outgoing edges, and mark as color
|
|
// for step 3.
|
|
// if not, pick any, and mark as spill for step 3.
|
|
let register_lt_k = find_vr_lt_k(interference_graph, k);
|
|
if let Some(vr) = register_lt_k {
|
|
let vr = vr.clone();
|
|
|
|
// remove edges; save outgoing to work_item
|
|
let edges = remove_vr_and_edges(interference_graph, &vr);
|
|
|
|
// push to work stack
|
|
WorkItem {
|
|
vr,
|
|
edges,
|
|
color: true,
|
|
}
|
|
} else {
|
|
// pick any
|
|
let vr = interference_graph.iter().last().unwrap().0.clone();
|
|
|
|
// remove edges
|
|
let edges = remove_vr_and_edges(interference_graph, &vr);
|
|
|
|
WorkItem {
|
|
vr,
|
|
edges,
|
|
color: false, // spill
|
|
}
|
|
}
|
|
}
|
|
|
|
fn rebuild_vr_and_edges(graph: &mut InterferenceGraph, work_item: &WorkItem) {
|
|
// init the vertex
|
|
graph.insert(work_item.vr.clone(), HashSet::new());
|
|
|
|
// outgoing
|
|
for neighbor in &work_item.edges {
|
|
// check if neighbor exists in the graph first; if it was marked spill earlier and could
|
|
// not optimistically color, it won't be in the graph
|
|
if graph.contains_key(neighbor) {
|
|
// get outgoing set and insert neighbor
|
|
graph
|
|
.get_mut(&work_item.vr)
|
|
.unwrap()
|
|
.insert(neighbor.clone());
|
|
}
|
|
}
|
|
|
|
// incoming
|
|
for neighbor in &work_item.edges {
|
|
// like above, neighbor may not have been added because of failure to optimistically
|
|
// color
|
|
if graph.contains_key(neighbor) {
|
|
graph
|
|
.get_mut(neighbor)
|
|
.unwrap()
|
|
.insert(work_item.vr.clone());
|
|
}
|
|
}
|
|
}
|
|
|
|
fn can_optimistically_color(
|
|
work_item: &WorkItem,
|
|
register_assignments: &HashMap<IrVrVariableDescriptor, usize>,
|
|
k: usize,
|
|
) -> bool {
|
|
// 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;
|
|
}
|
|
}
|
|
number_of_assigned_edges < k
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn line_graph() -> InterferenceGraph {
|
|
let mut graph: InterferenceGraph = HashMap::new();
|
|
let v0 = IrVrVariableDescriptor::new("v0".into(), 0);
|
|
let v1 = IrVrVariableDescriptor::new("v1".into(), 0);
|
|
let v2 = IrVrVariableDescriptor::new("v2".into(), 0);
|
|
|
|
// v1 -- v0 -- v2
|
|
graph.insert(v0.clone(), HashSet::from([v1.clone(), v2.clone()]));
|
|
graph.insert(v1.clone(), HashSet::from([v0.clone()]));
|
|
graph.insert(v2.clone(), HashSet::from([v0.clone()]));
|
|
graph
|
|
}
|
|
|
|
fn triangle_graph() -> InterferenceGraph {
|
|
let mut graph: InterferenceGraph = HashMap::new();
|
|
let v0 = IrVrVariableDescriptor::new("v0".into(), 0);
|
|
let v1 = IrVrVariableDescriptor::new("v1".into(), 0);
|
|
let v2 = IrVrVariableDescriptor::new("v2".into(), 0);
|
|
|
|
// triangle: each has two edges
|
|
// v0
|
|
// | \
|
|
// v1--v2
|
|
graph.insert(v0.clone(), HashSet::from([v1.clone(), v2.clone()]));
|
|
graph.insert(v1.clone(), HashSet::from([v0.clone(), v2.clone()]));
|
|
graph.insert(v2.clone(), HashSet::from([v0.clone(), v1.clone()]));
|
|
graph
|
|
}
|
|
|
|
fn get_vrs(graph: &InterferenceGraph) -> Vec<IrVrVariableDescriptor> {
|
|
let v0 = graph.keys().find(|k| k.name() == "v0").unwrap();
|
|
let v1 = graph.keys().find(|k| k.name() == "v1").unwrap();
|
|
let v2 = graph.keys().find(|k| k.name() == "v2").unwrap();
|
|
vec![v0.clone(), v1.clone(), v2.clone()]
|
|
}
|
|
|
|
#[test]
|
|
fn find_vr_lt_k_when_k_2() {
|
|
let graph = line_graph();
|
|
let found = find_vr_lt_k(&graph, 2);
|
|
assert!(found.is_some());
|
|
assert!(found.unwrap().name() == "v1" || found.unwrap().name() == "v2");
|
|
}
|
|
|
|
#[test]
|
|
fn find_vr_lt_k_when_k_1() {
|
|
let graph = line_graph();
|
|
let found = find_vr_lt_k(&graph, 1);
|
|
assert!(found.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn remove_edges_v0() {
|
|
let mut graph = line_graph();
|
|
let vrs = get_vrs(&graph);
|
|
|
|
let v0_outgoing = remove_vr_and_edges(&mut graph, &vrs[0]);
|
|
assert!(v0_outgoing.contains(&vrs[1]));
|
|
assert!(v0_outgoing.contains(&vrs[2]));
|
|
|
|
// check that incoming edges were removed
|
|
let v1_outgoing = graph.get(&vrs[1]).unwrap();
|
|
assert!(v1_outgoing.is_empty());
|
|
let v2_outgoing = graph.get(&vrs[2]).unwrap();
|
|
assert!(v2_outgoing.is_empty());
|
|
}
|
|
|
|
fn triangle_work_stack_k_2() -> Vec<WorkItem> {
|
|
let k = 2;
|
|
let mut graph = triangle_graph();
|
|
|
|
let mut work_stack = vec![];
|
|
|
|
// run three times, once for each register
|
|
work_stack.push(next_work_item(&mut graph, k));
|
|
work_stack.push(next_work_item(&mut graph, k));
|
|
work_stack.push(next_work_item(&mut graph, k));
|
|
|
|
work_stack
|
|
}
|
|
|
|
#[test]
|
|
fn next_work_item_k_2() {
|
|
let work_stack = triangle_work_stack_k_2();
|
|
|
|
// the actual edges may be different, depending on the underlying order in the sets
|
|
// (HashSet seems to use randomness in order)
|
|
// however, the bottommost item must be a spill, and the edge counts must be (from the
|
|
// bottom of the stack) 2-1-0
|
|
assert!(!work_stack[0].color);
|
|
assert_eq!(work_stack[0].edges.len(), 2);
|
|
assert_eq!(work_stack[1].edges.len(), 1);
|
|
assert_eq!(work_stack[2].edges.len(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn rebuild_graph_triangle_k_2() {
|
|
let mut work_stack = triangle_work_stack_k_2();
|
|
let mut rebuilt_graph: InterferenceGraph = HashMap::new();
|
|
|
|
// it should be possible to rebuild the graph from the stack, without yet worrying
|
|
// about spilling/etc.
|
|
while let Some(work_item) = work_stack.pop() {
|
|
rebuild_vr_and_edges(&mut rebuilt_graph, &work_item);
|
|
}
|
|
|
|
// we should have a triangle graph again
|
|
let vrs = get_vrs(&rebuilt_graph);
|
|
for vr in &vrs {
|
|
assert!(rebuilt_graph.contains_key(vr));
|
|
assert_eq!(rebuilt_graph.get(vr).unwrap().len(), 2);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn registers_and_spills_triangle_k_2() {
|
|
let mut graph = triangle_graph();
|
|
let (registers, spills) = registers_and_spills(&mut graph, 2);
|
|
// there should be one spill when k is 2
|
|
assert_eq!(registers.len(), 2);
|
|
assert_eq!(spills.len(), 1);
|
|
}
|
|
}
|