YJIT: Show Context stats on exit (#7327)

4 files changed, 63 insertions, 8 deletions

diff --git a/yjit.c b/yjit.c
index c4cf155cf6..e00eda4581 100644
--- a/yjit.c
+++ b/yjit.c
@@ -1100,7 +1100,7 @@ object_shape_count(rb_execution_context_t *ec, VALUE self)
 // Primitives used by yjit.rb
 VALUE rb_yjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_trace_exit_locations_enabled_p(rb_execution_context_t *ec, VALUE self);
-VALUE rb_yjit_get_stats(rb_execution_context_t *ec, VALUE self);
+VALUE rb_yjit_get_stats(rb_execution_context_t *ec, VALUE self, VALUE context);
 VALUE rb_yjit_reset_stats_bang(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_disasm_iseq(rb_execution_context_t *ec, VALUE self, VALUE iseq);
 VALUE rb_yjit_insns_compiled(rb_execution_context_t *ec, VALUE self, VALUE iseq);
diff --git a/yjit.rb b/yjit.rb
index b695a10cac..8387767958 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -145,8 +145,8 @@ module RubyVM::YJIT
 
   # Return a hash for statistics generated for the --yjit-stats command line option.
   # Return nil when option is not passed or unavailable.
-  def self.runtime_stats
-    stats = Primitive.rb_yjit_get_stats
+  def self.runtime_stats(context: false)
+    stats = Primitive.rb_yjit_get_stats(context)
     return stats if stats.nil?
 
     stats[:object_shape_count] = Primitive.object_shape_count
@@ -233,7 +233,7 @@ module RubyVM::YJIT
 
     # Format and print out counters
     def _print_stats # :nodoc:
-      stats = runtime_stats
+      stats = runtime_stats(context: true)
       return unless stats
 
       $stderr.puts("***YJIT: Printing YJIT statistics on exit***")
@@ -277,6 +277,8 @@ module RubyVM::YJIT
       $stderr.puts "freed_code_size:       " + format_number(13, stats[:freed_code_size])
       $stderr.puts "code_region_size:      " + format_number(13, stats[:code_region_size])
       $stderr.puts "yjit_alloc_size:       " + format_number(13, stats[:yjit_alloc_size]) if stats.key?(:yjit_alloc_size)
+      $stderr.puts "live_context_size:     " + format_number(13, stats[:live_context_size])
+      $stderr.puts "live_context_count:    " + format_number(13, stats[:live_context_count])
       $stderr.puts "live_page_count:       " + format_number(13, stats[:live_page_count])
       $stderr.puts "freed_page_count:      " + format_number(13, stats[:freed_page_count])
       $stderr.puts "code_gc_count:         " + format_number(13, stats[:code_gc_count])
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index bae99546d6..268555a87b 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -440,6 +440,18 @@ impl Branch {
     fn get_target_address(&self, target_idx: usize) -> Option<CodePtr> {
         self.targets[target_idx].as_ref().and_then(|target| target.get_address())
     }
+
+    fn get_stub_count(&self) -> usize {
+        let mut count = 0;
+        for target in self.targets.iter() {
+            if let Some(target) = target {
+                if let BranchTarget::Stub(_) = target.as_ref() {
+                    count += 1;
+                }
+            }
+        }
+        count
+    }
 }
 
 // In case a block is invalidated, this helps to remove all pointers to the block.
@@ -551,7 +563,7 @@ impl Eq for BlockRef {}
 #[derive(Default)]
 pub struct IseqPayload {
     // Basic block versions
-    version_map: VersionMap,
+    pub version_map: VersionMap,
 
     // Indexes of code pages used by this this ISEQ
     pub pages: HashSet<usize>,
@@ -621,6 +633,15 @@ pub fn for_each_iseq<F: FnMut(IseqPtr)>(mut callback: F) {
     unsafe { rb_yjit_for_each_iseq(Some(callback_wrapper), (&mut data) as *mut _ as *mut c_void) };
 }
 
+/// Iterate over all ISEQ payloads
+pub fn for_each_iseq_payload<F: FnMut(&IseqPayload)>(mut callback: F) {
+    for_each_iseq(|iseq| {
+        if let Some(iseq_payload) = get_iseq_payload(iseq) {
+            callback(iseq_payload);
+        }
+    });
+}
+
 /// Iterate over all on-stack ISEQs
 pub fn for_each_on_stack_iseq<F: FnMut(IseqPtr)>(mut callback: F) {
     unsafe extern "C" fn callback_wrapper(iseq: IseqPtr, data: *mut c_void) {
@@ -1032,6 +1053,14 @@ impl Block {
         self.ctx.clone()
     }
 
+    pub fn get_ctx_count(&self) -> usize {
+        let mut count = 1; // block.ctx
+        for branch in self.outgoing.iter() {
+            count += branch.borrow().get_stub_count();
+        }
+        count
+    }
+
     #[allow(unused)]
     pub fn get_start_addr(&self) -> CodePtr {
         self.start_addr
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 774f718318..49a4e43295 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -4,6 +4,8 @@
 #![allow(dead_code)] // Counters are only used with the stats features
 
 use crate::codegen::CodegenGlobals;
+use crate::core::Context;
+use crate::core::for_each_iseq_payload;
 use crate::cruby::*;
 use crate::options::*;
 use crate::yjit::yjit_enabled_p;
@@ -347,8 +349,8 @@ pub extern "C" fn rb_yjit_stats_enabled_p(_ec: EcPtr, _ruby_self: VALUE) -> VALU
 /// Primitive called in yjit.rb.
 /// Export all YJIT statistics as a Ruby hash.
 #[no_mangle]
-pub extern "C" fn rb_yjit_get_stats(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
-    with_vm_lock(src_loc!(), || rb_yjit_gen_stats_dict())
+pub extern "C" fn rb_yjit_get_stats(_ec: EcPtr, _ruby_self: VALUE, context: VALUE) -> VALUE {
+    with_vm_lock(src_loc!(), || rb_yjit_gen_stats_dict(context == Qtrue))
 }
 
 /// Primitive called in yjit.rb
@@ -403,7 +405,7 @@ pub extern "C" fn rb_yjit_get_exit_locations(_ec: EcPtr, _ruby_self: VALUE) -> V
 }
 
 /// Export all YJIT statistics as a Ruby hash.
-fn rb_yjit_gen_stats_dict() -> VALUE {
+fn rb_yjit_gen_stats_dict(context: bool) -> VALUE {
     // If YJIT is not enabled, return Qnil
     if !yjit_enabled_p() {
         return Qnil;
@@ -450,6 +452,13 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
         // Rust global allocations in bytes
         #[cfg(feature="stats")]
         hash_aset_usize!(hash, "yjit_alloc_size", global_allocation_size());
+
+        if context {
+            let live_context_count = get_live_context_count();
+            let context_size = std::mem::size_of::<Context>();
+            hash_aset_usize!(hash, "live_context_count", live_context_count);
+            hash_aset_usize!(hash, "live_context_size", live_context_count * context_size);
+        }
     }
 
     // If we're not generating stats, the hash is done
@@ -496,6 +505,21 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
     hash
 }
 
+fn get_live_context_count() -> usize {
+    let mut count = 0;
+    for_each_iseq_payload(|iseq_payload| {
+        for blocks in iseq_payload.version_map.iter() {
+            for block in blocks.iter() {
+                count += block.borrow().get_ctx_count();
+            }
+        }
+        for block in iseq_payload.dead_blocks.iter() {
+            count += block.borrow().get_ctx_count();
+        }
+    });
+    count
+}
+
 /// Record the backtrace when a YJIT exit occurs. This functionality requires
 /// that the stats feature is enabled as well as the --yjit-trace-exits option.
 ///