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 <takashikkbn@gmail.com>

* remove helper fn

* Update test/ruby/test_yjit.rb

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

* trailing white space

---------

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>
Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>
Co-authored-by: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>

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")

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

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

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 {