From 5085ec3ed90beb54125d5eb9329e202ae1542b5c Mon Sep 17 00:00:00 2001 From: annichai-stripe Date: Mon, 3 Mar 2025 12:45:39 -0800 Subject: [PATCH] Allow YJIT `mem-size` and `call-threshold` to be set at runtime via `YJIT.enable()` (#12505) * first commit * yjit.rb change * revert formatting * rename mem-size to exec-mem-size for correctness * wip, move setting into rb_yjit_enable directly * remove unused helper functions * add in call threshold * input validation with extensive eprintln * delete test script * exec-mem-size -> mem-size * handle input validation with asserts * add test cases related to input validation * modify test cases * move validation out of rs, into rb * add comments * remove trailing spaces * remove logging Co-authored-by: Takashi Kokubun * remove helper fn * Update test/ruby/test_yjit.rb Co-authored-by: Takashi Kokubun * trailing white space --------- Co-authored-by: Alan Wu Co-authored-by: Takashi Kokubun Co-authored-by: Maxime Chevalier-Boisvert --- test/ruby/test_yjit.rb | 25 +++++++++++++++++++++++++ yjit.c | 2 +- yjit.rb | 20 ++++++++++++++++---- yjit/src/yjit.rs | 18 +++++++++++++++++- 4 files changed, 59 insertions(+), 6 deletions(-) diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb index c0212987e7..450e6e5625 100644 --- a/test/ruby/test_yjit.rb +++ b/test/ruby/test_yjit.rb @@ -142,6 +142,31 @@ class TestYJIT < Test::Unit::TestCase RUBY end + def test_yjit_enable_with_valid_runtime_call_threshold_option + assert_in_out_err(['--yjit-disable', '-e', + 'RubyVM::YJIT.enable(call_threshold: 1); puts RubyVM::YJIT.enabled?']) do |stdout, stderr, _status| + assert_empty stderr + assert_include stdout.join, "true" + end + end + + def test_yjit_enable_with_invalid_runtime_call_threshold_option + assert_in_out_err(['--yjit-disable', '-e', 'RubyVM::YJIT.enable(mem_size: 0)']) do |stdout, stderr, status| + assert_not_empty stderr + assert_match(/ArgumentError/, stderr.join) + assert_equal 1, status.exitstatus + end + end + + def test_yjit_enable_with_invalid_runtime_mem_size_option + assert_in_out_err(['--yjit-disable', '-e', 'RubyVM::YJIT.enable(mem_size: 0)']) do |stdout, stderr, status| + assert_not_empty stderr + assert_match(/ArgumentError/, stderr.join) + assert_equal 1, status.exitstatus + end + end + + def test_yjit_stats_and_v_no_error _stdout, stderr, _status = invoke_ruby(%w(-v --yjit-stats), '', true, true) refute_includes(stderr, "NoMethodError") diff --git a/yjit.c b/yjit.c index 1f481aaa78..04231a0491 100644 --- a/yjit.c +++ b/yjit.c @@ -1250,7 +1250,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 gen_stats, VALUE print_stats, VALUE gen_compilation_log, VALUE print_compilation_log); +VALUE rb_yjit_enable(rb_execution_context_t *ec, VALUE self, VALUE gen_stats, VALUE print_stats, VALUE gen_compilation_log, VALUE print_compilation_log, VALUE mem_size, VALUE call_threshold); VALUE rb_yjit_c_builtin_p(rb_execution_context_t *ec, VALUE self); // Allow YJIT_C_BUILTIN macro to force --yjit-c-builtin diff --git a/yjit.rb b/yjit.rb index c87bc1f4af..9e9b112347 100644 --- a/yjit.rb +++ b/yjit.rb @@ -34,7 +34,8 @@ module RubyVM::YJIT end # Enable \YJIT compilation. `stats` option decides whether to enable \YJIT stats or not. `compilation_log` decides - # whether to enable \YJIT compilation logging or not. + # whether to enable \YJIT compilation logging or not. Optional `mem_size` and `call_threshold` can be + # provided to override default configuration. # # * `stats`: # * `false`: Don't enable stats. @@ -44,11 +45,22 @@ module RubyVM::YJIT # * `false`: Don't enable the log. # * `true`: Enable the log. Print log at exit. # * `:quiet`: Enable the log. Do not print log at exit. - def self.enable(stats: false, log: false) + def self.enable(stats: false, log: false, mem_size: nil, call_threshold: nil) return false if enabled? + + if mem_size + raise ArgumentError, "mem_size must be a Integer" unless mem_size.is_a?(Integer) + raise ArgumentError, "mem_size must be between 1 and 2048 MB" unless (1..2048).include?(mem_size) + end + + if call_threshold + raise ArgumentError, "call_threshold must be a Integer" unless call_threshold.is_a?(Integer) + raise ArgumentError, "call_threshold must be a positive integer" unless call_threshold.positive? + end + at_exit { print_and_dump_stats } if stats call_yjit_hooks - Primitive.rb_yjit_enable(stats, stats != :quiet, log, log != :quiet) + Primitive.rb_yjit_enable(stats, stats != :quiet, log, log != :quiet, mem_size, call_threshold) end # If --yjit-trace-exits is enabled parse the hashes from @@ -532,4 +544,4 @@ module RubyVM::YJIT # :startdoc: end -end +end \ No newline at end of file diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs index c3d851fdbc..4b2b5214f4 100644 --- a/yjit/src/yjit.rs +++ b/yjit/src/yjit.rs @@ -185,8 +185,24 @@ 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, gen_stats: VALUE, print_stats: VALUE, gen_log: VALUE, print_log: VALUE) -> VALUE { +pub extern "C" fn rb_yjit_enable(_ec: EcPtr, _ruby_self: VALUE, gen_stats: VALUE, print_stats: VALUE, gen_log: VALUE, print_log: VALUE, mem_size: VALUE, call_threshold: VALUE) -> VALUE { with_vm_lock(src_loc!(), || { + + if !mem_size.nil_p() { + let mem_size_mb = mem_size.as_isize() >> 1; + let mem_size_bytes = mem_size_mb * 1024 * 1024; + unsafe { + OPTIONS.mem_size = mem_size_bytes as usize; + } + } + + if !call_threshold.nil_p() { + let threshold = call_threshold.as_isize() >> 1; + unsafe { + rb_yjit_call_threshold = threshold as u64; + } + } + // Initialize and enable YJIT if gen_stats.test() { unsafe {