ZJIT: Add --zjit-stats (ruby#14034)

k0kubun · web-flow · commit b22eb0e468a8 · 2025-07-29T10:00:15.000-07:00
diff --git a/configure.ac b/configure.ac
@@ -3976,7 +3976,7 @@ AS_CASE(["${YJIT_SUPPORT}"],
 
 ZJIT_LIBS=
 AS_CASE(["${ZJIT_SUPPORT}"],
-[yes|dev|dev_nodebug], [
+[yes|dev|dev_nodebug|stats], [
     AS_IF([test x"$RUSTC" = "xno"],
         AC_MSG_ERROR([rustc is required. Installation instructions available at https://www.rust-lang.org/tools/install])
     )
@@ -3987,11 +3987,16 @@ AS_CASE(["${ZJIT_SUPPORT}"],
     [dev], [
         rb_cargo_features="$rb_cargo_features,disasm"
         JIT_CARGO_SUPPORT=dev
-	AC_DEFINE(RUBY_DEBUG, 1)
+        AC_DEFINE(RUBY_DEBUG, 1)
     ],
     [dev_nodebug], [
         rb_cargo_features="$rb_cargo_features,disasm"
         JIT_CARGO_SUPPORT=dev_nodebug
+        AC_DEFINE(ZJIT_STATS, 1)
+    ],
+    [stats], [
+        JIT_CARGO_SUPPORT=stats
+        AC_DEFINE(ZJIT_STATS, 1)
     ])
 
     ZJIT_LIBS="target/release/libzjit.a"
diff --git a/internal/vm.h b/internal/vm.h
@@ -79,7 +79,7 @@ VALUE rb_block_call2(VALUE obj, ID mid, int argc, const VALUE *argv, rb_block_ca
 struct vm_ifunc *rb_current_ifunc(void);
 VALUE rb_gccct_clear_table(VALUE);
 
-#if USE_YJIT
+#if USE_YJIT || USE_ZJIT
 /* vm_exec.c */
 extern uint64_t rb_vm_insns_count;
 #endif
diff --git a/vm.c b/vm.c
@@ -35,6 +35,8 @@
 #include "iseq.h"
 #include "symbol.h" // This includes a macro for a more performant rb_id2sym.
 #include "yjit.h"
+#include "insns.inc"
+#include "zjit.h"
 #include "ruby/st.h"
 #include "ruby/vm.h"
 #include "vm_core.h"
@@ -45,8 +47,6 @@
 #include "ractor_core.h"
 #include "vm_sync.h"
 #include "shape.h"
-#include "insns.inc"
-#include "zjit.h"
 
 #include "builtin.h"
 
diff --git a/vm_exec.c b/vm_exec.c
@@ -11,8 +11,8 @@
 
 #include <math.h>
 
-#if USE_YJIT
-// The number of instructions executed on vm_exec_core. --yjit-stats uses this.
+#if USE_YJIT || USE_ZJIT
+// The number of instructions executed on vm_exec_core. --yjit-stats and --zjit-stats use this.
 uint64_t rb_vm_insns_count = 0;
 #endif
 
diff --git a/vm_insnhelper.h b/vm_insnhelper.h
@@ -16,7 +16,7 @@ RUBY_EXTERN rb_serial_t ruby_vm_constant_cache_invalidations;
 RUBY_EXTERN rb_serial_t ruby_vm_constant_cache_misses;
 RUBY_EXTERN rb_serial_t ruby_vm_global_cvar_state;
 
-#if USE_YJIT && YJIT_STATS // We want vm_insns_count only on stats builds.
+#if YJIT_STATS || ZJIT_STATS // We want vm_insns_count only on stats builds.
 // Increment vm_insns_count for --yjit-stats. We increment this even when
 // --yjit or --yjit-stats is not used because branching to skip it is slower.
 // We also don't use ATOMIC_INC for performance, allowing inaccuracy on Ractors.
diff --git a/yjit.c b/yjit.c
@@ -23,13 +23,13 @@
 #include "insns_info.inc"
 #include "vm_sync.h"
 #include "yjit.h"
+#include "zjit.h"
 #include "vm_insnhelper.h"
 #include "probes.h"
 #include "probes_helper.h"
 #include "iseq.h"
 #include "ruby/debug.h"
 #include "internal/cont.h"
-#include "zjit.h"
 
 // For mmapp(), sysconf()
 #ifndef _WIN32
diff --git a/yjit.h b/yjit.h
@@ -9,10 +9,9 @@
 #include "vm_core.h"
 #include "method.h"
 
-// YJIT_STATS controls whether to support runtime counters in generated code
-// and in the interpreter.
+// YJIT_STATS controls whether to support runtime counters in the interpreter
 #ifndef YJIT_STATS
-# define YJIT_STATS RUBY_DEBUG
+# define YJIT_STATS (USE_YJIT && RUBY_DEBUG)
 #endif
 
 #if USE_YJIT
diff --git a/zjit.c b/zjit.c
@@ -15,14 +15,14 @@
 #include "builtin.h"
 #include "insns.inc"
 #include "insns_info.inc"
+#include "zjit.h"
 #include "vm_sync.h"
 #include "vm_insnhelper.h"
 #include "probes.h"
 #include "probes_helper.h"
 #include "iseq.h"
 #include "ruby/debug.h"
 #include "internal/cont.h"
-#include "zjit.h"
 
 // For mmapp(), sysconf()
 #ifndef _WIN32
@@ -331,9 +331,6 @@ rb_iseq_set_zjit_payload(const rb_iseq_t *iseq, void *payload)
     iseq->body->zjit_payload = payload;
 }
 
-// Primitives used by zjit.rb
-VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
-
 void
 rb_zjit_print_exception(void)
 {
@@ -349,5 +346,10 @@ rb_zjit_shape_obj_too_complex_p(VALUE obj)
     return rb_shape_obj_too_complex_p(obj);
 }
 
+// Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them.
+VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
+VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self);
+VALUE rb_zjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self);
+
 // Preprocessed zjit.rb generated during build
 #include "zjit.rbinc"
diff --git a/zjit.h b/zjit.h
@@ -4,6 +4,11 @@
 // This file contains definitions ZJIT exposes to the CRuby codebase
 //
 
+// ZJIT_STATS controls whether to support runtime counters in the interpreter
+#ifndef ZJIT_STATS
+# define ZJIT_STATS (USE_ZJIT && RUBY_DEBUG)
+#endif
+
 #if USE_ZJIT
 extern bool rb_zjit_enabled_p;
 extern uint64_t rb_zjit_call_threshold;
diff --git a/zjit.rb b/zjit.rb
@@ -1,6 +1,61 @@
+# frozen_string_literal: true
+
+# This module allows for introspection of \ZJIT, CRuby's just-in-time compiler.
+# Everything in the module is highly implementation specific and the API might
+# be less stable compared to the standard library.
+#
+# This module may not exist if \ZJIT does not support the particular platform
+# for which CRuby is built.
 module RubyVM::ZJIT
+  # Avoid calling a Ruby method here to avoid interfering with compilation tests
+  if Primitive.rb_zjit_stats_enabled_p
+    at_exit { print_stats }
+  end
+end
+
+class << RubyVM::ZJIT
+  # Return ZJIT statistics as a Hash
+  def stats
+    stats = Primitive.rb_zjit_stats
+
+    if stats.key?(:vm_insns_count) && stats.key?(:zjit_insns_count)
+      stats[:total_insns_count] = stats[:vm_insns_count] + stats[:zjit_insns_count]
+      stats[:ratio_in_zjit] = 100.0 * stats[:zjit_insns_count] / stats[:total_insns_count]
+    end
+
+    stats
+  end
+
+  # Get the summary of ZJIT statistics as a String
+  def stats_string
+    buf = +''
+    stats = self.stats
+
+    [
+      :total_insns_count,
+      :vm_insns_count,
+      :zjit_insns_count,
+      :ratio_in_zjit,
+    ].each do |key|
+      value = stats[key]
+      if key == :ratio_in_zjit
+        value = '%0.1f%%' % value
+      end
+      buf << "#{'%-18s' % "#{key}:"} #{value}\n"
+    end
+    buf
+  end
+
   # Assert that any future ZJIT compilation will return a function pointer
-  def self.assert_compiles
+  def assert_compiles # :nodoc:
     Primitive.rb_zjit_assert_compiles
   end
+
+  # :stopdoc:
+  private
+
+  # Print ZJIT stats
+  def print_stats
+    $stderr.write stats_string
+  end
 end
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
@@ -392,6 +392,9 @@ fn main() {
         .allowlist_function("rb_ivar_set")
         .allowlist_function("rb_mod_name")
 
+        // From internal/vm.h
+        .allowlist_var("rb_vm_insns_count")
+
         // From include/ruby/internal/intern/vm.h
         .allowlist_function("rb_get_alloc_func")
 
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -7,6 +7,7 @@ use crate::backend::current::{Reg, ALLOC_REGS};
 use crate::invariants::{track_bop_assumption, track_cme_assumption, track_stable_constant_names_assumption};
 use crate::gc::{get_or_create_iseq_payload, append_gc_offsets};
 use crate::state::ZJITState;
+use crate::stats::{counter_ptr, Counter};
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, asm_comment, asm_ccall, Assembler, Opnd, SideExitContext, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, NATIVE_STACK_PTR, NATIVE_BASE_PTR, SP};
 use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, Invariant, RangeType, SideExitReason, SideExitReason::*, SpecialObjectType, SELF_PARAM_IDX};
@@ -354,6 +355,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::PutSpecialObject { value_type } => gen_putspecialobject(asm, *value_type),
         Insn::AnyToString { val, str, state } => gen_anytostring(asm, opnd!(val), opnd!(str), &function.frame_state(*state))?,
         Insn::Defined { op_type, obj, pushval, v } => gen_defined(jit, asm, *op_type, *obj, *pushval, opnd!(v))?,
+        &Insn::IncrCounter(counter) => return Some(gen_incr_counter(asm, counter)),
         _ => {
             debug!("ZJIT: gen_function: unexpected insn {insn}");
             return None;
@@ -1072,6 +1074,16 @@ fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd,
     Some(val)
 }
 
+/// Generate code that increments a counter in 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);
+
+    // Increment and store the updated value
+    asm.incr_counter(counter_opnd, Opnd::UImm(1));
+}
+
 /// Save the incremented PC on the CFP.
 /// This is necessary when callees can raise or allocate.
 fn gen_save_pc(asm: &mut Assembler, state: &FrameState) {
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -4,7 +4,7 @@
 #![allow(non_upper_case_globals)]
 
 use crate::{
-    cast::IntoUsize, cruby::*, options::{get_option, DumpHIR}, gc::{get_or_create_iseq_payload, IseqPayload}, state::ZJITState
+    cast::IntoUsize, cruby::*, gc::{get_or_create_iseq_payload, IseqPayload}, options::{get_option, DumpHIR}, state::ZJITState, stats::Counter
 };
 use std::{
     cell::RefCell, collections::{HashMap, HashSet, VecDeque}, ffi::{c_int, c_void, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter
@@ -566,6 +566,9 @@ pub enum Insn {
 
     /// Side-exit into the interpreter.
     SideExit { state: InsnId, reason: SideExitReason },
+
+    /// Increment a counter in ZJIT stats
+    IncrCounter(Counter),
 }
 
 impl Insn {
@@ -576,7 +579,7 @@ impl Insn {
             | Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::Return { .. }
             | Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::ArrayExtend { .. }
             | Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetGlobal { .. }
-            | Insn::SetLocal { .. } | Insn::Throw { .. } => false,
+            | Insn::SetLocal { .. } | Insn::Throw { .. } | Insn::IncrCounter(_) => false,
             _ => true,
         }
     }
@@ -786,6 +789,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 write!(f, "Throw {state_string}, {val}")
             }
+            Insn::IncrCounter(counter) => write!(f, "IncrCounter {counter:?}"),
             insn => { write!(f, "{insn:?}") }
         }
     }
@@ -1096,7 +1100,8 @@ impl Function {
                     | PutSpecialObject {..}
                     | GetGlobal {..}
                     | GetLocal {..}
-                    | SideExit {..}) => result.clone(),
+                    | SideExit {..}
+                    | IncrCounter(_)) => result.clone(),
             Snapshot { state: FrameState { iseq, insn_idx, pc, stack, locals } } =>
                 Snapshot {
                     state: FrameState {
@@ -1217,7 +1222,7 @@ impl Function {
             Insn::SetGlobal { .. } | Insn::ArraySet { .. } | Insn::Jump(_)
             | Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::Return { .. } | Insn::Throw { .. }
             | Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::ArrayExtend { .. }
-            | Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetLocal { .. } =>
+            | Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetLocal { .. } | Insn::IncrCounter(_) =>
                 panic!("Cannot infer type of instruction with no output: {}", self.insns[insn.0]),
             Insn::Const { val: Const::Value(val) } => Type::from_value(*val),
             Insn::Const { val: Const::CBool(val) } => Type::from_cbool(*val),
@@ -1835,7 +1840,8 @@ impl Function {
             &Insn::Const { .. }
             | &Insn::Param { .. }
             | &Insn::GetLocal { .. }
-            | &Insn::PutSpecialObject { .. } =>
+            | &Insn::PutSpecialObject { .. }
+            | &Insn::IncrCounter(_) =>
                 {}
             &Insn::PatchPoint { state, .. }
             | &Insn::GetConstantPath { ic: _, state } => {
@@ -2622,6 +2628,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
             let exit_state = state.clone();
             profiles.profile_stack(&exit_state);
 
+            // Increment zjit_insns_count for each YARV instruction if --zjit-stats is enabled.
+            if get_option!(stats) {
+                fun.push_insn(block, Insn::IncrCounter(Counter::zjit_insns_count));
+            }
+
             // try_into() call below is unfortunate. Maybe pick i32 instead of usize for opcodes.
             let opcode: u32 = unsafe { rb_iseq_opcode_at_pc(iseq, pc) }
                 .try_into()
diff --git a/zjit/src/options.rs b/zjit/src/options.rs
diff --git a/zjit/src/state.rs b/zjit/src/state.rs
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
