ZJIT: Allow --zjit-perf to dump perf symbols for HIR (ruby#16501)

k0kubun · web-flow · commit 8a5fc14f7ed9 · 2026-03-23T13:14:07.000-07:00
diff --git a/misc/jit_perf.py b/misc/jit_perf.py
@@ -22,8 +22,7 @@ def categorize_symbol(dso, symbol):
         return '[sha256]'
     elif symbol.startswith('[JIT] gen_send'):
         return '[JIT send]'
-    # TODO: Stop using zjit:: as the prefix for JIT code. Rust modules and JIT code should use different namespaces.
-    elif symbol.startswith('[JIT]') or (symbol.startswith('zjit::') and '@') or symbol == 'zjit::ZJIT entry trampoline':
+    elif symbol.startswith('[JIT]') or symbol.startswith('ZJIT: ') or dso.startswith('perf-'):
         return '[JIT code]'
     elif '::' in symbol or symbol.startswith('_ZN4yjit') or symbol.startswith('_ZN4zjit'):
         return '[JIT compile]'
diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
@@ -1734,6 +1734,13 @@ mod tests {
         (asm, CodeBlock::new_dummy())
     }
 
+    fn setup_asm_with_scratch_reg() -> (Assembler, CodeBlock, Opnd) {
+        crate::options::rb_zjit_prepare_options(); // Allow `get_option!` in Assembler
+        let (mut asm, scratch_reg) = Assembler::new_with_scratch_reg();
+        asm.new_block_without_id("test");
+        (asm, CodeBlock::new_dummy(), scratch_reg)
+    }
+
     #[test]
     fn test_lir_string() {
         use crate::hir::SideExitReason;
@@ -2159,9 +2166,7 @@ mod tests {
 
     #[test]
     fn test_store_with_valid_scratch_reg() {
-        let (mut asm, scratch_reg) = Assembler::new_with_scratch_reg();
-        asm.new_block_without_id("test");
-        let mut cb = CodeBlock::new_dummy();
+        let (mut asm, mut cb, scratch_reg) = setup_asm_with_scratch_reg();
         asm.store(Opnd::mem(64, scratch_reg, 0), 0x83902.into());
 
         asm.compile_with_num_regs(&mut cb, 0);
diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
@@ -5,11 +5,11 @@ use std::panic;
 use std::rc::Rc;
 use std::sync::{Arc, Mutex};
 use crate::bitset::BitSet;
-use crate::codegen::local_size_and_idx_to_ep_offset;
+use crate::codegen::{local_size_and_idx_to_ep_offset, perf_symbol_range_start, perf_symbol_range_end};
 use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32, vm_stack_canary};
 use crate::hir::{Invariant, SideExitReason};
 use crate::hir;
-use crate::options::{TraceExits, get_option};
+use crate::options::{TraceExits, PerfMap, get_option};
 use crate::cruby::VALUE;
 use crate::payload::IseqVersionRef;
 use crate::stats::{exit_counter_ptr, exit_counter_ptr_for_opcode, side_exit_counter, CompileError};
@@ -2675,6 +2675,13 @@ impl Assembler
         // Map from SideExit to compiled Label. This table is used to deduplicate side exit code.
         let mut compiled_exits: HashMap<SideExit, Label> = HashMap::new();
 
+        // Start a new perf range for side exits
+        let perf_symbol = if get_option!(perf) == Some(PerfMap::HIR) {
+            Some(perf_symbol_range_start(self, "side exit"))
+        } else {
+            None
+        };
+
         // Mark the start of side-exit code so we can measure its size
         if !targets.is_empty() {
             self.pos_marker(move |start_pos, cb| {
@@ -2756,6 +2763,11 @@ impl Assembler
             crate::stats::incr_counter_by(crate::stats::Counter::compile_side_exit_time_ns, nanos as u64);
         }
 
+        // Close the current perf range for side exits
+        if let Some(perf_symbol) = &perf_symbol {
+            perf_symbol_range_end(self, perf_symbol);
+        }
+
         // Extract exit instructions and restore the previous current block
         let exit_insns = take(&mut self.basic_blocks[exit_block.0].insns);
         self.set_current_block(saved_block);
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -23,7 +23,7 @@ use crate::backend::lir::{self, Assembler, C_ARG_OPNDS, C_RET_OPND, CFP, EC, NAT
 use crate::hir::{iseq_to_hir, BlockId, Invariant, RangeType, SideExitReason::{self, *}, SpecialBackrefSymbol, SpecialObjectType};
 use crate::hir::{Const, FrameState, Function, Insn, InsnId, SendFallbackReason};
 use crate::hir_type::{types, Type};
-use crate::options::get_option;
+use crate::options::{get_option, PerfMap};
 use crate::cast::IntoUsize;
 
 /// At the moment, we support recompiling each ISEQ only once.
@@ -211,16 +211,16 @@ pub fn gen_iseq_call(cb: &mut CodeBlock, iseq_call: &IseqCallRef) -> Result<(),
 }
 
 /// Write an entry to the perf map in /tmp
-fn register_with_perf(iseq_name: String, start_ptr: usize, code_size: usize) {
+fn register_with_perf(symbol_name: String, start_ptr: usize, code_size: usize) {
     use std::io::Write;
     let perf_map = format!("/tmp/perf-{}.map", std::process::id());
     let Ok(file) = std::fs::OpenOptions::new().create(true).append(true).open(&perf_map) else {
         debug!("Failed to open perf map file: {perf_map}");
         return;
     };
     let mut file = std::io::BufWriter::new(file);
-    let Ok(_) = writeln!(file, "{start_ptr:#x} {code_size:#x} zjit::{iseq_name}") else {
-        debug!("Failed to write {iseq_name} to perf map file: {perf_map}");
+    let Ok(_) = writeln!(file, "{start_ptr:#x} {code_size:#x} ZJIT: {symbol_name}") else {
+        debug!("Failed to write {symbol_name} to perf map file: {perf_map}");
         return;
     };
 }
@@ -244,11 +244,11 @@ pub fn gen_entry_trampoline(cb: &mut CodeBlock) -> Result<CodePtr, CompileError>
 
     let (code_ptr, gc_offsets) = asm.compile(cb)?;
     assert!(gc_offsets.is_empty());
-    if get_option!(perf) {
+    if get_option!(perf).is_some() {
         let start_ptr = code_ptr.raw_addr(cb);
         let end_ptr = cb.get_write_ptr().raw_addr(cb);
         let code_size = end_ptr - start_ptr;
-        register_with_perf("ZJIT entry trampoline".into(), start_ptr, code_size);
+        register_with_perf("entry trampoline".into(), start_ptr, code_size);
     }
     Ok(code_ptr)
 }
@@ -448,7 +448,23 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, func
                     assert!(insn_idx == block.insns().len() - 1, "Jump must be the last instruction in HIR block");
                 },
                 _ => {
-                    if let Err(last_snapshot) = gen_insn(cb, &mut jit, &mut asm, function, insn_id, &insn) {
+                    // Start a new perf range for the HIR instruction. For now, we do this only for
+                    // non-terminator instructions because LIR blocks must end with a terminator instruction.
+                    let perf_symbol = if get_option!(perf) == Some(PerfMap::HIR) && !insn.is_terminator() {
+                        let insn_name = format!("{insn}").split_whitespace().next().unwrap().to_string();
+                        Some(perf_symbol_range_start(&mut asm, &insn_name))
+                    } else {
+                        None
+                    };
+
+                    let result = gen_insn(cb, &mut jit, &mut asm, function, insn_id, &insn);
+
+                    // Close the current perf range for the HIR instruction.
+                    if let Some(perf_symbol) = &perf_symbol {
+                        perf_symbol_range_end(&mut asm, perf_symbol);
+                    }
+
+                    if let Err(last_snapshot) = result {
                         debug!("ZJIT: gen_function: Failed to compile insn: {insn_id} {insn}. Generating side-exit.");
                         gen_incr_counter(&mut asm, exit_counter_for_unhandled_hir_insn(&insn));
                         gen_side_exit(&mut jit, &mut asm, &SideExitReason::UnhandledHIRInsn(insn_id), &function.frame_state(last_snapshot));
@@ -472,7 +488,7 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, func
     // Generate code if everything can be compiled
     let result = asm.compile(cb);
     if let Ok((start_ptr, _)) = result {
-        if get_option!(perf) {
+        if get_option!(perf) == Some(PerfMap::ISEQ) {
             let start_usize = start_ptr.raw_addr(cb);
             let end_usize = cb.get_write_ptr().raw_addr(cb);
             let code_size = end_usize - start_usize;
@@ -3248,6 +3264,34 @@ impl IseqCall {
     }
 }
 
+type PerfSymbol = Rc<RefCell<Option<(CodePtr, String)>>>;
+
+/// Mark the start of a perf symbol range via pos_marker.
+/// Returns a handle to pass to perf_symbol_range_end.
+pub fn perf_symbol_range_start(asm: &mut Assembler, symbol_name: &str) -> PerfSymbol {
+    let symbol_name = symbol_name.to_string();
+    let perf_symbol: PerfSymbol = Rc::new(RefCell::new(None));
+    let current = perf_symbol.clone();
+    asm.pos_marker(move |start, _| {
+        let mut current = current.borrow_mut();
+        assert!(current.is_none(), "perf symbol range already open");
+        *current = Some((start, symbol_name.clone()));
+    });
+    perf_symbol
+}
+
+/// Mark the end of a perf symbol range via pos_marker.
+pub fn perf_symbol_range_end(asm: &mut Assembler, perf_symbol: &PerfSymbol) {
+    let current = perf_symbol.clone();
+    asm.pos_marker(move |end, cb| {
+        if let Some((start, name)) = current.borrow_mut().take() {
+            let start_addr = start.raw_addr(cb);
+            let code_size = end.raw_addr(cb) - start_addr;
+            register_with_perf(name, start_addr, code_size);
+        }
+    });
+}
+
 #[cfg(test)]
 #[path = "codegen_tests.rs"]
 mod tests;
diff --git a/zjit/src/options.rs b/zjit/src/options.rs
@@ -6,6 +6,15 @@ use crate::cruby::*;
 use crate::stats::Counter;
 use std::collections::HashSet;
 
+/// Type of symbols to dump into /tmp/perf-{pid}.map
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum PerfMap {
+    /// Dump one symbol per ISEQ
+    ISEQ,
+    /// Dump one symbol per HIR instruction
+    HIR,
+}
+
 /// Default --zjit-num-profiles
 const DEFAULT_NUM_PROFILES: NumProfiles = 5;
 pub type NumProfiles = u16;
@@ -89,7 +98,7 @@ pub struct Options {
     pub trace_side_exits_sample_interval: usize,
 
     /// Dump code map to /tmp for performance profilers.
-    pub perf: bool,
+    pub perf: Option<PerfMap>,
 
     /// List of ISEQs that can be compiled, identified by their iseq_get_location()
     pub allowed_iseqs: Option<HashSet<String>>,
@@ -118,7 +127,7 @@ impl Default for Options {
             dump_disasm: None,
             trace_side_exits: None,
             trace_side_exits_sample_interval: 0,
-            perf: false,
+            perf: None,
             allowed_iseqs: None,
             log_compiled_iseqs: None,
         }
@@ -141,7 +150,8 @@ pub const ZJIT_OPTIONS: &[(&str, &str)] = &[
                      "Collect ZJIT stats (=file to write to a file)."),
     ("--zjit-disable",
                      "Disable ZJIT for lazily enabling it with RubyVM::ZJIT.enable."),
-    ("--zjit-perf",  "Dump ISEQ symbols into /tmp/perf-{}.map for Linux perf."),
+    ("--zjit-perf[=iseq|hir]",
+                     "Dump symbols for Linux perf /tmp/perf-{}.map (default: iseq)."),
     ("--zjit-log-compiled-iseqs=path",
                      "Log compiled ISEQs to the file. The file will be truncated."),
     ("--zjit-trace-exits[=counter]",
@@ -452,7 +462,8 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             }
         }
 
-        ("perf", "") => options.perf = true,
+        ("perf", "" | "iseq") => options.perf = Some(PerfMap::ISEQ),
+        ("perf", "hir") => options.perf = Some(PerfMap::HIR),
 
         ("allowed-iseqs", _) if !opt_val.is_empty() => options.allowed_iseqs = Some(parse_jit_list(opt_val)),
         ("log-compiled-iseqs", _) if !opt_val.is_empty() => {
