ZJIT: Write a callee frame on JIT-to-JIT calls (#13579)

k0kubun · tekknolagi · web-flow · commit 7fa3e1a1db6e · 2025-06-12T17:18:50.000-07:00
Co-authored-by: Max Bernstein &lt;tekknolagi@gmail.com&gt;
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
@@ -102,12 +102,39 @@ def test(a) = 1 + a
     }, call_threshold: 2
   end
 
+  def test_opt_plus_type_guard_exit_with_locals
+    assert_compiles '[6, 6.0]', %q{
+      def test(a)
+        local = 3
+        1 + a + local
+      end
+      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{
+    assert_compiles '[4, 4.0]', %q{
       def side_exit(n) = 1 + n
       def jit_frame(n) = 1 + side_exit(n)
       def entry(n) = jit_frame(n)
+      entry(2) # profile send
+      [entry(2), entry(2.0)]
+    }, call_threshold: 2
+  end
+
+  def test_opt_plus_type_guard_nested_exit_with_locals
+    assert_compiles '[9, 9.0]', %q{
+      def side_exit(n)
+        local = 2
+        1 + n + local
+      end
+      def jit_frame(n)
+        local = 3
+        1 + side_exit(n) + local
+      end
+      def entry(n) = jit_frame(n)
+      entry(2) # profile send
       [entry(2), entry(2.0)]
     }, call_threshold: 2
   end
@@ -130,7 +157,6 @@ def test(a, b) = a * b
   end
 
   def test_opt_mult_overflow
-    omit 'side exits are not implemented yet'
     assert_compiles '[6, -6, 9671406556917033397649408, -9671406556917033397649408, 21267647932558653966460912964485513216]', %q{
       def test(a, b)
         a * b
@@ -610,6 +636,22 @@ def test() = @foo = 1
     }
   end
 
+  def test_send_backtrace
+    backtrace = [
+      "-e:2:in 'Object#jit_frame1'",
+      "-e:3:in 'Object#entry'",
+      "-e:5:in 'block in <main>'",
+      "-e:6:in '<main>'",
+    ]
+    assert_compiles backtrace.inspect, %q{
+      def jit_frame2 = caller     # 1
+      def jit_frame1 = jit_frame2 # 2
+      def entry = jit_frame1      # 3
+      entry # profile send        # 4
+      entry                       # 5
+    }, call_threshold: 2
+  end
+
   # tool/ruby_vm/views/*.erb relies on the zjit instructions a) being contiguous and
   # b) being reliably ordered after all the other instructions.
   def test_instruction_order
@@ -631,11 +673,7 @@ def assert_compiles(expected, test_script, insns: [], **opts)
     pipe_fd = 3
 
     script = <<~RUBY
-      _test_proc = -> {
-        RubyVM::ZJIT.assert_compiles
-        #{test_script}
-      }
-      ret_val = _test_proc.call
+      ret_val = (_test_proc = -> { RubyVM::ZJIT.assert_compiles; #{test_script.lstrip} }).call
       result = {
         ret_val:,
         #{ unless insns.empty?
diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
@@ -211,11 +211,6 @@ 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.
diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
@@ -1,8 +1,8 @@
 use std::collections::HashMap;
 use std::fmt;
 use std::mem::take;
-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::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::VALUE};
 use crate::backend::current::*;
 use crate::virtualmem::CodePtr;
@@ -1797,7 +1797,7 @@ impl Assembler
                 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);
+                    self.store(Opnd::mem(64, SP, (-local_size_and_idx_to_ep_offset(locals.len(), idx) - 1) * SIZEOF_VALUE_I32), opnd);
                 }
 
                 asm_comment!(self, "save cfp->pc");
@@ -1809,10 +1809,6 @@ impl Assembler
                 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()));
@@ -1823,13 +1819,6 @@ impl Assembler
         }
         Some(())
     }
-
-    #[unsafe(no_mangle)]
-    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");
-        }
-    }
 }
 
 impl fmt::Debug for Assembler {
diff --git a/zjit/src/backend/x86_64/mod.rs b/zjit/src/backend/x86_64/mod.rs
@@ -109,11 +109,6 @@ 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
 
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -258,7 +258,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::IfTrue { val, target } => return gen_if_true(jit, asm, opnd!(val), target),
         Insn::IfFalse { val, target } => return gen_if_false(jit, asm, opnd!(val), target),
         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::SendWithoutBlockDirect { cme, iseq, self_val, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(self_val), args, &function.frame_state(*state))?,
         Insn::Return { val } => return Some(gen_return(asm, opnd!(val))?),
         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))?,
@@ -484,8 +484,16 @@ fn gen_send_without_block(
     self_val: &InsnId,
     args: &Vec<InsnId>,
 ) -> Option<lir::Opnd> {
-    // Spill the receiver and the arguments onto the stack. They need to be marked by GC and may be caller-saved registers.
+    // Spill locals onto the stack.
+    // TODO: Don't spill locals eagerly; lazily reify frames
+    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)?);
+    }
+    // Spill the receiver and the arguments onto the stack.
+    // They need to be on the interpreter stack to let the interpreter access them.
     // TODO: Avoid spilling operands that have been spilled before.
+    asm_comment!(asm, "spill receiver and arguments");
     for (idx, &insn_id) in [*self_val].iter().chain(args.iter()).enumerate() {
         // Currently, we don't move the SP register. So it's equal to the base pointer.
         let stack_opnd = Opnd::mem(64, SP, idx as i32  * SIZEOF_VALUE_I32);
@@ -515,10 +523,40 @@ fn gen_send_without_block_direct(
     cb: &mut CodeBlock,
     jit: &mut JITState,
     asm: &mut Assembler,
+    cme: *const rb_callable_method_entry_t,
     iseq: IseqPtr,
     recv: Opnd,
     args: &Vec<InsnId>,
+    state: &FrameState,
 ) -> Option<lir::Opnd> {
+    // Save cfp->pc and cfp->sp for the caller frame
+    gen_save_pc(asm, state);
+    gen_save_sp(asm, state.stack().len() - args.len() - 1); // -1 for receiver
+
+    // Spill the virtual stack and the locals of the caller onto the stack
+    // TODO: Lazily materialize caller frames on side exits or when needed
+    asm_comment!(asm, "spill locals and stack");
+    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)?);
+    }
+    for (idx, &insn_id) in state.stack().enumerate() {
+        asm.mov(Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32), jit.get_opnd(insn_id)?);
+    }
+
+    // Set up the new frame
+    // TODO: Lazily materialize caller frames on side exits or when needed
+    gen_push_frame(asm, args.len(), state, ControlFrame {
+        recv,
+        iseq,
+        cme,
+        frame_type: VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL,
+    });
+
+    asm_comment!(asm, "switch to new SP register");
+    let local_size = unsafe { get_iseq_body_local_table_size(iseq) } as usize;
+    let new_sp = asm.add(SP, ((state.stack().len() + local_size - args.len() + VM_ENV_DATA_SIZE as usize) * SIZEOF_VALUE).into());
+    asm.mov(SP, new_sp);
+
     asm_comment!(asm, "switch to new CFP");
     let new_cfp = asm.sub(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
     asm.mov(CFP, new_cfp);
@@ -537,7 +575,15 @@ fn gen_send_without_block_direct(
     jit.branch_iseqs.push((branch.clone(), iseq));
     // TODO(max): Add a PatchPoint here that can side-exit the function if the callee messed with
     // the frame's locals
-    Some(asm.ccall_with_branch(dummy_ptr, c_args, &branch))
+    let ret = asm.ccall_with_branch(dummy_ptr, c_args, &branch);
+
+    // If a callee side-exits, i.e. returns Qundef, propagate the return value to the caller.
+    // The caller will side-exit the callee into the interpreter.
+    // TODO: Let side exit code pop all JIT frames to optimize away this cmp + je.
+    asm.cmp(ret, Qundef.into());
+    asm.je(ZJITState::get_exit_trampoline().into());
+
+    Some(ret)
 }
 
 /// Compile an array duplication instruction
@@ -749,6 +795,45 @@ fn gen_save_sp(asm: &mut Assembler, stack_size: usize) {
     asm.mov(cfp_sp, sp_addr);
 }
 
+/// Frame metadata written by gen_push_frame()
+struct ControlFrame {
+    recv: Opnd,
+    iseq: IseqPtr,
+    cme: *const rb_callable_method_entry_t,
+    frame_type: u32,
+}
+
+/// Compile an interpreter frame
+fn gen_push_frame(asm: &mut Assembler, argc: usize, state: &FrameState, frame: ControlFrame) {
+    // Locals are written by the callee frame on side-exits or non-leaf calls
+
+    // See vm_push_frame() for details
+    asm_comment!(asm, "push cme, specval, frame type");
+    // ep[-2]: cref of cme
+    let local_size = unsafe { get_iseq_body_local_table_size(frame.iseq) } as i32;
+    let ep_offset = state.stack().len() as i32 + local_size - argc as i32 + VM_ENV_DATA_SIZE as i32 - 1;
+    asm.store(Opnd::mem(64, SP, (ep_offset - 2) * SIZEOF_VALUE_I32), VALUE::from(frame.cme).into());
+    // ep[-1]: block_handler or prev EP
+    // block_handler is not supported for now
+    asm.store(Opnd::mem(64, SP, (ep_offset - 1) * SIZEOF_VALUE_I32), VM_BLOCK_HANDLER_NONE.into());
+    // ep[0]: ENV_FLAGS
+    asm.store(Opnd::mem(64, SP, ep_offset * SIZEOF_VALUE_I32), frame.frame_type.into());
+
+    // Write to the callee CFP
+    fn cfp_opnd(offset: i32) -> Opnd {
+        Opnd::mem(64, CFP, offset - (RUBY_SIZEOF_CONTROL_FRAME as i32))
+    }
+
+    asm_comment!(asm, "push callee control frame");
+    // cfp_opnd(RUBY_OFFSET_CFP_PC): written by the callee frame on side-exits or non-leaf calls
+    // cfp_opnd(RUBY_OFFSET_CFP_SP): written by the callee frame on side-exits or non-leaf calls
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_ISEQ), VALUE::from(frame.iseq).into());
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), frame.recv);
+    let ep = asm.lea(Opnd::mem(64, SP, ep_offset * SIZEOF_VALUE_I32));
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_EP), ep);
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_BLOCK_CODE), 0.into());
+}
+
 /// Return a register we use for the basic block argument at a given index
 fn param_reg(idx: usize) -> Reg {
     // To simplify the implementation, allocate a fixed register for each basic block argument for now.
@@ -764,10 +849,13 @@ fn param_reg(idx: usize) -> Reg {
 
 /// Inverse of ep_offset_to_local_idx(). See ep_offset_to_local_idx() for details.
 fn local_idx_to_ep_offset(iseq: IseqPtr, local_idx: usize) -> i32 {
-    let local_table_size: i32 = unsafe { get_iseq_body_local_table_size(iseq) }
-        .try_into()
-        .unwrap();
-    local_table_size - local_idx as i32 - 1 + VM_ENV_DATA_SIZE as i32
+    let local_size = unsafe { get_iseq_body_local_table_size(iseq) };
+    local_size_and_idx_to_ep_offset(local_size as usize, local_idx)
+}
+
+/// Convert the number of locals and a local index to an offset in the EP
+pub fn local_size_and_idx_to_ep_offset(local_size: usize, local_idx: usize) -> i32 {
+    local_size as i32 - local_idx as i32 - 1 + VM_ENV_DATA_SIZE as i32
 }
 
 /// Convert ISEQ into High-level IR
@@ -816,9 +904,8 @@ impl Assembler {
             move |code_ptr, _| {
                 start_branch.start_addr.set(Some(code_ptr));
             },
-            move |code_ptr, cb| {
+            move |code_ptr, _| {
                 end_branch.end_addr.set(Some(code_ptr));
-                ZJITState::add_iseq_return_addr(code_ptr.raw_ptr(cb));
             },
         )
     }
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -426,7 +426,15 @@ pub enum Insn {
     /// Ignoring keyword arguments etc for now
     SendWithoutBlock { self_val: InsnId, call_info: CallInfo, cd: *const rb_call_data, args: Vec<InsnId>, state: InsnId },
     Send { self_val: InsnId, call_info: CallInfo, cd: *const rb_call_data, blockiseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
-    SendWithoutBlockDirect { self_val: InsnId, call_info: CallInfo, cd: *const rb_call_data, iseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
+    SendWithoutBlockDirect {
+        self_val: InsnId,
+        call_info: CallInfo,
+        cd: *const rb_call_data,
+        cme: *const rb_callable_method_entry_t,
+        iseq: IseqPtr,
+        args: Vec<InsnId>,
+        state: InsnId,
+    },
 
     /// Control flow instructions
     Return { val: InsnId },
@@ -957,10 +965,11 @@ impl Function {
                 args: args.iter().map(|arg| find!(*arg)).collect(),
                 state: *state,
             },
-            SendWithoutBlockDirect { self_val, call_info, cd, iseq, args, state } => SendWithoutBlockDirect {
+            SendWithoutBlockDirect { self_val, call_info, cd, cme, iseq, args, state } => SendWithoutBlockDirect {
                 self_val: find!(*self_val),
                 call_info: call_info.clone(),
                 cd: *cd,
+                cme: *cme,
                 iseq: *iseq,
                 args: args.iter().map(|arg| find!(*arg)).collect(),
                 state: *state,
@@ -1261,7 +1270,7 @@ impl Function {
                         if let Some(expected) = guard_equal_to {
                             self_val = self.push_insn(block, Insn::GuardBitEquals { val: self_val, expected, state });
                         }
-                        let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { self_val, call_info, cd, iseq, args, state });
+                        let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { self_val, call_info, cd, cme, iseq, args, state });
                         self.make_equal_to(insn_id, send_direct);
                     }
                     Insn::GetConstantPath { ic } => {
diff --git a/zjit/src/state.rs b/zjit/src/state.rs
