ZJIT: Stub JIT-to-JIT calls (#14052)

jit.c

@@ -442,3 +442,15 @@ rb_yarv_ary_entry_internal(VALUE ary, long offset)
 {
     return rb_ary_entry_internal(ary, offset);
 }
+
+void
+rb_set_cfp_pc(struct rb_control_frame_struct *cfp, const VALUE *pc)
+{
+    cfp->pc = pc;
+}
+
+void
+rb_set_cfp_sp(struct rb_control_frame_struct *cfp, VALUE *sp)
+{
+    cfp->sp = sp;
+}

yjit.c

@@ -499,18 +499,6 @@ rb_yjit_str_simple_append(VALUE str1, VALUE str2)
     return rb_str_cat(str1, RSTRING_PTR(str2), RSTRING_LEN(str2));
 }
 
-void
-rb_set_cfp_pc(struct rb_control_frame_struct *cfp, const VALUE *pc)
-{
-    cfp->pc = pc;
-}
-
-void
-rb_set_cfp_sp(struct rb_control_frame_struct *cfp, VALUE *sp)
-{
-    cfp->sp = sp;
-}
-
 extern VALUE *rb_vm_base_ptr(struct rb_control_frame_struct *cfp);
 
 // YJIT needs this function to never allocate and never raise

yjit/src/cruby_bindings.inc.rs

@@ -1202,8 +1202,6 @@ extern "C" {
     pub fn rb_yjit_iseq_builtin_attrs(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
     pub fn rb_yjit_builtin_function(iseq: *const rb_iseq_t) -> *const rb_builtin_function;
     pub fn rb_yjit_str_simple_append(str1: VALUE, str2: VALUE) -> VALUE;
-    pub fn rb_set_cfp_pc(cfp: *mut rb_control_frame_struct, pc: *const VALUE);
-    pub fn rb_set_cfp_sp(cfp: *mut rb_control_frame_struct, sp: *mut VALUE);
     pub fn rb_vm_base_ptr(cfp: *mut rb_control_frame_struct) -> *mut VALUE;
     pub fn rb_yarv_str_eql_internal(str1: VALUE, str2: VALUE) -> VALUE;
     pub fn rb_str_neq_internal(str1: VALUE, str2: VALUE) -> VALUE;
@@ -1330,4 +1328,6 @@ extern "C" {
     pub fn rb_IMEMO_TYPE_P(imemo: VALUE, imemo_type: imemo_type) -> ::std::os::raw::c_int;
     pub fn rb_assert_cme_handle(handle: VALUE);
     pub fn rb_yarv_ary_entry_internal(ary: VALUE, offset: ::std::os::raw::c_long) -> VALUE;
+    pub fn rb_set_cfp_pc(cfp: *mut rb_control_frame_struct, pc: *const VALUE);
+    pub fn rb_set_cfp_sp(cfp: *mut rb_control_frame_struct, sp: *mut VALUE);
 }

zjit/src/codegen.rs

@@ -1,6 +1,6 @@
 use std::cell::Cell;
 use std::rc::Rc;
-use std::ffi::{c_int};
+use std::ffi::{c_int, c_void};
 
 use crate::asm::Label;
 use crate::backend::current::{Reg, ALLOC_REGS};
@@ -105,53 +105,33 @@ fn gen_iseq_entry_point(iseq: IseqPtr) -> *const u8 {
     let code_ptr = gen_iseq_entry_point_body(cb, iseq);
 
     // Always mark the code region executable if asm.compile() has been used.
-    // We need to do this even if code_ptr is null because, whether gen_entry()
+    // We need to do this even if code_ptr is null because, whether gen_entry() or
-    // or gen_iseq() fails or not, gen_function() has already used asm.compile().
+    // gen_function_stub() fails or not, gen_function() has already used asm.compile().
     cb.mark_all_executable();
 
-    code_ptr
+    code_ptr.map_or(std::ptr::null(), |ptr| ptr.raw_ptr(cb))
 }
 
 /// Compile an entry point for a given ISEQ
-fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> *const u8 {
+fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<CodePtr> {
     // Compile ISEQ into High-level IR
-    let function = match compile_iseq(iseq) {
+    let function = compile_iseq(iseq)?;
-        Some(function) => function,
-        None => return std::ptr::null(),
-    };
 
     // Compile the High-level IR
     let Some((start_ptr, gc_offsets, jit)) = gen_function(cb, iseq, &function) else {
         debug!("Failed to compile iseq: gen_function failed: {}", iseq_get_location(iseq, 0));
-        return std::ptr::null();
+        return None;
     };
 
     // Compile an entry point to the JIT code
     let Some(entry_ptr) = gen_entry(cb, iseq, &function, start_ptr) else {
         debug!("Failed to compile iseq: gen_entry failed: {}", iseq_get_location(iseq, 0));
-        return std::ptr::null();
+        return None;
     };
 
-    let mut branch_iseqs = jit.branch_iseqs;
+    // Stub callee ISEQs for JIT-to-JIT calls
-
+    for (branch, callee_iseq) in jit.branch_iseqs.into_iter() {
-    // Recursively compile callee ISEQs
+        gen_iseq_branch(cb, callee_iseq, iseq, branch)?;
-    let caller_iseq = iseq;
-    while let Some((branch, iseq)) = branch_iseqs.pop() {
-        // Disable profiling. This will be the last use of the profiling information for the ISEQ.
-        unsafe { rb_zjit_profile_disable(iseq); }
-
-        // Compile the ISEQ
-        let Some((callee_ptr, callee_branch_iseqs)) = gen_iseq(cb, iseq) else {
-            // Failed to compile the callee. Bail out of compiling this graph of ISEQs.
-            debug!("Failed to compile iseq: could not compile callee: {} -> {}",
-                   iseq_get_location(caller_iseq, 0), iseq_get_location(iseq, 0));
-            return std::ptr::null();
-        };
-        let callee_addr = callee_ptr.raw_ptr(cb);
-        branch.regenerate(cb, |asm| {
-            asm.ccall(callee_addr, vec![]);
-        });
-        branch_iseqs.extend(callee_branch_iseqs);
     }
 
     // Remember the block address to reuse it later
@@ -160,7 +140,27 @@ fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> *const u8 {
     append_gc_offsets(iseq, &gc_offsets);
 
     // Return a JIT code address
-    entry_ptr.raw_ptr(cb)
+    Some(entry_ptr)
+}
+
+/// Stub a branch for a JIT-to-JIT call
+fn gen_iseq_branch(cb: &mut CodeBlock, iseq: IseqPtr, caller_iseq: IseqPtr, branch: Rc<Branch>) -> Option<()> {
+    // Compile a function stub
+    let Some((stub_ptr, gc_offsets)) = gen_function_stub(cb, iseq, branch.clone()) else {
+        // Failed to compile the stub. Bail out of compiling the caller ISEQ.
+        debug!("Failed to compile iseq: could not compile stub: {} -> {}",
+               iseq_get_location(caller_iseq, 0), iseq_get_location(iseq, 0));
+        return None;
+    };
+    append_gc_offsets(iseq, &gc_offsets);
+
+    // Update the JIT-to-JIT call to call the stub
+    let stub_addr = stub_ptr.raw_ptr(cb);
+    branch.regenerate(cb, |asm| {
+        asm_comment!(asm, "call function stub: {}", iseq_get_location(iseq, 0));
+        asm.ccall(stub_addr, vec![]);
+    });
+    Some(())
 }
 
 /// Write an entry to the perf map in /tmp
@@ -1263,6 +1263,127 @@ fn max_num_params(function: &Function) -> usize {
     }).max().unwrap_or(0)
 }
 
+#[cfg(target_arch = "x86_64")]
+macro_rules! c_callable {
+    ($(#[$outer:meta])*
+    fn $f:ident $args:tt $(-> $ret:ty)? $body:block) => {
+        $(#[$outer])*
+        extern "sysv64" fn $f $args $(-> $ret)? $body
+    };
+}
+#[cfg(target_arch = "aarch64")]
+macro_rules! c_callable {
+    ($(#[$outer:meta])*
+    fn $f:ident $args:tt $(-> $ret:ty)? $body:block) => {
+        $(#[$outer])*
+        extern "C" fn $f $args $(-> $ret)? $body
+    };
+}
+pub(crate) use c_callable;
+
+c_callable! {
+    /// Generated code calls this function with the SysV calling convention.
+    /// See [gen_function_stub].
+    fn function_stub_hit(iseq: IseqPtr, branch_ptr: *const c_void, ec: EcPtr, sp: *mut VALUE) -> *const u8 {
+        with_vm_lock(src_loc!(), || {
+            // Get a pointer to compiled code or the side-exit trampoline
+            let cb = ZJITState::get_code_block();
+            let code_ptr = if let Some(code_ptr) = function_stub_hit_body(cb, iseq, branch_ptr) {
+                code_ptr
+            } else {
+                // gen_push_frame() doesn't set PC and SP, so we need to set them for side-exit
+                // TODO: We could generate code that sets PC/SP. Note that we'd still need to handle OOM.
+                let cfp = unsafe { get_ec_cfp(ec) };
+                let pc = unsafe { rb_iseq_pc_at_idx(iseq, 0) }; // TODO: handle opt_pc once supported
+                unsafe { rb_set_cfp_pc(cfp, pc) };
+                unsafe { rb_set_cfp_sp(cfp, sp) };
+
+                // Exit to the interpreter
+                ZJITState::get_stub_exit()
+            };
+
+            cb.mark_all_executable();
+            code_ptr.raw_ptr(cb)
+        })
+    }
+}
+
+/// Compile an ISEQ for a function stub
+fn function_stub_hit_body(cb: &mut CodeBlock, iseq: IseqPtr, branch_ptr: *const c_void) -> Option<CodePtr> {
+    // Compile the stubbed ISEQ
+    let Some((code_ptr, branch_iseqs)) = gen_iseq(cb, iseq) else {
+        debug!("Failed to compile iseq: gen_iseq failed: {}", iseq_get_location(iseq, 0));
+        return None;
+    };
+
+    // Stub callee ISEQs for JIT-to-JIT calls
+    for (branch, callee_iseq) in branch_iseqs.into_iter() {
+        gen_iseq_branch(cb, callee_iseq, iseq, branch)?;
+    }
+
+    // Update the stub to call the code pointer
+    let branch = unsafe { Rc::from_raw(branch_ptr as *const Branch) };
+    let code_addr = code_ptr.raw_ptr(cb);
+    branch.regenerate(cb, |asm| {
+        asm_comment!(asm, "call compiled function: {}", iseq_get_location(iseq, 0));
+        asm.ccall(code_addr, vec![]);
+    });
+
+    Some(code_ptr)
+}
+
+/// Compile a stub for an ISEQ called by SendWithoutBlockDirect
+/// TODO: Consider creating a trampoline to share some of the code among function stubs
+fn gen_function_stub(cb: &mut CodeBlock, iseq: IseqPtr, branch: Rc<Branch>) -> Option<(CodePtr, Vec<CodePtr>)> {
+    let mut asm = Assembler::new();
+    asm_comment!(asm, "Stub: {}", iseq_get_location(iseq, 0));
+
+    // Maintain alignment for x86_64, and set up a frame for arm64 properly
+    asm.frame_setup(&[], 0);
+
+    asm_comment!(asm, "preserve argument registers");
+    for &reg in ALLOC_REGS.iter() {
+        asm.cpush(Opnd::Reg(reg));
+    }
+    const { assert!(ALLOC_REGS.len() % 2 == 0, "x86_64 would need to push one more if we push an odd number of regs"); }
+
+    // Compile the stubbed ISEQ
+    let branch_addr = Rc::into_raw(branch);
+    let jump_addr = asm_ccall!(asm, function_stub_hit,
+        Opnd::Value(iseq.into()),
+        Opnd::const_ptr(branch_addr as *const u8),
+        EC,
+        SP
+    );
+    asm.mov(Opnd::Reg(Assembler::SCRATCH_REG), jump_addr);
+
+    asm_comment!(asm, "restore argument registers");
+    for &reg in ALLOC_REGS.iter().rev() {
+        asm.cpop_into(Opnd::Reg(reg));
+    }
+
+    // Discard the current frame since the JIT function will set it up again
+    asm.frame_teardown(&[]);
+
+    // Jump to SCRATCH_REG so that cpop_all() doesn't clobber it
+    asm.jmp_opnd(Opnd::Reg(Assembler::SCRATCH_REG));
+    asm.compile(cb)
+}
+
+/// Generate a trampoline that is used when a function stub fails to compile the ISEQ
+pub fn gen_stub_exit(cb: &mut CodeBlock) -> Option<CodePtr> {
+    let mut asm = Assembler::new();
+
+    asm_comment!(asm, "exit from function stub");
+    asm.frame_teardown(lir::JIT_PRESERVED_REGS);
+    asm.cret(Qundef.into());
+
+    asm.compile(cb).map(|(code_ptr, gc_offsets)| {
+        assert_eq!(gc_offsets.len(), 0);
+        code_ptr
+    })
+}
+
 impl Assembler {
     /// Make a C call while marking the start and end positions of it
     fn ccall_with_branch(&mut self, fptr: *const u8, opnds: Vec<Opnd>, branch: &Rc<Branch>) -> Opnd {

zjit/src/cruby_bindings.inc.rs

@@ -1015,4 +1015,6 @@ unsafe extern "C" {
     pub fn rb_IMEMO_TYPE_P(imemo: VALUE, imemo_type: imemo_type) -> ::std::os::raw::c_int;
     pub fn rb_assert_cme_handle(handle: VALUE);
     pub fn rb_yarv_ary_entry_internal(ary: VALUE, offset: ::std::os::raw::c_long) -> VALUE;
+    pub fn rb_set_cfp_pc(cfp: *mut rb_control_frame_struct, pc: *const VALUE);
+    pub fn rb_set_cfp_sp(cfp: *mut rb_control_frame_struct, sp: *mut VALUE);
 }

zjit/src/state.rs

@@ -1,9 +1,11 @@
+use crate::codegen::gen_stub_exit;
 use crate::cruby::{self, rb_bug_panic_hook, rb_vm_insns_count, EcPtr, Qnil, VALUE};
 use crate::cruby_methods;
 use crate::invariants::Invariants;
 use crate::asm::CodeBlock;
 use crate::options::get_option;
 use crate::stats::Counters;
+use crate::virtualmem::CodePtr;
 
 #[allow(non_upper_case_globals)]
 #[unsafe(no_mangle)]
@@ -30,6 +32,9 @@ pub struct ZJITState {
 
     /// Properties of core library methods
     method_annotations: cruby_methods::Annotations,
+
+    /// Side-exit trampoline used when it fails to compile the ISEQ for a function stub
+    stub_exit: CodePtr,
 }
 
 /// Private singleton instance of the codegen globals
@@ -39,7 +44,7 @@ impl ZJITState {
     /// Initialize the ZJIT globals
     pub fn init() {
         #[cfg(not(test))]
-        let cb = {
+        let mut cb = {
             use crate::cruby::*;
             use crate::options::*;
 
@@ -76,7 +81,9 @@ impl ZJITState {
             CodeBlock::new(mem_block.clone(), get_option!(dump_disasm))
         };
         #[cfg(test)]
-        let cb = CodeBlock::new_dummy();
+        let mut cb = CodeBlock::new_dummy();
+
+        let stub_exit = gen_stub_exit(&mut cb).unwrap();
 
         // Initialize the codegen globals instance
         let zjit_state = ZJITState {
@@ -85,6 +92,7 @@ impl ZJITState {
             invariants: Invariants::default(),
             assert_compiles: false,
             method_annotations: cruby_methods::init(),
+            stub_exit,
         };
         unsafe { ZJIT_STATE = Some(zjit_state); }
     }
@@ -160,6 +168,11 @@ impl ZJITState {
             true // If no restrictions, allow all ISEQs
         }
     }
+
+    /// Return a code pointer to the side-exit trampoline for function stubs
+    pub fn get_stub_exit() -> CodePtr {
+        ZJITState::get_instance().stub_exit
+    }
 }
 
 /// Initialize ZJIT