Skip to content

Commit

Permalink
Merge pull request #1531 from ltratt/rev_analyse_change
Browse files Browse the repository at this point in the history 
Change and extend the reverse analyser
  • Loading branch information
@vext01
vext01 authored Dec 30, 2024
2 parents c4139c4 + af3277b commit a898029
Show file tree
Hide file tree
Showing 5 changed files with 297 additions and 266 deletions.
66 changes: 37 additions & 29 deletions ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,13 +27,14 @@
//! where it has spilled an instruction's value: it guarantees to spill an instruction to at most
//! one place on the stack.
use super::rev_analyse::RevAnalyse;
use crate::compile::jitc_yk::{
codegen::{
abs_stack::AbstractStack,
reg_alloc::{self, VarLocation},
x64::REG64_BYTESIZE,
},
jit_ir::{Const, ConstIdx, FloatTy, Inst, InstIdx, Module, Operand, Ty},
jit_ir::{Const, ConstIdx, FloatTy, Inst, InstIdx, Module, Operand, PtrAddInst, Ty},
};
use dynasmrt::{
dynasm,
Expand Down Expand Up @@ -106,6 +107,7 @@ static RESERVED_FP_REGS: [Rx; 0] = [];
/// A linear scan register allocator.
pub(crate) struct LSRegAlloc<'a> {
m: &'a Module,
rev_an: RevAnalyse<'a>,
/// Which general purpose registers are active?
gp_regset: RegSet<Rq>,
/// In what state are the general purpose registers?
Expand All @@ -114,28 +116,17 @@ pub(crate) struct LSRegAlloc<'a> {
fp_regset: RegSet<Rx>,
/// In what state are the floating point registers?
fp_reg_states: [RegState; FP_REGS_LEN],
/// Record the [InstIdx] of the last instruction that the value produced by an instruction is
/// used. By definition this must either be unused (if an instruction does not produce a value)
/// or `>=` the offset in this vector.
inst_vals_alive_until: Vec<InstIdx>,
/// Where on the stack is an instruction's value spilled? Set to `usize::MAX` if that offset is
/// currently unknown. Note: multiple instructions can alias to the same [SpillState].
spills: Vec<SpillState>,
/// The abstract stack: shared between general purpose and floating point registers.
stack: AbstractStack,
/// What [VarLocation] should an instruction aim to put its output to?
vloc_hints: Vec<Option<VarLocation>>,
}

impl<'a> LSRegAlloc<'a> {
/// Create a new register allocator, with the existing interpreter frame spanning
/// `interp_stack_len` bytes.
pub(crate) fn new(
m: &'a Module,
inst_vals_alive_until: Vec<InstIdx>,
vloc_hints: Vec<Option<VarLocation>>,
interp_stack_len: usize,
) -> Self {
pub(crate) fn new(m: &'a Module, interp_stack_len: usize) -> Self {
#[cfg(debug_assertions)]
{
// We rely on the registers in GP_REGS being numbered 0..15 (inc.) for correctness.
Expand All @@ -162,22 +153,23 @@ impl<'a> LSRegAlloc<'a> {
let mut stack = AbstractStack::default();
stack.grow(interp_stack_len);

let mut rev_an = RevAnalyse::new(m);
rev_an.analyse_header();
LSRegAlloc {
m,
rev_an,
gp_regset: RegSet::with_gp_reserved(),
gp_reg_states,
fp_regset: RegSet::with_fp_reserved(),
fp_reg_states,
inst_vals_alive_until,
vloc_hints,
spills: vec![SpillState::Empty; m.insts_len()],
stack,
}
}

/// Reset the register allocator. We use this when moving from the trace header into the trace
/// body.
pub(crate) fn reset(&mut self) {
pub(crate) fn reset(&mut self, header_end_vlocs: &[VarLocation]) {
for rs in self.gp_reg_states.iter_mut() {
*rs = RegState::Empty;
}
Expand All @@ -193,6 +185,8 @@ impl<'a> LSRegAlloc<'a> {
self.fp_reg_states[usize::from(reg.code())] = RegState::Reserved;
}
self.fp_regset = RegSet::with_fp_reserved();

self.rev_an.analyse_body(header_end_vlocs);
}

/// Before generating code for the instruction at `iidx`, see which registers are no longer
Expand Down Expand Up @@ -262,24 +256,37 @@ impl<'a> LSRegAlloc<'a> {
self.stack.size()
}

/// Is the instruction at [iidx] a tombstone or otherwise known to be dead (i.e. equivalent to
/// a tombstone)?
pub(crate) fn is_inst_tombstone(&self, iidx: InstIdx) -> bool {
self.rev_an.is_inst_tombstone(iidx)
}

/// Is the value produced by instruction `query_iidx` used after (but not including!)
/// instruction `cur_idx`?
pub(crate) fn is_inst_var_still_used_after(
&self,
cur_iidx: InstIdx,
query_iidx: InstIdx,
) -> bool {
usize::from(cur_iidx) < usize::from(self.inst_vals_alive_until[usize::from(query_iidx)])
self.rev_an
.is_inst_var_still_used_after(cur_iidx, query_iidx)
}

/// Is the value produced by instruction `query_iidx` used at or after instruction `cur_idx`?
fn is_inst_var_still_used_at(&self, cur_iidx: InstIdx, query_iidx: InstIdx) -> bool {
usize::from(cur_iidx) <= usize::from(self.inst_vals_alive_until[usize::from(query_iidx)])
usize::from(cur_iidx)
<= usize::from(self.rev_an.inst_vals_alive_until[usize::from(query_iidx)])
}

#[cfg(test)]
pub(crate) fn inst_vals_alive_until(&self) -> &Vec<InstIdx> {
&self.inst_vals_alive_until
&self.rev_an.inst_vals_alive_until
}

/// Return the inline [PtrAddInst] for a load/store, if there is one.
pub(crate) fn ptradd(&self, iidx: InstIdx) -> Option<PtrAddInst> {
self.rev_an.ptradds[usize::from(iidx)]
}
}

Expand Down Expand Up @@ -438,8 +445,7 @@ impl LSRegAlloc<'_> {
// Deal with `OutputCanBeSameAsInput`.
for i in 0..constraints.len() {
if let RegConstraint::OutputCanBeSameAsInput(search_op) = constraints[i].clone() {
if let Some(VarLocation::Register(reg_alloc::Register::GP(reg))) =
self.vloc_hints[usize::from(iidx)]
if let Some(reg_alloc::Register::GP(reg)) = self.rev_an.reg_hints[usize::from(iidx)]
{
if avoid.is_set(reg) {
continue;
Expand Down Expand Up @@ -469,8 +475,8 @@ impl LSRegAlloc<'_> {
RegConstraint::Output
| RegConstraint::OutputCanBeSameAsInput(_)
| RegConstraint::InputOutput(_) => {
if let Some(VarLocation::Register(reg_alloc::Register::GP(reg))) =
self.vloc_hints[usize::from(iidx)]
if let Some(reg_alloc::Register::GP(reg)) =
self.rev_an.reg_hints[usize::from(iidx)]
{
if !avoid.is_set(reg) {
*cnstr = match cnstr {
Expand Down Expand Up @@ -564,8 +570,9 @@ impl LSRegAlloc<'_> {
if furthest.is_none() {
furthest = Some((reg, from_iidx));
} else if let Some((_, furthest_iidx)) = furthest {
if self.inst_vals_alive_until[usize::from(from_iidx)]
>= self.inst_vals_alive_until[usize::from(furthest_iidx)]
if self.rev_an.inst_vals_alive_until[usize::from(from_iidx)]
>= self.rev_an.inst_vals_alive_until
[usize::from(furthest_iidx)]
{
furthest = Some((reg, from_iidx))
}
Expand Down Expand Up @@ -755,8 +762,8 @@ impl LSRegAlloc<'_> {
if self.is_inst_var_still_used_after(cur_iidx, query_iidx) {
let mut new_reg = None;
// Try to use `query_iidx`s hint, if there is one, and it's not in use...
if let Some(VarLocation::Register(reg_alloc::Register::GP(reg))) =
self.vloc_hints[usize::from(query_iidx)]
if let Some(reg_alloc::Register::GP(reg)) =
self.rev_an.reg_hints[usize::from(query_iidx)]
{
if !self.gp_regset.is_set(reg) && !avoid.is_set(reg) {
new_reg = Some(reg);
Expand Down Expand Up @@ -1094,8 +1101,9 @@ impl LSRegAlloc<'_> {
if furthest.is_none() {
furthest = Some((reg, from_iidx));
} else if let Some((_, furthest_iidx)) = furthest {
if self.inst_vals_alive_until[usize::from(from_iidx)]
>= self.inst_vals_alive_until[usize::from(furthest_iidx)]
if self.rev_an.inst_vals_alive_until[usize::from(from_iidx)]
>= self.rev_an.inst_vals_alive_until
[usize::from(furthest_iidx)]
{
furthest = Some((reg, from_iidx))
}
Expand Down
34 changes: 12 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 @@ -23,7 +23,7 @@
use super::{
super::{
int_signs::{SignExtend, Truncate},
jit_ir::{self, BinOp, FloatTy, Inst, InstIdx, Module, Operand, PtrAddInst, TraceKind, Ty},
jit_ir::{self, BinOp, FloatTy, Inst, InstIdx, Module, Operand, TraceKind, Ty},
CompilationError,
},
reg_alloc::{self, VarLocation},
Expand Down Expand Up @@ -60,7 +60,6 @@ use std::{
slice,
sync::{Arc, Weak},
};
use vob::Vob;
use ykaddr::addr::symbol_to_ptr;

mod deopt;
Expand Down Expand Up @@ -195,12 +194,6 @@ struct Assemble<'a> {
/// The stack pointer offset of the root trace's frame from the base pointer of the interpreter
/// frame. If this is the root trace, this will be None.
root_offset: Option<usize>,
/// Does this Load/Store instruction reference a [PtrAddInst]?
ptradds: Vec<Option<PtrAddInst>>,
/// For each instruction, does this code generator use its value? This is implicitly a second
/// layer of dead-code elimination: it doesn't cause JIT IR instructions to be removed, but
/// it will stop any code being (directly) generated for some of them.
used_insts: Vob,
/// The offset after the trace's prologue. This is the re-entry point when returning from
/// side-traces.
prologue_offset: AssemblyOffset,
Expand Down Expand Up @@ -251,20 +244,16 @@ impl<'a> Assemble<'a> {
}
};

let (inst_vals_alive_until, used_insts, ptradds, vloc_hints) = rev_analyse::rev_analyse(m)?;

Ok(Box::new(Self {
m,
ra: LSRegAlloc::new(m, inst_vals_alive_until, vloc_hints, sp_offset),
ra: LSRegAlloc::new(m, sp_offset),
asm,
header_start_locs: Vec::new(),
body_start_locs: Vec::new(),
deoptinfo: HashMap::new(),
comments: Cell::new(IndexMap::new()),
sp_offset,
root_offset,
used_insts,
ptradds,
prologue_offset: AssemblyOffset(0),
}))
}
Expand Down Expand Up @@ -468,7 +457,7 @@ impl<'a> Assemble<'a> {
let mut in_header = true;
while let Some((iidx, inst)) = next {
self.comment(self.asm.offset(), inst.display(self.m, iidx).to_string());
if !self.used_insts[usize::from(iidx)] {
if self.ra.is_inst_tombstone(iidx) {
next = iter.next();
continue;
}
Expand Down Expand Up @@ -1104,7 +1093,7 @@ impl<'a> Assemble<'a> {
/// Generate code for a [LoadInst], loading from a `register + off`. `off` should only be
/// non-zero if the [LoadInst] references a [PtrAddInst].
fn cg_load(&mut self, iidx: jit_ir::InstIdx, inst: &jit_ir::LoadInst) {
let (ptr_op, off) = match self.ptradds[usize::from(iidx)] {
let (ptr_op, off) = match self.ra.ptradd(iidx) {
Some(x) => (x.ptr(self.m), x.off()),
None => (inst.operand(self.m), 0),
};
Expand Down Expand Up @@ -1216,7 +1205,7 @@ impl<'a> Assemble<'a> {
/// Generate code for a [StoreInst], storing it at a `register + off`. `off` should only be
/// non-zero if the [StoreInst] references a [PtrAddInst].
fn cg_store(&mut self, iidx: InstIdx, inst: &jit_ir::StoreInst) {
let (tgt_op, off) = match self.ptradds[usize::from(iidx)] {
let (tgt_op, off) = match self.ra.ptradd(iidx) {
Some(x) => (x.ptr(self.m), x.off()),
None => (inst.tgt(self.m), 0),
};
Expand Down Expand Up @@ -1972,14 +1961,15 @@ impl<'a> Assemble<'a> {
// that have become constants during the trace header. So we will always have to either
// update the [ParamInst]s of the trace body, which isn't ideal since it requires the
// [Module] the be mutable. Or we do what we do below just for constants.
let mut varlocs = Vec::new();
for var in self.m.trace_header_end().iter() {
let varloc = self.op_to_var_location(var.unpack(self.m));
varlocs.push(varloc);
}
let varlocs = self
.m
.trace_header_end()
.iter()
.map(|pop| self.op_to_var_location(pop.unpack(self.m)))
.collect::<Vec<_>>();
// Reset the register allocator before priming it with information about the trace body
// inputs.
self.ra.reset();
self.ra.reset(varlocs.as_slice());
for (i, op) in self.m.trace_body_start().iter().enumerate() {
// By definition these can only be variables.
let iidx = match op.unpack(self.m) {
Expand Down
Loading

0 comments on commit a898029

Please sign in to comment.