ZJIT: Add Perfetto trace sub-phases for LIR codegen and stubs

Author: tekknolagi

zjit/src/backend/arm64/mod.rs

@@ -4,7 +4,7 @@ use crate::codegen::split_patch_point;
 use crate::cruby::*;
 use crate::backend::lir::*;
 use crate::options::asm_dump;
-use crate::stats::CompileError;
+use crate::stats::{CompileError, trace_compile_phase};
 use crate::virtualmem::CodePtr;
 use crate::cast::*;
 
@@ -1610,14 +1610,15 @@ impl Assembler {
         let use_scratch_reg = !self.accept_scratch_reg;
         asm_dump!(self, init);
 
-        let mut asm = self.arm64_split();
+        let mut asm = trace_compile_phase("split", || self.arm64_split());
 
         asm_dump!(asm, split);
 
-        asm.number_instructions(0);
+        let (intervals, assignments, num_stack_slots) = trace_compile_phase("regalloc", || {
+        trace_compile_phase("number_instructions", || asm.number_instructions(0));
 
-        let live_in = asm.analyze_liveness();
-        let intervals = asm.build_intervals(live_in);
+        let live_in = trace_compile_phase("analyze_liveness", || asm.analyze_liveness());
+        let intervals = trace_compile_phase("build_intervals", || asm.build_intervals(live_in));
 
         // Dump live intervals if requested
         if let Some(crate::options::Options { dump_lir: Some(dump_lirs), .. }) = unsafe { crate::options::OPTIONS.as_ref() } {
@@ -1627,7 +1628,9 @@ impl Assembler {
         }
 
         let preferred_registers = asm.preferred_register_assignments(&intervals);
-        let (assignments, num_stack_slots) = asm.linear_scan(intervals.clone(), regs.len(), &preferred_registers);
+        let (assignments, num_stack_slots) = trace_compile_phase("linear_scan", || asm.linear_scan(intervals.clone(), regs.len(), &preferred_registers));
+        (intervals, assignments, num_stack_slots)
+        });
 
         let total_stack_slots = asm.stack_base_idx + num_stack_slots;
         if total_stack_slots > Self::MAX_FRAME_STACK_SLOTS {
@@ -1665,36 +1668,41 @@ impl Assembler {
             }
         }
 
-        asm.handle_caller_saved_regs(&intervals, &assignments, &C_ARG_REGREGS);
-        asm.resolve_ssa(&intervals, &assignments);
+        trace_compile_phase("resolve_ssa", || {
+            asm.handle_caller_saved_regs(&intervals, &assignments, &C_ARG_REGREGS);
+            asm.resolve_ssa(&intervals, &assignments);
+        });
         asm_dump!(asm, alloc_regs);
 
         // We are moved out of SSA after resolve_ssa
 
         // We put compile_exits after alloc_regs to avoid extending live ranges for VRegs spilled on side exits.
         // Exit code is compiled into a separate list of instructions that we append
         // to the last reachable block before scratch_split, so it gets linearized and split.
-        let exit_insns = asm.compile_exits();
-        asm_dump!(asm, compile_exits);
-
-        // Append exit instructions to the last reachable block so they are
-        // included in linearize_instructions and processed by scratch_split.
-        if let Some(&last_block) = asm.block_order().last() {
-            for insn in exit_insns {
-                asm.basic_blocks[last_block.0].insns.push(insn);
-                asm.basic_blocks[last_block.0].insn_ids.push(None);
+        trace_compile_phase("compile_exits", || {
+            let exit_insns = asm.compile_exits();
+
+            // Append exit instructions to the last reachable block so they are
+            // included in linearize_instructions and processed by scratch_split.
+            if let Some(&last_block) = asm.block_order().last() {
+                for insn in exit_insns {
+                    asm.basic_blocks[last_block.0].insns.push(insn);
+                    asm.basic_blocks[last_block.0].insn_ids.push(None);
+                }
             }
-        }
+        });
+        asm_dump!(asm, compile_exits);
 
         if use_scratch_reg {
-            asm = asm.arm64_scratch_split();
+            asm = trace_compile_phase("scratch_split", || asm.arm64_scratch_split());
             asm_dump!(asm, scratch_split);
         } else {
             // For trampolines that use scratch registers, resolve ParallelMov without scratch_reg.
-            asm = asm.resolve_parallel_mov_pass();
+            asm = trace_compile_phase("resolve_parallel_mov", || asm.resolve_parallel_mov_pass());
             asm_dump!(asm, resolve_parallel_mov);
         }
 
+        trace_compile_phase("emit", || {
         // Create label instances in the code block
         for (idx, name) in asm.label_names.iter().enumerate() {
             let label = cb.new_label(name.to_string());
@@ -1711,6 +1719,7 @@ impl Assembler {
         unsafe { rb_jit_icache_invalidate(start_ptr.raw_ptr(cb) as _, cb.get_write_ptr().raw_ptr(cb) as _) };
 
         Ok((start_ptr, gc_offsets))
+        })
     }
 }
 

zjit/src/backend/x86_64/mod.rs

@@ -3,7 +3,7 @@ use std::mem;
 use crate::asm::*;
 use crate::asm::x86_64::*;
 use crate::codegen::split_patch_point;
-use crate::stats::CompileError;
+use crate::stats::{CompileError, trace_compile_phase};
 use crate::virtualmem::CodePtr;
 use crate::cruby::*;
 use crate::backend::lir::*;
@@ -1143,14 +1143,15 @@ impl Assembler {
         let use_scratch_regs = !self.accept_scratch_reg;
         asm_dump!(self, init);
 
-        let mut asm = self.x86_split();
+        let mut asm = trace_compile_phase("split", || self.x86_split());
 
         asm_dump!(asm, split);
 
-        asm.number_instructions(0);
+        let (intervals, assignments, num_stack_slots) = trace_compile_phase("regalloc", || {
+        trace_compile_phase("number_instructions", || asm.number_instructions(0));
 
-        let live_in = asm.analyze_liveness();
-        let intervals = asm.build_intervals(live_in);
+        let live_in = trace_compile_phase("analyze_liveness", || asm.analyze_liveness());
+        let intervals = trace_compile_phase("build_intervals", || asm.build_intervals(live_in));
 
         // Dump live intervals if requested
         if let Some(crate::options::Options { dump_lir: Some(dump_lirs), .. }) = unsafe { crate::options::OPTIONS.as_ref() } {
@@ -1160,7 +1161,9 @@ impl Assembler {
         }
 
         let preferred_registers = asm.preferred_register_assignments(&intervals);
-        let (assignments, num_stack_slots) = asm.linear_scan(intervals.clone(), regs.len(), &preferred_registers);
+        let (assignments, num_stack_slots) = trace_compile_phase("linear_scan", || asm.linear_scan(intervals.clone(), regs.len(), &preferred_registers));
+        (intervals, assignments, num_stack_slots)
+        });
 
         let total_stack_slots = asm.stack_base_idx + num_stack_slots;
         if total_stack_slots > Self::MAX_FRAME_STACK_SLOTS {
@@ -1198,36 +1201,41 @@ impl Assembler {
             }
         }
 
-        asm.handle_caller_saved_regs(&intervals, &assignments, &C_ARG_REGREGS);
-        asm.resolve_ssa(&intervals, &assignments);
+        trace_compile_phase("resolve_ssa", || {
+            asm.handle_caller_saved_regs(&intervals, &assignments, &C_ARG_REGREGS);
+            asm.resolve_ssa(&intervals, &assignments);
+        });
         asm_dump!(asm, alloc_regs);
 
         // We are moved out of SSA after resolve_ssa
 
         // We put compile_exits after alloc_regs to avoid extending live ranges for VRegs spilled on side exits.
         // Exit code is compiled into a separate list of instructions that we append
         // to the last reachable block before scratch_split, so it gets linearized and split.
-        let exit_insns = asm.compile_exits();
-        asm_dump!(asm, compile_exits);
-
-        // Append exit instructions to the last reachable block so they are
-        // included in linearize_instructions and processed by scratch_split.
-        if let Some(&last_block) = asm.block_order().last() {
-            for insn in exit_insns {
-                asm.basic_blocks[last_block.0].insns.push(insn);
-                asm.basic_blocks[last_block.0].insn_ids.push(None);
+        trace_compile_phase("compile_exits", || {
+            let exit_insns = asm.compile_exits();
+
+            // Append exit instructions to the last reachable block so they are
+            // included in linearize_instructions and processed by scratch_split.
+            if let Some(&last_block) = asm.block_order().last() {
+                for insn in exit_insns {
+                    asm.basic_blocks[last_block.0].insns.push(insn);
+                    asm.basic_blocks[last_block.0].insn_ids.push(None);
+                }
             }
-        }
+        });
+        asm_dump!(asm, compile_exits);
 
         if use_scratch_regs {
-            asm = asm.x86_scratch_split();
+            asm = trace_compile_phase("scratch_split", || asm.x86_scratch_split());
             asm_dump!(asm, scratch_split);
         } else {
             // For trampolines that use scratch registers, resolve ParallelMov without scratch_reg.
-            asm = asm.resolve_parallel_mov_pass();
+            asm = trace_compile_phase("resolve_parallel_mov", || asm.resolve_parallel_mov_pass());
             asm_dump!(asm, resolve_parallel_mov);
         }
 
+        trace_compile_phase("emit", || {
         // Create label instances in the code block
         for (idx, name) in asm.label_names.iter().enumerate() {
             let label = cb.new_label(name.to_string());
@@ -1240,6 +1248,7 @@ impl Assembler {
 
         cb.link_labels().or(Err(CompileError::LabelLinkingFailure))?;
         Ok((start_ptr, gc_offsets))
+        })
     }
 }
 

zjit/src/codegen.rs

@@ -242,6 +242,7 @@ pub fn invalidate_iseq_version(cb: &mut CodeBlock, iseq: IseqPtr, version: &mut
 
 /// Stub a branch for a JIT-to-JIT call
 pub fn gen_iseq_call(cb: &mut CodeBlock, iseq_call: &IseqCallRef) -> Result<(), CompileError> {
+    trace_compile_phase("compile_stub", || {
     // Compile a function stub
     let stub_ptr = gen_function_stub(cb, iseq_call.clone()).inspect_err(|err| {
         debug!("{err:?}: gen_function_stub failed: {}", iseq_get_location(iseq_call.iseq.get(), 0));
@@ -255,6 +256,7 @@ pub fn gen_iseq_call(cb: &mut CodeBlock, iseq_call: &IseqCallRef) -> Result<(),
         asm.ccall_into(C_RET_OPND, stub_addr, vec![]);
     });
     Ok(())
+    })
 }
 
 /// Write an entry to the perf map in /tmp
@@ -352,14 +354,20 @@ 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) =
-        trace_compile_phase("codegen", ||
-            crate::stats::with_time_stat(Counter::compile_lir_time_ns, || gen_function(cb, iseq, version, function))
-        )?;
+        trace_compile_phase("codegen", || {
+            let (iseq_code_ptrs, gc_offsets, iseq_calls) =
+                crate::stats::with_time_stat(Counter::compile_lir_time_ns, || gen_function(cb, iseq, version, function))?;
+
+            // Stub callee ISEQs for JIT-to-JIT calls
+            trace_compile_phase("generate_jit_jit_stubs", || {
+                for iseq_call in iseq_calls.iter() {
+                    gen_iseq_call(cb, iseq_call)?;
+                }
+                Ok::<(), CompileError>(())
+            })?;
 
-    // Stub callee ISEQs for JIT-to-JIT calls
-    for iseq_call in iseq_calls.iter() {
-        gen_iseq_call(cb, iseq_call)?;
-    }
+            Ok((iseq_code_ptrs, gc_offsets, iseq_calls))
+        })?;
 
     // Prepare for GC
     unsafe { version.as_mut() }.outgoing.extend(iseq_calls);
@@ -369,6 +377,7 @@ fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, mut version: IseqVersionRef,
 
 /// Compile a function
 fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, function: &Function) -> Result<(IseqCodePtrs, Vec<CodePtr>, Vec<IseqCallRef>), CompileError> {
+    let (mut jit, asm) = trace_compile_phase("lowering", || {
     let num_spilled_params = max_num_params(function).saturating_sub(ALLOC_REGS.len());
     let mut jit = JITState::new(iseq, version, function.num_insns(), function.num_blocks());
     let mut asm = Assembler::new_with_stack_slots(num_spilled_params);
@@ -546,6 +555,9 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, func
     // Validate CFG invariants after HIR to LIR lowering
     asm.validate_jump_positions();
 
+    (jit, asm)
+    });
+
     // Generate code if everything can be compiled
     let result = asm.compile(cb);
     if let Ok((start_ptr, _)) = result {
@@ -3270,9 +3282,11 @@ fn function_stub_hit_body(cb: &mut CodeBlock, iseq_call: &IseqCallRef) -> Result
     let jit_entry_ptr = jit_entry_ptrs[iseq_call.jit_entry_idx.to_usize()];
     let code_addr = jit_entry_ptr.raw_ptr(cb);
     let iseq = iseq_call.iseq.get();
-    iseq_call.regenerate(cb, |asm| {
-        asm_comment!(asm, "call compiled function: {}", iseq_get_location(iseq, 0));
-        asm.ccall_into(C_RET_OPND, code_addr, vec![]);
+    trace_compile_phase("compile_stub", || {
+        iseq_call.regenerate(cb, |asm| {
+            asm_comment!(asm, "call compiled function: {}", iseq_get_location(iseq, 0));
+            asm.ccall_into(C_RET_OPND, code_addr, vec![]);
+        });
     });
 
     Ok(jit_entry_ptr)

zjit/src/hir.rs

@@ -5933,7 +5933,7 @@ impl Function {
                 crate::stats::trace_compile_phase(stringify!($name), ||
                     crate::stats::with_time_stat(counter, || self.$name())
                 );
-                #[cfg(debug_assertions)] self.assert_validates();
+                #[cfg(debug_assertions)] crate::stats::trace_compile_phase("validate", || self.assert_validates());
                 if should_dump {
                     passes.push(
                         self.to_iongraph_pass(stringify!($name))
@@ -8414,7 +8414,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
     }
 
     fun.profiles = Some(profiles);
-    if let Err(err) = fun.validate() {
+    if let Err(err) = crate::stats::trace_compile_phase("validate", || fun.validate()) {
         debug!("ZJIT: {err:?}: Initial HIR:\n{}", FunctionPrinter::without_snapshot(&fun));
         return Err(ParseError::Validation(err));
     }

zjit/src/state.rs

@@ -300,6 +300,7 @@ impl ZJITState {
 
     /// Get a mutable reference to the Perfetto tracer
     pub fn get_tracer() -> Option<&'static mut PerfettoTracer> {
+        if !ZJITState::has_instance() { return None; }
         ZJITState::get_instance().perfetto_tracer.as_mut()
     }
 }