From f8ee06901cfec2ebb7340087f039b103e8ab51b3 Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Thu, 20 Nov 2025 17:25:53 -0500 Subject: ZJIT: For JIT-to-JIT send, avoid loading uninitialized local through EP JIT-to-JIT sends don't blit locals to nil in the callee's EP memory region because HIR is aware of this initial state and memory ops are only done when necessary. Previously, we read from this initialized memory by emitting `GetLocal` in e.g. BBs that are immediate successor to an entrypoint. The entry points sets up the frame state properly and we also reload locals if necessary after an operation that potentially makes the environment escape. So, listen to the frame state when it's supposed to be up-to-date (`!local_inval`). --- zjit/src/hir.rs | 21 ++++++++++------ zjit/src/hir/opt_tests.rs | 62 +++++++++++++++++++++++++++++++++++++++++++---- zjit/src/hir/tests.rs | 7 +++--- 3 files changed, 73 insertions(+), 17 deletions(-) diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index d6ccf9160e..dd7bc953b9 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -5230,19 +5230,24 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result { } YARVINSN_getlocal_WC_0 => { let ep_offset = get_arg(pc, 0).as_u32(); - if ep_escaped || has_blockiseq { // TODO: figure out how to drop has_blockiseq here + if !local_inval { + // The FrameState is the source of truth for locals until invalidated. + // In case of JIT-to-JIT send locals might never end up in EP memory. + let val = state.getlocal(ep_offset); + state.stack_push(val); + } else if ep_escaped || has_blockiseq { // TODO: figure out how to drop has_blockiseq here // Read the local using EP let val = fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0, use_sp: false, rest_param: false }); state.setlocal(ep_offset, val); // remember the result to spill on side-exits state.stack_push(val); } else { - if local_inval { - // If there has been any non-leaf call since JIT entry or the last patch point, - // add a patch point to make sure locals have not been escaped. - let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state.without_locals() }); // skip spilling locals - fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: exit_id }); - local_inval = false; - } + assert!(local_inval); // if check above + // There has been some non-leaf call since JIT entry or the last patch point, + // so add a patch point to make sure locals have not been escaped. + let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state.without_locals() }); // skip spilling locals + fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: exit_id }); + local_inval = false; + // Read the local from FrameState let val = state.getlocal(ep_offset); state.stack_push(val); diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs index c99f01d088..866d0ec06d 100644 --- a/zjit/src/hir/opt_tests.rs +++ b/zjit/src/hir/opt_tests.rs @@ -780,14 +780,13 @@ mod hir_opt_tests { EntryPoint JIT(0) Jump bb2(v5, v6) bb2(v8:BasicObject, v9:BasicObject): - v13:BasicObject = GetLocal l0, EP@3 PatchPoint MethodRedefined(C@0x1000, fun_new_map@0x1008, cme:0x1010) PatchPoint NoSingletonClass(C@0x1000) - v24:ArraySubclass[class_exact:C] = GuardType v13, ArraySubclass[class_exact:C] - v25:BasicObject = CCallWithFrame C#fun_new_map@0x1038, v24, block=0x1040 - v16:BasicObject = GetLocal l0, EP@3 + v23:ArraySubclass[class_exact:C] = GuardType v9, ArraySubclass[class_exact:C] + v24:BasicObject = CCallWithFrame C#fun_new_map@0x1038, v23, block=0x1040 + v15:BasicObject = GetLocal l0, EP@3 CheckInterrupts - Return v25 + Return v24 "); } @@ -8425,4 +8424,57 @@ mod hir_opt_tests { Return v32 "); } + + #[test] + fn no_load_from_ep_right_after_entrypoint() { + let formatted = eval(" + def read_nil_local(a, _b, _c) + formatted ||= a + @formatted = formatted + -> { formatted } # the environment escapes + end + + def call + puts [], [], [], [] # fill VM stack with junk + read_nil_local(true, 1, 1) # expected direct send + end + + call # profile + call # compile + @formatted + "); + assert_eq!(Qtrue, formatted, "{}", formatted.obj_info()); + assert_snapshot!(hir_string("read_nil_local"), @r" + fn read_nil_local@:3: + bb0(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + v2:BasicObject = GetLocal l0, SP@7 + v3:BasicObject = GetLocal l0, SP@6 + v4:BasicObject = GetLocal l0, SP@5 + v5:NilClass = Const Value(nil) + Jump bb2(v1, v2, v3, v4, v5) + bb1(v8:BasicObject, v9:BasicObject, v10:BasicObject, v11:BasicObject): + EntryPoint JIT(0) + v12:NilClass = Const Value(nil) + Jump bb2(v8, v9, v10, v11, v12) + bb2(v14:BasicObject, v15:BasicObject, v16:BasicObject, v17:BasicObject, v18:NilClass): + CheckInterrupts + v27:BasicObject = GetLocal l0, EP@6 + SetLocal l0, EP@3, v27 + v39:BasicObject = GetLocal l0, EP@3 + PatchPoint SingleRactorMode + SetIvar v14, :@formatted, v39 + v45:Class[VMFrozenCore] = Const Value(VALUE(0x1000)) + PatchPoint MethodRedefined(Class@0x1008, lambda@0x1010, cme:0x1018) + PatchPoint NoSingletonClass(Class@0x1008) + v59:BasicObject = CCallWithFrame RubyVM::FrozenCore.lambda@0x1040, v45, block=0x1048 + v48:BasicObject = GetLocal l0, EP@6 + v49:BasicObject = GetLocal l0, EP@5 + v50:BasicObject = GetLocal l0, EP@4 + v51:BasicObject = GetLocal l0, EP@3 + CheckInterrupts + Return v59 + "); + } } diff --git a/zjit/src/hir/tests.rs b/zjit/src/hir/tests.rs index 76a74c75eb..f8a6abc2bc 100644 --- a/zjit/src/hir/tests.rs +++ b/zjit/src/hir/tests.rs @@ -1500,11 +1500,10 @@ pub mod hir_build_tests { EntryPoint JIT(0) Jump bb2(v5, v6) bb2(v8:BasicObject, v9:BasicObject): - v13:BasicObject = GetLocal l0, EP@3 - v15:BasicObject = Send v13, 0x1000, :each - v16:BasicObject = GetLocal l0, EP@3 + v14:BasicObject = Send v9, 0x1000, :each + v15:BasicObject = GetLocal l0, EP@3 CheckInterrupts - Return v15 + Return v14 "); } -- cgit v1.2.3