ZJIT: Add stats for cfuncs that are not optimized

Author: st0012

zjit.c

@@ -211,6 +211,8 @@ rb_zjit_class_has_default_allocator(VALUE klass)
     return alloc == rb_class_allocate_instance;
 }
 
+void rb_zjit_count_unoptimized_cfunc(const struct rb_callable_method_entry_struct *cme);
+
 // Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them.
 VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self, VALUE target_key);

zjit.h

@@ -24,6 +24,7 @@ void rb_zjit_iseq_mark(void *payload);
 void rb_zjit_iseq_update_references(void *payload);
 void rb_zjit_before_ractor_spawn(void);
 void rb_zjit_tracing_invalidate_all(void);
+void rb_zjit_count_unoptimized_cfunc(const struct rb_callable_method_entry_struct *cme);
 #else
 #define rb_zjit_enabled_p false
 static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, bool jit_exception) {}
@@ -33,6 +34,7 @@ static inline void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_ope
 static inline void rb_zjit_cme_invalidate(const rb_callable_method_entry_t *cme) {}
 static inline void rb_zjit_invalidate_no_ep_escape(const rb_iseq_t *iseq) {}
 static inline void rb_zjit_constant_state_changed(ID id) {}
+static inline void rb_zjit_count_unoptimized_cfunc(const struct rb_callable_method_entry_struct *cme) {}
 static inline void rb_zjit_before_ractor_spawn(void) {}
 static inline void rb_zjit_tracing_invalidate_all(void) {}
 #endif // #if USE_ZJIT

zjit.rb

@@ -43,6 +43,7 @@ def stats_string
     print_counters_with_prefix(prefix: 'dynamic_send_type_', prompt: 'dynamic send types', buf:, stats:, limit: 20)
     print_counters_with_prefix(prefix: 'unspecialized_def_type_', prompt: 'send fallback unspecialized def_types', buf:, stats:, limit: 20)
     print_counters_with_prefix(prefix: 'send_fallback_', prompt: 'dynamic send types', buf:, stats:, limit: 20)
+    print_counters_with_prefix(prefix: 'not_optimized_cfuncs', prompt: 'Unoptimized C functions', buf:, stats:, limit: 20)
 
     # Show exit counters, ordered by the typical amount of exits for the prefix at the time
     print_counters_with_prefix(prefix: 'unhandled_yarv_insn_', prompt: 'unhandled YARV insns', buf:, stats:, limit: 20)

zjit/src/codegen.rs

@@ -415,6 +415,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::GetSpecialSymbol { symbol_type, state: _ } => gen_getspecial_symbol(asm, *symbol_type),
         Insn::GetSpecialNumber { nth, state } => gen_getspecial_number(asm, *nth, &function.frame_state(*state)),
         &Insn::IncrCounter(counter) => no_output!(gen_incr_counter(asm, counter)),
+        &Insn::CountUnoptimizedCFunc { cme } => no_output!(gen_count_unoptimized_cfunc(asm, cme)),
         Insn::ObjToString { val, cd, state, .. } => gen_objtostring(jit, asm, opnd!(val), *cd, &function.frame_state(*state)),
         &Insn::CheckInterrupts { state } => no_output!(gen_check_interrupts(jit, asm, &function.frame_state(state))),
         &Insn::HashDup { val, state } => { gen_hash_dup(asm, opnd!(val), &function.frame_state(state)) },
@@ -1575,6 +1576,17 @@ fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd,
     val
 }
 
+/// Generate code that records unoptimized C functions if --zjit-stats is enabled
+fn gen_count_unoptimized_cfunc(asm: &mut Assembler, cme: *const rb_callable_method_entry_struct) {
+    if get_option!(stats) {
+        unsafe extern "C" {
+            fn rb_zjit_count_unoptimized_cfunc(cme: *const rb_callable_method_entry_struct);
+        }
+        asm_comment!(asm, "count unoptimized cfunc");
+        asm_ccall!(asm, rb_zjit_count_unoptimized_cfunc, Opnd::const_ptr(cme));
+    }
+}
+
 /// Generate code that increments a counter if --zjit-stats
 fn gen_incr_counter(asm: &mut Assembler, counter: Counter) {
     if get_option!(stats) {

zjit/src/hir.rs

@@ -697,6 +697,9 @@ pub enum Insn {
     /// Increment a counter in ZJIT stats
     IncrCounter(Counter),
 
+    /// Increment a counter in ZJIT stats for the given C function
+    CountUnoptimizedCFunc { cme: *const rb_callable_method_entry_struct },
+
     /// Equivalent of RUBY_VM_CHECK_INTS. Automatically inserted by the compiler before jumps and
     /// return instructions.
     CheckInterrupts { state: InsnId },
@@ -710,7 +713,7 @@ impl Insn {
             | Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::Return { .. }
             | Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::ArrayExtend { .. }
             | Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetGlobal { .. }
-            | Insn::SetLocal { .. } | Insn::Throw { .. } | Insn::IncrCounter(_)
+            | Insn::SetLocal { .. } | Insn::Throw { .. } | Insn::IncrCounter(_) | Insn::CountUnoptimizedCFunc { .. }
             | Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. } => false,
             _ => true,
         }
@@ -956,6 +959,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             },
+            Insn::CountUnoptimizedCFunc { .. } => write!(f, "CountUnoptimizedCFunc"),
             Insn::Snapshot { state } => write!(f, "Snapshot {}", state.print(self.ptr_map)),
             Insn::Defined { op_type, v, .. } => {
                 // op_type (enum defined_type) printing logic from iseq.c.
@@ -1347,6 +1351,7 @@ impl Function {
                     | GetGlobal {..}
                     | GetLocal {..}
                     | SideExit {..}
+                    | CountUnoptimizedCFunc {..}
                     | IncrCounter(_)) => result.clone(),
             &Snapshot { state: FrameState { iseq, insn_idx, pc, ref stack, ref locals } } =>
                 Snapshot {
@@ -1495,7 +1500,7 @@ impl Function {
             | Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::Return { .. } | Insn::Throw { .. }
             | Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::ArrayExtend { .. }
             | Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetLocal { .. } | Insn::IncrCounter(_)
-            | Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. } =>
+            | Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. } | Insn::CountUnoptimizedCFunc { .. } =>
                 panic!("Cannot infer type of instruction with no output: {}", self.insns[insn.0]),
             Insn::Const { val: Const::Value(val) } => Type::from_value(*val),
             Insn::Const { val: Const::CBool(val) } => Type::from_cbool(*val),
@@ -2083,9 +2088,9 @@ impl Function {
             self_type: Type,
             send: Insn,
             send_insn_id: InsnId,
-        ) -> Result<(), ()> {
+        ) -> Result<(), Option<*const rb_callable_method_entry_struct>> {
             let Insn::SendWithoutBlock { mut recv, cd, mut args, state, .. } = send else {
-                return Err(());
+                return Err(None);
             };
 
             let call_info = unsafe { (*cd).ci };
@@ -2097,20 +2102,20 @@ impl Function {
                 (class, None)
             } else {
                 let iseq_insn_idx = fun.frame_state(state).insn_idx;
-                let Some(recv_type) = fun.profiled_type_of_at(recv, iseq_insn_idx) else { return Err(()) };
+                let Some(recv_type) = fun.profiled_type_of_at(recv, iseq_insn_idx) else { return Err(None) };
                 (recv_type.class(), Some(recv_type))
             };
 
             // Do method lookup
-            let method = unsafe { rb_callable_method_entry(recv_class, method_id) };
+            let method: *const rb_callable_method_entry_struct = unsafe { rb_callable_method_entry(recv_class, method_id) };
             if method.is_null() {
-                return Err(());
+                return Err(None);
             }
 
             // Filter for C methods
             let def_type = unsafe { get_cme_def_type(method) };
             if def_type != VM_METHOD_TYPE_CFUNC {
-                return Err(());
+                return Err(None);
             }
 
             // Find the `argc` (arity) of the C method, which describes the parameters it expects
@@ -2122,15 +2127,18 @@ impl Function {
                     //
                     // Bail on argc mismatch
                     if argc != cfunc_argc as u32 {
-                        return Err(());
+                        return Err(Some(method));
                     }
 
                     // Filter for a leaf and GC free function
                     use crate::cruby_methods::FnProperties;
-                    let Some(FnProperties { leaf: true, no_gc: true, return_type, elidable }) =
-                        ZJITState::get_method_annotations().get_cfunc_properties(method)
-                    else {
-                        return Err(());
+                    let cfunc_properties = ZJITState::get_method_annotations().get_cfunc_properties(method);
+
+                    let (return_type, elidable) = match cfunc_properties {
+                        Some(FnProperties { leaf: true, no_gc: true, return_type, elidable }) => {
+                            (return_type, elidable)
+                        }
+                        _ => return Err(Some(method))
                     };
 
                     let ci_flags = unsafe { vm_ci_flag(call_info) };
@@ -2147,13 +2155,15 @@ impl Function {
                         cfunc_args.append(&mut args);
                         let ccall = fun.push_insn(block, Insn::CCall { cfun, args: cfunc_args, name: method_id, return_type, elidable });
                         fun.make_equal_to(send_insn_id, ccall);
-                        return Ok(());
+                        Ok(())
+                    } else {
+                        Err(Some(method))
                     }
                 }
                 // Variadic method
                 -1 => {
                     if unsafe { rb_zjit_method_tracing_currently_enabled() } {
-                        return Err(());
+                        return Err(None);
                     }
                     // The method gets a pointer to the first argument
                     // func(int argc, VALUE *argv, VALUE recv)
@@ -2185,18 +2195,19 @@ impl Function {
                         });
 
                         fun.make_equal_to(send_insn_id, ccall);
-                        return Ok(());
+                        Ok(())
+                    } else {
+                        // Fall through for complex cases (splat, kwargs, etc.)
+                        Err(Some(method))
                     }
-                    // Fall through for complex cases (splat, kwargs, etc.)
                 }
                 -2 => {
                     // (self, args_ruby_array) parameter form
                     // Falling through for now
+                    Err(Some(method))
                 }
                 _ => unreachable!("unknown cfunc kind: argc={argc}")
             }
-
-            Err(())
         }
 
         for block in self.rpo() {
@@ -2205,8 +2216,12 @@ impl Function {
             for insn_id in old_insns {
                 if let send @ Insn::SendWithoutBlock { recv, .. } = self.find(insn_id) {
                     let recv_type = self.type_of(recv);
-                    if reduce_to_ccall(self, block, recv_type, send, insn_id).is_ok() {
-                        continue;
+                    match reduce_to_ccall(self, block, recv_type, send, insn_id) {
+                        Ok(()) => continue,
+                        Err(Some(cme)) => {
+                            self.push_insn(block, Insn::CountUnoptimizedCFunc { cme });
+                        }
+                        _ => {}
                     }
                 }
                 self.push_insn_id(block, insn_id);
@@ -2349,7 +2364,8 @@ impl Function {
             | &Insn::Param { .. }
             | &Insn::GetLocal { .. }
             | &Insn::PutSpecialObject { .. }
-            | &Insn::IncrCounter(_) =>
+            | &Insn::IncrCounter(_)
+            | &Insn::CountUnoptimizedCFunc { .. } =>
                 {}
             &Insn::PatchPoint { state, .. }
             | &Insn::CheckInterrupts { state }
@@ -8072,6 +8088,7 @@ mod opt_tests {
         bb0(v0:BasicObject):
           v4:Fixnum[1] = Const Value(1)
           v5:Fixnum[0] = Const Value(0)
+          CountUnoptimizedCFunc
           v7:BasicObject = SendWithoutBlock v4, :itself, v5
           CheckInterrupts
           Return v7
@@ -8709,6 +8726,7 @@ mod opt_tests {
         fn test@<compiled>:2:
         bb0(v0:BasicObject):
           v5:HashExact = NewHash
+          CountUnoptimizedCFunc
           v7:BasicObject = SendWithoutBlock v5, :dup
           v9:BasicObject = SendWithoutBlock v7, :freeze
           CheckInterrupts
@@ -8726,6 +8744,7 @@ mod opt_tests {
         bb0(v0:BasicObject):
           v5:HashExact = NewHash
           v6:NilClass = Const Value(nil)
+          CountUnoptimizedCFunc
           v8:BasicObject = SendWithoutBlock v5, :freeze, v6
           CheckInterrupts
           Return v8
@@ -8772,6 +8791,7 @@ mod opt_tests {
         fn test@<compiled>:2:
         bb0(v0:BasicObject):
           v5:ArrayExact = NewArray
+          CountUnoptimizedCFunc
           v7:BasicObject = SendWithoutBlock v5, :dup
           v9:BasicObject = SendWithoutBlock v7, :freeze
           CheckInterrupts
@@ -8789,6 +8809,7 @@ mod opt_tests {
         bb0(v0:BasicObject):
           v5:ArrayExact = NewArray
           v6:NilClass = Const Value(nil)
+          CountUnoptimizedCFunc
           v8:BasicObject = SendWithoutBlock v5, :freeze, v6
           CheckInterrupts
           Return v8
@@ -8836,6 +8857,7 @@ mod opt_tests {
         bb0(v0:BasicObject):
           v4:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
           v6:StringExact = StringCopy v4
+          CountUnoptimizedCFunc
           v8:BasicObject = SendWithoutBlock v6, :dup
           v10:BasicObject = SendWithoutBlock v8, :freeze
           CheckInterrupts
@@ -8854,6 +8876,7 @@ mod opt_tests {
           v4:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
           v6:StringExact = StringCopy v4
           v7:NilClass = Const Value(nil)
+          CountUnoptimizedCFunc
           v9:BasicObject = SendWithoutBlock v6, :freeze, v7
           CheckInterrupts
           Return v9
@@ -8901,6 +8924,7 @@ mod opt_tests {
         bb0(v0:BasicObject):
           v4:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
           v6:StringExact = StringCopy v4
+          CountUnoptimizedCFunc
           v8:BasicObject = SendWithoutBlock v6, :dup
           v10:BasicObject = SendWithoutBlock v8, :-@
           CheckInterrupts
@@ -9000,6 +9024,7 @@ mod opt_tests {
         bb0(v0:BasicObject, v1:BasicObject):
           v5:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
           v17:BasicObject = GuardTypeNot v1, String
+          CountUnoptimizedCFunc
           v18:BasicObject = SendWithoutBlock v1, :to_s
           v9:String = AnyToString v1, str: v18
           v11:StringExact = StringConcat v5, v9

zjit/src/stats.rs

@@ -2,6 +2,8 @@
 
 use std::time::Instant;
 use std::sync::atomic::Ordering;
+use std::sync::{LazyLock, Mutex};
+use std::collections::HashMap;
 
 #[cfg(feature = "stats_allocator")]
 #[path = "../../jit/src/lib.rs"]
@@ -183,6 +185,29 @@ pub(crate) use incr_counter;
 /// The number of side exits from each YARV instruction
 pub type ExitCounters = [u64; VM_INSTRUCTION_SIZE as usize];
 
+static SEND_FALLBACK_REASONS: LazyLock<Mutex<HashMap<String, u64>>> = LazyLock::new(|| Mutex::new(HashMap::new()));
+
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_count_unoptimized_cfunc(cme: *const rb_callable_method_entry_struct) {
+    unsafe extern "C" {
+        fn rb_class2name(klass: VALUE) -> *const std::ffi::c_char;
+        fn rb_id2name(id: ID) -> *const std::ffi::c_char;
+    }
+
+    // Get class name and method name directly from cme
+    let defined_class = unsafe { (*cme).defined_class };
+    let called_id = unsafe { (*cme).called_id };
+
+    // Convert to C strings and then to Rust strings
+    let class_str = unsafe { cstr_to_rust_string(rb_class2name(defined_class)) }.unwrap_or_else(|| "UnknownOwner".to_string());
+    let method_str = unsafe { cstr_to_rust_string(rb_id2name(called_id)) }.unwrap_or_else(|| "UnknownMethod".to_string());
+
+    let reason_str = format!("{}#{}", class_str, method_str);
+
+    let mut map = SEND_FALLBACK_REASONS.lock().unwrap();
+    *map.entry(reason_str).or_insert(0) += 1;
+}
+
 /// Return a raw pointer to the exit counter for a given YARV opcode
 pub fn exit_counter_ptr_for_opcode(opcode: u32) -> *mut u64 {
     let exit_counters = ZJITState::get_exit_counters();
@@ -370,6 +395,13 @@ pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE, target_key: VALUE) ->
         set_stat_f64!(hash, "ratio_in_zjit", 100.0 * zjit_insn_count as f64 / total_insn_count as f64);
     }
 
+    // Set send fallback reasons
+    let fallback_reasons = SEND_FALLBACK_REASONS.lock().unwrap();
+    for (reason, count) in fallback_reasons.iter() {
+        let key_string = "not_optimized_cfuncs".to_owned() + reason;
+        set_stat_usize!(hash, &key_string, *count);
+    }
+
     hash
 }