ZJIT: Add polymorphism counters (#14608)

* ZJIT: Add polymorphism counters

* .

* .

zjit.rb

@@ -41,7 +41,8 @@ class << RubyVM::ZJIT
 
     # Show non-exit counters
     print_counters_with_prefix(prefix: 'dynamic_send_type_', prompt: 'dynamic send types', buf:, stats:, limit: 20)
-    print_counters_with_prefix(prefix: 'send_fallback_', prompt: 'send fallback def_types', buf:, stats:, limit: 20)
+    print_counters_with_prefix(prefix: 'unspecialized_def_type_', prompt: 'send fallback unspecialized def_types', buf:, stats:, limit: 20)
+    print_counters_with_prefix(prefix: 'send_fallback_', prompt: 'dynamic send types', buf:, stats:, limit: 20)
 
     # Show exit counters, ordered by the typical amount of exits for the prefix at the time
     print_counters_with_prefix(prefix: 'unhandled_yarv_insn_', prompt: 'unhandled YARV insns', buf:, stats:, limit: 20)

zjit/src/hir.rs

@@ -1690,6 +1690,28 @@ impl Function {
         None
     }
 
+    /// Return whether a given HIR instruction as profiled by the interpreter is polymorphic or
+    /// whether it lacks a profile entirely.
+    ///
+    /// * `Some(true)` if polymorphic
+    /// * `Some(false)` if monomorphic
+    /// * `None` if no profiled information so far
+    fn is_polymorphic_at(&self, insn: InsnId, iseq_insn_idx: usize) -> Option<bool> {
+        let profiles = self.profiles.as_ref()?;
+        let entries = profiles.types.get(&iseq_insn_idx)?;
+        let insn = self.chase_insn(insn);
+        for (entry_insn, entry_type_summary) in entries {
+            if self.union_find.borrow().find_const(*entry_insn) == insn {
+                if !entry_type_summary.is_monomorphic() && !entry_type_summary.is_skewed_polymorphic() {
+                    return Some(true);
+                } else {
+                    return Some(false);
+                }
+            }
+        }
+        None
+    }
+
     fn likely_is_fixnum(&self, val: InsnId, profiled_type: ProfiledType) -> bool {
         self.is_a(val, types::Fixnum) || profiled_type.is_fixnum()
     }
@@ -1840,6 +1862,15 @@ impl Function {
                             // If we know that self is reasonably monomorphic from profile information, guard and use it to fold the lookup at compile-time.
                             // TODO(max): Figure out how to handle top self?
                             let Some(recv_type) = self.profiled_type_of_at(recv, frame_state.insn_idx) else {
+                                if get_option!(stats) {
+                                    match self.is_polymorphic_at(recv, frame_state.insn_idx) {
+                                        Some(true) => self.push_insn(block, Insn::IncrCounter(Counter::send_fallback_polymorphic)),
+                                        // If the class isn't known statically, then it should not also be monomorphic
+                                        Some(false) => panic!("Should not have monomorphic profile at this point in this branch"),
+                                        None => self.push_insn(block, Insn::IncrCounter(Counter::send_fallback_no_profiles)),
+
+                                    };
+                                }
                                 self.push_insn_id(block, insn_id); continue;
                             };
                             (recv_type.class(), Some(recv_type))

zjit/src/stats.rs

@@ -138,19 +138,22 @@ make_counters! {
     dynamic_send_type_invokesuper,
 
     // Method call def_type related to fallback to dynamic dispatch
-    send_fallback_iseq,
-    send_fallback_cfunc,
-    send_fallback_attrset,
-    send_fallback_ivar,
-    send_fallback_bmethod,
-    send_fallback_zsuper,
-    send_fallback_alias,
-    send_fallback_undef,
-    send_fallback_not_implemented,
-    send_fallback_optimized,
-    send_fallback_missing,
-    send_fallback_refined,
-    send_fallback_null,
+    unspecialized_def_type_iseq,
+    unspecialized_def_type_cfunc,
+    unspecialized_def_type_attrset,
+    unspecialized_def_type_ivar,
+    unspecialized_def_type_bmethod,
+    unspecialized_def_type_zsuper,
+    unspecialized_def_type_alias,
+    unspecialized_def_type_undef,
+    unspecialized_def_type_not_implemented,
+    unspecialized_def_type_optimized,
+    unspecialized_def_type_missing,
+    unspecialized_def_type_refined,
+    unspecialized_def_type_null,
+
+    send_fallback_polymorphic,
+    send_fallback_no_profiles,
 
     // Writes to the VM frame
     vm_write_pc_count,
@@ -252,19 +255,19 @@ pub fn send_fallback_counter(def_type: crate::hir::MethodType) -> Counter {
     use crate::stats::Counter::*;
 
     match def_type {
-        Iseq => send_fallback_iseq,
-        Cfunc => send_fallback_cfunc,
-        Attrset => send_fallback_attrset,
-        Ivar => send_fallback_ivar,
-        Bmethod => send_fallback_bmethod,
-        Zsuper => send_fallback_zsuper,
-        Alias => send_fallback_alias,
-        Undefined => send_fallback_undef,
-        NotImplemented => send_fallback_not_implemented,
-        Optimized => send_fallback_optimized,
-        Missing => send_fallback_missing,
-        Refined => send_fallback_refined,
-        Null => send_fallback_null,
+        Iseq => unspecialized_def_type_iseq,
+        Cfunc => unspecialized_def_type_cfunc,
+        Attrset => unspecialized_def_type_attrset,
+        Ivar => unspecialized_def_type_ivar,
+        Bmethod => unspecialized_def_type_bmethod,
+        Zsuper => unspecialized_def_type_zsuper,
+        Alias => unspecialized_def_type_alias,
+        Undefined => unspecialized_def_type_undef,
+        NotImplemented => unspecialized_def_type_not_implemented,
+        Optimized => unspecialized_def_type_optimized,
+        Missing => unspecialized_def_type_missing,
+        Refined => unspecialized_def_type_refined,
+        Null => unspecialized_def_type_null,
     }
 }