llvm_codegen/
regalloc_gc.rs

1//! Graph-coloring register allocator (Chaitin-Briggs style).
2
3use crate::isel::{PReg, VReg};
4use crate::regalloc::{LiveInterval, RegAllocResult};
5use std::collections::{HashMap, HashSet};
6
7/// Public API for `InterferenceGraph`.
8pub type InterferenceGraph = HashMap<VReg, HashSet<VReg>>;
9
10fn overlaps(a: &LiveInterval, b: &LiveInterval) -> bool {
11    a.start < b.end && b.start < a.end
12}
13
14/// Public API for `build_interference_graph`.
15pub fn build_interference_graph(intervals: &[LiveInterval]) -> InterferenceGraph {
16    let mut graph: InterferenceGraph = HashMap::new();
17    for iv in intervals {
18        graph.entry(iv.vreg).or_default();
19    }
20
21    for i in 0..intervals.len() {
22        for j in (i + 1)..intervals.len() {
23            let a = &intervals[i];
24            let b = &intervals[j];
25            if overlaps(a, b) {
26                graph.entry(a.vreg).or_default().insert(b.vreg);
27                graph.entry(b.vreg).or_default().insert(a.vreg);
28            }
29        }
30    }
31
32    graph
33}
34
35/// Compute spill costs from interval length and use density.
36/// Longer intervals are considered more expensive to spill.
37pub fn spill_costs(intervals: &[LiveInterval]) -> HashMap<VReg, usize> {
38    let mut costs = HashMap::new();
39    for iv in intervals {
40        let len = iv.end.saturating_sub(iv.start).max(1);
41        costs.insert(iv.vreg, len);
42    }
43    costs
44}
45
46/// Chaitin-Briggs style graph-coloring allocator.
47///
48/// The returned type matches linear-scan, so spill/reload insertion and
49/// allocation application can remain unchanged.
50pub fn graph_color(intervals: &[LiveInterval], allocatable: &[PReg]) -> RegAllocResult {
51    if allocatable.is_empty() {
52        return RegAllocResult {
53            vreg_to_preg: HashMap::new(),
54            spilled: intervals.iter().map(|iv| iv.vreg).collect(),
55        };
56    }
57    if intervals.is_empty() {
58        return RegAllocResult::default();
59    }
60
61    let k = allocatable.len();
62    let graph = build_interference_graph(intervals);
63    let costs = spill_costs(intervals);
64
65    let mut work = graph.clone();
66    let mut stack: Vec<VReg> = Vec::with_capacity(work.len());
67    let mut potential_spills: HashSet<VReg> = HashSet::new();
68
69    while !work.is_empty() {
70        let low_degree = work
71            .iter()
72            .filter(|(_, neigh)| neigh.len() < k)
73            .map(|(vr, neigh)| (*vr, neigh.len()))
74            .min_by_key(|(vr, degree)| (*degree, vr.0));
75
76        let chosen = if let Some((vr, _)) = low_degree {
77            vr
78        } else {
79            let spill_vr = work
80                .iter()
81                .map(|(vr, neigh)| {
82                    let cost = *costs.get(vr).unwrap_or(&1);
83                    (*vr, cost, neigh.len())
84                })
85                .min_by_key(|(vr, cost, degree)| (*cost, std::cmp::Reverse(*degree), vr.0))
86                .map(|(vr, _, _)| vr)
87                .expect("non-empty graph must pick a node");
88            potential_spills.insert(spill_vr);
89            spill_vr
90        };
91
92        if let Some(neighs) = work.remove(&chosen) {
93            for n in neighs {
94                if let Some(nset) = work.get_mut(&n) {
95                    nset.remove(&chosen);
96                }
97            }
98        }
99        stack.push(chosen);
100    }
101
102    let mut assigned: HashMap<VReg, PReg> = HashMap::new();
103    let mut spilled: HashSet<VReg> = HashSet::new();
104
105    while let Some(vr) = stack.pop() {
106        let forbidden: HashSet<PReg> = graph
107            .get(&vr)
108            .into_iter()
109            .flat_map(|neigh| neigh.iter())
110            .filter_map(|n| assigned.get(n).copied())
111            .collect();
112
113        if let Some(&pr) = allocatable.iter().find(|pr| !forbidden.contains(pr)) {
114            assigned.insert(vr, pr);
115        } else {
116            spilled.insert(vr);
117        }
118    }
119
120    // Keep potential spills that could not be assigned as actual spills.
121    for vr in potential_spills {
122        if !assigned.contains_key(&vr) {
123            spilled.insert(vr);
124        }
125    }
126
127    RegAllocResult {
128        vreg_to_preg: assigned,
129        spilled: {
130            let mut out: Vec<VReg> = spilled.into_iter().collect();
131            out.sort_unstable_by_key(|vr| vr.0);
132            out
133        },
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use crate::regalloc::{allocate_registers, linear_scan, RegAllocStrategy};
141
142    fn iv(vreg: u32, start: usize, end: usize) -> LiveInterval {
143        LiveInterval {
144            vreg: VReg(vreg),
145            start,
146            end,
147        }
148    }
149
150    #[test]
151    fn graph_builds_empty() {
152        let g = build_interference_graph(&[]);
153        assert!(g.is_empty());
154    }
155
156    #[test]
157    fn graph_adds_overlap_edge() {
158        let g = build_interference_graph(&[iv(0, 0, 4), iv(1, 2, 5)]);
159        assert!(g[&VReg(0)].contains(&VReg(1)));
160        assert!(g[&VReg(1)].contains(&VReg(0)));
161    }
162
163    #[test]
164    fn graph_omits_non_overlap_edge() {
165        let g = build_interference_graph(&[iv(0, 0, 2), iv(1, 2, 4)]);
166        assert!(g[&VReg(0)].is_empty());
167        assert!(g[&VReg(1)].is_empty());
168    }
169
170    #[test]
171    fn graph_touching_endpoints_do_not_interfere() {
172        let a = iv(0, 3, 7);
173        let b = iv(1, 7, 9);
174        assert!(!overlaps(&a, &b));
175    }
176
177    #[test]
178    fn spill_cost_favors_longer_ranges() {
179        let costs = spill_costs(&[iv(0, 0, 2), iv(1, 0, 10)]);
180        assert!(costs[&VReg(1)] > costs[&VReg(0)]);
181    }
182
183    #[test]
184    fn graph_color_empty_intervals() {
185        let r = graph_color(&[], &[PReg(0)]);
186        assert!(r.vreg_to_preg.is_empty());
187        assert!(r.spilled.is_empty());
188    }
189
190    #[test]
191    fn graph_color_no_registers_spills_all() {
192        let intervals = vec![iv(0, 0, 4), iv(1, 1, 3)];
193        let r = graph_color(&intervals, &[]);
194        assert_eq!(r.vreg_to_preg.len(), 0);
195        assert_eq!(r.spilled.len(), 2);
196    }
197
198    #[test]
199    fn graph_color_non_overlapping_share_one_reg() {
200        let intervals = vec![iv(0, 0, 2), iv(1, 2, 4), iv(2, 4, 6)];
201        let r = graph_color(&intervals, &[PReg(0)]);
202        assert!(r.spilled.is_empty());
203        assert_eq!(r.vreg_to_preg.len(), 3);
204    }
205
206    #[test]
207    fn graph_color_triangle_two_regs_spills_one() {
208        let intervals = vec![iv(0, 0, 6), iv(1, 1, 7), iv(2, 2, 8)];
209        let r = graph_color(&intervals, &[PReg(0), PReg(1)]);
210        assert_eq!(r.spilled.len(), 1);
211    }
212
213    #[test]
214    fn graph_color_triangle_three_regs_spills_none() {
215        let intervals = vec![iv(0, 0, 6), iv(1, 1, 7), iv(2, 2, 8)];
216        let r = graph_color(&intervals, &[PReg(0), PReg(1), PReg(2)]);
217        assert!(r.spilled.is_empty());
218    }
219
220    #[test]
221    fn graph_color_deterministic_for_same_input() {
222        let intervals = vec![iv(3, 0, 5), iv(0, 1, 4), iv(2, 2, 6), iv(1, 5, 7)];
223        let regs = vec![PReg(0), PReg(1)];
224        let a = graph_color(&intervals, &regs);
225        let b = graph_color(&intervals, &regs);
226        assert_eq!(a.spilled, b.spilled);
227        assert_eq!(a.vreg_to_preg, b.vreg_to_preg);
228    }
229
230    #[test]
231    fn graph_color_respects_interference_assignments() {
232        let intervals = vec![iv(0, 0, 5), iv(1, 0, 5), iv(2, 5, 8)];
233        let r = graph_color(&intervals, &[PReg(0), PReg(1)]);
234        let p0 = r.vreg_to_preg[&VReg(0)];
235        let p1 = r.vreg_to_preg[&VReg(1)];
236        assert_ne!(p0, p1);
237    }
238
239    #[test]
240    fn graph_color_can_recover_potential_spill() {
241        // 0 and 2 do not interfere, so one color can be reused.
242        let intervals = vec![iv(0, 0, 3), iv(1, 1, 4), iv(2, 4, 6)];
243        let r = graph_color(&intervals, &[PReg(0), PReg(1)]);
244        assert!(r.spilled.is_empty());
245        assert_eq!(r.vreg_to_preg.len(), 3);
246    }
247
248    #[test]
249    fn strategy_dispatch_uses_graph_color() {
250        let intervals = vec![iv(0, 0, 6), iv(1, 1, 7), iv(2, 2, 8)];
251        let regs = vec![PReg(0), PReg(1)];
252        let r = allocate_registers(&intervals, &regs, RegAllocStrategy::GraphColor);
253        assert_eq!(r.spilled.len(), 1);
254    }
255
256    #[test]
257    fn strategy_dispatch_linear_scan_default() {
258        let intervals = vec![iv(0, 0, 4), iv(1, 2, 6)];
259        let regs = vec![PReg(0)];
260        let a = allocate_registers(&intervals, &regs, RegAllocStrategy::default());
261        let b = linear_scan(&intervals, &regs);
262        assert_eq!(a.spilled, b.spilled);
263        assert_eq!(a.vreg_to_preg, b.vreg_to_preg);
264    }
265
266    #[test]
267    fn graph_color_spills_no_more_than_linear_scan_case1() {
268        let intervals = vec![
269            iv(0, 0, 8),
270            iv(1, 1, 5),
271            iv(2, 2, 6),
272            iv(3, 6, 10),
273            iv(4, 8, 11),
274        ];
275        let regs = vec![PReg(0), PReg(1)];
276        let gc = graph_color(&intervals, &regs);
277        let ls = linear_scan(&intervals, &regs);
278        assert!(gc.spilled.len() <= ls.spilled.len());
279    }
280
281    #[test]
282    fn graph_color_spills_no_more_than_linear_scan_case2() {
283        let intervals = vec![
284            iv(0, 0, 7),
285            iv(1, 0, 7),
286            iv(2, 1, 3),
287            iv(3, 3, 5),
288            iv(4, 5, 9),
289            iv(5, 6, 10),
290        ];
291        let regs = vec![PReg(0), PReg(1), PReg(2)];
292        let gc = graph_color(&intervals, &regs);
293        let ls = linear_scan(&intervals, &regs);
294        assert!(gc.spilled.len() <= ls.spilled.len());
295    }
296}
llvm_codegen/regalloc_gc.rs

llvm_codegen/
regalloc_gc.rs
