ZJIT: Add Perfetto duration events for compilation phases

tekknolagi · tekknolagi · commit 52169ef4787b · 2026-04-06T15:18:52.000-04:00
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -17,7 +17,7 @@ use crate::gc::append_gc_offsets;
 use crate::payload::{IseqCodePtrs, IseqStatus, IseqVersion, IseqVersionRef, JITFrame, get_or_create_iseq_payload};
 use crate::state::ZJITState;
 use crate::stats::{CompileError, exit_counter_for_compile_error, exit_counter_for_unhandled_hir_insn, incr_counter, incr_counter_by, send_fallback_counter, send_fallback_counter_for_method_type, send_fallback_counter_for_super_method_type, send_fallback_counter_ptr_for_opcode, send_without_block_fallback_counter_for_method_type, send_without_block_fallback_counter_for_optimized_method_type};
-use crate::stats::{counter_ptr, with_time_stat, Counter, Counter::{compile_time_ns, exit_compile_error}};
+use crate::stats::{counter_ptr, with_time_stat, trace_compile_phase, Counter, Counter::{compile_time_ns, exit_compile_error}};
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, Assembler, C_ARG_OPNDS, C_RET_OPND, CFP, EC, NATIVE_STACK_PTR, Opnd, SP, SideExit, SideExitRecompile, Target, asm_ccall, asm_comment};
 use crate::hir::{iseq_to_hir, BlockId, Invariant, RangeType, SideExitReason::{self, *}, SpecialBackrefSymbol, SpecialObjectType};
@@ -200,17 +200,20 @@ fn gen_iseq_entry_point(cb: &mut CodeBlock, iseq: IseqPtr, jit_exception: bool)
         return Err(CompileError::ExceptionHandler);
     }
 
-    // Compile ISEQ into High-level IR
-    let function = crate::stats::with_time_stat(Counter::compile_hir_time_ns, || compile_iseq(iseq).inspect_err(|_| {
-        incr_counter!(failed_iseq_count);
-    }))?;
+    let iseq_name = iseq_get_location(iseq, 0);
+    trace_compile_phase(&iseq_name, || {
+        // Compile ISEQ into High-level IR
+        let function = crate::stats::with_time_stat(Counter::compile_hir_time_ns, || compile_iseq(iseq).inspect_err(|_| {
+            incr_counter!(failed_iseq_count);
+        }))?;
 
-    // Compile the High-level IR
-    let IseqCodePtrs { start_ptr, .. } = gen_iseq(cb, iseq, Some(&function)).inspect_err(|err| {
-        debug!("{err:?}: gen_iseq failed: {}", iseq_get_location(iseq, 0));
-    })?;
+        // Compile the High-level IR
+        let IseqCodePtrs { start_ptr, .. } = gen_iseq(cb, iseq, Some(&function)).inspect_err(|err| {
+            debug!("{err:?}: gen_iseq failed: {}", iseq_get_location(iseq, 0));
+        })?;
 
-    Ok(start_ptr)
+        Ok(start_ptr)
+    })
 }
 
 /// Invalidate an ISEQ version and allow it to be recompiled on the next call.
@@ -312,9 +315,14 @@ fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>) -> R
         return Err(CompileError::IseqVersionLimitReached);
     }
 
-    // Compile the ISEQ
+    // Compile the ISEQ. When function is None, this is a lazy compile
+    // from a stub hit — wrap in a trace event covering the full compile.
     let mut version = IseqVersion::new(iseq);
-    let code_ptrs = gen_iseq_body(cb, iseq, version, function);
+    let code_ptrs = if function.is_none() {
+        trace_compile_phase(&iseq_get_location(iseq, 0), || gen_iseq_body(cb, iseq, version, function))
+    } else {
+        gen_iseq_body(cb, iseq, version, function)
+    };
     match &code_ptrs {
         Ok(code_ptrs) => {
             unsafe { version.as_mut() }.status = IseqStatus::Compiled(code_ptrs.clone());
@@ -344,7 +352,9 @@ fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, mut version: IseqVersionRef,
 
     // Compile the High-level IR
     let (iseq_code_ptrs, gc_offsets, iseq_calls) =
-        crate::stats::with_time_stat(Counter::compile_lir_time_ns, || gen_function(cb, iseq, version, function))?;
+        trace_compile_phase("codegen", ||
+            crate::stats::with_time_stat(Counter::compile_lir_time_ns, || gen_function(cb, iseq, version, function))
+        )?;
 
     // Stub callee ISEQs for JIT-to-JIT calls
     for iseq_call in iseq_calls.iter() {
@@ -2955,7 +2965,9 @@ fn compile_iseq(iseq: IseqPtr) -> Result<Function, CompileError> {
         return Err(CompileError::IseqStackTooLarge);
     }
 
-    let hir = crate::stats::with_time_stat(Counter::compile_hir_build_time_ns, || iseq_to_hir(iseq));
+    let hir = trace_compile_phase("build_hir", ||
+        crate::stats::with_time_stat(Counter::compile_hir_build_time_ns, || iseq_to_hir(iseq))
+    );
     let mut function = match hir {
         Ok(function) => function,
         Err(err) => {
@@ -2964,7 +2976,7 @@ fn compile_iseq(iseq: IseqPtr) -> Result<Function, CompileError> {
         }
     };
     if !get_option!(disable_hir_opt) {
-        function.optimize();
+        trace_compile_phase("optimize", || function.optimize());
     }
     function.dump_hir();
     Ok(function)
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -5930,7 +5930,9 @@ impl Function {
         macro_rules! run_pass {
             ($name:ident) => {
                 let counter = counter_for!($name);
-                crate::stats::with_time_stat(counter, || self.$name());
+                crate::stats::trace_compile_phase(stringify!($name), ||
+                    crate::stats::with_time_stat(counter, || self.$name())
+                );
                 #[cfg(debug_assertions)] self.assert_validates();
                 if should_dump {
                     passes.push(
diff --git a/zjit/src/options.rs b/zjit/src/options.rs
@@ -97,6 +97,9 @@ pub struct Options {
     /// Frequency of tracing side exits.
     pub trace_side_exits_sample_interval: usize,
 
+    /// Trace compilation phases as Perfetto duration events.
+    pub trace_compiles: bool,
+
     /// Dump code map to /tmp for performance profilers.
     pub perf: Option<PerfMap>,
 
@@ -130,6 +133,7 @@ impl Default for Options {
             dump_disasm: None,
             trace_side_exits: None,
             trace_side_exits_sample_interval: 0,
+            trace_compiles: false,
             perf: None,
             allowed_iseqs: None,
             log_compiled_iseqs: None,
@@ -161,7 +165,9 @@ pub const ZJIT_OPTIONS: &[(&str, &str)] = &[
     ("--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`.")
+                     "Frequency at which to record side exits. Must be `usize`."),
+    ("--zjit-trace-compiles",
+                     "Record compilation phases as Perfetto trace events.")
 ];
 
 #[derive(Copy, Clone, Debug)]
@@ -388,6 +394,8 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             options.trace_side_exits_sample_interval = sample_interval.parse::<usize>().ok()?;
         }
 
+        ("trace-compiles", "") => options.trace_compiles = true,
+
         ("debug", "") => options.debug = true,
 
         ("disable", "") => options.disable = true,
diff --git a/zjit/src/state.rs b/zjit/src/state.rs
@@ -128,7 +128,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 perfetto_tracer = if get_option!(trace_side_exits).is_some() {
+        let perfetto_tracer = if get_option!(trace_side_exits).is_some() || get_option!(trace_compiles) {
             Some(PerfettoTracer::new())
         } else {
             None
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
@@ -993,6 +993,23 @@ pub fn zjit_alloc_bytes() -> usize {
     jit::GLOBAL_ALLOCATOR.alloc_size.load(Ordering::SeqCst)
 }
 
+/// Record a Perfetto duration-complete event spanning the execution of `func`.
+/// Nested calls produce nested events on the same thread in the trace viewer.
+pub fn trace_compile_phase<F, R>(name: &str, func: F) -> R where F: FnOnce() -> R {
+    if !get_option!(trace_compiles) {
+        return func();
+    }
+    let start_ns = ZJITState::get_tracer().map(|t| t.elapsed_ns());
+    let result = func();
+    if let Some(start_ns) = start_ns {
+        if let Some(tracer) = ZJITState::get_tracer() {
+            let end_ns = tracer.elapsed_ns();
+            tracer.write_duration_complete("compile", name, start_ns, end_ns);
+        }
+    }
+    result
+}
+
 /// Fuchsia Trace Format (FXT) binary writer for --zjit-trace-exits.
 /// Produces .fxt files that can be opened directly in Perfetto UI.
 /// Uses the string table for deduplication of repeated reason/frame strings.
@@ -1060,6 +1077,7 @@ impl PerfettoTracer {
 
         // Pre-intern common strings
         tracer.intern_string("side_exit");
+        tracer.intern_string("compile");
         // Pre-intern argument names "0".."14" for per-frame arguments
         for i in 0..15u32 {
             tracer.intern_string(&i.to_string());
@@ -1107,6 +1125,36 @@ impl PerfettoTracer {
         idx
     }
 
+    /// Return nanoseconds elapsed since tracer creation.
+    pub fn elapsed_ns(&self) -> u64 {
+        self.start_time.elapsed().as_nanos() as u64
+    }
+
+    /// Write a Duration Complete event (FXT event type 4).
+    /// Records an event spanning from `start_ns` to `end_ns` (both in ticks
+    /// relative to `self.start_time`).
+    pub fn write_duration_complete(&mut self, category: &str, name: &str, start_ns: u64, end_ns: u64) {
+        let category_ref = self.intern_string(category);
+        let name_ref = self.intern_string(name);
+
+        // Duration Complete: header + start_ts + end_ts = 3 words, no args
+        let event_words: u64 = 3;
+        let header: u64 = 4u64                           // record type = event
+            | (event_words << 4)                          // record size
+            | (4u64 << 16)                                // event type = duration complete
+            | (1u64 << 24)                                // thread_ref = 1
+            | ((category_ref as u64) << 32)
+            | ((name_ref as u64) << 48);
+        self.write_word(header);
+        self.write_word(start_ns);
+        self.write_word(end_ns);
+
+        self.event_count += 1;
+
+        use std::io::Write;
+        let _ = self.writer.flush();
+    }
+
     pub fn write_event(&mut self, reason: &str, frames: &[String]) {
         let ts_nanos = self.start_time.elapsed().as_nanos() as u64;
 
