mir register allcoator

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"

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)]
 

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<u32, amd64::Register>,
         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, &reg_alloc, node as u32)?;
         }
         Ok(())
     }
@@ -518,6 +567,171 @@ impl Mir {
     }
 }
 
+impl Mir {
+    pub fn build_liveness(&self) -> HashMap<u32, amd64::Register> {
+        struct Interval {
+            start: u32,
+            end: u32,
+        }
+
+        let mut references = BTreeMap::<u32, Vec<u32>>::new();
+        let mut prefered_colors = BTreeMap::<u32, amd64::Register>::new();
+        use amd64::Register::*;
+        let mut param_registers = [rsi, rdi, rdx, rcx, r8, r9]
+            .into_iter()
+            .rev()
+            .collect::<Vec<_>>();
+
+        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::<BTreeMap<u32, Interval>>();
+
+        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::<u32, amd64::Register>::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::<BTreeSet<_>>();
+            let clique_preferred_colors = inference_graph
+                .neighbors(node.into())
+                .filter_map(|e| prefered_colors.get(&(e.index() as u32)).cloned())
+                .collect::<BTreeSet<_>>();
+
+            let color = prefered_colors
+                .get(&node)
+                .into_iter()
+                .chain(gprs.iter())
+                .filter(|reg| !clique_colors.contains(reg))
+                .find_or_first(|&reg| !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,

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)]