From 52b22f815fc11c3f52ef7e045c93c7bfebfb31d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9=20Luiz=20Tiago=20Soares?= Date: Tue, 16 Sep 2025 19:38:53 -0400 Subject: [PATCH] ZJIT: Debug mechanism to verify leafness for `gen_prepare_call_with_gc` (#14553) * functional debug mechanism; small refactor * make all tests pass * clean up implementation of debug mechanism for gen_prepare_call_with_gc * revert unnecessary change to gen_object_alloc * make ObjectAlloc non-leaf * fix merge error; reintroduce accidentally deleted counter * remove outdated comment * changes as per review comments * make set_stack_canary more stable * add todo comment to specialize object_alloc * revert whitespace changes * create gen_prepare_leaf_call_with_gc helper * fix spacing --- configure.ac | 2 +- zjit/src/backend/lir.rs | 30 ++++++++++++++- zjit/src/codegen.rs | 82 ++++++++++++++++------------------------- 3 files changed, 61 insertions(+), 53 deletions(-) diff --git a/configure.ac b/configure.ac index da5afc0018..b035e1688a 100644 --- a/configure.ac +++ b/configure.ac @@ -3974,7 +3974,7 @@ AS_CASE(["${ZJIT_SUPPORT}"], [yes], [ ], [dev], [ - rb_cargo_features="$rb_cargo_features,disasm" + rb_cargo_features="$rb_cargo_features,disasm,runtime_checks" JIT_CARGO_SUPPORT=dev AC_DEFINE(RUBY_DEBUG, 1) ], diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs index 5f96d0718a..21adc42cd1 100644 --- a/zjit/src/backend/lir.rs +++ b/zjit/src/backend/lir.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use std::fmt; use std::mem::take; use crate::codegen::local_size_and_idx_to_ep_offset; -use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32}; +use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32, vm_stack_canary}; use crate::hir::SideExitReason; use crate::options::{debug, get_option}; use crate::cruby::VALUE; @@ -1171,6 +1171,9 @@ pub struct Assembler { /// Names of labels pub(super) label_names: Vec, + + /// If Some, the next ccall should verify its leafness + leaf_ccall_stack_size: Option } impl Assembler @@ -1190,9 +1193,32 @@ impl Assembler insns: Vec::with_capacity(ASSEMBLER_INSNS_CAPACITY), live_ranges, label_names, + leaf_ccall_stack_size: None, } } + pub fn expect_leaf_ccall(&mut self, stack_size: usize) { + self.leaf_ccall_stack_size = Some(stack_size); + } + + fn set_stack_canary(&mut self) -> Option { + if cfg!(feature = "runtime_checks") { + if let Some(stack_size) = self.leaf_ccall_stack_size.take() { + let canary_addr = self.lea(Opnd::mem(64, SP, (stack_size as i32) * SIZEOF_VALUE_I32)); + let canary_opnd = Opnd::mem(64, canary_addr, 0); + self.mov(canary_opnd, vm_stack_canary().into()); + return Some(canary_opnd) + } + } + None + } + + fn clear_stack_canary(&mut self, canary_opnd: Option){ + if let Some(canary_opnd) = canary_opnd { + self.store(canary_opnd, 0.into()); + }; + } + /// Build an Opnd::VReg and initialize its LiveRange pub(super) fn new_vreg(&mut self, num_bits: u8) -> Opnd { let vreg = Opnd::VReg { idx: self.live_ranges.len(), num_bits }; @@ -1657,8 +1683,10 @@ impl Assembler { /// Call a C function without PosMarkers pub fn ccall(&mut self, fptr: *const u8, opnds: Vec) -> Opnd { + let canary_opnd = self.set_stack_canary(); let out = self.new_vreg(Opnd::match_num_bits(&opnds)); self.push_insn(Insn::CCall { fptr, opnds, start_marker: None, end_marker: None, out }); + self.clear_stack_canary(canary_opnd); out } diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs index 1f71ad9db3..12bd6d4aa3 100644 --- a/zjit/src/codegen.rs +++ b/zjit/src/codegen.rs @@ -350,7 +350,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio Insn::NewRange { low, high, flag, state } => gen_new_range(jit, asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)), Insn::NewRangeFixnum { low, high, flag, state } => gen_new_range_fixnum(asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)), Insn::ArrayDup { val, state } => gen_array_dup(asm, opnd!(val), &function.frame_state(*state)), - Insn::ObjectAlloc { val, state } => gen_object_alloc(asm, opnd!(val), &function.frame_state(*state)), + Insn::ObjectAlloc { val, state } => gen_object_alloc(jit, asm, opnd!(val), &function.frame_state(*state)), Insn::StringCopy { val, chilled, state } => gen_string_copy(asm, opnd!(val), *chilled, &function.frame_state(*state)), // concatstrings shouldn't have 0 strings // If it happens we abort the compilation for now @@ -660,7 +660,7 @@ fn gen_getglobal(jit: &mut JITState, asm: &mut Assembler, id: ID, state: &FrameS /// Intern a string fn gen_intern(asm: &mut Assembler, val: Opnd, state: &FrameState) -> Opnd { - gen_prepare_call_with_gc(asm, state); + gen_prepare_leaf_call_with_gc(asm, state); asm_ccall!(asm, rb_str_intern, val) } @@ -713,7 +713,7 @@ fn gen_getspecial_number(asm: &mut Assembler, nth: u64, state: &FrameState) -> O let backref = asm_ccall!(asm, rb_backref_get,); - gen_prepare_call_with_gc(asm, state); + gen_prepare_leaf_call_with_gc(asm, state); asm_ccall!(asm, rb_reg_nth_match, Opnd::Imm((nth >> 1).try_into().unwrap()), backref) } @@ -730,12 +730,12 @@ fn gen_check_interrupts(jit: &mut JITState, asm: &mut Assembler, state: &FrameSt } fn gen_hash_dup(asm: &mut Assembler, val: Opnd, state: &FrameState) -> lir::Opnd { - gen_prepare_call_with_gc(asm, state); + gen_prepare_leaf_call_with_gc(asm, state); asm_ccall!(asm, rb_hash_resurrect, val) } fn gen_array_push(asm: &mut Assembler, array: Opnd, val: Opnd, state: &FrameState) { - gen_prepare_call_with_gc(asm, state); + gen_prepare_leaf_call_with_gc(asm, state); asm_ccall!(asm, rb_ary_push, array, val); } @@ -983,14 +983,7 @@ fn gen_send( gen_incr_counter(asm, Counter::dynamic_send_count); gen_incr_counter(asm, Counter::dynamic_send_type_send); - // Save PC and SP - gen_prepare_call_with_gc(asm, state); - gen_save_sp(asm, state.stack().len()); - - // Spill locals and stack - gen_spill_locals(jit, asm, state); - gen_spill_stack(jit, asm, state); - + gen_prepare_non_leaf_call(jit, asm, state); asm_comment!(asm, "call #{} with dynamic dispatch", ruby_call_method_name(cd)); unsafe extern "C" { fn rb_vm_send(ec: EcPtr, cfp: CfpPtr, cd: VALUE, blockiseq: IseqPtr) -> VALUE; @@ -1010,20 +1003,7 @@ fn gen_send_without_block( ) -> lir::Opnd { gen_incr_counter(asm, Counter::dynamic_send_count); gen_incr_counter(asm, Counter::dynamic_send_type_send_without_block); - - // Note that it's incorrect to use this frame state to side exit because - // the state might not be on the boundary of an interpreter instruction. - // For example, `opt_str_uminus` pushes to the stack and then sends. - asm_comment!(asm, "spill frame state"); - - // Save PC and SP - gen_prepare_call_with_gc(asm, state); - gen_save_sp(asm, state.stack().len()); - - // Spill locals and stack - gen_spill_locals(jit, asm, state); - gen_spill_stack(jit, asm, state); - + gen_prepare_non_leaf_call(jit, asm, state); asm_comment!(asm, "call #{} with dynamic dispatch", ruby_call_method_name(cd)); unsafe extern "C" { fn rb_vm_opt_send_without_block(ec: EcPtr, cfp: CfpPtr, cd: VALUE) -> VALUE; @@ -1059,7 +1039,8 @@ fn gen_send_without_block_direct( asm.jbe(side_exit(jit, state, StackOverflow)); // Save cfp->pc and cfp->sp for the caller frame - gen_prepare_call_with_gc(asm, state); + gen_prepare_call_with_gc(asm, state, false); + // Special SP math. Can't use gen_prepare_non_leaf_call gen_save_sp(asm, state.stack().len() - args.len() - 1); // -1 for receiver gen_spill_locals(jit, asm, state); @@ -1119,11 +1100,7 @@ fn gen_invokeblock( gen_incr_counter(asm, Counter::dynamic_send_count); gen_incr_counter(asm, Counter::dynamic_send_type_invokeblock); - // Save PC and SP, spill locals and stack - gen_prepare_call_with_gc(asm, state); - gen_save_sp(asm, state.stack().len()); - gen_spill_locals(jit, asm, state); - gen_spill_stack(jit, asm, state); + gen_prepare_non_leaf_call(jit, asm, state); asm_comment!(asm, "call invokeblock"); unsafe extern "C" { @@ -1146,14 +1123,7 @@ fn gen_invokesuper( gen_incr_counter(asm, Counter::dynamic_send_count); gen_incr_counter(asm, Counter::dynamic_send_type_invokesuper); - // Save PC and SP - gen_prepare_call_with_gc(asm, state); - gen_save_sp(asm, state.stack().len()); - - // Spill locals and stack - gen_spill_locals(jit, asm, state); - gen_spill_stack(jit, asm, state); - + gen_prepare_non_leaf_call(jit, asm, state); asm_comment!(asm, "call super with dynamic dispatch"); unsafe extern "C" { fn rb_vm_invokesuper(ec: EcPtr, cfp: CfpPtr, cd: VALUE, blockiseq: IseqPtr) -> VALUE; @@ -1167,7 +1137,7 @@ fn gen_invokesuper( /// Compile a string resurrection fn gen_string_copy(asm: &mut Assembler, recv: Opnd, chilled: bool, state: &FrameState) -> Opnd { // TODO: split rb_ec_str_resurrect into separate functions - gen_prepare_call_with_gc(asm, state); + gen_prepare_leaf_call_with_gc(asm, state); let chilled = if chilled { Opnd::Imm(1) } else { Opnd::Imm(0) }; asm_ccall!(asm, rb_ec_str_resurrect, EC, recv, chilled) } @@ -1178,7 +1148,7 @@ fn gen_array_dup( val: lir::Opnd, state: &FrameState, ) -> lir::Opnd { - gen_prepare_call_with_gc(asm, state); + gen_prepare_leaf_call_with_gc(asm, state); asm_ccall!(asm, rb_ary_resurrect, val) } @@ -1189,7 +1159,7 @@ fn gen_new_array( elements: Vec, state: &FrameState, ) -> lir::Opnd { - gen_prepare_call_with_gc(asm, state); + gen_prepare_leaf_call_with_gc(asm, state); let length: c_long = elements.len().try_into().expect("Unable to fit length of elements into c_long"); @@ -1256,12 +1226,14 @@ fn gen_new_range_fixnum( flag: RangeType, state: &FrameState, ) -> lir::Opnd { - gen_prepare_call_with_gc(asm, state); + gen_prepare_leaf_call_with_gc(asm, state); asm_ccall!(asm, rb_range_new, low, high, (flag as i64).into()) } -fn gen_object_alloc(asm: &mut Assembler, val: lir::Opnd, state: &FrameState) -> lir::Opnd { - gen_prepare_call_with_gc(asm, state); +fn gen_object_alloc(jit: &JITState, asm: &mut Assembler, val: lir::Opnd, state: &FrameState) -> lir::Opnd { + // TODO: this is leaf in the vast majority of cases, + // Should specialize to avoid `gen_prepare_non_leaf_call` (Shopify#747) + gen_prepare_non_leaf_call(jit, asm, state); asm_ccall!(asm, rb_obj_alloc, val) } @@ -1375,7 +1347,7 @@ fn gen_is_method_cfunc(jit: &JITState, asm: &mut Assembler, val: lir::Opnd, cd: } fn gen_anytostring(asm: &mut Assembler, val: lir::Opnd, str: lir::Opnd, state: &FrameState) -> lir::Opnd { - gen_prepare_call_with_gc(asm, state); + gen_prepare_leaf_call_with_gc(asm, state); asm_ccall!(asm, rb_obj_as_string_result, str, val) } @@ -1532,12 +1504,20 @@ fn gen_incr_counter(asm: &mut Assembler, counter: Counter) { /// /// Unlike YJIT, we don't need to save the stack slots to protect them from GC /// because the backend spills all live registers onto the C stack on CCall. -fn gen_prepare_call_with_gc(asm: &mut Assembler, state: &FrameState) { +fn gen_prepare_call_with_gc(asm: &mut Assembler, state: &FrameState, leaf: bool) { let opcode: usize = state.get_opcode().try_into().unwrap(); let next_pc: *const VALUE = unsafe { state.pc.offset(insn_len(opcode) as isize) }; asm_comment!(asm, "save PC to CFP"); asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_PC), Opnd::const_ptr(next_pc)); + + if leaf { + asm.expect_leaf_ccall(state.stack_size()); + } +} + +fn gen_prepare_leaf_call_with_gc(asm: &mut Assembler, state: &FrameState) { + gen_prepare_call_with_gc(asm, state, true); } /// Save the current SP on the CFP @@ -1572,11 +1552,11 @@ fn gen_spill_stack(jit: &JITState, asm: &mut Assembler, state: &FrameState) { } /// Prepare for calling a C function that may call an arbitrary method. -/// Use gen_prepare_call_with_gc() if the method is leaf but allocates objects. +/// Use gen_prepare_leaf_call_with_gc() if the method is leaf but allocates objects. fn gen_prepare_non_leaf_call(jit: &JITState, asm: &mut Assembler, state: &FrameState) { // TODO: Lazily materialize caller frames when needed // Save PC for backtraces and allocation tracing - gen_prepare_call_with_gc(asm, state); + gen_prepare_call_with_gc(asm, state, false); // Save SP and spill the virtual stack in case it raises an exception // and the interpreter uses the stack for handling the exception