YJIT: Allow tracing a counted exit (#9890)

* YJIT: Allow tracing a counted exit

* Avoid clobbering caller-saved registers

doc/yjit/yjit.md

@@ -177,8 +177,9 @@ YJIT supports all command-line options supported by upstream CRuby, but also add
   It will cause all machine code to be discarded when the executable memory size limit is hit, meaning JIT compilation will then start over.
   This can allow you to use a lower executable memory size limit, but may cause a slight drop in performance when the limit is hit.
 - `--yjit-perf`: enable frame pointers and profiling with the `perf` tool
-- `--yjit-trace-exits`: produce a Marshal dump of backtraces from specific exits. Automatically enables `--yjit-stats`
-- `--yjit-trace-exits-sample-rate=N`: trace exit locations only every Nth occurrence
+- `--yjit-trace-exits`: produce a Marshal dump of backtraces from all exits. Automatically enables `--yjit-stats`
+- `--yjit-trace-exits=COUNTER`: produce a Marshal dump of backtraces from specified exits. Automatically enables `--yjit-stats`
+- `--yjit-trace-exits-sample-rate=N`: trace exit locations only every Nth occurrence. Automatically enables `--yjit-trace-exits`
 
 Note that there is also an environment variable `RUBY_YJIT_ENABLE` which can be used to enable YJIT.
 This can be useful for some deployment scripts where specifying an extra command-line option to Ruby is not practical.

yjit/src/backend/ir.rs

@@ -1117,7 +1117,7 @@ impl Assembler
         };
 
         // Wrap a counter if needed
-        gen_counted_exit(side_exit, ocb, counter)
+        gen_counted_exit(side_exit_context.pc, side_exit, ocb, counter)
     }
 
     /// Create a new label instance that we can jump to

yjit/src/codegen.rs

@@ -532,9 +532,9 @@ fn gen_exit(exit_pc: *mut VALUE, asm: &mut Assembler) {
             vec![Opnd::const_ptr(exit_pc as *const u8)]
         );
 
-        // If --yjit-trace-exits option is enabled, record the exit stack
-        // while recording the side exits.
-        if get_option!(gen_trace_exits) {
+        // If --yjit-trace-exits is enabled, record the exit stack while recording
+        // the side exits. TraceExits::Counter is handled by gen_counted_exit().
+        if get_option!(trace_exits) == Some(TraceExits::All) {
             asm.ccall(
                 rb_yjit_record_exit_stack as *const u8,
                 vec![Opnd::const_ptr(exit_pc as *const u8)]
@@ -575,7 +575,7 @@ pub fn gen_outlined_exit(exit_pc: *mut VALUE, ctx: &Context, ocb: &mut OutlinedC
 }
 
 /// Get a side exit. Increment a counter in it if --yjit-stats is enabled.
-pub fn gen_counted_exit(side_exit: CodePtr, ocb: &mut OutlinedCb, counter: Option<Counter>) -> Option<CodePtr> {
+pub fn gen_counted_exit(exit_pc: *mut VALUE, side_exit: CodePtr, ocb: &mut OutlinedCb, counter: Option<Counter>) -> Option<CodePtr> {
     // The counter is only incremented when stats are enabled
     if !get_option!(gen_stats) {
         return Some(side_exit);
@@ -587,13 +587,16 @@ pub fn gen_counted_exit(side_exit: CodePtr, ocb: &mut OutlinedCb, counter: Optio
 
     let mut asm = Assembler::new();
 
-    // Load the pointer into a register
-    asm_comment!(asm, "increment counter {}", counter.get_name());
-    let ptr_reg = asm.load(Opnd::const_ptr(get_counter_ptr(&counter.get_name()) as *const u8));
-    let counter_opnd = Opnd::mem(64, ptr_reg, 0);
+    // Increment a counter
+    gen_counter_incr(&mut asm, counter);
 
-    // Increment and store the updated value
-    asm.incr_counter(counter_opnd, Opnd::UImm(1));
+    // Trace a counted exit if --yjit-trace-exits=counter is given.
+    // TraceExits::All is handled by gen_exit().
+    if get_option!(trace_exits) == Some(TraceExits::CountedExit(counter)) {
+        with_caller_saved_temp_regs(&mut asm, |asm| {
+            asm.ccall(rb_yjit_record_exit_stack as *const u8, vec![Opnd::const_ptr(exit_pc as *const u8)]);
+        });
+    }
 
     // Jump to the existing side exit
     asm.jmp(Target::CodePtr(side_exit));
@@ -602,6 +605,18 @@ pub fn gen_counted_exit(side_exit: CodePtr, ocb: &mut OutlinedCb, counter: Optio
     asm.compile(ocb, None).map(|(code_ptr, _)| code_ptr)
 }
 
+/// Preserve caller-saved stack temp registers during the call of a given block
+fn with_caller_saved_temp_regs<F, R>(asm: &mut Assembler, block: F) -> R where F: FnOnce(&mut Assembler) -> R {
+    for &reg in caller_saved_temp_regs() {
+        asm.cpush(Opnd::Reg(reg)); // save stack temps
+    }
+    let ret = block(asm);
+    for &reg in caller_saved_temp_regs().rev() {
+        asm.cpop_into(Opnd::Reg(reg)); // restore stack temps
+    }
+    ret
+}
+
 // Ensure that there is an exit for the start of the block being compiled.
 // Block invalidation uses this exit.
 #[must_use]

yjit/src/core.rs

@@ -2960,7 +2960,7 @@ pub fn gen_branch_stub_hit_trampoline(ocb: &mut OutlinedCb) -> Option<CodePtr> {
 }
 
 /// Return registers to be pushed and popped on branch_stub_hit.
-fn caller_saved_temp_regs() -> impl Iterator<Item = &'static Reg> + DoubleEndedIterator {
+pub fn caller_saved_temp_regs() -> impl Iterator<Item = &'static Reg> + DoubleEndedIterator {
     let temp_regs = Assembler::get_temp_regs().iter();
     let len = temp_regs.len();
     // The return value gen_leave() leaves in C_RET_REG

yjit/src/options.rs

@@ -1,5 +1,5 @@
 use std::{ffi::{CStr, CString}, ptr::null, fs::File};
-use crate::backend::current::TEMP_REGS;
+use crate::{backend::current::TEMP_REGS, stats::Counter};
 use std::os::raw::{c_char, c_int, c_uint};
 
 // Call threshold for small deployments and command-line apps
@@ -48,7 +48,7 @@ pub struct Options {
     pub print_stats: bool,
 
     // Trace locations of exits
-    pub gen_trace_exits: bool,
+    pub trace_exits: Option<TraceExits>,
 
     // how often to sample exit trace data
     pub trace_exits_sample_rate: usize,
@@ -86,7 +86,7 @@ pub static mut OPTIONS: Options = Options {
     max_versions: 4,
     num_temp_regs: 5,
     gen_stats: false,
-    gen_trace_exits: false,
+    trace_exits: None,
     print_stats: true,
     trace_exits_sample_rate: 0,
     disable: false,
@@ -112,6 +112,14 @@ static YJIT_OPTIONS: [(&str, &str); 9] = [
     ("--yjit-trace-exits-sample-rate=num", "Trace exit locations only every Nth occurrence"),
 ];
 
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum TraceExits {
+    // Trace all exits
+    All,
+    // Trace a specific counted exit
+    CountedExit(Counter),
+}
+
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub enum DumpDisasm {
     // Dump to stdout
@@ -267,8 +275,23 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
                 return None;
             }
         },
-        ("trace-exits", "") => unsafe { OPTIONS.gen_trace_exits = true; OPTIONS.gen_stats = true; OPTIONS.trace_exits_sample_rate = 0 },
-        ("trace-exits-sample-rate", sample_rate) => unsafe { OPTIONS.gen_trace_exits = true; OPTIONS.gen_stats = true; OPTIONS.trace_exits_sample_rate = sample_rate.parse().unwrap(); },
+        ("trace-exits", _) => unsafe {
+            OPTIONS.gen_stats = true;
+            OPTIONS.trace_exits = match opt_val {
+                "" => Some(TraceExits::All),
+                name => match Counter::get(name) {
+                    Some(counter) => Some(TraceExits::CountedExit(counter)),
+                    None => return None,
+                },
+            };
+        },
+        ("trace-exits-sample-rate", sample_rate) => unsafe {
+            OPTIONS.gen_stats = true;
+            if OPTIONS.trace_exits.is_none() {
+                OPTIONS.trace_exits = Some(TraceExits::All);
+            }
+            OPTIONS.trace_exits_sample_rate = sample_rate.parse().unwrap();
+        },
         ("dump-insns", "") => unsafe { OPTIONS.dump_insns = true },
         ("verify-ctx", "") => unsafe { OPTIONS.verify_ctx = true },
 

yjit/src/stats.rs

@@ -128,7 +128,7 @@ impl YjitExitLocations {
     /// Initialize the yjit exit locations
     pub fn init() {
         // Return if --yjit-trace-exits isn't enabled
-        if !get_option!(gen_trace_exits) {
+        if get_option!(trace_exits).is_none() {
             return;
         }
 
@@ -177,7 +177,7 @@ impl YjitExitLocations {
         }
 
         // Return if --yjit-trace-exits isn't enabled
-        if !get_option!(gen_trace_exits) {
+        if get_option!(trace_exits).is_none() {
             return;
         }
 
@@ -219,6 +219,14 @@ macro_rules! make_counters {
         pub enum Counter { $($counter_name),+ }
 
         impl Counter {
+            /// Map a counter name string to a counter enum
+            pub fn get(name: &str) -> Option<Counter> {
+                match name {
+                    $( stringify!($counter_name) => { Some(Counter::$counter_name) } ),+
+                    _ => None,
+                }
+            }
+
             /// Get a counter name string
             pub fn get_name(&self) -> String {
                 match self {
@@ -636,7 +644,7 @@ pub extern "C" fn rb_yjit_get_stats(_ec: EcPtr, _ruby_self: VALUE, context: VALU
 /// to be enabled.
 #[no_mangle]
 pub extern "C" fn rb_yjit_trace_exit_locations_enabled_p(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
-    if get_option!(gen_trace_exits) {
+    if get_option!(trace_exits).is_some() {
         return Qtrue;
     }
 
@@ -653,7 +661,7 @@ pub extern "C" fn rb_yjit_get_exit_locations(_ec: EcPtr, _ruby_self: VALUE) -> V
     }
 
     // Return if --yjit-trace-exits isn't enabled
-    if !get_option!(gen_trace_exits) {
+    if get_option!(trace_exits).is_none() {
         return Qnil;
     }
 
@@ -834,7 +842,7 @@ pub extern "C" fn rb_yjit_record_exit_stack(_exit_pc: *const VALUE)
     }
 
     // Return if --yjit-trace-exits isn't enabled
-    if !get_option!(gen_trace_exits) {
+    if get_option!(trace_exits).is_none() {
         return;
     }