YJIT: Encode doubles to VALUE objects and move stat generation to rust (#11388)

* YJIT: Encode doubles to VALUE objects and move stat generation to rust

Stats that can now be generated from rust have been moved there.

* Move object_shape_count call for runtime_stats to rust

This reduces the ruby method to a single primitive.

* Change hash_aset_usize from macro to function

test/ruby/test_yjit.rb

@@ -1623,6 +1623,21 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_runtime_stats_types
+    assert_compiles(<<~'RUBY', exits: :any, result: true)
+      def test = :ok
+      3.times { test }
+
+      stats = RubyVM::YJIT.runtime_stats
+      return true unless stats[:all_stats]
+
+      [
+        stats[:object_shape_count].is_a?(Integer),
+        stats[:ratio_in_yjit].is_a?(Float),
+      ].all?
+    RUBY
+  end
+
   private
 
   def code_gc_helpers

yjit.c

@@ -1165,8 +1165,8 @@ rb_yjit_invalidate_all_method_lookup_assumptions(void)
 }
 
 // Number of object shapes, which might be useful for investigating YJIT exit reasons.
-static VALUE
-object_shape_count(rb_execution_context_t *ec, VALUE self)
+VALUE
+rb_object_shape_count(void)
 {
     // next_shape_id starts from 0, so it's the same as the count
     return ULONG2NUM((unsigned long)GET_SHAPE_TREE()->next_shape_id);

yjit.rb

@@ -156,34 +156,7 @@ 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()
-    return stats if stats.nil?
-
-    stats[:object_shape_count] = Primitive.object_shape_count
-    return stats unless Primitive.rb_yjit_stats_enabled_p
-
-    side_exits = total_exit_count(stats)
-    total_exits = side_exits + stats[:leave_interp_return]
-
-    # Number of instructions that finish executing in YJIT.
-    # See :count-placement: about the subtraction.
-    retired_in_yjit = stats[:yjit_insns_count] - side_exits
-
-    # Average length of instruction sequences executed by YJIT
-    avg_len_in_yjit = total_exits > 0 ? retired_in_yjit.to_f / total_exits : 0
-
-    # Proportion of instructions that retire in YJIT
-    total_insns_count = retired_in_yjit + stats[:vm_insns_count]
-    yjit_ratio_pct = 100.0 * retired_in_yjit.to_f / total_insns_count
-    stats[:total_insns_count] = total_insns_count
-    stats[:ratio_in_yjit] = yjit_ratio_pct
-
-    # Make those stats available in RubyVM::YJIT.runtime_stats as well
-    stats[:side_exit_count]  = side_exits
-    stats[:total_exit_count] = total_exits
-    stats[:avg_len_in_yjit]  = avg_len_in_yjit
-
-    stats
+    Primitive.rb_yjit_get_stats
   end
 
   # Format and print out counters as a String. This returns a non-empty
@@ -437,7 +410,7 @@ module RubyVM::YJIT
     end
 
     def print_sorted_exit_counts(stats, out:, prefix:, how_many: 20, left_pad: 4) # :nodoc:
-      total_exits = total_exit_count(stats)
+      total_exits = stats[:side_exit_count]
 
       if total_exits > 0
         exits = []
@@ -464,14 +437,6 @@ module RubyVM::YJIT
       end
     end
 
-    def total_exit_count(stats, prefix: "exit_") # :nodoc:
-      total = 0
-      stats.each do |k,v|
-        total += v if k.start_with?(prefix)
-      end
-      total
-    end
-
     def print_counters(counters, out:, prefix:, prompt:, optional: false) # :nodoc:
       counters = counters.filter { |key, _| key.start_with?(prefix) }
       counters.filter! { |_, value| value != 0 }

yjit/bindgen/src/main.rs

@@ -264,6 +264,9 @@ fn main() {
         .blocklist_type("rb_method_definition_.*") // Large struct with a bitfield and union of many types - don't import (yet?)
         .opaque_type("rb_method_definition_.*")
 
+        // From numeric.c
+        .allowlist_function("rb_float_new")
+
         // From vm_core.h
         .allowlist_var("rb_mRubyVMFrozenCore")
         .allowlist_var("VM_BLOCK_HANDLER_NONE")
@@ -304,6 +307,7 @@ fn main() {
         .allowlist_type("rb_iseq_type")
 
         // From yjit.c
+        .allowlist_function("rb_object_shape_count")
         .allowlist_function("rb_iseq_(get|set)_yjit_payload")
         .allowlist_function("rb_iseq_pc_at_idx")
         .allowlist_function("rb_iseq_opcode_at_pc")

yjit/src/cruby.rs

@@ -116,6 +116,11 @@ extern "C" {
         me: *const rb_callable_method_entry_t,
         ci: *const rb_callinfo,
     ) -> *const rb_callable_method_entry_t;
+
+    // Floats within range will be encoded without creating objects in the heap.
+    // (Range is 0x3000000000000001 to 0x4fffffffffffffff (1.7272337110188893E-77 to 2.3158417847463237E+77).
+    pub fn rb_float_new(d: f64) -> VALUE;
+
     pub fn rb_hash_empty_p(hash: VALUE) -> VALUE;
     pub fn rb_yjit_str_concat_codepoint(str: VALUE, codepoint: VALUE);
     pub fn rb_str_setbyte(str: VALUE, index: VALUE, value: VALUE) -> VALUE;

yjit/src/cruby_bindings.inc.rs

@@ -976,6 +976,7 @@ extern "C" {
     pub fn rb_get_alloc_func(klass: VALUE) -> rb_alloc_func_t;
     pub fn rb_method_basic_definition_p(klass: VALUE, mid: ID) -> ::std::os::raw::c_int;
     pub fn rb_bug(fmt: *const ::std::os::raw::c_char, ...) -> !;
+    pub fn rb_float_new(d: f64) -> VALUE;
     pub fn rb_gc_mark(obj: VALUE);
     pub fn rb_gc_mark_movable(obj: VALUE);
     pub fn rb_gc_location(obj: VALUE) -> VALUE;
@@ -1269,6 +1270,7 @@ extern "C" {
         file: *const ::std::os::raw::c_char,
         line: ::std::os::raw::c_int,
     );
+    pub fn rb_object_shape_count() -> VALUE;
     pub fn rb_yjit_assert_holding_vm_lock();
     pub fn rb_yjit_sendish_sp_pops(ci: *const rb_callinfo) -> usize;
     pub fn rb_yjit_invokeblock_sp_pops(ci: *const rb_callinfo) -> usize;

yjit/src/stats.rs

@@ -709,6 +709,17 @@ pub extern "C" fn rb_yjit_incr_counter(counter_name: *const std::os::raw::c_char
     unsafe { *counter_ptr += 1 };
 }
 
+fn hash_aset_usize(hash: VALUE, key: &str, value: usize) {
+    let rb_key = rust_str_to_sym(key);
+    let rb_value = VALUE::fixnum_from_usize(value);
+    unsafe { rb_hash_aset(hash, rb_key, rb_value); }
+}
+
+fn hash_aset_double(hash: VALUE, key: &str, value: f64) {
+    let rb_key = rust_str_to_sym(key);
+    unsafe { rb_hash_aset(hash, rb_key, rb_float_new(value)); }
+}
+
 /// Export all YJIT statistics as a Ruby hash.
 fn rb_yjit_gen_stats_dict() -> VALUE {
     // If YJIT is not enabled, return Qnil
@@ -716,14 +727,6 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
         return Qnil;
     }
 
-    macro_rules! hash_aset_usize {
-        ($hash:ident, $counter_name:expr, $value:expr) => {
-            let key = rust_str_to_sym($counter_name);
-            let value = VALUE::fixnum_from_usize($value);
-            rb_hash_aset($hash, key, value);
-        }
-    }
-
     let hash = unsafe { rb_hash_new() };
 
     unsafe {
@@ -732,37 +735,39 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
         let ocb = CodegenGlobals::get_outlined_cb();
 
         // Inline code size
-        hash_aset_usize!(hash, "inline_code_size", cb.code_size());
+        hash_aset_usize(hash, "inline_code_size", cb.code_size());
 
         // Outlined code size
-        hash_aset_usize!(hash, "outlined_code_size", ocb.unwrap().code_size());
+        hash_aset_usize(hash, "outlined_code_size", ocb.unwrap().code_size());
 
         // GCed pages
         let freed_page_count = cb.num_freed_pages();
-        hash_aset_usize!(hash, "freed_page_count", freed_page_count);
+        hash_aset_usize(hash, "freed_page_count", freed_page_count);
 
         // GCed code size
-        hash_aset_usize!(hash, "freed_code_size", freed_page_count * cb.page_size());
+        hash_aset_usize(hash, "freed_code_size", freed_page_count * cb.page_size());
 
         // Live pages
-        hash_aset_usize!(hash, "live_page_count", cb.num_mapped_pages() - freed_page_count);
+        hash_aset_usize(hash, "live_page_count", cb.num_mapped_pages() - freed_page_count);
 
         // Size of memory region allocated for JIT code
-        hash_aset_usize!(hash, "code_region_size", cb.mapped_region_size());
+        hash_aset_usize(hash, "code_region_size", cb.mapped_region_size());
 
         // Rust global allocations in bytes
-        hash_aset_usize!(hash, "yjit_alloc_size", GLOBAL_ALLOCATOR.alloc_size.load(Ordering::SeqCst));
+        hash_aset_usize(hash, "yjit_alloc_size", GLOBAL_ALLOCATOR.alloc_size.load(Ordering::SeqCst));
 
         // How many bytes we are using to store context data
         let context_data = CodegenGlobals::get_context_data();
-        hash_aset_usize!(hash, "context_data_bytes", context_data.num_bytes());
-        hash_aset_usize!(hash, "context_cache_bytes", crate::core::CTX_CACHE_BYTES);
+        hash_aset_usize(hash, "context_data_bytes", context_data.num_bytes());
+        hash_aset_usize(hash, "context_cache_bytes", crate::core::CTX_CACHE_BYTES);
 
         // VM instructions count
-        hash_aset_usize!(hash, "vm_insns_count", rb_vm_insns_count as usize);
+        hash_aset_usize(hash, "vm_insns_count", rb_vm_insns_count as usize);
 
-        hash_aset_usize!(hash, "live_iseq_count", rb_yjit_live_iseq_count as usize);
-        hash_aset_usize!(hash, "iseq_alloc_count", rb_yjit_iseq_alloc_count as usize);
+        hash_aset_usize(hash, "live_iseq_count", rb_yjit_live_iseq_count as usize);
+        hash_aset_usize(hash, "iseq_alloc_count", rb_yjit_iseq_alloc_count as usize);
+
+        rb_hash_aset(hash, rust_str_to_sym("object_shape_count"), rb_object_shape_count());
     }
 
     // If we're not generating stats, put only default counters
@@ -797,16 +802,44 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
             rb_hash_aset(hash, key, value);
         }
 
+        let mut side_exits = 0;
+
         // For each entry in exit_op_count, add a stats entry with key "exit_INSTRUCTION_NAME"
         // and the value is the count of side exits for that instruction.
         for op_idx in 0..VM_INSTRUCTION_SIZE_USIZE {
             let op_name = insn_name(op_idx);
             let key_string = "exit_".to_owned() + &op_name;
             let key = rust_str_to_sym(&key_string);
-            let value = VALUE::fixnum_from_usize(EXIT_OP_COUNT[op_idx] as usize);
+            let count = EXIT_OP_COUNT[op_idx];
+            side_exits += count;
+            let value = VALUE::fixnum_from_usize(count as usize);
             rb_hash_aset(hash, key, value);
         }
 
+        hash_aset_usize(hash, "side_exit_count", side_exits as usize);
+
+        let total_exits = side_exits + *get_counter_ptr(&Counter::leave_interp_return.get_name());
+        hash_aset_usize(hash, "total_exit_count", total_exits as usize);
+
+        // Number of instructions that finish executing in YJIT.
+        // See :count-placement: about the subtraction.
+        let retired_in_yjit = *get_counter_ptr(&Counter::yjit_insns_count.get_name()) - side_exits;
+
+        // Average length of instruction sequences executed by YJIT
+        let avg_len_in_yjit: f64 = if total_exits > 0 {
+            retired_in_yjit as f64 / total_exits as f64
+        } else {
+            0_f64
+        };
+        hash_aset_double(hash, "avg_len_in_yjit", avg_len_in_yjit);
+
+        // Proportion of instructions that retire in YJIT
+        let total_insns_count = retired_in_yjit + rb_vm_insns_count;
+        hash_aset_usize(hash, "total_insns_count", total_insns_count as usize);
+
+        let ratio_in_yjit: f64 = 100.0 * retired_in_yjit as f64 / total_insns_count as f64;
+        hash_aset_double(hash, "ratio_in_yjit", ratio_in_yjit);
+
         // Set method call counts in a Ruby dict
         fn set_call_counts(
             calls_hash: VALUE,