YJIT: Add RubyVM::YJIT.enable (#8705)

11 files changed, 110 insertions, 108 deletions

diff --git a/cont.c b/cont.c
index db8508e724..2ba6fe216b 100644
--- a/cont.c
+++ b/cont.c
@@ -71,8 +71,6 @@ static VALUE rb_cFiberPool;
 #define FIBER_POOL_ALLOCATION_FREE
 #endif
 
-#define jit_cont_enabled (rb_rjit_enabled || rb_yjit_enabled_p())
-
 enum context_type {
     CONTINUATION_CONTEXT = 0,
     FIBER_CONTEXT = 1
@@ -1062,10 +1060,8 @@ cont_free(void *ptr)
 
     RUBY_FREE_UNLESS_NULL(cont->saved_vm_stack.ptr);
 
-    if (jit_cont_enabled) {
-        VM_ASSERT(cont->jit_cont != NULL);
-        jit_cont_free(cont->jit_cont);
-    }
+    VM_ASSERT(cont->jit_cont != NULL);
+    jit_cont_free(cont->jit_cont);
     /* free rb_cont_t or rb_fiber_t */
     ruby_xfree(ptr);
     RUBY_FREE_LEAVE("cont");
@@ -1311,9 +1307,6 @@ rb_jit_cont_each_iseq(rb_iseq_callback callback, void *data)
 void
 rb_jit_cont_finish(void)
 {
-    if (!jit_cont_enabled)
-        return;
-
     struct rb_jit_cont *cont, *next;
     for (cont = first_jit_cont; cont != NULL; cont = next) {
         next = cont->next;
@@ -1326,9 +1319,8 @@ static void
 cont_init_jit_cont(rb_context_t *cont)
 {
     VM_ASSERT(cont->jit_cont == NULL);
-    if (jit_cont_enabled) {
-        cont->jit_cont = jit_cont_new(&(cont->saved_ec));
-    }
+    // We always allocate this since YJIT may be enabled later
+    cont->jit_cont = jit_cont_new(&(cont->saved_ec));
 }
 
 struct rb_execution_context_struct *
@@ -1375,15 +1367,11 @@ rb_fiberptr_blocking(struct rb_fiber_struct *fiber)
     return fiber->blocking;
 }
 
-// Start working with jit_cont.
+// Initialize the jit_cont_lock
 void
 rb_jit_cont_init(void)
 {
-    if (!jit_cont_enabled)
-        return;
-
     rb_native_mutex_initialize(&jit_cont_lock);
-    cont_init_jit_cont(&GET_EC()->fiber_ptr->cont);
 }
 
 #if 0
@@ -2564,10 +2552,9 @@ rb_threadptr_root_fiber_setup(rb_thread_t *th)
     fiber->killed = 0;
     fiber_status_set(fiber, FIBER_RESUMED); /* skip CREATED */
     th->ec = &fiber->cont.saved_ec;
-    // When rb_threadptr_root_fiber_setup is called for the first time, rb_rjit_enabled and
-    // rb_yjit_enabled_p() are still false. So this does nothing and rb_jit_cont_init() that is
-    // called later will take care of it. However, you still have to call cont_init_jit_cont()
-    // here for other Ractors, which are not initialized by rb_jit_cont_init().
+    // This is the first fiber. Hence it's the first jit_cont_new() as well.
+    // Initialize the mutex for jit_cont_new() in cont_init_jit_cont().
+    rb_jit_cont_init();
     cont_init_jit_cont(&fiber->cont);
 }
 
diff --git a/ruby.c b/ruby.c
index 0efe4efbf4..b4f78b3415 100644
--- a/ruby.c
+++ b/ruby.c
@@ -1796,10 +1796,6 @@ ruby_opt_init(ruby_cmdline_options_t *opt)
     if (opt->yjit)
         rb_yjit_init();
 #endif
-    // rb_threadptr_root_fiber_setup for the initial thread is called before rb_yjit_enabled_p()
-    // or rjit_enabled becomes true, meaning jit_cont_new is skipped for the initial root fiber.
-    // Therefore we need to call this again here to set the initial root fiber's jit_cont.
-    rb_jit_cont_init(); // must be after rjit_enabled = true and rb_yjit_init()
 
     ruby_set_script_name(opt->script_name);
     require_libraries(&opt->req_list);
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 4066664600..fa357b4977 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -51,27 +51,36 @@ 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 test_yjit_enable
+    args = []
+    args << "--disable=yjit" if RubyVM::YJIT.enabled?
+    assert_separately(args, <<~RUBY)
+      assert_false RubyVM::YJIT.enabled?
+      assert_false RUBY_DESCRIPTION.include?("+YJIT")
+
+      RubyVM::YJIT.enable
+
+      assert_true RubyVM::YJIT.enabled?
+      assert_true RUBY_DESCRIPTION.include?("+YJIT")
+    RUBY
+  end
+
+  def test_yjit_enable_with_call_threshold
+    assert_separately(%w[--yjit-disable --yjit-call-threshold=1], <<~RUBY)
       def not_compiled = nil
       def will_compile = nil
-      def compiled_counts = RubyVM::YJIT.runtime_stats[:compiled_iseq_count]
-      counts = []
+      def compiled_counts = RubyVM::YJIT.runtime_stats&.dig(:compiled_iseq_count)
+
       not_compiled
-      counts << compiled_counts
+      assert_nil compiled_counts
+      assert_false RubyVM::YJIT.enabled?
 
-      RubyVM::YJIT.resume
+      RubyVM::YJIT.enable
 
       will_compile
-      counts << compiled_counts
-
-      if counts[0] == 0 && counts[1] > 0
-        p :ok
-      end
+      assert compiled_counts > 0
+      assert_true RubyVM::YJIT.enabled?
     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
diff --git a/version.c b/version.c
index 5f4db2a622..b5b0bfda31 100644
--- a/version.c
+++ b/version.c
@@ -141,8 +141,8 @@ Init_version(void)
 
 int ruby_mn_threads_enabled;
 
-void
-Init_ruby_description(ruby_cmdline_options_t *opt)
+static void
+define_ruby_description(const char *const jit_opt)
 {
     static char desc[
         sizeof(ruby_description)
@@ -150,11 +150,6 @@ Init_ruby_description(ruby_cmdline_options_t *opt)
         + rb_strlen_lit(" +MN")
         ];
 
-    const char *const jit_opt =
-        RJIT_OPTS_ON ? " +RJIT" :
-        YJIT_OPTS_ON ? YJIT_DESCRIPTION :
-        "";
-
     const char *const threads_opt = ruby_mn_threads_enabled ? " +MN" : "";
 
     int n = snprintf(desc, sizeof(desc),
@@ -177,6 +172,23 @@ Init_ruby_description(ruby_cmdline_options_t *opt)
 }
 
 void
+Init_ruby_description(ruby_cmdline_options_t *opt)
+{
+    const char *const jit_opt =
+        RJIT_OPTS_ON ? " +RJIT" :
+        YJIT_OPTS_ON ? YJIT_DESCRIPTION :
+        "";
+    define_ruby_description(jit_opt);
+}
+
+void
+ruby_set_yjit_description(void)
+{
+    rb_const_remove(rb_cObject, rb_intern("RUBY_DESCRIPTION"));
+    define_ruby_description(YJIT_DESCRIPTION);
+}
+
+void
 ruby_show_version(void)
 {
     puts(rb_dynamic_description);
diff --git a/vm.c b/vm.c
index fd4d7f933a..456af58bfe 100644
--- a/vm.c
+++ b/vm.c
@@ -426,15 +426,14 @@ jit_compile(rb_execution_context_t *ec)
 {
     const rb_iseq_t *iseq = ec->cfp->iseq;
     struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
-    bool yjit_enabled = rb_yjit_compile_new_iseqs();
-    if (!(yjit_enabled || rb_rjit_call_p)) {
+    if (!(rb_yjit_enabled_p || rb_rjit_call_p)) {
         return NULL;
     }
 
     // Increment the ISEQ's call counter and trigger JIT compilation if not compiled
     if (body->jit_entry == NULL) {
         body->jit_entry_calls++;
-        if (yjit_enabled) {
+        if (rb_yjit_enabled_p) {
             if (rb_yjit_threshold_hit(iseq, body->jit_entry_calls)) {
                 rb_yjit_compile_iseq(iseq, ec, false);
             }
@@ -476,7 +475,7 @@ jit_compile_exception(rb_execution_context_t *ec)
 {
     const rb_iseq_t *iseq = ec->cfp->iseq;
     struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
-    if (!rb_yjit_compile_new_iseqs()) {
+    if (!rb_yjit_enabled_p) {
         return NULL;
     }
 
diff --git a/vm_method.c b/vm_method.c
index e2a3907c33..d2b642e6e1 100644
--- a/vm_method.c
+++ b/vm_method.c
@@ -200,7 +200,7 @@ clear_method_cache_by_id_in_class(VALUE klass, ID mid)
         struct rb_id_table *cm_tbl;
         if ((cm_tbl = RCLASS_CALLABLE_M_TBL(klass)) != NULL) {
             VALUE cme;
-            if (rb_yjit_enabled_p() && rb_id_table_lookup(cm_tbl, mid, &cme)) {
+            if (rb_yjit_enabled_p && rb_id_table_lookup(cm_tbl, mid, &cme)) {
                 rb_yjit_cme_invalidate((rb_callable_method_entry_t *)cme);
             }
             if (rb_rjit_enabled && rb_id_table_lookup(cm_tbl, mid, &cme)) {
diff --git a/yjit.c b/yjit.c
index 7c9a933c20..baefed41cf 100644
--- a/yjit.c
+++ b/yjit.c
@@ -1171,20 +1171,15 @@ 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);
+VALUE rb_yjit_enable(rb_execution_context_t *ec, VALUE self);
 
 // Preprocessed yjit.rb generated during build
 #include "yjit.rbinc"
 
-// Can raise RuntimeError
+// Initialize the GC hooks
 void
-rb_yjit_init(void)
+rb_yjit_init_gc_hooks(void)
 {
-    // Call the Rust initialization code
-    void rb_yjit_init_rust(void);
-    rb_yjit_init_rust();
-
-    // Initialize the GC hooks. Do this second as some code depend on Rust initialization.
     struct yjit_root_struct *root;
     VALUE yjit_root = TypedData_Make_Struct(0, struct yjit_root_struct, &yjit_root_type, root);
     rb_gc_register_mark_object(yjit_root);
diff --git a/yjit.h b/yjit.h
index 33b9214952..f683c9c1f7 100644
--- a/yjit.h
+++ b/yjit.h
@@ -28,9 +28,8 @@
 extern uint64_t rb_yjit_call_threshold;
 extern uint64_t rb_yjit_cold_threshold;
 extern uint64_t rb_yjit_live_iseq_count;
+extern bool rb_yjit_enabled_p;
 void rb_yjit_incr_counter(const char *counter_name);
-bool rb_yjit_enabled_p(void);
-bool rb_yjit_compile_new_iseqs(void);
 void rb_yjit_invalidate_all_method_lookup_assumptions(void);
 void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme);
 void rb_yjit_collect_binding_alloc(void);
@@ -51,9 +50,8 @@ void rb_yjit_show_usage(int help, int highlight, unsigned int width, int columns
 // !USE_YJIT
 // In these builds, YJIT could never be turned on. Provide dummy implementations.
 
+#define rb_yjit_enabled_p false
 static inline void rb_yjit_incr_counter(const char *counter_name) {}
-static inline bool rb_yjit_enabled_p(void) { return false; }
-static inline bool rb_yjit_compile_new_iseqs(void) { return false; }
 static inline void rb_yjit_invalidate_all_method_lookup_assumptions(void) {}
 static inline void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme) {}
 static inline void rb_yjit_collect_binding_alloc(void) {}
diff --git a/yjit.rb b/yjit.rb
index 3374e4419c..295e88ef34 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -11,7 +11,7 @@
 module RubyVM::YJIT
   # Check if YJIT is enabled
   def self.enabled?
-    Primitive.cexpr! 'RBOOL(rb_yjit_enabled_p())'
+    Primitive.cexpr! 'RBOOL(rb_yjit_enabled_p)'
   end
 
   # Check if --yjit-stats is used.
@@ -29,9 +29,9 @@ 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
+  # Enable YJIT compilation.
+  def self.enable
+    Primitive.rb_yjit_enable
   end
 
   # If --yjit-trace-exits is enabled parse the hashes from
diff --git a/yjit/src/options.rs b/yjit/src/options.rs
index e5e0552d7e..c4f3e8df3a 100644
--- a/yjit/src/options.rs
+++ b/yjit/src/options.rs
@@ -47,9 +47,9 @@ pub struct Options {
     // how often to sample exit trace data
     pub trace_exits_sample_rate: usize,
 
-    // Whether to start YJIT in paused state (initialize YJIT but don't
-    // compile anything)
-    pub pause: bool,
+    // Whether to enable YJIT at boot. This option prevents other
+    // YJIT tuning options from enabling YJIT at boot.
+    pub disable: bool,
 
     /// Dump compiled and executed instructions for debugging
     pub dump_insns: bool,
@@ -81,7 +81,7 @@ pub static mut OPTIONS: Options = Options {
     gen_trace_exits: false,
     print_stats: true,
     trace_exits_sample_rate: 0,
-    pause: false,
+    disable: false,
     dump_insns: false,
     dump_disasm: None,
     verify_ctx: false,
@@ -186,8 +186,8 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             }
         },
 
-        ("pause", "") => unsafe {
-            OPTIONS.pause = true;
+        ("disable", "") => unsafe {
+            OPTIONS.disable = true;
         },
 
         ("temp-regs", _) => match opt_val.parse() {
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index 515fa75ce8..813918b4bc 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -8,16 +8,12 @@ use crate::stats::incr_counter;
 use crate::stats::with_compile_time;
 
 use std::os::raw;
-use std::sync::atomic::{AtomicBool, Ordering};
 
-/// For tracking whether the user enabled YJIT through command line arguments or environment
-/// variables. AtomicBool to avoid `unsafe`. On x86 it compiles to simple movs.
-/// See <https://doc.rust-lang.org/std/sync/atomic/enum.Ordering.html>
-/// 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);
+/// Is YJIT on? The interpreter uses this variable to decide whether to trigger
+/// compilation. See jit_exec() and jit_compile().
+#[allow(non_upper_case_globals)]
+#[no_mangle]
+pub static mut rb_yjit_enabled_p: bool = false;
 
 /// Parse one command-line option.
 /// This is called from ruby.c
@@ -26,29 +22,22 @@ pub extern "C" fn rb_yjit_parse_option(str_ptr: *const raw::c_char) -> bool {
     return parse_option(str_ptr).is_some();
 }
 
-/// Is YJIT on? The interpreter uses this function to decide whether to increment
-/// ISEQ call counters. See jit_exec().
-/// This is used frequently since it's used on every method call in the interpreter.
-#[no_mangle]
-pub extern "C" fn rb_yjit_enabled_p() -> raw::c_int {
-    // Note that we might want to call this function from signal handlers so
-    // might need to ensure signal-safety(7).
-    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)
+    unsafe { rb_yjit_enabled_p }
 }
 
 /// This function is called from C code
 #[no_mangle]
-pub extern "C" fn rb_yjit_init_rust() {
+pub extern "C" fn rb_yjit_init() {
+    // If --yjit-disable, yjit_init() will not be called until RubyVM::YJIT.enable.
+    if !get_option!(disable) {
+        yjit_init();
+    }
+}
+
+/// Initialize and enable YJIT. You should call this at boot or with GVL.
+fn yjit_init() {
     // TODO: need to make sure that command-line options have been
     // initialized by CRuby
 
@@ -63,13 +52,12 @@ pub extern "C" fn rb_yjit_init_rust() {
         rb_bug_panic_hook();
 
         // YJIT enabled and initialized successfully
-        YJIT_ENABLED.store(true, Ordering::Release);
-
-        COMPILE_NEW_ISEQS.store(!get_option!(pause), Ordering::Release);
+        assert!(unsafe{ !rb_yjit_enabled_p });
+        unsafe { rb_yjit_enabled_p = true; }
     });
 
     if let Err(_) = result {
-        println!("YJIT: rb_yjit_init_rust() panicked. Aborting.");
+        println!("YJIT: yjit_init() panicked. Aborting.");
         std::process::abort();
     }
 
@@ -79,6 +67,12 @@ pub extern "C" fn rb_yjit_init_rust() {
         let _ = std::fs::remove_file(&perf_map);
         println!("YJIT perf map: {perf_map}");
     }
+
+    // Initialize the GC hooks. Do this at last as some code depend on Rust initialization.
+    extern "C" {
+        fn rb_yjit_init_gc_hooks();
+    }
+    unsafe { rb_yjit_init_gc_hooks() }
 }
 
 /// At the moment, we abort in all cases we panic.
@@ -161,13 +155,25 @@ pub extern "C" fn rb_yjit_code_gc(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
     Qnil
 }
 
+/// Enable YJIT compilation, returning true if YJIT was previously disabled
 #[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);
-    }
+pub extern "C" fn rb_yjit_enable(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+    with_vm_lock(src_loc!(), || {
+        if yjit_enabled_p() {
+            return Qfalse;
+        }
 
-    Qnil
+        // Initialize and enable YJIT if currently disabled
+        yjit_init();
+
+        // Add "+YJIT" to RUBY_DESCRIPTION
+        extern "C" {
+            fn ruby_set_yjit_description();
+        }
+        unsafe { ruby_set_yjit_description(); }
+
+        Qtrue
+    })
 }
 
 /// Simulate a situation where we are out of executable memory