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

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"

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

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"
 

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
 

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.

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

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

zjit.c

@@ -15,6 +15,7 @@
 #include "builtin.h"
 #include "insns.inc"
 #include "insns_info.inc"
+#include "zjit.h"
 #include "vm_sync.h"
 #include "vm_insnhelper.h"
 #include "probes.h"
@@ -22,7 +23,6 @@
 #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"

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;

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
-  # Assert that any future ZJIT compilation will return a function pointer
-  def self.assert_compiles
-    Primitive.rb_zjit_assert_compiles
+  # 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 assert_compiles # :nodoc:
+    Primitive.rb_zjit_assert_compiles
+  end
+
+  # :stopdoc:
+  private
+
+  # Print ZJIT stats
+  def print_stats
+    $stderr.write stats_string
   end
 end

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")
 

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) {

zjit/src/cruby_bindings.inc.rs

@@ -832,6 +832,7 @@ unsafe extern "C" {
         elts: *const VALUE,
     ) -> VALUE;
     pub fn rb_vm_top_self() -> VALUE;
+    pub static mut rb_vm_insns_count: u64;
     pub fn rb_method_entry_at(obj: VALUE, id: ID) -> *const rb_method_entry_t;
     pub fn rb_callable_method_entry(klass: VALUE, id: ID) -> *const rb_callable_method_entry_t;
     pub fn rb_callable_method_entry_or_negative(

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()

zjit/src/options.rs

@@ -1,5 +1,6 @@
 use std::{ffi::{CStr, CString}, ptr::null};
 use std::os::raw::{c_char, c_int, c_uint};
+use crate::cruby::*;
 
 /// Number of calls to start profiling YARV instructions.
 /// They are profiled `rb_zjit_call_threshold - rb_zjit_profile_threshold` times,
@@ -14,11 +15,18 @@ pub static mut rb_zjit_profile_threshold: u64 = 1;
 #[allow(non_upper_case_globals)]
 pub static mut rb_zjit_call_threshold: u64 = 2;
 
+/// True if --zjit-stats is enabled.
+#[allow(non_upper_case_globals)]
+static mut zjit_stats_enabled_p: bool = false;
+
 #[derive(Clone, Copy, Debug)]
 pub struct Options {
     /// Number of times YARV instructions should be profiled.
     pub num_profiles: u8,
 
+    /// Enable YJIT statsitics
+    pub stats: bool,
+
     /// Enable debug logging
     pub debug: bool,
 
@@ -42,6 +50,7 @@ pub struct Options {
 pub fn init_options() -> Options {
     Options {
         num_profiles: 1,
+        stats: false,
         debug: false,
         dump_hir_init: None,
         dump_hir_opt: None,
@@ -56,6 +65,8 @@ pub fn init_options() -> Options {
 pub const ZJIT_OPTIONS: &'static [(&str, &str)] = &[
     ("--zjit-call-threshold=num", "Number of calls to trigger JIT (default: 2)."),
     ("--zjit-num-profiles=num",   "Number of profiled calls before JIT (default: 1, max: 255)."),
+    ("--zjit-stats",              "Enable collecting ZJIT statistics."),
+    ("--zjit-perf",               "Dump ISEQ symbols into /tmp/perf-{}.map for Linux perf."),
 ];
 
 #[derive(Clone, Copy, Debug)]
@@ -132,6 +143,11 @@ fn parse_option(options: &mut Options, str_ptr: *const std::os::raw::c_char) ->
             Err(_) => return None,
         },
 
+        ("stats", "") => {
+            unsafe { zjit_stats_enabled_p = true; }
+            options.stats = true;
+        }
+
         ("debug", "") => options.debug = true,
 
         // --zjit-dump-hir dumps the actual input to the codegen, which is currently the same as --zjit-dump-hir-opt.
@@ -193,3 +209,15 @@ macro_rules! debug {
     };
 }
 pub(crate) use debug;
+
+/// Return Qtrue if --zjit-stats has been enabled
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_stats_enabled_p(_ec: EcPtr, _self: VALUE) -> VALUE {
+    // ZJITState is not initialized yet when loading builtins, so this relies
+    // on a separate global variable.
+    if unsafe { zjit_stats_enabled_p } {
+        Qtrue
+    } else {
+        Qfalse
+    }
+}

zjit/src/state.rs

@@ -1,8 +1,9 @@
-use crate::cruby::{self, rb_bug_panic_hook, EcPtr, Qnil, VALUE};
+use crate::cruby::{self, rb_bug_panic_hook, rb_vm_insns_count, EcPtr, Qnil, VALUE};
 use crate::cruby_methods;
 use crate::invariants::Invariants;
 use crate::options::Options;
 use crate::asm::CodeBlock;
+use crate::stats::Counters;
 
 #[allow(non_upper_case_globals)]
 #[unsafe(no_mangle)]
@@ -21,6 +22,9 @@ pub struct ZJITState {
     /// ZJIT command-line options
     options: Options,
 
+    /// ZJIT statistics
+    counters: Counters,
+
     /// Assumptions that require invalidation
     invariants: Invariants,
 
@@ -80,6 +84,7 @@ impl ZJITState {
         let zjit_state = ZJITState {
             code_block: cb,
             options,
+            counters: Counters::default(),
             invariants: Invariants::default(),
             assert_compiles: false,
             method_annotations: cruby_methods::init(),
@@ -126,6 +131,11 @@ impl ZJITState {
         let instance = ZJITState::get_instance();
         instance.assert_compiles = true;
     }
+
+    /// Get a mutable reference to counters for ZJIT stats
+    pub fn get_counters() -> &'static mut Counters {
+        &mut ZJITState::get_instance().counters
+    }
 }
 
 /// Initialize ZJIT, given options allocated by rb_zjit_init_options()
@@ -142,6 +152,9 @@ pub extern "C" fn rb_zjit_init(options: *const u8) {
 
         rb_bug_panic_hook();
 
+        // Discard the instruction count for boot which we never compile
+        unsafe { rb_vm_insns_count = 0; }
+
         // ZJIT enabled and initialized successfully
         assert!(unsafe{ !rb_zjit_enabled_p });
         unsafe { rb_zjit_enabled_p = true; }

zjit/src/stats.rs

@@ -4,6 +4,59 @@
 //
 // Comptime vs Runtime stats?
 
+use crate::{cruby::*, options::get_option, state::{zjit_enabled_p, ZJITState}};
+
+macro_rules! make_counters {
+    ($($counter_name:ident,)+) => {
+        /// Struct containing the counter values
+        #[derive(Default, Debug)]
+        pub struct Counters { $(pub $counter_name: u64),+ }
+
+        /// Enum to represent a counter
+        #[allow(non_camel_case_types)]
+        #[derive(Clone, Copy, PartialEq, Eq, Debug)]
+        pub enum Counter { $($counter_name),+ }
+
+        /// Map a counter to a pointer
+        pub fn counter_ptr(counter: Counter) -> *mut u64 {
+            let counters = $crate::state::ZJITState::get_counters();
+            match counter {
+                $( Counter::$counter_name => std::ptr::addr_of_mut!(counters.$counter_name) ),+
+            }
+        }
+    }
+}
+
+// Declare all the counters we track
+make_counters! {
+    // The number of times YARV instructions are executed on JIT code
+    zjit_insns_count,
+}
+
 pub fn zjit_alloc_size() -> usize {
     0 // TODO: report the actual memory usage
 }
+
+/// Return a Hash object that contains ZJIT statistics
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE) -> VALUE {
+    if !zjit_enabled_p() {
+        return Qnil;
+    }
+
+    fn set_stat(hash: VALUE, key: &str, value: u64) {
+        unsafe { rb_hash_aset(hash, rust_str_to_sym(key), VALUE::fixnum_from_usize(value as usize)); }
+    }
+
+    let hash = unsafe { rb_hash_new() };
+    // TODO: Set counters that are always available here
+
+    // Set counters that are enabled when --zjit-stats is enabled
+    if get_option!(stats) {
+        let counters = ZJITState::get_counters();
+        set_stat(hash, "zjit_insns_count", counters.zjit_insns_count);
+        set_stat(hash, "vm_insns_count", unsafe { rb_vm_insns_count });
+    }
+
+    hash
+}