archive/ruby
1
0
Fork 0
mirror of https://github.com/ruby/ruby.git synced 2025-08-15 05:29:10 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

ZJIT: Prepare for sharing JIT hooks with ZJIT (#14044)

Browse source
This commit is contained in:
Takashi Kokubun 2025-07-30 10:11:10 -07:00 • committed by GitHub
parent 4263c49d1c
commit 2cd10de330
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 142 additions and 118 deletions
Download patch file Download diff file Expand all files Collapse all files

2
array.rb
View file

@ -212,7 +212,7 @@ class Array
indexes
end
with_yjit do
with_jit do
if Primitive.rb_builtin_basic_definition_p(:each)
undef :each

3
common.mk
View file

@ -1236,8 +1236,9 @@ BUILTIN_RB_SRCS = \
$(srcdir)/nilclass.rb \
$(srcdir)/prelude.rb \
$(srcdir)/gem_prelude.rb \
$(srcdir)/jit_hook.rb \
$(srcdir)/jit_undef.rb \
$(srcdir)/yjit.rb \
$(srcdir)/yjit_hook.rb \
$(srcdir)/zjit.rb \
$(empty)
BUILTIN_RB_INCS = $(BUILTIN_RB_SRCS:.rb=.rbinc)

6
depend
View file

@ -9196,6 +9196,8 @@ miniinit.$(OBJEXT): {$(VPATH)}internal/warning_push.h
miniinit.$(OBJEXT): {$(VPATH)}internal/xmalloc.h
miniinit.$(OBJEXT): {$(VPATH)}io.rb
miniinit.$(OBJEXT): {$(VPATH)}iseq.h
miniinit.$(OBJEXT): {$(VPATH)}jit_hook.rb
miniinit.$(OBJEXT): {$(VPATH)}jit_undef.rb
miniinit.$(OBJEXT): {$(VPATH)}kernel.rb
miniinit.$(OBJEXT): {$(VPATH)}marshal.rb
miniinit.$(OBJEXT): {$(VPATH)}method.h
@ -9232,7 +9234,6 @@ miniinit.$(OBJEXT): {$(VPATH)}vm_core.h
miniinit.$(OBJEXT): {$(VPATH)}vm_opts.h
miniinit.$(OBJEXT): {$(VPATH)}warning.rb
miniinit.$(OBJEXT): {$(VPATH)}yjit.rb
miniinit.$(OBJEXT): {$(VPATH)}yjit_hook.rb
miniinit.$(OBJEXT): {$(VPATH)}zjit.rb
namespace.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
namespace.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
@ -18755,6 +18756,8 @@ vm.$(OBJEXT): {$(VPATH)}internal/variable.h
vm.$(OBJEXT): {$(VPATH)}internal/warning_push.h
vm.$(OBJEXT): {$(VPATH)}internal/xmalloc.h
vm.$(OBJEXT): {$(VPATH)}iseq.h
vm.$(OBJEXT): {$(VPATH)}jit_hook.rbinc
vm.$(OBJEXT): {$(VPATH)}jit_undef.rbinc
vm.$(OBJEXT): {$(VPATH)}method.h
vm.$(OBJEXT): {$(VPATH)}missing.h
vm.$(OBJEXT): {$(VPATH)}node.h
@ -18797,7 +18800,6 @@ vm.$(OBJEXT): {$(VPATH)}vm_opts.h
vm.$(OBJEXT): {$(VPATH)}vm_sync.h
vm.$(OBJEXT): {$(VPATH)}vmtc.inc
vm.$(OBJEXT): {$(VPATH)}yjit.h
vm.$(OBJEXT): {$(VPATH)}yjit_hook.rbinc
vm.$(OBJEXT): {$(VPATH)}zjit.h
vm_backtrace.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
vm_backtrace.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h

7
inits.c
View file

@ -88,8 +88,10 @@ void
rb_call_builtin_inits(void)
{
#define BUILTIN(n) CALL(builtin_##n)
BUILTIN(kernel);
BUILTIN(jit_hook);
BUILTIN(yjit);
BUILTIN(zjit);
BUILTIN(kernel);
BUILTIN(gc);
BUILTIN(ractor);
BUILTIN(numeric);
@ -107,8 +109,7 @@ rb_call_builtin_inits(void)
BUILTIN(thread_sync);
BUILTIN(nilclass);
BUILTIN(marshal);
BUILTIN(zjit);
BUILTIN(yjit_hook);
BUILTIN(jit_undef);
Init_builtin_prelude();
}
#undef CALL

6
internal/cmdlineopt.h
View file

@ -23,9 +23,6 @@ typedef struct ruby_cmdline_options {
ruby_features_t warn;
unsigned int dump;
long backtrace_length_limit;
#if USE_ZJIT
void *zjit;
#endif
const char *crash_report;
@ -42,6 +39,9 @@ typedef struct ruby_cmdline_options {
#if USE_YJIT
unsigned int yjit: 1;
#endif
#if USE_ZJIT
unsigned int zjit: 1;
#endif
} ruby_cmdline_options_t;
struct ruby_opt_message {

13
jit_hook.rb Normal file
View file

@ -0,0 +1,13 @@
class Module
# Internal helper for built-in initializations to define methods only when JIT is enabled.
# 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
elsif defined?(RubyVM::YJIT)
RubyVM::YJIT.send(:add_jit_hook, block)
end
end
end

4
jit_undef.rb Normal file
View file

@ -0,0 +1,4 @@
# Remove the helper defined in jit_hook.rb
class Module
undef :with_jit
end

10
kernel.rb
View file

@ -291,13 +291,3 @@ module Kernel
end
end
end
class Module
# Internal helper for built-in initializations to define methods only when YJIT is enabled.
# This method is removed in yjit_hook.rb.
private def with_yjit(&block) # :nodoc:
if defined?(RubyVM::YJIT)
RubyVM::YJIT.send(:add_yjit_hook, block)
end
end
end

2
numeric.rb
View file

@ -322,7 +322,7 @@ class Integer
1
end
with_yjit do
with_jit do
if Primitive.rb_builtin_basic_definition_p(:downto)
undef :downto

19
ruby.c
View file

@ -1196,14 +1196,12 @@ setup_yjit_options(const char *s)
#if USE_ZJIT
static void
setup_zjit_options(ruby_cmdline_options_t *opt, const char *s)
setup_zjit_options(const char *s)
{
// The option parsing is done in zjit/src/options.rs
extern void *rb_zjit_init_options(void);
extern bool rb_zjit_parse_option(void *options, const char *s);
extern bool rb_zjit_parse_option(const char *s);
if (!opt->zjit) opt->zjit = rb_zjit_init_options();
if (!rb_zjit_parse_option(opt->zjit, s)) {
if (!rb_zjit_parse_option(s)) {
rb_raise(rb_eRuntimeError, "invalid ZJIT option '%s' (--help will show valid zjit options)", s);
}
}
@ -1481,7 +1479,7 @@ proc_long_options(ruby_cmdline_options_t *opt, const char *s, long argc, char **
else if (is_option_with_optarg("zjit", '-', true, false, false)) {
#if USE_ZJIT
FEATURE_SET(opt->features, FEATURE_BIT(zjit));
setup_zjit_options(opt, s);
setup_zjit_options(s);
#else
rb_warn("Ruby was built without ZJIT support."
" You may need to install rustc to build Ruby with ZJIT.");
@ -1828,8 +1826,8 @@ ruby_opt_init(ruby_cmdline_options_t *opt)
#endif
#if USE_ZJIT
if (opt->zjit) {
extern void rb_zjit_init(void *options);
rb_zjit_init(opt->zjit);
extern void rb_zjit_init(void);
rb_zjit_init();
}
#endif
@ -2370,8 +2368,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_init_options(void);
opt->zjit = rb_zjit_init_options();
extern void rb_zjit_prepare_options(void);
rb_zjit_prepare_options();
opt->zjit = true;
}
#endif

7
test/ruby/test_zjit.rb
View file

@ -1084,6 +1084,13 @@ class TestZJIT < Test::Unit::TestCase
}, call_threshold: 2
end
def test_zjit_option_uses_array_each_in_ruby
omit 'ZJIT wrongly compiles Array#each, so it is disabled for now'
assert_runs '"<internal:array>"', %q{
Array.instance_method(:each).source_location&.first
}
end
def test_profile_under_nested_jit_call
assert_compiles '[nil, nil, 3]', %q{
def profile

13
vm.c
View file

@ -4509,14 +4509,21 @@ Init_vm_objects(void)
vm->cc_refinement_table = rb_set_init_numtable();
}
#if USE_ZJIT
extern VALUE rb_zjit_option_enabled_p(rb_execution_context_t *ec, VALUE self);
#else
static VALUE rb_zjit_option_enabled_p(rb_execution_context_t *ec, VALUE self) { return Qfalse; }
#endif
// Whether JIT is enabled or not, we need to load/undef `#with_jit` for other builtins.
#include "jit_hook.rbinc"
#include "jit_undef.rbinc"
// Stub for builtin function when not building YJIT units
#if !USE_YJIT
void Init_builtin_yjit(void) {}
#endif
// Whether YJIT is enabled or not, we load yjit_hook.rb to remove Kernel#with_yjit.
#include "yjit_hook.rbinc"
// Stub for builtin function when not building ZJIT units
#if !USE_ZJIT
void Init_builtin_zjit(void) {}

2
vm_method.c
View file

@ -775,7 +775,7 @@ rb_method_definition_set(const rb_method_entry_t *me, rb_method_definition_t *de
/* setup iseq first (before invoking GC) */
RB_OBJ_WRITE(me, &def->body.iseq.iseqptr, iseq);
// Methods defined in `with_yjit` should be considered METHOD_ENTRY_BASIC
// Methods defined in `with_jit` should be considered METHOD_ENTRY_BASIC
if (rb_iseq_attr_p(iseq, BUILTIN_ATTR_C_TRACE)) {
METHOD_ENTRY_BASIC_SET((rb_method_entry_t *)me, TRUE);
}

14
yjit.rb
View file

@ -264,23 +264,23 @@ module RubyVM::YJIT
end
# Blocks that are called when YJIT is enabled
@yjit_hooks = []
@jit_hooks = []
class << self
# :stopdoc:
private
# Register a block to be called when YJIT is enabled
def add_yjit_hook(hook)
@yjit_hooks << hook
def add_jit_hook(hook)
@jit_hooks << hook
end
# Run YJIT hooks registered by RubyVM::YJIT.with_yjit
def call_yjit_hooks
# Run YJIT hooks registered by `#with_jit`
def call_jit_hooks
# Skip using builtin methods in Ruby if --yjit-c-builtin is given
return if Primitive.yjit_c_builtin_p
@yjit_hooks.each(&:call)
@yjit_hooks.clear
@jit_hooks.each(&:call)
@jit_hooks.clear
end
# Print stats and dump exit locations

2
yjit/src/options.rs
View file

@ -46,7 +46,7 @@ pub struct Options {
// The number of registers allocated for stack temps
pub num_temp_regs: usize,
// Disable Ruby builtin methods defined by `with_yjit` hooks, e.g. Array#each in Ruby
// Disable Ruby builtin methods defined by `with_jit` hooks, e.g. Array#each in Ruby
pub c_builtin: bool,
// Capture stats

2
yjit/src/yjit.rs
View file

@ -57,7 +57,7 @@ fn yjit_init() {
// Call YJIT hooks before enabling YJIT to avoid compiling the hooks themselves
unsafe {
let yjit = rb_const_get(rb_cRubyVM, rust_str_to_id("YJIT"));
rb_funcall(yjit, rust_str_to_id("call_yjit_hooks"), 0);
rb_funcall(yjit, rust_str_to_id("call_jit_hooks"), 0);
}
// Catch panics to avoid UB for unwinding into C frames.

4
yjit_hook.rb
View file

@ -1,4 +0,0 @@
# Remove the helper defined in kernel.rb
class Module
undef :with_yjit
end

5
zjit/src/cruby.rs
View file

@ -957,7 +957,7 @@ pub use manual_defs::*;
pub mod test_utils {
use std::{ptr::null, sync::Once};
use crate::{options::init_options, state::rb_zjit_enabled_p, state::ZJITState};
use crate::{options::rb_zjit_prepare_options, state::rb_zjit_enabled_p, state::ZJITState};
use super::*;
@ -979,6 +979,7 @@ pub mod test_utils {
// <https://github.com/Shopify/zjit/pull/37>, though
let mut var: VALUE = Qnil;
ruby_init_stack(&mut var as *mut VALUE as *mut _);
rb_zjit_prepare_options(); // enable `#with_jit` on builtins
ruby_init();
// Pass command line options so the VM loads core library methods defined in
@ -994,7 +995,7 @@ pub mod test_utils {
}
// Set up globals for convenience
ZJITState::init(init_options());
ZJITState::init();
// Enable zjit_* instructions
unsafe { rb_zjit_enabled_p = true; }

77
zjit/src/options.rs
View file

@ -16,9 +16,9 @@ pub static mut rb_zjit_profile_threshold: u64 = 1;
#[allow(non_upper_case_globals)]
pub static mut rb_zjit_call_threshold: u64 = 2;
/// True if --zjit-stats is enabled.
#[allow(non_upper_case_globals)]
static mut zjit_stats_enabled_p: bool = false;
/// ZJIT command-line options. This is set before rb_zjit_init() sets
/// ZJITState so that we can query some options while loading builtins.
pub static mut OPTIONS: Option<Options> = None;
#[derive(Clone, Debug)]
pub struct Options {
@ -53,8 +53,8 @@ pub struct Options {
pub log_compiled_iseqs: Option<String>,
}
/// Return an Options with default values
pub fn init_options() -> Options {
impl Default for Options {
fn default() -> Self {
Options {
num_profiles: 1,
stats: false,
@ -68,6 +68,7 @@ pub fn init_options() -> Options {
log_compiled_iseqs: None,
}
}
}
/// `ruby --help` descriptions for user-facing options. Do not add options for ZJIT developers.
/// Note that --help allows only 80 chars per line, including indentation. 80-char limit --> |
@ -95,28 +96,26 @@ macro_rules! get_option {
// Unsafe is ok here because options are initialized
// once before any Ruby code executes
($option_name:ident) => {
{
use crate::state::ZJITState;
ZJITState::get_options().$option_name
}
unsafe { crate::options::OPTIONS.as_ref() }.unwrap().$option_name
};
}
pub(crate) use get_option;
/// Allocate Options on the heap, initialize it, and return the address of it.
/// The return value will be modified by rb_zjit_parse_option() and then
/// passed to rb_zjit_init() for initialization.
/// Set default values to ZJIT options. Setting Some to OPTIONS will make `#with_jit`
/// enable the JIT hook while not enabling compilation yet.
#[unsafe(no_mangle)]
pub extern "C" fn rb_zjit_init_options() -> *const u8 {
let options = init_options();
Box::into_raw(Box::new(options)) as *const u8
pub extern "C" fn rb_zjit_prepare_options() {
// rb_zjit_prepare_options() could be called for feature flags or $RUBY_ZJIT_ENABLE
// after rb_zjit_parse_option() is called, so we need to handle the already-initialized case.
if unsafe { OPTIONS.is_none() } {
unsafe { OPTIONS = Some(Options::default()); }
}
}
/// Parse a --zjit* command-line flag
#[unsafe(no_mangle)]
pub extern "C" fn rb_zjit_parse_option(options: *const u8, str_ptr: *const c_char) -> bool {
let options = unsafe { &mut *(options as *mut Options) };
parse_option(options, str_ptr).is_some()
pub extern "C" fn rb_zjit_parse_option(str_ptr: *const c_char) -> bool {
parse_option(str_ptr).is_some()
}
fn parse_jit_list(path_like: &str) -> HashSet<String> {
@ -142,7 +141,10 @@ fn parse_jit_list(path_like: &str) -> HashSet<String> {
/// Expected to receive what comes after the third dash in "--zjit-*".
/// Empty string means user passed only "--zjit". C code rejects when
/// they pass exact "--zjit-".
fn parse_option(options: &mut Options, str_ptr: *const std::os::raw::c_char) -> Option<()> {
fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
rb_zjit_prepare_options();
let options = unsafe { OPTIONS.as_mut().unwrap() };
let c_str: &CStr = unsafe { CStr::from_ptr(str_ptr) };
let opt_str: &str = c_str.to_str().ok()?;
@ -161,7 +163,7 @@ fn parse_option(options: &mut Options, str_ptr: *const std::os::raw::c_char) ->
("call-threshold", _) => match opt_val.parse() {
Ok(n) => {
unsafe { rb_zjit_call_threshold = n; }
update_profile_threshold(options);
update_profile_threshold();
},
Err(_) => return None,
},
@ -169,13 +171,12 @@ fn parse_option(options: &mut Options, str_ptr: *const std::os::raw::c_char) ->
("num-profiles", _) => match opt_val.parse() {
Ok(n) => {
options.num_profiles = n;
update_profile_threshold(options);
update_profile_threshold();
},
Err(_) => return None,
},
("stats", "") => {
unsafe { zjit_stats_enabled_p = true; }
options.stats = true;
}
@ -217,15 +218,14 @@ fn parse_option(options: &mut Options, str_ptr: *const std::os::raw::c_char) ->
}
/// Update rb_zjit_profile_threshold based on rb_zjit_call_threshold and options.num_profiles
fn update_profile_threshold(options: &Options) {
unsafe {
if rb_zjit_call_threshold == 1 {
fn update_profile_threshold() {
if unsafe { rb_zjit_call_threshold == 1 } {
// If --zjit-call-threshold=1, never rewrite ISEQs to profile instructions.
rb_zjit_profile_threshold = 0;
unsafe { rb_zjit_profile_threshold = 0; }
} else {
// Otherwise, profile instructions at least once.
rb_zjit_profile_threshold = rb_zjit_call_threshold.saturating_sub(options.num_profiles as u64).max(1);
}
let num_profiles = get_option!(num_profiles) as u64;
unsafe { rb_zjit_profile_threshold = rb_zjit_call_threshold.saturating_sub(num_profiles).max(1) };
}
}
@ -254,12 +254,23 @@ macro_rules! debug {
}
pub(crate) use debug;
/// Return Qtrue if --zjit-stats has been enabled
/// Return Qtrue if --zjit* has been specified. For the `#with_jit` hook,
/// this becomes Qtrue before ZJIT is actually initialized and enabled.
#[unsafe(no_mangle)]
pub extern "C" fn rb_zjit_stats_enabled_p(_ec: EcPtr, _self: VALUE) -> VALUE {
// ZJITState is not initialized yet when loading builtins, so this relies
// on a separate global variable.
if unsafe { zjit_stats_enabled_p } {
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
} else {
Qfalse
}
}
/// Return Qtrue if --zjit-stats has been specified.
#[unsafe(no_mangle)]
pub extern "C" fn rb_zjit_stats_enabled_p(_ec: EcPtr, _self: VALUE) -> VALUE {
// Builtin zjit.rb calls this even if ZJIT is disabled, so OPTIONS may not be set.
if unsafe { OPTIONS.as_ref() }.map_or(false, |opts| opts.stats) {
Qtrue
} else {
Qfalse

34
zjit/src/state.rs
View file

@ -1,8 +1,8 @@
use crate::cruby::{self, rb_bug_panic_hook, rb_vm_insns_count, EcPtr, Qnil, VALUE};
use crate::cruby_methods;
use crate::invariants::Invariants;
use crate::options::Options;
use crate::asm::CodeBlock;
use crate::options::get_option;
use crate::stats::Counters;
#[allow(non_upper_case_globals)]
@ -19,9 +19,6 @@ pub struct ZJITState {
/// Inline code block (fast path)
code_block: CodeBlock,
/// ZJIT command-line options
options: Options,
/// ZJIT statistics
counters: Counters,
@ -39,11 +36,12 @@ pub struct ZJITState {
static mut ZJIT_STATE: Option<ZJITState> = None;
impl ZJITState {
/// Initialize the ZJIT globals, given options allocated by rb_zjit_init_options()
pub fn init(options: Options) {
/// Initialize the ZJIT globals
pub fn init() {
#[cfg(not(test))]
let cb = {
use crate::cruby::*;
use crate::options::*;
let exec_mem_size: usize = 64 * 1024 * 1024; // TODO: implement the option
let virt_block: *mut u8 = unsafe { rb_zjit_reserve_addr_space(64 * 1024 * 1024) };
@ -75,7 +73,7 @@ impl ZJITState {
);
let mem_block = Rc::new(RefCell::new(mem_block));
CodeBlock::new(mem_block.clone(), options.dump_disasm)
CodeBlock::new(mem_block.clone(), get_option!(dump_disasm))
};
#[cfg(test)]
let cb = CodeBlock::new_dummy();
@ -83,7 +81,6 @@ impl ZJITState {
// Initialize the codegen globals instance
let zjit_state = ZJITState {
code_block: cb,
options,
counters: Counters::default(),
invariants: Invariants::default(),
assert_compiles: false,
@ -107,11 +104,6 @@ impl ZJITState {
&mut ZJITState::get_instance().code_block
}
/// Get a mutable reference to the options
pub fn get_options() -> &'static mut Options {
&mut ZJITState::get_instance().options
}
/// Get a mutable reference to the invariants
pub fn get_invariants() -> &'static mut Invariants {
&mut ZJITState::get_instance().invariants
@ -139,13 +131,13 @@ impl ZJITState {
/// Was --zjit-save-compiled-iseqs specified?
pub fn should_log_compiled_iseqs() -> bool {
ZJITState::get_instance().options.log_compiled_iseqs.is_some()
get_option!(log_compiled_iseqs).is_some()
}
/// Log the name of a compiled ISEQ to the file specified in options.log_compiled_iseqs
pub fn log_compile(iseq_name: String) {
assert!(ZJITState::should_log_compiled_iseqs());
let filename = ZJITState::get_instance().options.log_compiled_iseqs.as_ref().unwrap();
let filename = get_option!(log_compiled_iseqs).as_ref().unwrap();
use std::io::Write;
let mut file = match std::fs::OpenOptions::new().create(true).append(true).open(filename) {
Ok(f) => f,
@ -161,7 +153,7 @@ impl ZJITState {
/// Check if we are allowed to compile a given ISEQ based on --zjit-allowed-iseqs
pub fn can_compile_iseq(iseq: cruby::IseqPtr) -> bool {
if let Some(ref allowed_iseqs) = ZJITState::get_instance().options.allowed_iseqs {
if let Some(ref allowed_iseqs) = get_option!(allowed_iseqs) {
let name = cruby::iseq_get_location(iseq, 0);
allowed_iseqs.contains(&name)
} else {
@ -170,17 +162,17 @@ impl ZJITState {
}
}
/// Initialize ZJIT, given options allocated by rb_zjit_init_options()
/// Initialize ZJIT
#[unsafe(no_mangle)]
pub extern "C" fn rb_zjit_init(options: *const u8) {
pub extern "C" fn rb_zjit_init() {
// 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();
ZJITState::init();
let options = unsafe { Box::from_raw(options as *mut Options) };
ZJITState::init(*options);
// Install a panic hook for ZJIT
rb_bug_panic_hook();
// Discard the instruction count for boot which we never compile