Skip to content

Commit

Permalink
Assign registers in one go on loop jump.
Browse files Browse the repository at this point in the history 
This requires us to introduce a `None` constraint: we can then build up
all GP and FP constraints in one go, gradually turning `None`s into
other constraints as necessary.
  • Loading branch information
@ltratt
ltratt committed Dec 17, 2024
1 parent 1bc4fa5 commit 2ae5594
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 27 deletions.
28 changes: 23 additions & 5 deletions ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -397,7 +397,8 @@ impl LSRegAlloc<'_> {
| RegConstraint::OutputCanBeSameAsInput(_)
| RegConstraint::OutputFromReg(_)
| RegConstraint::InputOutput(_) => true,
RegConstraint::Clobber(_) | RegConstraint::Temporary => false,
RegConstraint::Clobber(_) | RegConstraint::Temporary | RegConstraint::None =>
false,
})
.count()
<= 1
Expand Down Expand Up @@ -427,7 +428,8 @@ impl LSRegAlloc<'_> {
| RegConstraint::Output
| RegConstraint::OutputCanBeSameAsInput(_)
| RegConstraint::Input(_)
| RegConstraint::Temporary => {}
| RegConstraint::Temporary
| RegConstraint::None => {}
}
}

Expand Down Expand Up @@ -524,6 +526,9 @@ impl LSRegAlloc<'_> {
| RegConstraint::OutputCanBeSameAsInput(_)
| RegConstraint::OutputFromReg(_)
| RegConstraint::Temporary => (),
RegConstraint::None => {
asgn[i] = Some(GP_REGS[0]);
}
}
}

Expand Down Expand Up @@ -613,7 +618,8 @@ impl LSRegAlloc<'_> {
| RegConstraint::OutputCanBeSameAsInput(_)
| RegConstraint::OutputFromReg(_)
| RegConstraint::Clobber(_)
| RegConstraint::Temporary => (),
| RegConstraint::Temporary
| RegConstraint::None => (),
}
}

Expand Down Expand Up @@ -659,6 +665,7 @@ impl LSRegAlloc<'_> {
self.gp_regset.unset(reg);
self.gp_reg_states[usize::from(reg.code())] = RegState::Empty;
}
RegConstraint::None => (),
}
}
asgn.map(|x| x.unwrap())
Expand Down Expand Up @@ -967,7 +974,8 @@ impl LSRegAlloc<'_> {
| RegConstraint::OutputCanBeSameAsInput(_)
| RegConstraint::OutputFromReg(_)
| RegConstraint::InputOutput(_) => true,
RegConstraint::Clobber(_) | RegConstraint::Temporary => false,
RegConstraint::Clobber(_) | RegConstraint::Temporary | RegConstraint::None =>
false,
})
.count()
<= 1
Expand Down Expand Up @@ -998,6 +1006,9 @@ impl LSRegAlloc<'_> {
| RegConstraint::OutputCanBeSameAsInput(_)
| RegConstraint::Output
| RegConstraint::Temporary => {}
RegConstraint::None => {
asgn[i] = Some(FP_REGS[0]);
}
}
}

Expand Down Expand Up @@ -1044,6 +1055,7 @@ impl LSRegAlloc<'_> {
| RegConstraint::OutputFromReg(_)
| RegConstraint::Temporary => (),
RegConstraint::OutputCanBeSameAsInput(_) => todo!(),
RegConstraint::None => (),
}
}

Expand Down Expand Up @@ -1134,6 +1146,7 @@ impl LSRegAlloc<'_> {
| RegConstraint::Clobber(_)
| RegConstraint::Temporary => (),
RegConstraint::OutputCanBeSameAsInput(_) => todo!(),
RegConstraint::None => (),
}
}

Expand Down Expand Up @@ -1178,6 +1191,7 @@ impl LSRegAlloc<'_> {
self.fp_reg_states[usize::from(reg.code())] = RegState::Empty;
}
RegConstraint::OutputCanBeSameAsInput(_) => todo!(),
RegConstraint::None => (),
}
}
asgn.map(|x| x.unwrap())
Expand Down Expand Up @@ -1421,6 +1435,9 @@ pub(crate) enum RegConstraint<R: Register> {
Clobber(R),
/// A temporary register *x* that the instruction will clobber.
Temporary,
/// A no-op register constraint. A random register will be assigned to this: using this
/// register for any purposes leads to undefined behaviour.
None,
}

#[cfg(debug_assertions)]
Expand All @@ -1437,7 +1454,8 @@ impl<R: dynasmrt::Register> RegConstraint<R> {
| Self::OutputCanBeSameAsInput(_)
| Self::OutputFromReg(_)
| Self::Clobber(_)
| Self::Temporary => None,
| Self::Temporary
| Self::None => None,
}
}
}
Expand Down
44 changes: 22 additions & 22 deletions ykrt/src/compile/jitc_yk/codegen/x64/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1792,6 +1792,14 @@ impl<'a> Assemble<'a> {
// If we pass in `None` use `self.loop_start_locs` instead. We need to do this since we
// can't pass in `&self.loop_start_locs` directly due to borrowing restrictions.
let tgt_vars = tgt_vars.unwrap_or(self.loop_start_locs.as_slice());
let mut gp_regs = lsregalloc::GP_REGS
.iter()
.map(|_| RegConstraint::None)
.collect::<Vec<_>>();
let mut fp_regs = lsregalloc::FP_REGS
.iter()
.map(|_| RegConstraint::None)
.collect::<Vec<_>>();
for (i, op) in self.m.trace_body_end().iter().enumerate() {
// FIXME: This is completely broken: see the FIXME later.
let op = op.unpack(self.m);
Expand Down Expand Up @@ -1846,32 +1854,24 @@ impl<'a> Assemble<'a> {
// somewhere else (register/normal stack) so dst and src no longer
// match. But since the value can't change we can safely ignore this.
}
VarLocation::Register(reg) => {
// FIXME: This is completely broken, only works by accident and, probably,
// doesn't even always work. Continually running the register allocator in this
// way means we can end up clobbering values and, because this is the last
// instruction in the trace, none of the values will be used after this, so
// they don't have to be spilled.
match reg {
reg_alloc::Register::GP(r) => {
let [_] = self.ra.assign_gp_regs(
&mut self.asm,
iidx,
[RegConstraint::InputIntoReg(op.clone(), r)],
);
}
reg_alloc::Register::FP(r) => {
let [_] = self.ra.assign_fp_regs(
&mut self.asm,
iidx,
[RegConstraint::InputIntoReg(op.clone(), r)],
);
}
VarLocation::Register(reg) => match reg {
reg_alloc::Register::GP(r) => {
gp_regs[usize::from(r.code())] = RegConstraint::InputIntoReg(op.clone(), r);
}
}
reg_alloc::Register::FP(r) => {
fp_regs[usize::from(r.code())] = RegConstraint::InputIntoReg(op.clone(), r);
}
},
_ => todo!(),
}
}

let _: [_; lsregalloc::GP_REGS.len()] =
self.ra
.assign_gp_regs(&mut self.asm, iidx, gp_regs.try_into().unwrap());
let _: [_; lsregalloc::FP_REGS.len()] =
self.ra
.assign_fp_regs(&mut self.asm, iidx, fp_regs.try_into().unwrap());
}

fn cg_traceloopjump(&mut self, iidx: InstIdx) {
Expand Down

0 comments on commit 2ae5594

Please sign in to comment.