YJIT: add top C function call counts to `--yjit-stats` (#9047)

* YJIT: gather call counts for individual cfuncs

Co-authored by Takashi Kokubun

2 files changed, 97 insertions, 1 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 7846635e71..51a0ba3a31 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -5398,6 +5398,7 @@ fn gen_send_cfunc(
         return None;
     }
 
+    // Increment total cfunc send count
     gen_counter_incr(asm, Counter::num_send_cfunc);
 
     // Delegate to codegen for C methods if we have it.
@@ -5416,6 +5417,32 @@ fn gen_send_cfunc(
         }
     }
 
+    // Log the name of the method we're calling to,
+    // note that we intentionally don't do this for inlined cfuncs
+    if get_option!(gen_stats) {
+        // TODO: extract code to get method name string into its own function
+
+        // Assemble the method name string
+        let mid = unsafe { vm_ci_mid(ci) };
+        let class_name = if recv_known_klass != ptr::null() {
+            unsafe { cstr_to_rust_string(rb_class2name(*recv_known_klass)) }.unwrap()
+        } else {
+            "Unknown".to_string()
+        };
+        let method_name = if mid != 0 {
+            unsafe { cstr_to_rust_string(rb_id2name(mid)) }.unwrap()
+        } else {
+            "Unknown".to_string()
+        };
+        let name_str = format!("{}#{}", class_name, method_name);
+
+        // Get an index for this cfunc name
+        let cfunc_idx = get_cfunc_idx(&name_str);
+
+        // Increment the counter for this cfunc
+        asm.ccall(incr_cfunc_counter as *const u8, vec![cfunc_idx.into()]);
+    }
+
     // Check for interrupts
     gen_check_ints(asm, Counter::guard_send_interrupted);
 
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index e3b6103299..7a52fc0e23 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -6,6 +6,7 @@
 use std::alloc::{GlobalAlloc, Layout, System};
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::time::Instant;
+use std::collections::HashMap;
 
 use crate::codegen::CodegenGlobals;
 use crate::core::Context;
@@ -52,6 +53,58 @@ unsafe impl GlobalAlloc for StatsAlloc {
     }
 }
 
+/// Mapping of C function name to integer indices
+/// This is accessed at compilation time only (protected by a lock)
+static mut CFUNC_NAME_TO_IDX: Option<HashMap<String, usize>> = None;
+
+/// Vector of call counts for each C function index
+/// This is modified (but not resized) by JITted code
+static mut CFUNC_CALL_COUNT: Option<Vec<u64>> = None;
+
+/// Assign an index to a given cfunc name string
+pub fn get_cfunc_idx(name: &str) -> usize
+{
+    //println!("{}", name);
+
+    unsafe {
+        if CFUNC_NAME_TO_IDX.is_none() {
+            CFUNC_NAME_TO_IDX = Some(HashMap::default());
+        }
+
+        if CFUNC_CALL_COUNT.is_none() {
+            CFUNC_CALL_COUNT = Some(Vec::default());
+        }
+
+        let name_to_idx = CFUNC_NAME_TO_IDX.as_mut().unwrap();
+
+        match name_to_idx.get(name) {
+            Some(idx) => *idx,
+            None => {
+                let idx = name_to_idx.len();
+                name_to_idx.insert(name.to_string(), idx);
+
+                // Resize the call count vector
+                let cfunc_call_count = CFUNC_CALL_COUNT.as_mut().unwrap();
+                if idx >= cfunc_call_count.len() {
+                    cfunc_call_count.resize(idx + 1, 0);
+                }
+
+                idx
+            }
+        }
+    }
+}
+
+// Increment the counter for a C function
+pub extern "C" fn incr_cfunc_counter(idx: usize)
+{
+    unsafe {
+        let cfunc_call_count = CFUNC_CALL_COUNT.as_mut().unwrap();
+        assert!(idx < cfunc_call_count.len());
+        cfunc_call_count[idx] += 1;
+    }
+}
+
 // YJIT exit counts for each instruction type
 const VM_INSTRUCTION_SIZE_USIZE: usize = VM_INSTRUCTION_SIZE as usize;
 static mut EXIT_OP_COUNT: [u64; VM_INSTRUCTION_SIZE_USIZE] = [0; VM_INSTRUCTION_SIZE_USIZE];
@@ -663,7 +716,6 @@ fn rb_yjit_gen_stats_dict(context: bool) -> VALUE {
         return hash;
     }
 
-    // If the stats feature is enabled
     unsafe {
         // Indicate that the complete set of stats is available
         rb_hash_aset(hash, rust_str_to_sym("all_stats"), Qtrue);
@@ -689,6 +741,23 @@ fn rb_yjit_gen_stats_dict(context: bool) -> VALUE {
             let value = VALUE::fixnum_from_usize(EXIT_OP_COUNT[op_idx] as usize);
             rb_hash_aset(hash, key, value);
         }
+
+        // Create a hash for the cfunc call counts
+        if let Some(cfunc_name_to_idx) = CFUNC_NAME_TO_IDX.as_mut() {
+            let call_counts = CFUNC_CALL_COUNT.as_mut().unwrap();
+            let calls_hash = rb_hash_new();
+
+            for (name, idx) in cfunc_name_to_idx {
+                let count = call_counts[*idx];
+                println!("{}: {}", name, count);
+
+                let key = rust_str_to_sym(name);
+                let value = VALUE::fixnum_from_usize(count as usize);
+                rb_hash_aset(calls_hash, key, value);
+            }
+
+            rb_hash_aset(hash, rust_str_to_sym("cfunc_calls"), calls_hash);
+        }
     }
 
     hash