From 10be86157829df54f9cd04c98538f7ba5255e5eb Mon Sep 17 00:00:00 2001 From: Andrew Brown Date: Wed, 17 Sep 2025 11:23:15 -0700 Subject: [PATCH] Refactor `RegTraversalIter` This is a non-functional change to refactor `RegTraversalIter` to make it more amenable for future commits adding limits to which registers can be allocated. This change does change the API of `RegTraversalIter`, though. It boils down to several parts: - using `Cursor` internally: much of the cursor-advancing logic is duplicated between the preferred and non-preferred register list. This change adds a new, private struct--`Cursor`--to handle this. This refactoring also reduces the scope of the borrow we take down to only what is necessary (i.e., the actual register lists). - changing parameter order: this also passes through the arguments constructing a `RegTraversalIter` in the order they will be checked during iteration: first it emits the fixed register, then the hinted register, then the offset-based lookup. This is a clarity tweak. - only allow one hint register: `RegTraversalIter` allowed for two hint registers but the library never uses the second one. This change switches to only using a single register, with no other change in behavior. - minor simplifications: e.g., we can use `Option::take` to avoid a more verbose access (and removal) of the fixed register. Another possible change here would be to pass the hinted register as an `Option` with `None` to indicate no hint. This seems more conventional and a quick audit of the code supports this but I avoided this out of an abundance of caution. --- src/ion/moves.rs | 10 +-- src/ion/process.rs | 9 +- src/ion/reg_traversal.rs | 176 +++++++++++++++++++++------------------ src/ion/spill.rs | 8 +- 4 files changed, 101 insertions(+), 102 deletions(-) diff --git a/src/ion/moves.rs b/src/ion/moves.rs index 22685e53..c2f56ac2 100644 --- a/src/ion/moves.rs +++ b/src/ion/moves.rs @@ -892,14 +892,8 @@ impl<'a, F: Function> Env<'a, F> { } let resolved = parallel_moves.resolve(); - let mut scratch_iter = RegTraversalIter::new( - self.env, - regclass, - PReg::invalid(), - PReg::invalid(), - 0, - None, - ); + let mut scratch_iter = + RegTraversalIter::new(self.env, regclass, None, PReg::invalid(), 0); let mut dedicated_scratch = self.env.scratch_by_class[regclass as usize]; let key = LiveRangeKey::from_range(&CodeRange { from: pos_prio.pos, diff --git a/src/ion/process.rs b/src/ion/process.rs index 2d174815..982f7e93 100644 --- a/src/ion/process.rs +++ b/src/ion/process.rs @@ -1078,14 +1078,7 @@ impl<'a, F: Function> Env<'a, F> { + bundle.index(); self.ctx.output.stats.process_bundle_reg_probe_start_any += 1; - for preg in RegTraversalIter::new( - self.env, - class, - hint_reg, - PReg::invalid(), - scan_offset, - fixed_preg, - ) { + for preg in RegTraversalIter::new(self.env, class, fixed_preg, hint_reg, scan_offset) { self.ctx.output.stats.process_bundle_reg_probes_any += 1; let preg_idx = PRegIndex::new(preg.index()); trace!("trying preg {:?}", preg_idx); diff --git a/src/ion/reg_traversal.rs b/src/ion/reg_traversal.rs index 729fd33e..eb4ef6b9 100644 --- a/src/ion/reg_traversal.rs +++ b/src/ion/reg_traversal.rs @@ -1,79 +1,104 @@ +//! Iterate over available registers. + use crate::{MachineEnv, PReg, RegClass}; -/// This iterator represents a traversal through all allocatable -/// registers of a given class, in a certain order designed to -/// minimize allocation contention. +/// Keep track of where we are in the register traversal. +struct Cursor<'a> { + registers: &'a [PReg], + index: usize, + offset: usize, +} + +impl<'a> Cursor<'a> { + #[inline] + fn new(registers: &'a [PReg], offset_hint: usize) -> Self { + let offset = if registers.len() > 0 { + offset_hint % registers.len() + } else { + 0 + }; + Self { + registers, + index: 0, + offset, + } + } + + /// Wrap around the end of the register list; [`Cursor::done`] guarantees we + /// do not see the same register twice. + #[inline] + fn wrap(index: usize, end: usize) -> usize { + if index >= end { + index - end + } else { + index + } + } + + /// Advance to the next register and return it. + #[inline] + fn advance(&mut self) -> PReg { + let loc = Self::wrap(self.index + self.offset, self.registers.len()); + let reg = self.registers[loc]; + self.index += 1; + reg + } + + /// Return `true` if we have seen all registers. + #[inline] + fn done(&self) -> bool { + self.index >= self.registers.len() + } +} + +/// This iterator represents a traversal through all allocatable registers of a +/// given class, in a certain order designed to minimize allocation contention. /// /// The order in which we try registers is somewhat complex: -/// - First, if there is a hint, we try that. -/// - Then, we try registers in a traversal order that is based on an -/// "offset" (usually the bundle index) spreading pressure evenly -/// among registers to reduce commitment-map contention. -/// - Within that scan, we try registers in two groups: first, -/// prferred registers; then, non-preferred registers. (In normal -/// usage, these consist of caller-save and callee-save registers -/// respectively, to minimize clobber-saves; but they need not.) - +/// - First, if the register is fixed (i.e., pre-assigned), return that and stop +/// iteration. +/// - Then, if there is a hint, try that one. +/// - Next we try registers in a traversal order that is based on an "offset" +/// (usually the bundle index) spreading pressure evenly among registers to +/// reduce commitment-map contention. +/// - Within that scan, we try registers in two groups: first, preferred +/// registers; then, non-preferred registers. (In normal usage, these consist +/// of caller-save and callee-save registers respectively, to minimize +/// clobber-saves; but they need not.) pub struct RegTraversalIter<'a> { - env: &'a MachineEnv, - class: usize, - hints: [Option; 2], - hint_idx: usize, - pref_idx: usize, - non_pref_idx: usize, - offset_pref: usize, - offset_non_pref: usize, is_fixed: bool, fixed: Option, + use_hint: bool, + hint: Option, + preferred: Cursor<'a>, + non_preferred: Cursor<'a>, } impl<'a> RegTraversalIter<'a> { pub fn new( env: &'a MachineEnv, class: RegClass, - hint_reg: PReg, - hint2_reg: PReg, - offset: usize, fixed: Option, + hint: PReg, + offset: usize, ) -> Self { - let mut hint_reg = if hint_reg != PReg::invalid() { - Some(hint_reg) - } else { - None - }; - let mut hint2_reg = if hint2_reg != PReg::invalid() { - Some(hint2_reg) + let hint = if hint != PReg::invalid() { + Some(hint) } else { None }; - if hint_reg.is_none() { - hint_reg = hint2_reg; - hint2_reg = None; - } - let hints = [hint_reg, hint2_reg]; let class = class as u8 as usize; - let offset_pref = if env.preferred_regs_by_class[class].len() > 0 { - offset % env.preferred_regs_by_class[class].len() - } else { - 0 - }; - let offset_non_pref = if env.non_preferred_regs_by_class[class].len() > 0 { - offset % env.non_preferred_regs_by_class[class].len() - } else { - 0 - }; + let preferred = Cursor::new(&env.preferred_regs_by_class[class], offset); + let non_preferred = Cursor::new(&env.non_preferred_regs_by_class[class], offset); + Self { - env, - class, - hints, - hint_idx: 0, - pref_idx: 0, - non_pref_idx: 0, - offset_pref, - offset_non_pref, is_fixed: fixed.is_some(), fixed, + use_hint: hint.is_some(), + hint, + preferred, + non_preferred, } } } @@ -83,41 +108,30 @@ impl<'a> core::iter::Iterator for RegTraversalIter<'a> { fn next(&mut self) -> Option { if self.is_fixed { - let ret = self.fixed; - self.fixed = None; - return ret; + return self.fixed.take(); } - fn wrap(idx: usize, limit: usize) -> usize { - if idx >= limit { - idx - limit - } else { - idx - } + if self.use_hint { + self.use_hint = false; + return self.hint; } - if self.hint_idx < 2 && self.hints[self.hint_idx].is_some() { - let h = self.hints[self.hint_idx]; - self.hint_idx += 1; - return h; - } - while self.pref_idx < self.env.preferred_regs_by_class[self.class].len() { - let arr = &self.env.preferred_regs_by_class[self.class][..]; - let r = arr[wrap(self.pref_idx + self.offset_pref, arr.len())]; - self.pref_idx += 1; - if Some(r) == self.hints[0] || Some(r) == self.hints[1] { - continue; + + while !self.preferred.done() { + let reg = self.preferred.advance(); + if Some(reg) == self.hint { + continue; // Try again; we already tried the hint. } - return Some(r); + return Some(reg); } - while self.non_pref_idx < self.env.non_preferred_regs_by_class[self.class].len() { - let arr = &self.env.non_preferred_regs_by_class[self.class][..]; - let r = arr[wrap(self.non_pref_idx + self.offset_non_pref, arr.len())]; - self.non_pref_idx += 1; - if Some(r) == self.hints[0] || Some(r) == self.hints[1] { - continue; + + while !self.non_preferred.done() { + let reg = self.non_preferred.advance(); + if Some(reg) == self.hint { + continue; // Try again; we already tried the hint. } - return Some(r); + return Some(reg); } + None } } diff --git a/src/ion/spill.rs b/src/ion/spill.rs index c40efc80..61d8b70f 100644 --- a/src/ion/spill.rs +++ b/src/ion/spill.rs @@ -13,8 +13,8 @@ //! Spillslot allocation. use super::{ - AllocRegResult, Env, LiveRangeKey, PReg, PRegIndex, RegTraversalIter, SpillSetIndex, - SpillSlotData, SpillSlotIndex, + AllocRegResult, Env, LiveRangeKey, PRegIndex, RegTraversalIter, SpillSetIndex, SpillSlotData, + SpillSlotIndex, }; use crate::{Allocation, Function, SpillSlot}; @@ -40,9 +40,7 @@ impl<'a, F: Function> Env<'a, F> { let mut success = false; self.ctx.output.stats.spill_bundle_reg_probes += 1; - for preg in - RegTraversalIter::new(self.env, class, hint, PReg::invalid(), bundle.index(), None) - { + for preg in RegTraversalIter::new(self.env, class, None, hint, bundle.index()) { trace!("trying bundle {:?} to preg {:?}", bundle, preg); let preg_idx = PRegIndex::new(preg.index()); if let AllocRegResult::Allocated(_) =