YJIT: Instrument global allocations on stats build (#6712)

* YJIT: Instrument global allocations on stats build

* Just use GLOVAL_ALLOCATOR.stats()

yjit.rb

@@ -262,6 +262,7 @@ module RubyVM::YJIT
       $stderr.puts "inline_code_size:      " + ("%10d" % stats[:inline_code_size])
       $stderr.puts "outlined_code_size:    " + ("%10d" % stats[:outlined_code_size])
       $stderr.puts "freed_code_size:       " + ("%10d" % stats[:freed_code_size])
+      $stderr.puts "yjit_alloc_size:       " + ("%10d" % stats[:yjit_alloc_size]) if stats.key?(:yjit_alloc_size)
       $stderr.puts "live_page_count:       " + ("%10d" % stats[:live_page_count])
       $stderr.puts "freed_page_count:      " + ("%10d" % stats[:freed_page_count])
       $stderr.puts "code_gc_count:         " + ("%10d" % stats[:code_gc_count])

yjit/Cargo.lock

@@ -34,9 +34,16 @@ version = "0.2.124"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "21a41fed9d98f27ab1c6d161da622a4fa35e8a54a8adc24bbf3ddd0ef70b0e50"
 
+[[package]]
+name = "stats_alloc"
+version = "0.1.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c0e04424e733e69714ca1bbb9204c1a57f09f5493439520f9f68c132ad25eec"
+
 [[package]]
 name = "yjit"
 version = "0.1.0"
 dependencies = [
  "capstone",
+ "stats_alloc",
 ]

yjit/Cargo.toml

@@ -16,12 +16,13 @@ crate-type = ["staticlib"]
 # No required dependencies to simplify build process. TODO: Link to yet to be
 # written rationale. Optional For development and testing purposes
 capstone = { version = "0.10.0", optional = true }
+stats_alloc = { version = "0.1.10", optional = true }
 
 [features]
 # NOTE: Development builds select a set of these via configure.ac
 # For debugging, `make V=1` shows exact cargo invocation.
 disasm = ["capstone"]
-stats = []
+stats = ["stats_alloc"]
 
 [profile.dev]
 opt-level = 0

yjit/src/stats.rs

@@ -8,6 +8,11 @@ use crate::cruby::*;
 use crate::options::*;
 use crate::yjit::yjit_enabled_p;
 
+// stats_alloc is a middleware to instrument global allocations in Rust.
+#[cfg(feature="stats")]
+#[global_allocator]
+static GLOBAL_ALLOCATOR: &stats_alloc::StatsAlloc<std::alloc::System> = &stats_alloc::INSTRUMENTED_SYSTEM;
+
 // 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];
@@ -396,6 +401,10 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
 
         // Code GC count
         hash_aset_usize!(hash, "code_gc_count", CodegenGlobals::get_code_gc_count());
+
+        // Rust global allocations in bytes
+        #[cfg(feature="stats")]
+        hash_aset_usize!(hash, "yjit_alloc_size", global_allocation_size());
     }
 
     // If we're not generating stats, the hash is done
@@ -606,3 +615,10 @@ pub extern "C" fn rb_yjit_count_side_exit_op(exit_pc: *const VALUE) -> *const VA
     // This function must return exit_pc!
     return exit_pc;
 }
+
+// Get the size of global allocations in Rust.
+#[cfg(feature="stats")]
+fn global_allocation_size() -> usize {
+    let stats = GLOBAL_ALLOCATOR.stats();
+    stats.bytes_allocated.saturating_sub(stats.bytes_deallocated)
+}