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

* YJIT: gather call counts for individual cfuncs

Co-authored by Takashi Kokubun

yjit.rb

@@ -369,6 +369,35 @@ module RubyVM::YJIT
       out.puts "avg_len_in_yjit:       " + ("%13.1f" % stats[:avg_len_in_yjit])
 
       print_sorted_exit_counts(stats, out: out, prefix: "exit_")
+
+      print_sorted_cfunc_calls(stats, out:out)
+    end
+
+    def print_sorted_cfunc_calls(stats, out:, how_many: 20, left_pad: 4) # :nodoc:
+      calls = stats[:cfunc_calls]
+      #puts calls
+
+      # Total number of cfunc calls
+      num_send_cfunc = stats[:num_send_cfunc]
+
+      # Sort calls by decreasing frequency and keep the top N
+      pairs = calls.map { |k,v| [k, v] }
+      pairs.sort_by! {|pair| pair[1] }
+      pairs.reverse!
+      pairs = pairs[0...how_many]
+
+      top_n_total = pairs.sum { |name, count| count }
+      top_n_pct = 100.0 * top_n_total / num_send_cfunc
+      longest_name_len = pairs.max_by { |name, count| name.length }.first.length
+
+      out.puts "Top-#{pairs.size} most frequent C calls (#{"%.1f" % top_n_pct}% of C calls):"
+
+      pairs.each do |name, count|
+        padding = longest_name_len + left_pad
+        padded_name = "%#{padding}s" % name
+        padded_count = format_number_pct(10, count, num_send_cfunc)
+        out.puts("#{padded_name}: #{padded_count}")
+      end
     end
 
     def print_sorted_exit_counts(stats, out:, prefix:, how_many: 20, left_pad: 4) # :nodoc:

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);
 

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