wave_compiler/regalloc/
coloring.rs

1// Copyright 2026 Ojima Abraham
2// SPDX-License-Identifier: Apache-2.0
3
4//! Graph coloring register allocator following Chaitin-Briggs.
5//!
6//! Assigns physical registers to virtual registers by coloring the
7//! interference graph. Pre-colors parameter registers per Rule 9:
8//! VReg(0)→PhysReg(0), VReg(1)→PhysReg(1), etc.
9
10use std::collections::{HashMap, HashSet};
11
12use super::interference::InterferenceGraph;
13use crate::lir::operand::{PhysReg, VReg};
14
15/// Result of graph coloring.
16pub struct ColoringResult {
17    /// Mapping from `VReg` to assigned `PhysReg`.
18    pub assignment: HashMap<VReg, PhysReg>,
19    /// `VReg`s that could not be colored and need spilling.
20    pub spilled: Vec<VReg>,
21}
22
23/// Perform graph coloring with pre-colored parameter registers.
24///
25/// Parameters VReg(0)..VReg(num_params-1) are pre-colored to
26/// PhysReg(0)..PhysReg(num_params-1) per Rule 9.
27#[must_use]
28pub fn color(ig: &InterferenceGraph, max_regs: u32, num_params: u32) -> ColoringResult {
29    let mut assignment: HashMap<VReg, PhysReg> = HashMap::new();
30    let mut spilled = Vec::new();
31
32    for i in 0..num_params {
33        let vreg = VReg(i);
34        if ig.adj.contains_key(&vreg) {
35            #[allow(clippy::cast_possible_truncation)]
36            assignment.insert(vreg, PhysReg(i as u8));
37        }
38    }
39
40    let mut nodes: Vec<VReg> = ig.nodes();
41    nodes.sort_by_key(|v| std::cmp::Reverse(ig.degree(*v)));
42
43    let mut stack: Vec<VReg> = Vec::new();
44    let mut removed: HashSet<VReg> = HashSet::new();
45
46    for &vreg in &nodes {
47        if !assignment.contains_key(&vreg) {
48            stack.push(vreg);
49            removed.insert(vreg);
50        }
51    }
52
53    while let Some(vreg) = stack.pop() {
54        let mut used_colors: HashSet<u8> = HashSet::new();
55        if let Some(neighbors) = ig.adj.get(&vreg) {
56            for &neighbor in neighbors {
57                if let Some(&phys) = assignment.get(&neighbor) {
58                    used_colors.insert(phys.0);
59                }
60            }
61        }
62
63        let mut color_found = false;
64        for c in 0..max_regs.min(32) as u8 {
65            if !used_colors.contains(&c) {
66                assignment.insert(vreg, PhysReg(c));
67                color_found = true;
68                break;
69            }
70        }
71
72        if !color_found {
73            spilled.push(vreg);
74        }
75    }
76
77    ColoringResult {
78        assignment,
79        spilled,
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use crate::lir::instruction::LirInst;
87
88    #[test]
89    fn test_coloring_simple() {
90        let insts = vec![
91            LirInst::MovImm {
92                dest: VReg(0),
93                value: 1,
94            },
95            LirInst::MovImm {
96                dest: VReg(1),
97                value: 2,
98            },
99            LirInst::Iadd {
100                dest: VReg(2),
101                src1: VReg(0),
102                src2: VReg(1),
103            },
104            LirInst::Halt,
105        ];
106
107        let ig = InterferenceGraph::build(&insts);
108        let result = color(&ig, 32, 0);
109
110        assert!(result.spilled.is_empty());
111        assert!(result.assignment.contains_key(&VReg(0)));
112        assert!(result.assignment.contains_key(&VReg(1)));
113        assert!(result.assignment.contains_key(&VReg(2)));
114
115        let r0 = result.assignment[&VReg(0)];
116        let r1 = result.assignment[&VReg(1)];
117        if ig.interferes(VReg(0), VReg(1)) {
118            assert_ne!(r0, r1);
119        }
120    }
121
122    #[test]
123    fn test_coloring_precolors_params() {
124        let insts = vec![
125            LirInst::Iadd {
126                dest: VReg(2),
127                src1: VReg(0),
128                src2: VReg(1),
129            },
130            LirInst::Halt,
131        ];
132
133        let ig = InterferenceGraph::build(&insts);
134        let result = color(&ig, 32, 2);
135
136        assert_eq!(result.assignment[&VReg(0)], PhysReg(0));
137        assert_eq!(result.assignment[&VReg(1)], PhysReg(1));
138    }
139}
wave_compiler/regalloc/coloring.rs

wave_compiler/regalloc/
coloring.rs
