ZJIT: Add stats for cfuncs that are not optimized

Author: st0012

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

@@ -419,6 +419,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 { signature, counter_ptr } => no_output!(gen_count_unoptimized_cfunc(asm, signature, *counter_ptr)),
         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)) },
@@ -1528,6 +1529,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, signature: &str, counter_ptr: *mut u64) {
+    if get_option!(stats) {
+        unsafe extern "C" {
+            fn rb_zjit_count_unoptimized_cfunc(counter_ptr: *mut u64);
+        }
+        asm_comment!(asm, "count unoptimized cfunc: {}", signature);
+        asm_ccall!(asm, rb_zjit_count_unoptimized_cfunc, Opnd::const_ptr(counter_ptr as *const u8));
+    }
+}
+
 /// 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

@@ -14,7 +14,7 @@ use std::{
 use crate::hir_type::{Type, types};
 use crate::bitset::BitSet;
 use crate::profile::{TypeDistributionSummary, ProfiledType};
-use crate::stats::Counter;
+use crate::stats::{Counter, get_or_create_unoptimized_cfunc_counter_ptr};
 
 /// An index of an [`Insn`] in a [`Function`]. This is a popular
 /// type since this effectively acts as a pointer to an [`Insn`].
@@ -711,6 +711,9 @@ pub enum Insn {
     /// Increment a counter in ZJIT stats
     IncrCounter(Counter),
 
+    /// Increment a counter in ZJIT stats for the given unoptimized C function
+    CountUnoptimizedCFunc { signature: String, counter_ptr: *mut u64 },
+
     /// Equivalent of RUBY_VM_CHECK_INTS. Automatically inserted by the compiler before jumps and
     /// return instructions.
     CheckInterrupts { state: InsnId },
@@ -724,7 +727,7 @@ impl Insn {
             | Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::EntryPoint { .. } | 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,
         }
@@ -978,6 +981,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             },
+            Insn::CountUnoptimizedCFunc { signature, .. } => write!(f, "CountUnoptimizedCFunc {}", signature),
             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.
@@ -1385,6 +1389,7 @@ impl Function {
                     | EntryPoint {..}
                     | LoadPC
                     | LoadSelf
+                    | CountUnoptimizedCFunc {..}
                     | IncrCounter(_)) => result.clone(),
             &Snapshot { state: FrameState { iseq, insn_idx, pc, ref stack, ref locals } } =>
                 Snapshot {
@@ -1534,7 +1539,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),
@@ -2154,9 +2159,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 };
@@ -2168,20 +2173,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
@@ -2193,15 +2198,15 @@ 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(());
+                        return Err(Some(method));
                     };
 
                     let ci_flags = unsafe { vm_ci_flag(call_info) };
@@ -2218,13 +2223,13 @@ 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(());
+                        return Ok(())
                     }
                 }
                 // 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)
@@ -2256,8 +2261,9 @@ impl Function {
                         });
 
                         fun.make_equal_to(send_insn_id, ccall);
-                        return Ok(());
+                        return Ok(())
                     }
+
                     // Fall through for complex cases (splat, kwargs, etc.)
                 }
                 -2 => {
@@ -2267,7 +2273,7 @@ impl Function {
                 _ => unreachable!("unknown cfunc kind: argc={argc}")
             }
 
-            Err(())
+            Err(Some(method))
         }
 
         for block in self.rpo() {
@@ -2276,8 +2282,19 @@ 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)) => {
+                            let owner = unsafe { (*cme).owner };
+                            let called_id = unsafe { (*cme).called_id };
+                            let class_name = get_class_name(owner);
+                            let method_name = called_id.contents_lossy();
+                            let signature = format!("{}#{}", class_name, method_name);
+                            let counter_ptr = get_or_create_unoptimized_cfunc_counter_ptr(signature.clone());
+
+                            self.push_insn(block, Insn::CountUnoptimizedCFunc { signature, counter_ptr });
+                        }
+                        _ => {}
                     }
                 }
                 self.push_insn_id(block, insn_id);
@@ -2423,7 +2440,8 @@ impl Function {
             | &Insn::LoadSelf
             | &Insn::GetLocal { .. }
             | &Insn::PutSpecialObject { .. }
-            | &Insn::IncrCounter(_) =>
+            | &Insn::IncrCounter(_)
+            | &Insn::CountUnoptimizedCFunc { .. } =>
                 {}
             &Insn::PatchPoint { state, .. }
             | &Insn::CheckInterrupts { state }
@@ -9841,6 +9859,7 @@ mod opt_tests {
         bb2(v6:BasicObject):
           v10:Fixnum[1] = Const Value(1)
           v11:Fixnum[0] = Const Value(0)
+          CountUnoptimizedCFunc Kernel#itself
           v13:BasicObject = SendWithoutBlock v10, :itself, v11
           CheckInterrupts
           Return v13
@@ -10720,6 +10739,7 @@ mod opt_tests {
           Jump bb2(v4)
         bb2(v6:BasicObject):
           v11:HashExact = NewHash
+          CountUnoptimizedCFunc Kernel#dup
           v13:BasicObject = SendWithoutBlock v11, :dup
           v15:BasicObject = SendWithoutBlock v13, :freeze
           CheckInterrupts
@@ -10744,6 +10764,7 @@ mod opt_tests {
         bb2(v6:BasicObject):
           v11:HashExact = NewHash
           v12:NilClass = Const Value(nil)
+          CountUnoptimizedCFunc Hash#freeze
           v14:BasicObject = SendWithoutBlock v11, :freeze, v12
           CheckInterrupts
           Return v14
@@ -10811,6 +10832,7 @@ mod opt_tests {
           Jump bb2(v4)
         bb2(v6:BasicObject):
           v11:ArrayExact = NewArray
+          CountUnoptimizedCFunc Kernel#dup
           v13:BasicObject = SendWithoutBlock v11, :dup
           v15:BasicObject = SendWithoutBlock v13, :freeze
           CheckInterrupts
@@ -10835,6 +10857,7 @@ mod opt_tests {
         bb2(v6:BasicObject):
           v11:ArrayExact = NewArray
           v12:NilClass = Const Value(nil)
+          CountUnoptimizedCFunc Array#freeze
           v14:BasicObject = SendWithoutBlock v11, :freeze, v12
           CheckInterrupts
           Return v14
@@ -10903,6 +10926,7 @@ mod opt_tests {
         bb2(v6:BasicObject):
           v10:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
           v12:StringExact = StringCopy v10
+          CountUnoptimizedCFunc String#dup
           v14:BasicObject = SendWithoutBlock v12, :dup
           v16:BasicObject = SendWithoutBlock v14, :freeze
           CheckInterrupts
@@ -10928,6 +10952,7 @@ mod opt_tests {
           v10:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
           v12:StringExact = StringCopy v10
           v13:NilClass = Const Value(nil)
+          CountUnoptimizedCFunc String#freeze
           v15:BasicObject = SendWithoutBlock v12, :freeze, v13
           CheckInterrupts
           Return v15
@@ -10996,6 +11021,7 @@ mod opt_tests {
         bb2(v6:BasicObject):
           v10:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
           v12:StringExact = StringCopy v10
+          CountUnoptimizedCFunc String#dup
           v14:BasicObject = SendWithoutBlock v12, :dup
           v16:BasicObject = SendWithoutBlock v14, :-@
           CheckInterrupts
@@ -11133,6 +11159,7 @@ mod opt_tests {
         bb2(v8:BasicObject, v9:BasicObject):
           v13:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
           v25:BasicObject = GuardTypeNot v9, String
+          CountUnoptimizedCFunc Array#to_s
           v26:BasicObject = SendWithoutBlock v9, :to_s
           v17:String = AnyToString v9, str: v26
           v19:StringExact = StringConcat v13, v17

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"]
@@ -188,6 +190,21 @@ pub(crate) use incr_counter;
 /// The number of side exits from each YARV instruction
 pub type ExitCounters = [u64; VM_INSTRUCTION_SIZE as usize];
 
+/// Store signature ("Klass#method") to counter pointer mappings for unoptimized cfuncs
+pub static UNOPTIMIZED_CFUNC_COUNTER_POINTERS: LazyLock<Mutex<HashMap<String, Box<u64>>>> = LazyLock::new(|| Mutex::new(HashMap::new()));
+
+/// Get or create a counter pointer for the given signature
+pub fn get_or_create_unoptimized_cfunc_counter_ptr(signature: String) -> *mut u64 {
+    let mut map = UNOPTIMIZED_CFUNC_COUNTER_POINTERS.lock().unwrap();
+    let counter = map.entry(signature).or_insert_with(|| Box::new(0));
+    &mut **counter as *mut u64
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_count_unoptimized_cfunc(counter_ptr: *mut u64) {
+    unsafe { *counter_ptr += 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();
@@ -384,6 +401,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 unoptimized cfunc counters
+    let unoptimized_cfuncs = UNOPTIMIZED_CFUNC_COUNTER_POINTERS.lock().unwrap();
+    for (signature, counter) in unoptimized_cfuncs.iter() {
+        let key_string = format!("not_optimized_cfuncs_{}", signature);
+        set_stat_usize!(hash, &key_string, **counter);
+    }
+
     hash
 }