YJIT: Add a sampling option to exit tracing (#7693)

Add a sampling option to trace exits

Running YJIT with trace exits enabled can make very large metrics files.
This allows us to configure a sample rate to make tracing exits possible
on larger tests. This also updates the documented YJIT options.

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>
Co-authored-by: John Hawthorn <john@hawthorn.email>
Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>

ruby.c

@@ -334,11 +334,13 @@ usage(const char *name, int help, int highlight, int columns)
     };
 #if USE_YJIT
     static const struct ruby_opt_message yjit_options[] = {
-        M("--yjit-stats",              "", "Enable collecting YJIT statistics"),
-        M("--yjit-exec-mem-size=num",  "", "Size of executable memory block in MiB (default: 64)"),
-        M("--yjit-call-threshold=num", "", "Number of calls to trigger JIT (default: 30)"),
-        M("--yjit-max-versions=num",   "", "Maximum number of versions per basic block (default: 4)"),
-        M("--yjit-greedy-versioning",  "", "Greedy versioning mode (default: disabled)"),
+        M("--yjit-stats",                    "", "Enable collecting YJIT statistics"),
+        M("--yjit-trace-exits",              "", "Record Ruby source location when exiting from generated code"),
+        M("--yjit-trace-exits-sample-rate",  "", "Trace exit locations only every Nth occurrence"),
+        M("--yjit-exec-mem-size=num",        "", "Size of executable memory block in MiB (default: 64)"),
+        M("--yjit-call-threshold=num",       "", "Number of calls to trigger JIT (default: 30)"),
+        M("--yjit-max-versions=num",         "", "Maximum number of versions per basic block (default: 4)"),
+        M("--yjit-greedy-versioning",        "", "Greedy versioning mode (default: disabled)"),
     };
 #endif
 #if USE_RJIT

yjit/src/options.rs

@@ -32,6 +32,9 @@ pub struct Options {
     // Trace locations of exits
     pub gen_trace_exits: bool,
 
+    // how often to sample exit trace data
+    pub trace_exits_sample_rate: usize,
+
     // Whether to start YJIT in paused state (initialize YJIT but don't
     // compile anything)
     pub pause: bool,
@@ -59,6 +62,7 @@ pub static mut OPTIONS: Options = Options {
     num_temp_regs: 5,
     gen_stats: false,
     gen_trace_exits: false,
+    trace_exits_sample_rate: 0,
     pause: false,
     dump_insns: false,
     dump_disasm: None,
@@ -173,7 +177,8 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
         ("greedy-versioning", "") => unsafe { OPTIONS.greedy_versioning = true },
         ("no-type-prop", "") => unsafe { OPTIONS.no_type_prop = true },
         ("stats", "") => unsafe { OPTIONS.gen_stats = true },
-        ("trace-exits", "") => unsafe { OPTIONS.gen_trace_exits = true; OPTIONS.gen_stats = true },
+        ("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(); },
         ("dump-insns", "") => unsafe { OPTIONS.dump_insns = true },
         ("verify-ctx", "") => unsafe { OPTIONS.verify_ctx = true },
 
@@ -183,6 +188,19 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
         }
     }
 
+    // before we continue, check that sample_rate is either 0 or a prime number
+    let trace_sample_rate = unsafe { OPTIONS.trace_exits_sample_rate };
+    if trace_sample_rate > 1 {
+        let mut i = 2;
+        while i*i <= trace_sample_rate {
+            if trace_sample_rate % i == 0 {
+                println!("Warning: using a non-prime number as your sampling rate can result in less accurate sampling data");
+                return Some(());
+            }
+            i += 1;
+        }
+    }
+
     // dbg!(unsafe {OPTIONS});
 
     // Option successfully parsed

yjit/src/stats.rs

@@ -26,7 +26,9 @@ pub struct YjitExitLocations {
     raw_samples: Vec<VALUE>,
     /// Vec to hold line_samples which represent line numbers of
     /// the iseq caller.
-    line_samples: Vec<i32>
+    line_samples: Vec<i32>,
+    /// Number of samples skipped when sampling
+    skipped_samples: usize
 }
 
 /// Private singleton instance of yjit exit locations
@@ -47,7 +49,8 @@ impl YjitExitLocations {
 
         let yjit_exit_locations = YjitExitLocations {
             raw_samples: Vec::new(),
-            line_samples: Vec::new()
+            line_samples: Vec::new(),
+            skipped_samples: 0
         };
 
         // Initialize the yjit exit locations instance
@@ -71,6 +74,11 @@ impl YjitExitLocations {
         &mut YjitExitLocations::get_instance().line_samples
     }
 
+    /// Get the number of samples skipped
+    pub fn get_skipped_samples() -> &'static mut usize {
+        &mut YjitExitLocations::get_instance().skipped_samples
+    }
+
     /// Mark the data stored in YjitExitLocations::get_raw_samples that needs to be used by
     /// rb_yjit_add_frame. YjitExitLocations::get_raw_samples are an array of
     /// VALUE pointers, exit instruction, and number of times we've seen this stack row
@@ -573,6 +581,15 @@ pub extern "C" fn rb_yjit_record_exit_stack(exit_pc: *const VALUE)
         return;
     }
 
+    if get_option!(trace_exits_sample_rate) > 0 {
+        if get_option!(trace_exits_sample_rate) <= *YjitExitLocations::get_skipped_samples() {
+            YjitExitLocations::get_instance().skipped_samples = 0;
+        } else {
+            YjitExitLocations::get_instance().skipped_samples += 1;
+            return;
+        }
+    }
+
     // rb_vm_insn_addr2opcode won't work in cargo test --all-features
     // because it's a C function. Without insn call, this function is useless
     // so wrap the whole thing in a not test check.