lamina_codegen/aarch64/
regalloc.rs

1use std::collections::{HashMap, HashSet, VecDeque};
2
3use crate::regalloc::{Allocation, LocalRegisterAllocator as MirRegisterAllocator};
4use lamina_mir::{Register, RegisterClass, VirtualReg};
5
6// AArch64 register allocator for MIR virtual registers.
7
8/// AArch64 register allocator for MIR virtual registers.
9///
10/// Tracks a fixed pool of available caller-saved GPRs and gives
11/// stable mappings from virtual registers to physical registers.
12/// Returns None when no register is available (caller should spill to stack).
13pub struct A64RegAlloc {
14    free_gprs: VecDeque<&'static str>,
15    used_gprs: HashSet<&'static str>,
16    global_mappings: HashMap<VirtualReg, &'static str>,
17    vreg_to_preg: HashMap<VirtualReg, &'static str>,
18    scratch_free: VecDeque<&'static str>,
19    scratch_used: HashSet<&'static str>,
20}
21
22const MAP_GPRS: &[&str] = &["x13", "x14", "x15"];
23
24/// Callee-saved registers a global plan may draw on, once the prologue saves them.
25///
26/// x16 and x17 are the emitter's own scratch, and x18 is the platform register on
27/// Darwin and Windows.
28const CALLEE_SAVED_GPRS: &[&str] = &[
29    "x19", "x20", "x21", "x22", "x23", "x24", "x25", "x26", "x27", "x28",
30];
31
32const SCRATCH_GPRS: &[&str] = &["x9", "x10", "x11", "x12"];
33
34impl Default for A64RegAlloc {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl A64RegAlloc {
41    /// Creates a new register allocator with default register pools.
42    pub fn new() -> Self {
43        let mut free_gprs = VecDeque::new();
44        for &r in MAP_GPRS {
45            free_gprs.push_back(r);
46        }
47
48        let mut scratch_free = VecDeque::new();
49        for &r in SCRATCH_GPRS {
50            scratch_free.push_back(r);
51        }
52
53        Self {
54            free_gprs,
55            used_gprs: HashSet::new(),
56            global_mappings: HashMap::new(),
57            vreg_to_preg: HashMap::new(),
58            scratch_free,
59            scratch_used: HashSet::new(),
60        }
61    }
62
63    /// Sets conservative mode for complex functions, using fewer registers to reduce pressure.
64    pub fn set_conservative_mode(&mut self) {
65        self.free_gprs.clear();
66        self.free_gprs.push_back("x13");
67        self.free_gprs.push_back("x14");
68        self.free_gprs.push_back("x15");
69        self.used_gprs.retain(|r| self.free_gprs.contains(r));
70        let free_set: HashSet<&str> = self.free_gprs.iter().copied().collect();
71        self.vreg_to_preg.retain(|_, preg| free_set.contains(preg));
72    }
73
74    #[inline]
75    pub fn alloc_scratch(&mut self) -> Option<&'static str> {
76        MirRegisterAllocator::alloc_scratch(self)
77    }
78
79    #[inline]
80    pub fn free_scratch(&mut self, phys: &'static str) {
81        MirRegisterAllocator::free_scratch(self, phys);
82    }
83
84    #[inline]
85    pub fn is_occupied(&self, phys: &'static str) -> bool {
86        MirRegisterAllocator::is_occupied(self, phys)
87    }
88
89    #[inline]
90    pub fn occupy(&mut self, phys: &'static str) {
91        MirRegisterAllocator::occupy(self, phys);
92    }
93
94    #[inline]
95    pub fn release(&mut self, phys: &'static str) {
96        MirRegisterAllocator::release(self, phys);
97    }
98
99    #[inline]
100    pub fn get_mapping_for(&self, v: &VirtualReg) -> Option<&'static str> {
101        MirRegisterAllocator::get_mapping(self, v)
102    }
103
104    #[inline]
105    pub fn ensure_mapping(&mut self, v: VirtualReg) -> Option<&'static str> {
106        MirRegisterAllocator::ensure_mapping(self, v)
107    }
108
109    #[inline]
110    pub fn ensure_mapping_for_gpr(&mut self, v: VirtualReg) -> Option<&'static str> {
111        MirRegisterAllocator::ensure_mapping(self, v)
112    }
113
114    #[inline]
115    pub fn mapped_for_register(&self, r: &Register) -> Option<&'static str> {
116        MirRegisterAllocator::mapped_for_register(self, r)
117    }
118
119    /// GPR candidates for global allocation. `callee_saved` widens the pool with
120    /// x21-x28, which the prologue then has to save and the epilogue restore.
121    pub fn gpr_pool_for_global_allocation(callee_saved: bool) -> Vec<&'static str> {
122        let mut pool = MAP_GPRS.to_vec();
123        if callee_saved {
124            pool.extend_from_slice(CALLEE_SAVED_GPRS);
125        }
126        pool
127    }
128
129    /// The caller-saved members of the global pool. A call destroys each of these, so a
130    /// value live across one may not be planned into them.
131    pub fn caller_saved_in_global_pool() -> &'static [&'static str] {
132        MAP_GPRS
133    }
134
135    /// The callee-saved registers a plan hands out, in register-number order.
136    pub fn callee_saved_in_plan(
137        plan: &HashMap<VirtualReg, Allocation<&'static str>>,
138    ) -> Vec<&'static str> {
139        CALLEE_SAVED_GPRS
140            .iter()
141            .copied()
142            .filter(|reg| {
143                plan.iter().any(|(vreg, alloc)| {
144                    vreg.class == RegisterClass::Gpr
145                        && matches!(alloc, Allocation::Register(phys) if phys == reg)
146                })
147            })
148            .collect()
149    }
150
151    pub fn from_global_plan(plan: &HashMap<VirtualReg, Allocation<&'static str>>) -> Self {
152        let mut s = Self::new();
153        s.free_gprs.clear();
154        for (&vreg, alloc) in plan {
155            if vreg.class != RegisterClass::Gpr {
156                continue;
157            }
158            if let Allocation::Register(phys) = alloc
159                && (MAP_GPRS.contains(phys) || CALLEE_SAVED_GPRS.contains(phys))
160            {
161                s.global_mappings.insert(vreg, *phys);
162                s.used_gprs.insert(*phys);
163            }
164        }
165        s
166    }
167}
168
169impl MirRegisterAllocator for A64RegAlloc {
170    type PhysReg = &'static str;
171
172    fn alloc_scratch(&mut self) -> Option<Self::PhysReg> {
173        if let Some(phys) = self.scratch_free.pop_front() {
174            self.scratch_used.insert(phys);
175            Some(phys)
176        } else {
177            None
178        }
179    }
180
181    fn free_scratch(&mut self, phys: Self::PhysReg) {
182        if self.scratch_used.remove(&phys) {
183            self.scratch_free.push_back(phys);
184        }
185    }
186
187    fn get_mapping(&self, vreg: &VirtualReg) -> Option<Self::PhysReg> {
188        if let Some(&phys) = self.global_mappings.get(vreg) {
189            return Some(phys);
190        }
191        self.vreg_to_preg.get(vreg).copied()
192    }
193
194    fn ensure_mapping(&mut self, vreg: VirtualReg) -> Option<Self::PhysReg> {
195        if let Some(phys) = MirRegisterAllocator::get_mapping(self, &vreg) {
196            return Some(phys);
197        }
198
199        if vreg.class != RegisterClass::Gpr {
200            return None;
201        }
202
203        if let Some(phys) = self.free_gprs.pop_front() {
204            self.used_gprs.insert(phys);
205            self.vreg_to_preg.insert(vreg, phys);
206            Some(phys)
207        } else if let Some((vreg_to_replace, &phys)) = self.vreg_to_preg.iter().next() {
208            let vreg_to_replace = *vreg_to_replace;
209            self.vreg_to_preg.remove(&vreg_to_replace);
210            self.vreg_to_preg.insert(vreg, phys);
211            Some(phys)
212        } else {
213            None
214        }
215    }
216
217    fn mapped_for_register(&self, reg: &Register) -> Option<Self::PhysReg> {
218        match reg {
219            Register::Virtual(v) => MirRegisterAllocator::get_mapping(self, v),
220            Register::Physical(p) => Some(p.name),
221        }
222    }
223
224    fn occupy(&mut self, phys: Self::PhysReg) {
225        if self.used_gprs.insert(phys)
226            && let Some(pos) = self.free_gprs.iter().position(|&p| p == phys)
227        {
228            self.free_gprs.remove(pos);
229        }
230    }
231
232    fn release(&mut self, phys: Self::PhysReg) {
233        if self.used_gprs.remove(&phys) {
234            self.free_gprs.push_back(phys);
235        }
236    }
237
238    fn is_occupied(&self, phys: Self::PhysReg) -> bool {
239        self.used_gprs.contains(&phys)
240    }
241}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246
247    #[test]
248    fn global_plan_mappings_survive_and_are_shared() {
249        let first = VirtualReg::gpr(0);
250        let second = VirtualReg::gpr(1);
251        let third = VirtualReg::gpr(2);
252        let plan = HashMap::from([
253            (first, Allocation::Register("x21")),
254            (second, Allocation::Register("x21")),
255            (third, Allocation::Spill(-8)),
256        ]);
257        let mut allocator = A64RegAlloc::from_global_plan(&plan);
258        assert_eq!(allocator.get_mapping_for(&first), Some("x21"));
259        assert_eq!(allocator.get_mapping_for(&second), Some("x21"));
260        assert_eq!(allocator.get_mapping_for(&third), None);
261        assert_eq!(allocator.ensure_mapping(VirtualReg::gpr(9)), None);
262        assert_eq!(allocator.get_mapping_for(&first), Some("x21"));
263    }
264
265    #[test]
266    fn callee_saved_pool_is_opt_in() {
267        let narrow = A64RegAlloc::gpr_pool_for_global_allocation(false);
268        let wide = A64RegAlloc::gpr_pool_for_global_allocation(true);
269        assert_eq!(narrow, vec!["x13", "x14", "x15"]);
270        assert!(wide.starts_with(&["x13", "x14", "x15"]));
271        assert_eq!(wide.len(), 13);
272        for reserved in ["x16", "x17", "x18", "x29", "x30"] {
273            assert!(
274                !wide.contains(&reserved),
275                "{reserved} must stay unallocated"
276            );
277        }
278        for scratch in SCRATCH_GPRS {
279            assert!(!wide.contains(scratch));
280        }
281    }
282
283    #[test]
284    fn callee_saved_in_plan_reports_register_order() {
285        let plan = HashMap::from([
286            (VirtualReg::gpr(0), Allocation::Register("x24")),
287            (VirtualReg::gpr(1), Allocation::Register("x21")),
288            (VirtualReg::gpr(2), Allocation::Register("x24")),
289            (VirtualReg::gpr(3), Allocation::Register("x13")),
290            (VirtualReg::gpr(4), Allocation::Spill(-16)),
291        ]);
292        assert_eq!(A64RegAlloc::callee_saved_in_plan(&plan), vec!["x21", "x24"]);
293    }
294
295    #[test]
296    fn local_allocation_stays_on_caller_saved_registers() {
297        let mut allocator = A64RegAlloc::new();
298        let mut seen = Vec::new();
299        for index in 0..3 {
300            seen.push(allocator.ensure_mapping(VirtualReg::gpr(index)).unwrap());
301        }
302        assert_eq!(seen, vec!["x13", "x14", "x15"]);
303    }
304}
lamina_codegen/aarch64/regalloc.rs

lamina_codegen/aarch64/
regalloc.rs
