1berry/ruby
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

ZJIT: Implement side exits for entry frames (#13469)

Browse Source 
Co-authored-by: Max Bernstein <tekknolagi@gmail.com>
Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>
This commit is contained in:
Takashi Kokubun 2025-06-04 17:14:40 -04:00 committed by GitHub
parent 99cc100cdf
commit 112c34252d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
Notes: git 2025-06-04 21:14:53 +00:00 
Merged-By: k0kubun <takashikkbn@gmail.com>
9 changed files with 302 additions and 112 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
test/ruby/test_zjit.rb
View File

@ -94,6 +94,24 @@ class TestZJIT < Test::Unit::TestCase
}, call_threshold: 2
end
def test_opt_plus_type_guard_exit
assert_compiles '[3, 3.0]', %q{
def test(a) = 1 + a
test(1) # profile opt_plus
[test(2), test(2.0)]
}, call_threshold: 2
end
def test_opt_plus_type_guard_nested_exit
omit 'rewind_caller_frames is not implemented yet'
assert_compiles '[3, 3.0]', %q{
def side_exit(n) = 1 + n
def jit_frame(n) = 1 + side_exit(n)
def entry(n) = jit_frame(n)
[entry(2), entry(2.0)]
}, call_threshold: 2
end
# Test argument ordering
def test_opt_minus
assert_compiles '2', %q{

4
zjit/src/asm/arm64/mod.rs
View File

@ -644,7 +644,7 @@ pub fn mov(cb: &mut CodeBlock, rd: A64Opnd, rm: A64Opnd) {
LogicalImm::mov(rd.reg_no, bitmask_imm, rd.num_bits).into()
},
_ => panic!("Invalid operand combination to mov instruction")
_ => panic!("Invalid operand combination to mov instruction: {rd:?}, {rm:?}")
};
cb.write_bytes(&bytes);
@ -940,7 +940,7 @@ pub fn stur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
LoadStore::stur(rt.reg_no, rn.base_reg_no, rn.disp as i16, rn.num_bits).into()
},
_ => panic!("Invalid operand combination to stur instruction.")
_ => panic!("Invalid operand combination to stur instruction: {rt:?}, {rn:?}")
};
cb.write_bytes(&bytes);

5
zjit/src/asm/arm64/opnd.rs
View File

@ -119,6 +119,9 @@ pub const X20_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 20 };
pub const X21_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 21 };
pub const X22_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 22 };
// link register
pub const X30_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 30 };
// zero register
pub const XZR_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 31 };
@ -153,7 +156,7 @@ pub const X26: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 26 });
pub const X27: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 27 });
pub const X28: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 28 });
pub const X29: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 29 });
pub const X30: A64Opnd = A64Opnd::Reg(A64Reg { num_bits: 64, reg_no: 30 });
pub const X30: A64Opnd = A64Opnd::Reg(X30_REG);
pub const X31: A64Opnd = A64Opnd::Reg(XZR_REG);
// 32-bit registers

18
zjit/src/backend/arm64/mod.rs
View File

@ -211,6 +211,11 @@ impl Assembler
vec![X1_REG, X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
}
/// Get the address that the current frame returns to
pub fn return_addr_opnd() -> Opnd {
Opnd::Reg(X30_REG)
}
/// Split platform-specific instructions
/// The transformations done here are meant to make our lives simpler in later
/// stages of the compilation pipeline.
@ -757,7 +762,7 @@ impl Assembler
/// called when lowering any of the conditional jump instructions.
fn emit_conditional_jump<const CONDITION: u8>(cb: &mut CodeBlock, target: Target) {
match target {
Target::CodePtr(dst_ptr) | Target::SideExitPtr(dst_ptr) => {
Target::CodePtr(dst_ptr) => {
let dst_addr = dst_ptr.as_offset();
let src_addr = cb.get_write_ptr().as_offset();
@ -829,8 +834,10 @@ impl Assembler
}
/// Emit a CBZ or CBNZ which branches when a register is zero or non-zero
fn emit_cmp_zero_jump(cb: &mut CodeBlock, reg: A64Opnd, branch_if_zero: bool, target: Target) {
if let Target::SideExitPtr(dst_ptr) = target {
fn emit_cmp_zero_jump(_cb: &mut CodeBlock, _reg: A64Opnd, _branch_if_zero: bool, target: Target) {
if let Target::Label(_) = target {
unimplemented!("this should be re-implemented with Label for side exits");
/*
let dst_addr = dst_ptr.as_offset();
let src_addr = cb.get_write_ptr().as_offset();
@ -862,6 +869,7 @@ impl Assembler
br(cb, Assembler::SCRATCH0);
}
*/
} else {
unreachable!("We should only generate Joz/Jonz with side-exit targets");
}
@ -1162,9 +1170,6 @@ impl Assembler
Target::CodePtr(dst_ptr) => {
emit_jmp_ptr(cb, dst_ptr, true);
},
Target::SideExitPtr(dst_ptr) => {
emit_jmp_ptr(cb, dst_ptr, false);
},
Target::Label(label_idx) => {
// Here we're going to save enough space for
// ourselves and then come back and write the
@ -1297,6 +1302,7 @@ impl Assembler
pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Option<(CodePtr, Vec<u32>)> {
let asm = self.arm64_split();
let mut asm = asm.alloc_regs(regs);
asm.compile_side_exits()?;
// Create label instances in the code block
for (idx, name) in asm.label_names.iter().enumerate() {

265
zjit/src/backend/lir.rs
View File

@ -1,6 +1,9 @@
use std::collections::HashMap;
use std::fmt;
use std::mem::take;
use crate::{cruby::VALUE, hir::FrameState};
use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32, VM_ENV_DATA_SIZE};
use crate::state::ZJITState;
use crate::{cruby::VALUE};
use crate::backend::current::*;
use crate::virtualmem::CodePtr;
use crate::asm::{CodeBlock, Label};
@ -273,9 +276,7 @@ pub enum Target
/// Pointer to a piece of ZJIT-generated code
CodePtr(CodePtr),
// Side exit with a counter
SideExit(FrameState),
/// Pointer to a side exit code
SideExitPtr(CodePtr),
SideExit { pc: *const VALUE, stack: Vec<Opnd>, locals: Vec<Opnd> },
/// A label within the generated code
Label(Label),
}
@ -292,7 +293,6 @@ impl Target
pub fn unwrap_code_ptr(&self) -> CodePtr {
match self {
Target::CodePtr(ptr) => *ptr,
Target::SideExitPtr(ptr) => *ptr,
_ => unreachable!("trying to unwrap {:?} into code ptr", self)
}
}
@ -539,11 +539,11 @@ impl Insn {
Insn::Jne(target) |
Insn::Jnz(target) |
Insn::Jo(target) |
Insn::Jz(target) |
Insn::Label(target) |
Insn::JoMul(target) |
Insn::Jz(target) |
Insn::Joz(_, target) |
Insn::Jonz(_, target) |
Insn::Label(target) |
Insn::LeaJumpTarget { target, .. } => {
Some(target)
}
@ -697,7 +697,11 @@ impl Insn {
Insn::Jne(target) |
Insn::Jnz(target) |
Insn::Jo(target) |
Insn::JoMul(target) |
Insn::Jz(target) |
Insn::Joz(_, target) |
Insn::Jonz(_, target) |
Insn::Label(target) |
Insn::LeaJumpTarget { target, .. } => Some(target),
_ => None
}
@ -731,6 +735,63 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
fn next(&mut self) -> Option<Self::Item> {
match self.insn {
Insn::Jbe(target) |
Insn::Jb(target) |
Insn::Je(target) |
Insn::Jl(target) |
Insn::Jg(target) |
Insn::Jge(target) |
Insn::Jmp(target) |
Insn::Jne(target) |
Insn::Jnz(target) |
Insn::Jo(target) |
Insn::JoMul(target) |
Insn::Jz(target) |
Insn::Label(target) |
Insn::LeaJumpTarget { target, .. } => {
if let Target::SideExit { stack, locals, .. } = target {
let stack_idx = self.idx;
if stack_idx < stack.len() {
let opnd = &stack[stack_idx];
self.idx += 1;
return Some(opnd);
}
let local_idx = self.idx - stack.len();
if local_idx < locals.len() {
let opnd = &locals[local_idx];
self.idx += 1;
return Some(opnd);
}
}
None
}
Insn::Joz(opnd, target) |
Insn::Jonz(opnd, target) => {
if self.idx == 0 {
self.idx += 1;
return Some(opnd);
}
if let Target::SideExit { stack, locals, .. } = target {
let stack_idx = self.idx - 1;
if stack_idx < stack.len() {
let opnd = &stack[stack_idx];
self.idx += 1;
return Some(opnd);
}
let local_idx = stack_idx - stack.len();
if local_idx < locals.len() {
let opnd = &locals[local_idx];
self.idx += 1;
return Some(opnd);
}
}
None
}
Insn::BakeString(_) |
Insn::Breakpoint |
Insn::Comment(_) |
@ -739,20 +800,6 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
Insn::CPushAll |
Insn::FrameSetup |
Insn::FrameTeardown |
Insn::Jbe(_) |
Insn::Jb(_) |
Insn::Je(_) |
Insn::Jl(_) |
Insn::Jg(_) |
Insn::Jge(_) |
Insn::Jmp(_) |
Insn::Jne(_) |
Insn::Jnz(_) |
Insn::Jo(_) |
Insn::JoMul(_) |
Insn::Jz(_) |
Insn::Label(_) |
Insn::LeaJumpTarget { .. } |
Insn::PadInvalPatch |
Insn::PosMarker(_) => None,
@ -764,8 +811,6 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
Insn::LiveReg { opnd, .. } |
Insn::Load { opnd, .. } |
Insn::LoadSExt { opnd, .. } |
Insn::Joz(opnd, _) |
Insn::Jonz(opnd, _) |
Insn::Not { opnd, .. } => {
match self.idx {
0 => {
@ -845,6 +890,63 @@ impl<'a> InsnOpndMutIterator<'a> {
pub(super) fn next(&mut self) -> Option<&mut Opnd> {
match self.insn {
Insn::Jbe(target) |
Insn::Jb(target) |
Insn::Je(target) |
Insn::Jl(target) |
Insn::Jg(target) |
Insn::Jge(target) |
Insn::Jmp(target) |
Insn::Jne(target) |
Insn::Jnz(target) |
Insn::Jo(target) |
Insn::JoMul(target) |
Insn::Jz(target) |
Insn::Label(target) |
Insn::LeaJumpTarget { target, .. } => {
if let Target::SideExit { stack, locals, .. } = target {
let stack_idx = self.idx;
if stack_idx < stack.len() {
let opnd = &mut stack[stack_idx];
self.idx += 1;
return Some(opnd);
}
let local_idx = self.idx - stack.len();
if local_idx < locals.len() {
let opnd = &mut locals[local_idx];
self.idx += 1;
return Some(opnd);
}
}
None
}
Insn::Joz(opnd, target) |
Insn::Jonz(opnd, target) => {
if self.idx == 0 {
self.idx += 1;
return Some(opnd);
}
if let Target::SideExit { stack, locals, .. } = target {
let stack_idx = self.idx - 1;
if stack_idx < stack.len() {
let opnd = &mut stack[stack_idx];
self.idx += 1;
return Some(opnd);
}
let local_idx = stack_idx - stack.len();
if local_idx < locals.len() {
let opnd = &mut locals[local_idx];
self.idx += 1;
return Some(opnd);
}
}
None
}
Insn::BakeString(_) |
Insn::Breakpoint |
Insn::Comment(_) |
@ -853,20 +955,6 @@ impl<'a> InsnOpndMutIterator<'a> {
Insn::CPushAll |
Insn::FrameSetup |
Insn::FrameTeardown |
Insn::Jbe(_) |
Insn::Jb(_) |
Insn::Je(_) |
Insn::Jl(_) |
Insn::Jg(_) |
Insn::Jge(_) |
Insn::Jmp(_) |
Insn::Jne(_) |
Insn::Jnz(_) |
Insn::Jo(_) |
Insn::JoMul(_) |
Insn::Jz(_) |
Insn::Label(_) |
Insn::LeaJumpTarget { .. } |
Insn::PadInvalPatch |
Insn::PosMarker(_) => None,
@ -878,8 +966,6 @@ impl<'a> InsnOpndMutIterator<'a> {
Insn::LiveReg { opnd, .. } |
Insn::Load { opnd, .. } |
Insn::LoadSExt { opnd, .. } |
Insn::Joz(opnd, _) |
Insn::Jonz(opnd, _) |
Insn::Not { opnd, .. } => {
match self.idx {
0 => {
@ -1649,10 +1735,8 @@ impl Assembler
/// Compile the instructions down to machine code.
/// Can fail due to lack of code memory and inopportune code placement, among other reasons.
#[must_use]
pub fn compile(mut self, cb: &mut CodeBlock) -> Option<(CodePtr, Vec<u32>)>
pub fn compile(self, cb: &mut CodeBlock) -> Option<(CodePtr, Vec<u32>)>
{
self.compile_side_exits(cb)?;
#[cfg(feature = "disasm")]
let start_addr = cb.get_write_ptr();
let alloc_regs = Self::get_alloc_regs();
@ -1669,47 +1753,74 @@ impl Assembler
/// Compile Target::SideExit and convert it into Target::CodePtr for all instructions
#[must_use]
pub fn compile_side_exits(&mut self, cb: &mut CodeBlock) -> Option<()> {
for insn in self.insns.iter_mut() {
if let Some(target) = insn.target_mut() {
if let Target::SideExit(state) = target {
let side_exit_ptr = cb.get_write_ptr();
let mut asm = Assembler::new();
asm_comment!(asm, "side exit: {state}");
asm.ccall(Self::rb_zjit_side_exit as *const u8, vec![]);
asm.compile(cb)?;
*target = Target::SideExitPtr(side_exit_ptr);
pub fn compile_side_exits(&mut self) -> Option<()> {
let mut targets = HashMap::new();
for (idx, insn) in self.insns.iter().enumerate() {
if let Some(target @ Target::SideExit { .. }) = insn.target() {
targets.insert(idx, target.clone());
}
}
for (idx, target) in targets {
// Compile a side exit. Note that this is past the split pass and alloc_regs(),
// so you can't use a VReg or an instruction that needs to be split.
if let Target::SideExit { pc, stack, locals } = target {
let side_exit_label = self.new_label("side_exit".into());
self.write_label(side_exit_label.clone());
// Load an operand that cannot be used as a source of Insn::Store
fn split_store_source(asm: &mut Assembler, opnd: Opnd) -> Opnd {
if matches!(opnd, Opnd::Mem(_) | Opnd::Value(_)) ||
(cfg!(target_arch = "aarch64") && matches!(opnd, Opnd::UImm(_))) {
asm.load_into(Opnd::Reg(Assembler::SCRATCH_REG), opnd);
Opnd::Reg(Assembler::SCRATCH_REG)
} else {
opnd
}
}
asm_comment!(self, "write stack slots: {stack:?}");
for (idx, &opnd) in stack.iter().enumerate() {
let opnd = split_store_source(self, opnd);
self.store(Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32), opnd);
}
asm_comment!(self, "write locals: {locals:?}");
for (idx, &opnd) in locals.iter().enumerate() {
let opnd = split_store_source(self, opnd);
self.store(Opnd::mem(64, SP, (-(VM_ENV_DATA_SIZE as i32) - locals.len() as i32 + idx as i32) * SIZEOF_VALUE_I32), opnd);
}
asm_comment!(self, "save cfp->pc");
self.load_into(Opnd::Reg(Assembler::SCRATCH_REG), Opnd::const_ptr(pc as *const u8));
self.store(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_PC), Opnd::Reg(Assembler::SCRATCH_REG));
asm_comment!(self, "save cfp->sp");
self.lea_into(Opnd::Reg(Assembler::SCRATCH_REG), Opnd::mem(64, SP, stack.len() as i32 * SIZEOF_VALUE_I32));
let cfp_sp = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SP);
self.store(cfp_sp, Opnd::Reg(Assembler::SCRATCH_REG));
asm_comment!(self, "rewind caller frames");
self.mov(C_ARG_OPNDS[0], Assembler::return_addr_opnd());
self.ccall(Self::rewind_caller_frames as *const u8, vec![]);
asm_comment!(self, "exit to the interpreter");
self.frame_teardown();
self.mov(C_RET_OPND, Opnd::UImm(Qundef.as_u64()));
self.cret(C_RET_OPND);
*self.insns[idx].target_mut().unwrap() = side_exit_label;
}
}
Some(())
}
#[unsafe(no_mangle)]
extern "C" fn rb_zjit_side_exit() {
unimplemented!("side exits are not implemented yet");
extern "C" fn rewind_caller_frames(addr: *const u8) {
if ZJITState::is_iseq_return_addr(addr) {
unimplemented!("Can't side-exit from JIT-JIT call: rewind_caller_frames is not implemented yet");
}
}
/*
/// Compile with a limited number of registers. Used only for unit tests.
#[cfg(test)]
pub fn compile_with_num_regs(self, cb: &mut CodeBlock, num_regs: usize) -> (CodePtr, Vec<u32>)
{
let mut alloc_regs = Self::get_alloc_regs();
let alloc_regs = alloc_regs.drain(0..num_regs).collect();
self.compile_with_regs(cb, None, alloc_regs).unwrap()
}
/// Return true if the next ccall() is expected to be leaf.
pub fn get_leaf_ccall(&mut self) -> bool {
self.leaf_ccall
}
/// Assert that the next ccall() is going to be leaf.
pub fn expect_leaf_ccall(&mut self) {
self.leaf_ccall = true;
}
*/
}
impl fmt::Debug for Assembler {
@ -1970,6 +2081,10 @@ impl Assembler {
out
}
pub fn lea_into(&mut self, out: Opnd, opnd: Opnd) {
self.push_insn(Insn::Lea { opnd, out });
}
#[must_use]
pub fn lea_jump_target(&mut self, target: Target) -> Opnd {
let out = self.new_vreg(Opnd::DEFAULT_NUM_BITS);

28
zjit/src/backend/x86_64/mod.rs
View File

@ -109,6 +109,11 @@ impl Assembler
vec![RAX_REG, RCX_REG, RDX_REG, RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG, R11_REG]
}
/// Get the address that the current frame returns to
pub fn return_addr_opnd() -> Opnd {
Opnd::mem(64, Opnd::Reg(RSP_REG), 0)
}
// These are the callee-saved registers in the x86-64 SysV ABI
// RBX, RSP, RBP, and R12–R15
@ -665,7 +670,7 @@ impl Assembler
// Conditional jump to a label
Insn::Jmp(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jmp_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jmp_ptr(cb, code_ptr),
Target::Label(label) => jmp_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -673,7 +678,7 @@ impl Assembler
Insn::Je(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => je_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => je_ptr(cb, code_ptr),
Target::Label(label) => je_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -681,7 +686,7 @@ impl Assembler
Insn::Jne(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jne_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jne_ptr(cb, code_ptr),
Target::Label(label) => jne_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -689,7 +694,7 @@ impl Assembler
Insn::Jl(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jl_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jl_ptr(cb, code_ptr),
Target::Label(label) => jl_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -697,7 +702,7 @@ impl Assembler
Insn::Jg(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jg_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jg_ptr(cb, code_ptr),
Target::Label(label) => jg_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -705,7 +710,7 @@ impl Assembler
Insn::Jge(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jge_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jge_ptr(cb, code_ptr),
Target::Label(label) => jge_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -713,7 +718,7 @@ impl Assembler
Insn::Jbe(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jbe_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jbe_ptr(cb, code_ptr),
Target::Label(label) => jbe_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -721,7 +726,7 @@ impl Assembler
Insn::Jb(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jb_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jb_ptr(cb, code_ptr),
Target::Label(label) => jb_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -729,7 +734,7 @@ impl Assembler
Insn::Jz(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jz_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jz_ptr(cb, code_ptr),
Target::Label(label) => jz_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -737,7 +742,7 @@ impl Assembler
Insn::Jnz(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jnz_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jnz_ptr(cb, code_ptr),
Target::Label(label) => jnz_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -746,7 +751,7 @@ impl Assembler
Insn::Jo(target) |
Insn::JoMul(target) => {
match *target {
Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jo_ptr(cb, code_ptr),
Target::CodePtr(code_ptr) => jo_ptr(cb, code_ptr),
Target::Label(label) => jo_label(cb, label),
Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exits"),
}
@ -831,6 +836,7 @@ impl Assembler
pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Option<(CodePtr, Vec<u32>)> {
let asm = self.x86_split();
let mut asm = asm.alloc_regs(regs);
asm.compile_side_exits()?;
// Create label instances in the code block
for (idx, name) in asm.label_names.iter().enumerate() {

53
zjit/src/codegen.rs
View File

@ -260,9 +260,9 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
Insn::SendWithoutBlock { call_info, cd, state, self_val, args, .. } => gen_send_without_block(jit, asm, call_info, *cd, &function.frame_state(*state), self_val, args)?,
Insn::SendWithoutBlockDirect { iseq, self_val, args, .. } => gen_send_without_block_direct(cb, jit, asm, *iseq, opnd!(self_val), args)?,
Insn::Return { val } => return Some(gen_return(asm, opnd!(val))?),
Insn::FixnumAdd { left, right, state } => gen_fixnum_add(asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
Insn::FixnumSub { left, right, state } => gen_fixnum_sub(asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
Insn::FixnumMult { left, right, state } => gen_fixnum_mult(asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
Insn::FixnumAdd { left, right, state } => gen_fixnum_add(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
Insn::FixnumSub { left, right, state } => gen_fixnum_sub(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
Insn::FixnumMult { left, right, state } => gen_fixnum_mult(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
Insn::FixnumEq { left, right } => gen_fixnum_eq(asm, opnd!(left), opnd!(right))?,
Insn::FixnumNeq { left, right } => gen_fixnum_neq(asm, opnd!(left), opnd!(right))?,
Insn::FixnumLt { left, right } => gen_fixnum_lt(asm, opnd!(left), opnd!(right))?,
@ -270,8 +270,8 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
Insn::FixnumGt { left, right } => gen_fixnum_gt(asm, opnd!(left), opnd!(right))?,
Insn::FixnumGe { left, right } => gen_fixnum_ge(asm, opnd!(left), opnd!(right))?,
Insn::Test { val } => gen_test(asm, opnd!(val))?,
Insn::GuardType { val, guard_type, state } => gen_guard_type(asm, opnd!(val), *guard_type, &function.frame_state(*state))?,
Insn::GuardBitEquals { val, expected, state } => gen_guard_bit_equals(asm, opnd!(val), *expected, &function.frame_state(*state))?,
Insn::GuardType { val, guard_type, state } => gen_guard_type(jit, asm, opnd!(val), *guard_type, &function.frame_state(*state))?,
Insn::GuardBitEquals { val, expected, state } => gen_guard_bit_equals(jit, asm, opnd!(val), *expected, &function.frame_state(*state))?,
Insn::PatchPoint(_) => return Some(()), // For now, rb_zjit_bop_redefined() panics. TODO: leave a patch point and fix rb_zjit_bop_redefined()
Insn::CCall { cfun, args, name: _, return_type: _, elidable: _ } => gen_ccall(jit, asm, *cfun, args)?,
_ => {
@ -569,27 +569,27 @@ fn gen_return(asm: &mut Assembler, val: lir::Opnd) -> Option<()> {
}
/// Compile Fixnum + Fixnum
fn gen_fixnum_add(asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd, state: &FrameState) -> Option<lir::Opnd> {
fn gen_fixnum_add(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd, state: &FrameState) -> Option<lir::Opnd> {
// Add left + right and test for overflow
let left_untag = asm.sub(left, Opnd::Imm(1));
let out_val = asm.add(left_untag, right);
asm.jo(Target::SideExit(state.clone()));
asm.jo(side_exit(jit, state)?);
Some(out_val)
}
/// Compile Fixnum - Fixnum
fn gen_fixnum_sub(asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd, state: &FrameState) -> Option<lir::Opnd> {
fn gen_fixnum_sub(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd, state: &FrameState) -> Option<lir::Opnd> {
// Subtract left - right and test for overflow
let val_untag = asm.sub(left, right);
asm.jo(Target::SideExit(state.clone()));
asm.jo(side_exit(jit, state)?);
let out_val = asm.add(val_untag, Opnd::Imm(1));
Some(out_val)
}
/// Compile Fixnum * Fixnum
fn gen_fixnum_mult(asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd, state: &FrameState) -> Option<lir::Opnd> {
fn gen_fixnum_mult(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd, state: &FrameState) -> Option<lir::Opnd> {
// Do some bitwise gymnastics to handle tag bits
// x * y is translated to (x >> 1) * (y - 1) + 1
let left_untag = asm.rshift(left, Opnd::UImm(1));
@ -597,7 +597,7 @@ fn gen_fixnum_mult(asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd, state
let out_val = asm.mul(left_untag, right_untag);
// Test for overflow
asm.jo_mul(Target::SideExit(state.clone()));
asm.jo_mul(side_exit(jit, state)?);
let out_val = asm.add(out_val, Opnd::UImm(1));
Some(out_val)
@ -651,11 +651,11 @@ fn gen_test(asm: &mut Assembler, val: lir::Opnd) -> Option<lir::Opnd> {
}
/// Compile a type check with a side exit
fn gen_guard_type(asm: &mut Assembler, val: lir::Opnd, guard_type: Type, state: &FrameState) -> Option<lir::Opnd> {
fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard_type: Type, state: &FrameState) -> Option<lir::Opnd> {
if guard_type.is_subtype(Fixnum) {
// Check if opnd is Fixnum
asm.test(val, Opnd::UImm(RUBY_FIXNUM_FLAG as u64));
asm.jz(Target::SideExit(state.clone()));
asm.jz(side_exit(jit, state)?);
} else {
unimplemented!("unsupported type: {guard_type}");
}
@ -663,9 +663,9 @@ fn gen_guard_type(asm: &mut Assembler, val: lir::Opnd, guard_type: Type, state:
}
/// Compile an identity check with a side exit
fn gen_guard_bit_equals(asm: &mut Assembler, val: lir::Opnd, expected: VALUE, state: &FrameState) -> Option<lir::Opnd> {
fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, expected: VALUE, state: &FrameState) -> Option<lir::Opnd> {
asm.cmp(val, Opnd::UImm(expected.into()));
asm.jnz(Target::SideExit(state.clone()));
asm.jnz(side_exit(jit, state)?);
Some(val)
}
@ -731,6 +731,26 @@ fn compile_iseq(iseq: IseqPtr) -> Option<Function> {
Some(function)
}
/// Build a Target::SideExit out of a FrameState
fn side_exit(jit: &mut JITState, state: &FrameState) -> Option<Target> {
let mut stack = Vec::new();
for &insn_id in state.stack() {
stack.push(jit.get_opnd(insn_id)?);
}
let mut locals = Vec::new();
for &insn_id in state.locals() {
locals.push(jit.get_opnd(insn_id)?);
}
let target = Target::SideExit {
pc: state.pc,
stack,
locals,
};
Some(target)
}
impl Assembler {
/// Make a C call while marking the start and end positions of it
fn ccall_with_branch(&mut self, fptr: *const u8, opnds: Vec<Opnd>, branch: &Rc<Branch>) -> Opnd {
@ -744,8 +764,9 @@ impl Assembler {
move |code_ptr, _| {
start_branch.start_addr.set(Some(code_ptr));
},
move |code_ptr, _| {
move |code_ptr, cb| {
end_branch.end_addr.set(Some(code_ptr));
ZJITState::add_iseq_return_addr(code_ptr.raw_ptr(cb));
},
)
}

5
zjit/src/hir.rs
View File

@ -1719,6 +1719,11 @@ impl FrameState {
self.stack.iter()
}
/// Iterate over all local variables
pub fn locals(&self) -> Iter<InsnId> {
self.locals.iter()
}
/// Push a stack operand
fn stack_push(&mut self, opnd: InsnId) {
self.stack.push(opnd);

18
zjit/src/state.rs
View File

@ -1,3 +1,5 @@
use std::collections::HashSet;
use crate::cruby::{self, rb_bug_panic_hook, EcPtr, Qnil, VALUE};
use crate::cruby_methods;
use crate::invariants::Invariants;
@ -29,6 +31,9 @@ pub struct ZJITState {
/// Properties of core library methods
method_annotations: cruby_methods::Annotations,
/// The address of the instruction that JIT-to-JIT calls return to
iseq_return_addrs: HashSet<*const u8>,
}
/// Private singleton instance of the codegen globals
@ -82,7 +87,8 @@ impl ZJITState {
options,
invariants: Invariants::default(),
assert_compiles: false,
method_annotations: cruby_methods::init()
method_annotations: cruby_methods::init(),
iseq_return_addrs: HashSet::new(),
};
unsafe { ZJIT_STATE = Some(zjit_state); }
}
@ -126,6 +132,16 @@ impl ZJITState {
let instance = ZJITState::get_instance();
instance.assert_compiles = true;
}
/// Record an address that a JIT-to-JIT call returns to
pub fn add_iseq_return_addr(addr: *const u8) {
ZJITState::get_instance().iseq_return_addrs.insert(addr);
}
/// Returns true if a JIT-to-JIT call returns to a given address
pub fn is_iseq_return_addr(addr: *const u8) -> bool {
ZJITState::get_instance().iseq_return_addrs.contains(&addr)
}
}
/// Initialize ZJIT, given options allocated by rb_zjit_init_options()