diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs index 72068c1e5b1775..4aa0208252ec18 100644 --- a/zjit/src/codegen.rs +++ b/zjit/src/codegen.rs @@ -52,9 +52,6 @@ const JIT_RETURN_POISON: Option = if cfg!(feature = "runtime_checks") { /// Ephemeral code generation state struct JITState { - /// Instruction sequence for the method being compiled - iseq: IseqPtr, - /// ISEQ version that is being compiled, which will be used by PatchPoint version: IseqVersionRef, @@ -73,9 +70,8 @@ struct JITState { impl JITState { /// Create a new JITState instance - fn new(iseq: IseqPtr, version: IseqVersionRef, num_insns: usize, num_blocks: usize) -> Self { + fn new(version: IseqVersionRef, num_insns: usize, num_blocks: usize) -> Self { JITState { - iseq, version, opnds: vec![None; num_insns], labels: vec![None; num_blocks], @@ -379,7 +375,7 @@ fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, mut version: IseqVersionRef, fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, function: &Function) -> Result<(IseqCodePtrs, Vec, Vec), CompileError> { let (mut jit, asm) = trace_compile_phase("codegen", || { let num_spilled_params = max_num_params(function).saturating_sub(ALLOC_REGS.len()); - let mut jit = JITState::new(iseq, version, function.num_insns(), function.num_blocks()); + let mut jit = JITState::new(version, function.num_insns(), function.num_blocks()); let mut asm = Assembler::new_with_stack_slots(num_spilled_params); // Mapping from HIR block IDs to LIR block IDs. @@ -706,7 +702,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio &Insn::FixnumMod { left, right, state } => gen_fixnum_mod(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state)), &Insn::FixnumAref { recv, index } => gen_fixnum_aref(asm, opnd!(recv), opnd!(index)), Insn::IsNil { val } => gen_isnil(asm, opnd!(val)), - &Insn::IsMethodCfunc { val, cd, cfunc, state: _ } => gen_is_method_cfunc(jit, asm, opnd!(val), cd, cfunc), + &Insn::IsMethodCfunc { val, cd, cfunc, state } => gen_is_method_cfunc(asm, opnd!(val), cd, cfunc, &function.frame_state(state)), &Insn::IsBitEqual { left, right } => gen_is_bit_equal(asm, opnd!(left), opnd!(right)), &Insn::IsBitNotEqual { left, right } => gen_is_bit_not_equal(asm, opnd!(left), opnd!(right)), &Insn::BoxBool { val } => gen_box_bool(asm, opnd!(val)), @@ -733,7 +729,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio Insn::CCallVariadic { cfunc, recv, name, args, cme, state, block, return_type: _, elidable: _ } => { gen_ccall_variadic(jit, asm, *cfunc, *name, opnd!(recv), opnds!(args), *cme, *block, &function.frame_state(*state)) } - Insn::GetIvar { self_val, id, ic, state: _ } => gen_getivar(jit, asm, opnd!(self_val), *id, *ic), + Insn::GetIvar { self_val, id, ic, state } => gen_getivar(asm, opnd!(self_val), *id, *ic, &function.frame_state(*state)), Insn::SetGlobal { id, val, state } => no_output!(gen_setglobal(jit, asm, *id, opnd!(val), &function.frame_state(*state))), Insn::GetGlobal { id, state } => gen_getglobal(jit, asm, *id, &function.frame_state(*state)), &Insn::IsBlockParamModified { flags } => gen_is_block_param_modified(asm, opnd!(flags)), @@ -748,7 +744,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio Insn::SideExit { state, reason, recompile } => no_output!(gen_side_exit(jit, asm, reason, *recompile, &function.frame_state(*state))), Insn::PutSpecialObject { value_type } => gen_putspecialobject(asm, *value_type), Insn::AnyToString { val, str, state } => gen_anytostring(asm, opnd!(val), opnd!(str), &function.frame_state(*state)), - Insn::Defined { op_type, obj, pushval, v, state } => gen_defined(jit, asm, *op_type, *obj, *pushval, opnd!(v), &function.frame_state(*state)), + Insn::Defined { op_type, obj, pushval, v, lep_level, state } => gen_defined(jit, asm, *op_type, *obj, *pushval, opnd!(v), *lep_level, &function.frame_state(*state)), Insn::CheckMatch { target, pattern, flag, state } => gen_checkmatch(jit, asm, opnd!(target), opnd!(pattern), *flag, &function.frame_state(*state)), Insn::GetSpecialSymbol { symbol_type, state: _ } => gen_getspecial_symbol(asm, *symbol_type), Insn::GetSpecialNumber { nth, state } => gen_getspecial_number(asm, *nth, &function.frame_state(*state)), @@ -794,13 +790,6 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio Ok(()) } -/// Gets the EP of the ISeq of the containing method, or "local level". -/// Equivalent of GET_LEP() macro. -fn gen_get_lep(jit: &JITState, asm: &mut Assembler) -> Opnd { - let level = get_lvar_level(jit.iseq); - gen_get_ep(asm, level) -} - // Get EP at `level` from CFP fn gen_get_ep(asm: &mut Assembler, level: u32) -> Opnd { // Load environment pointer EP from CFP into a register @@ -835,18 +824,12 @@ fn gen_objtostring(jit: &mut JITState, asm: &mut Assembler, val: Opnd, cd: *cons ret } -fn gen_defined(jit: &JITState, asm: &mut Assembler, op_type: usize, obj: VALUE, pushval: VALUE, tested_value: Opnd, state: &FrameState) -> Opnd { +fn gen_defined(jit: &JITState, asm: &mut Assembler, op_type: usize, obj: VALUE, pushval: VALUE, tested_value: Opnd, lep_level: u32, state: &FrameState) -> Opnd { match op_type as defined_type { DEFINED_YIELD => { - // `yield` goes to the block handler stowed in the "local" iseq which is - // the current iseq or a parent. Only the "method" iseq type can be passed a - // block handler. (e.g. `yield` in the top level script is a syntax error.) - // - // Similar to gen_is_block_given - let local_iseq = unsafe { rb_get_iseq_body_local_iseq(jit.iseq) }; - assert_eq!(unsafe { rb_get_iseq_body_type(local_iseq) }, ISEQ_TYPE_METHOD, - "defined?(yield) in non-method iseq should be handled by HIR construction"); - let lep = gen_get_lep(jit, asm); + // `lep_level` was precomputed at HIR construction so we can materialize the local EP + // inline without walking the parent iseq chain here. + let lep = gen_get_ep(asm, lep_level); let block_handler = asm.load(Opnd::mem(64, lep, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL)); let pushval = asm.load(pushval.into()); asm.cmp(block_handler, VM_BLOCK_HANDLER_NONE.into()); @@ -1210,11 +1193,11 @@ fn gen_ccall_variadic( } /// Emit an uncached instance variable lookup -fn gen_getivar(jit: &mut JITState, asm: &mut Assembler, recv: Opnd, id: ID, ic: *const iseq_inline_iv_cache_entry) -> Opnd { +fn gen_getivar(asm: &mut Assembler, recv: Opnd, id: ID, ic: *const iseq_inline_iv_cache_entry, state: &FrameState) -> Opnd { if ic.is_null() { asm_ccall!(asm, rb_ivar_get, recv, id.0.into()) } else { - let iseq = Opnd::Value(jit.iseq.into()); + let iseq = Opnd::Value(state.iseq.into()); asm_ccall!(asm, rb_vm_getinstancevariable, iseq, recv, id.0.into(), Opnd::const_ptr(ic)) } } @@ -1226,19 +1209,19 @@ fn gen_setivar(jit: &mut JITState, asm: &mut Assembler, recv: Opnd, id: ID, ic: if ic.is_null() { asm_ccall!(asm, rb_ivar_set, recv, id.0.into(), val); } else { - let iseq = Opnd::Value(jit.iseq.into()); + let iseq = Opnd::Value(state.iseq.into()); asm_ccall!(asm, rb_vm_setinstancevariable, iseq, recv, id.0.into(), val, Opnd::const_ptr(ic)); } } fn gen_getclassvar(jit: &mut JITState, asm: &mut Assembler, id: ID, ic: *const iseq_inline_cvar_cache_entry, state: &FrameState) -> Opnd { gen_prepare_non_leaf_call(jit, asm, state); - asm_ccall!(asm, rb_vm_getclassvariable, VALUE::from(jit.iseq).into(), CFP, id.0.into(), Opnd::const_ptr(ic)) + asm_ccall!(asm, rb_vm_getclassvariable, VALUE::from(state.iseq).into(), CFP, id.0.into(), Opnd::const_ptr(ic)) } fn gen_setclassvar(jit: &mut JITState, asm: &mut Assembler, id: ID, val: Opnd, ic: *const iseq_inline_cvar_cache_entry, state: &FrameState) { gen_prepare_non_leaf_call(jit, asm, state); - asm_ccall!(asm, rb_vm_setclassvariable, VALUE::from(jit.iseq).into(), CFP, id.0.into(), val, Opnd::const_ptr(ic)); + asm_ccall!(asm, rb_vm_setclassvariable, VALUE::from(state.iseq).into(), CFP, id.0.into(), val, Opnd::const_ptr(ic)); } /// Look up global variables @@ -2435,11 +2418,11 @@ fn gen_isnil(asm: &mut Assembler, val: lir::Opnd) -> lir::Opnd { asm.csel_e(Opnd::Imm(1), Opnd::Imm(0)) } -fn gen_is_method_cfunc(jit: &JITState, asm: &mut Assembler, val: lir::Opnd, cd: *const rb_call_data, cfunc: *const u8) -> lir::Opnd { +fn gen_is_method_cfunc(asm: &mut Assembler, val: lir::Opnd, cd: *const rb_call_data, cfunc: *const u8, state: &FrameState) -> lir::Opnd { unsafe extern "C" { fn rb_vm_method_cfunc_is(iseq: IseqPtr, cd: *const rb_call_data, recv: VALUE, cfunc: *const u8) -> VALUE; } - asm_ccall!(asm, rb_vm_method_cfunc_is, VALUE::from(jit.iseq).into(), Opnd::const_ptr(cd), val, Opnd::const_ptr(cfunc)) + asm_ccall!(asm, rb_vm_method_cfunc_is, VALUE::from(state.iseq).into(), Opnd::const_ptr(cd), val, Opnd::const_ptr(cfunc)) } fn gen_is_bit_equal(asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd) -> lir::Opnd { @@ -2867,7 +2850,7 @@ fn gen_spill_locals(jit: &JITState, asm: &mut Assembler, state: &FrameState) { gen_incr_counter(asm, Counter::vm_write_locals_count); asm_comment!(asm, "spill locals"); for (idx, &insn_id) in state.locals().enumerate() { - asm.mov(Opnd::mem(64, SP, (-local_idx_to_ep_offset(jit.iseq, idx) - 1) * SIZEOF_VALUE_I32), jit.get_opnd(insn_id)); + asm.mov(Opnd::mem(64, SP, (-local_idx_to_ep_offset(state.iseq, idx) - 1) * SIZEOF_VALUE_I32), jit.get_opnd(insn_id)); } } @@ -3041,7 +3024,7 @@ fn side_exit(jit: &JITState, state: &FrameState, reason: SideExitReason) -> Targ fn side_exit_with_recompile(jit: &JITState, state: &FrameState, reason: SideExitReason, recompile: Option) -> Target { let mut exit = build_side_exit(jit, state); exit.recompile = recompile.map(|strategy| SideExitRecompile { - iseq: Opnd::Value(VALUE::from(jit.iseq)), + iseq: Opnd::Value(VALUE::from(state.iseq)), insn_idx: state.insn_idx() as u32, strategy, }); @@ -3064,7 +3047,7 @@ fn build_side_exit(jit: &JITState, state: &FrameState) -> SideExit { pc: Opnd::const_ptr(state.pc), stack, locals, - iseq: jit.iseq, + iseq: state.iseq, recompile: None, } } diff --git a/zjit/src/codegen_tests.rs b/zjit/src/codegen_tests.rs index b2cc401694f083..ac5f80358d10f4 100644 --- a/zjit/src/codegen_tests.rs +++ b/zjit/src/codegen_tests.rs @@ -23,7 +23,6 @@ fn test_breakpoint_hir_codegen() { let breakpoint = function.push_insn(function.entries_block, Insn::BreakPoint); let mut jit = JITState::new( - iseq, IseqVersion::new(iseq), function.num_insns(), function.num_blocks(), diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs index 40df1fab030a6f..6728d884f93a79 100644 --- a/zjit/src/cruby.rs +++ b/zjit/src/cruby.rs @@ -261,6 +261,9 @@ pub struct ID(pub ::std::os::raw::c_ulong); /// Pointer to an ISEQ pub type IseqPtr = *const rb_iseq_t; +/// Index of a YARV instruction within an ISEQ's bytecode array. +pub type YarvInsnIdx = usize; + #[derive(Clone, Copy, PartialEq, Eq, Debug)] pub struct ShapeId(pub u32); diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index 671249f7010cca..5314b07c20e815 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -883,7 +883,9 @@ pub enum Insn { UnboxFixnum { val: InsnId }, FixnumAref { recv: InsnId, index: InsnId }, // TODO(max): In iseq body types that are not ISEQ_TYPE_METHOD, rewrite to Constant false. - Defined { op_type: usize, obj: VALUE, pushval: VALUE, v: InsnId, state: InsnId }, + // `lep_level` is the lexical distance from this insn's iseq up to its local_iseq, used only + // for the DEFINED_YIELD op_type to materialize the local EP inline. Zero for other op_types. + Defined { op_type: usize, obj: VALUE, pushval: VALUE, v: InsnId, lep_level: u32, state: InsnId }, GetConstant { klass: InsnId, id: ID, allow_nil: InsnId, state: InsnId }, GetConstantPath { ic: *const iseq_inline_constant_cache, state: InsnId }, /// Kernel#block_given? but without pushing a frame. Similar to [`Insn::Defined`] with @@ -2982,7 +2984,7 @@ impl Function { cfunc, recv: find!(recv), args: find_vec!(args), cme, name, state, return_type, elidable, block }, &CheckMatch { target, pattern, flag, state } => CheckMatch { target: find!(target), pattern: find!(pattern), flag, state: find!(state) }, - &Defined { op_type, obj, pushval, v, state } => Defined { op_type, obj, pushval, v: find!(v), state: find!(state) }, + &Defined { op_type, obj, pushval, v, lep_level, state } => Defined { op_type, obj, pushval, v: find!(v), lep_level, state: find!(state) }, &DefinedIvar { self_val, pushval, id, state } => DefinedIvar { self_val: find!(self_val), pushval, id, state }, &GetConstant { klass, id, allow_nil, state } => GetConstant { klass: find!(klass), id, allow_nil: find!(allow_nil), state }, &NewArray { ref elements, state } => NewArray { elements: find_vec!(elements), state: find!(state) }, @@ -3323,7 +3325,7 @@ impl Function { /// Return the profiled type of the HIR instruction at the given ISEQ instruction /// index, if it is known to be monomorphic or skewed polymorphic. This historical type /// record is not a guarantee and must be checked with a GuardType or similar. - fn profiled_type_of_at(&self, insn: InsnId, iseq_insn_idx: usize) -> Option { + fn profiled_type_of_at(&self, insn: InsnId, iseq_insn_idx: YarvInsnIdx) -> Option { match self.resolve_receiver_type_from_profile(insn, iseq_insn_idx) { ReceiverTypeResolution::Monomorphic { profiled_type } | ReceiverTypeResolution::SkewedPolymorphic { profiled_type } => Some(profiled_type), @@ -3492,14 +3494,14 @@ impl Function { /// Returns: /// - `StaticallyKnown` if the receiver's exact class is known at compile-time /// - Result of [`Self::resolve_receiver_type_from_profile`] if we need to check profile data - fn resolve_receiver_type(&self, recv: InsnId, recv_type: Type, insn_idx: usize) -> ReceiverTypeResolution { + fn resolve_receiver_type(&self, recv: InsnId, recv_type: Type, insn_idx: YarvInsnIdx) -> ReceiverTypeResolution { if let Some(class) = recv_type.runtime_exact_ruby_class() { return ReceiverTypeResolution::StaticallyKnown { class }; } self.resolve_receiver_type_from_profile(recv, insn_idx) } - fn polymorphic_summary(&self, profiles: &ProfileOracle, recv: InsnId, insn_idx: usize) -> Option { + fn polymorphic_summary(&self, profiles: &ProfileOracle, recv: InsnId, insn_idx: YarvInsnIdx) -> Option { let Some(entries) = profiles.types.get(&insn_idx) else { return None; }; @@ -3523,7 +3525,7 @@ impl Function { /// - `Megamorphic`/`SkewedMegamorphic` if the receiver has too many types to optimize /// (SkewedMegamorphic may be optimized in the future, but for now we don't) /// - `NoProfile` if we have no type information - fn resolve_receiver_type_from_profile(&self, recv: InsnId, insn_idx: usize) -> ReceiverTypeResolution { + fn resolve_receiver_type_from_profile(&self, recv: InsnId, insn_idx: YarvInsnIdx) -> ReceiverTypeResolution { let Some(profiles) = self.profiles.as_ref() else { return ReceiverTypeResolution::NoProfile; }; @@ -6412,7 +6414,7 @@ impl<'a> std::fmt::Display for FunctionPrinter<'a> { #[derive(Debug, Clone, PartialEq)] pub struct FrameState { pub iseq: IseqPtr, - insn_idx: usize, + insn_idx: YarvInsnIdx, // Ruby bytecode instruction pointer pub pc: *const VALUE, @@ -6422,7 +6424,7 @@ pub struct FrameState { impl FrameState { /// Get the YARV instruction index for the current instruction - pub fn insn_idx(&self) -> usize { + pub fn insn_idx(&self) -> YarvInsnIdx { self.insn_idx } @@ -6702,7 +6704,7 @@ struct ProfileOracle { /// instruction index. At a given ISEQ instruction, the interpreter has profiled the stack /// operands to a given ISEQ instruction, and this list of pairs of (InsnId, Type) map that /// profiling information into HIR instructions. - types: HashMap>, + types: HashMap>, } impl ProfileOracle { @@ -7174,7 +7176,18 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result { // Similar to gen_is_block_given Insn::Const { val: Const::Value(Qnil) } } else { - Insn::Defined { op_type, obj, pushval, v, state: exit_id } + // For DEFINED_YIELD, codegen materializes the local EP inline (similar to + // gen_is_block_given) to check for a block handler. Precompute the lexical + // distance from this iseq up to local_iseq so codegen does not have to + // walk the parent chain. Any DEFINED_YIELD reaching this branch has a + // method local_iseq by construction — the above branch has already + // diverted the non-method case to Qnil. + let lep_level = if op_type == DEFINED_YIELD as usize { + get_lvar_level(iseq) + } else { + 0 + }; + Insn::Defined { op_type, obj, pushval, v, lep_level, state: exit_id } }; state.stack_push(fun.push_insn(block, insn)); } diff --git a/zjit/src/profile.rs b/zjit/src/profile.rs index ae496745bc41be..d8bcd50d96c031 100644 --- a/zjit/src/profile.rs +++ b/zjit/src/profile.rs @@ -13,7 +13,7 @@ use crate::stats::with_time_stat; struct Profiler { cfp: CfpPtr, iseq: IseqPtr, - insn_idx: usize, + insn_idx: YarvInsnIdx, } impl Profiler { @@ -424,7 +424,7 @@ pub struct IseqProfile { entries: Vec, /// Method entries for `super` calls (stored as VALUE to be GC-safe) - super_cme: HashMap + super_cme: HashMap } impl IseqProfile { @@ -436,7 +436,7 @@ impl IseqProfile { } /// Get or create a mutable profile entry for the given instruction index. - fn entry_mut(&mut self, insn_idx: usize) -> &mut ProfileEntry { + fn entry_mut(&mut self, insn_idx: YarvInsnIdx) -> &mut ProfileEntry { let idx = insn_idx as u32; match self.entries.binary_search_by_key(&idx, |e| e.insn_idx) { Ok(i) => &mut self.entries[i], @@ -452,14 +452,14 @@ impl IseqProfile { } /// Get a profile entry for the given instruction index (read-only). - fn entry(&self, insn_idx: usize) -> Option<&ProfileEntry> { + fn entry(&self, insn_idx: YarvInsnIdx) -> Option<&ProfileEntry> { let idx = insn_idx as u32; self.entries.binary_search_by_key(&idx, |e| e.insn_idx) .ok().map(|i| &self.entries[i]) } /// Check if enough profiles have been gathered for this instruction. - pub fn done_profiling_at(&self, insn_idx: usize) -> bool { + pub fn done_profiling_at(&self, insn_idx: YarvInsnIdx) -> bool { self.entry(insn_idx).map_or(false, |e| e.profiles_remaining == 0) } @@ -467,7 +467,7 @@ impl IseqProfile { /// `sp` is the current stack pointer (after the args and receiver). /// `argc` is the number of arguments (not counting receiver). /// Returns true if enough profiles have been gathered and the ISEQ should be recompiled. - pub fn profile_send_at(&mut self, iseq: IseqPtr, insn_idx: usize, sp: *const VALUE, argc: usize) -> bool { + pub fn profile_send_at(&mut self, iseq: IseqPtr, insn_idx: YarvInsnIdx, sp: *const VALUE, argc: usize) -> bool { let n = argc + 1; // args + receiver let entry = self.entry_mut(insn_idx); if entry.opnd_types.is_empty() { @@ -487,7 +487,7 @@ impl IseqProfile { /// This may be called on an instruction that was already profiled by YARV, /// so we reset the counter to re-profile with the new shapes seen at runtime. /// Returns true if enough profiles have been gathered and the ISEQ should be recompiled. - pub fn profile_self_at(&mut self, iseq: IseqPtr, insn_idx: usize, self_val: VALUE) -> bool { + pub fn profile_self_at(&mut self, iseq: IseqPtr, insn_idx: YarvInsnIdx, self_val: VALUE) -> bool { let entry = self.entry_mut(insn_idx); // Reset profiling if the previous round already finished (stale YARV profiles). // This ensures we collect num_profiles samples of the new shapes before recompiling. @@ -505,11 +505,11 @@ impl IseqProfile { } /// Get profiled operand types for a given instruction index - pub fn get_operand_types(&self, insn_idx: usize) -> Option<&[TypeDistribution]> { + pub fn get_operand_types(&self, insn_idx: YarvInsnIdx) -> Option<&[TypeDistribution]> { self.entry(insn_idx).map(|e| e.opnd_types.as_slice()).filter(|s| !s.is_empty()) } - pub fn get_super_method_entry(&self, insn_idx: usize) -> Option<*const rb_callable_method_entry_t> { + pub fn get_super_method_entry(&self, insn_idx: YarvInsnIdx) -> Option<*const rb_callable_method_entry_t> { let Some(entry) = self.super_cme.get(&insn_idx) else { return None }; let summary = TypeDistributionSummary::new(entry);