armfortas Public

Watch 0 Fork 0 Star 0
Rust · 10463 bytes Raw Blame History
1 //! Naive register allocation — spill everything.
2 //!
3 //! Every virtual register gets a stack slot. Before each use, load from
4 //! the slot into a scratch register. After each def, store from the
5 //! scratch register to the slot. This produces correct but slow code.
6 //!
7 //! Scratch registers used:
8 //! - x9, x10, x11 for integer operands (caller-saved temporaries)
9 //! - d16, d17, d18 for float operands (caller-saved temporaries)
10 //!
11 //! This will be replaced by linear scan allocation in Sprint 21.
12
13 use super::mir::*;
14 use std::collections::HashMap;
15
16 /// Integer scratch registers (caller-saved, safe to clobber).
17 const GP_SCRATCH: [u8; 3] = [9, 10, 11];
18 /// Float scratch registers (caller-saved, safe to clobber).
19 const FP_SCRATCH: [u8; 3] = [16, 17, 18];
20
21 /// Allocate registers for a machine function using spill-everything strategy.
22 /// Modifies the function in place: replaces VReg operands with PhysReg,
23 /// inserts loads/stores around each instruction.
24 pub fn regalloc_naive(mf: &mut MachineFunction) {
25 // Phase 1: assign a stack slot to every vreg.
26 let mut vreg_slots: HashMap<VRegId, i32> = HashMap::new();
27 for vreg in &mf.vregs {
28 let size = match vreg.class {
29 RegClass::Gp64 => 8,
30 RegClass::Gp32 => 4,
31 RegClass::Fp64 => 8,
32 RegClass::Fp32 => 4,
33 RegClass::V128 => 16,
34 };
35 let offset = mf.frame.alloc_local(size);
36 vreg_slots.insert(vreg.id, offset);
37 }
38
39 // Build class map for quick lookup.
40 let vreg_classes: HashMap<VRegId, RegClass> =
41 mf.vregs.iter().map(|v| (v.id, v.class)).collect();
42
43 // Phase 2: rewrite each block's instructions.
44 for block_idx in 0..mf.blocks.len() {
45 let mut new_insts = Vec::new();
46
47 let insts = std::mem::take(&mut mf.blocks[block_idx].insts);
48 for inst in insts {
49 let mut rewritten = inst.clone();
50
51 // Collect which vreg operands need loading (inputs) and which
52 // need storing (output/def).
53 let mut loads: Vec<(usize, VRegId)> = Vec::new(); // (operand_idx, vreg)
54 let def_vreg = inst.def;
55
56 // Identify vreg operands that are inputs (not the def).
57 for (i, op) in inst.operands.iter().enumerate() {
58 if let MachineOperand::VReg(vid) = op {
59 // If this operand is the same as the def and it's the first
60 // operand (destination), it's an output — don't load.
61 if Some(*vid) == def_vreg && i == 0 {
62 continue;
63 }
64 loads.push((i, *vid));
65 }
66 }
67
68 // Emit loads for input operands.
69 let mut gp_scratch_idx = 0;
70 let mut fp_scratch_idx = 0;
71
72 for (op_idx, vid) in &loads {
73 if let Some(&offset) = vreg_slots.get(vid) {
74 let class = vreg_classes.get(vid).copied().unwrap_or(RegClass::Gp64);
75 let (scratch, load_op) = match class {
76 RegClass::Fp64 => {
77 let s = FP_SCRATCH[fp_scratch_idx % FP_SCRATCH.len()];
78 fp_scratch_idx += 1;
79 (PhysReg::Fp(s), ArmOpcode::LdrFpImm)
80 }
81 RegClass::Fp32 => {
82 let s = FP_SCRATCH[fp_scratch_idx % FP_SCRATCH.len()];
83 fp_scratch_idx += 1;
84 (PhysReg::Fp32(s), ArmOpcode::LdrFpImm)
85 }
86 RegClass::V128 => {
87 // 128-bit vector spill/fill uses LdrQ / StrQ
88 // off the FP-scratch pool — same physical
89 // register bank as Fp{32,64}, just at 128b
90 // width.
91 let s = FP_SCRATCH[fp_scratch_idx % FP_SCRATCH.len()];
92 fp_scratch_idx += 1;
93 (PhysReg::Fp(s), ArmOpcode::LdrQ)
94 }
95 RegClass::Gp32 => {
96 let s = GP_SCRATCH[gp_scratch_idx % GP_SCRATCH.len()];
97 gp_scratch_idx += 1;
98 (PhysReg::Gp32(s), ArmOpcode::LdrImm)
99 }
100 RegClass::Gp64 => {
101 let s = GP_SCRATCH[gp_scratch_idx % GP_SCRATCH.len()];
102 gp_scratch_idx += 1;
103 (PhysReg::Gp(s), ArmOpcode::LdrImm)
104 }
105 };
106
107 // LDR scratch, [FP, #offset]
108 new_insts.push(MachineInst {
109 opcode: load_op,
110 operands: vec![
111 MachineOperand::PhysReg(scratch),
112 MachineOperand::PhysReg(PhysReg::FP),
113 MachineOperand::Imm(offset as i64),
114 ],
115 def: None,
116 });
117
118 // Replace the vreg operand with the scratch register.
119 rewritten.operands[*op_idx] = MachineOperand::PhysReg(scratch);
120 }
121 }
122
123 // Replace the def operand (first operand if it matches def).
124 let def_scratch = if let Some(def_vid) = def_vreg {
125 if let Some(&offset) = vreg_slots.get(&def_vid) {
126 let class = vreg_classes
127 .get(&def_vid)
128 .copied()
129 .unwrap_or(RegClass::Gp64);
130 let scratch = match class {
131 RegClass::Fp64 => PhysReg::Fp(FP_SCRATCH[0]),
132 RegClass::Fp32 => PhysReg::Fp32(FP_SCRATCH[0]),
133 RegClass::V128 => PhysReg::Fp(FP_SCRATCH[0]),
134 RegClass::Gp32 => PhysReg::Gp32(GP_SCRATCH[0]),
135 RegClass::Gp64 => PhysReg::Gp(GP_SCRATCH[0]),
136 };
137
138 // Replace def operand.
139 if let Some(MachineOperand::VReg(vid)) = rewritten.operands.first() {
140 if *vid == def_vid {
141 rewritten.operands[0] = MachineOperand::PhysReg(scratch);
142 }
143 }
144
145 Some((scratch, offset, class))
146 } else {
147 None
148 }
149 } else {
150 None
151 };
152
153 // Emit the rewritten instruction.
154 rewritten.def = None; // physical regs don't track defs
155 new_insts.push(rewritten);
156
157 // Emit store for the def.
158 if let Some((scratch, offset, class)) = def_scratch {
159 let store_op = match class {
160 RegClass::Fp32 | RegClass::Fp64 => ArmOpcode::StrFpImm,
161 RegClass::V128 => ArmOpcode::StrQ,
162 RegClass::Gp32 | RegClass::Gp64 => ArmOpcode::StrImm,
163 };
164 new_insts.push(MachineInst {
165 opcode: store_op,
166 operands: vec![
167 MachineOperand::PhysReg(scratch),
168 MachineOperand::PhysReg(PhysReg::FP),
169 MachineOperand::Imm(offset as i64),
170 ],
171 def: None,
172 });
173 }
174 }
175
176 mf.blocks[block_idx].insts = new_insts;
177 }
178 }
179
180 #[cfg(test)]
181 mod tests {
182 use super::*;
183 use crate::codegen::isel::select_function;
184 use crate::ir::builder::FuncBuilder;
185 use crate::ir::inst::*;
186 use crate::ir::types::*;
187
188 #[test]
189 fn regalloc_replaces_vregs() {
190 let mut func = Function::new("test".into(), vec![], IrType::Void);
191 {
192 let mut b = FuncBuilder::new(&mut func);
193 let x = b.const_i32(42);
194 let y = b.const_i32(10);
195 let _z = b.iadd(x, y);
196 b.ret_void();
197 }
198 let mut mf = select_function(&func);
199 regalloc_naive(&mut mf);
200
201 // After regalloc, no VReg operands should remain.
202 for block in &mf.blocks {
203 for inst in &block.insts {
204 for op in &inst.operands {
205 assert!(
206 !matches!(op, MachineOperand::VReg(_)),
207 "vreg still present after regalloc: {:?}",
208 inst
209 );
210 }
211 }
212 }
213 }
214
215 #[test]
216 fn regalloc_frame_grows() {
217 let mut func = Function::new("test".into(), vec![], IrType::Void);
218 {
219 let mut b = FuncBuilder::new(&mut func);
220 b.const_i32(1);
221 b.const_i32(2);
222 b.const_i32(3);
223 b.ret_void();
224 }
225 let mut mf = select_function(&func);
226 let before = mf.frame.size;
227 regalloc_naive(&mut mf);
228 assert!(
229 mf.frame.size >= before,
230 "frame should grow to accommodate spill slots"
231 );
232 }
233
234 #[test]
235 fn regalloc_uses_scratch_registers() {
236 let mut func = Function::new("test".into(), vec![], IrType::Void);
237 {
238 let mut b = FuncBuilder::new(&mut func);
239 let x = b.const_i32(10);
240 let y = b.const_i32(20);
241 let _z = b.iadd(x, y);
242 b.ret_void();
243 }
244 let mut mf = select_function(&func);
245 regalloc_naive(&mut mf);
246
247 // Should use x9, x10, x11 as scratch registers.
248 let uses_scratch = mf.blocks.iter().any(|b| {
249 b.insts.iter().any(|i| {
250 i.operands.iter().any(|op| {
251 matches!(
252 op,
253 MachineOperand::PhysReg(PhysReg::Gp(9))
254 | MachineOperand::PhysReg(PhysReg::Gp(10))
255 | MachineOperand::PhysReg(PhysReg::Gp(11))
256 | MachineOperand::PhysReg(PhysReg::Gp32(9))
257 | MachineOperand::PhysReg(PhysReg::Gp32(10))
258 | MachineOperand::PhysReg(PhysReg::Gp32(11))
259 )
260 })
261 })
262 });
263 assert!(
264 uses_scratch,
265 "should use scratch registers x9-x11 or w9-w11"
266 );
267 }
268 }
269