From d91062cfb23bcbe8d4da39e8cf278152e8258bc6 Mon Sep 17 00:00:00 2001 From: Takashi Kokubun Date: Wed, 19 Apr 2023 10:15:36 -0700 Subject: [PATCH] YJIT: Remove Insn::RegTemps --- yjit/src/backend/arm64/mod.rs | 5 +- yjit/src/backend/ir.rs | 140 +++++++++++++++------------------ yjit/src/backend/x86_64/mod.rs | 5 +- yjit/src/codegen.rs | 11 +-- yjit/src/core.rs | 17 ++-- 5 files changed, 81 insertions(+), 97 deletions(-) diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs index 3795062da8ab6e..ee7e0edfa6ce10 100644 --- a/yjit/src/backend/arm64/mod.rs +++ b/yjit/src/backend/arm64/mod.rs @@ -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 diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs index 2e85a660ec1b8a..cb6391bfcf153f 100644 --- a/yjit/src/backend/ir.rs +++ b/yjit/src/backend/ir.rs @@ -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 + }, // Low-level operands, for lowering Imm(i64), // Raw signed immediate @@ -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, } } @@ -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 }, @@ -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", @@ -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) | @@ -816,7 +822,6 @@ impl<'a> InsnOpndMutIterator<'a> { Insn::Jz(_) | Insn::Label(_) | Insn::LeaLabel { .. } | - Insn::RegTemps(_) | Insn::PadInvalPatch | Insn::PosMarker(_) => None, Insn::CPopInto(opnd) | @@ -929,10 +934,6 @@ pub struct Assembler { /// Index of the last insn using the output of this insn pub(super) live_ranges: Vec, - /// Parallel vec with insns - /// Bitmap of which temps are in a register for this insn - pub(super) reg_temps: Vec, - /// Names of labels pub(super) label_names: Vec, @@ -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, @@ -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; @@ -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() { @@ -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 @@ -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 { @@ -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, ®s)); @@ -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, ®s) } else { mem_opnd(opnd) }; } } + asm.push_insn(insn); } } - asm.push_insn(insn); iterator.map_insn_index(&mut asm); } @@ -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. @@ -1572,7 +1569,7 @@ impl Assembler { } pub fn ccall(&mut self, fptr: *const u8, opnds: Vec) -> 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 @@ -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) { diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs index 0fa947d1ba2ad1..69e8d24ed7b71c 100644 --- a/yjit/src/backend/x86_64/mod.rs +++ b/yjit/src/backend/x86_64/mod.rs @@ -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 diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs index 640940777c803b..599de226a43f36 100644 --- a/yjit/src/codegen.rs +++ b/yjit/src/codegen.rs @@ -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 @@ -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 { @@ -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()); } @@ -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(); diff --git a/yjit/src/core.rs b/yjit/src/core.rs index 5362618a285d27..5f14914472d67e 100644 --- a/yjit/src/core.rs +++ b/yjit/src/core.rs @@ -1623,9 +1623,11 @@ impl Context { /// Stop using a register for a given stack temp. pub fn dealloc_temp_reg(&mut self, stack_idx: u8) { - let mut reg_temps = self.get_reg_temps(); - reg_temps.set(stack_idx, false); - self.set_reg_temps(reg_temps); + if stack_idx < MAX_REG_TEMPS { + let mut reg_temps = self.get_reg_temps(); + reg_temps.set(stack_idx, false); + self.set_reg_temps(reg_temps); + } } /// Get the type of an instruction operand @@ -1911,7 +1913,6 @@ impl Assembler { } // Allocate a register to the stack operand - assert_eq!(self.ctx.reg_temps, self.get_reg_temps()); if self.ctx.stack_size < MAX_REG_TEMPS { self.alloc_temp_reg(self.ctx.stack_size); } @@ -1982,7 +1983,13 @@ impl Assembler { /// Get an operand pointing to a slot on the temp stack pub fn stack_opnd(&self, idx: i32) -> Opnd { - Opnd::Stack { idx, stack_size: self.ctx.stack_size, sp_offset: self.ctx.sp_offset, num_bits: 64 } + Opnd::Stack { + idx, + num_bits: 64, + stack_size: self.ctx.stack_size, + sp_offset: self.ctx.sp_offset, + reg_temps: None, // push_insn will set this + } } }