ZJIT: Add compilation failure stats (ruby#14310)

Author: k0kubun

zjit.rb

@@ -39,19 +39,31 @@ def stats(key = nil)
 
   # Get the summary of ZJIT statistics as a String
   def stats_string
-    buf = +''
+    buf = +"***ZJIT: Printing ZJIT statistics on exit***\n"
     stats = self.stats
 
-    [
+    print_counters_with_prefix(prefix: 'failed_', prompt: 'compilation failure reasons', buf:, stats:)
+    print_counters([
+      :compiled_iseq_count,
+      :compilation_failure,
+
       :compile_time_ns,
       :profile_time_ns,
       :gc_time_ns,
       :invalidation_time_ns,
+
       :total_insns_count,
       :vm_insns_count,
       :zjit_insns_count,
       :ratio_in_zjit,
-    ].each do |key|
+    ], buf:, stats:)
+
+    buf
+  end
+
+  def print_counters(keys, buf:, stats:)
+    left_pad = keys.map(&:size).max + 1
+    keys.each do |key|
       # Some stats like vm_insns_count and ratio_in_zjit are not supported on the release build
       next unless stats.key?(key)
       value = stats[key]
@@ -66,9 +78,16 @@ def stats_string
         value = number_with_delimiter(value)
       end
 
-      buf << "#{'%-18s' % "#{key}:"} #{value}\n"
+      buf << "#{"%-#{left_pad}s" % "#{key}:"} #{value}\n"
+    end
+  end
+
+  def print_counters_with_prefix(buf:, stats:, prefix:, prompt:)
+    keys = stats.keys.select { |key| key.start_with?(prefix) && stats[key] > 0 }
+    unless keys.empty?
+      buf << "#{prompt}:\n"
+      print_counters(keys, buf:, stats:)
     end
-    buf
   end
 
   # Assert that any future ZJIT compilation will return a function pointer

zjit/src/codegen.rs

@@ -6,8 +6,9 @@ use std::slice;
 use crate::asm::Label;
 use crate::backend::current::{Reg, ALLOC_REGS};
 use crate::invariants::{track_bop_assumption, track_cme_assumption, track_single_ractor_assumption, track_stable_constant_names_assumption};
-use crate::gc::{append_gc_offsets, get_or_create_iseq_payload, get_or_create_iseq_payload_ptr, IseqStatus};
+use crate::gc::{append_gc_offsets, get_or_create_iseq_payload, get_or_create_iseq_payload_ptr, IseqPayload, IseqStatus};
 use crate::state::ZJITState;
+use crate::stats::incr_counter;
 use crate::stats::{counter_ptr, with_time_stat, Counter, Counter::compile_time_ns};
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, asm_comment, asm_ccall, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, NATIVE_STACK_PTR, NATIVE_BASE_PTR, SCRATCH_OPND, SP};
@@ -112,11 +113,14 @@ fn gen_iseq_entry_point(iseq: IseqPtr) -> *const u8 {
 /// Compile an entry point for a given ISEQ
 fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<CodePtr> {
     // Compile ISEQ into High-level IR
-    let function = compile_iseq(iseq)?;
+    let Some(function) = compile_iseq(iseq) else {
+        incr_counter!(compilation_failure);
+        return None;
+    };
 
     // Compile the High-level IR
-    let Some((start_ptr, gc_offsets, jit)) = gen_function(cb, iseq, &function) else {
-        debug!("Failed to compile iseq: gen_function failed: {}", iseq_get_location(iseq, 0));
+    let Some(start_ptr) = gen_iseq(cb, iseq, Some(&function)) else {
+        debug!("Failed to compile iseq: gen_iseq failed: {}", iseq_get_location(iseq, 0));
         return None;
     };
 
@@ -126,17 +130,6 @@ fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<CodePt
         return None;
     };
 
-    // Stub callee ISEQs for JIT-to-JIT calls
-    for iseq_call in jit.iseq_calls.iter() {
-        gen_iseq_call(cb, iseq, iseq_call)?;
-    }
-
-    // Remember the block address to reuse it later
-    let payload = get_or_create_iseq_payload(iseq);
-    payload.status = IseqStatus::Compiled(start_ptr);
-    payload.iseq_calls.extend(jit.iseq_calls);
-    append_gc_offsets(iseq, &gc_offsets);
-
     // Return a JIT code address
     Some(entry_ptr)
 }
@@ -194,53 +187,64 @@ fn gen_entry(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function, function_pt
         println!("LIR:\nJIT entry for {}:\n{:?}", iseq_name(iseq), asm);
     }
 
-    let result = asm.compile(cb).map(|(start_ptr, _)| start_ptr);
-    if let Some(start_addr) = result {
-        if get_option!(perf) {
-            let start_ptr = start_addr.raw_ptr(cb) as usize;
-            let end_ptr = cb.get_write_ptr().raw_ptr(cb) as usize;
-            let code_size = end_ptr - start_ptr;
-            let iseq_name = iseq_get_location(iseq, 0);
-            register_with_perf(format!("entry for {iseq_name}"), start_ptr, code_size);
-        }
+    let (code_ptr, gc_offsets) = asm.compile(cb)?;
+    assert!(gc_offsets.is_empty());
+    if get_option!(perf) {
+        let start_ptr = code_ptr.raw_ptr(cb) as usize;
+        let end_ptr = cb.get_write_ptr().raw_ptr(cb) as usize;
+        let code_size = end_ptr - start_ptr;
+        let iseq_name = iseq_get_location(iseq, 0);
+        register_with_perf(format!("entry for {iseq_name}"), start_ptr, code_size);
     }
-    result
+    Some(code_ptr)
 }
 
-/// Compile an ISEQ into machine code
-fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<(CodePtr, Vec<Rc<RefCell<IseqCall>>>)> {
+/// Compile an ISEQ into machine code if not compiled yet
+fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>) -> Option<CodePtr> {
     // Return an existing pointer if it's already compiled
     let payload = get_or_create_iseq_payload(iseq);
     match payload.status {
-        IseqStatus::Compiled(start_ptr) => return Some((start_ptr, vec![])),
+        IseqStatus::Compiled(start_ptr) => return Some(start_ptr),
         IseqStatus::CantCompile => return None,
         IseqStatus::NotCompiled => {},
     }
 
-    // Convert ISEQ into High-level IR and optimize HIR
-    let function = match compile_iseq(iseq) {
+    // Compile the ISEQ
+    let code_ptr = gen_iseq_body(cb, iseq, function, payload);
+    if let Some(start_ptr) = code_ptr {
+        payload.status = IseqStatus::Compiled(start_ptr);
+        incr_counter!(compiled_iseq_count);
+    } else {
+        payload.status = IseqStatus::CantCompile;
+        incr_counter!(compilation_failure);
+    }
+    code_ptr
+}
+
+/// Compile an ISEQ into machine code
+fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>, payload: &mut IseqPayload) -> Option<CodePtr> {
+    // Convert ISEQ into optimized High-level IR if not given
+    let function = match function {
         Some(function) => function,
-        None => {
-            payload.status = IseqStatus::CantCompile;
-            return None;
-        }
+        None => &compile_iseq(iseq)?,
     };
 
     // Compile the High-level IR
-    let result = gen_function(cb, iseq, &function);
-    if let Some((start_ptr, gc_offsets, jit)) = result {
-        payload.status = IseqStatus::Compiled(start_ptr);
-        payload.iseq_calls.extend(jit.iseq_calls.clone());
-        append_gc_offsets(iseq, &gc_offsets);
-        Some((start_ptr, jit.iseq_calls))
-    } else {
-        payload.status = IseqStatus::CantCompile;
-        None
+    let (start_ptr, gc_offsets, iseq_calls) = gen_function(cb, iseq, function)?;
+
+    // Stub callee ISEQs for JIT-to-JIT calls
+    for iseq_call in iseq_calls.iter() {
+        gen_iseq_call(cb, iseq, iseq_call)?;
     }
+
+    // Prepare for GC
+    payload.iseq_calls.extend(iseq_calls.clone());
+    append_gc_offsets(iseq, &gc_offsets);
+    Some(start_ptr)
 }
 
 /// Compile a function
-fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Option<(CodePtr, Vec<CodePtr>, JITState)> {
+fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Option<(CodePtr, Vec<CodePtr>, Vec<Rc<RefCell<IseqCall>>>)> {
     let c_stack_slots = max_num_params(function).saturating_sub(ALLOC_REGS.len());
     let mut jit = JITState::new(iseq, function.num_insns(), function.num_blocks(), c_stack_slots);
     let mut asm = Assembler::new();
@@ -279,6 +283,7 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Optio
             let insn = function.find(insn_id);
             if gen_insn(cb, &mut jit, &mut asm, function, insn_id, &insn).is_none() {
                 debug!("Failed to compile insn: {insn_id} {insn}");
+                incr_counter!(failed_gen_insn);
                 return None;
             }
         }
@@ -291,8 +296,8 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Optio
     }
 
     // Generate code if everything can be compiled
-    let result = asm.compile(cb).map(|(start_ptr, gc_offsets)| (start_ptr, gc_offsets, jit));
-    if let Some((start_ptr, _, _)) = result {
+    let result = asm.compile(cb);
+    if let Some((start_ptr, _)) = result {
         if get_option!(perf) {
             let start_usize = start_ptr.raw_ptr(cb) as usize;
             let end_usize = cb.get_write_ptr().raw_ptr(cb) as usize;
@@ -304,8 +309,10 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Optio
             let iseq_name = iseq_get_location(iseq, 0);
             ZJITState::log_compile(iseq_name);
         }
+    } else {
+        incr_counter!(failed_asm_compile);
     }
-    result
+    result.map(|(start_ptr, gc_offsets)| (start_ptr, gc_offsets, jit.iseq_calls))
 }
 
 /// Compile an instruction
@@ -410,6 +417,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         | Insn::Const { .. }
         => {
             debug!("ZJIT: gen_function: unexpected insn {insn}");
+            incr_counter!(failed_gen_insn_unexpected);
             return None;
         }
     };
@@ -1372,6 +1380,7 @@ fn compile_iseq(iseq: IseqPtr) -> Option<Function> {
         Err(err) => {
             let name = crate::cruby::iseq_get_location(iseq, 0);
             debug!("ZJIT: iseq_to_hir: {err:?}: {name}");
+            incr_counter!(failed_hir_compile);
             return None;
         }
     };
@@ -1382,6 +1391,7 @@ fn compile_iseq(iseq: IseqPtr) -> Option<Function> {
     #[cfg(debug_assertions)]
     if let Err(err) = function.validate() {
         debug!("ZJIT: compile_iseq: {err:?}");
+        incr_counter!(failed_hir_optimize);
         return None;
     }
     Some(function)
@@ -1516,16 +1526,11 @@ c_callable! {
 /// Compile an ISEQ for a function stub
 fn function_stub_hit_body(cb: &mut CodeBlock, iseq_call: &Rc<RefCell<IseqCall>>) -> Option<CodePtr> {
     // Compile the stubbed ISEQ
-    let Some((code_ptr, iseq_calls)) = gen_iseq(cb, iseq_call.borrow().iseq) else {
+    let Some(code_ptr) = gen_iseq(cb, iseq_call.borrow().iseq, None) else {
         debug!("Failed to compile iseq: gen_iseq failed: {}", iseq_get_location(iseq_call.borrow().iseq, 0));
         return None;
     };
 
-    // Stub callee ISEQs for JIT-to-JIT calls
-    for callee_iseq_call in iseq_calls.iter() {
-        gen_iseq_call(cb, iseq_call.borrow().iseq, callee_iseq_call)?;
-    }
-
     // Update the stub to call the code pointer
     let code_addr = code_ptr.raw_ptr(cb);
     let iseq = iseq_call.borrow().iseq;

zjit/src/hir.rs

@@ -4,14 +4,15 @@
 #![allow(non_upper_case_globals)]
 
 use crate::{
-    cast::IntoUsize, cruby::*, gc::{get_or_create_iseq_payload, IseqPayload}, options::{get_option, DumpHIR}, state::ZJITState, stats::Counter
+    cast::IntoUsize, cruby::*, gc::{get_or_create_iseq_payload, IseqPayload}, options::{get_option, DumpHIR}, state::ZJITState
 };
 use std::{
     cell::RefCell, collections::{HashMap, HashSet, VecDeque}, ffi::{c_int, c_void, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter
 };
 use crate::hir_type::{Type, types};
 use crate::bitset::BitSet;
 use crate::profile::{TypeDistributionSummary, ProfiledType};
+use crate::stats::{incr_counter, Counter};
 
 /// An index of an [`Insn`] in a [`Function`]. This is a popular
 /// type since this effectively acts as a pointer to an [`Insn`].
@@ -3494,6 +3495,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
 
     fun.profiles = Some(profiles);
     if let Err(e) = fun.validate() {
+        incr_counter!(failed_hir_compile_validate);
         return Err(ParseError::Validation(e));
     }
     Ok(fun)