wave_compiler/regalloc/
mod.rs

1// Copyright 2026 Ojima Abraham
2// SPDX-License-Identifier: Apache-2.0
3
4//! Register allocator using graph coloring (Chaitin-Briggs).
5//!
6//! Assigns physical registers to virtual registers. Handles interference
7//! graph construction, move coalescing, graph coloring with pre-colored
8//! parameters, and spill code generation. After allocation, rewrites
9//! `VReg`s to `PhysReg`s in the instruction stream.
10
11pub mod coalesce;
12pub mod coloring;
13pub mod interference;
14pub mod live_range;
15pub mod spill;
16
17use std::collections::HashMap;
18
19use crate::emit::wave_emit::RegMap;
20use crate::lir::instruction::LirInst;
21use crate::lir::operand::VReg;
22
23/// Perform register allocation on LIR instructions.
24///
25/// Returns the register mapping from `VReg`s to `PhysReg`s.
26/// Pre-colors parameter registers: VReg(0)→PhysReg(0), etc.
27#[must_use]
28pub fn allocate_registers(instructions: &[LirInst], num_params: u32, max_regs: u32) -> RegMap {
29    let ig = interference::InterferenceGraph::build_with_params(instructions, num_params);
30    let coalesce_result = coalesce::coalesce(&ig, max_regs);
31    let coloring_result = coloring::color(&ig, max_regs, num_params);
32
33    let mut reg_map: RegMap = HashMap::new();
34
35    for (vreg, phys) in &coloring_result.assignment {
36        let final_vreg = find_representative(&coalesce_result.mapping, *vreg);
37        reg_map.insert(final_vreg, *phys);
38        reg_map.insert(*vreg, *phys);
39    }
40
41    for (coalesced, representative) in &coalesce_result.mapping {
42        if let Some(&phys) = reg_map.get(representative) {
43            reg_map.insert(*coalesced, phys);
44        }
45    }
46
47    reg_map
48}
49
50fn find_representative(mapping: &HashMap<VReg, VReg>, vreg: VReg) -> VReg {
51    let mut current = vreg;
52    while let Some(&next) = mapping.get(&current) {
53        if next == current {
54            break;
55        }
56        current = next;
57    }
58    current
59}
60
61/// Count the maximum physical register index used.
62#[must_use]
63pub fn max_register_used(reg_map: &RegMap) -> u32 {
64    reg_map
65        .values()
66        .map(|p| u32::from(p.0) + 1)
67        .max()
68        .unwrap_or(0)
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use crate::lir::operand::PhysReg;
75
76    #[test]
77    fn test_allocate_registers_simple() {
78        let insts = vec![
79            LirInst::MovImm {
80                dest: VReg(0),
81                value: 1,
82            },
83            LirInst::MovImm {
84                dest: VReg(1),
85                value: 2,
86            },
87            LirInst::Iadd {
88                dest: VReg(2),
89                src1: VReg(0),
90                src2: VReg(1),
91            },
92            LirInst::Halt,
93        ];
94
95        let reg_map = allocate_registers(&insts, 0, 32);
96        assert!(reg_map.contains_key(&VReg(0)));
97        assert!(reg_map.contains_key(&VReg(1)));
98        assert!(reg_map.contains_key(&VReg(2)));
99    }
100
101    #[test]
102    fn test_allocate_with_params() {
103        let insts = vec![
104            LirInst::Iadd {
105                dest: VReg(2),
106                src1: VReg(0),
107                src2: VReg(1),
108            },
109            LirInst::Halt,
110        ];
111
112        let reg_map = allocate_registers(&insts, 2, 32);
113        assert_eq!(reg_map[&VReg(0)], PhysReg(0));
114        assert_eq!(reg_map[&VReg(1)], PhysReg(1));
115    }
116
117    #[test]
118    fn test_max_register_used() {
119        let mut map = RegMap::new();
120        map.insert(VReg(0), PhysReg(5));
121        map.insert(VReg(1), PhysReg(10));
122        assert_eq!(max_register_used(&map), 11);
123    }
124}
wave_compiler/regalloc/mod.rs

wave_compiler/regalloc/
mod.rs
