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
This commit is contained in:
André Luiz Tiago Soares 2025-09-16 19:38:53 -04:00 committed by GitHub
parent c7c6bcc9c8
commit 52b22f815f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 61 additions and 53 deletions

View File

@ -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)
],

View File

@ -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<String>,
/// If Some, the next ccall should verify its leafness
leaf_ccall_stack_size: Option<usize>
}
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<Opnd> {
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<Opnd>){
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>) -> 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
}

View File

@ -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<Opnd>,
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