ZJIT: Add trace exit counter (#14831)

4 files changed, 64 insertions, 21 deletions

diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
index aad5600f56..6efb3e1259 100644
--- a/zjit/src/backend/lir.rs
+++ b/zjit/src/backend/lir.rs
@@ -4,9 +4,9 @@ use std::mem::take;
 use crate::codegen::local_size_and_idx_to_ep_offset;
 use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32, vm_stack_canary};
 use crate::hir::SideExitReason;
-use crate::options::{debug, get_option};
+use crate::options::{debug, get_option, TraceExits};
 use crate::cruby::VALUE;
-use crate::stats::{exit_counter_ptr, exit_counter_ptr_for_opcode, CompileError};
+use crate::stats::{exit_counter_ptr, exit_counter_ptr_for_opcode, side_exit_counter, CompileError};
 use crate::virtualmem::CodePtr;
 use crate::asm::{CodeBlock, Label};
 use crate::state::rb_zjit_record_exit_stack;
@@ -1644,8 +1644,22 @@ impl Assembler
                     }
                 }
 
-                if get_option!(trace_side_exits) {
-                    asm_ccall!(self, rb_zjit_record_exit_stack, Opnd::const_ptr(pc as *const u8));
+                if get_option!(trace_side_exits).is_some() {
+                    // Get the corresponding `Counter` for the current `SideExitReason`.
+                    let side_exit_counter = side_exit_counter(reason);
+
+                    // Only record the exit if `trace_side_exits` is defined and the counter is either the one specified
+                    let should_record_exit = get_option!(trace_side_exits)
+                        .map(|trace| match trace {
+                            TraceExits::All => true,
+                            TraceExits::Counter(counter) if counter == side_exit_counter => true,
+                            _ => false,
+                        })
+                        .unwrap_or(false);
+
+                    if should_record_exit {
+                        asm_ccall!(self, rb_zjit_record_exit_stack, Opnd::const_ptr(pc as *const u8));
+                    }
                 }
 
                 asm_comment!(self, "exit to the interpreter");
diff --git a/zjit/src/options.rs b/zjit/src/options.rs
index 4ef998aced..b7b20e63c4 100644
--- a/zjit/src/options.rs
+++ b/zjit/src/options.rs
@@ -3,6 +3,7 @@
 use std::{ffi::{CStr, CString}, ptr::null};
 use std::os::raw::{c_char, c_int, c_uint};
 use crate::cruby::*;
+use crate::stats::Counter;
 use std::collections::HashSet;
 
 /// Default --zjit-num-profiles
@@ -72,7 +73,7 @@ pub struct Options {
     pub dump_disasm: bool,
 
     /// Trace and write side exit source maps to /tmp for stackprof.
-    pub trace_side_exits: bool,
+    pub trace_side_exits: Option<TraceExits>,
 
     /// Frequency of tracing side exits.
     pub trace_side_exits_sample_interval: usize,
@@ -102,7 +103,7 @@ impl Default for Options {
             dump_hir_graphviz: None,
             dump_lir: false,
             dump_disasm: false,
-            trace_side_exits: false,
+            trace_side_exits: None,
             trace_side_exits_sample_interval: 0,
             perf: false,
             allowed_iseqs: None,
@@ -125,12 +126,20 @@ pub const ZJIT_OPTIONS: &[(&str, &str)] = &[
     ("--zjit-perf",  "Dump ISEQ symbols into /tmp/perf-{}.map for Linux perf."),
     ("--zjit-log-compiled-iseqs=path",
                      "Log compiled ISEQs to the file. The file will be truncated."),
-    ("--zjit-trace-exits",
-                     "Record Ruby source location when side-exiting."),
-    ("--zjit-trace-exits-sample-rate",
+    ("--zjit-trace-exits[=counter]",
+                     "Record source on side-exit. `Counter` picks specific counter."),
+    ("--zjit-trace-exits-sample-rate=num",
                      "Frequency at which to record side exits. Must be `usize`.")
 ];
 
+#[derive(Copy, Clone, Debug)]
+pub enum TraceExits {
+    // Trace all exits
+    All,
+    // Trace exits for a specific `Counter`
+    Counter(Counter),
+}
+
 #[derive(Clone, Copy, Debug)]
 pub enum DumpHIR {
     // Dump High-level IR without Snapshot
@@ -249,13 +258,18 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             options.print_stats = false;
         }
 
-        ("trace-exits", "") => {
-            options.trace_side_exits = true;
+        ("trace-exits", exits) => {
+            options.trace_side_exits = match exits {
+                "" => Some(TraceExits::All),
+                name => Counter::get(name).map(TraceExits::Counter),
+            }
         }
 
         ("trace-exits-sample-rate", sample_interval) => {
-            // Even if `trace_side_exits` is already set, set it.
-            options.trace_side_exits = true;
+            // If not already set, then set it to `TraceExits::All` by default.
+            if options.trace_side_exits.is_none() {
+                options.trace_side_exits = Some(TraceExits::All);
+            }
             // `sample_interval ` must provide a string that can be validly parsed to a `usize`.
             options.trace_side_exits_sample_interval = sample_interval.parse::<usize>().ok()?;
         }
diff --git a/zjit/src/state.rs b/zjit/src/state.rs
index 1b766d5bc4..c0e9e0b77c 100644
--- a/zjit/src/state.rs
+++ b/zjit/src/state.rs
@@ -82,7 +82,7 @@ impl ZJITState {
         let exit_trampoline = gen_exit_trampoline(&mut cb).unwrap();
         let function_stub_hit_trampoline = gen_function_stub_hit_trampoline(&mut cb).unwrap();
 
-        let exit_locations = if get_option!(trace_side_exits) {
+        let exit_locations = if get_option!(trace_side_exits).is_some() {
             Some(SideExitLocations::default())
         } else {
             None
@@ -369,7 +369,7 @@ fn try_increment_existing_stack(
 /// Record a backtrace with ZJIT side exits
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_record_exit_stack(exit_pc: *const VALUE) {
-    if !zjit_enabled_p() || !get_option!(trace_side_exits) {
+    if !zjit_enabled_p() || get_option!(trace_side_exits).is_none() {
         return;
     }
 
@@ -425,7 +425,7 @@ pub extern "C" fn rb_zjit_record_exit_stack(exit_pc: *const VALUE) {
 /// Mark `raw_samples` so they can be used by rb_zjit_add_frame.
 pub fn gc_mark_raw_samples() {
     // Return if ZJIT is not enabled
-    if !zjit_enabled_p() || !get_option!(trace_side_exits) {
+    if !zjit_enabled_p() || get_option!(trace_side_exits).is_none() {
         return;
     }
 
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
index d902c69b79..6faa328a1c 100644
--- a/zjit/src/stats.rs
+++ b/zjit/src/stats.rs
@@ -57,6 +57,17 @@ macro_rules! make_counters {
                     $( Counter::$counter_name => stringify!($counter_name), )+
                 }
             }
+
+            pub fn get(name: &str) -> Option<Counter> {
+                match name {
+                    $( stringify!($default_counter_name) => Some(Counter::$default_counter_name), )+
+                    $( stringify!($exit_counter_name) => Some(Counter::$exit_counter_name), )+
+                    $( stringify!($dynamic_send_counter_name) => Some(Counter::$dynamic_send_counter_name), )+
+                    $( stringify!($optimized_send_counter_name) => Some(Counter::$optimized_send_counter_name), )+
+                    $( stringify!($counter_name) => Some(Counter::$counter_name), )+
+                    _ => None,
+                }
+            }
         }
 
         /// Map a counter to a pointer
@@ -298,11 +309,11 @@ pub fn exit_counter_for_compile_error(compile_error: &CompileError) -> Counter {
     }
 }
 
-pub fn exit_counter_ptr(reason: crate::hir::SideExitReason) -> *mut u64 {
+pub fn side_exit_counter(reason: crate::hir::SideExitReason) -> Counter {
     use crate::hir::SideExitReason::*;
     use crate::hir::CallType::*;
     use crate::stats::Counter::*;
-    let counter = match reason {
+    match reason {
         UnknownNewarraySend(_)        => exit_unknown_newarray_send,
         UnhandledCallType(Tailcall)   => exit_unhandled_tailcall,
         UnhandledCallType(Splat)      => exit_unhandled_splat,
@@ -324,7 +335,11 @@ pub fn exit_counter_ptr(reason: crate::hir::SideExitReason) -> *mut u64 {
         StackOverflow                 => exit_stackoverflow,
         BlockParamProxyModified       => exit_block_param_proxy_modified,
         BlockParamProxyNotIseqOrIfunc => exit_block_param_proxy_not_iseq_or_ifunc,
-    };
+    }
+}
+
+pub fn exit_counter_ptr(reason: crate::hir::SideExitReason) -> *mut u64 {
+    let counter = side_exit_counter(reason);
     counter_ptr(counter)
 }
 
@@ -563,7 +578,7 @@ pub struct SideExitLocations {
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_trace_exit_locations_enabled_p(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
     // Builtin zjit.rb calls this even if ZJIT is disabled, so OPTIONS may not be set.
-    if unsafe { OPTIONS.as_ref() }.is_some_and(|opts| opts.trace_side_exits) {
+    if unsafe { OPTIONS.as_ref() }.is_some_and(|opts| opts.trace_side_exits.is_some()) {
         Qtrue
     } else {
         Qfalse
@@ -574,7 +589,7 @@ pub extern "C" fn rb_zjit_trace_exit_locations_enabled_p(_ec: EcPtr, _ruby_self:
 /// into raw, lines, and frames hash for RubyVM::YJIT.exit_locations.
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_get_exit_locations(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
-    if !zjit_enabled_p() || !get_option!(trace_side_exits) {
+    if !zjit_enabled_p() || get_option!(trace_side_exits).is_none() {
         return Qnil;
     }