YJIT: Add `--yjit-pause` and `RubyVM::YJIT.resume` (#7609)

* YJIT: Add --yjit-pause and RubyVM::YJIT.resume This allows booting YJIT in a suspended state. We chose to add a new command line option as opposed to simply allowing YJIT.resume to work without any command line option because it allows for combining with YJIT tuning command line options. It also simpifies implementation. Paired with Kokubun and Maxime. * Update yjit.rb Co-authored-by: Takashi Kokubun <[email protected]> --------- Co-authored-by: Alan Wu <[email protected]> Co-authored-by: Takashi Kokubun <[email protected]>
author: Maxime Chevalier-Boisvert <[email protected]> 2023-03-28 15:21:19 -0400
committer: GitHub <[email protected]> 2023-03-28 15:21:19 -0400
commit: 39a34694a0e33e18b4ac6e43cb8042e2d818ecd4 (patch)
tree: 3d70bdeb6fdf015b95d3306130eb52bbff0802a0
parent: 2f8a598dc598b4faaab5d8fd4740811d52fece96 (diff)
7 files changed, 60 insertions, 1 deletions
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 925efdf39f..af7309385f 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -51,6 +51,29 @@ class TestYJIT < Test::Unit::TestCase
     #assert_in_out_err('--yjit-call-threshold=', '', [], /--yjit-call-threshold needs an argument/)
   end
 
+  def test_starting_paused
+    program = <<~RUBY
+      def not_compiled = nil
+      def will_compile = nil
+      def compiled_counts = RubyVM::YJIT.runtime_stats[:compiled_iseq_count]
+      counts = []
+      not_compiled
+      counts << compiled_counts
+
+      RubyVM::YJIT.resume
+
+      will_compile
+      counts << compiled_counts
+
+      if counts[0] == 0 && counts[1] > 0
+        p :ok
+      end
+    RUBY
+    assert_in_out_err(%w[--yjit-pause --yjit-stats --yjit-call-threshold=1], program, success: true) do |stdout, stderr|
+      assert_equal([":ok"], stdout)
+    end
+  end
+
   def test_yjit_stats_and_v_no_error
     _stdout, stderr, _status = EnvUtil.invoke_ruby(%w(-v --yjit-stats), '', true, true)
     refute_includes(stderr, "NoMethodError")
diff --git a/vm.c b/vm.c
index 11fc54095b..31725ae2d0 100644
--- a/vm.c
+++ b/vm.c
@@ -376,7 +376,7 @@ jit_compile(rb_execution_context_t *ec)
     // Increment the ISEQ's call counter
     const rb_iseq_t *iseq = ec->cfp->iseq;
     struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
-    bool yjit_enabled = rb_yjit_enabled_p();
+    bool yjit_enabled = rb_yjit_compile_new_iseqs();
     if (yjit_enabled || rb_rjit_call_p) {
         body->total_calls++;
     }
diff --git a/yjit.c b/yjit.c
index 97ad722270..3740b9003d 100644
--- a/yjit.c
+++ b/yjit.c
@@ -1114,6 +1114,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_resume(rb_execution_context_t *ec, VALUE self);
 
 // Preprocessed yjit.rb generated during build
 #include "yjit.rbinc"
diff --git a/yjit.h b/yjit.h
index 765e735b9d..a640d5982a 100644
--- a/yjit.h
+++ b/yjit.h
@@ -26,6 +26,7 @@
 
 // Expose these as declarations since we are building YJIT.
 bool rb_yjit_enabled_p(void);
+bool rb_yjit_compile_new_iseqs(void);
 unsigned rb_yjit_call_threshold(void);
 void rb_yjit_invalidate_all_method_lookup_assumptions(void);
 void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme);
@@ -48,6 +49,7 @@ void rb_yjit_tracing_invalidate_all(void);
 // In these builds, YJIT could never be turned on. Provide dummy implementations.
 
 static inline bool rb_yjit_enabled_p(void) { return false; }
+static inline bool rb_yjit_compile_new_iseqs(void) { return false; }
 static inline unsigned rb_yjit_call_threshold(void) { return UINT_MAX; }
 static inline void rb_yjit_invalidate_all_method_lookup_assumptions(void) {}
 static inline void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme) {}
diff --git a/yjit.rb b/yjit.rb
index 7d56bf0517..343f5f0ed8 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -29,6 +29,11 @@ module RubyVM::YJIT
     Primitive.rb_yjit_reset_stats_bang
   end
 
+  # Resume YJIT compilation after paused on startup with --yjit-pause
+  def self.resume
+    Primitive.rb_yjit_resume
+  end
+
   # If --yjit-trace-exits is enabled parse the hashes from
   # Primitive.rb_yjit_get_exit_locations into a format readable
   # by Stackprof. This will allow us to find the exact location of a
diff --git a/yjit/src/options.rs b/yjit/src/options.rs
index 9a8205d933..dfae06d1e7 100644
--- a/yjit/src/options.rs
+++ b/yjit/src/options.rs
@@ -28,6 +28,10 @@ pub struct Options {
     // Trace locations of exits
     pub gen_trace_exits: bool,
 
+    // Whether to start YJIT in paused state (initialize YJIT but don't
+    // compile anything)
+    pub pause: bool,
+
     /// Dump compiled and executed instructions for debugging
     pub dump_insns: bool,
 
@@ -50,6 +54,7 @@ pub static mut OPTIONS: Options = Options {
     max_versions: 4,
     gen_stats: false,
     gen_trace_exits: false,
+    pause: false,
     dump_insns: false,
     dump_disasm: None,
     verify_ctx: false,
@@ -132,6 +137,10 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             }
         },
 
+        ("pause", "") => unsafe {
+            OPTIONS.pause = true;
+        },
+
         ("dump-disasm", _) => match opt_val.to_string().as_str() {
             "" => unsafe { OPTIONS.dump_disasm = Some(DumpDisasm::Stdout) },
             directory => {
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index e133049ee8..10bddcc600 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -15,6 +15,9 @@ use std::sync::atomic::{AtomicBool, Ordering};
 /// See [rb_yjit_enabled_p]
 static YJIT_ENABLED: AtomicBool = AtomicBool::new(false);
 
+/// When false, we don't compile new iseqs, but might still service existing branch stubs.
+static COMPILE_NEW_ISEQS: AtomicBool = AtomicBool::new(false);
+
 /// Parse one command-line option.
 /// This is called from ruby.c
 #[no_mangle]
@@ -32,6 +35,11 @@ pub extern "C" fn rb_yjit_enabled_p() -> raw::c_int {
     YJIT_ENABLED.load(Ordering::Acquire).into()
 }
 
+#[no_mangle]
+pub extern "C" fn rb_yjit_compile_new_iseqs() -> bool {
+    COMPILE_NEW_ISEQS.load(Ordering::Acquire).into()
+}
+
 /// Like rb_yjit_enabled_p, but for Rust code.
 pub fn yjit_enabled_p() -> bool {
     YJIT_ENABLED.load(Ordering::Acquire)
@@ -60,6 +68,8 @@ pub extern "C" fn rb_yjit_init_rust() {
 
         // YJIT enabled and initialized successfully
         YJIT_ENABLED.store(true, Ordering::Release);
+
+        COMPILE_NEW_ISEQS.store(!get_option!(pause), Ordering::Release);
     });
 
     if let Err(_) = result {
@@ -134,6 +144,15 @@ pub extern "C" fn rb_yjit_code_gc(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
     Qnil
 }
 
+#[no_mangle]
+pub extern "C" fn rb_yjit_resume(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+    if yjit_enabled_p() {
+        COMPILE_NEW_ISEQS.store(true, Ordering::Release);
+    }
+
+    Qnil
+}
+
 /// Simulate a situation where we are out of executable memory
 #[no_mangle]
 pub extern "C" fn rb_yjit_simulate_oom_bang(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
author	Maxime Chevalier-Boisvert <[email protected]>	2023-03-28 15:21:19 -0400
committer	GitHub <[email protected]>	2023-03-28 15:21:19 -0400
commit	39a34694a0e33e18b4ac6e43cb8042e2d818ecd4 (patch)
tree	3d70bdeb6fdf015b95d3306130eb52bbff0802a0
parent	2f8a598dc598b4faaab5d8fd4740811d52fece96 (diff)