From 16399192b3d3b3e4eac2e19197c1ca0cdfe63b45 Mon Sep 17 00:00:00 2001 From: Janis Date: Mon, 26 Aug 2024 13:47:44 +0200 Subject: [PATCH] mir register allcoator --- Cargo.toml | 1 + src/lib.rs | 3 +- src/mir.rs | 216 +++++++++++++++++++++++++++++++++++++++++++- src/string_table.rs | 3 + 4 files changed, 221 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 39741c0..34257e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,5 +8,6 @@ ansi_term = "0.12.1" clap = "4.5.14" itertools = "0.13.0" log = "0.4.22" +petgraph = "0.6.5" thiserror = "1.0.63" unicode-xid = "0.2.4" diff --git a/src/lib.rs b/src/lib.rs index fa145b0..2679f9d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,7 +3,8 @@ iter_advance_by, box_into_inner, hash_extract_if, - bigint_helper_methods + bigint_helper_methods, + map_try_insert )] #![allow(unused_macros)] diff --git a/src/mir.rs b/src/mir.rs index 5b5bb37..cd6d47e 100644 --- a/src/mir.rs +++ b/src/mir.rs @@ -1,5 +1,10 @@ //! Machine-level Intermediate Representation +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; + +use itertools::Itertools; +use petgraph::visit::{IntoEdges, NodeRef}; + use crate::string_table::{Index as StringsIndex, StringTable}; #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)] @@ -150,6 +155,43 @@ pub enum Inst { //FunctionStart, } +impl Inst { + fn has_value(&self) -> bool { + // basically, when an arithmetic instruction has two immediates, then just replace it with a mov into the dst reg + match self { + Inst::Label + | Inst::ConstantBytes + | Inst::ConstantByte + | Inst::ConstantWord + | Inst::ConstantDWord + | Inst::ConstantQWord + | Inst::ExternRef + | Inst::Alloca + | Inst::Store(_) + | Inst::ReturnValue + | Inst::Return => false, + Inst::GetElementPtr(_) + | Inst::Load(_) + | Inst::Parameter(_) + | Inst::Add(_) + | Inst::Sub(_) + | Inst::Mul(_) + | Inst::MulSigned(_) + | Inst::Div(_) + | Inst::DivSigned(_) + | Inst::Rem(_) + | Inst::RemSigned(_) + | Inst::BitAnd(_) + | Inst::BitOr(_) + | Inst::BitXOr(_) + | Inst::Negate(_) + | Inst::ShiftLeft(_) + | Inst::ShiftRightSigned(_) + | Inst::ShiftRightUnsigned(_) => true, + } + } +} + #[derive(Clone, Copy)] pub union Data { none: (), @@ -397,11 +439,17 @@ impl Mir { &self, w: &mut W, strings: &StringTable, + reg_alloc: &HashMap, node: u32, ) -> core::fmt::Result { let idx = node as usize; let inst = self.nodes[idx]; let data = self.data[idx]; + + if let Some(reg) = reg_alloc.get(&node) { + write!(w, "({reg}) ")?; + } + match inst { Inst::Label => writeln!(w, "%{node} = {}", strings.get_str(data.as_index())), Inst::ConstantBytes => writeln!( @@ -507,8 +555,9 @@ impl Mir { w: &mut W, strings: &StringTable, ) -> core::fmt::Result { + let reg_alloc = self.build_liveness(); for node in 0..self.nodes.len() { - self.render_node(w, strings, node as u32)?; + self.render_node(w, strings, ®_alloc, node as u32)?; } Ok(()) } @@ -518,6 +567,171 @@ impl Mir { } } +impl Mir { + pub fn build_liveness(&self) -> HashMap { + struct Interval { + start: u32, + end: u32, + } + + let mut references = BTreeMap::>::new(); + let mut prefered_colors = BTreeMap::::new(); + use amd64::Register::*; + let mut param_registers = [rsi, rdi, rdx, rcx, r8, r9] + .into_iter() + .rev() + .collect::>(); + + for i in 0..self.nodes.len() { + let inst = self.nodes[i]; + let data = self.data[i]; + + let node = i as u32; + references.insert(node, Vec::new()); + + match inst { + Inst::Parameter(_) => { + if let Some(reg) = param_registers.pop() { + println!("prefering {reg} for param"); + prefered_colors.insert(node, reg); + } + } + Inst::Negate(_) | Inst::Load(_) => { + references.get_mut(&data.as_node()).unwrap().push(node); + } + // return is thru rax. + Inst::ReturnValue => { + let val = data.as_node(); + references.get_mut(&val).unwrap().push(node); + _ = prefered_colors.try_insert(val, amd64::Register::rax); + } + Inst::GetElementPtr(_) => { + let (src, _) = data.as_binary(); + references.get_mut(&src).unwrap().push(node); + } + // mul wants lhs (or rhs) to be rax, imul doesn't care as much. + // note that it also clobers rdx + Inst::Mul(ty) | Inst::MulSigned(ty) => { + let (lhs, rhs) = data.as_binary(); + references.get_mut(&lhs).unwrap().push(node); + references.get_mut(&rhs).unwrap().push(node); + + if !ty.is_floating() { + _ = prefered_colors.try_insert(lhs, amd64::Register::rax); + _ = prefered_colors.try_insert(rhs, amd64::Register::rax); + } + } + // mul wants lhs to be rax, imul can do either. + // note that it also clobers rdx + Inst::Div(_) | Inst::DivSigned(_) | Inst::Rem(_) | Inst::RemSigned(_) => { + let (lhs, rhs) = data.as_binary(); + references.get_mut(&lhs).unwrap().push(node); + references.get_mut(&rhs).unwrap().push(node); + _ = prefered_colors.try_insert(lhs, amd64::Register::rax); + } + // shr,shl,sar want the shift to be in cl. + Inst::ShiftLeft(_) | Inst::ShiftRightSigned(_) | Inst::ShiftRightUnsigned(_) => { + let (lhs, rhs) = data.as_binary(); + references.get_mut(&lhs).unwrap().push(node); + references.get_mut(&rhs).unwrap().push(node); + _ = prefered_colors.try_insert(rhs, amd64::Register::rcx); + } + // add,adc,sub,sbb,or,and,xor and mov don't care much about their source registers + Inst::Add(_) + | Inst::Sub(_) + | Inst::BitAnd(_) + | Inst::BitOr(_) + | Inst::BitXOr(_) + | Inst::Store(_) => { + let (lhs, rhs) = data.as_binary(); + references.get_mut(&lhs).unwrap().push(node); + references.get_mut(&rhs).unwrap().push(node); + } + _ => {} + } + } + + references.retain(|&node, refs| !refs.is_empty() && self.nodes[node as usize].has_value()); + prefered_colors.retain(|&node, _| self.nodes[node as usize].has_value()); + + let intervals = references + .iter() + .map(|(node, refs)| { + ( + *node, + Interval { + start: *node, + end: refs.last().cloned().unwrap(), + }, + ) + }) + .collect::>(); + + let mut edges = HashSet::<(u32, u32)>::new(); + for (&node, interval) in &intervals { + for (&other, _) in references.range(interval.start + 1..interval.end) { + edges.insert((node, other)); + } + } + let inference_graph = petgraph::graph::UnGraph::<(), ()>::from_edges(edges.into_iter()); + + let gprs = amd64::Register::gp_registers(); + let sses = amd64::Register::sse_registers(); + + let mut assigned_colors = HashMap::::new(); + for &node in references.keys().rev() { + if matches!(self.nodes[node as usize], Inst::Alloca) { + continue; + } + + let clique_colors = inference_graph + .neighbors(node.into()) + .filter_map(|e| assigned_colors.get(&(e.index() as u32)).cloned()) + .collect::>(); + let clique_preferred_colors = inference_graph + .neighbors(node.into()) + .filter_map(|e| prefered_colors.get(&(e.index() as u32)).cloned()) + .collect::>(); + + let color = prefered_colors + .get(&node) + .into_iter() + .chain(gprs.iter()) + .filter(|reg| !clique_colors.contains(reg)) + .find_or_first(|®| !clique_preferred_colors.contains(reg)) + .cloned() + .expect("ran out of registers"); + + println!("%{node} wants {:?}\n\tclique: {clique_colors:?}\n\tclique prefs: {clique_preferred_colors:?}\n\t-> {color}", prefered_colors.get(&node)); + + prefered_colors.remove(&node); + + assigned_colors.insert(node, color); + } + + // eprintln!( + // "Inference Graph:\n{:?}", + // petgraph::dot::Dot::with_attr_getters( + // &inference_graph, + // &[ + // petgraph::dot::Config::EdgeNoLabel, + // petgraph::dot::Config::NodeNoLabel, + // ], + // &Box::new(|_, _| { "".to_owned() }), + // &Box::new(|_, node: (petgraph::graph::NodeIndex, &())| { + // let (node, _) = node; + // if let Some(reg) = assigned_colors.get(&(node.index() as u32)) { + // format!("label=\"{} ({})\"", node.index(), reg) + // } else { + // format!("label=\"{}\"", node.index()) + // } + // }), + // ), + // ); + + assigned_colors + } +} pub struct DisplayMir<'a, 'b> { mir: &'a Mir, strings: &'b StringTable, diff --git a/src/string_table.rs b/src/string_table.rs index ee91107..bc837cd 100644 --- a/src/string_table.rs +++ b/src/string_table.rs @@ -17,6 +17,9 @@ impl Index { pub fn new(start: u32, end: u32) -> Self { Self { start, end } } + pub fn none() -> Self { + Self::new(u32::MAX, u32::MAX) + } } #[derive(Clone)]