From 1171ce75f99be7521f148d09b509c418ee53fa96 Mon Sep 17 00:00:00 2001 From: Jesse Brault Date: Sat, 7 Mar 2026 18:23:45 -0600 Subject: [PATCH] Propagate register assignments. --- dm/src/main.rs | 2 +- dmc-lib/src/ir/ir_add.rs | 10 ++++++- dmc-lib/src/ir/ir_assign.rs | 17 ++++++++++- dmc-lib/src/ir/ir_block.rs | 42 ++++++++++++--------------- dmc-lib/src/ir/ir_call.rs | 11 ++++++- dmc-lib/src/ir/ir_expression.rs | 29 +++++++++++++++++- dmc-lib/src/ir/ir_function.rs | 6 ++-- dmc-lib/src/ir/ir_operation.rs | 19 +++++++++++- dmc-lib/src/ir/ir_return.rs | 11 ++++++- dmc-lib/src/ir/ir_statement.rs | 19 +++++++++++- dmc-lib/src/ir/ir_variable.rs | 8 +++++ dmc-lib/src/ir/register_allocation.rs | 31 ++++++++++++++++++++ 12 files changed, 171 insertions(+), 34 deletions(-) diff --git a/dm/src/main.rs b/dm/src/main.rs index 97c3f45..0f9b865 100644 --- a/dm/src/main.rs +++ b/dm/src/main.rs @@ -61,7 +61,7 @@ fn main() { for declaration in compilation_unit.declarations() { if let ModuleLevelDeclaration::Function(function) = declaration { let mut ir_function = function.to_ir(&symbol_table); - let register_assignments = ir_function.register_assignments(args.register_count); + let register_assignments = ir_function.assign_registers(args.register_count); println!("{}", ir_function); println!("{:?}", register_assignments); } diff --git a/dmc-lib/src/ir/ir_add.rs b/dmc-lib/src/ir/ir_add.rs index 13695f9..ff0d4f3 100644 --- a/dmc-lib/src/ir/ir_add.rs +++ b/dmc-lib/src/ir/ir_add.rs @@ -1,7 +1,7 @@ use crate::ir::ir_expression::IrExpression; use crate::ir::ir_variable::IrVrVariableDescriptor; use crate::ir::register_allocation::VrUser; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; pub struct IrAdd { @@ -43,4 +43,12 @@ impl VrUser for IrAdd { self.left.propagate_spills(spills); self.right.propagate_spills(spills); } + + fn propagate_register_assignments( + &mut self, + assignments: &HashMap, + ) { + self.left.propagate_register_assignments(assignments); + self.right.propagate_register_assignments(assignments); + } } diff --git a/dmc-lib/src/ir/ir_assign.rs b/dmc-lib/src/ir/ir_assign.rs index cc7be49..903aaf6 100644 --- a/dmc-lib/src/ir/ir_assign.rs +++ b/dmc-lib/src/ir/ir_assign.rs @@ -2,7 +2,7 @@ use crate::ir::ir_operation::IrOperation; use crate::ir::ir_variable::{IrVariable, IrVariableDescriptor, IrVrVariableDescriptor}; use crate::ir::register_allocation::VrUser; use std::cell::RefCell; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; use std::rc::Rc; @@ -50,6 +50,21 @@ impl VrUser for IrAssign { } } } + + fn propagate_register_assignments( + &mut self, + assignments: &HashMap, + ) { + let mut borrowed_destination = self.destination.borrow_mut(); + if let IrVariableDescriptor::VirtualRegister(vr_variable) = + borrowed_destination.descriptor_mut() + { + if assignments.contains_key(vr_variable) { + vr_variable.set_assigned_register(assignments[vr_variable]); + } + } + self.initializer.propagate_register_assignments(assignments); + } } impl Display for IrAssign { diff --git a/dmc-lib/src/ir/ir_block.rs b/dmc-lib/src/ir/ir_block.rs index 93e84c5..c858f34 100644 --- a/dmc-lib/src/ir/ir_block.rs +++ b/dmc-lib/src/ir/ir_block.rs @@ -1,6 +1,6 @@ use crate::ir::ir_statement::IrStatement; use crate::ir::ir_variable::IrVrVariableDescriptor; -use crate::ir::register_allocation::{HasVrUsers, VrUser, registers_and_spills}; +use crate::ir::register_allocation::{HasVrUsers, VrUser}; use std::cell::RefCell; use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; @@ -32,34 +32,19 @@ impl IrBlock { pub fn statements(&self) -> &[IrStatement] { &self.statements } - - pub fn register_assignments( - &mut self, - register_count: usize, - ) -> HashMap { - let mut spills: HashSet = 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; - // mutate all IrVirtualRegisters to constituent statements - for statement in &mut self.statements { - statement.propagate_spills(&spills); - } - } else { - return registers; - } - } - } } impl HasVrUsers for IrBlock { fn vr_users(&self) -> Vec<&dyn VrUser> { self.statements.iter().map(|s| s as &dyn VrUser).collect() } + + fn vr_users_mut(&mut self) -> Vec<&mut dyn VrUser> { + self.statements + .iter_mut() + .map(|s| s as &mut dyn VrUser) + .collect() + } } impl VrUser for IrBlock { @@ -79,6 +64,15 @@ impl VrUser for IrBlock { statement.propagate_spills(spills); } } + + fn propagate_register_assignments( + &mut self, + assignments: &HashMap, + ) { + for statement in &mut self.statements { + statement.propagate_register_assignments(assignments); + } + } } impl Display for IrBlock { @@ -141,7 +135,7 @@ mod tests { if let ModuleLevelDeclaration::Function(main) = main { let mut main_ir = main.to_ir(&symbol_table); - let register_assignments = main_ir.register_assignments(2); + let register_assignments = main_ir.assign_registers(2); assert_eq!(register_assignments.len(), 4); } else { unreachable!() diff --git a/dmc-lib/src/ir/ir_call.rs b/dmc-lib/src/ir/ir_call.rs index 134d9ae..449d9a1 100644 --- a/dmc-lib/src/ir/ir_call.rs +++ b/dmc-lib/src/ir/ir_call.rs @@ -1,7 +1,7 @@ use crate::ir::ir_expression::IrExpression; use crate::ir::ir_variable::IrVrVariableDescriptor; use crate::ir::register_allocation::VrUser; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; use std::rc::Rc; @@ -33,6 +33,15 @@ impl VrUser for IrCall { argument.propagate_spills(spills); } } + + fn propagate_register_assignments( + &mut self, + assignments: &HashMap, + ) { + for argument in &mut self.arguments { + argument.propagate_register_assignments(assignments); + } + } } impl Display for IrCall { diff --git a/dmc-lib/src/ir/ir_expression.rs b/dmc-lib/src/ir/ir_expression.rs index 9e955cd..05bc390 100644 --- a/dmc-lib/src/ir/ir_expression.rs +++ b/dmc-lib/src/ir/ir_expression.rs @@ -4,7 +4,7 @@ use crate::ir::ir_variable::{ }; use crate::ir::register_allocation::VrUser; use std::cell::RefCell; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; use std::rc::Rc; @@ -85,4 +85,31 @@ impl VrUser for IrExpression { } } } + + fn propagate_register_assignments( + &mut self, + assignments: &HashMap, + ) { + match self { + IrExpression::Parameter(_) => { + // no-op + } + IrExpression::Variable(ir_variable) => { + let mut borrowed = ir_variable.borrow_mut(); + if let IrVariableDescriptor::VirtualRegister(vr_variable) = + borrowed.descriptor_mut() + { + if assignments.contains_key(vr_variable) { + vr_variable.set_assigned_register(assignments[vr_variable]); + } + } + } + IrExpression::Int(_) => { + // no-op + } + IrExpression::String(_) => { + // no-op + } + } + } } diff --git a/dmc-lib/src/ir/ir_function.rs b/dmc-lib/src/ir/ir_function.rs index 7610a15..35ea071 100644 --- a/dmc-lib/src/ir/ir_function.rs +++ b/dmc-lib/src/ir/ir_function.rs @@ -1,6 +1,7 @@ use crate::ir::ir_block::IrBlock; use crate::ir::ir_parameter::IrParameter; use crate::ir::ir_variable::IrVrVariableDescriptor; +use crate::ir::register_allocation::{HasVrUsers, VrUser}; use crate::type_info::TypeInfo; use std::cell::RefCell; use std::collections::HashMap; @@ -29,11 +30,12 @@ impl IrFunction { } } - pub fn register_assignments( + /// Returns the register assignments for debugging purposes + pub fn assign_registers( &mut self, register_count: usize, ) -> HashMap { - self.entry.borrow_mut().register_assignments(register_count) + self.entry.borrow_mut().assign_registers(register_count) } } diff --git a/dmc-lib/src/ir/ir_operation.rs b/dmc-lib/src/ir/ir_operation.rs index de1bbcd..e7dc8a7 100644 --- a/dmc-lib/src/ir/ir_operation.rs +++ b/dmc-lib/src/ir/ir_operation.rs @@ -3,7 +3,7 @@ use crate::ir::ir_call::IrCall; use crate::ir::ir_expression::IrExpression; use crate::ir::ir_variable::IrVrVariableDescriptor; use crate::ir::register_allocation::VrUser; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; pub enum IrOperation { @@ -58,4 +58,21 @@ impl VrUser for IrOperation { } } } + + fn propagate_register_assignments( + &mut self, + assignments: &HashMap, + ) { + match self { + IrOperation::Load(ir_expression) => { + ir_expression.propagate_register_assignments(assignments); + } + IrOperation::Add(ir_add) => { + ir_add.propagate_register_assignments(assignments); + } + IrOperation::Call(ir_call) => { + ir_call.propagate_register_assignments(assignments); + } + } + } } diff --git a/dmc-lib/src/ir/ir_return.rs b/dmc-lib/src/ir/ir_return.rs index caaa754..7ad8674 100644 --- a/dmc-lib/src/ir/ir_return.rs +++ b/dmc-lib/src/ir/ir_return.rs @@ -1,7 +1,7 @@ use crate::ir::ir_expression::IrExpression; use crate::ir::ir_variable::IrVrVariableDescriptor; use crate::ir::register_allocation::VrUser; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; pub struct IrReturn { @@ -32,6 +32,15 @@ impl VrUser for IrReturn { ir_expression.propagate_spills(spills); } } + + fn propagate_register_assignments( + &mut self, + assignments: &HashMap, + ) { + if let Some(ir_expression) = self.value.as_mut() { + ir_expression.propagate_register_assignments(assignments); + } + } } impl Display for IrReturn { diff --git a/dmc-lib/src/ir/ir_statement.rs b/dmc-lib/src/ir/ir_statement.rs index e5d16b8..5676671 100644 --- a/dmc-lib/src/ir/ir_statement.rs +++ b/dmc-lib/src/ir/ir_statement.rs @@ -3,7 +3,7 @@ use crate::ir::ir_call::IrCall; use crate::ir::ir_return::IrReturn; use crate::ir::ir_variable::IrVrVariableDescriptor; use crate::ir::register_allocation::VrUser; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; pub enum IrStatement { @@ -42,6 +42,23 @@ impl VrUser for IrStatement { } } } + + fn propagate_register_assignments( + &mut self, + assignments: &HashMap, + ) { + match self { + IrStatement::Assign(ir_assign) => { + ir_assign.propagate_register_assignments(assignments); + } + IrStatement::Call(ir_call) => { + ir_call.propagate_register_assignments(assignments); + } + IrStatement::Return(ir_return) => { + ir_return.propagate_register_assignments(assignments); + } + } + } } impl Display for IrStatement { diff --git a/dmc-lib/src/ir/ir_variable.rs b/dmc-lib/src/ir/ir_variable.rs index 3ad9753..1705776 100644 --- a/dmc-lib/src/ir/ir_variable.rs +++ b/dmc-lib/src/ir/ir_variable.rs @@ -33,6 +33,10 @@ impl IrVariable { &self.descriptor } + pub fn descriptor_mut(&mut self) -> &mut IrVariableDescriptor { + &mut self.descriptor + } + pub fn set_descriptor(&mut self, descriptor: IrVariableDescriptor) { self.descriptor = descriptor; } @@ -94,6 +98,10 @@ impl IrVrVariableDescriptor { pub fn block_id(&self) -> usize { self.block_id } + + pub fn set_assigned_register(&mut self, register: usize) { + self.assigned_register = Some(register); + } } impl Display for IrVrVariableDescriptor { diff --git a/dmc-lib/src/ir/register_allocation.rs b/dmc-lib/src/ir/register_allocation.rs index cae0c9e..349956d 100644 --- a/dmc-lib/src/ir/register_allocation.rs +++ b/dmc-lib/src/ir/register_allocation.rs @@ -6,6 +6,7 @@ pub type LivenessMap = HashMap>; 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(); @@ -99,12 +100,42 @@ pub trait HasVrUsers { graph } + + fn assign_registers( + &mut self, + register_count: usize, + ) -> HashMap { + let mut spills: HashSet = 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); + } + return registers; + } + } + } } pub trait VrUser { fn vr_definitions(&self) -> HashSet; fn vr_uses(&self) -> HashSet; fn propagate_spills(&mut self, spills: &HashSet); + fn propagate_register_assignments( + &mut self, + assignments: &HashMap, + ); } #[derive(Debug)]