YJIT: Add stats option to RubyVM::YJIT.enable (#9297)

ruby · Dec 19, 2023 · bd91c51 · bd91c51
1 parent 084b44e
commit bd91c51
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 17 deletions.
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
@@ -67,6 +67,29 @@ def test_yjit_enable
     RUBY
   end
 
+  def test_yjit_enable_stats
+    args = []
+    args << "--disable=yjit" if RubyVM::YJIT.enabled?
+    assert_separately(args, <<~RUBY, ignore_stderr: true)
+      assert_false RubyVM::YJIT.enabled?
+      assert_nil RubyVM::YJIT.runtime_stats
+
+      RubyVM::YJIT.enable(stats: true)
+
+      assert_true RubyVM::YJIT.enabled?
+      assert_true RubyVM::YJIT.runtime_stats[:all_stats]
+    RUBY
+  end
+
+  def test_yjit_enable_stats_quiet
+    assert_in_out_err(['--yjit-disable', '-e', 'RubyVM::YJIT.enable(stats: true)']) do |_stdout, stderr, _status|
+      assert_not_empty stderr
+    end
+    assert_in_out_err(['--yjit-disable', '-e', 'RubyVM::YJIT.enable(stats: :quiet)']) do |_stdout, stderr, _status|
+      assert_empty stderr
+    end
+  end
+
   def test_yjit_enable_with_call_threshold
     assert_separately(%w[--yjit-disable --yjit-call-threshold=1], <<~RUBY)
       def not_compiled = nil

diff --git a/yjit.c b/yjit.c
@@ -1171,7 +1171,7 @@ VALUE rb_yjit_insns_compiled(rb_execution_context_t *ec, VALUE self, VALUE iseq)
 VALUE rb_yjit_code_gc(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_simulate_oom_bang(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_get_exit_locations(rb_execution_context_t *ec, VALUE self);
-VALUE rb_yjit_enable(rb_execution_context_t *ec, VALUE self);
+VALUE rb_yjit_enable(rb_execution_context_t *ec, VALUE self, VALUE gen_stats, VALUE print_stats);
 
 // Preprocessed yjit.rb generated during build
 #include "yjit.rbinc"

diff --git a/yjit.rb b/yjit.rb
@@ -29,8 +29,10 @@ def self.reset_stats!
   end
 
   # Enable \YJIT compilation.
-  def self.enable
-    Primitive.rb_yjit_enable
+  def self.enable(stats: false)
+    return false if enabled?
+    at_exit { print_and_dump_stats } if stats
+    Primitive.rb_yjit_enable(stats, stats != :quiet)
   end
 
   # If --yjit-trace-exits is enabled parse the hashes from
@@ -223,18 +225,21 @@ def self.simulate_oom! # :nodoc:
     Primitive.rb_yjit_simulate_oom_bang
   end
 
-  # Avoid calling a method here to not interfere with compilation tests
+  # Avoid calling a Ruby method here to not interfere with compilation tests
   if Primitive.rb_yjit_stats_enabled_p
-    at_exit do
+    at_exit { print_and_dump_stats }
+  end
+
+  class << self # :stopdoc:
+    private
+
+    # Print stats and dump exit locations
+    def print_and_dump_stats
       if Primitive.rb_yjit_print_stats_p
         _print_stats
       end
       _dump_locations
     end
-  end
-
-  class << self # :stopdoc:
-    private
 
     def _dump_locations # :nodoc:
       return unless trace_exit_locations_enabled?

diff --git a/yjit/src/options.rs b/yjit/src/options.rs
@@ -249,9 +249,9 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
         ("no-type-prop", "") => unsafe { OPTIONS.no_type_prop = true },
         ("stats", _) => match opt_val {
             "" => unsafe { OPTIONS.gen_stats = true },
-            "quiet" => {
-                unsafe { OPTIONS.gen_stats = true }
-                unsafe { OPTIONS.print_stats = false }
+            "quiet" => unsafe {
+                OPTIONS.gen_stats = true;
+                OPTIONS.print_stats = false;
             },
             _ => {
                 return None;

diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
@@ -168,13 +168,13 @@ pub extern "C" fn rb_yjit_code_gc(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
 
 /// Enable YJIT compilation, returning true if YJIT was previously disabled
 #[no_mangle]
-pub extern "C" fn rb_yjit_enable(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+pub extern "C" fn rb_yjit_enable(_ec: EcPtr, _ruby_self: VALUE, gen_stats: VALUE, print_stats: VALUE) -> VALUE {
     with_vm_lock(src_loc!(), || {
-        if yjit_enabled_p() {
-            return Qfalse;
+        // Initialize and enable YJIT
+        unsafe {
+            OPTIONS.gen_stats = gen_stats.test();
+            OPTIONS.print_stats = print_stats.test();
         }
-
-        // Initialize and enable YJIT if currently disabled
         yjit_init();
 
         // Add "+YJIT" to RUBY_DESCRIPTION