440 lines
14 KiB
Rust
440 lines
14 KiB
Rust
use crate::ir::ir_block::IrBlock;
|
|
use crate::ir::ir_variable::IrVariableId;
|
|
use crate::ir::variable_locations::VariableLocations;
|
|
use crate::offset_counter::OffsetCounter;
|
|
use dvm_lib::instruction::{Register, StackFrameOffset};
|
|
use std::collections::{HashMap, HashSet};
|
|
|
|
pub type RegisterAssignment = Register;
|
|
pub type InterferenceGraph = HashMap<IrVariableId, HashSet<IrVariableId>>;
|
|
pub type LivenessSets = Vec<HashSet<IrVariableId>>;
|
|
|
|
pub fn block_live_in_live_out(ir_block: &IrBlock) -> (LivenessSets, LivenessSets) {
|
|
// init
|
|
let n_statements = ir_block.statements().len();
|
|
let mut live_in: LivenessSets = vec![HashSet::new(); n_statements];
|
|
let mut live_out: LivenessSets = vec![HashSet::new(); n_statements];
|
|
|
|
loop {
|
|
let mut did_work = false;
|
|
// Go backwards for efficiency
|
|
for (statement_index, ir_statement) in ir_block.statements().iter().enumerate().rev() {
|
|
// out (union of successors ins)
|
|
// for now, a statement can only have one successor
|
|
// this will need to be updated when we add jumps
|
|
let statement_live_out = &mut live_out[statement_index];
|
|
let successor_live_in = &live_in[statement_index + 1];
|
|
for ir_variable_id in successor_live_in {
|
|
if statement_live_out.insert(*ir_variable_id) {
|
|
did_work = true;
|
|
}
|
|
}
|
|
|
|
// in: use(s) U ( out(s) - def(s) )
|
|
let use_s = ir_statement
|
|
.vr_uses()
|
|
.iter()
|
|
.map(|u| *u)
|
|
.collect::<HashSet<_>>();
|
|
let out_s = &live_out[statement_index];
|
|
let def_s = ir_statement
|
|
.vr_definitions()
|
|
.iter()
|
|
.map(|d| *d)
|
|
.collect::<HashSet<_>>();
|
|
let rhs = out_s - &def_s;
|
|
let new_ins = use_s.union(&rhs).map(|v| *v).collect::<HashSet<_>>();
|
|
|
|
// add new ins to statement's live in
|
|
let statement_live_in = &mut live_in[statement_index];
|
|
for new_in in new_ins {
|
|
if statement_live_in.insert(new_in) {
|
|
// if we added a previously unadded variable, we did work
|
|
did_work = true;
|
|
}
|
|
}
|
|
}
|
|
// break only if we did nothing!
|
|
if !did_work {
|
|
break;
|
|
}
|
|
}
|
|
(live_in, live_out)
|
|
}
|
|
|
|
pub fn block_interference_graph(
|
|
ir_block: &IrBlock,
|
|
spilled: &HashSet<IrVariableId>,
|
|
) -> InterferenceGraph {
|
|
// create a set of all variables used in the block
|
|
let mut all_vr_variables: HashSet<IrVariableId> = HashSet::new();
|
|
for statement in ir_block.statements() {
|
|
all_vr_variables.extend(statement.vr_definitions());
|
|
all_vr_variables.extend(statement.vr_uses());
|
|
}
|
|
|
|
// create graph and init all variables' outgoing-edge sets
|
|
let mut graph: InterferenceGraph = HashMap::new();
|
|
for variable in all_vr_variables {
|
|
graph.insert(variable, HashSet::new());
|
|
}
|
|
|
|
let (_, live_out) = block_live_in_live_out(ir_block);
|
|
|
|
for (statement_index, ir_statement) in ir_block.statements().iter().enumerate() {
|
|
let statement_live_out = &live_out[statement_index];
|
|
for definition_vr_variable in ir_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 {
|
|
// add edges to both sets (two-way)
|
|
let definition_edges = graph.get_mut(&definition_vr_variable).unwrap();
|
|
definition_edges.insert(*live_out_variable);
|
|
let live_out_variable_edges = graph.get_mut(live_out_variable).unwrap();
|
|
live_out_variable_edges.insert(definition_vr_variable);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
graph
|
|
}
|
|
|
|
pub fn block_assign_registers(ir_block: &IrBlock, register_count: usize) -> VariableLocations {
|
|
let mut spilled: HashSet<IrVariableId> = HashSet::new();
|
|
loop {
|
|
let mut interference_graph = block_interference_graph(ir_block, &spilled);
|
|
let (registers, new_spills) = registers_and_spills(&mut interference_graph, register_count);
|
|
|
|
if spilled != new_spills {
|
|
spilled = new_spills; // todo: figure out if this works algorithmically
|
|
} else {
|
|
let mut spills_to_stack_offsets: HashMap<IrVariableId, StackFrameOffset> =
|
|
HashMap::new();
|
|
let mut offset_counter = 0isize;
|
|
for spill in spilled {
|
|
spills_to_stack_offsets.insert(spill, offset_counter);
|
|
offset_counter += 1;
|
|
}
|
|
return VariableLocations::new(registers, spills_to_stack_offsets);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub trait HasVrUsers {
|
|
fn vr_users(&self) -> Vec<&dyn VrUser>;
|
|
fn vr_users_mut(&mut self) -> Vec<&mut dyn VrUser>;
|
|
}
|
|
|
|
pub trait VrUser {
|
|
fn vr_definitions(&self) -> HashSet<IrVariableId> {
|
|
HashSet::new()
|
|
}
|
|
|
|
#[deprecated]
|
|
fn vr_uses(&self) -> HashSet<IrVariableId>;
|
|
|
|
#[deprecated]
|
|
fn propagate_spills(&mut self, _spills: &HashSet<IrVariableId>) {
|
|
// default no-op
|
|
}
|
|
|
|
#[deprecated]
|
|
fn propagate_register_assignments(
|
|
&mut self,
|
|
_assignments: &HashMap<IrVariableId, RegisterAssignment>,
|
|
) {
|
|
// default no-op
|
|
}
|
|
|
|
#[deprecated]
|
|
fn propagate_stack_offsets(&mut self, _counter: &mut OffsetCounter) {
|
|
// default no-op
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct WorkItem {
|
|
vr: IrVariableId,
|
|
edges: HashSet<IrVariableId>,
|
|
color: bool,
|
|
}
|
|
|
|
pub fn registers_and_spills(
|
|
interference_graph: &mut InterferenceGraph,
|
|
k: usize,
|
|
) -> (
|
|
HashMap<IrVariableId, RegisterAssignment>,
|
|
HashSet<IrVariableId>,
|
|
) {
|
|
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<IrVariableId, RegisterAssignment> = HashMap::new();
|
|
let mut spills: HashSet<IrVariableId> = 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<IrVariableId, RegisterAssignment>,
|
|
) {
|
|
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<IrVariableId> {
|
|
interference_graph.iter().find_map(
|
|
|(vr, neighbors)| {
|
|
if neighbors.len() < k { Some(*vr) } else { None }
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Returns the (removed) outgoing edges for the given vr
|
|
fn remove_vr_and_edges(
|
|
interference_graph: &mut InterferenceGraph,
|
|
vr: &IrVariableId,
|
|
) -> HashSet<IrVariableId> {
|
|
// 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<IrVariableId, 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();
|
|
|
|
// v1 -- v0 -- v2
|
|
graph.insert(0, HashSet::from([1, 2]));
|
|
graph.insert(1, HashSet::from([0]));
|
|
graph.insert(2, HashSet::from([0]));
|
|
graph
|
|
}
|
|
|
|
fn triangle_graph() -> InterferenceGraph {
|
|
let mut graph: InterferenceGraph = HashMap::new();
|
|
|
|
// triangle: each has two edges
|
|
// v0
|
|
// | \
|
|
// v1--v2
|
|
graph.insert(0, HashSet::from([1, 2]));
|
|
graph.insert(1, HashSet::from([0, 2]));
|
|
graph.insert(2, HashSet::from([0, 1]));
|
|
graph
|
|
}
|
|
|
|
fn get_vrs() -> Vec<IrVariableId> {
|
|
vec![0, 1, 2]
|
|
}
|
|
|
|
#[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() == 1 || found.unwrap() == 2);
|
|
}
|
|
|
|
#[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();
|
|
|
|
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();
|
|
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);
|
|
}
|
|
}
|