armfortas Public

Watch 0 Fork 0 Star 0
Rust · 16575 bytes Raw Blame History
1 //! Liveness analysis for Machine IR.
2 //!
3 //! Computes live intervals for each virtual register using backward
4 //! dataflow. Used by the linear scan register allocator.
5
6 use super::mir::*;
7 use std::collections::HashMap;
8
9 /// A live interval: the range of instruction positions where a vreg is live.
10 #[derive(Debug, Clone)]
11 pub struct LiveInterval {
12 pub vreg: VRegId,
13 pub class: RegClass,
14 pub start: u32, // first instruction position (definition)
15 pub end: u32, // last instruction position (final use)
16 pub crosses_call: bool, // true if a BL/BLR falls within [start, end]
17 /// Sorted positions of every BL/BLR strictly inside (start, end).
18 /// Empty when `crosses_call` is false. Populated as the
19 /// foundation for live-range splitting: knowing exactly where
20 /// an interval crosses calls lets the allocator split it at
21 /// those boundaries rather than spilling the whole range.
22 pub call_crossings: Vec<u32>,
23 /// Preferred physical register (hint). If set, the allocator tries this register first.
24 /// Used to avoid unnecessary moves (e.g., arg values prefer xN, return values prefer x0).
25 pub hint: Option<u8>,
26 }
27
28 /// Result of liveness analysis.
29 pub struct LivenessResult {
30 pub intervals: Vec<LiveInterval>,
31 /// Total number of instruction positions.
32 pub num_positions: u32,
33 }
34
35 #[derive(Clone, PartialEq, Eq)]
36 struct LiveSet {
37 words: Vec<u64>,
38 }
39
40 impl LiveSet {
41 fn new(num_vregs: usize) -> Self {
42 let num_words = num_vregs.div_ceil(64);
43 Self {
44 words: vec![0; num_words],
45 }
46 }
47
48 fn clear(&mut self) {
49 self.words.fill(0);
50 }
51
52 fn copy_from(&mut self, other: &Self) {
53 self.words.copy_from_slice(&other.words);
54 }
55
56 fn insert(&mut self, vreg: VRegId) {
57 let idx = vreg.0 as usize;
58 let word = idx / 64;
59 let bit = idx % 64;
60 if let Some(slot) = self.words.get_mut(word) {
61 *slot |= 1u64 << bit;
62 }
63 }
64
65 fn remove(&mut self, vreg: &VRegId) {
66 let idx = vreg.0 as usize;
67 let word = idx / 64;
68 let bit = idx % 64;
69 if let Some(slot) = self.words.get_mut(word) {
70 *slot &= !(1u64 << bit);
71 }
72 }
73
74 fn union_with(&mut self, other: &Self) {
75 for (dst, src) in self.words.iter_mut().zip(&other.words) {
76 *dst |= *src;
77 }
78 }
79
80 fn iter(&self) -> LiveSetIter<'_> {
81 LiveSetIter {
82 words: &self.words,
83 word_idx: 0,
84 active_word: self.words.first().copied().unwrap_or(0),
85 }
86 }
87 }
88
89 struct LiveSetIter<'a> {
90 words: &'a [u64],
91 word_idx: usize,
92 active_word: u64,
93 }
94
95 impl Iterator for LiveSetIter<'_> {
96 type Item = VRegId;
97
98 fn next(&mut self) -> Option<Self::Item> {
99 loop {
100 if self.active_word != 0 {
101 let bit = self.active_word.trailing_zeros() as usize;
102 self.active_word &= self.active_word - 1;
103 let idx = self.word_idx * 64 + bit;
104 return Some(VRegId(idx as u32));
105 }
106
107 self.word_idx += 1;
108 if self.word_idx >= self.words.len() {
109 return None;
110 }
111 self.active_word = self.words[self.word_idx];
112 }
113 }
114 }
115
116 fn machine_vreg_capacity(mf: &MachineFunction) -> usize {
117 let mut max_idx = 0usize;
118 for vreg in &mf.vregs {
119 max_idx = max_idx.max(vreg.id.0 as usize + 1);
120 }
121 for block in &mf.blocks {
122 for inst in &block.insts {
123 if let Some(def) = inst.def {
124 max_idx = max_idx.max(def.0 as usize + 1);
125 }
126 for op in &inst.operands {
127 if let MachineOperand::VReg(vreg) = op {
128 max_idx = max_idx.max(vreg.0 as usize + 1);
129 }
130 }
131 }
132 }
133 max_idx
134 }
135
136 /// Compute live intervals for all vregs in a machine function.
137 pub fn compute_liveness(mf: &MachineFunction) -> LivenessResult {
138 let num_vregs = machine_vreg_capacity(mf);
139 let block_indices: HashMap<MBlockId, usize> = mf
140 .blocks
141 .iter()
142 .enumerate()
143 .map(|(idx, block)| (block.id, idx))
144 .collect();
145
146 // Phase 1: assign a linear position to each instruction.
147 // Positions increase by 2 (leave gaps for inserted moves).
148 let mut inst_positions: Vec<Vec<u32>> = Vec::with_capacity(mf.blocks.len());
149 let mut block_start: Vec<u32> = Vec::with_capacity(mf.blocks.len());
150 let mut block_end: Vec<u32> = Vec::with_capacity(mf.blocks.len());
151 let mut pos: u32 = 0;
152
153 for block in &mf.blocks {
154 block_start.push(pos);
155 let mut block_positions = Vec::with_capacity(block.insts.len());
156 for _inst in &block.insts {
157 block_positions.push(pos);
158 pos += 2;
159 }
160 inst_positions.push(block_positions);
161 block_end.push(pos);
162 }
163 let num_positions = pos;
164
165 // Phase 3: backward dataflow to compute live-in and live-out per block.
166 let mut live_in: Vec<LiveSet> = (0..mf.blocks.len())
167 .map(|_| LiveSet::new(num_vregs))
168 .collect();
169 let mut live_out: Vec<LiveSet> = (0..mf.blocks.len())
170 .map(|_| LiveSet::new(num_vregs))
171 .collect();
172
173 // Build successor map from branch targets.
174 // Our codegen emits conditional branches as: CmpImm, BCond, B (unconditional fallback).
175 // So a block can have both BCond and B targets.
176 let mut successors: Vec<Vec<usize>> = vec![Vec::new(); mf.blocks.len()];
177 for (block_idx, block) in mf.blocks.iter().enumerate() {
178 let mut succs = Vec::new();
179 // Scan ALL instructions in the block for branch targets (not just the last).
180 for inst in &block.insts {
181 match inst.opcode {
182 ArmOpcode::B => {
183 if let Some(MachineOperand::BlockRef(target)) = inst.operands.first() {
184 if let Some(&succ_idx) = block_indices.get(target) {
185 if !succs.contains(&succ_idx) {
186 succs.push(succ_idx);
187 }
188 }
189 }
190 }
191 ArmOpcode::BCond => {
192 if let Some(MachineOperand::BlockRef(target)) = inst.operands.get(1) {
193 if let Some(&succ_idx) = block_indices.get(target) {
194 if !succs.contains(&succ_idx) {
195 succs.push(succ_idx);
196 }
197 }
198 }
199 }
200 _ => {}
201 }
202 }
203 successors[block_idx] = succs;
204 }
205
206 // Iterate until fixed point.
207 let mut changed = true;
208 let mut scratch_out = LiveSet::new(num_vregs);
209 let mut scratch_in = LiveSet::new(num_vregs);
210 while changed {
211 changed = false;
212 for block_idx in (0..mf.blocks.len()).rev() {
213 let block = &mf.blocks[block_idx];
214 // live_out = union of live_in of all successors.
215 scratch_out.clear();
216 for &succ_idx in &successors[block_idx] {
217 scratch_out.union_with(&live_in[succ_idx]);
218 }
219
220 // live_in = (live_out - defs) ∪ uses
221 scratch_in.copy_from(&scratch_out);
222 // Walk instructions backward.
223 for inst in block.insts.iter().rev() {
224 // Remove defs.
225 if let Some(def) = &inst.def {
226 scratch_in.remove(def);
227 }
228 // Add uses.
229 for op in &inst.operands {
230 if let MachineOperand::VReg(vid) = op {
231 scratch_in.insert(*vid);
232 }
233 }
234 }
235
236 if live_in[block_idx] != scratch_in {
237 live_in[block_idx].copy_from(&scratch_in);
238 changed = true;
239 }
240 if live_out[block_idx] != scratch_out {
241 live_out[block_idx].copy_from(&scratch_out);
242 changed = true;
243 }
244 }
245 }
246
247 // Phase 4: build live intervals from def/use positions.
248 let mut starts: Vec<Option<u32>> = vec![None; num_vregs];
249 let mut ends: Vec<Option<u32>> = vec![None; num_vregs];
250
251 for (block_idx, block) in mf.blocks.iter().enumerate() {
252 let b_start = block_start[block_idx];
253 let b_end = block_end[block_idx];
254
255 // Vregs live-in to this block: extend their interval to block start.
256 for vreg in live_in[block_idx].iter() {
257 let idx = vreg.0 as usize;
258 ends[idx] = Some(ends[idx].map_or(b_end, |end| end.max(b_end)));
259 starts[idx] = Some(starts[idx].map_or(b_start, |start| start.min(b_start)));
260 }
261
262 // Vregs live-out of this block: extend their interval to block end.
263 for vreg in live_out[block_idx].iter() {
264 let idx = vreg.0 as usize;
265 ends[idx] = Some(ends[idx].map_or(b_end, |end| end.max(b_end)));
266 starts[idx] = Some(starts[idx].map_or(b_start, |start| start.min(b_start)));
267 }
268
269 // Walk instructions for precise def/use positions.
270 for (i, inst) in block.insts.iter().enumerate() {
271 let p = inst_positions[block_idx][i];
272 if let Some(def) = &inst.def {
273 let idx = def.0 as usize;
274 starts[idx] = Some(starts[idx].map_or(p, |start| start.min(p)));
275 ends[idx] = Some(ends[idx].map_or(p, |end| end.max(p)));
276 }
277 for op in &inst.operands {
278 if let MachineOperand::VReg(vid) = op {
279 let idx = vid.0 as usize;
280 ends[idx] = Some(ends[idx].map_or(p, |end| end.max(p)));
281 starts[idx] = Some(starts[idx].map_or(p, |start| start.min(p)));
282 }
283 }
284 }
285 }
286
287 // Collect positions of all call instructions. Both direct
288 // (`Bl`) and indirect (`Blr`) calls clobber every caller-saved
289 // register and must contribute to `crosses_call` so the
290 // allocator pins those intervals to callee-saved.
291 let mut call_positions: Vec<u32> = Vec::new();
292 for (block_idx, block) in mf.blocks.iter().enumerate() {
293 for (i, inst) in block.insts.iter().enumerate() {
294 if matches!(inst.opcode, ArmOpcode::Bl | ArmOpcode::Blr) {
295 call_positions.push(inst_positions[block_idx][i]);
296 }
297 }
298 }
299 call_positions.sort();
300
301 // Build intervals.
302 let mut vreg_classes = vec![RegClass::Gp64; num_vregs];
303 for vreg in &mf.vregs {
304 vreg_classes[vreg.id.0 as usize] = vreg.class;
305 }
306
307 let mut intervals: Vec<LiveInterval> = starts
308 .iter()
309 .enumerate()
310 .filter_map(|(idx, start)| {
311 let start = (*start)?;
312 let end = ends[idx]?;
313 let vreg = VRegId(idx as u32);
314 let class = vreg_classes[idx];
315 // Collect every call within (start, end). These are
316 // candidate split points for the live-range splitter
317 // when callee-saved is exhausted; for today's allocator
318 // the boolean flag is what's consumed, but the
319 // splitter (sprint 18) needs the per-position list.
320 let call_crossings: Vec<u32> = call_positions
321 .iter()
322 .copied()
323 .filter(|&cp| cp > start && cp < end)
324 .collect();
325 let crosses_call = !call_crossings.is_empty();
326 Some(LiveInterval {
327 vreg,
328 class,
329 start,
330 end,
331 crosses_call,
332 call_crossings,
333 hint: None,
334 })
335 })
336 .collect();
337
338 // Sort by start position, breaking ties on the vreg ID. The
339 // tie-break is critical: linear-scan register allocation
340 // processes intervals in vector order, and `starts` is a
341 // HashMap whose iteration order varies between runs. Without
342 // a deterministic tie-break, two compiles of the same source
343 // could pick different registers and produce different
344 // assembly — surfaced as the long-standing select_type.f90
345 // intermittent failure during the optimizer audit rounds and
346 // exposed clearly when mem2reg started perturbing vreg counts.
347 intervals.sort_by_key(|i| (i.start, i.vreg.0));
348
349 LivenessResult {
350 intervals,
351 num_positions,
352 }
353 }
354
355 #[cfg(test)]
356 mod tests {
357 use super::*;
358 use crate::codegen::isel::select_function;
359 use crate::ir::builder::FuncBuilder;
360 use crate::ir::inst::*;
361 use crate::ir::types::*;
362
363 #[test]
364 fn liveness_simple_function() {
365 let mut func = Function::new("test".into(), vec![], IrType::Void);
366 {
367 let mut b = FuncBuilder::new(&mut func);
368 let x = b.const_i32(10);
369 let y = b.const_i32(20);
370 let _z = b.iadd(x, y);
371 b.ret_void();
372 }
373 let mf = select_function(&func);
374 let result = compute_liveness(&mf);
375 assert!(!result.intervals.is_empty(), "should have live intervals");
376 assert!(result.num_positions > 0);
377 }
378
379 #[test]
380 fn liveness_intervals_have_correct_class() {
381 let mut func = Function::new("test".into(), vec![], IrType::Void);
382 {
383 let mut b = FuncBuilder::new(&mut func);
384 let a = b.const_f64(3.14);
385 let c = b.const_f64(2.0);
386 let _d = b.fadd(a, c);
387 b.ret_void();
388 }
389 let mf = select_function(&func);
390 let result = compute_liveness(&mf);
391 // Should have some Fp64 intervals.
392 assert!(
393 result.intervals.iter().any(|i| i.class == RegClass::Fp64),
394 "should have Fp64 intervals"
395 );
396 }
397
398 #[test]
399 fn intervals_crossing_blr_are_marked_call_crossing() {
400 // Construct: def vreg %0 → BLR %0 → use %0. Since %0 is
401 // both the BLR target and a value live across the call,
402 // its interval must be marked crosses_call so the
403 // allocator pins it to a callee-saved register (or
404 // accepts the spill rather than burying the value in a
405 // caller-saved that the BLR will clobber).
406 let mut mf = MachineFunction::new("test".into());
407 let v0 = mf.new_vreg(RegClass::Gp64);
408 let v1 = mf.new_vreg(RegClass::Gp64);
409 // Define %0.
410 mf.blocks[0].insts.push(MachineInst {
411 opcode: ArmOpcode::Movz,
412 operands: vec![MachineOperand::VReg(v0), MachineOperand::Imm(42)],
413 def: Some(v0),
414 });
415 // BLR through %0 (an indirect call).
416 mf.blocks[0].insts.push(MachineInst {
417 opcode: ArmOpcode::Blr,
418 operands: vec![MachineOperand::VReg(v0)],
419 def: None,
420 });
421 // Use %0 again post-call (forces the interval to cross
422 // the BLR position).
423 mf.blocks[0].insts.push(MachineInst {
424 opcode: ArmOpcode::MovReg,
425 operands: vec![MachineOperand::VReg(v1), MachineOperand::VReg(v0)],
426 def: Some(v1),
427 });
428 mf.blocks[0].insts.push(MachineInst {
429 opcode: ArmOpcode::Ret,
430 operands: vec![],
431 def: None,
432 });
433 let result = compute_liveness(&mf);
434 let interval = result
435 .intervals
436 .iter()
437 .find(|i| i.vreg == v0)
438 .expect("vreg 0 must have an interval");
439 assert!(
440 interval.crosses_call,
441 "interval crossing BLR must be marked crosses_call"
442 );
443 assert_eq!(
444 interval.call_crossings.len(),
445 1,
446 "interval should record exactly one BLR crossing, got {:?}",
447 interval.call_crossings
448 );
449 }
450
451 #[test]
452 fn liveness_branching() {
453 let mut func = Function::new("test".into(), vec![], IrType::Void);
454 {
455 let mut b = FuncBuilder::new(&mut func);
456 let cond = b.const_bool(true);
457 let bb_t = b.create_block("then");
458 let bb_f = b.create_block("else");
459 b.cond_branch(cond, bb_t, vec![], bb_f, vec![]);
460 b.set_block(bb_t);
461 b.ret_void();
462 b.set_block(bb_f);
463 b.ret_void();
464 }
465 let mf = select_function(&func);
466 let result = compute_liveness(&mf);
467 assert!(result.num_positions > 0);
468 }
469 }
470