Propagate register assignments.

This commit is contained in:
Jesse Brault 2026-03-07 18:23:45 -06:00
parent d8bd826cb0
commit 1171ce75f9
12 changed files with 171 additions and 34 deletions

View File

@ -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);
}

View File

@ -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<IrVrVariableDescriptor, usize>,
) {
self.left.propagate_register_assignments(assignments);
self.right.propagate_register_assignments(assignments);
}
}

View File

@ -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<IrVrVariableDescriptor, usize>,
) {
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 {

View File

@ -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<IrVrVariableDescriptor, usize> {
let mut spills: HashSet<IrVrVariableDescriptor> = 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<IrVrVariableDescriptor, usize>,
) {
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!()

View File

@ -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<IrVrVariableDescriptor, usize>,
) {
for argument in &mut self.arguments {
argument.propagate_register_assignments(assignments);
}
}
}
impl Display for IrCall {

View File

@ -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<IrVrVariableDescriptor, usize>,
) {
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
}
}
}
}

View File

@ -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<IrVrVariableDescriptor, usize> {
self.entry.borrow_mut().register_assignments(register_count)
self.entry.borrow_mut().assign_registers(register_count)
}
}

View File

@ -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<IrVrVariableDescriptor, usize>,
) {
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);
}
}
}
}

View File

@ -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<IrVrVariableDescriptor, usize>,
) {
if let Some(ir_expression) = self.value.as_mut() {
ir_expression.propagate_register_assignments(assignments);
}
}
}
impl Display for IrReturn {

View File

@ -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<IrVrVariableDescriptor, usize>,
) {
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 {

View File

@ -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 {

View File

@ -6,6 +6,7 @@ pub type LivenessMap = HashMap<usize, HashSet<IrVrVariableDescriptor>>;
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<IrVrVariableDescriptor, usize> {
let mut spills: HashSet<IrVrVariableDescriptor> = 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(&registers);
}
return registers;
}
}
}
}
pub trait VrUser {
fn vr_definitions(&self) -> HashSet<IrVrVariableDescriptor>;
fn vr_uses(&self) -> HashSet<IrVrVariableDescriptor>;
fn propagate_spills(&mut self, spills: &HashSet<IrVrVariableDescriptor>);
fn propagate_register_assignments(
&mut self,
assignments: &HashMap<IrVrVariableDescriptor, usize>,
);
}
#[derive(Debug)]