ZJIT: Side-exit-driven recompilation

When a compiled method's type guards fail repeatedly (100 side exits by
default), invalidate the V1 code and recompile. The recompilation resets
interpreter profiles and re-enables profiling so the method gets fresh
type data before V2 compilation.

The mechanism is bounded by a global cap of 50 ISEQs to prevent code
size inflation. The highest-traffic ISEQs reach the threshold first and
get recompiled; the rest stay at V1.

New runtime options:
  --zjit-recompile-threshold=num  Side exits to trigger recompile
                                  (default: 100, 0 to disable)

Key additions:
- rb_zjit_count_side_exit: called from exit stubs to count side exits
- trigger_recompilation: shared logic to invalidate V1 and re-profile
- compile_exit in LIR emits the exit counter when recompilation is enabled
- IseqPayload.side_exit_count tracks per-ISEQ exit count
- DeferredForReprofiling compile error variant (used by later commits)

Author: nmeans

zjit.rb

@@ -249,6 +249,7 @@ def stats_string
     print_counters([
       :compiled_iseq_count,
       :failed_iseq_count,
+      :recompile_count,
 
       :compile_time_ns,
       :compile_side_exit_time_ns,

zjit/src/backend/lir.rs

@@ -1640,6 +1640,11 @@ pub struct Assembler {
 
     /// Current instruction index, incremented for each instruction pushed
     idx: usize,
+
+    /// Payload pointer for the ISEQ being compiled (stored as usize for Send/Sync).
+    /// Points to the IseqPayload on the Rust heap (not moved by Ruby GC compaction).
+    /// Used by compile_exits to call rb_zjit_count_side_exit from generated exit code.
+    pub payload_ptr: Option<usize>,
 }
 
 impl Assembler
@@ -1655,6 +1660,7 @@ impl Assembler
             current_block_id: BlockId(0),
             num_vregs: 0,
             idx: 0,
+            payload_ptr: None,
         }
     }
 
@@ -1688,6 +1694,7 @@ impl Assembler
         // Initialize num_vregs to match the old assembler's size
         // This allows reusing VRegs from the old assembler
         asm.num_vregs = old_asm.num_vregs;
+        asm.payload_ptr = old_asm.payload_ptr;
 
         asm
     }
@@ -2632,8 +2639,18 @@ impl Assembler
 
         /// Compile the main side-exit code. This function takes only SideExit so
         /// that it can be safely deduplicated by using SideExit as a dedup key.
-        fn compile_exit(asm: &mut Assembler, exit: &SideExit) {
+        /// When recompilation is enabled, calls rb_zjit_count_side_exit(payload_ptr)
+        /// after saving VM state (ccall is safe here since state is already saved).
+        fn compile_exit(asm: &mut Assembler, exit: &SideExit, payload_ptr: Option<usize>) {
             compile_exit_save_state(asm, exit);
+            // Recompilation: count exits in the shared (deduplicated) exit stub.
+            if get_option!(recompile_threshold) > 0 {
+                if let Some(payload_ptr) = payload_ptr {
+                    use crate::codegen::rb_zjit_count_side_exit;
+                    asm_comment!(asm, "count side exit for recompilation");
+                    asm_ccall!(asm, rb_zjit_count_side_exit, Opnd::UImm(payload_ptr as u64));
+                }
+            }
             compile_exit_return(asm);
         }
 
@@ -2729,7 +2746,7 @@ impl Assembler
                     let new_exit = self.new_label("side_exit");
                     self.write_label(new_exit.clone());
                     asm_comment!(self, "Exit: {pc}");
-                    compile_exit(self, &exit);
+                    compile_exit(self, &exit, self.payload_ptr);
                     compiled_exits.insert(exit, new_exit.unwrap_label());
                     new_exit
                 };

zjit/src/codegen.rs

@@ -6,6 +6,7 @@ use std::cell::{Cell, RefCell};
 use std::rc::Rc;
 use std::ffi::{c_int, c_long, c_void};
 use std::slice;
+use std::sync::atomic::{AtomicU64, Ordering};
 
 use crate::backend::current::ALLOC_REGS;
 use crate::invariants::{
@@ -14,7 +15,7 @@ use crate::invariants::{
     track_root_box_assumption
 };
 use crate::gc::append_gc_offsets;
-use crate::payload::{get_or_create_iseq_payload, IseqCodePtrs, IseqVersion, IseqVersionRef, IseqStatus};
+use crate::payload::{get_or_create_iseq_payload, IseqCodePtrs, IseqPayload, IseqVersion, IseqVersionRef, IseqStatus};
 use crate::state::ZJITState;
 use crate::stats::{CompileError, exit_counter_for_compile_error, exit_counter_for_unhandled_hir_insn, incr_counter, incr_counter_by, send_fallback_counter, send_fallback_counter_for_method_type, send_fallback_counter_for_super_method_type, send_fallback_counter_ptr_for_opcode, send_without_block_fallback_counter_for_method_type, send_without_block_fallback_counter_for_optimized_method_type};
 use crate::stats::{counter_ptr, with_time_stat, Counter, Counter::{compile_time_ns, exit_compile_error}};
@@ -29,6 +30,52 @@ use crate::cast::IntoUsize;
 /// At the moment, we support recompiling each ISEQ only once.
 pub const MAX_ISEQ_VERSIONS: usize = 2;
 
+/// Called from side-exit stubs to count exits for recompilation.
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_count_side_exit(payload_raw: *mut std::ffi::c_void) {
+    if payload_raw.is_null() { return; }
+    let payload = unsafe { &mut *(payload_raw as *mut IseqPayload) };
+    let threshold = get_option!(recompile_threshold) as u64;
+    if threshold == 0 || payload.side_exit_count >= threshold { return; }
+    payload.side_exit_count += 1;
+    if payload.side_exit_count == threshold && payload.versions.len() < MAX_ISEQ_VERSIONS {
+        let iseq = match payload.versions.last() {
+            Some(version_ref) => unsafe { version_ref.as_ref() }.iseq,
+            None => return,
+        };
+        with_vm_lock(src_loc!(), || {
+            trigger_recompilation(payload_raw, iseq);
+        });
+    }
+}
+
+static GLOBAL_RECOMPILE_COUNT: AtomicU64 = AtomicU64::new(0);
+const MAX_GLOBAL_RECOMPILATIONS: u64 = 50;
+
+fn trigger_recompilation(payload_raw: *mut std::ffi::c_void, iseq: IseqPtr) {
+    if MAX_GLOBAL_RECOMPILATIONS > 0 {
+        let prev = GLOBAL_RECOMPILE_COUNT.fetch_add(1, Ordering::Relaxed);
+        if prev >= MAX_GLOBAL_RECOMPILATIONS {
+            GLOBAL_RECOMPILE_COUNT.fetch_sub(1, Ordering::Relaxed);
+            return;
+        }
+    }
+    let payload = unsafe { &mut *(payload_raw as *mut IseqPayload) };
+    debug!("trigger_recompilation: recompiling {}", iseq_get_location(iseq, 0));
+    incr_counter!(recompile_count);
+    payload.profile.reset_for_recompile();
+    if let Some(version) = payload.versions.last_mut() {
+        let version = unsafe { version.as_mut() };
+        version.status = IseqStatus::Invalidated;
+    }
+    unsafe { rb_iseq_reset_jit_func(iseq) };
+    unsafe { rb_zjit_profile_enable(iseq) };
+}
+
+unsafe extern "C" {
+    fn rb_zjit_profile_enable(iseq: IseqPtr);
+}
+
 /// Sentinel program counter stored in C frames when runtime checks are enabled.
 const PC_POISON: Option<*const VALUE> = if cfg!(feature = "runtime_checks") {
     Some(usize::MAX as *const VALUE)
@@ -55,18 +102,24 @@ struct JITState {
 
     /// ISEQ calls that need to be compiled later
     iseq_calls: Vec<IseqCallRef>,
+    payload_ptr: usize,
+    has_version_budget: bool,
 }
 
 impl JITState {
-    /// Create a new JITState instance
     fn new(iseq: IseqPtr, version: IseqVersionRef, num_insns: usize, num_blocks: usize) -> Self {
+        let payload_ptr = get_or_create_iseq_payload(iseq) as *const _ as usize;
+        let payload = unsafe { &*(payload_ptr as *const IseqPayload) };
+        let has_version_budget = payload.versions.len() < MAX_ISEQ_VERSIONS;
         JITState {
             iseq,
             version,
             opnds: vec![None; num_insns],
             labels: vec![None; num_blocks],
             jit_entries: Vec::default(),
             iseq_calls: Vec::default(),
+            payload_ptr,
+            has_version_budget,
         }
     }
 
@@ -152,16 +205,15 @@ pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, jit_exception: boo
         let mut code_ptr = with_time_stat(compile_time_ns, || gen_iseq_entry_point(cb, iseq, jit_exception));
 
         if let Err(err) = &code_ptr {
-            // Assert that the ISEQ compiles if RubyVM::ZJIT.assert_compiles is enabled.
-            // We assert only `jit_exception: false` cases until we support exception handlers.
-            if ZJITState::assert_compiles_enabled() && !jit_exception {
-                let iseq_location = iseq_get_location(iseq, 0);
-                panic!("Failed to compile: {iseq_location}");
-            }
-
-            // For --zjit-stats, generate an entry that just increments exit_compilation_failure and exits
-            if get_option!(stats) {
-                code_ptr = gen_compile_error_counter(cb, err);
+            // DeferredForReprofiling is not a real failure
+            if *err != CompileError::DeferredForReprofiling {
+                if ZJITState::assert_compiles_enabled() && !jit_exception {
+                    let iseq_location = iseq_get_location(iseq, 0);
+                    panic!("Failed to compile: {iseq_location}");
+                }
+                if get_option!(stats) {
+                    code_ptr = gen_compile_error_counter(cb, err);
+                }
             }
         }
 
@@ -319,6 +371,10 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, func
     let mut jit = JITState::new(iseq, version, function.num_insns(), function.num_blocks());
     let mut asm = Assembler::new_with_stack_slots(num_spilled_params);
 
+    if get_option!(recompile_threshold) > 0 && jit.payload_ptr != 0 {
+        asm.payload_ptr = Some(jit.payload_ptr);
+    }
+
     // Mapping from HIR block IDs to LIR block IDs.
     // This is is a one-to-one mapping from HIR to LIR blocks used for finding
     // jump targets in LIR (LIR should always jump to the head of an HIR block)

zjit/src/options.rs

@@ -22,6 +22,9 @@ pub type CallThreshold = u64;
 #[allow(non_upper_case_globals)]
 pub static mut rb_zjit_profile_threshold: CallThreshold = DEFAULT_CALL_THRESHOLD - DEFAULT_NUM_PROFILES as CallThreshold;
 
+/// Default --zjit-recompile-threshold. Number of side exits before triggering recompilation.
+pub const DEFAULT_RECOMPILE_THRESHOLD: u64 = 100;
+
 /// Number of calls to compile ISEQ with ZJIT at jit_compile() in vm.c.
 /// --zjit-call-threshold=1 compiles on first execution without profiling information.
 #[unsafe(no_mangle)]
@@ -96,6 +99,10 @@ pub struct Options {
 
     /// Path to a file where compiled ISEQs will be saved.
     pub log_compiled_iseqs: Option<std::path::PathBuf>,
+
+    /// Number of side exits before triggering recompilation of an ISEQ.
+    /// Set to 0 to disable recompilation.
+    pub recompile_threshold: u64,
 }
 
 impl Default for Options {
@@ -121,6 +128,7 @@ impl Default for Options {
             perf: false,
             allowed_iseqs: None,
             log_compiled_iseqs: None,
+            recompile_threshold: DEFAULT_RECOMPILE_THRESHOLD,
         }
     }
 }
@@ -144,6 +152,8 @@ pub const ZJIT_OPTIONS: &[(&str, &str)] = &[
     ("--zjit-perf",  "Dump ISEQ symbols into /tmp/perf-{}.map for Linux perf."),
     ("--zjit-log-compiled-iseqs=path",
                      "Log compiled ISEQs to the file. The file will be truncated."),
+    ("--zjit-recompile-threshold=num",
+                     "Side exits to trigger recompile (default: 100, 0=off)."),
     ("--zjit-trace-exits[=counter]",
                      "Record source on side-exit. `Counter` picks specific counter."),
     ("--zjit-trace-exits-sample-rate=num",
@@ -331,6 +341,13 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
         },
 
 
+        ("recompile-threshold", _) => match opt_val.parse() {
+            Ok(n) => {
+                options.recompile_threshold = n;
+            }
+            Err(_) => return None,
+        },
+
         ("stats-quiet", _) => {
             options.stats = true;
             options.print_stats = false;

zjit/src/payload.rs

@@ -11,13 +11,16 @@ pub struct IseqPayload {
     pub profile: IseqProfile,
     /// JIT code versions. Different versions should have different assumptions.
     pub versions: Vec<IseqVersionRef>,
+    /// Number of side exits observed for this ISEQ
+    pub side_exit_count: u64,
 }
 
 impl IseqPayload {
     fn new(iseq_size: u32) -> Self {
         Self {
             profile: IseqProfile::new(iseq_size),
             versions: vec![],
+            side_exit_count: 0,
         }
     }
 }

zjit/src/profile.rs

@@ -447,6 +447,17 @@ impl IseqProfile {
             }
         }
     }
+    /// Reset profiling state so all instructions get re-profiled from scratch.
+    /// Clears both the profiling counters AND the type distribution data.
+    pub fn reset_for_recompile(&mut self) {
+        for count in self.num_profiles.iter_mut() {
+            *count = 0;
+        }
+        for operands in self.opnd_types.iter_mut() {
+            operands.clear();
+        }
+        self.super_cme.clear();
+    }
 }
 
 #[cfg(test)]

zjit/src/stats.rs

@@ -155,6 +155,7 @@ make_counters! {
     default {
         compiled_iseq_count,
         failed_iseq_count,
+        recompile_count,
 
         compile_time_ns,
         profile_time_ns,
@@ -330,6 +331,7 @@ make_counters! {
     compile_error_exception_handler,
     compile_error_out_of_memory,
     compile_error_label_linking_failure,
+    compile_error_deferred_for_reprofiling,
     compile_error_jit_to_jit_optional,
     compile_error_register_spill_on_ccall,
     compile_error_register_spill_on_alloc,
@@ -505,6 +507,8 @@ pub enum CompileError {
     /// offsets that don't fit in one instruction. We error in
     /// error that case.
     LabelLinkingFailure,
+    /// Compilation deferred for re-profiling of cold branches.
+    DeferredForReprofiling,
 }
 
 /// Return a raw pointer to the exit counter for a given CompileError
@@ -519,6 +523,7 @@ pub fn exit_counter_for_compile_error(compile_error: &CompileError) -> Counter {
         ExceptionHandler        => compile_error_exception_handler,
         OutOfMemory             => compile_error_out_of_memory,
         LabelLinkingFailure     => compile_error_label_linking_failure,
+        DeferredForReprofiling  => compile_error_deferred_for_reprofiling,
         ParseError(parse_error) => match parse_error {
             StackUnderflow(_)       => compile_error_parse_stack_underflow,
             MalformedIseq(_)        => compile_error_parse_malformed_iseq,