Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

YJIT: Stack temp register allocation #7651

Merged
merged 1 commit into from
Apr 4, 2023
Merged

YJIT: Stack temp register allocation #7651

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
4 changes: 4 additions & 0 deletions .github/workflows/yjit-ubuntu.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,10 @@ jobs:
configure: "--enable-yjit=dev"
yjit_opts: "--yjit-call-threshold=1 --yjit-verify-ctx"

- test_task: "check"
configure: "--enable-yjit=dev"
yjit_opts: "--yjit-call-threshold=1 --yjit-temp-regs=5"

- test_task: "test-all TESTS=--repeat-count=2"
configure: "--enable-yjit=dev"

Expand Down
3 changes: 3 additions & 0 deletions yjit.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,6 +268,9 @@ def _print_stats # :nodoc:

$stderr.puts "iseq_stack_too_large: " + format_number(13, stats[:iseq_stack_too_large])
$stderr.puts "iseq_too_long: " + format_number(13, stats[:iseq_too_long])
$stderr.puts "temp_reg_opnd: " + format_number(13, stats[:temp_reg_opnd])
$stderr.puts "temp_mem_opnd: " + format_number(13, stats[:temp_mem_opnd])
$stderr.puts "temp_spill: " + format_number(13, stats[:temp_spill])
$stderr.puts "bindings_allocations: " + format_number(13, stats[:binding_allocations])
$stderr.puts "bindings_set: " + format_number(13, stats[:binding_set])
$stderr.puts "compilation_failure: " + format_number(13, compilation_failure) if compilation_failure != 0
Expand Down
11 changes: 10 additions & 1 deletion yjit/src/backend/arm64/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,13 @@ impl Assembler
vec![X11_REG, X12_REG, X13_REG]
}

/// Get the list of registers that can be used for stack temps.
pub fn get_temp_regs() -> Vec<Reg> {
// FIXME: arm64 is not supported yet. Insn::Store doesn't support registers
// in its dest operand. Currently crashing at split_memory_address.
vec![]
}

/// Get a list of all of the caller-saved registers
pub fn get_caller_save_regs() -> Vec<Reg> {
vec![X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
Expand Down Expand Up @@ -1046,7 +1053,9 @@ impl Assembler
Insn::CSelGE { truthy, falsy, out } => {
csel(cb, out.into(), truthy.into(), falsy.into(), Condition::GE);
}
Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
Insn::LiveReg { .. } |
Insn::RegTemps(_) |
Insn::SpillTemp(_) => (), // 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()))) < JMP_PTR_BYTES && !cb.has_dropped_bytes() {
nop(cb);
Expand Down
189 changes: 179 additions & 10 deletions yjit/src/backend/ir.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,9 @@ use std::mem::take;
use crate::cruby::{VALUE, SIZEOF_VALUE_I32};
use crate::virtualmem::{CodePtr};
use crate::asm::{CodeBlock, uimm_num_bits, imm_num_bits};
use crate::core::{Context, Type, TempMapping};
use crate::core::{Context, Type, TempMapping, RegTemps, MAX_REG_TEMPS, MAX_TEMP_TYPES};
use crate::options::*;
use crate::stats::*;

#[cfg(target_arch = "x86_64")]
use crate::backend::x86_64::*;
Expand Down Expand Up @@ -73,7 +74,7 @@ pub enum Opnd
InsnOut{ idx: usize, num_bits: u8 },

// Pointer to a slot on the VM stack
Stack { idx: i32, sp_offset: i8, num_bits: u8 },
Stack { idx: i32, stack_size: u8, sp_offset: i8, num_bits: u8 },

// Low-level operands, for lowering
Imm(i64), // Raw signed immediate
Expand Down Expand Up @@ -162,7 +163,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, sp_offset, .. } => Some(Opnd::Stack { idx, sp_offset, num_bits }),
Opnd::Stack { idx, stack_size, sp_offset, .. } => Some(Opnd::Stack { idx, stack_size, sp_offset, num_bits }),
_ => None,
}
}
Expand Down Expand Up @@ -216,6 +217,26 @@ impl Opnd
pub fn match_num_bits(opnds: &[Opnd]) -> u8 {
Self::match_num_bits_iter(opnds.iter())
}

/// Calculate Opnd::Stack's index from the stack bottom.
pub fn stack_idx(&self) -> u8 {
match self {
Opnd::Stack { idx, stack_size, .. } => {
(*stack_size as isize - *idx as isize - 1) as u8
},
_ => unreachable!(),
}
}

/// Get the index for stack temp registers.
pub fn reg_idx(&self) -> usize {
match self {
Opnd::Stack { .. } => {
self.stack_idx() as usize % get_option!(num_temp_regs)
},
_ => unreachable!(),
}
}
}

impl From<usize> for Opnd {
Expand Down Expand Up @@ -408,6 +429,9 @@ 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 @@ -443,6 +467,9 @@ pub enum Insn {
/// Shift a value right by a certain amount (signed).
RShift { opnd: Opnd, shift: Opnd, out: Opnd },

/// Spill a stack temp from a register into memory
SpillTemp(Opnd),

// Low-level instruction to store a value to memory.
Store { dest: Opnd, src: Opnd },

Expand Down Expand Up @@ -514,6 +541,7 @@ impl Insn {
Insn::LeaLabel { .. } => "LeaLabel",
Insn::Lea { .. } => "Lea",
Insn::LiveReg { .. } => "LiveReg",
Insn::RegTemps(_) => "RegTemps",
Insn::Load { .. } => "Load",
Insn::LoadInto { .. } => "LoadInto",
Insn::LoadSExt { .. } => "LoadSExt",
Expand All @@ -524,6 +552,7 @@ impl Insn {
Insn::PadInvalPatch => "PadEntryExit",
Insn::PosMarker(_) => "PosMarker",
Insn::RShift { .. } => "RShift",
Insn::SpillTemp(_) => "SpillTemp",
Insn::Store { .. } => "Store",
Insn::Sub { .. } => "Sub",
Insn::Test { .. } => "Test",
Expand Down Expand Up @@ -658,6 +687,7 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
Insn::Jz(_) |
Insn::Label(_) |
Insn::LeaLabel { .. } |
Insn::RegTemps(_) |
Insn::PadInvalPatch |
Insn::PosMarker(_) => None,
Insn::CPopInto(opnd) |
Expand All @@ -668,7 +698,8 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
Insn::LiveReg { opnd, .. } |
Insn::Load { opnd, .. } |
Insn::LoadSExt { opnd, .. } |
Insn::Not { opnd, .. } => {
Insn::Not { opnd, .. } |
Insn::SpillTemp(opnd) => {
match self.idx {
0 => {
self.idx += 1;
Expand Down Expand Up @@ -755,6 +786,7 @@ impl<'a> InsnOpndMutIterator<'a> {
Insn::Jz(_) |
Insn::Label(_) |
Insn::LeaLabel { .. } |
Insn::RegTemps(_) |
Insn::PadInvalPatch |
Insn::PosMarker(_) => None,
Insn::CPopInto(opnd) |
Expand All @@ -765,7 +797,8 @@ impl<'a> InsnOpndMutIterator<'a> {
Insn::LiveReg { opnd, .. } |
Insn::Load { opnd, .. } |
Insn::LoadSExt { opnd, .. } |
Insn::Not { opnd, .. } => {
Insn::Not { opnd, .. } |
Insn::SpillTemp(opnd) => {
match self.idx {
0 => {
self.idx += 1;
Expand Down Expand Up @@ -857,6 +890,10 @@ 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 @@ -871,6 +908,7 @@ impl Assembler
Self {
insns: Vec::default(),
live_ranges: Vec::default(),
reg_temps: Vec::default(),
label_names
}
}
Expand Down Expand Up @@ -905,8 +943,33 @@ impl Assembler
}
}

// 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));
}
}

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())
}

/// Create a new label instance that we can jump to
Expand All @@ -922,22 +985,113 @@ impl Assembler
/// Convert Stack operands to memory operands
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 {
incr_counter!(temp_mem_opnd);
Opnd::mem(num_bits, SP, (sp_offset as i32 - idx - 1) * SIZEOF_VALUE_I32)
} else {
unreachable!()
}
}

// Convert Opnd::Stack to Opnd::Reg
fn reg_opnd(opnd: &Opnd, regs: &Vec<Reg>) -> Opnd {
if let Opnd::Stack { num_bits, .. } = *opnd {
incr_counter!(temp_reg_opnd);
Opnd::Reg(regs[opnd.reg_idx()]).with_num_bits(num_bits).unwrap()
} else {
unreachable!()
}
}

let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
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_unmapped() {
let mut opnd_iter = insn.opnd_iter_mut();
while let Some(opnd) = opnd_iter.next() {
if let Opnd::Stack { idx, sp_offset, num_bits } = *opnd {
*opnd = Opnd::mem(num_bits, SP, (sp_offset as i32 - idx - 1) * SIZEOF_VALUE_I32);
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));
}
_ => {
// next_mapped() doesn't map out_opnd. So we need to map it here.
if insn.out_opnd().is_some() {
let out_num_bits = Opnd::match_num_bits_iter(insn.opnd_iter());
let out = insn.out_opnd_mut().unwrap();
*out = asm.next_opnd_out(out_num_bits);
}

// 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()) {
reg_opnd(opnd, &regs)
} else {
mem_opnd(opnd)
};
}
}
}
}
asm.push_insn(insn);
iterator.map_insn_index(&mut asm);
}

asm
}

/// Allocate a register to a stack temp if available.
pub fn alloc_temp_reg(&mut self, ctx: &mut Context, stack_idx: u8) {
if get_option!(num_temp_regs) == 0 {
return;
}

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

// Allocate a register if there's no conflict.
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);
ctx.set_reg_temps(reg_temps);
}
}

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

// Forget registers above the stack top
let mut reg_temps = self.get_reg_temps();
for stack_idx in 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()));
for stack_idx in 0..u8::min(MAX_REG_TEMPS, ctx.get_stack_size()) {
if self.get_reg_temps().get(stack_idx) {
let idx = ctx.get_stack_size() - 1 - stack_idx;
self.spill_temp(ctx.stack_opnd(idx.into()));
}
}
}

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

/// 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 @@ -1318,6 +1472,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");
let out = self.next_opnd_out(Opnd::match_num_bits(&opnds));
self.push_insn(Insn::CCall { fptr, opnds, out });
out
Expand Down Expand Up @@ -1545,6 +1700,20 @@ 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));
}
}

/// Spill a stack temp from a register to the stack
pub fn spill_temp(&mut self, opnd: Opnd) {
assert!(self.get_reg_temps().get(opnd.stack_idx()));
self.push_insn(Insn::SpillTemp(opnd));
}

pub fn store(&mut self, dest: Opnd, src: Opnd) {
self.push_insn(Insn::Store { dest, src });
}
Expand Down