ZJIT: Weigh an ivar respecialization by the traffic it would take off the fallback

A frozen ivar dispatch samples what its fallback path is handed and earns an extra
compiled version once a 64-sample window names a shape the dispatch has no arm for.
One missing shape is weak evidence. At a site that is genuinely polymorphic the
fallback sees a long tail of receivers, so a window nearly always contains something
new, and the arm the recompile adds covers a few percent of the traffic while the
version and the code for the whole dispatch are paid in full. On railsbench that is
what happens: the mechanism fires 119 times and dynamic_getivar_count does not move.

Weigh the vote by the share of the window the shapes a recompile would specialize
account for *together*. A per-shape threshold is the wrong shape of test -- it was
tried at 25% and cost rubyboy its win, because its windows are split across two or
three shapes that are individually small and collectively everything. The share is
read off the bucket indices: ProfileEntry now remembers how many shapes the running
dispatch was compiled from, and `Distribution::observe_stable` records a shape without
counting or reordering, so any index at or past that mark is a shape the running code
has no arm for. Shapes the distribution has no room left for, and too-complex shapes
the dispatch would filter back out, are counted as samples but not as fixable ones:
they are exactly the evidence that a recompile will not help. The first of those was
also a plain bug -- a full distribution silently dropped the shape on the floor and
granted a version for an arm the recompile could not add.

Declining a window has a cost the old code never paid, because it used to spend its
two-version budget almost immediately and the next compile dropped the sampling call.
A site that keeps declining keeps sampling, so railsbench closed 31,407 windows --
2M calls on an otherwise exit-free path. Give each compiled version four windows to
spend on evidence that does not pan out, checked with one load off the version the
call already has in hand, and leave the sampling out of anything the ISEQ compiles
after that.

railsbench: ivar_respecialize_count 119 -> 2 with dynamic_getivar_count unchanged
(3.739M), so the removed respecializations bought nothing. Against the same tree they
take compile_time from 5,394ms to 3,945ms (-27%), code_region_bytes from 13.24MB to
12.53MB (-5%), and compiled_iseq_count from 3,473 to 3,311. Declined windows
31,407 -> 222 across 50 versions that gave up. rubyboy is unchanged in every respect:
the same 10 respecializations, the same 23,281,169 dynamic_getivar_count, the same
4,567,040 code_region_bytes, and one declined window in the whole program.

Author: k0kubun

zjit.rb

@@ -147,6 +147,8 @@ def stats_string
       :compiled_side_exit_count,
       :failed_iseq_count,
       :ivar_respecialize_count,
+      :ivar_respecialize_declined_count,
+      :ivar_respecialize_giveup_count,
 
       :compile_time_ns,
       :compile_side_exit_time_ns,

zjit/src/distribution.rs

@@ -44,6 +44,29 @@ impl<T: Copy + PartialEq + Default, const N: usize> Distribution<T, N> {
         }
     }
 
+    /// Look `item` up and, if `matches` finds no bucket holding it and a bucket is free, record
+    /// it there. Unlike [`Self::observe`] this does not count the sighting and does not reorder
+    /// the buckets, so the index it reports stays valid for as long as the distribution is only
+    /// updated this way: a caller that remembers how many buckets were occupied at some earlier
+    /// point can tell from the index alone whether an item was already known then.
+    ///
+    /// New buckets always start at a count of 1, which cannot exceed `counts[0]`, so this
+    /// preserves "`buckets[0]` is the most common item" without bubbling.
+    pub fn observe_stable(&mut self, item: T, matches: impl Fn(T, T) -> bool) -> StableBucket {
+        for (index, (bucket, count)) in self.buckets.iter_mut().zip(self.counts.iter_mut()).enumerate() {
+            if *count == 0 {
+                *bucket = item;
+                *count = 1;
+                return StableBucket::Inserted(index);
+            }
+            if matches(*bucket, item) {
+                return StableBucket::Existing(index);
+            }
+        }
+        self.other = self.other.saturating_add(1);
+        StableBucket::Full
+    }
+
     pub fn each_item(&self) -> impl Iterator<Item = T> + '_ {
         self.buckets.iter().zip(self.counts.iter())
             .filter_map(|(&bucket, &count)| if count > 0 { Some(bucket) } else { None })
@@ -55,6 +78,17 @@ impl<T: Copy + PartialEq + Default, const N: usize> Distribution<T, N> {
     }
 }
 
+/// Where an item sits in a [`Distribution`] after [`Distribution::observe_stable`].
+#[derive(PartialEq, Eq, Debug, Clone, Copy)]
+pub enum StableBucket {
+    /// The item already had this bucket.
+    Existing(usize),
+    /// The item was recorded in this bucket, which was empty until now.
+    Inserted(usize),
+    /// Every bucket belongs to some other item; this one was only counted as `other`.
+    Full,
+}
+
 #[derive(PartialEq, Debug, Clone, Copy)]
 enum DistributionKind {
     /// No types seen
@@ -200,6 +234,54 @@ mod distribution_tests {
         assert_eq!(dist.other, 1);
     }
 
+    #[test]
+    fn observe_stable_reports_new_and_known_buckets() {
+        let mut dist = Distribution::<usize, 4>::new();
+        dist.observe(10);
+        // A bucket that already exists is reported without being counted again.
+        assert_eq!(dist.observe_stable(10, |a, b| a == b), StableBucket::Existing(0));
+        assert_eq!(dist.counts[0], 1);
+        // New items land past the buckets that were occupied, which is how a caller tells
+        // them apart from the ones the compiled code already knows about.
+        assert_eq!(dist.observe_stable(11, |a, b| a == b), StableBucket::Inserted(1));
+        assert_eq!(dist.observe_stable(11, |a, b| a == b), StableBucket::Existing(1));
+        assert_eq!(dist.observe_stable(12, |a, b| a == b), StableBucket::Inserted(2));
+    }
+
+    #[test]
+    fn observe_stable_does_not_reorder_buckets() {
+        let mut dist = Distribution::<usize, 4>::new();
+        dist.observe(10);
+        for _ in 0..10 {
+            dist.observe_stable(11, |a, b| a == b);
+        }
+        // `observe` would have bubbled 11 to the front; a stable observation may not, or the
+        // index it reported earlier would now name a different item.
+        assert_eq!(dist.buckets[0], 10);
+        assert_eq!(dist.buckets[1], 11);
+        // Counts stay ordered, so DistributionSummary's invariant still holds.
+        assert!(dist.counts[0] >= dist.counts[1]);
+    }
+
+    #[test]
+    fn observe_stable_reports_full_without_recording() {
+        let mut dist = Distribution::<usize, 2>::new();
+        assert_eq!(dist.observe_stable(10, |a, b| a == b), StableBucket::Inserted(0));
+        assert_eq!(dist.observe_stable(11, |a, b| a == b), StableBucket::Inserted(1));
+        assert_eq!(dist.observe_stable(12, |a, b| a == b), StableBucket::Full);
+        assert_eq!(dist.other, 1);
+        assert_eq!(dist.buckets, [10, 11]);
+    }
+
+    #[test]
+    fn observe_stable_matches_with_the_given_predicate() {
+        let mut dist = Distribution::<(usize, usize), 4>::new();
+        dist.observe_stable((1, 100), |a, b| a.0 == b.0);
+        // Equal under the predicate, different as a whole: still the same bucket.
+        assert_eq!(dist.observe_stable((1, 999), |a, b| a.0 == b.0), StableBucket::Existing(0));
+        assert_eq!(dist.buckets[0], (1, 100));
+    }
+
     #[test]
     fn empty_distribution_returns_empty_summary() {
         let dist = Distribution::<usize, 4>::new();

zjit/src/hir.rs

@@ -8453,7 +8453,9 @@ impl Function {
             return;
         }
         let payload = get_or_create_iseq_payload(self.iseq);
-        if payload.ivar_respecializations >= crate::payload::MAX_IVAR_RESPECIALIZATIONS {
+        if payload.ivar_respecializations >= crate::payload::MAX_IVAR_RESPECIALIZATIONS
+            || payload.ivar_reprofile_giveup
+        {
             return;
         }
         self.push_insn(block, Insn::IvarReprofile { self_val: self_param, state });

zjit/src/payload.rs

@@ -38,6 +38,11 @@ pub struct IseqPayload {
     /// [`crate::profile::rb_zjit_ivar_reprofile`] grants these, and only against evidence
     /// from the fallback path. Capped at [`MAX_IVAR_RESPECIALIZATIONS`].
     pub ivar_respecializations: u8,
+    /// Whether an ivar fallback in this ISEQ has spent a compiled version's worth of
+    /// re-profiling windows without earning a recompile. Sampling costs a non-leaf call on the
+    /// fallback path, so once the evidence says a recompile would not help, later compiles of
+    /// this ISEQ leave the sampling out.
+    pub ivar_reprofile_giveup: bool,
 }
 
 /// The interpreter state observed at one exception-handler entry.
@@ -73,6 +78,7 @@ impl IseqPayload {
             self_is_heap_object: false,
             num_exits_until_invalidate: get_option!(num_exits_until_invalidate),
             ivar_respecializations: 0,
+            ivar_reprofile_giveup: false,
         }
     }
 
@@ -120,8 +126,20 @@ pub struct IseqVersion {
 
     /// JIT-to-JIT calls to the ISEQ. The IseqPayload's ISEQ is the callee of it.
     pub incoming: Vec<IseqCallRef>,
+
+    /// Re-profiling windows this version's ivar fallback paths may still close without earning a
+    /// recompile. See [`crate::profile::rb_zjit_ivar_reprofile`]: sampling is a C call on a path
+    /// that is otherwise exit-free, so a version whose fallbacks keep failing to make the case
+    /// for a recompile stops paying for the evidence.
+    pub ivar_reprofile_windows: u8,
 }
 
+/// How many windows an ivar fallback may close without earning a recompile before the version
+/// stops sampling. A fallback that has handed the same unspecializable mix of shapes to this
+/// many windows in a row is not about to change its mind, and every sample after that is a call
+/// on a hot path buying nothing.
+pub const MAX_IVAR_REPROFILE_WINDOWS: u8 = 4;
+
 /// We use a raw pointer instead of Rc to save space for refcount
 pub type IseqVersionRef = NonNull<IseqVersion>;
 
@@ -139,6 +157,7 @@ impl IseqVersion {
             gc_offsets: vec![],
             outgoing: vec![],
             incoming: vec![],
+            ivar_reprofile_windows: MAX_IVAR_REPROFILE_WINDOWS,
         };
         let version_ptr = Box::into_raw(Box::new(version));
         NonNull::new(version_ptr).expect("no null from Box")

zjit/src/profile.rs

@@ -5,7 +5,7 @@
 
 use std::collections::HashMap;
 use crate::{cruby::*, payload::get_or_create_iseq_payload, options::{get_option, NumProfiles}};
-use crate::distribution::{Distribution, DistributionSummary};
+use crate::distribution::{Distribution, DistributionSummary, StableBucket};
 use crate::stats::Counter::profile_time_ns;
 use crate::stats::with_time_stat;
 
@@ -216,19 +216,32 @@ fn profile_self(profiler: &mut Profiler, profile: &mut IseqProfile) {
 /// unusual receiver does not spend a version.
 const IVAR_REPROFILE_WINDOW: u32 = 64;
 
+/// Share of a window that the shapes a recompile would add an arm for must account for,
+/// in percent, before the recompile is worth a version.
+///
+/// A single missing shape is not evidence that a recompile helps: at a site that is genuinely
+/// polymorphic the fallback sees a long tail of shapes, and an arm for one of them removes a
+/// few percent of the calls while costing a version and the code for every other arm. What
+/// matters is the share the arms a recompile would add cover *together*, which is why this is
+/// not a per-shape threshold -- a window split evenly between three missing shapes is a good
+/// bet, and a window with one missing shape in it is not.
+const IVAR_REPROFILE_MIN_SHARE_PERCENT: u32 = 50;
+
 /// Result of [`IseqProfile::observe_ivar_fallback`].
 #[derive(PartialEq, Eq, Debug)]
 enum IvarReprofiled {
-    /// Recorded. Nothing to do.
+    /// Recorded, mid-window. Nothing to do.
     Sampled,
-    /// The window closed on a shape the compiled dispatch does not have an arm for.
-    /// The shape is now in the instruction's profile; recompile to pick it up.
+    /// The window closed without making the case for a recompile.
+    Declined,
+    /// The window closed with most of its samples on shapes the compiled dispatch has no arm
+    /// for. They are now in the instruction's profile; recompile to pick them up.
     Recompile,
 }
 
 impl IseqProfile {
     /// Record the shape of a receiver that reached an ivar site's fallback path, and report
-    /// whether a recompile would now give that receiver an arm of its own.
+    /// whether a recompile would now take most of that traffic off the fallback.
     ///
     /// The window is kept in the instruction's own profile entry, so a site that falls back
     /// rarely costs one distribution slot and nothing else.
@@ -237,27 +250,58 @@ impl IseqProfile {
         if entry.opnd_types.is_empty() {
             entry.opnd_types.resize(1, TypeDistribution::new());
         }
+        // The dispatch that is running was built from the shapes the profile held when the ISEQ
+        // was last compiled. Remember how many that was the first time we sample: because this
+        // path only ever grows the distribution through `observe_stable`, every bucket at or
+        // past that index is a shape folded in since, which the running code has no arm for.
+        let dispatch_shapes = match entry.ivar_dispatch_shapes {
+            Some(shapes) => shapes,
+            None => {
+                let shapes = entry.opnd_types[0].each_item().count() as u8;
+                entry.ivar_dispatch_shapes = Some(shapes);
+                shapes
+            }
+        };
         let ty = ProfiledType::new(recv);
-        if !entry.opnd_types[0].each_item().any(|seen| seen.shape() == ty.shape()) {
-            // A shape the profile the compiled dispatch was built from has never seen. Fold it
-            // in so a recompile can give it an arm, and remember that this window found one.
-            // The distribution keeps the class alive for as long as the profile does.
-            VALUE::from(iseq).write_barrier(ty.class());
-            entry.opnd_types[0].observe(ty);
-            entry.ivar_fallback_new_shape = true;
-        }
+        let fixable = if ty.shape().is_complex() {
+            // Too-complex shapes keep their ivars in a hash table, so the dispatch would filter
+            // this one back out. Recording it would only cost a bucket a fixable shape could use.
+            false
+        } else {
+            // De-duplicate by shape, the way the dispatch itself does.
+            let bucket = entry.opnd_types[0].observe_stable(ty, |seen, ty| seen.shape() == ty.shape());
+            if let StableBucket::Inserted(_) = bucket {
+                // The distribution keeps the class alive for as long as the profile does.
+                VALUE::from(iseq).write_barrier(ty.class());
+            }
+            match bucket {
+                StableBucket::Existing(index) | StableBucket::Inserted(index) => index >= dispatch_shapes as usize,
+                // No bucket left to record this shape in, so no recompile can specialize it.
+                // It argues against spending a version on this site, not for one.
+                StableBucket::Full => false,
+            }
+        };
         entry.ivar_fallback_samples = entry.ivar_fallback_samples.saturating_add(1);
+        if fixable {
+            entry.ivar_fallback_fixable = entry.ivar_fallback_fixable.saturating_add(1);
+        }
         if entry.ivar_fallback_samples < IVAR_REPROFILE_WINDOW {
             return IvarReprofiled::Sampled;
         }
         entry.ivar_fallback_samples = 0;
-        if std::mem::take(&mut entry.ivar_fallback_new_shape) {
+        let fixable = std::mem::take(&mut entry.ivar_fallback_fixable);
+        if fixable * 100 >= IVAR_REPROFILE_WINDOW * IVAR_REPROFILE_MIN_SHARE_PERCENT {
+            // The recompile rebuilds the dispatch from the whole profile, so start the next
+            // window measuring against all of it.
+            entry.ivar_dispatch_shapes = None;
             IvarReprofiled::Recompile
         } else {
-            // Everything this window saw is already in the profile, so the fallback is being
-            // taken for a reason a recompile cannot fix (a too-complex shape, or an arm the
-            // dispatch dropped). Keep sampling in case that changes, but do not spend a version.
-            IvarReprofiled::Sampled
+            // Most of what the fallback handles is something a recompile cannot take away: a
+            // shape already in the dispatch (whose arm was dropped, or which is unspecializable),
+            // or one of so many shapes that the profile has no room left for them. Do not spend a
+            // version to move a few percent of it.
+            crate::stats::incr_counter!(ivar_respecialize_declined_count);
+            IvarReprofiled::Declined
         }
     }
 }
@@ -283,7 +327,9 @@ impl IseqProfile {
 /// Testing it also disarms frames still running the invalidated code.
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_ivar_reprofile(version: *mut crate::payload::IseqVersion, frame_iseq: VALUE, insn_idx: u32, recv: VALUE) {
-    if unsafe { (*version).is_invalidated() } {
+    // Both of these are one load off `version`, which keeps the give-up path from paying for the
+    // payload lookup and the profile timer below.
+    if unsafe { (*version).is_invalidated() || (*version).ivar_reprofile_windows == 0 } {
         return;
     }
     // Immediates have no shape to specialize and would poison the profile with a bucket the
@@ -296,8 +342,26 @@ pub extern "C" fn rb_zjit_ivar_reprofile(version: *mut crate::payload::IseqVersi
     let reprofiled = with_time_stat(profile_time_ns, || {
         get_or_create_iseq_payload(frame_iseq).profile.observe_ivar_fallback(frame_iseq, insn_idx, recv)
     });
-    if reprofiled == IvarReprofiled::Sampled {
-        return;
+    match reprofiled {
+        IvarReprofiled::Sampled => return,
+        IvarReprofiled::Declined => {
+            // Spend one of the version's windows. When they run out the sampling call stays in
+            // the code, but every execution of it stops at the check above.
+            let windows = unsafe { &mut (*version).ivar_reprofile_windows };
+            *windows -= 1;
+            if *windows == 0 {
+                crate::stats::incr_counter!(ivar_respecialize_giveup_count);
+                // Leave the sampling out of whatever this ISEQ compiles next. Invalidating just
+                // to drop it is not worth a compile -- the call is already a no-op after the
+                // check above -- but a version compiled for any other reason should not pay for
+                // evidence this ISEQ has already gathered and rejected.
+                // The flag belongs to the unit that compiled the call, which is what
+                // `Function::emit_ivar_reprofile` tests, not the frame the site came from.
+                get_or_create_iseq_payload(unsafe { (*version).iseq }).ivar_reprofile_giveup = true;
+            }
+            return;
+        }
+        IvarReprofiled::Recompile => {}
     }
     // Read the compiled unit's ISEQ out before taking the lock: `version` points into the
     // payload, and holding a reference to it across the lock's unwind boundary is not allowed.
@@ -507,8 +571,11 @@ pub struct ProfileEntry {
     /// Receivers seen on this ivar site's fallback path in the current re-profiling window.
     /// See [`rb_zjit_ivar_reprofile`].
     ivar_fallback_samples: u32,
-    /// Whether the current re-profiling window has seen a shape the profile did not have.
-    ivar_fallback_new_shape: bool,
+    /// How many of those a recompile would give an arm of its own.
+    ivar_fallback_fixable: u32,
+    /// Number of shapes the dispatch now running was compiled from, i.e. how many buckets of
+    /// `opnd_types[0]` it has arms for. `None` until the first sample after a compile.
+    ivar_dispatch_shapes: Option<u8>,
 }
 
 #[derive(Debug)]
@@ -551,7 +618,8 @@ impl IseqProfile {
                     opnd_types: Vec::new(),
                     profiles_remaining: get_option!(num_profiles),
                     ivar_fallback_samples: 0,
-                    ivar_fallback_new_shape: false,
+                    ivar_fallback_fixable: 0,
+                    ivar_dispatch_shapes: None,
                 });
                 &mut self.entries[i]
             }

zjit/src/stats.rs

@@ -159,6 +159,8 @@ make_counters! {
         jit_frame_heap_bytes,
         failed_jit_frame_bytes,
         ivar_respecialize_count,
+        ivar_respecialize_declined_count,
+        ivar_respecialize_giveup_count,
         skipped_native_stack_full,
         skipped_exceptional_entry_escaped_env,