ZJIT: Per-ISEQ feedback for singleton class speculation

Instead of globally disabling NoSingletonClass optimization for a class
once any singleton class has been seen, track invalidations per-ISEQ.
When an ISEQ's NoSingletonClass patchpoint is invalidated, set a flag
on its payload so only that ISEQ avoids re-speculating on recompilation.

ISEQs that were never invalidated continue to optimize with
NoSingletonClass even if other ISEQs were invalidated for the same
class. This reduces singleton_class_seen fallbacks from 9.7M (26.9% of
all fallbacks) to 1.3M (4.8%) on the lobsters benchmark.

Author: XrXr

zjit/src/hir.rs

@@ -7,7 +7,7 @@
 #![allow(clippy::match_like_matches_macro)]
 use crate::{
     backend::lir::C_ARG_OPNDS,
-    cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, invariants::{self, has_singleton_class_of}, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState, json::Json,
+    cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, invariants::{self}, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState, json::Json,
     state,
 };
 use std::{
@@ -2302,6 +2302,9 @@ fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq
 pub struct Function {
     // ISEQ this function refers to
     iseq: *const rb_iseq_t,
+    /// Whether previously, a function for this ISEQ was invalidated due to
+    /// singleton class creation (violation of NoSingletonClass invariant).
+    was_invalidated_for_singleton_class_creation: bool,
     /// The types for the parameters of this function. They are copied to the type
     /// of entry block params after infer_types() fills Empty to all insn_types.
     param_types: Vec<Type>,
@@ -2408,6 +2411,7 @@ impl Function {
     fn new(iseq: *const rb_iseq_t) -> Function {
         Function {
             iseq,
+            was_invalidated_for_singleton_class_creation: false,
             insns: vec![],
             insn_types: vec![],
             union_find: UnionFind::new().into(),
@@ -2533,9 +2537,9 @@ impl Function {
             // No patchpoint needed.
             return true;
         }
-        if has_singleton_class_of(klass) {
-            // We've seen a singleton class for this klass. Disable the optimization
-            // to avoid an invalidation loop.
+        if self.was_invalidated_for_singleton_class_creation {
+            // A previous compilation of this ISEQ was invalidated for singleton class
+            // creation. Avoid a repeating the invalidation.
             return false;
         }
         self.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoSingletonClass { klass }, state });
@@ -6700,6 +6704,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
     let payload = get_or_create_iseq_payload(iseq);
     let mut profiles = ProfileOracle::new(payload);
     let mut fun = Function::new(iseq);
+    fun.was_invalidated_for_singleton_class_creation = payload.was_invalidated_for_singleton_class_creation;
 
     // Compute a map of PC->Block by finding jump targets
     let jit_entry_insns = jit_entry_insns(iseq);

zjit/src/hir/opt_tests.rs

@@ -13227,44 +13227,82 @@ mod hir_opt_tests {
     // NoSingletonClass optimization to avoid an invalidation loop.
     #[test]
     fn test_skip_optimization_after_singleton_class_seen() {
-        // First, trigger the singleton class callback for String by creating a singleton class.
-        // This should mark String as having had a singleton class seen.
+        // First, compile a function that uses the NoSingletonClass assumption
         eval(r#"
-            "hello".singleton_class
+            def test(s, proc)
+              s.length
+              proc.call
+              s.length
+            end
+            test("hi", -> {})
+            test("hi", -> {})
         "#);
+        assert!(hir_string("test").contains("NoSingletonClass"));
 
-        // Now define and compile a method that would normally be optimized with NoSingletonClass.
-        // Since String has had a singleton class, the optimization should be skipped and we
-        // should fall back to SendWithoutBlock.
+        // Now we break the assumption by defining a singleton method on a string.
         eval(r#"
-            def test(s)
-              s.length
-            end
-            test("asdf")
+            special_string = +""
+            test(special_string, -> { def special_string.length = -1 })
         "#);
 
         // The output should NOT have NoSingletonClass patchpoint for String, and should
         // fall back to SendWithoutBlock instead of the optimized CCall path.
-        assert_snapshot!(hir_string("test"), @"
+        let hir = hir_string("test");
+        assert!(! hir.contains("NoSingletonClass"));
+        assert_snapshot!(hir, @"
         fn test@<compiled>:3:
         bb1():
           EntryPoint interpreter
           v1:BasicObject = LoadSelf
           v2:CPtr = LoadSP
           v3:BasicObject = LoadField v2, :s@0x1000
-          Jump bb3(v1, v3)
+          v4:BasicObject = LoadField v2, :proc@0x1001
+          Jump bb3(v1, v3, v4)
         bb2():
           EntryPoint JIT(0)
-          v6:BasicObject = LoadArg :self@0
-          v7:BasicObject = LoadArg :s@1
-          Jump bb3(v6, v7)
-        bb3(v9:BasicObject, v10:BasicObject):
-          v16:BasicObject = Send v10, :length # SendFallbackReason: Singleton class previously created for receiver class
+          v7:BasicObject = LoadArg :self@0
+          v8:BasicObject = LoadArg :s@1
+          v9:BasicObject = LoadArg :proc@2
+          Jump bb3(v7, v8, v9)
+        bb3(v11:BasicObject, v12:BasicObject, v13:BasicObject):
+          v19:BasicObject = Send v12, :length # SendFallbackReason: Singleton class previously created for receiver class
+          v24:BasicObject = Send v13, :call # SendFallbackReason: Singleton class previously created for receiver class
+          PatchPoint NoEPEscape(test)
+          v32:BasicObject = Send v12, :length # SendFallbackReason: Singleton class previously created for receiver class
           CheckInterrupts
-          Return v16
+          Return v32
         ");
     }
 
+    #[test]
+    fn test_no_singleton_class_busts_isolated_per_iseq() {
+        // First, compile a function that uses the NoSingletonClass assumption
+        eval(r#"
+            def will_bust(s, proc)
+              s.length
+              proc.call
+              s.length
+            end
+
+            def call_length(s) = s.length
+
+            will_bust("hi", -> {})
+            will_bust("hi", -> {})
+        "#);
+        assert!(hir_string("will_bust").contains("NoSingletonClass"));
+
+        // Now we break the assumption by defining a singleton method on a string.
+        eval(r#"
+            special_string = +""
+            will_bust(special_string, -> { def special_string.length = -1 })
+        "#);
+        assert!(! hir_string("will_bust").contains("NoSingletonClass"));
+
+        // But, the unrelated call_length() should still use NoSingletonClass
+        eval("call_length('profile')");
+        assert!(hir_string("call_length").contains("NoSingletonClass"));
+    }
+
     #[test]
     fn test_invokesuper_to_iseq_optimizes_to_direct() {
         eval("

zjit/src/invariants.rs

@@ -11,7 +11,7 @@ use crate::stats::Counter::invalidation_time_ns;
 use crate::gc::remove_gc_offsets;
 
 macro_rules! compile_patch_points {
-    ($cb:expr, $patch_points:expr, $($comment_args:tt)*) => {
+    ($cb:expr, $patch_points:expr, $cause:ident, $($comment_args:tt)*) => {
         with_time_stat(invalidation_time_ns, || {
             for patch_point in $patch_points {
                 let written_range = $cb.with_write_ptr(patch_point.patch_point_ptr, |cb| {
@@ -24,11 +24,11 @@ macro_rules! compile_patch_points {
                 // Stop marking GC offsets corrupted by the jump instruction
                 remove_gc_offsets(patch_point.version, &written_range);
 
-                // If the ISEQ doesn't have max versions, invalidate this version.
                 let mut version = patch_point.version;
                 let iseq = unsafe { version.as_ref() }.iseq;
                 if !iseq.is_null() {
                     let payload = get_or_create_iseq_payload(iseq);
+                    // If the ISEQ doesn't have max versions, invalidate this version.
                     if unsafe { version.as_ref() }.status != IseqStatus::Invalidated && payload.versions.len() < MAX_ISEQ_VERSIONS {
                         unsafe { version.as_mut() }.status = IseqStatus::Invalidated;
                         unsafe { rb_iseq_reset_jit_func(version.as_ref().iseq) };
@@ -40,12 +40,21 @@ macro_rules! compile_patch_points {
                             }
                         }
                     }
+                    // Remember NoSingletonClass busts on the payload
+                    if is_no_singleton_class!($cause) {
+                        payload.was_invalidated_for_singleton_class_creation = true;
+                    }
                 }
             }
         });
     };
 }
 
+macro_rules! is_no_singleton_class {
+    (NoSingletonClass) => { true };
+    ($_:ident) => { false };
+}
+
 /// When a PatchPoint is invalidated, it generates a jump instruction from `from` to `to`.
 #[derive(Debug, Eq, Hash, PartialEq)]
 struct PatchPoint {
@@ -196,7 +205,7 @@ pub extern "C" fn rb_zjit_bop_redefined(klass: RedefinitionFlag, bop: ruby_basic
             debug!("BOP is redefined: {}", bop);
 
             // Invalidate all patch points for this BOP
-            compile_patch_points!(cb, patch_points, "BOP is redefined: {}", bop);
+            compile_patch_points!(cb, patch_points, BOP, "BOP is redefined: {}", bop);
 
             cb.mark_all_executable();
         }
@@ -223,7 +232,7 @@ pub extern "C" fn rb_zjit_invalidate_no_ep_escape(iseq: IseqPtr) {
 
             // Invalidate the patch points for this ISEQ
             let cb = ZJITState::get_code_block();
-            compile_patch_points!(cb, patch_points, "EP is escaped: {}", iseq_name(iseq));
+            compile_patch_points!(cb, patch_points, EP, "EP is escaped: {}", iseq_name(iseq));
 
             cb.mark_all_executable();
         }
@@ -338,7 +347,7 @@ pub extern "C" fn rb_zjit_cme_invalidate(cme: *const rb_callable_method_entry_t)
             debug!("CME is invalidated: {:?}", cme);
 
             // Invalidate all patch points for this CME
-            compile_patch_points!(cb, patch_points, "CME is invalidated: {:?}", cme);
+            compile_patch_points!(cb, patch_points, CME, "CME is invalidated: {:?}", cme);
 
             cb.mark_all_executable();
         }
@@ -360,7 +369,7 @@ pub extern "C" fn rb_zjit_constant_state_changed(id: ID) {
             debug!("Constant state changed: {:?}", id);
 
             // Invalidate all patch points for this constant ID
-            compile_patch_points!(cb, patch_points, "Constant state changed: {:?}", id);
+            compile_patch_points!(cb, patch_points, Const, "Constant state changed: {:?}", id);
 
             cb.mark_all_executable();
         }
@@ -395,7 +404,7 @@ pub extern "C" fn rb_zjit_before_ractor_spawn() {
         let patch_points = mem::take(&mut ZJITState::get_invariants().single_ractor_patch_points);
 
         // Invalidate all patch points for single ractor mode
-        compile_patch_points!(cb, patch_points, "Another ractor spawned, invalidating single ractor mode assumption");
+        compile_patch_points!(cb, patch_points, Ractor, "Another ractor spawned, invalidating single ractor mode assumption");
 
         cb.mark_all_executable();
     });
@@ -439,7 +448,7 @@ pub extern "C" fn rb_zjit_tracing_invalidate_all() {
         let cb = ZJITState::get_code_block();
         let patch_points = mem::take(&mut ZJITState::get_invariants().no_trace_point_patch_points);
 
-        compile_patch_points!(cb, patch_points, "TracePoint is enabled, invalidating no TracePoint assumption");
+        compile_patch_points!(cb, patch_points, TracePoint, "TracePoint is enabled, invalidating no TracePoint assumption");
 
         cb.mark_all_executable();
     });
@@ -481,7 +490,7 @@ pub extern "C" fn rb_zjit_invalidate_root_box() {
         let patch_points = mem::take(&mut invariants.root_box_patch_points);
 
         // Invalidate all patch points for root box mode
-        compile_patch_points!(cb, patch_points, "Non-root box created, invalidating root box assumption");
+        compile_patch_points!(cb, patch_points, Box, "Non-root box created, invalidating root box assumption");
 
         cb.mark_all_executable();
     });
@@ -513,7 +522,7 @@ pub extern "C" fn rb_zjit_invalidate_no_singleton_class(klass: VALUE) {
                 if !patch_points.is_empty() {
                     let cb = ZJITState::get_code_block();
                     debug!("Singleton class created for {:?}", klass);
-                    compile_patch_points!(cb, patch_points, "Singleton class created for {:?}", klass);
+                    compile_patch_points!(cb, patch_points, NoSingletonClass, "Singleton class created for {:?}", klass);
                     cb.mark_all_executable();
                 }
             }

zjit/src/payload.rs

@@ -11,13 +11,17 @@ pub struct IseqPayload {
     pub profile: IseqProfile,
     /// JIT code versions. Different versions should have different assumptions.
     pub versions: Vec<IseqVersionRef>,
+    /// Whether a previous compilation of this ISEQ was invalidated due to
+    /// singleton class creation (violation of [`crate::hir::Invariant::NoSingletonClass`]).
+    pub was_invalidated_for_singleton_class_creation: bool,
 }
 
 impl IseqPayload {
     fn new(iseq_size: u32) -> Self {
         Self {
             profile: IseqProfile::new(iseq_size),
             versions: vec![],
+            was_invalidated_for_singleton_class_creation: false,
         }
     }
 }