From 9c81860d848bf4c50582c37955dc80d520bc8c2b Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Mon, 19 Jan 2026 17:21:51 -0800 Subject: [PATCH 1/2] Store `PRegSet` in `MachineEnv` Helps to make the `MachineEnv`s in Wasmtime `const`-allocatable. --- src/fastalloc/lru.rs | 7 ++-- src/fastalloc/mod.rs | 24 ++++--------- src/fastalloc/tests.rs | 6 ++-- src/fuzzing/func.rs | 8 ++--- src/ion/liveranges.rs | 5 ++- src/ion/process.rs | 4 +-- src/ion/reg_traversal.rs | 78 ++++++++++++++-------------------------- src/lib.rs | 64 ++++++++++++++++++++++++++++----- 8 files changed, 104 insertions(+), 92 deletions(-) diff --git a/src/fastalloc/lru.rs b/src/fastalloc/lru.rs index 392539db..d86598a3 100644 --- a/src/fastalloc/lru.rs +++ b/src/fastalloc/lru.rs @@ -28,7 +28,8 @@ pub struct LruNode { } impl Lru { - pub fn new(regclass: RegClass, regs: &[PReg]) -> Self { + pub fn new(regclass: RegClass, regs: &PRegSet) -> Self { + let regs = regs.into_iter().collect::>(); let mut data = vec![ LruNode { prev: u8::MAX, @@ -248,7 +249,7 @@ impl fmt::Debug for Lru { while node != self.head { if seen.contains(&node) { panic!( - "The {:?} LRU is messed up: + "The {:?} LRU is messed up: head: {:?}, {:?} -> p{node}, actual data: {:?}", self.regclass, self.head, data_str, self.data ); @@ -298,7 +299,7 @@ impl PartialEq for PartedByRegClass { pub type Lrus = PartedByRegClass; impl Lrus { - pub fn new(int_regs: &[PReg], float_regs: &[PReg], vec_regs: &[PReg]) -> Self { + pub fn new(int_regs: &PRegSet, float_regs: &PRegSet, vec_regs: &PRegSet) -> Self { Self { items: [ Lru::new(RegClass::Int, int_regs), diff --git a/src/fastalloc/mod.rs b/src/fastalloc/mod.rs index e8ea1ad9..91fdaa75 100644 --- a/src/fastalloc/mod.rs +++ b/src/fastalloc/mod.rs @@ -444,21 +444,9 @@ impl<'a, F: Function> Env<'a, F> { env.preferred_regs_by_class[RegClass::Float as usize].clone(), env.preferred_regs_by_class[RegClass::Vector as usize].clone(), ]; - regs[0].extend( - env.non_preferred_regs_by_class[RegClass::Int as usize] - .iter() - .cloned(), - ); - regs[1].extend( - env.non_preferred_regs_by_class[RegClass::Float as usize] - .iter() - .cloned(), - ); - regs[2].extend( - env.non_preferred_regs_by_class[RegClass::Vector as usize] - .iter() - .cloned(), - ); + regs[0].union_from(env.non_preferred_regs_by_class[RegClass::Int as usize]); + regs[1].union_from(env.non_preferred_regs_by_class[RegClass::Float as usize]); + regs[2].union_from(env.non_preferred_regs_by_class[RegClass::Vector as usize]); let allocatable_regs = PRegSet::from(env); let num_available_pregs: PartedByRegClass = PartedByRegClass { items: [ @@ -508,9 +496,9 @@ impl<'a, F: Function> Env<'a, F> { ], preferred_victim: PartedByRegClass { items: [ - regs[0].last().cloned().unwrap_or(PReg::invalid()), - regs[1].last().cloned().unwrap_or(PReg::invalid()), - regs[2].last().cloned().unwrap_or(PReg::invalid()), + regs[0].max_preg().unwrap_or(PReg::invalid()), + regs[1].max_preg().unwrap_or(PReg::invalid()), + regs[2].max_preg().unwrap_or(PReg::invalid()), ], }, reused_input_to_reuse_op: vec![usize::MAX; max_operand_len as usize], diff --git a/src/fastalloc/tests.rs b/src/fastalloc/tests.rs index 41be3771..26767aa7 100644 --- a/src/fastalloc/tests.rs +++ b/src/fastalloc/tests.rs @@ -146,10 +146,10 @@ fn mach_env(no_of_regs: usize) -> MachineEnv { (0..no_of_regs) .map(|no| PReg::new(no, RegClass::Int)) .collect(), - vec![], - vec![], + PRegSet::empty(), + PRegSet::empty(), ], - non_preferred_regs_by_class: [vec![], vec![], vec![]], + non_preferred_regs_by_class: [PRegSet::empty(); 3], scratch_by_class: [None, None, None], fixed_stack_slots: vec![], } diff --git a/src/fuzzing/func.rs b/src/fuzzing/func.rs index 1ba234ba..b2a4817f 100644 --- a/src/fuzzing/func.rs +++ b/src/fuzzing/func.rs @@ -696,20 +696,20 @@ impl core::fmt::Debug for Func { } pub fn machine_env() -> MachineEnv { - fn regs(r: core::ops::Range, c: RegClass) -> Vec { + fn regs(r: core::ops::Range, c: RegClass) -> PRegSet { r.map(|i| PReg::new(i, c)).collect() } - let preferred_regs_by_class: [Vec; 3] = [ + let preferred_regs_by_class = [ regs(0..24, RegClass::Int), regs(0..24, RegClass::Float), regs(0..24, RegClass::Vector), ]; - let non_preferred_regs_by_class: [Vec; 3] = [ + let non_preferred_regs_by_class = [ regs(24..32, RegClass::Int), regs(24..32, RegClass::Float), regs(24..32, RegClass::Vector), ]; - let scratch_by_class: [Option; 3] = [None, None, None]; + let scratch_by_class = [None, None, None]; let fixed_stack_slots = (32..63) .flat_map(|i| { [ diff --git a/src/ion/liveranges.rs b/src/ion/liveranges.rs index 1162740f..449f2626 100644 --- a/src/ion/liveranges.rs +++ b/src/ion/liveranges.rs @@ -115,9 +115,8 @@ impl<'a, F: Function> Env<'a, F> { } for class in 0..self.preferred_victim_by_class.len() { self.preferred_victim_by_class[class] = self.env.non_preferred_regs_by_class[class] - .last() - .or(self.env.preferred_regs_by_class[class].last()) - .cloned() + .max_preg() + .or(self.env.preferred_regs_by_class[class].max_preg()) .unwrap_or(PReg::invalid()); } // Create VRegs from the vreg count. diff --git a/src/ion/process.rs b/src/ion/process.rs index a3cbc05a..94009191 100644 --- a/src/ion/process.rs +++ b/src/ion/process.rs @@ -1219,8 +1219,8 @@ impl<'a, F: Function> Env<'a, F> { let mut fixed_assigned = 0; let mut total_regs = 0; for preg in self.env.preferred_regs_by_class[class as u8 as usize] - .iter() - .chain(self.env.non_preferred_regs_by_class[class as u8 as usize].iter()) + .into_iter() + .chain(self.env.non_preferred_regs_by_class[class as u8 as usize]) { trace!(" -> PR {:?}", preg); let start = LiveRangeKey::from_range(&CodeRange { diff --git a/src/ion/reg_traversal.rs b/src/ion/reg_traversal.rs index 3659fe97..4c968772 100644 --- a/src/ion/reg_traversal.rs +++ b/src/ion/reg_traversal.rs @@ -1,53 +1,31 @@ //! Iterate over available registers. -use crate::{MachineEnv, PReg, RegClass}; +use crate::{MachineEnv, PReg, PRegSet, PRegSetIter, RegClass}; /// Keep track of where we are in the register traversal. -struct Cursor<'a> { - registers: &'a [PReg], - index: usize, - offset: usize, +struct Cursor { + first: PRegSetIter, + second: PRegSetIter, } -impl<'a> Cursor<'a> { +impl Cursor { #[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 - } + fn new(registers: &PRegSet, offset_hint: usize) -> Self { + let first = registers + .into_iter() + .skip(offset_hint) + .collect::() + .into_iter(); + let second = registers + .into_iter() + .take(offset_hint) + .collect::() + .into_iter(); + Self { first, second } } - /// 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() + fn next(&mut self) -> Option { + self.first.next().or_else(|| self.second.next()) } } @@ -65,19 +43,19 @@ impl<'a> Cursor<'a> { /// 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> { +pub struct RegTraversalIter { is_fixed: bool, fixed: Option, use_hint: bool, hint: Option, - preferred: Cursor<'a>, - non_preferred: Cursor<'a>, + preferred: Cursor, + non_preferred: Cursor, limit: Option, } -impl<'a> RegTraversalIter<'a> { +impl RegTraversalIter { pub fn new( - env: &'a MachineEnv, + env: &MachineEnv, class: RegClass, fixed: Option, hint: Option, @@ -103,7 +81,7 @@ impl<'a> RegTraversalIter<'a> { } } -impl<'a> core::iter::Iterator for RegTraversalIter<'a> { +impl core::iter::Iterator for RegTraversalIter { type Item = PReg; fn next(&mut self) -> Option { @@ -118,16 +96,14 @@ impl<'a> core::iter::Iterator for RegTraversalIter<'a> { } } - while !self.preferred.done() { - let reg = self.preferred.advance(); + while let Some(reg) = self.preferred.next() { if Some(reg) == self.hint || reg.hw_enc() >= self.limit.unwrap_or(usize::MAX) { continue; // Try again; we already tried the hint or we are outside of the register range limit. } return Some(reg); } - while !self.non_preferred.done() { - let reg = self.non_preferred.advance(); + while let Some(reg) = self.non_preferred.next() { if Some(reg) == self.hint || reg.hw_enc() >= self.limit.unwrap_or(usize::MAX) { continue; // Try again; we already tried the hint or we are outside of the register range limit. } diff --git a/src/lib.rs b/src/lib.rs index 30a4bcc8..75d977fc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -293,6 +293,16 @@ impl PRegSet { pub fn is_empty(&self, regclass: RegClass) -> bool { self.bits[regclass as usize] == 0 } + + /// Returns the number of register in this set. + pub fn len(&self) -> u32 { + self.bits.iter().map(|s| s.count_ones()).sum() + } + + /// Returns the maximum register in this set, with the highest hw_enc value. + pub fn max_preg(&self) -> Option { + self.into_iter().last() + } } impl core::ops::BitAnd for PRegSet { @@ -315,6 +325,14 @@ impl core::ops::BitOr for PRegSet { } } +impl IntoIterator for &PRegSet { + type Item = PReg; + type IntoIter = PRegSetIter; + fn into_iter(self) -> PRegSetIter { + (*self).into_iter() + } +} + impl IntoIterator for PRegSet { type Item = PReg; type IntoIter = PRegSetIter; @@ -352,15 +370,11 @@ impl From<&MachineEnv> for PRegSet { let mut res = Self::default(); for class in env.preferred_regs_by_class.iter() { - for preg in class { - res.add(*preg) - } + res.union_from(*class) } for class in env.non_preferred_regs_by_class.iter() { - for preg in class { - res.add(*preg) - } + res.union_from(*class) } res @@ -1483,7 +1497,7 @@ pub struct MachineEnv { /// /// If an explicit scratch register is provided in `scratch_by_class` then /// it must not appear in this list. - pub preferred_regs_by_class: [Vec; 3], + pub preferred_regs_by_class: [PRegSet; 3], /// Non-preferred physical registers for each class. These are the /// registers that will be allocated if a preferred register is @@ -1492,7 +1506,7 @@ pub struct MachineEnv { /// /// If an explicit scratch register is provided in `scratch_by_class` then /// it must not appear in this list. - pub non_preferred_regs_by_class: [Vec; 3], + pub non_preferred_regs_by_class: [PRegSet; 3], /// Optional dedicated scratch register per class. This is needed to perform /// moves between registers when cyclic move patterns occur. The @@ -1774,3 +1788,37 @@ unsafe impl allocator_api2::alloc::Allocator for Bump { self.0.deref().shrink(ptr, old_layout, new_layout) } } + +#[cfg(test)] +mod tests { + use super::{PReg, PRegSet, RegClass::Int}; + + #[test] + fn preg_set_len() { + let mut set = PRegSet::empty(); + assert_eq!(set.len(), 0); + + set.add(PReg::new(3, Int)); + assert_eq!(set.len(), 1); + set.add(PReg::new(3, Int)); + assert_eq!(set.len(), 1); + + set.add(PReg::new(4, Int)); + assert_eq!(set.len(), 2); + } + + #[test] + fn preg_set_max_preg() { + let mut set = PRegSet::empty(); + assert_eq!(set.max_preg(), None); + + set.add(PReg::new(3, Int)); + assert_eq!(set.max_preg(), Some(PReg::new(3, Int))); + + set.add(PReg::new(4, Int)); + assert_eq!(set.max_preg(), Some(PReg::new(4, Int))); + + set.add(PReg::new(2, Int)); + assert_eq!(set.max_preg(), Some(PReg::new(4, Int))); + } +} From da5fb89930d1948d37f14fbc3ce992840b2750e9 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Tue, 20 Jan 2026 11:58:44 -0800 Subject: [PATCH 2/2] Tweak `Cursor::new` a bit Use `offset_hint` as a mask start rather than looking for the Nth bit set. Should in theory have the same overall desired balancing semantics but enables more bit-tricks. --- src/ion/reg_traversal.rs | 28 +++++++++++++--------------- src/lib.rs | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 15 deletions(-) diff --git a/src/ion/reg_traversal.rs b/src/ion/reg_traversal.rs index 4c968772..33193e8c 100644 --- a/src/ion/reg_traversal.rs +++ b/src/ion/reg_traversal.rs @@ -10,18 +10,15 @@ struct Cursor { impl Cursor { #[inline] - fn new(registers: &PRegSet, offset_hint: usize) -> Self { - let first = registers - .into_iter() - .skip(offset_hint) - .collect::() - .into_iter(); - let second = registers - .into_iter() - .take(offset_hint) - .collect::() - .into_iter(); - Self { first, second } + fn new(registers: &PRegSet, class: RegClass, offset_hint: usize) -> Self { + let mut mask = PRegSet::empty(); + mask.add_up_to(PReg::new(offset_hint % PReg::MAX, class)); + let first = *registers & mask.invert(); + let second = *registers & mask; + Self { + first: first.into_iter(), + second: second.into_iter(), + } } fn next(&mut self) -> Option { @@ -65,9 +62,10 @@ impl RegTraversalIter { debug_assert!(fixed != Some(PReg::invalid())); debug_assert!(hint != Some(PReg::invalid())); - let class = class as u8 as usize; - let preferred = Cursor::new(&env.preferred_regs_by_class[class], offset); - let non_preferred = Cursor::new(&env.non_preferred_regs_by_class[class], offset); + let class_index = class as u8 as usize; + let preferred = Cursor::new(&env.preferred_regs_by_class[class_index], class, offset); + let non_preferred = + Cursor::new(&env.non_preferred_regs_by_class[class_index], class, offset); Self { is_fixed: fixed.is_some(), diff --git a/src/lib.rs b/src/lib.rs index 75d977fc..8ee6139b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -303,6 +303,15 @@ impl PRegSet { pub fn max_preg(&self) -> Option { self.into_iter().last() } + + /// Add all registers from `0..reg` to this set, not including `reg` itself. + pub fn add_up_to(&mut self, reg: PReg) { + let (index, bit) = Self::split_index(reg); + for i in 0..index { + self.bits[i] = !0; + } + self.bits[index] = (1 << bit) - 1; + } } impl core::ops::BitAnd for PRegSet { @@ -1821,4 +1830,31 @@ mod tests { set.add(PReg::new(2, Int)); assert_eq!(set.max_preg(), Some(PReg::new(4, Int))); } + + #[test] + fn preg_set_new_up_to() { + let p0 = PReg::new(0, Int); + let p1 = PReg::new(1, Int); + let p2 = PReg::new(2, Int); + let p3 = PReg::new(3, Int); + { + let mut set = PRegSet::empty(); + set.add_up_to(p1); + assert!(set.contains(p0)); + assert!(!set.contains(p1)); + } + { + let mut set = PRegSet::empty(); + set.add_up_to(p0); + assert!(!set.contains(p0)); + } + { + let mut set = PRegSet::empty(); + set.add_up_to(p3); + assert!(set.contains(p0)); + assert!(set.contains(p1)); + assert!(set.contains(p2)); + assert!(!set.contains(p3)); + } + } }