diff --git a/jit_hook.rb b/jit_hook.rb index 487361c049..346b716948 100644 --- a/jit_hook.rb +++ b/jit_hook.rb @@ -3,9 +3,8 @@ class Module # This method is removed in jit_undef.rb. private def with_jit(&block) # :nodoc: # ZJIT currently doesn't compile Array#each properly, so it's disabled for now. - if defined?(RubyVM::ZJIT) && Primitive.rb_zjit_option_enabled_p && false # TODO: remove `&& false` (Shopify/ruby#667) - # We don't support lazily enabling ZJIT yet, so we can call the block right away. - block.call + if defined?(RubyVM::ZJIT) && false # TODO: remove `&& false` (Shopify/ruby#667) + RubyVM::ZJIT.send(:add_jit_hook, block) elsif defined?(RubyVM::YJIT) RubyVM::YJIT.send(:add_jit_hook, block) end diff --git a/ruby.c b/ruby.c index 872a317e3b..f1089ca41e 100644 --- a/ruby.c +++ b/ruby.c @@ -1842,10 +1842,8 @@ ruby_opt_init(ruby_cmdline_options_t *opt) rb_yjit_init(opt->yjit); #endif #if USE_ZJIT - if (opt->zjit) { - extern void rb_zjit_init(void); - rb_zjit_init(); - } + extern void rb_zjit_init(bool); + rb_zjit_init(opt->zjit); #endif ruby_set_script_name(opt->script_name); @@ -2368,6 +2366,12 @@ process_options(int argc, char **argv, ruby_cmdline_options_t *opt) #if USE_ZJIT if (!FEATURE_USED_P(opt->features, zjit) && env_var_truthy("RUBY_ZJIT_ENABLE")) { FEATURE_SET(opt->features, FEATURE_BIT(zjit)); + + // When the --zjit flag is specified, we would have call setup_zjit_options(""), + // which would have called rb_zjit_prepare_options() internally. This ensures we + // go through the same set up but with less overhead than setup_zjit_options(""). + extern void rb_zjit_prepare_options(); + rb_zjit_prepare_options(); } #endif } @@ -2383,10 +2387,9 @@ process_options(int argc, char **argv, ruby_cmdline_options_t *opt) } #endif #if USE_ZJIT - if (FEATURE_SET_P(opt->features, zjit) && !opt->zjit) { - extern void rb_zjit_prepare_options(void); - rb_zjit_prepare_options(); - opt->zjit = true; + if (FEATURE_SET_P(opt->features, zjit)) { + bool rb_zjit_option_enable(void); + opt->zjit = rb_zjit_option_enable(); // set opt->zjit for Init_ruby_description() and calling rb_zjit_init() } #endif diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb index 13ee4a45c8..2e04dbddd5 100644 --- a/test/ruby/test_zjit.rb +++ b/test/ruby/test_zjit.rb @@ -59,6 +59,44 @@ class TestZJIT < Test::Unit::TestCase end end + def test_zjit_enable + assert_separately([], <<~'RUBY') + refute_predicate RubyVM::ZJIT, :enabled? + refute_predicate RubyVM::ZJIT, :stats_enabled? + refute_includes RUBY_DESCRIPTION, "+ZJIT" + + RubyVM::ZJIT.enable + + assert_predicate RubyVM::ZJIT, :enabled? + refute_predicate RubyVM::ZJIT, :stats_enabled? + assert_includes RUBY_DESCRIPTION, "+ZJIT" + RUBY + end + + def test_zjit_disable + assert_separately(["--zjit", "--zjit-disable"], <<~'RUBY') + refute_predicate RubyVM::ZJIT, :enabled? + refute_includes RUBY_DESCRIPTION, "+ZJIT" + + RubyVM::ZJIT.enable + + assert_predicate RubyVM::ZJIT, :enabled? + assert_includes RUBY_DESCRIPTION, "+ZJIT" + RUBY + end + + def test_zjit_enable_respects_existing_options + assert_separately(['--zjit-disable', '--zjit-stats=quiet'], <<~RUBY) + refute_predicate RubyVM::ZJIT, :enabled? + assert_predicate RubyVM::ZJIT, :stats_enabled? + + RubyVM::ZJIT.enable + + assert_predicate RubyVM::ZJIT, :enabled? + assert_predicate RubyVM::ZJIT, :stats_enabled? + RUBY + end + def test_call_itself assert_compiles '42', <<~RUBY, call_threshold: 2 def test = 42.itself diff --git a/version.c b/version.c index b9c57a71b8..366ba1b2b5 100644 --- a/version.c +++ b/version.c @@ -276,6 +276,15 @@ ruby_set_yjit_description(void) define_ruby_description(YJIT_DESCRIPTION); } +void +ruby_set_zjit_description(void) +{ + VALUE mRuby = rb_path2class("Ruby"); + rb_const_remove(rb_cObject, rb_intern("RUBY_DESCRIPTION")); + rb_const_remove(mRuby, rb_intern("DESCRIPTION")); + define_ruby_description(ZJIT_DESCRIPTION); +} + void ruby_show_version(void) { diff --git a/zjit.c b/zjit.c index d1f192801a..b8a89aad1a 100644 --- a/zjit.c +++ b/zjit.c @@ -305,6 +305,7 @@ rb_zjit_class_has_default_allocator(VALUE klass) VALUE rb_vm_get_untagged_block_handler(rb_control_frame_t *reg_cfp); // Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them. +VALUE rb_zjit_enable(rb_execution_context_t *ec, VALUE self); VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self); VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self, VALUE target_key); VALUE rb_zjit_reset_stats_bang(rb_execution_context_t *ec, VALUE self); diff --git a/zjit.rb b/zjit.rb index 7cdf84cfbe..f5ed347f2c 100644 --- a/zjit.rb +++ b/zjit.rb @@ -7,6 +7,8 @@ # This module may not exist if ZJIT does not support the particular platform # for which CRuby is built. module RubyVM::ZJIT + # Blocks that are called when YJIT is enabled + @jit_hooks = [] # Avoid calling a Ruby method here to avoid interfering with compilation tests if Primitive.rb_zjit_print_stats_p at_exit { print_stats } @@ -22,6 +24,18 @@ class << RubyVM::ZJIT Primitive.cexpr! 'RBOOL(rb_zjit_enabled_p)' end + # Enable ZJIT compilation. + def enable + return false if enabled? + + if Primitive.cexpr! 'RBOOL(rb_yjit_enabled_p)' + warn("Only one JIT can be enabled at the same time.") + return false + end + + Primitive.rb_zjit_enable + end + # Check if `--zjit-trace-exits` is used def trace_exit_locations_enabled? Primitive.rb_zjit_trace_exit_locations_enabled_p @@ -234,6 +248,17 @@ class << RubyVM::ZJIT # :stopdoc: private + # Register a block to be called when ZJIT is enabled + def add_jit_hook(hook) + @jit_hooks << hook + end + + # Run ZJIT hooks registered by `#with_jit` + def call_jit_hooks + @jit_hooks.each(&:call) + @jit_hooks.clear + end + def print_counters(keys, buf:, stats:, right_align: false, base: nil) key_pad = keys.map { |key| key.to_s.sub(/_time_ns\z/, '_time').size }.max + 1 key_align = '-' unless right_align diff --git a/zjit/src/options.rs b/zjit/src/options.rs index cd3a643971..c165035eaa 100644 --- a/zjit/src/options.rs +++ b/zjit/src/options.rs @@ -45,7 +45,7 @@ pub struct Options { /// Number of times YARV instructions should be profiled. pub num_profiles: NumProfiles, - /// Enable YJIT statsitics + /// Enable ZJIT statistics pub stats: bool, /// Print stats on exit (when stats is also true) @@ -54,6 +54,10 @@ pub struct Options { /// Enable debug logging pub debug: bool, + // Whether to enable JIT at boot. This option prevents other + // ZJIT tuning options from enabling ZJIT at boot. + pub disable: bool, + /// Turn off the HIR optimizer pub disable_hir_opt: bool, @@ -97,6 +101,7 @@ impl Default for Options { stats: false, print_stats: false, debug: false, + disable: false, disable_hir_opt: false, dump_hir_init: None, dump_hir_opt: None, @@ -123,6 +128,8 @@ pub const ZJIT_OPTIONS: &[(&str, &str)] = &[ ("--zjit-num-profiles=num", "Number of profiled calls before JIT (default: 5)."), ("--zjit-stats[=quiet]", "Enable collecting ZJIT statistics (=quiet to suppress output)."), + ("--zjit-disable", + "Disable ZJIT for lazily enabling it with RubyVM::ZJIT.enable."), ("--zjit-perf", "Dump ISEQ symbols into /tmp/perf-{}.map for Linux perf."), ("--zjit-log-compiled-iseqs=path", "Log compiled ISEQs to the file. The file will be truncated."), @@ -175,7 +182,7 @@ const DUMP_LIR_ALL: &[DumpLIR] = &[ DumpLIR::scratch_split, ]; -/// Mamximum value for --zjit-mem-size/--zjit-exec-mem-size in MiB. +/// Maximum value for --zjit-mem-size/--zjit-exec-mem-size in MiB. /// We set 1TiB just to avoid overflow. We could make it smaller. const MAX_MEM_MIB: usize = 1024 * 1024; @@ -319,6 +326,8 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> { ("debug", "") => options.debug = true, + ("disable", "") => options.disable = true, + ("disable-hir-opt", "") => options.disable_hir_opt = true, // --zjit-dump-hir dumps the actual input to the codegen, which is currently the same as --zjit-dump-hir-opt. @@ -442,15 +451,13 @@ macro_rules! debug { } pub(crate) use debug; -/// Return Qtrue if --zjit* has been specified. For the `#with_jit` hook, -/// this becomes Qtrue before ZJIT is actually initialized and enabled. +/// Return true if ZJIT should be enabled at boot. #[unsafe(no_mangle)] -pub extern "C" fn rb_zjit_option_enabled_p(_ec: EcPtr, _self: VALUE) -> VALUE { - // If any --zjit* option is specified, OPTIONS becomes Some. - if unsafe { OPTIONS.is_some() } { - Qtrue +pub extern "C" fn rb_zjit_option_enable() -> bool { + if unsafe { OPTIONS.as_ref() }.is_some_and(|opts| !opts.disable) { + true } else { - Qfalse + false } } diff --git a/zjit/src/state.rs b/zjit/src/state.rs index 3cb60cffcb..06296eb8f2 100644 --- a/zjit/src/state.rs +++ b/zjit/src/state.rs @@ -1,11 +1,11 @@ //! Runtime state of ZJIT. use crate::codegen::{gen_entry_trampoline, gen_exit_trampoline, gen_exit_trampoline_with_counter, gen_function_stub_hit_trampoline}; -use crate::cruby::{self, rb_bug_panic_hook, rb_vm_insn_count, EcPtr, Qnil, rb_vm_insn_addr2opcode, rb_profile_frames, VALUE, VM_INSTRUCTION_SIZE, size_t, rb_gc_mark}; +use crate::cruby::{self, rb_bug_panic_hook, rb_vm_insn_count, src_loc, EcPtr, Qnil, Qtrue, rb_vm_insn_addr2opcode, rb_profile_frames, VALUE, VM_INSTRUCTION_SIZE, size_t, rb_gc_mark, with_vm_lock}; use crate::cruby_methods; use crate::invariants::Invariants; use crate::asm::CodeBlock; -use crate::options::get_option; +use crate::options::{get_option, rb_zjit_prepare_options}; use crate::stats::{Counters, InsnCounters, SideExitLocations}; use crate::virtualmem::CodePtr; use std::collections::HashMap; @@ -63,12 +63,42 @@ pub struct ZJITState { exit_locations: Option, } +/// Tracks the initialization progress +enum InitializationState { + Uninitialized, + + /// At boot time, rb_zjit_init will be called regardless of whether + /// ZJIT is enabled, in this phase we initialize any states that must + /// be captured at during boot. + Initialized(cruby_methods::Annotations), + + /// When ZJIT is enabled, either during boot with `--zjit`, or lazily + /// at a later time with `RubyVM::ZJIT.enable`, we perform the rest + /// of the initialization steps and produce the `ZJITState` instance. + Enabled(ZJITState), + + /// Indicates that ZJITState::init has panicked. Should never be + /// encountered in practice since we abort immediately when that + /// happens. + Panicked, +} + /// Private singleton instance of the codegen globals -static mut ZJIT_STATE: Option = None; +static mut ZJIT_STATE: InitializationState = InitializationState::Uninitialized; impl ZJITState { /// Initialize the ZJIT globals. Return the address of the JIT entry trampoline. pub fn init() -> *const u8 { + use InitializationState::*; + + let initialization_state = unsafe { + std::mem::replace(&mut ZJIT_STATE, Panicked) + }; + + let Initialized(method_annotations) = initialization_state else { + panic!("rb_zjit_init was never called"); + }; + let mut cb = { use crate::options::*; use crate::virtualmem::*; @@ -99,7 +129,7 @@ impl ZJITState { send_fallback_counters: [0; VM_INSTRUCTION_SIZE as usize], invariants: Invariants::default(), assert_compiles: false, - method_annotations: cruby_methods::init(), + method_annotations, exit_trampoline, function_stub_hit_trampoline, exit_trampoline_with_counter: exit_trampoline, @@ -107,7 +137,7 @@ impl ZJITState { not_annotated_frame_cfunc_counter_pointers: HashMap::new(), exit_locations, }; - unsafe { ZJIT_STATE = Some(zjit_state); } + unsafe { ZJIT_STATE = Enabled(zjit_state); } // With --zjit-stats, use a different trampoline on function stub exits // to count exit_compilation_failure. Note that the trampoline code depends @@ -123,12 +153,16 @@ impl ZJITState { /// Return true if zjit_state has been initialized pub fn has_instance() -> bool { - unsafe { ZJIT_STATE.as_mut().is_some() } + matches!(unsafe { &ZJIT_STATE }, InitializationState::Enabled(_)) } /// Get a mutable reference to the codegen globals instance fn get_instance() -> &'static mut ZJITState { - unsafe { ZJIT_STATE.as_mut().unwrap() } + if let InitializationState::Enabled(instance) = unsafe { &mut ZJIT_STATE } { + instance + } else { + panic!("ZJITState::get_instance called when ZJIT is not enabled") + } } /// Get a mutable reference to the inline code block @@ -249,14 +283,39 @@ impl ZJITState { } } -/// Initialize ZJIT +/// Initialize ZJIT at boot. This is called even if ZJIT is disabled. #[unsafe(no_mangle)] -pub extern "C" fn rb_zjit_init() { +pub extern "C" fn rb_zjit_init(zjit_enabled: bool) { + use InitializationState::*; + + debug_assert!( + matches!(unsafe { &ZJIT_STATE }, Uninitialized), + "rb_zjit_init should only be called once during boot", + ); + + // Initialize IDs and method annotations. + // cruby_methods::init() must be called at boot, + // as cmes could have been re-defined after boot. + cruby::ids::init(); + + let method_annotations = cruby_methods::init(); + + unsafe { ZJIT_STATE = Initialized(method_annotations); } + + // If --zjit, enable ZJIT immediately + if zjit_enabled { + zjit_enable(); + } +} + +/// Enable ZJIT compilation. +fn zjit_enable() { + // TODO: call RubyVM::ZJIT::call_jit_hooks here + // Catch panics to avoid UB for unwinding into C frames. // See https://doc.rust-lang.org/nomicon/exception-safety.html let result = std::panic::catch_unwind(|| { // Initialize ZJIT states - cruby::ids::init(); let zjit_entry = ZJITState::init(); // Install a panic hook for ZJIT @@ -271,11 +330,33 @@ pub extern "C" fn rb_zjit_init() { }); if result.is_err() { - println!("ZJIT: zjit_init() panicked. Aborting."); + println!("ZJIT: zjit_enable() panicked. Aborting."); std::process::abort(); } } +/// Enable ZJIT compilation, returning Qtrue if ZJIT was previously disabled +#[unsafe(no_mangle)] +pub extern "C" fn rb_zjit_enable(_ec: EcPtr, _self: VALUE) -> VALUE { + with_vm_lock(src_loc!(), || { + // Options would not have been initialized during boot if no flags were specified + rb_zjit_prepare_options(); + + // Initialize and enable ZJIT + zjit_enable(); + + // Add "+ZJIT" to RUBY_DESCRIPTION + unsafe { + unsafe extern "C" { + fn ruby_set_zjit_description(); + } + ruby_set_zjit_description(); + } + + Qtrue + }) +} + /// Assert that any future ZJIT compilation will return a function pointer (not fail to compile) #[unsafe(no_mangle)] pub extern "C" fn rb_zjit_assert_compiles(_ec: EcPtr, _self: VALUE) -> VALUE {