lamina_codegen/riscv/
regalloc.rs

1use std::collections::HashMap;
2
3use crate::{
4    regalloc::{Allocation, LocalRegisterAllocator as MirRegisterAllocator},
5    riscv::{
6        frame::RiscVFrame,
7        target::{RiscVTarget, Xlen},
8    },
9};
10use lamina_mir::{Register, RegisterClass, VirtualReg};
11use lamina_platform::TargetOperatingSystem;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum RiscVRegisterConvention {
15    CallerSaved,
16    CalleeSaved,
17    Reserved,
18}
19
20/// RISC-V register allocator with platform-aware register selection.
21///
22/// RISC-V has 32 general-purpose registers (x0-x31):
23/// - x0: zero (hardwired zero)
24/// - x1: ra (return address)
25/// - x2: sp (stack pointer)
26/// - x3: gp (global pointer)
27/// - x4: tp (thread pointer)
28/// - x5-x7, x28-x31: temporaries
29/// - x8: fp/s0 (frame pointer/saved register)
30/// - x9-x15: s1-s7 (saved registers)
31/// - x16-x27: a0-a7 (argument registers), t0-t6 (temporaries)
32///
33pub struct RiscVRegAlloc {
34    target: RiscVTarget,
35    #[allow(dead_code)]
36    target_os: TargetOperatingSystem,
37    available_gprs: Vec<&'static str>,
38    global_mappings: HashMap<VirtualReg, &'static str>,
39    allocated_gprs: HashMap<&'static str, VirtualReg>,
40    stack_slots: HashMap<VirtualReg, i32>,
41    next_stack_slot: i32,
42}
43
44impl Default for RiscVRegAlloc {
45    fn default() -> Self {
46        Self::new(TargetOperatingSystem::Linux)
47    }
48}
49
50impl RiscVRegAlloc {
51    pub const CALLER_SAVED_REGISTERS: &'static [&'static str] = &[
52        "t0", "t1", "t2", "t3", "t4", "t5", "t6", "a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7",
53    ];
54    pub const CALLEE_SAVED_REGISTERS: &'static [&'static str] = &[
55        "s0", "s1", "s2", "s3", "s4", "s5", "s6", "s7", "s8", "s9", "s10", "s11",
56    ];
57    pub const RESERVED_REGISTERS: &'static [&'static str] = &["zero", "ra", "sp", "gp", "tp"];
58
59    /// Registers this allocator may hand out, by ABI name.
60    ///
61    /// ABI names, not `x` numbers, because the emitter writes `a0`, `a1` and `t0`
62    /// directly. While this list said `x10` and the emitter said `a0`, they were the
63    /// same register under two spellings and nothing could see the clash.
64    ///
65    /// Excluded and why:
66    ///   a0-a7 (x10-x17)  arguments and return values, and the emitter names a0/a1
67    ///   t0     (x5)      the emitter's own scratch
68    ///   t1-t3  (x6, x7, x28)  scratch in the float-to-integer clamp and the indirect call
69    ///   s0/fp  (x8)      the frame pointer
70    ///   ra, sp, gp, tp, zero     reserved
71    const AVAILABLE_REGISTERS: &'static [&'static str] = &[
72        "t4", "t5", "t6", "s1", "s2", "s3", "s4", "s5", "s6", "s7", "s8", "s9", "s10", "s11",
73    ];
74
75    pub fn new(target_os: TargetOperatingSystem) -> Self {
76        Self::with_target(target_os, RiscVTarget::general(Xlen::Rv64))
77    }
78
79    /// The ISA this allocator is emitting for. Held here because every helper that needs
80    /// the word size already receives the allocator.
81    pub fn target(&self) -> RiscVTarget {
82        self.target
83    }
84
85    pub fn register_convention(reg: &str) -> Option<RiscVRegisterConvention> {
86        if Self::CALLER_SAVED_REGISTERS.contains(&reg) {
87            Some(RiscVRegisterConvention::CallerSaved)
88        } else if Self::CALLEE_SAVED_REGISTERS.contains(&reg) || reg == "fp" {
89            Some(RiscVRegisterConvention::CalleeSaved)
90        } else if Self::RESERVED_REGISTERS.contains(&reg) {
91            Some(RiscVRegisterConvention::Reserved)
92        } else {
93            None
94        }
95    }
96
97    pub fn caller_saved_mappings(
98        &self,
99        live_vregs: &[VirtualReg],
100    ) -> Vec<(VirtualReg, &'static str)> {
101        let mut mappings: Vec<_> = live_vregs
102            .iter()
103            .filter_map(|vreg| {
104                let physical = self.get_mapping(vreg)?;
105                (Self::register_convention(physical) == Some(RiscVRegisterConvention::CallerSaved))
106                    .then_some((*vreg, physical))
107            })
108            .collect();
109        mappings.sort_unstable_by_key(|(_, physical)| *physical);
110        mappings
111    }
112
113    pub fn used_callee_saved_registers(&self) -> Vec<&'static str> {
114        let mut registers: Vec<_> = self
115            .allocated_gprs
116            .keys()
117            .copied()
118            .filter(|reg| {
119                Self::register_convention(reg) == Some(RiscVRegisterConvention::CalleeSaved)
120            })
121            .collect();
122        registers.sort_unstable_by_key(|reg| {
123            Self::CALLEE_SAVED_REGISTERS
124                .iter()
125                .position(|candidate| candidate == reg)
126        });
127        registers
128    }
129
130    pub fn with_target(target_os: TargetOperatingSystem, target: RiscVTarget) -> Self {
131        Self {
132            target,
133            target_os,
134            available_gprs: Self::AVAILABLE_REGISTERS.to_vec(),
135            global_mappings: HashMap::new(),
136            allocated_gprs: HashMap::new(),
137            stack_slots: HashMap::new(),
138            next_stack_slot: RiscVFrame::calculate_stack_offset(0, target),
139        }
140    }
141
142    /// Set conservative mode (limit to fewer registers)
143    pub fn set_conservative_mode(&mut self) {
144        self.available_gprs = vec!["t4", "t5"];
145    }
146
147    /// Get stack slot for a virtual register
148    pub fn get_stack_slot(&self, vreg: &VirtualReg) -> Option<i32> {
149        self.stack_slots.get(vreg).copied()
150    }
151
152    pub fn gpr_pool_for_global_allocation() -> Vec<&'static str> {
153        Self::AVAILABLE_REGISTERS.to_vec()
154    }
155
156    pub fn from_global_plan(
157        target_os: TargetOperatingSystem,
158        target: RiscVTarget,
159        plan: &HashMap<VirtualReg, Allocation<&'static str>>,
160        stack_slots: &HashMap<VirtualReg, i32>,
161    ) -> Self {
162        let mut s = Self::with_target(target_os, target);
163        let mut min_spill = 0i32;
164        for (&vreg, alloc) in plan {
165            if vreg.class != RegisterClass::Gpr {
166                continue;
167            }
168            match alloc {
169                Allocation::Register(phys) => {
170                    if s.available_gprs.contains(phys) {
171                        s.global_mappings.insert(vreg, *phys);
172                        s.allocated_gprs.insert(*phys, vreg);
173                    }
174                }
175                Allocation::Spill(_) => {
176                    if let Some(off) = stack_slots.get(&vreg) {
177                        s.stack_slots.insert(vreg, *off);
178                        if *off < min_spill {
179                            min_spill = *off;
180                        }
181                    }
182                }
183            }
184        }
185        let word_bytes = target.word_bytes();
186        s.next_stack_slot = if min_spill == 0 {
187            RiscVFrame::calculate_stack_offset(0, target)
188        } else {
189            min_spill - word_bytes
190        };
191        s
192    }
193}
194
195impl MirRegisterAllocator for RiscVRegAlloc {
196    type PhysReg = &'static str;
197
198    fn alloc_scratch(&mut self) -> Option<Self::PhysReg> {
199        self.available_gprs
200            .iter()
201            .copied()
202            .find(|reg| !self.allocated_gprs.contains_key(reg))
203    }
204
205    fn free_scratch(&mut self, phys: Self::PhysReg) {
206        self.allocated_gprs.remove(phys);
207    }
208
209    fn get_mapping(&self, vreg: &VirtualReg) -> Option<Self::PhysReg> {
210        if let Some(reg) = self.global_mappings.get(vreg) {
211            return Some(*reg);
212        }
213        for (reg, allocated_vreg) in &self.allocated_gprs {
214            if allocated_vreg == vreg {
215                return Some(*reg);
216            }
217        }
218        None
219    }
220
221    fn ensure_mapping(&mut self, vreg: VirtualReg) -> Option<Self::PhysReg> {
222        if vreg.class != RegisterClass::Gpr {
223            return None;
224        }
225
226        if let Some(phys) = self.get_mapping(&vreg) {
227            return Some(phys);
228        }
229
230        if let Some(phys) = self.alloc_scratch() {
231            self.allocated_gprs.insert(phys, vreg);
232            return Some(phys);
233        }
234
235        let stack_slot = self.next_stack_slot;
236        self.stack_slots.insert(vreg, stack_slot);
237        self.next_stack_slot -= self.target.word_bytes();
238        None
239    }
240
241    fn mapped_for_register(&self, reg: &Register) -> Option<Self::PhysReg> {
242        match reg {
243            Register::Virtual(v) => self.get_mapping(v),
244            Register::Physical(p) => Some(p.name),
245        }
246    }
247
248    fn occupy(&mut self, _phys: Self::PhysReg) {}
249
250    fn release(&mut self, phys: Self::PhysReg) {
251        self.allocated_gprs.remove(phys);
252    }
253
254    fn is_occupied(&self, phys: Self::PhysReg) -> bool {
255        self.allocated_gprs.contains_key(phys)
256    }
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    /// Registers the RISC-V emitter writes by name, plus those the ABI reserves.
264    const OFF_LIMITS: &[&str] = &[
265        "a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7", "t0", "t1", "t2", "t3", "ra", "sp", "gp",
266        "tp", "zero", "fp",
267    ];
268
269    #[test]
270    fn global_plan_preserves_reused_register_mappings() {
271        let first = VirtualReg::gpr(0);
272        let second = VirtualReg::gpr(1);
273        let plan = HashMap::from([
274            (first, Allocation::Register("t4")),
275            (second, Allocation::Register("t4")),
276        ]);
277        let allocator = RiscVRegAlloc::from_global_plan(
278            TargetOperatingSystem::Linux,
279            RiscVTarget::general(Xlen::Rv64),
280            &plan,
281            &HashMap::new(),
282        );
283        assert_eq!(allocator.get_mapping(&first), Some("t4"));
284        assert_eq!(allocator.get_mapping(&second), Some("t4"));
285    }
286
287    #[test]
288    fn pool_avoids_registers_the_emitter_uses_by_name() {
289        for reg in RiscVRegAlloc::AVAILABLE_REGISTERS {
290            assert!(
291                !OFF_LIMITS.contains(reg),
292                "{reg} is in the pool but the emitter or the ABI owns it"
293            );
294        }
295        let mut conservative = RiscVRegAlloc::default();
296        conservative.set_conservative_mode();
297        for reg in &conservative.available_gprs {
298            assert!(!OFF_LIMITS.contains(reg), "{reg} in conservative pool");
299        }
300    }
301
302    #[test]
303    fn pool_is_written_in_abi_names() {
304        for reg in RiscVRegAlloc::AVAILABLE_REGISTERS {
305            assert!(
306                !reg.starts_with('x'),
307                "{reg} is an x-number; the emitter uses ABI names"
308            );
309        }
310    }
311
312    #[test]
313    fn abi_register_file_is_fully_classified() {
314        let mut classified: Vec<&str> = Vec::new();
315        classified.extend(RiscVRegAlloc::CALLER_SAVED_REGISTERS.iter().copied());
316        classified.extend(RiscVRegAlloc::CALLEE_SAVED_REGISTERS.iter().copied());
317        classified.extend(RiscVRegAlloc::RESERVED_REGISTERS.iter().copied());
318        classified.sort_unstable();
319        classified.dedup();
320        assert_eq!(classified.len(), 32);
321        assert_eq!(
322            RiscVRegAlloc::register_convention("t3"),
323            Some(RiscVRegisterConvention::CallerSaved)
324        );
325        assert_eq!(
326            RiscVRegAlloc::register_convention("s7"),
327            Some(RiscVRegisterConvention::CalleeSaved)
328        );
329        assert_eq!(
330            RiscVRegAlloc::register_convention("ra"),
331            Some(RiscVRegisterConvention::Reserved)
332        );
333    }
334
335    #[test]
336    fn global_pool_has_caller_and_callee_saved_registers() {
337        assert!(RiscVRegAlloc::AVAILABLE_REGISTERS.iter().any(|reg| {
338            RiscVRegAlloc::register_convention(reg) == Some(RiscVRegisterConvention::CallerSaved)
339        }));
340        assert!(RiscVRegAlloc::AVAILABLE_REGISTERS.iter().any(|reg| {
341            RiscVRegAlloc::register_convention(reg) == Some(RiscVRegisterConvention::CalleeSaved)
342        }));
343    }
344
345    #[test]
346    fn rv32_spills_use_word_slots_below_the_saved_pair() {
347        let target = RiscVTarget::general(Xlen::Rv32);
348        let mut allocator = RiscVRegAlloc::with_target(TargetOperatingSystem::Linux, target);
349        for id in 0..RiscVRegAlloc::AVAILABLE_REGISTERS.len() {
350            assert!(
351                allocator
352                    .ensure_mapping(VirtualReg::gpr(id as u32))
353                    .is_some()
354            );
355        }
356        let first_spill = VirtualReg::gpr(RiscVRegAlloc::AVAILABLE_REGISTERS.len() as u32);
357        let second_spill = VirtualReg::gpr(first_spill.id + 1);
358        assert!(allocator.ensure_mapping(first_spill).is_none());
359        assert!(allocator.ensure_mapping(second_spill).is_none());
360        assert_eq!(allocator.get_stack_slot(&first_spill), Some(-20));
361        assert_eq!(allocator.get_stack_slot(&second_spill), Some(-24));
362    }
363
364    #[test]
365    fn exhausted_pool_reports_none_rather_than_a_live_register() {
366        let mut alloc = RiscVRegAlloc::default();
367        let total = alloc.available_gprs.len();
368        let mut handed_out = Vec::new();
369        for i in 0..total {
370            let phys = alloc.alloc_scratch().expect("pool should still have room");
371            assert!(!handed_out.contains(&phys), "{phys} handed out twice");
372            alloc.allocated_gprs.insert(phys, VirtualReg::gpr(i as u32));
373            handed_out.push(phys);
374        }
375        assert_eq!(alloc.alloc_scratch(), None);
376    }
377}
lamina_codegen/riscv/regalloc.rs

lamina_codegen/riscv/
regalloc.rs
