Skip to content

Commit

Permalink
YJIT: Remove Insn::RegTemps
Browse files Browse the repository at this point in the history
  • Loading branch information
@k0kubun
k0kubun committed Apr 19, 2023
1 parent 7477284 commit d91062c
Show file tree
Hide file tree
Showing 5 changed files with 81 additions and 97 deletions.
5 changes: 2 additions & 3 deletions yjit/src/backend/arm64/mod.rs
View file
Expand Up @@ -1113,14 +1113,13 @@ impl Assembler
Insn::CSelGE { truthy, falsy, out } => {
csel(cb, out.into(), truthy.into(), falsy.into(), Condition::GE);
}
Insn::LiveReg { .. } |
Insn::RegTemps(_) |
Insn::SpillTemp(_) => (), // just a reg alloc signal, no code
Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
Insn::PadInvalPatch => {
while (cb.get_write_pos().saturating_sub(std::cmp::max(start_write_pos, cb.page_start_pos()))) < cb.jmp_ptr_bytes() && !cb.has_dropped_bytes() {
nop(cb);
}
}
Insn::SpillTemp(_) => unreachable!("Insn::SpillTemp should have been lowered by lower_stack"),
};

// On failure, jump to the next page and retry the current insn
Expand Down
140 changes: 62 additions & 78 deletions yjit/src/backend/ir.rs
View file
Expand Up @@ -75,8 +75,19 @@ pub enum Opnd
// Output of a preceding instruction in this block
InsnOut{ idx: usize, num_bits: u8 },

// Pointer to a slot on the VM stack
Stack { idx: i32, stack_size: u8, sp_offset: i8, num_bits: u8 },
/// Pointer to a slot on the VM stack
Stack {
/// Index from stack top. Used for conversion to StackOpnd.
idx: i32,
/// Number of bits for Opnd::Reg and Opnd::Mem.
num_bits: u8,
/// ctx.stack_size when this operand is made. Used with idx for Opnd::Reg.
stack_size: u8,
/// ctx.sp_offset when this operand is made. Used with idx for Opnd::Mem.
sp_offset: i8,
/// ctx.reg_temps when this operand is pushed. Used for register allocation.
reg_temps: Option<RegTemps>
},

// Low-level operands, for lowering
Imm(i64), // Raw signed immediate
Expand Down Expand Up @@ -165,7 +176,7 @@ impl Opnd
Opnd::Reg(reg) => Some(Opnd::Reg(reg.with_num_bits(num_bits))),
Opnd::Mem(Mem { base, disp, .. }) => Some(Opnd::Mem(Mem { base, disp, num_bits })),
Opnd::InsnOut { idx, .. } => Some(Opnd::InsnOut { idx, num_bits }),
Opnd::Stack { idx, stack_size, sp_offset, .. } => Some(Opnd::Stack { idx, stack_size, sp_offset, num_bits }),
Opnd::Stack { idx, stack_size, sp_offset, reg_temps, .. } => Some(Opnd::Stack { idx, num_bits, stack_size, sp_offset, reg_temps }),
_ => None,
}
}
Expand Down Expand Up @@ -440,9 +451,6 @@ pub enum Insn {
/// Take a specific register. Signal the register allocator to not use it.
LiveReg { opnd: Opnd, out: Opnd },

/// Update live stack temps without spill
RegTemps(RegTemps),

// A low-level instruction that loads a value into a register.
Load { opnd: Opnd, out: Opnd },

Expand Down Expand Up @@ -571,7 +579,6 @@ impl Insn {
Insn::LeaLabel { .. } => "LeaLabel",
Insn::Lea { .. } => "Lea",
Insn::LiveReg { .. } => "LiveReg",
Insn::RegTemps(_) => "RegTemps",
Insn::Load { .. } => "Load",
Insn::LoadInto { .. } => "LoadInto",
Insn::LoadSExt { .. } => "LoadSExt",
Expand Down Expand Up @@ -717,7 +724,6 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
Insn::Jz(_) |
Insn::Label(_) |
Insn::LeaLabel { .. } |
Insn::RegTemps(_) |
Insn::PadInvalPatch |
Insn::PosMarker(_) => None,
Insn::CPopInto(opnd) |
Expand Down Expand Up @@ -816,7 +822,6 @@ impl<'a> InsnOpndMutIterator<'a> {
Insn::Jz(_) |
Insn::Label(_) |
Insn::LeaLabel { .. } |
Insn::RegTemps(_) |
Insn::PadInvalPatch |
Insn::PosMarker(_) => None,
Insn::CPopInto(opnd) |
Expand Down Expand Up @@ -929,10 +934,6 @@ pub struct Assembler {
/// Index of the last insn using the output of this insn
pub(super) live_ranges: Vec<usize>,

/// Parallel vec with insns
/// Bitmap of which temps are in a register for this insn
pub(super) reg_temps: Vec<RegTemps>,

/// Names of labels
pub(super) label_names: Vec<String>,

Expand All @@ -959,7 +960,6 @@ impl Assembler
Self {
insns: Vec::default(),
live_ranges: Vec::default(),
reg_temps: Vec::default(),
label_names,
ctx: Context::default(),
side_exits,
Expand Down Expand Up @@ -994,11 +994,12 @@ impl Assembler
// Index of this instruction
let insn_idx = self.insns.len();

// If we find any InsnOut from previous instructions, we're going to
// update the live range of the previous instruction to point to this
// one.
for opnd in insn.opnd_iter() {
let mut insn = insn;
let mut opnd_iter = insn.opnd_iter_mut();
while let Some(opnd) = opnd_iter.next() {
match opnd {
// If we find any InsnOut from previous instructions, we're going to update
// the live range of the previous instruction to point to this one.
Opnd::InsnOut { idx, .. } => {
assert!(*idx < self.insns.len());
self.live_ranges[*idx] = insn_idx;
Expand All @@ -1007,29 +1008,20 @@ impl Assembler
assert!(*idx < self.insns.len());
self.live_ranges[*idx] = insn_idx;
}
// Set current ctx.reg_temps to Opnd::Stack.
Opnd::Stack { idx, num_bits, stack_size, sp_offset, reg_temps: None } => {
*opnd = Opnd::Stack {
idx: *idx,
num_bits: *num_bits,
stack_size: *stack_size,
sp_offset: *sp_offset,
reg_temps: Some(self.ctx.get_reg_temps()),
};
}
_ => {}
}
}

// Update live stack temps for this instruction
let mut reg_temps = self.get_reg_temps();
match insn {
Insn::RegTemps(next_temps) => {
reg_temps = next_temps;
}
Insn::SpillTemp(opnd) => {
assert_eq!(reg_temps.get(opnd.stack_idx()), true);
reg_temps.set(opnd.stack_idx(), false);
}
_ => {}
}
// Assert no conflict
for stack_idx in 0..MAX_REG_TEMPS {
if reg_temps.get(stack_idx) {
assert!(!reg_temps.conflicts_with(stack_idx));
}
}

// Set a side exit context to Target::SideExit
let mut insn = insn;
if let Some(Target::SideExit { context, .. }) = insn.target_mut() {
Expand All @@ -1044,12 +1036,6 @@ impl Assembler

self.insns.push(insn);
self.live_ranges.push(insn_idx);
self.reg_temps.push(reg_temps);
}

/// Get stack temps that are currently in a register
pub fn get_reg_temps(&self) -> RegTemps {
*self.reg_temps.last().unwrap_or(&RegTemps::default())
}

/// Get a cached side exit, wrapping a counter if specified
Expand Down Expand Up @@ -1083,8 +1069,7 @@ impl Assembler
}

/// Convert Stack operands to memory operands
pub fn lower_stack(mut self) -> Assembler
{
pub fn lower_stack(mut self) -> Assembler {
// Convert Opnd::Stack to Opnd::Mem
fn mem_opnd(opnd: &Opnd) -> Opnd {
if let Opnd::Stack { idx, sp_offset, num_bits, .. } = *opnd {
Expand All @@ -1107,13 +1092,11 @@ impl Assembler

let mut asm = Assembler::new_with_label_names(take(&mut self.label_names), take(&mut self.side_exits));
let regs = Assembler::get_temp_regs();
let reg_temps = take(&mut self.reg_temps);
let mut iterator = self.into_draining_iter();

while let Some((index, mut insn)) = iterator.next_mapped() {
match &insn {
// The original insn is pushed to the new asm to satisfy ccall's reg_temps assertion.
Insn::RegTemps(_) => {} // noop
Insn::SpillTemp(opnd) => {
incr_counter!(temp_spill);
asm.mov(mem_opnd(opnd), reg_opnd(opnd, &regs));
Expand All @@ -1129,17 +1112,17 @@ impl Assembler
// Lower Opnd::Stack to Opnd::Reg or Opnd::Mem
let mut opnd_iter = insn.opnd_iter_mut();
while let Some(opnd) = opnd_iter.next() {
if let Opnd::Stack { idx, stack_size, sp_offset, num_bits } = *opnd {
*opnd = if opnd.stack_idx() < MAX_REG_TEMPS && reg_temps[index].get(opnd.stack_idx()) {
if let Opnd::Stack { idx, num_bits, stack_size, sp_offset, reg_temps } = *opnd {
*opnd = if opnd.stack_idx() < MAX_REG_TEMPS && reg_temps.unwrap().get(opnd.stack_idx()) {
reg_opnd(opnd, &regs)
} else {
mem_opnd(opnd)
};
}
}
asm.push_insn(insn);
}
}
asm.push_insn(insn);
iterator.map_insn_index(&mut asm);
}

Expand All @@ -1152,46 +1135,60 @@ impl Assembler
return;
}

assert_eq!(self.get_reg_temps(), self.ctx.get_reg_temps());
let mut reg_temps = self.get_reg_temps();

// Allocate a register if there's no conflict.
let mut reg_temps = self.ctx.get_reg_temps();
if reg_temps.conflicts_with(stack_idx) {
assert!(!reg_temps.get(stack_idx));
} else {
reg_temps.set(stack_idx, true);
self.set_reg_temps(reg_temps);
self.ctx.set_reg_temps(reg_temps);
}
}

/// Spill all live stack temps from registers to the stack
pub fn spill_temps(&mut self) {
assert_eq!(self.get_reg_temps(), self.ctx.get_reg_temps());

// Forget registers above the stack top
let mut reg_temps = self.get_reg_temps();
let mut reg_temps = self.ctx.get_reg_temps();
for stack_idx in self.ctx.get_stack_size()..MAX_REG_TEMPS {
reg_temps.set(stack_idx, false);
}
self.set_reg_temps(reg_temps);

// Spill live stack temps
if self.get_reg_temps() != RegTemps::default() {
self.comment(&format!("spill_temps: {:08b} -> {:08b}", self.get_reg_temps().as_u8(), RegTemps::default().as_u8()));
if self.ctx.get_reg_temps() != RegTemps::default() {
self.comment(&format!("spill_temps: {:08b} -> {:08b}", self.ctx.get_reg_temps().as_u8(), RegTemps::default().as_u8()));
for stack_idx in 0..u8::min(MAX_REG_TEMPS, self.ctx.get_stack_size()) {
if self.get_reg_temps().get(stack_idx) {
if self.ctx.get_reg_temps().get(stack_idx) {
let idx = self.ctx.get_stack_size() - 1 - stack_idx;
self.spill_temp(self.stack_opnd(idx.into()));
reg_temps.set(stack_idx, false);
}
}
self.ctx.set_reg_temps(reg_temps);
}

// Every stack temp should have been spilled
assert_eq!(self.get_reg_temps(), RegTemps::default());
assert_eq!(self.ctx.get_reg_temps(), RegTemps::default());
}

/// Update which stack temps are in a register
pub fn set_reg_temps(&mut self, reg_temps: RegTemps) {
if self.ctx.get_reg_temps() != reg_temps {
self.comment(&format!("reg_temps: {:08b} -> {:08b}", self.ctx.get_reg_temps().as_u8(), reg_temps.as_u8()));
self.ctx.set_reg_temps(reg_temps);
self.verify_reg_temps();
}
}

/// Assert there's no conflict in stack temp register allocation
fn verify_reg_temps(&self) {
for stack_idx in 0..MAX_REG_TEMPS {
if self.ctx.get_reg_temps().get(stack_idx) {
assert!(!self.ctx.get_reg_temps().conflicts_with(stack_idx));
}
}
}

/// Sets the out field on the various instructions that require allocated
/// registers because their output is used as the operand on a subsequent
/// instruction. This is our implementation of the linear scan algorithm.
Expand Down Expand Up @@ -1572,7 +1569,7 @@ impl Assembler {
}

pub fn ccall(&mut self, fptr: *const u8, opnds: Vec<Opnd>) -> Opnd {
assert_eq!(self.get_reg_temps(), RegTemps::default(), "temps must be spilled before ccall");
assert_eq!(self.ctx.get_reg_temps(), RegTemps::default(), "temps must be spilled before ccall");
let out = self.next_opnd_out(Opnd::match_num_bits(&opnds));
self.push_insn(Insn::CCall { fptr, opnds, out });
out
Expand Down Expand Up @@ -1810,23 +1807,10 @@ impl Assembler {
out
}

/// Update which stack temps are in a register
pub fn set_reg_temps(&mut self, reg_temps: RegTemps) {
if self.get_reg_temps() != reg_temps {
self.comment(&format!("reg_temps: {:08b} -> {:08b}", self.get_reg_temps().as_u8(), reg_temps.as_u8()));
self.push_insn(Insn::RegTemps(reg_temps));
self.ctx.set_reg_temps(self.get_reg_temps());
}
}

/// Spill a stack temp from a register to the stack
pub fn spill_temp(&mut self, opnd: Opnd) {
assert_eq!(self.get_reg_temps(), self.ctx.get_reg_temps());

if opnd.stack_idx() < MAX_REG_TEMPS && self.get_reg_temps().get(opnd.stack_idx()) {
self.push_insn(Insn::SpillTemp(opnd));
self.ctx.set_reg_temps(self.get_reg_temps());
}
fn spill_temp(&mut self, opnd: Opnd) {
assert!(self.ctx.get_reg_temps().get(opnd.stack_idx()));
self.push_insn(Insn::SpillTemp(opnd));
}

pub fn store(&mut self, dest: Opnd, src: Opnd) {
Expand Down
5 changes: 2 additions & 3 deletions yjit/src/backend/x86_64/mod.rs
View file
Expand Up @@ -737,15 +737,14 @@ impl Assembler
Insn::CSelGE { truthy, falsy, out } => {
emit_csel(cb, *truthy, *falsy, *out, cmovl);
}
Insn::LiveReg { .. } |
Insn::RegTemps(_) |
Insn::SpillTemp(_) => (), // just a reg alloc signal, no code
Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
Insn::PadInvalPatch => {
let code_size = cb.get_write_pos().saturating_sub(std::cmp::max(start_write_pos, cb.page_start_pos()));
if code_size < cb.jmp_ptr_bytes() {
nop(cb, (cb.jmp_ptr_bytes() - code_size) as u32);
}
}
Insn::SpillTemp(_) => unreachable!("Insn::SpillTemp should have been lowered by lower_stack"),
};

// On failure, jump to the next page and retry the current insn
Expand Down
11 changes: 3 additions & 8 deletions yjit/src/codegen.rs
View file
Expand Up @@ -805,7 +805,6 @@ pub fn gen_single_block(
let chain_depth = if asm.ctx.get_chain_depth() > 0 { format!(", chain_depth: {}", asm.ctx.get_chain_depth()) } else { "".to_string() };
asm.comment(&format!("Block: {} (ISEQ offset: {}{})", iseq_get_location(blockid.iseq, blockid_idx), blockid_idx, chain_depth));
}
asm.set_reg_temps(asm.ctx.get_reg_temps());

// For each instruction to compile
// NOTE: could rewrite this loop with a std::iter::Iterator
Expand Down Expand Up @@ -834,11 +833,9 @@ pub fn gen_single_block(

// stack_pop doesn't immediately deallocate a register for stack temps,
// but it's safe to do so at this instruction boundary.
assert_eq!(asm.get_reg_temps(), asm.ctx.get_reg_temps());
for stack_idx in asm.ctx.get_stack_size()..MAX_REG_TEMPS {
asm.ctx.dealloc_temp_reg(stack_idx);
}
asm.set_reg_temps(asm.ctx.get_reg_temps());

// If previous instruction requested to record the boundary
if jit.record_boundary_patch_point {
Expand Down Expand Up @@ -6113,7 +6110,7 @@ fn gen_send_iseq(
// pushed onto the stack that represents the parameters that weren't
// explicitly given a value and have a non-constant default.
let unspec_opnd = VALUE::fixnum_from_usize(unspecified_bits).as_u64();
asm.spill_temp(asm.stack_opnd(-1)); // avoid using a register for unspecified_bits
asm.ctx.dealloc_temp_reg(asm.stack_opnd(-1).stack_idx()); // avoid using a register for unspecified_bits
asm.mov(asm.stack_opnd(-1), unspec_opnd.into());
}

Expand Down Expand Up @@ -6243,10 +6240,8 @@ fn gen_send_iseq(
return_asm.ctx = asm.ctx.clone();
return_asm.stack_pop(sp_offset.try_into().unwrap());
let return_val = return_asm.stack_push(Type::Unknown);
if return_val.stack_idx() < MAX_REG_TEMPS {
// The callee writes a return value on stack. Update reg_temps accordingly.
return_asm.ctx.dealloc_temp_reg(return_val.stack_idx());
}
// The callee writes a return value on stack. Update reg_temps accordingly.
return_asm.ctx.dealloc_temp_reg(return_val.stack_idx());
return_asm.ctx.set_sp_offset(1);
return_asm.ctx.reset_chain_depth();

Expand Down

0 comments on commit d91062c

Please sign in to comment.