YJIT: Allow tracing a counted exit (#9890)

* YJIT: Allow tracing a counted exit

* Avoid clobbering caller-saved registers
This commit is contained in:
Takashi Kokubun 2024-02-08 15:47:02 -08:00 committed by GitHub
parent 4a40364c62
commit 5cbca9110c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 71 additions and 24 deletions

View File

@ -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. 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. 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-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`: produce a Marshal dump of backtraces from all exits. Automatically enables `--yjit-stats`
- `--yjit-trace-exits-sample-rate=N`: trace exit locations only every Nth occurrence - `--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. 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. This can be useful for some deployment scripts where specifying an extra command-line option to Ruby is not practical.

View File

@ -1117,7 +1117,7 @@ impl Assembler
}; };
// Wrap a counter if needed // 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 /// Create a new label instance that we can jump to

View File

@ -532,9 +532,9 @@ fn gen_exit(exit_pc: *mut VALUE, asm: &mut Assembler) {
vec![Opnd::const_ptr(exit_pc as *const u8)] vec![Opnd::const_ptr(exit_pc as *const u8)]
); );
// If --yjit-trace-exits option is enabled, record the exit stack // If --yjit-trace-exits is enabled, record the exit stack while recording
// while recording the side exits. // the side exits. TraceExits::Counter is handled by gen_counted_exit().
if get_option!(gen_trace_exits) { if get_option!(trace_exits) == Some(TraceExits::All) {
asm.ccall( asm.ccall(
rb_yjit_record_exit_stack as *const u8, rb_yjit_record_exit_stack as *const u8,
vec![Opnd::const_ptr(exit_pc 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. /// 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 // The counter is only incremented when stats are enabled
if !get_option!(gen_stats) { if !get_option!(gen_stats) {
return Some(side_exit); 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(); let mut asm = Assembler::new();
// Load the pointer into a register // Increment a counter
asm_comment!(asm, "increment counter {}", counter.get_name()); gen_counter_incr(&mut asm, counter);
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 and store the updated value // Trace a counted exit if --yjit-trace-exits=counter is given.
asm.incr_counter(counter_opnd, Opnd::UImm(1)); // 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 // Jump to the existing side exit
asm.jmp(Target::CodePtr(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) 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. // Ensure that there is an exit for the start of the block being compiled.
// Block invalidation uses this exit. // Block invalidation uses this exit.
#[must_use] #[must_use]

View File

@ -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. /// 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 temp_regs = Assembler::get_temp_regs().iter();
let len = temp_regs.len(); let len = temp_regs.len();
// The return value gen_leave() leaves in C_RET_REG // The return value gen_leave() leaves in C_RET_REG

View File

@ -1,5 +1,5 @@
use std::{ffi::{CStr, CString}, ptr::null, fs::File}; 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}; use std::os::raw::{c_char, c_int, c_uint};
// Call threshold for small deployments and command-line apps // Call threshold for small deployments and command-line apps
@ -48,7 +48,7 @@ pub struct Options {
pub print_stats: bool, pub print_stats: bool,
// Trace locations of exits // Trace locations of exits
pub gen_trace_exits: bool, pub trace_exits: Option<TraceExits>,
// how often to sample exit trace data // how often to sample exit trace data
pub trace_exits_sample_rate: usize, pub trace_exits_sample_rate: usize,
@ -86,7 +86,7 @@ pub static mut OPTIONS: Options = Options {
max_versions: 4, max_versions: 4,
num_temp_regs: 5, num_temp_regs: 5,
gen_stats: false, gen_stats: false,
gen_trace_exits: false, trace_exits: None,
print_stats: true, print_stats: true,
trace_exits_sample_rate: 0, trace_exits_sample_rate: 0,
disable: false, 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"), ("--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)] #[derive(Clone, PartialEq, Eq, Debug)]
pub enum DumpDisasm { pub enum DumpDisasm {
// Dump to stdout // Dump to stdout
@ -267,8 +275,23 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
return None; return None;
} }
}, },
("trace-exits", "") => unsafe { OPTIONS.gen_trace_exits = true; OPTIONS.gen_stats = true; OPTIONS.trace_exits_sample_rate = 0 }, ("trace-exits", _) => unsafe {
("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(); }, 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 }, ("dump-insns", "") => unsafe { OPTIONS.dump_insns = true },
("verify-ctx", "") => unsafe { OPTIONS.verify_ctx = true }, ("verify-ctx", "") => unsafe { OPTIONS.verify_ctx = true },

View File

@ -128,7 +128,7 @@ impl YjitExitLocations {
/// Initialize the yjit exit locations /// Initialize the yjit exit locations
pub fn init() { pub fn init() {
// Return if --yjit-trace-exits isn't enabled // Return if --yjit-trace-exits isn't enabled
if !get_option!(gen_trace_exits) { if get_option!(trace_exits).is_none() {
return; return;
} }
@ -177,7 +177,7 @@ impl YjitExitLocations {
} }
// Return if --yjit-trace-exits isn't enabled // Return if --yjit-trace-exits isn't enabled
if !get_option!(gen_trace_exits) { if get_option!(trace_exits).is_none() {
return; return;
} }
@ -219,6 +219,14 @@ macro_rules! make_counters {
pub enum Counter { $($counter_name),+ } pub enum Counter { $($counter_name),+ }
impl Counter { 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 /// Get a counter name string
pub fn get_name(&self) -> String { pub fn get_name(&self) -> String {
match self { match self {
@ -636,7 +644,7 @@ pub extern "C" fn rb_yjit_get_stats(_ec: EcPtr, _ruby_self: VALUE, context: VALU
/// to be enabled. /// to be enabled.
#[no_mangle] #[no_mangle]
pub extern "C" fn rb_yjit_trace_exit_locations_enabled_p(_ec: EcPtr, _ruby_self: VALUE) -> VALUE { 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; 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 // Return if --yjit-trace-exits isn't enabled
if !get_option!(gen_trace_exits) { if get_option!(trace_exits).is_none() {
return Qnil; 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 // Return if --yjit-trace-exits isn't enabled
if !get_option!(gen_trace_exits) { if get_option!(trace_exits).is_none() {
return; return;
} }