Skip to content

Commit fbe1825

Browse files
authored
Store PRegSet in MachineEnv (#254)
Helps to make the `MachineEnv`s in Wasmtime `const`-allocatable. Closes #252
1 parent 33c6242 commit fbe1825

File tree

8 files changed

+148
-93
lines changed

8 files changed

+148
-93
lines changed

src/fastalloc/lru.rs

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@ pub struct LruNode {
2828
}
2929

3030
impl Lru {
31-
pub fn new(regclass: RegClass, regs: &[PReg]) -> Self {
31+
pub fn new(regclass: RegClass, regs: &PRegSet) -> Self {
32+
let regs = regs.into_iter().collect::<Vec<_>>();
3233
let mut data = vec![
3334
LruNode {
3435
prev: u8::MAX,
@@ -248,7 +249,7 @@ impl fmt::Debug for Lru {
248249
while node != self.head {
249250
if seen.contains(&node) {
250251
panic!(
251-
"The {:?} LRU is messed up:
252+
"The {:?} LRU is messed up:
252253
head: {:?}, {:?} -> p{node}, actual data: {:?}",
253254
self.regclass, self.head, data_str, self.data
254255
);
@@ -298,7 +299,7 @@ impl<T: PartialEq> PartialEq for PartedByRegClass<T> {
298299
pub type Lrus = PartedByRegClass<Lru>;
299300

300301
impl Lrus {
301-
pub fn new(int_regs: &[PReg], float_regs: &[PReg], vec_regs: &[PReg]) -> Self {
302+
pub fn new(int_regs: &PRegSet, float_regs: &PRegSet, vec_regs: &PRegSet) -> Self {
302303
Self {
303304
items: [
304305
Lru::new(RegClass::Int, int_regs),

src/fastalloc/mod.rs

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -444,21 +444,9 @@ impl<'a, F: Function> Env<'a, F> {
444444
env.preferred_regs_by_class[RegClass::Float as usize].clone(),
445445
env.preferred_regs_by_class[RegClass::Vector as usize].clone(),
446446
];
447-
regs[0].extend(
448-
env.non_preferred_regs_by_class[RegClass::Int as usize]
449-
.iter()
450-
.cloned(),
451-
);
452-
regs[1].extend(
453-
env.non_preferred_regs_by_class[RegClass::Float as usize]
454-
.iter()
455-
.cloned(),
456-
);
457-
regs[2].extend(
458-
env.non_preferred_regs_by_class[RegClass::Vector as usize]
459-
.iter()
460-
.cloned(),
461-
);
447+
regs[0].union_from(env.non_preferred_regs_by_class[RegClass::Int as usize]);
448+
regs[1].union_from(env.non_preferred_regs_by_class[RegClass::Float as usize]);
449+
regs[2].union_from(env.non_preferred_regs_by_class[RegClass::Vector as usize]);
462450
let allocatable_regs = PRegSet::from(env);
463451
let num_available_pregs: PartedByRegClass<i16> = PartedByRegClass {
464452
items: [
@@ -508,9 +496,9 @@ impl<'a, F: Function> Env<'a, F> {
508496
],
509497
preferred_victim: PartedByRegClass {
510498
items: [
511-
regs[0].last().cloned().unwrap_or(PReg::invalid()),
512-
regs[1].last().cloned().unwrap_or(PReg::invalid()),
513-
regs[2].last().cloned().unwrap_or(PReg::invalid()),
499+
regs[0].max_preg().unwrap_or(PReg::invalid()),
500+
regs[1].max_preg().unwrap_or(PReg::invalid()),
501+
regs[2].max_preg().unwrap_or(PReg::invalid()),
514502
],
515503
},
516504
reused_input_to_reuse_op: vec![usize::MAX; max_operand_len as usize],

src/fastalloc/tests.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -146,10 +146,10 @@ fn mach_env(no_of_regs: usize) -> MachineEnv {
146146
(0..no_of_regs)
147147
.map(|no| PReg::new(no, RegClass::Int))
148148
.collect(),
149-
vec![],
150-
vec![],
149+
PRegSet::empty(),
150+
PRegSet::empty(),
151151
],
152-
non_preferred_regs_by_class: [vec![], vec![], vec![]],
152+
non_preferred_regs_by_class: [PRegSet::empty(); 3],
153153
scratch_by_class: [None, None, None],
154154
fixed_stack_slots: vec![],
155155
}

src/fuzzing/func.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -696,20 +696,20 @@ impl core::fmt::Debug for Func {
696696
}
697697

698698
pub fn machine_env() -> MachineEnv {
699-
fn regs(r: core::ops::Range<usize>, c: RegClass) -> Vec<PReg> {
699+
fn regs(r: core::ops::Range<usize>, c: RegClass) -> PRegSet {
700700
r.map(|i| PReg::new(i, c)).collect()
701701
}
702-
let preferred_regs_by_class: [Vec<PReg>; 3] = [
702+
let preferred_regs_by_class = [
703703
regs(0..24, RegClass::Int),
704704
regs(0..24, RegClass::Float),
705705
regs(0..24, RegClass::Vector),
706706
];
707-
let non_preferred_regs_by_class: [Vec<PReg>; 3] = [
707+
let non_preferred_regs_by_class = [
708708
regs(24..32, RegClass::Int),
709709
regs(24..32, RegClass::Float),
710710
regs(24..32, RegClass::Vector),
711711
];
712-
let scratch_by_class: [Option<PReg>; 3] = [None, None, None];
712+
let scratch_by_class = [None, None, None];
713713
let fixed_stack_slots = (32..63)
714714
.flat_map(|i| {
715715
[

src/ion/liveranges.rs

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -115,9 +115,8 @@ impl<'a, F: Function> Env<'a, F> {
115115
}
116116
for class in 0..self.preferred_victim_by_class.len() {
117117
self.preferred_victim_by_class[class] = self.env.non_preferred_regs_by_class[class]
118-
.last()
119-
.or(self.env.preferred_regs_by_class[class].last())
120-
.cloned()
118+
.max_preg()
119+
.or(self.env.preferred_regs_by_class[class].max_preg())
121120
.unwrap_or(PReg::invalid());
122121
}
123122
// Create VRegs from the vreg count.

src/ion/process.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1219,8 +1219,8 @@ impl<'a, F: Function> Env<'a, F> {
12191219
let mut fixed_assigned = 0;
12201220
let mut total_regs = 0;
12211221
for preg in self.env.preferred_regs_by_class[class as u8 as usize]
1222-
.iter()
1223-
.chain(self.env.non_preferred_regs_by_class[class as u8 as usize].iter())
1222+
.into_iter()
1223+
.chain(self.env.non_preferred_regs_by_class[class as u8 as usize])
12241224
{
12251225
trace!(" -> PR {:?}", preg);
12261226
let start = LiveRangeKey::from_range(&CodeRange {

src/ion/reg_traversal.rs

Lines changed: 26 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -1,53 +1,28 @@
11
//! Iterate over available registers.
22
3-
use crate::{MachineEnv, PReg, RegClass};
3+
use crate::{MachineEnv, PReg, PRegSet, PRegSetIter, RegClass};
44

55
/// Keep track of where we are in the register traversal.
6-
struct Cursor<'a> {
7-
registers: &'a [PReg],
8-
index: usize,
9-
offset: usize,
6+
struct Cursor {
7+
first: PRegSetIter,
8+
second: PRegSetIter,
109
}
1110

12-
impl<'a> Cursor<'a> {
11+
impl Cursor {
1312
#[inline]
14-
fn new(registers: &'a [PReg], offset_hint: usize) -> Self {
15-
let offset = if registers.len() > 0 {
16-
offset_hint % registers.len()
17-
} else {
18-
0
19-
};
13+
fn new(registers: &PRegSet, class: RegClass, offset_hint: usize) -> Self {
14+
let mut mask = PRegSet::empty();
15+
mask.add_up_to(PReg::new(offset_hint % PReg::MAX, class));
16+
let first = *registers & mask.invert();
17+
let second = *registers & mask;
2018
Self {
21-
registers,
22-
index: 0,
23-
offset,
19+
first: first.into_iter(),
20+
second: second.into_iter(),
2421
}
2522
}
2623

27-
/// Wrap around the end of the register list; [`Cursor::done`] guarantees we
28-
/// do not see the same register twice.
29-
#[inline]
30-
fn wrap(index: usize, end: usize) -> usize {
31-
if index >= end {
32-
index - end
33-
} else {
34-
index
35-
}
36-
}
37-
38-
/// Advance to the next register and return it.
39-
#[inline]
40-
fn advance(&mut self) -> PReg {
41-
let loc = Self::wrap(self.index + self.offset, self.registers.len());
42-
let reg = self.registers[loc];
43-
self.index += 1;
44-
reg
45-
}
46-
47-
/// Return `true` if we have seen all registers.
48-
#[inline]
49-
fn done(&self) -> bool {
50-
self.index >= self.registers.len()
24+
fn next(&mut self) -> Option<PReg> {
25+
self.first.next().or_else(|| self.second.next())
5126
}
5227
}
5328

@@ -65,19 +40,19 @@ impl<'a> Cursor<'a> {
6540
/// registers; then, non-preferred registers. (In normal usage, these consist
6641
/// of caller-save and callee-save registers respectively, to minimize
6742
/// clobber-saves; but they need not.)
68-
pub struct RegTraversalIter<'a> {
43+
pub struct RegTraversalIter {
6944
is_fixed: bool,
7045
fixed: Option<PReg>,
7146
use_hint: bool,
7247
hint: Option<PReg>,
73-
preferred: Cursor<'a>,
74-
non_preferred: Cursor<'a>,
48+
preferred: Cursor,
49+
non_preferred: Cursor,
7550
limit: Option<usize>,
7651
}
7752

78-
impl<'a> RegTraversalIter<'a> {
53+
impl RegTraversalIter {
7954
pub fn new(
80-
env: &'a MachineEnv,
55+
env: &MachineEnv,
8156
class: RegClass,
8257
fixed: Option<PReg>,
8358
hint: Option<PReg>,
@@ -87,9 +62,10 @@ impl<'a> RegTraversalIter<'a> {
8762
debug_assert!(fixed != Some(PReg::invalid()));
8863
debug_assert!(hint != Some(PReg::invalid()));
8964

90-
let class = class as u8 as usize;
91-
let preferred = Cursor::new(&env.preferred_regs_by_class[class], offset);
92-
let non_preferred = Cursor::new(&env.non_preferred_regs_by_class[class], offset);
65+
let class_index = class as u8 as usize;
66+
let preferred = Cursor::new(&env.preferred_regs_by_class[class_index], class, offset);
67+
let non_preferred =
68+
Cursor::new(&env.non_preferred_regs_by_class[class_index], class, offset);
9369

9470
Self {
9571
is_fixed: fixed.is_some(),
@@ -103,7 +79,7 @@ impl<'a> RegTraversalIter<'a> {
10379
}
10480
}
10581

106-
impl<'a> core::iter::Iterator for RegTraversalIter<'a> {
82+
impl core::iter::Iterator for RegTraversalIter {
10783
type Item = PReg;
10884

10985
fn next(&mut self) -> Option<PReg> {
@@ -118,16 +94,14 @@ impl<'a> core::iter::Iterator for RegTraversalIter<'a> {
11894
}
11995
}
12096

121-
while !self.preferred.done() {
122-
let reg = self.preferred.advance();
97+
while let Some(reg) = self.preferred.next() {
12398
if Some(reg) == self.hint || reg.hw_enc() >= self.limit.unwrap_or(usize::MAX) {
12499
continue; // Try again; we already tried the hint or we are outside of the register range limit.
125100
}
126101
return Some(reg);
127102
}
128103

129-
while !self.non_preferred.done() {
130-
let reg = self.non_preferred.advance();
104+
while let Some(reg) = self.non_preferred.next() {
131105
if Some(reg) == self.hint || reg.hw_enc() >= self.limit.unwrap_or(usize::MAX) {
132106
continue; // Try again; we already tried the hint or we are outside of the register range limit.
133107
}

0 commit comments

Comments
 (0)