diff --git a/src/asm/amd64.rs b/src/asm/amd64.rs index 1baca03..90c014c 100644 --- a/src/asm/amd64.rs +++ b/src/asm/amd64.rs @@ -314,6 +314,11 @@ impl Register { [r12, r13, r14, r15, rbx] }; + pub const SYSV_SCRATCH_GPRS: [Register; 8] = { + use Register::*; + [rax, rsi, rdx, rcx, r8, r9, r10, r11] + }; + pub const GPR: [Register; 14] = { use Register::*; [ diff --git a/src/bin/main.rs b/src/bin/main.rs index dcd6f4d..371047d 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -1,6 +1,6 @@ use std::{io::Read, path::PathBuf}; -use clap::{Arg, Command}; +use clap::Command; use compiler::{ lexer::Tokenizer, parser::Tree, diff --git a/src/mir.rs b/src/mir.rs index f2ea2f0..5d950c0 100644 --- a/src/mir.rs +++ b/src/mir.rs @@ -2,11 +2,14 @@ use std::collections::btree_map::Entry; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::fmt::Display; use std::hash::{Hash, Hasher}; +use std::u32; use itertools::Itertools; +use liveness::LivenessBuilder; -use crate::asm::amd64::{self, Register}; +use crate::asm::amd64::Register; use crate::ast::IntegralType; use crate::string_table::{Index as StringsIndex, StringTable}; @@ -155,6 +158,12 @@ pub enum Inst { /// lhs, rhs Div(Type), /// lhs, rhs + MulSSE(Type), + /// lhs, rhs + DivSSE(Type), + /// lhs, rhs + RemFP(Type), + /// lhs, rhs DivSigned(Type), /// lhs, rhs Rem(Type), @@ -215,6 +224,9 @@ impl Inst { | Inst::DivSigned(ty) | Inst::Rem(ty) | Inst::RemSigned(ty) + | Inst::MulSSE(ty) + | Inst::DivSSE(ty) + | Inst::RemFP(ty) | Inst::BitAnd(ty) | Inst::BitOr(ty) | Inst::BitXOr(ty) @@ -227,6 +239,7 @@ impl Inst { Inst::IsZero(_) => Some(Type::Byte), } } + #[allow(dead_code)] fn result_is_register(&self) -> bool { // basically, when an arithmetic instruction has two immediates, then just replace it with a mov into the dst reg // TODO: need to account for spilled values eventually; probably move this to `Mir`. @@ -256,6 +269,9 @@ impl Inst { | Inst::DivSigned(_) | Inst::Rem(_) | Inst::RemSigned(_) + | Inst::MulSSE(_) + | Inst::DivSSE(_) + | Inst::RemFP(_) | Inst::BitAnd(_) | Inst::BitOr(_) | Inst::BitXOr(_) @@ -347,6 +363,13 @@ impl Data { pub fn as_node(self) -> u32 { unsafe { self.node } } + pub fn as_noderef(self) -> NodeRef { + NodeRef(self.as_node()) + } + pub fn as_binary_noderefs(self) -> (NodeRef, NodeRef) { + let (a, b) = self.as_binary(); + (a.into(), b.into()) + } pub fn as_binary(self) -> (u32, u32) { unsafe { self.binary } } @@ -397,6 +420,9 @@ impl OperandKinds { fn imul() -> Self { Self::RegImm | Self::MemImm | Self::RegMem | Self::RegReg } + fn any_mem_reg() -> Self { + Self::MemMem | Self::MemReg | Self::RegMem | Self::RegReg + } /// works for: div,idiv,mul fn mul() -> Self { Self::RegMem | Self::RegReg @@ -407,7 +433,7 @@ impl OperandKinds { } /// works for: shl,shr,sar,sal,test fn shift() -> Self { - Self::RegReg | Self::RegImm | Self::MemImm | Self::MemReg + Self::RegReg | Self::RegImm // | Self::MemImm | Self::MemReg } const fn to_rhs_binop(self) -> BinaryOperandFlags { let reg = if self.intersects(Self::RegReg.union(Self::MemReg)) { @@ -637,6 +663,64 @@ pub struct Mir { pub data: Vec, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +#[repr(transparent)] +pub struct NodeRef(pub u32); + +impl From for NodeRef { + fn from(value: u32) -> Self { + Self(value) + } +} + +impl NodeRef { + fn index(self) -> usize { + self.0 as usize + } + + fn into_reference_range(self) -> std::ops::Range<(NodeRef, NodeRef)> { + (self, Self::MIN)..(self, Self::MAX) + } + #[allow(dead_code)] + fn inclusive_start(self) -> (Self, Self) { + (self, Self::MIN) + } + fn exclusive_start(self) -> (Self, Self) { + (self, Self::MAX) + } + fn inclusive_end(self) -> (Self, Self) { + (self, Self::MAX) + } + #[allow(dead_code)] + fn exclusive_end(self) -> (Self, Self) { + (self, Self::MIN) + } + + /// invalid pseudo-handle to the past-the-end node. + const MAX: Self = NodeRef(u32::MAX); + const MIN: Self = NodeRef(0); +} + +impl Display for NodeRef { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "%{}", self.0) + } +} + +impl Mir { + fn get_node(&self, node: NodeRef) -> (Inst, Data) { + (self.nodes[node.index()], self.data[node.index()]) + } + #[allow(dead_code)] + fn get_node_mut(&mut self, node: NodeRef) -> (&mut Inst, &mut Data) { + (&mut self.nodes[node.index()], &mut self.data[node.index()]) + } + + fn indices(&self) -> impl Iterator { + (0..self.nodes.len() as u32).map(|n| NodeRef::from(n)) + } +} + impl Mir { pub fn new(name: StringsIndex) -> Mir { Self { @@ -682,14 +766,32 @@ impl Mir { | Inst::DivSigned(_) | Inst::Rem(_) | Inst::RemSigned(_) + | Inst::MulSSE(_) + | Inst::DivSSE(_) + | Inst::RemFP(_) | Inst::BitAnd(_) | Inst::BitOr(_) | Inst::BitXOr(_) | Inst::Negate(_) + | Inst::SignExtend(_) + | Inst::ZeroExtend(_) + | Inst::IsZero(_) | Inst::ShiftLeft(_) | Inst::ShiftRightSigned(_) | Inst::ShiftRightUnsigned(_) => true, - _ => false, + Inst::Label + | Inst::ConstantBytes + | Inst::ConstantByte + | Inst::ConstantWord + | Inst::ConstantDWord + | Inst::ConstantQWord + | Inst::ConstantSinglePrecision + | Inst::ConstantDoublePrecision + | Inst::ExternRef + | Inst::Alloca + | Inst::Store(_) + | Inst::ReturnValue(_) + | Inst::Return => false, } } @@ -855,7 +957,7 @@ impl Mir { } pub fn gen_mul_sse(&mut self, ty: Type, lhs: u32, rhs: u32) -> u32 { let (lhs, rhs) = BinaryOperands::new_sse(self).wrangle(lhs, ty, rhs, ty); - self.push(Inst::Mul(ty), Data::binary(lhs, rhs)) + self.push(Inst::MulSSE(ty), Data::binary(lhs, rhs)) } pub fn gen_mul_unsigned(&mut self, ty: Type, lhs: u32, rhs: u32) -> u32 { let (lhs, rhs) = BinaryOperands::new_mul(self).wrangle(lhs, ty, rhs, ty); @@ -876,7 +978,7 @@ impl Mir { } pub fn gen_div_sse(&mut self, ty: Type, lhs: u32, rhs: u32) -> u32 { let (lhs, rhs) = BinaryOperands::new_sse(self).wrangle(lhs, ty, rhs, ty); - self.push(Inst::Div(ty), Data::binary(lhs, rhs)) + self.push(Inst::DivSSE(ty), Data::binary(lhs, rhs)) } pub fn gen_div_unsigned(&mut self, ty: Type, lhs: u32, rhs: u32) -> u32 { let (lhs, rhs) = BinaryOperands::new_div_idiv_rem_irem(self).wrangle(lhs, ty, rhs, ty); @@ -896,8 +998,9 @@ impl Mir { } } pub fn gen_rem_fp(&mut self, ty: Type, lhs: u32, rhs: u32) -> u32 { - _ = (ty, lhs, rhs); - todo!() + let (lhs, rhs) = + BinaryOperands::new(self, false, OperandKinds::any_mem_reg()).wrangle(lhs, ty, rhs, ty); + self.push(Inst::RemFP(ty), Data::binary(lhs, rhs)) } pub fn gen_rem_unsigned(&mut self, ty: Type, lhs: u32, rhs: u32) -> u32 { let (lhs, rhs) = BinaryOperands::new_div_idiv_rem_irem(self).wrangle(lhs, ty, rhs, ty); @@ -953,14 +1056,14 @@ impl Mir { &self, w: &mut W, strings: &StringTable, - reg_alloc: &BTreeMap, + liveness: &Liveness, 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) { + if let Some(reg) = liveness.get_register(node.into()) { write!(w, "({reg}) ")?; } @@ -1040,6 +1143,18 @@ impl Mir { let (lhs, rhs) = data.as_binary(); writeln!(w, "%{node} = signed rem {ty} %{lhs}, {ty} %{rhs}") } + Inst::MulSSE(ty) => { + let (lhs, rhs) = data.as_binary(); + writeln!(w, "%{node} = mulss {ty} %{lhs}, {ty} %{rhs}") + } + Inst::DivSSE(ty) => { + let (lhs, rhs) = data.as_binary(); + writeln!(w, "%{node} = divss {ty} %{lhs}, {ty} %{rhs}") + } + Inst::RemFP(ty) => { + let (lhs, rhs) = data.as_binary(); + writeln!(w, "%{node} = fp rem {ty} %{lhs}, {ty} %{rhs}") + } Inst::BitAnd(ty) => { let (lhs, rhs) = data.as_binary(); writeln!(w, "%{node} = bitand {ty} %{lhs}, {ty} %{rhs}") @@ -1107,309 +1222,457 @@ impl Mir { } } -impl Mir { - pub fn build_liveness(&self) -> BTreeMap { - struct Interval { - start: u32, - end: u32, +pub use liveness::Liveness; + +pub mod liveness { + use super::*; + pub struct Liveness { + register_map: BTreeMap, + inference_graph: petgraph::graph::UnGraph<(), ()>, + } + + impl Liveness { + pub fn get_register(&self, node: NodeRef) -> Option { + self.register_map.get(&node).cloned() + } + pub fn dirty_registers(&self) -> std::collections::btree_map::Values { + self.register_map.values() + } + pub fn get_scratch_register_at_node(&self, node: NodeRef) -> Option { + let dirty = self + .inference_graph + .neighbors(node.0.into()) + .filter_map(|n| self.register_map.get(&NodeRef(n.index() as u32))) + .cloned() + .collect::>(); + + Register::SYSV_SCRATCH_GPRS + .into_iter() + .filter(|c| !dirty.contains(c)) + .next() + } + pub fn is_register_in_use_at_node(&self, node: NodeRef, reg: Register) -> bool { + self.inference_graph + .neighbors(node.0.into()) + .filter_map(|n| self.register_map.get(&NodeRef(n.index() as u32))) + .find(|&&r| r.parent_reg() == reg.parent_reg()) + .is_some() + } + } + + pub struct LivenessBuilder<'a> { + mir: &'a Mir, + // tree of (node, referenced_by) pairs. + // a full range for each node is (node, NodeRef::MIN)..(node, NodeRef::MAX) + // QUESTION: do I want to treat every interval in this tree a separate node to color? + // or just the (node, NodeRef::MIN)..(node, NodeRef::MAX) range? + // references: BTreeSet<(NodeRef, NodeRef)>, + inference_graph: petgraph::graph::UnGraph<(), ()>, + // list of preferred colors by nodes, either because they output to some + // register like mul/div or because the write to one of their inputs. + // interesting to consider optimisations like i >>= s being turned into shl + // mem, cl, while i >> s is turned into shl reg, cl ... + // preferred_color: BTreeMap, + } + + #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] + enum Color { + #[default] + Unassigned, + Tentative(Register), + Final(Register), + } + + impl Color { + fn color(self) -> Option { + match self { + Color::Unassigned => None, + Color::Tentative(color) | Color::Final(color) => Some(color), + } + } + } + + impl<'a> LivenessBuilder<'a> { + pub fn new(mir: &'a Mir) -> Self { + let references = Self::build_references(mir); + + let intervals = mir.indices().filter_map(|node| { + let interval = references + .range(node.into_reference_range()) + .map(|&(_, to)| to) + .reduce(|acc, to| acc.max(to)); + + interval.map(|max| (node, max)) + }); + + eprintln!("intervals: ["); + + let mut edges = HashSet::<(u32, u32)>::new(); + for (from, to) in intervals { + eprint!("({from}..{to}), "); + for &(other, _) in references.range(from.exclusive_start()..to.inclusive_end()) { + edges.insert((from.0, other.0)); + } + } + eprintln!("]"); + + let inference_graph = petgraph::graph::UnGraph::<(), ()>::from_edges(edges.into_iter()); + + Self { + mir, + inference_graph, + // references, + } } - let mut references = BTreeMap::>::new(); - let mut preferred_colors = BTreeMap::::new(); - use amd64::Register::*; - let mut in_colors = [r9, r8, rcx, rdx, rdi, rsi].to_vec(); - let mut in_colors_sse = [xmm7, xmm6, xmm5, xmm4, xmm3, xmm2, xmm1, xmm0].to_vec(); - let mut inouts = Vec::::new(); + pub fn build(self) -> Liveness { + let preferred = self.pre_color_colors(); + let Self { + mir, + inference_graph, + .. + } = self; - for i in 0..self.nodes.len() { - let inst = self.nodes[i]; - let data = self.data[i]; + let mut colorizer = Colorizer { + mir, + preferred, + graph: inference_graph, + colors: BTreeMap::new(), + }; - let node = i as u32; - references.insert(node, Vec::new()); + // // prepass: assign preferred colours for in/out values and specific + // // instructions like mul/div which require one operand to be in rax + colorizer.prepass(); + + // for &node in references.keys().rev() { + // if !self.nodes[node as usize].result_is_register() { + // continue; + // } + // colorizer.color_node(node.into()); + // } + + let register_map = colorizer.colorise(); + + let Colorizer { graph, .. } = colorizer; + + Liveness { + register_map, + inference_graph: graph, + } + } + + fn pre_color_colors(&self) -> BTreeMap { + /* do i want to find colors from the bottom up? + * - on one hand, consumers with in/outs like mul/div are below and want their inputs in some specific register + * - on the other hand, producers like mul/div are above and want to dictate their output registers + * i think the answer is coloring mul/div first, specifically, then collapsing from there. + */ + let mut want_colors = BTreeMap::::new(); + + use Register::*; + let mut in_colors = [r9, r8, rcx, rdx, rdi, rsi].to_vec(); + let mut in_colors_sse = [xmm7, xmm6, xmm5, xmm4, xmm3, xmm2, xmm1, xmm0].to_vec(); + + for node in self.mir.indices() { + let want_color = match self.mir.get_node(node).0 { + // + Inst::Parameter(ty) => { + if ty.is_floating() { + in_colors_sse.pop() + } else { + in_colors.pop() + } + } + Inst::ShiftLeft(_) + | Inst::ShiftRightSigned(_) + | Inst::ShiftRightUnsigned(_) => Some(rcx), + Inst::Mul(_) | Inst::Div(_) | Inst::DivSigned(_) => Some(rax), + Inst::Rem(_) | Inst::RemSigned(_) => Some(rdx), + Inst::ReturnValue(ty) => { + if ty.is_floating() { + Some(xmm0) + } else { + Some(rax) + } + } + _ => { + if let Some(dst) = self.mir.dst_node(node) { + // check if there is interference + if self + .inference_graph + .find_edge(dst.0.into(), node.0.into()) + .is_none() + { + // want the same color as our dst node to avoid copying + want_colors.get(&dst).cloned() + } else { + None + } + } else { + None + } + } + }; + + if let Some(color) = want_color { + want_colors.insert(node, color); + } + } + + want_colors + } + + fn build_references(mir: &Mir) -> BTreeSet<(NodeRef, NodeRef)> { + let mut references = BTreeSet::new(); + for node in mir.indices() { + Self::reference_node_operands(mir, &mut references, node); + } + + references + } + + fn reference_node_operands( + mir: &Mir, + references: &mut BTreeSet<(NodeRef, NodeRef)>, + node: NodeRef, + ) { + let (inst, data) = mir.get_node(node); match inst { - Inst::Parameter(ty) => { - if let Some(reg) = if ty.is_floating() { - in_colors_sse.pop() - } else { - in_colors.pop() - } { - preferred_colors.insert(node, reg); - }; - inouts.push(node); - } - // return is thru rax. - Inst::ReturnValue(ty) => { - let val = data.as_node(); - inouts.push(val); - references.get_mut(&val).unwrap().push(node); - if ty.is_floating() { - _ = preferred_colors.try_insert(val, Register::xmm0); - } else { - _ = preferred_colors.try_insert(val, Register::rax); - } - } - Inst::SignExtend(_) + Inst::ReturnValue(_) + | Inst::SignExtend(_) | Inst::ZeroExtend(_) - | Inst::IsZero(_) | Inst::Negate(_) - | Inst::Load(_) => { - references.get_mut(&data.as_node()).unwrap().push(node); + | Inst::IsZero(_) + | Inst::Load(_) + | Inst::LoadRegister(_) => { + references.insert((data.as_noderef(), node)); } - 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() { - _ = preferred_colors.try_insert(lhs, Register::rax); - } - } - // div wants lhs to be rax, idiv can do either. - // note that it also clobers rdx - Inst::Div(ty) | Inst::DivSigned(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() { - _ = preferred_colors.try_insert(lhs, Register::rax); - } - } - // div wants lhs to be rax, idiv can do either. - // note that it also clobers rax - Inst::Rem(ty) | Inst::RemSigned(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() { - _ = preferred_colors.try_insert(lhs, Register::rdx); - } - } - // 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); - _ = preferred_colors.try_insert(rhs, Register::rcx); - } - // add,adc,sub,sbb,or,and,xor and mov don't care much about their source registers - Inst::Add(_) + Inst::Store(_) + | Inst::Add(_) | Inst::Sub(_) + | Inst::Mul(_) + | Inst::MulSigned(_) + | Inst::Div(_) + | Inst::DivSigned(_) + | Inst::Rem(_) + | Inst::RemSigned(_) + | Inst::MulSSE(_) + | Inst::DivSSE(_) + | Inst::RemFP(_) | 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); + | Inst::ShiftLeft(_) + | Inst::ShiftRightSigned(_) + | Inst::ShiftRightUnsigned(_) => { + let (lhs, rhs) = data.as_binary_noderefs(); + references.insert((lhs, node)); + references.insert((rhs, node)); } - _ => {} - } - } - - references.retain(|&node, refs| { - !refs.is_empty() && self.nodes[node as usize].result_is_register() - }); - preferred_colors.retain(|&node, _| self.nodes[node as usize].result_is_register()); - - 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()); - - #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] - enum Color { - #[default] - Unassigned, - Tentative(Register), - Final(Register), - } - - impl Color { - fn color(self) -> Option { - match self { - Color::Unassigned => None, - Color::Tentative(color) | Color::Final(color) => Some(color), + Inst::GetElementPtr(_) => { + let (src, _) = data.as_binary_noderefs(); + references.insert((src, node)); } + // these instructions have no inputs + // don't want a wildcard match here to make sure new instructions + // are handled here when they are added. + Inst::Return + | Inst::Parameter(_) + | Inst::Label + | Inst::ConstantBytes + | Inst::ConstantByte + | Inst::ConstantWord + | Inst::ConstantDWord + | Inst::ConstantQWord + | Inst::ConstantSinglePrecision + | Inst::ConstantDoublePrecision + | Inst::ExternRef + | Inst::Alloca => {} + } + } + } + + struct Colorizer<'a> { + mir: &'a Mir, + graph: petgraph::graph::UnGraph<(), ()>, + colors: BTreeMap, + preferred: BTreeMap, + } + + impl<'a> Colorizer<'a> { + fn node_colors(&self, node: petgraph::graph::NodeIndex) -> &[Register] { + let colors = if self.mir.nodes[node.index()] + .value_type() + .map(|t| t.is_floating()) + == Some(true) + { + &Register::SSE[..] + } else { + &Register::GPR[..] + }; + colors + } + + fn prepass(&mut self) { + // parameters are first in line + let keys = self.preferred.keys().cloned().collect::>(); + for node in keys { + self.precolor_node(node.0.into()); } } - struct Colorizer<'a> { - mir: &'a Mir, - graph: petgraph::graph::UnGraph<(), ()>, - colors: BTreeMap, - preferred: BTreeMap, - } - - impl<'a> Colorizer<'a> { - fn node_colors(&self, node: petgraph::graph::NodeIndex) -> &[Register] { - let colors = if self.mir.nodes[node.index()] - .value_type() - .map(|t| t.is_floating()) - == Some(true) - { - &Register::SSE[..] - } else { - &Register::GPR[..] - }; - colors - } - - fn prepass>(&mut self, inouts: I) { - for node in inouts.into_iter() { - self.precolor_node(node.into()); - } - - let keys = self.preferred.keys().cloned().collect::>(); - for node in keys { - self.precolor_node(node.into()); - } - } - - fn precolor_node(&mut self, node: petgraph::graph::NodeIndex) { - // prepass: assign preferred colours for in/out values and specific - // instructions like mul/div which require one operand to be in rax - let node_u32 = node.index() as u32; - // only apply color here if we have a preference - if let Some(preferred_color) = self.preferred.remove(&node_u32) { - let mut clique_colors = self - .graph - .neighbors(node) - .filter_map(|n| self.colors.get(&(n.index() as u32)).cloned()); - - if clique_colors - .find(|color| color.color() == Some(preferred_color)) - .is_none() - { - self.colors - .insert(node_u32, Color::Tentative(preferred_color)); - } - }; - // .chain(self.node_colors(node).into_iter().cloned()); - } - - fn color_node(&mut self, node: petgraph::graph::NodeIndex) { - // final pass: - // look at clique colors and prefer to steal colors from - // tentatively colored nodes. this results in preferential - // coloring depending on the order of the prepass. - let node_u32 = node.index() as u32; - let clique_colors = self + fn precolor_node(&mut self, node: petgraph::graph::NodeIndex) { + // prepass: assign preferred colours for in/out values and specific + // instructions like mul/div which require one operand to be in rax + let node_u32 = node.index() as u32; + let noderef = NodeRef(node_u32); + // only apply color here if we have a preference + if let Some(preferred_color) = self.preferred.remove(&noderef) { + let mut clique_colors = self .graph .neighbors(node) - .filter_map(|n| self.colors.get(&(n.index() as u32)).cloned()) - .collect::>(); + .filter_map(|n| self.colors.get(&NodeRef(n.index() as u32)).cloned()); - let colors = self - .node_colors(node) - .into_iter() - .filter(|&&r| !clique_colors.contains(&Color::Final(r))) - .cloned() - .collect::>(); - - // eprintln!("coloring %{node_u32}:"); - // eprintln!("\twants: {:?}", self.colors.get(&node_u32)); - // eprintln!("\tclique: {clique_colors:?}"); - // eprintln!("\tcandidates: {colors:?}"); - - match self.colors.entry(node_u32) { - Entry::Vacant(v) => { - // here we want to first check clique_colors with tentative coloring. - let color = colors - .into_iter() - .find_or_first(|&c| !clique_colors.contains(&Color::Tentative(c))) - .expect("ran out of registers :("); - v.insert(Color::Final(color)); - } - Entry::Occupied(mut e) => { - // we prefer to steal - variant!(e.get() => &Color::Tentative(reg)); - let color = colors - .into_iter() - .find_or_first(|&c| c == reg) - .expect("ran out of registers :("); - e.insert(Color::Final(color)); - } + if clique_colors + .find(|color| color.color() == Some(preferred_color)) + .is_none() + { + self.colors + .insert(noderef, Color::Tentative(preferred_color)); } - } + }; + // .chain(self.node_colors(node).into_iter().cloned()); + } - fn finalise(self) -> BTreeMap { - self.colors - .into_iter() - .filter_map(|(node, c)| match c { - Color::Final(reg) => Some((node, reg)), - _ => None, - }) - .collect() + fn color_node(&mut self, node: petgraph::graph::NodeIndex) { + // final pass: + // look at clique colors and prefer to steal colors from + // tentatively colored nodes. this results in preferential + // coloring depending on the order of the prepass. + let node_u32 = node.index() as u32; + let noderef = NodeRef(node_u32); + let clique_colors = self + .graph + .neighbors(node) + .filter_map(|n| self.colors.get(&NodeRef(n.index() as u32)).cloned()) + .collect::>(); + + let colors = self + .node_colors(node) + .into_iter() + .filter(|&&r| !clique_colors.contains(&Color::Final(r))) + .cloned() + .collect::>(); + + // eprintln!("coloring %{node_u32}:"); + // eprintln!("\twants: {:?}", self.colors.get(&node_u32)); + // eprintln!("\tclique: {clique_colors:?}"); + // eprintln!("\tcandidates: {colors:?}"); + + let color = match self.colors.entry(noderef) { + Entry::Vacant(v) => { + // here we want to first check clique_colors with tentative coloring. + let color = colors + .into_iter() + .find_or_first(|&c| !clique_colors.contains(&Color::Tentative(c))) + .expect("ran out of registers :("); + v.insert(Color::Final(color)); + color + } + Entry::Occupied(mut e) => { + // we prefer to steal + variant!(e.get() => &Color::Tentative(reg)); + let color = colors + .into_iter() + .find_or_first(|&c| c == reg) + .expect("ran out of registers :("); + e.insert(Color::Final(color)); + color + } + }; + + // if this node has a dst node (an operand that is both source and + // destination), give it our tentative color + if let Some(dst) = self.mir.dst_node(noderef) { + _ = self.colors.try_insert(dst, Color::Tentative(color)); } } - let mut colorizer = Colorizer { - mir: self, - graph: inference_graph, - preferred: preferred_colors, - colors: BTreeMap::new(), - }; - - // prepass: assign preferred colours for in/out values and specific - // instructions like mul/div which require one operand to be in rax - colorizer.prepass(inouts); - - for &node in references.keys().rev() { - if !self.nodes[node as usize].result_is_register() { - continue; + fn colorise(&mut self) -> BTreeMap { + for node in self.mir.indices() { + if !self.mir.is_register(node.0) { + continue; + } + self.color_node(node.0.into()); } - colorizer.color_node(node.into()); + self.colors + .iter() + .filter_map(|(&node, &c)| match c { + Color::Final(reg) => Some((node, reg)), + _ => None, + }) + .collect() } + } +} - let colors = colorizer.finalise(); +impl Mir { + fn dst_node(&self, node: NodeRef) -> Option { + // for each node, look at the dst node and see if it has preferred + // color, then also prefer that color. + let (inst, data) = self.get_node(node); - // 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()) - // } - // }), - // ), - // ); + match inst { + Inst::Add(_) + | Inst::Sub(_) + | Inst::MulSigned(_) + | Inst::MulSSE(_) + | Inst::DivSSE(_) + | Inst::BitAnd(_) + | Inst::BitOr(_) + | Inst::BitXOr(_) + | Inst::ShiftLeft(_) + | Inst::ShiftRightSigned(_) + | Inst::ShiftRightUnsigned(_) => { + let (lhs, _) = data.as_binary_noderefs(); + Some(lhs) + } + Inst::Negate(_) => { + let lhs = data.as_noderef(); + Some(lhs) + } + Inst::Parameter(_) + | Inst::GetElementPtr(_) + | Inst::LoadRegister(_) + | Inst::Load(_) + | Inst::Label + | Inst::ConstantBytes + | Inst::ConstantByte + | Inst::ConstantWord + | Inst::ConstantDWord + | Inst::ConstantQWord + | Inst::ConstantSinglePrecision + | Inst::ConstantDoublePrecision + | Inst::ExternRef + | Inst::Alloca + | Inst::Store(_) + | Inst::ReturnValue(_) + | Inst::Return + | Inst::SignExtend(_) + | Inst::ZeroExtend(_) + | Inst::Mul(_) + | Inst::Div(_) + | Inst::RemFP(_) + | Inst::DivSigned(_) + | Inst::Rem(_) + | Inst::RemSigned(_) + | Inst::IsZero(_) => None, + } + } - colors + pub fn build_liveness(&self) -> Liveness { + LivenessBuilder::new(self).build() } } @@ -1656,8 +1919,6 @@ impl Function { } } -type Liveness = BTreeMap; - #[allow(dead_code, unused)] impl Mir { fn node_as_operand( @@ -1697,7 +1958,7 @@ impl Mir { let label = func.get_constant_label(node as usize); ImmRegMem::Rip(RipRelative::Label(Type::QWord, label)) } - Inst::GetElementPtr(ty) => liveness.get(&node).unwrap().clone().into(), + Inst::GetElementPtr(ty) => liveness.get_register(node.into()).unwrap().into(), Inst::Parameter(ty) | Inst::Add(ty) | Inst::Sub(ty) @@ -1707,6 +1968,9 @@ impl Mir { | Inst::DivSigned(ty) | Inst::Rem(ty) | Inst::RemSigned(ty) + | Inst::MulSSE(ty) + | Inst::DivSSE(ty) + | Inst::RemFP(ty) | Inst::BitAnd(ty) | Inst::BitOr(ty) | Inst::BitXOr(ty) @@ -1718,7 +1982,9 @@ impl Mir { | Inst::ZeroExtend(ty) | Inst::IsZero(ty) | Inst::Load(ty) - | Inst::LoadRegister(ty) => ty.register_width(*liveness.get(&node).unwrap()).into(), + | Inst::LoadRegister(ty) => ty + .register_width(liveness.get_register(node.into()).unwrap()) + .into(), Inst::Alloca => { let (offset, size) = *mapping.get(&(node as usize)).unwrap(); ImmRegMem::Mem(StackMem::new(offset, size)) @@ -1737,7 +2003,7 @@ impl Mir { let mut func = Function::new(self.name); let name = strings.get_str(self.name).to_owned(); let liveness = self.build_liveness(); - func.dirty_registers.extend(liveness.values().cloned()); + func.dirty_registers.extend(liveness.dirty_registers()); let mut float_params = 0; let mut int_params = 0; @@ -1769,7 +2035,7 @@ impl Mir { } Inst::LoadRegister(ty) => { let src = data.as_node(); - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let src = self.node_as_operand(&liveness, &mapping, &mut func, strings, src); if ty.is_floating() { match src { @@ -1777,19 +2043,32 @@ impl Mir { | ImmRegMem::Word(_) | ImmRegMem::DWord(_) | ImmRegMem::QWord(_) => { - writeln!(func.current_branch(), "push rax")?; + let mut spill_rax = false; + let scratch = liveness + .get_scratch_register_at_node(node.into()) + .unwrap_or_else(|| { + spill_rax = true; + Register::rax + }); + + if spill_rax { + writeln!(func.current_branch(), "push rax")?; + } writeln!( func.current_branch(), "mov {}, {src}", - ty.register_width(Register::rax) + ty.register_width(scratch) )?; writeln!( func.current_branch(), "movd {dst}, {}", - ty.register_width(Register::rax) + ty.register_width(scratch) )?; - writeln!(func.current_branch(), "pop rax")?; + + if spill_rax { + writeln!(func.current_branch(), "pop rax")?; + } } ImmRegMem::Mem(_) | ImmRegMem::Rip(_) => { writeln!(func.current_branch(), "movss {dst}, {src}",)?; @@ -1809,7 +2088,7 @@ impl Mir { mapping.insert(i, (offset, size)); } Inst::Load(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let src = data.as_node(); let src = self.node_as_operand(&liveness, &mapping, &mut func, strings, src); match src { @@ -1817,15 +2096,22 @@ impl Mir { writeln!(func.current_branch(), "mov {}, {ty} ptr [{}]", dst, src,)?; } ImmRegMem::Mem(ref mem) => { - let tmp_reg = if dst.parent_reg() == Register::rax { - Register::rcx - } else { - Register::rax - }; - writeln!(func.current_branch(), "push {tmp_reg}")?; - writeln!(func.current_branch(), "mov {}, {}", tmp_reg, src)?; - writeln!(func.current_branch(), "mov {}, [{}]", dst, tmp_reg)?; - writeln!(func.current_branch(), "pop {tmp_reg}")?; + let mut spill_rax = false; + let scratch = liveness + .get_scratch_register_at_node(node.into()) + .unwrap_or_else(|| { + spill_rax = true; + Register::rax + }); + + if spill_rax { + writeln!(func.current_branch(), "push rax")?; + } + writeln!(func.current_branch(), "mov {}, {}", scratch, src)?; + writeln!(func.current_branch(), "mov {}, [{}]", dst, scratch)?; + if spill_rax { + writeln!(func.current_branch(), "pop rax")?; + } } _ => {} } @@ -1843,7 +2129,7 @@ impl Mir { } } Inst::GetElementPtr(ty) => { - let dst = *liveness.get(&node).unwrap(); + let dst = liveness.get_register(node.into()).unwrap(); let (src, idx) = data.as_binary(); let src = self.node_as_operand(&liveness, &mapping, &mut func, strings, src); @@ -1870,7 +2156,7 @@ impl Mir { } } Inst::Add(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); @@ -1888,7 +2174,7 @@ impl Mir { } } Inst::Sub(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); @@ -1905,69 +2191,66 @@ impl Mir { } } Inst::Mul(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); - if ty.is_floating() { - writeln!(func.current_branch(), "mulss {lhs}, {rhs}")?; - if !lhs.occupy_same_register(dst) { - writeln!(func.current_branch(), "movss {dst}, {lhs}")?; - } - } else { - let spill_rax = !(lhs.occupy_same_register(Register::rax) + let spill_rax = liveness.is_register_in_use_at_node(node.into(), Register::rax) + && !(lhs.occupy_same_register(Register::rax) || rhs.occupy_same_register(Register::rax) || dst.parent_reg() == Register::rax); - let spill_rdx = !(lhs.occupy_same_register(Register::rdx) + let spill_rdx = liveness.is_register_in_use_at_node(node.into(), Register::rdx) + && !(lhs.occupy_same_register(Register::rdx) || rhs.occupy_same_register(Register::rdx) || dst.parent_reg() == Register::rdx); - if spill_rax { - writeln!(func.current_branch(), "push rax")?; - } - if spill_rdx { - writeln!(func.current_branch(), "push rdx")?; - } + if spill_rax { + writeln!(func.current_branch(), "push rax")?; + } + if spill_rdx { + writeln!(func.current_branch(), "push rdx")?; + } - if !lhs.occupy_same_register(Register::rax) { - writeln!( - func.current_branch(), - "mov {}, {lhs}", - ty.register_width(Register::rax) - )?; - } + if !lhs.occupy_same_register(Register::rax) { + writeln!( + func.current_branch(), + "mov {}, {lhs}", + ty.register_width(Register::rax) + )?; + } - writeln!(func.current_branch(), "mul {rhs}")?; + writeln!(func.current_branch(), "mul {rhs}")?; - if dst.parent_reg() != Register::rax { - writeln!( - func.current_branch(), - "mov {dst}, {}", - ty.register_width(Register::rax) - )?; - } + if dst.parent_reg() != Register::rax { + writeln!( + func.current_branch(), + "mov {dst}, {}", + ty.register_width(Register::rax) + )?; + } - if spill_rdx { - writeln!(func.current_branch(), "pop rdx")?; - } - if spill_rax { - writeln!(func.current_branch(), "pop rax")?; - } + if spill_rdx { + writeln!(func.current_branch(), "pop rdx")?; + } + if spill_rax { + writeln!(func.current_branch(), "pop rax")?; } } Inst::MulSigned(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); - let spill_rax = !(lhs.occupy_same_register(Register::rax) - || rhs.occupy_same_register(Register::rax) - || dst.parent_reg() == Register::rax); - let spill_rdx = !(lhs.occupy_same_register(Register::rdx) - || rhs.occupy_same_register(Register::rdx) - || dst.parent_reg() == Register::rdx); + let spill_rax = liveness.is_register_in_use_at_node(node.into(), Register::rax) + && !(lhs.occupy_same_register(Register::rax) + || rhs.occupy_same_register(Register::rax) + || dst.parent_reg() == Register::rax); + let spill_rdx = liveness.is_register_in_use_at_node(node.into(), Register::rdx) + && !(lhs.occupy_same_register(Register::rdx) + || rhs.occupy_same_register(Register::rdx) + || dst.parent_reg() == Register::rdx); if spill_rax { writeln!(func.current_branch(), "push rax")?; @@ -1990,69 +2273,66 @@ impl Mir { } } Inst::Div(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); - if ty.is_floating() { - writeln!(func.current_branch(), "divss {lhs}, {rhs}")?; - if lhs.occupy_same_register(dst) { - writeln!(func.current_branch(), "movss {dst}, {lhs}")?; - } - } else { - let spill_rax = !(lhs.occupy_same_register(Register::rax) + let spill_rax = liveness.is_register_in_use_at_node(node.into(), Register::rax) + && !(lhs.occupy_same_register(Register::rax) || rhs.occupy_same_register(Register::rax) || dst.parent_reg() == Register::rax); - let spill_rdx = !(lhs.occupy_same_register(Register::rdx) + let spill_rdx = liveness.is_register_in_use_at_node(node.into(), Register::rdx) + && !(lhs.occupy_same_register(Register::rdx) || rhs.occupy_same_register(Register::rdx) || dst.parent_reg() == Register::rdx); - if spill_rax { - writeln!(func.current_branch(), "push rax")?; - } - if spill_rdx { - writeln!(func.current_branch(), "push rdx")?; - } + if spill_rax { + writeln!(func.current_branch(), "push rax")?; + } + if spill_rdx { + writeln!(func.current_branch(), "push rdx")?; + } - if !lhs.occupy_same_register(Register::rax) { - writeln!( - func.current_branch(), - "mov {}, {lhs}", - ty.register_width(Register::rax) - )?; - } + if !lhs.occupy_same_register(Register::rax) { + writeln!( + func.current_branch(), + "mov {}, {lhs}", + ty.register_width(Register::rax) + )?; + } - writeln!(func.current_branch(), "div {rhs}")?; + writeln!(func.current_branch(), "div {rhs}")?; - if dst.parent_reg() != Register::rax { - writeln!( - func.current_branch(), - "mov {dst}, {}", - ty.register_width(Register::rax) - )?; - } + if dst.parent_reg() != Register::rax { + writeln!( + func.current_branch(), + "mov {dst}, {}", + ty.register_width(Register::rax) + )?; + } - if spill_rdx { - writeln!(func.current_branch(), "pop rdx")?; - } - if spill_rax { - writeln!(func.current_branch(), "pop rax")?; - } + if spill_rdx { + writeln!(func.current_branch(), "pop rdx")?; + } + if spill_rax { + writeln!(func.current_branch(), "pop rax")?; } } Inst::DivSigned(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); - let spill_rax = !(lhs.occupy_same_register(Register::rax) - || rhs.occupy_same_register(Register::rax) - || dst.parent_reg() == Register::rax); - let spill_rdx = !(lhs.occupy_same_register(Register::rdx) - || rhs.occupy_same_register(Register::rdx) - || dst.parent_reg() == Register::rdx); + let spill_rax = liveness.is_register_in_use_at_node(node.into(), Register::rax) + && !(lhs.occupy_same_register(Register::rax) + || rhs.occupy_same_register(Register::rax) + || dst.parent_reg() == Register::rax); + let spill_rdx = liveness.is_register_in_use_at_node(node.into(), Register::rdx) + && !(lhs.occupy_same_register(Register::rdx) + || rhs.occupy_same_register(Register::rdx) + || dst.parent_reg() == Register::rdx); if spill_rax { writeln!(func.current_branch(), "push rax")?; @@ -2087,79 +2367,66 @@ impl Mir { } } Inst::Rem(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); - if ty.is_floating() { - let ty = ty.int_repr(); - let size = ty.bytes(); - writeln!(func.current_branch(), "sub rsp, 0x{:x}", size * 2)?; - writeln!(func.current_branch(), "movss {ty} ptr [rsp], {lhs}")?; - writeln!( - func.current_branch(), - "movss {ty} ptr [rsp + {size}], {rhs}" - )?; - writeln!(func.current_branch(), "fld {ty} ptr [rsp + {size}]")?; - writeln!(func.current_branch(), "fld {ty} ptr [rsp]")?; - writeln!(func.current_branch(), "fprem")?; - writeln!(func.current_branch(), "fst {ty} ptr [rsp]")?; - writeln!(func.current_branch(), "movss {dst}, {ty} ptr [rsp]")?; - writeln!(func.current_branch(), "add rsp, 0x{:x}", size * 2)?; - } else { - let spill_rax = !(lhs.occupy_same_register(Register::rax) + let spill_rax = liveness.is_register_in_use_at_node(node.into(), Register::rax) + && !(lhs.occupy_same_register(Register::rax) || rhs.occupy_same_register(Register::rax) || dst.parent_reg() == Register::rax); - let spill_rdx = !(lhs.occupy_same_register(Register::rdx) + let spill_rdx = liveness.is_register_in_use_at_node(node.into(), Register::rdx) + && !(lhs.occupy_same_register(Register::rdx) || rhs.occupy_same_register(Register::rdx) || dst.parent_reg() == Register::rdx); - if spill_rax { - writeln!(func.current_branch(), "push rax")?; - } - if spill_rdx { - writeln!(func.current_branch(), "push rdx")?; - } + if spill_rax { + writeln!(func.current_branch(), "push rax")?; + } + if spill_rdx { + writeln!(func.current_branch(), "push rdx")?; + } - if !lhs.occupy_same_register(Register::rax) { - writeln!( - func.current_branch(), - "mov {}, {lhs}", - ty.register_width(Register::rax) - )?; - } + if !lhs.occupy_same_register(Register::rax) { + writeln!( + func.current_branch(), + "mov {}, {lhs}", + ty.register_width(Register::rax) + )?; + } - writeln!(func.current_branch(), "div {rhs}")?; + writeln!(func.current_branch(), "div {rhs}")?; - if dst.parent_reg() != Register::rdx { - writeln!( - func.current_branch(), - "mov {dst}, {}", - ty.register_width(Register::rdx) - )?; - } + if dst.parent_reg() != Register::rdx { + writeln!( + func.current_branch(), + "mov {dst}, {}", + ty.register_width(Register::rdx) + )?; + } - if spill_rdx { - writeln!(func.current_branch(), "pop rdx")?; - } - if spill_rax { - writeln!(func.current_branch(), "pop rax")?; - } + if spill_rdx { + writeln!(func.current_branch(), "pop rdx")?; + } + if spill_rax { + writeln!(func.current_branch(), "pop rax")?; } } Inst::RemSigned(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); - let spill_rax = !(lhs.occupy_same_register(Register::rax) - || rhs.occupy_same_register(Register::rax) - || dst.parent_reg() == Register::rax); - let spill_rdx = !(lhs.occupy_same_register(Register::rdx) - || rhs.occupy_same_register(Register::rdx) - || dst.parent_reg() == Register::rdx); + let spill_rax = liveness.is_register_in_use_at_node(node.into(), Register::rax) + && !(lhs.occupy_same_register(Register::rax) + || rhs.occupy_same_register(Register::rax) + || dst.parent_reg() == Register::rax); + let spill_rdx = liveness.is_register_in_use_at_node(node.into(), Register::rdx) + && !(lhs.occupy_same_register(Register::rdx) + || rhs.occupy_same_register(Register::rdx) + || dst.parent_reg() == Register::rdx); if spill_rax { writeln!(func.current_branch(), "push rax")?; @@ -2193,8 +2460,51 @@ impl Mir { writeln!(func.current_branch(), "pop rax")?; } } + Inst::MulSSE(ty) => { + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); + let (lhs, rhs) = data.as_binary(); + let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); + let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); + + writeln!(func.current_branch(), "mulss {lhs}, {rhs}")?; + if lhs.occupy_same_register(dst) { + writeln!(func.current_branch(), "movss {dst}, {lhs}")?; + } + } + Inst::DivSSE(ty) => { + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); + let (lhs, rhs) = data.as_binary(); + let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); + let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); + + writeln!(func.current_branch(), "divss {lhs}, {rhs}")?; + if lhs.occupy_same_register(dst) { + writeln!(func.current_branch(), "movss {dst}, {lhs}")?; + } + } + Inst::RemFP(ty) => { + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); + let (lhs, rhs) = data.as_binary(); + let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); + let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); + + let ty = ty.int_repr(); + let size = ty.bytes(); + writeln!(func.current_branch(), "sub rsp, 0x{:x}", size * 2)?; + writeln!(func.current_branch(), "movss {ty} ptr [rsp], {lhs}")?; + writeln!( + func.current_branch(), + "movss {ty} ptr [rsp + {size}], {rhs}" + )?; + writeln!(func.current_branch(), "fld {ty} ptr [rsp + {size}]")?; + writeln!(func.current_branch(), "fld {ty} ptr [rsp]")?; + writeln!(func.current_branch(), "fprem")?; + writeln!(func.current_branch(), "fst {ty} ptr [rsp]")?; + writeln!(func.current_branch(), "movss {dst}, {ty} ptr [rsp]")?; + writeln!(func.current_branch(), "add rsp, 0x{:x}", size * 2)?; + } Inst::BitAnd(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); @@ -2204,7 +2514,7 @@ impl Mir { } } Inst::BitOr(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); @@ -2214,7 +2524,7 @@ impl Mir { } } Inst::BitXOr(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); @@ -2224,7 +2534,7 @@ impl Mir { } } Inst::ShiftLeft(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); @@ -2251,7 +2561,9 @@ impl Mir { // if lhs is in rcx, lhs needs to move to rax and we spill rax temporarily // if neither lhs nor rhx nor dst are rcx, spill rcx temporarily - let spill_rcx = !(lhs_is_rcx || rhs_is_rcx || dst_is_rcx); + let spill_rcx = liveness + .is_register_in_use_at_node(node.into(), Register::rcx) + && !(lhs_is_rcx || rhs_is_rcx || dst_is_rcx); if spill_rcx { writeln!(func.current_branch(), "push rcx")?; @@ -2267,7 +2579,10 @@ impl Mir { } if lhs_is_rcx { - writeln!(func.current_branch(), "push rax")?; + if liveness.is_register_in_use_at_node(node.into(), Register::rax) { + writeln!(func.current_branch(), "push rax")?; + } + writeln!(func.current_branch(), "test rax,rax")?; writeln!( func.current_branch(), @@ -2277,7 +2592,10 @@ impl Mir { )?; writeln!(func.current_branch(), "shl rax, cl")?; writeln!(func.current_branch(), "mov {dst}, rax")?; - writeln!(func.current_branch(), "pop rax")?; + + if liveness.is_register_in_use_at_node(node.into(), Register::rax) { + writeln!(func.current_branch(), "pop rax")?; + } } else { writeln!(func.current_branch(), "shl {lhs}, cl")?; if lhs.occupy_same_register(dst) { @@ -2296,7 +2614,7 @@ impl Mir { } } Inst::ShiftRightSigned(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); @@ -2323,7 +2641,9 @@ impl Mir { // if lhs is in rcx, lhs needs to move to rax and we spill rax temporarily // if neither lhs nor rhx nor dst are rcx, spill rcx temporarily - let spill_rcx = !(lhs_is_rcx || rhs_is_rcx || dst_is_rcx); + let spill_rcx = liveness + .is_register_in_use_at_node(node.into(), Register::rcx) + && !(lhs_is_rcx || rhs_is_rcx || dst_is_rcx); if spill_rcx { writeln!(func.current_branch(), "push rcx")?; @@ -2339,7 +2659,9 @@ impl Mir { } if lhs_is_rcx { - writeln!(func.current_branch(), "push rax")?; + if liveness.is_register_in_use_at_node(node.into(), Register::rax) { + writeln!(func.current_branch(), "push rax")?; + } writeln!(func.current_branch(), "test rax,rax")?; writeln!( func.current_branch(), @@ -2349,7 +2671,10 @@ impl Mir { )?; writeln!(func.current_branch(), "sar rax, cl")?; writeln!(func.current_branch(), "mov {dst}, rax")?; - writeln!(func.current_branch(), "pop rax")?; + + if liveness.is_register_in_use_at_node(node.into(), Register::rax) { + writeln!(func.current_branch(), "pop rax")?; + } } else { writeln!(func.current_branch(), "sar {lhs}, cl")?; if lhs.occupy_same_register(dst) { @@ -2368,7 +2693,7 @@ impl Mir { } } Inst::ShiftRightUnsigned(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let (lhs, rhs) = data.as_binary(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); let rhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, rhs); @@ -2395,7 +2720,9 @@ impl Mir { // if lhs is in rcx, lhs needs to move to rax and we spill rax temporarily // if neither lhs nor rhx nor dst are rcx, spill rcx temporarily - let spill_rcx = !(lhs_is_rcx || rhs_is_rcx || dst_is_rcx); + let spill_rcx = liveness + .is_register_in_use_at_node(node.into(), Register::rcx) + && !(lhs_is_rcx || rhs_is_rcx || dst_is_rcx); if spill_rcx { writeln!(func.current_branch(), "push rcx")?; @@ -2411,7 +2738,10 @@ impl Mir { } if lhs_is_rcx { - writeln!(func.current_branch(), "push rax")?; + if liveness.is_register_in_use_at_node(node.into(), Register::rax) { + writeln!(func.current_branch(), "push rax")?; + } + writeln!(func.current_branch(), "test rax,rax")?; writeln!( func.current_branch(), @@ -2421,7 +2751,10 @@ impl Mir { )?; writeln!(func.current_branch(), "shr rax, cl")?; writeln!(func.current_branch(), "mov {dst}, rax")?; - writeln!(func.current_branch(), "pop rax")?; + + if liveness.is_register_in_use_at_node(node.into(), Register::rax) { + writeln!(func.current_branch(), "pop rax")?; + } } else { writeln!(func.current_branch(), "shr {lhs}, cl")?; if lhs.occupy_same_register(dst) { @@ -2440,7 +2773,7 @@ impl Mir { } } Inst::Negate(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let lhs = data.as_node(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); @@ -2457,7 +2790,7 @@ impl Mir { } } Inst::SignExtend(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let lhs = data.as_node(); let lhs_ty = self.type_of_node(lhs).unwrap(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); @@ -2527,7 +2860,7 @@ impl Mir { } } Inst::ZeroExtend(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let lhs = data.as_node(); let lhs_ty = self.type_of_node(lhs).unwrap(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); @@ -2569,7 +2902,7 @@ impl Mir { } } Inst::IsZero(ty) => { - let dst = ty.register_width(*liveness.get(&node).unwrap()); + let dst = ty.register_width(liveness.get_register(node.into()).unwrap()); let lhs = data.as_node(); let lhs = self.node_as_operand(&liveness, &mapping, &mut func, strings, lhs); writeln!(func.current_branch(), "mov {dst}, {lhs}")?; diff --git a/src/parser.rs b/src/parser.rs index a404e7b..8c8a1fe 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -82,6 +82,7 @@ impl Nodes { self.push_tag(Tag::Undefined) } + #[allow(dead_code)] fn swap_nodes(&mut self, lhs: Node, rhs: Node) { self.inner.swap(lhs.get() as usize, rhs.get() as usize); } diff --git a/tests/legal/inverse_sqrt.sea b/tests/legal/inverse_sqrt.sea index f704fcd..eba5b71 100644 --- a/tests/legal/inverse_sqrt.sea +++ b/tests/legal/inverse_sqrt.sea @@ -1,7 +1,7 @@ fn inverse_sqrt(n: f32) -> f32 { let x = n; var i = *(&x as *i32); - i = 0x5f3759dfi32 - (i >> 1u8); + i = 0x5f3759dfi32 - (i >> 1); let y = *(&i as *f32); y * (1.5f32 - (x * 0.5f32 * y * y)) } \ No newline at end of file