ZJIT: Refactor stats implementations (#14378)

* ZJIT: Refactor stats implementations

* s/num_send_dynamic/dynamic_send_count/

6 files changed, 103 insertions, 72 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 9296cd3522..52009d9d0a 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -1485,7 +1485,7 @@ class TestZJIT < Test::Unit::TestCase
     }, call_threshold: 2
   end
 
-  def test_stats
+  def test_stats_availability
     assert_runs '[true, true]', %q{
       def test = 1
       test
@@ -1496,6 +1496,21 @@ class TestZJIT < Test::Unit::TestCase
     }, stats: true
   end
 
+  def test_stats_consistency
+    assert_runs '[]', %q{
+      def test = 1
+      test # increment some counters
+
+      RubyVM::ZJIT.stats.to_a.filter_map do |key, value|
+        # The value may be incremented, but the class should stay the same
+        other_value = RubyVM::ZJIT.stats(key)
+        if value.class != other_value.class
+          [key, value, other_value]
+        end
+      end
+    }, stats: true
+  end
+
   def test_zjit_option_uses_array_each_in_ruby
     omit 'ZJIT wrongly compiles Array#each, so it is disabled for now'
     assert_runs '"<internal:array>"', %q{
diff --git a/zjit.rb b/zjit.rb
index 86d6643df2..41ce371e7e 100644
--- a/zjit.rb
+++ b/zjit.rb
@@ -25,16 +25,8 @@ class << RubyVM::ZJIT
   end
 
   # Return ZJIT statistics as a Hash
-  def stats(key = nil)
-    stats = Primitive.rb_zjit_stats(key)
-    return stats if stats.nil? || !key.nil?
-
-    if stats.key?(:vm_insn_count) && stats.key?(:zjit_insn_count)
-      stats[:total_insn_count] = stats[:vm_insn_count] + stats[:zjit_insn_count]
-      stats[:ratio_in_zjit] = 100.0 * stats[:zjit_insn_count] / stats[:total_insn_count]
-    end
-
-    stats
+  def stats(target_key = nil)
+    Primitive.rb_zjit_stats(target_key)
   end
 
   # Get the summary of ZJIT statistics as a String
@@ -44,6 +36,8 @@ class << RubyVM::ZJIT
 
     print_counters_with_prefix(prefix: 'failed_', prompt: 'compilation failure reasons', buf:, stats:)
     print_counters([
+      :dynamic_send_count,
+
       :compiled_iseq_count,
       :compilation_failure,
 
@@ -56,7 +50,6 @@ class << RubyVM::ZJIT
       :total_insn_count,
       :vm_insn_count,
       :zjit_insn_count,
-      :zjit_dynamic_dispatch,
       :ratio_in_zjit,
     ], buf:, stats:)
     print_counters_with_prefix(prefix: 'exit_', prompt: 'side exit reasons', buf:, stats:, limit: 20)
diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
index 0639176fd6..b372acccea 100644
--- a/zjit/src/backend/lir.rs
+++ b/zjit/src/backend/lir.rs
@@ -6,7 +6,7 @@ use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_
 use crate::hir::SideExitReason;
 use crate::options::{debug, get_option};
 use crate::cruby::VALUE;
-use crate::stats::exit_counter_ptr;
+use crate::stats::{exit_counter_ptr, specific_exit_counter_ptr};
 use crate::virtualmem::CodePtr;
 use crate::asm::{CodeBlock, Label};
 
@@ -1593,29 +1593,15 @@ impl Assembler
                 let cfp_sp = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SP);
                 self.store(cfp_sp, SCRATCH_OPND);
 
+                // Using C_RET_OPND as an additional scratch register, which is no longer used
                 if get_option!(stats) {
                     asm_comment!(self, "increment an exit counter");
                     self.load_into(SCRATCH_OPND, Opnd::const_ptr(exit_counter_ptr(pc)));
-                    let counter_opnd = if cfg!(target_arch = "aarch64") { // See arm64_split()
-                        // Using C_CRET_OPND since arm64_emit uses both SCRATCH0 and SCRATCH1 for IncrCounter.
-                        self.lea_into(C_RET_OPND, Opnd::mem(64, SCRATCH_OPND, 0));
-                        C_RET_OPND
-                    } else { // x86_emit expects Opnd::Mem
-                        Opnd::mem(64, SCRATCH_OPND, 0)
-                    };
-                    self.incr_counter(counter_opnd, 1.into());
+                    self.incr_counter_with_reg(Opnd::mem(64, SCRATCH_OPND, 0), 1.into(), C_RET_OPND);
 
                     asm_comment!(self, "increment a specific exit counter");
-                    let counter = crate::stats::side_exit_reason_counter(reason);
-                    self.load_into(SCRATCH_OPND, Opnd::const_ptr(crate::stats::counter_ptr(counter)));
-                    let counter_opnd = if cfg!(target_arch = "aarch64") { // See arm64_split()
-                        // Using C_CRET_OPND since arm64_emit uses both SCRATCH0 and SCRATCH1 for IncrCounter.
-                        self.lea_into(C_RET_OPND, Opnd::mem(64, SCRATCH_OPND, 0));
-                        C_RET_OPND
-                    } else { // x86_emit expects Opnd::Mem
-                        Opnd::mem(64, SCRATCH_OPND, 0)
-                    };
-                    self.incr_counter(counter_opnd, 1.into());
+                    self.load_into(SCRATCH_OPND, Opnd::const_ptr(specific_exit_counter_ptr(reason)));
+                    self.incr_counter_with_reg(Opnd::mem(64, SCRATCH_OPND, 0), 1.into(), C_RET_OPND);
                 }
 
                 asm_comment!(self, "exit to the interpreter");
@@ -1799,6 +1785,19 @@ impl Assembler {
         self.push_insn(Insn::IncrCounter { mem, value });
     }
 
+    /// incr_counter() but uses a specific register to split Insn::Lea
+    pub fn incr_counter_with_reg(&mut self, mem: Opnd, value: Opnd, reg: Opnd) {
+        assert!(matches!(reg, Opnd::Reg(_)), "incr_counter_with_reg should take a register, got: {reg:?}");
+        let counter_opnd = if cfg!(target_arch = "aarch64") { // See arm64_split()
+            assert_ne!(reg, SCRATCH_OPND, "SCRATCH_REG should be reserved for IncrCounter");
+            self.lea_into(reg, mem);
+            reg
+        } else { // x86_emit() expects Opnd::Mem
+            mem
+        };
+        self.incr_counter(counter_opnd, value);
+    }
+
     pub fn jbe(&mut self, target: Target) {
         self.push_insn(Insn::Jbe(target));
     }
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index fd804e5a42..ba28af73b1 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -858,9 +858,7 @@ fn gen_send_without_block(
     cd: *const rb_call_data,
     state: &FrameState,
 ) -> lir::Opnd {
-    if get_option!(stats) {
-        gen_incr_counter(asm, Counter::zjit_dynamic_dispatch);
-    }
+    gen_incr_counter(asm, Counter::dynamic_send_count);
 
     // Note that it's incorrect to use this frame state to side exit because
     // the state might not be on the boundary of an interpreter instruction.
@@ -1223,14 +1221,16 @@ fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd,
     val
 }
 
-/// Generate code that increments a counter in ZJIT stats
+/// Generate code that increments a counter if --zjit-stats
 fn gen_incr_counter(asm: &mut Assembler, counter: Counter) {
-    let ptr = counter_ptr(counter);
-    let ptr_reg = asm.load(Opnd::const_ptr(ptr as *const u8));
-    let counter_opnd = Opnd::mem(64, ptr_reg, 0);
+    if get_option!(stats) {
+        let ptr = counter_ptr(counter);
+        let ptr_reg = asm.load(Opnd::const_ptr(ptr as *const u8));
+        let counter_opnd = Opnd::mem(64, ptr_reg, 0);
 
-    // Increment and store the updated value
-    asm.incr_counter(counter_opnd, Opnd::UImm(1));
+        // Increment and store the updated value
+        asm.incr_counter(counter_opnd, Opnd::UImm(1));
+    }
 }
 
 /// Save the incremented PC on the CFP.
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 039ba6de5b..d1f5a3af84 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -431,7 +431,9 @@ pub enum SideExitReason {
     UnknownNewarraySend(vm_opt_newarray_send_type),
     UnknownCallType,
     UnknownOpcode(u32),
+    UnknownSpecialVariable(u64),
     UnhandledInstruction(InsnId),
+    UnhandledDefinedType(usize),
     FixnumAddOverflow,
     FixnumSubOverflow,
     FixnumMultOverflow,
@@ -440,8 +442,6 @@ pub enum SideExitReason {
     PatchPoint(Invariant),
     CalleeSideExit,
     ObjToStringFallback,
-    UnknownSpecialVariable(u64),
-    UnhandledDefinedType(usize),
     Interrupt,
 }
 
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
index 10692d83f5..5d0a66ff2c 100644
--- a/zjit/src/stats.rs
+++ b/zjit/src/stats.rs
@@ -73,7 +73,7 @@ make_counters! {
     zjit_insn_count,
 
     // The number of times we do a dynamic dispatch from JIT code
-    zjit_dynamic_dispatch,
+    dynamic_send_count,
 
     // failed_: Compilation failure reasons
     failed_iseq_stack_too_large,
@@ -89,7 +89,9 @@ make_counters! {
     specific_exit_unknown_newarray_send,
     specific_exit_unknown_call_type,
     specific_exit_unknown_opcode,
+    specific_exit_unknown_special_variable,
     specific_exit_unhandled_instruction,
+    specific_exit_unhandled_defined_type,
     specific_exit_fixnum_add_overflow,
     specific_exit_fixnum_sub_overflow,
     specific_exit_fixnum_mult_overflow,
@@ -98,8 +100,6 @@ make_counters! {
     specific_exit_patchpoint,
     specific_exit_callee_side_exit,
     specific_exit_obj_to_string_fallback,
-    specific_exit_unknown_special_variable,
-    specific_exit_unhandled_defined_type,
     specific_exit_interrupt,
 }
 
@@ -127,25 +127,27 @@ pub fn exit_counter_ptr(pc: *const VALUE) -> *mut u64 {
     unsafe { exit_counters.get_unchecked_mut(opcode as usize) }
 }
 
-pub fn side_exit_reason_counter(reason: crate::hir::SideExitReason) -> Counter {
-    use crate::hir::SideExitReason;
-    match reason {
-        SideExitReason::UnknownNewarraySend(_) => Counter::specific_exit_unknown_newarray_send,
-        SideExitReason::UnknownCallType => Counter::specific_exit_unknown_call_type,
-        SideExitReason::UnknownOpcode(_) => Counter::specific_exit_unknown_opcode,
-        SideExitReason::UnhandledInstruction(_) => Counter::specific_exit_unhandled_instruction,
-        SideExitReason::FixnumAddOverflow => Counter::specific_exit_fixnum_add_overflow,
-        SideExitReason::FixnumSubOverflow => Counter::specific_exit_fixnum_sub_overflow,
-        SideExitReason::FixnumMultOverflow => Counter::specific_exit_fixnum_mult_overflow,
-        SideExitReason::GuardType(_) => Counter::specific_exit_guard_type_failure,
-        SideExitReason::GuardBitEquals(_) => Counter::specific_exit_guard_bit_equals_failure,
-        SideExitReason::PatchPoint(_) => Counter::specific_exit_patchpoint,
-        SideExitReason::CalleeSideExit => Counter::specific_exit_callee_side_exit,
-        SideExitReason::ObjToStringFallback => Counter::specific_exit_obj_to_string_fallback,
-        SideExitReason::UnknownSpecialVariable(_) => Counter::specific_exit_unknown_special_variable,
-        SideExitReason::UnhandledDefinedType(_) => Counter::specific_exit_unhandled_defined_type,
-        SideExitReason::Interrupt => Counter::specific_exit_interrupt,
-    }
+pub fn specific_exit_counter_ptr(reason: crate::hir::SideExitReason) -> *mut u64 {
+    use crate::hir::SideExitReason::*;
+    use crate::stats::Counter::*;
+    let counter = match reason {
+        UnknownNewarraySend(_)    => specific_exit_unknown_newarray_send,
+        UnknownCallType           => specific_exit_unknown_call_type,
+        UnknownOpcode(_)          => specific_exit_unknown_opcode,
+        UnknownSpecialVariable(_) => specific_exit_unknown_special_variable,
+        UnhandledInstruction(_)   => specific_exit_unhandled_instruction,
+        UnhandledDefinedType(_)   => specific_exit_unhandled_defined_type,
+        FixnumAddOverflow         => specific_exit_fixnum_add_overflow,
+        FixnumSubOverflow         => specific_exit_fixnum_sub_overflow,
+        FixnumMultOverflow        => specific_exit_fixnum_mult_overflow,
+        GuardType(_)              => specific_exit_guard_type_failure,
+        GuardBitEquals(_)         => specific_exit_guard_bit_equals_failure,
+        PatchPoint(_)             => specific_exit_patchpoint,
+        CalleeSideExit            => specific_exit_callee_side_exit,
+        ObjToStringFallback       => specific_exit_obj_to_string_fallback,
+        Interrupt                 => specific_exit_interrupt,
+    };
+    counter_ptr(counter)
 }
 
 /// Return a Hash object that contains ZJIT statistics
@@ -158,13 +160,24 @@ pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE, target_key: VALUE) ->
     macro_rules! set_stat {
         ($hash:ident, $key:expr, $value:expr) => {
             let key = rust_str_to_sym($key);
-            // Evaluate $value only when it's needed
             if key == target_key {
-                return VALUE::fixnum_from_usize($value as usize);
+                return $value;
             } else if $hash != Qnil {
                 #[allow(unused_unsafe)]
-                unsafe { rb_hash_aset($hash, key, VALUE::fixnum_from_usize($value as usize)); }
+                unsafe { rb_hash_aset($hash, key, $value); }
             }
+        };
+    }
+
+    macro_rules! set_stat_usize {
+        ($hash:ident, $key:expr, $value:expr) => {
+            set_stat!($hash, $key, VALUE::fixnum_from_usize($value as usize))
+        }
+    }
+
+    macro_rules! set_stat_f64 {
+        ($hash:ident, $key:expr, $value:expr) => {
+            set_stat!($hash, $key, unsafe { rb_float_new($value) })
         }
     }
 
@@ -177,14 +190,14 @@ pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE, target_key: VALUE) ->
     // If not --zjit-stats, set only default counters
     if !get_option!(stats) {
         for &counter in DEFAULT_COUNTERS {
-            set_stat!(hash, &counter.name(), unsafe { *counter_ptr(counter) });
+            set_stat_usize!(hash, &counter.name(), unsafe { *counter_ptr(counter) });
         }
         return hash;
     }
 
     // Set all counters for --zjit-stats
     for &counter in ALL_COUNTERS {
-        set_stat!(hash, &counter.name(), unsafe { *counter_ptr(counter) });
+        set_stat_usize!(hash, &counter.name(), unsafe { *counter_ptr(counter) });
     }
 
     // Set side exit stats
@@ -195,12 +208,23 @@ pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE, target_key: VALUE) ->
         let key_string = "exit_".to_owned() + &op_name;
         let count = exit_counters[op_idx];
         side_exit_count += count;
-        set_stat!(hash, &key_string, count);
+        set_stat_usize!(hash, &key_string, count);
     }
-    set_stat!(hash, "side_exit_count", side_exit_count);
+    set_stat_usize!(hash, "side_exit_count", side_exit_count);
+
+    // Share compilation_failure among both prefixes for side-exit stats
+    let counters = ZJITState::get_counters();
+    set_stat_usize!(hash, "specific_exit_compilation_failure", counters.exit_compilation_failure);
 
+    // Only ZJIT_STATS builds support rb_vm_insn_count
     if unsafe { rb_vm_insn_count } > 0 {
-        set_stat!(hash, "vm_insn_count", unsafe { rb_vm_insn_count });
+        let vm_insn_count = unsafe { rb_vm_insn_count };
+        set_stat_usize!(hash, "vm_insn_count", vm_insn_count);
+
+        let total_insn_count = vm_insn_count + counters.zjit_insn_count;
+        set_stat_usize!(hash, "total_insn_count", total_insn_count);
+
+        set_stat_f64!(hash, "ratio_in_yjit", 100.0 * vm_insn_count as f64 / total_insn_count as f64);
     }
 
     hash