ZJIT: Avoid compiling failed ISEQs repeatedly (#14195)

zjit/src/codegen.rs

@@ -5,7 +5,7 @@ use std::ffi::{c_int, c_void};
 use crate::asm::Label;
 use crate::backend::current::{Reg, ALLOC_REGS};
 use crate::invariants::{track_bop_assumption, track_cme_assumption, track_single_ractor_assumption, track_stable_constant_names_assumption};
-use crate::gc::{append_gc_offsets, get_or_create_iseq_payload, get_or_create_iseq_payload_ptr};
+use crate::gc::{append_gc_offsets, get_or_create_iseq_payload, get_or_create_iseq_payload_ptr, IseqStatus};
 use crate::state::ZJITState;
 use crate::stats::{counter_ptr, with_time_stat, Counter, Counter::compile_time_ns};
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
@@ -136,7 +136,7 @@ fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<CodePt
 
     // Remember the block address to reuse it later
     let payload = get_or_create_iseq_payload(iseq);
-    payload.start_ptr = Some(start_ptr);
+    payload.status = IseqStatus::Compiled(start_ptr);
     append_gc_offsets(iseq, &gc_offsets);
 
     // Return a JIT code address
@@ -213,23 +213,29 @@ fn gen_entry(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function, function_pt
 fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<(CodePtr, Vec<(Rc<Branch>, IseqPtr)>)> {
     // Return an existing pointer if it's already compiled
     let payload = get_or_create_iseq_payload(iseq);
-    if let Some(start_ptr) = payload.start_ptr {
+    match payload.status {
-        return Some((start_ptr, vec![]));
+        IseqStatus::Compiled(start_ptr) => return Some((start_ptr, vec![])),
+        IseqStatus::CantCompile => return None,
+        IseqStatus::NotCompiled => {},
     }
 
     // Convert ISEQ into High-level IR and optimize HIR
     let function = match compile_iseq(iseq) {
         Some(function) => function,
-        None => return None,
+        None => {
+            payload.status = IseqStatus::CantCompile;
+            return None;
+        }
     };
 
     // Compile the High-level IR
     let result = gen_function(cb, iseq, &function);
     if let Some((start_ptr, gc_offsets, jit)) = result {
-        payload.start_ptr = Some(start_ptr);
+        payload.status = IseqStatus::Compiled(start_ptr);
         append_gc_offsets(iseq, &gc_offsets);
         Some((start_ptr, jit.branch_iseqs))
     } else {
+        payload.status = IseqStatus::CantCompile;
         None
     }
 }
@@ -1356,27 +1362,40 @@ macro_rules! c_callable {
 pub(crate) use c_callable;
 
 c_callable! {
-    /// Generated code calls this function with the SysV calling convention.
+    /// Generated code calls this function with the SysV calling convention. See [gen_function_stub].
-    /// See [gen_function_stub].
+    /// This function is expected to be called repeatedly when ZJIT fails to compile the stub.
+    /// We should be able to compile most (if not all) function stubs by side-exiting at unsupported
+    /// instructions, so this should be used primarily for cb.has_dropped_bytes() situations.
     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
+            /// gen_push_frame() doesn't set PC and SP, so we need to set them before exit
-            let cb = ZJITState::get_code_block();
+            fn set_pc_and_sp(iseq: IseqPtr, ec: EcPtr, sp: *mut VALUE) {
-            let code_ptr = with_time_stat(compile_time_ns, || function_stub_hit_body(cb, iseq, branch_ptr));
-            let code_ptr = if let Some(code_ptr) = code_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) };
+            }
 
+            // If we already know we can't compile the ISEQ, fail early without cb.mark_all_executable().
+            // TODO: Alan thinks the payload status part of this check can happen without the VM lock, since the whole
+            // code path can be made read-only. But you still need the check as is while holding the VM lock in any case.
+            let cb = ZJITState::get_code_block();
+            let payload = get_or_create_iseq_payload(iseq);
+            if cb.has_dropped_bytes() || payload.status == IseqStatus::CantCompile {
                 // Exit to the interpreter
+                set_pc_and_sp(iseq, ec, sp);
+                return ZJITState::get_stub_exit().raw_ptr(cb);
+            }
+
+            // Otherwise, attempt to compile the ISEQ. We have to mark_all_executable() beyond this point.
+            let code_ptr = with_time_stat(compile_time_ns, || function_stub_hit_body(cb, iseq, branch_ptr));
+            let code_ptr = if let Some(code_ptr) = code_ptr {
+                code_ptr
+            } else {
+                // Exit to the interpreter
+                set_pc_and_sp(iseq, ec, sp);
                 ZJITState::get_stub_exit()
             };
-
             cb.mark_all_executable();
             code_ptr.raw_ptr(cb)
         })

zjit/src/gc.rs

@@ -7,12 +7,12 @@ use crate::stats::Counter::gc_time_ns;
 /// This is all the data ZJIT stores on an ISEQ. We mark objects in this struct on GC.
 #[derive(Debug)]
 pub struct IseqPayload {
+    /// Compilation status of the ISEQ. It has the JIT code address of the first block if Compiled.
+    pub status: IseqStatus,
+
     /// Type information of YARV instruction operands
     pub profile: IseqProfile,
 
-    /// JIT code address of the first block
-    pub start_ptr: Option<CodePtr>,
-
     /// GC offsets of the JIT code. These are the addresses of objects that need to be marked.
     pub gc_offsets: Vec<CodePtr>,
 }
@@ -20,13 +20,21 @@ pub struct IseqPayload {
 impl IseqPayload {
     fn new(iseq_size: u32) -> Self {
         Self {
+            status: IseqStatus::NotCompiled,
             profile: IseqProfile::new(iseq_size),
-            start_ptr: None,
             gc_offsets: vec![],
         }
     }
 }
 
+#[derive(Debug, PartialEq)]
+pub enum IseqStatus {
+    /// CodePtr has the JIT code address of the first block
+    Compiled(CodePtr),
+    CantCompile,
+    NotCompiled,
+}
+
 /// Get a pointer to the payload object associated with an ISEQ. Create one if none exists.
 pub fn get_or_create_iseq_payload_ptr(iseq: IseqPtr) -> *mut IseqPayload {
     type VoidPtr = *mut c_void;