YJIT: Replace Array#each only when YJIT is enabled (#11955)

* YJIT: Replace Array#each only when YJIT is enabled * Add comments about BUILTIN_ATTR_C_TRACE * Make Ruby Array#each available with --yjit as well * Fix all paths that expect a C location * Use method_basic_definition_p to detect patches * Copy a comment about C_TRACE flag to compilers * Rephrase a comment about add_yjit_hook * Give METHOD_ENTRY_BASIC flag to Array#each * Add --yjit-c-builtin option * Allow inconsistent source_location in test-spec * Refactor a check of BUILTIN_ATTR_C_TRACE * Set METHOD_ENTRY_BASIC without touching vm->running
2024-11-04 08:14:28 -08:00 · 2024-11-04 08:14:28 -08:00 · 478e0fc710
--- a/array.c
+++ b/array.c
@ -2604,6 +2604,39 @@ ary_fetch_next(VALUE self, VALUE *index, VALUE *value)
    return Qtrue;
 }

+/*
+ *  call-seq:
+ *    each {|element| ... } -> self
+ *    each -> new_enumerator
+ *
+ *  With a block given, iterates over the elements of +self+,
+ *  passing each element to the block;
+ *  returns +self+:
+ *
+ *    a = [:foo, 'bar', 2]
+ *    a.each {|element|  puts "#{element.class} #{element}" }
+ *
+ *  Output:
+ *
+ *    Symbol foo
+ *    String bar
+ *    Integer 2
+ *
+ *  Allows the array to be modified during iteration:
+ *
+ *    a = [:foo, 'bar', 2]
+ *    a.each {|element| puts element; a.clear if element.to_s.start_with?('b') }
+ *
+ *  Output:
+ *
+ *    foo
+ *    bar
+ *
+ *  With no block given, returns a new Enumerator.
+ *
+ *  Related: see {Methods for Iterating}[rdoc-ref:Array@Methods+for+Iterating].
+ */
+
 VALUE
 rb_ary_each(VALUE ary)
 {
@ -8634,6 +8667,7 @@ Init_Array(void)
    rb_define_method(rb_cArray, "unshift", rb_ary_unshift_m, -1);
    rb_define_alias(rb_cArray,  "prepend", "unshift");
    rb_define_method(rb_cArray, "insert", rb_ary_insert, -1);
+    rb_define_method(rb_cArray, "each", rb_ary_each, 0);
    rb_define_method(rb_cArray, "each_index", rb_ary_each_index, 0);
    rb_define_method(rb_cArray, "reverse_each", rb_ary_reverse_each, 0);
    rb_define_method(rb_cArray, "length", rb_ary_length, 0);
--- a/array.rb
+++ b/array.rb
@ -1,49 +1,4 @@
 class Array
-  # call-seq:
-  #   each {|element| ... } -> self
-  #   each -> new_enumerator
-  #
-  # With a block given, iterates over the elements of +self+,
-  # passing each element to the block;
-  # returns +self+:
-  #
-  #   a = [:foo, 'bar', 2]
-  #   a.each {|element|  puts "#{element.class} #{element}" }
-  #
-  # Output:
-  #
-  #   Symbol foo
-  #   String bar
-  #   Integer 2
-  #
-  # Allows the array to be modified during iteration:
-  #
-  #   a = [:foo, 'bar', 2]
-  #   a.each {|element| puts element; a.clear if element.to_s.start_with?('b') }
-  #
-  # Output:
-  #
-  #   foo
-  #   bar
-  #
-  # With no block given, returns a new Enumerator.
-  #
-  # Related: see {Methods for Iterating}[rdoc-ref:Array@Methods+for+Iterating].
-
-  def each
-    Primitive.attr! :inline_block
-
-    unless defined?(yield)
-      return Primitive.cexpr! 'SIZED_ENUMERATOR(self, 0, 0, ary_enum_length)'
-    end
-    _i = 0
-    value = nil
-    while Primitive.cexpr!(%q{ ary_fetch_next(self, LOCAL_PTR(_i), LOCAL_PTR(value)) })
-      yield value
-    end
-    self
-  end
-
  # call-seq:
  #   shuffle!(random: Random) -> self
  #
@ -258,4 +213,24 @@ class Array
    indexes.map! { |i| fetch(i, &block) }
    indexes
  end
+
+  with_yjit do
+    if Primitive.rb_builtin_basic_definition_p(:each)
+      undef :each
+
+      def each # :nodoc:
+        Primitive.attr! :inline_block, :c_trace
+
+        unless defined?(yield)
+          return Primitive.cexpr! 'SIZED_ENUMERATOR(self, 0, 0, ary_enum_length)'
+        end
+        _i = 0
+        value = nil
+        while Primitive.cexpr!(%q{ ary_fetch_next(self, LOCAL_PTR(_i), LOCAL_PTR(value)) })
+          yield value
+        end
+        self
+      end
+    end
+  end
 end
--- a/builtin.h
+++ b/builtin.h
@ -106,6 +106,12 @@ rb_vm_lvar(rb_execution_context_t *ec, int index)
 #endif
 }

+static inline VALUE
+rb_builtin_basic_definition_p(rb_execution_context_t *ec, VALUE klass, VALUE id_sym)
+{
+    return rb_method_basic_definition_p(klass, rb_sym2id(id_sym)) ? Qtrue : Qfalse;
+}
+
 #define LOCAL_PTR(local) local ## __ptr

 // dump/load
--- a/common.mk
+++ b/common.mk
@ -1212,6 +1212,7 @@ BUILTIN_RB_SRCS = \
 		$(srcdir)/prelude.rb \
 		$(srcdir)/gem_prelude.rb \
 		$(srcdir)/yjit.rb \
+		$(srcdir)/yjit_hook.rb \
 		$(empty)
 BUILTIN_RB_INCS = $(BUILTIN_RB_SRCS:.rb=.rbinc)

@ -10674,6 +10675,7 @@ 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
 node.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 node.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 node.$(OBJEXT): $(CCAN_DIR)/list/list.h
@ -20036,6 +20038,7 @@ 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_backtrace.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 vm_backtrace.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 vm_backtrace.$(OBJEXT): $(CCAN_DIR)/list/list.h
--- a/compile.c
+++ b/compile.c
@ -8960,6 +8960,10 @@ compile_builtin_attr(rb_iseq_t *iseq, const NODE *node)
        else if (strcmp(RSTRING_PTR(string), "use_block") == 0) {
            iseq_set_use_block(iseq);
        }
+        else if (strcmp(RSTRING_PTR(string), "c_trace") == 0) {
+            // Let the iseq act like a C method in backtraces
+            ISEQ_BODY(iseq)->builtin_attrs |= BUILTIN_ATTR_C_TRACE;
+        }
        else {
            goto unknown_arg;
        }
--- a/inits.c
+++ b/inits.c
@ -84,6 +84,9 @@ void
 rb_call_builtin_inits(void)
 {
 #define BUILTIN(n) CALL(builtin_##n)
+    BUILTIN(kernel);
+    BUILTIN(yjit);
+    // BUILTIN(yjit_hook) is called after rb_yjit_init()
    BUILTIN(gc);
    BUILTIN(ractor);
    BUILTIN(numeric);
@ -95,11 +98,9 @@ rb_call_builtin_inits(void)
    BUILTIN(warning);
    BUILTIN(array);
    BUILTIN(hash);
-    BUILTIN(kernel);
    BUILTIN(symbol);
    BUILTIN(timev);
    BUILTIN(thread_sync);
-    BUILTIN(yjit);
    BUILTIN(nilclass);
    BUILTIN(marshal);
    BUILTIN(rjit_c);
--- a/kernel.rb
+++ b/kernel.rb
@ -290,4 +290,12 @@ module Kernel
      Primitive.rb_f_integer(arg, base, exception);
    end
  end
+
+  # Internal helper for builtin inits to define methods only when YJIT is enabled.
+  # This method is removed in yjit_hook.rb.
+  def with_yjit(&block) # :nodoc:
+    if defined?(RubyVM::YJIT)
+      RubyVM::YJIT.send(:add_yjit_hook, block)
+    end
+  end
 end
--- a/prism_compile.c
+++ b/prism_compile.c
@ -3387,6 +3387,10 @@ pm_compile_builtin_attr(rb_iseq_t *iseq, const pm_scope_node_t *scope_node, cons
        else if (strcmp(RSTRING_PTR(string), "use_block") == 0) {
            iseq_set_use_block(iseq);
        }
+        else if (strcmp(RSTRING_PTR(string), "c_trace") == 0) {
+            // Let the iseq act like a C method in backtraces
+            ISEQ_BODY(iseq)->builtin_attrs |= BUILTIN_ATTR_C_TRACE;
+        }
        else {
            COMPILE_ERROR(iseq, node_location->line, "unknown argument to attr!: %s", RSTRING_PTR(string));
            return COMPILE_NG;
--- a/ruby.c
+++ b/ruby.c
@ -1816,6 +1816,10 @@ ruby_opt_init(ruby_cmdline_options_t *opt)
    rb_yjit_init(opt->yjit);
 #endif

+    // Call yjit_hook.rb after rb_yjit_init() to use `RubyVM::YJIT.enabled?`
+    void Init_builtin_yjit_hook();
+    Init_builtin_yjit_hook();
+
    ruby_set_script_name(opt->script_name);
    require_libraries(&opt->req_list);
 }
--- a/spec/ruby/optional/capi/debug_spec.rb
+++ b/spec/ruby/optional/capi/debug_spec.rb
@ -37,9 +37,12 @@ describe "C-API Debug function" do

    it "matches the locations in rb_debug_inspector_backtrace_locations" do
      frames = @o.rb_debug_inspector_open(42)
-      frames.each do |_s, _klass, binding, _iseq, backtrace_location|
+      frames.each do |_s, klass, binding, iseq, backtrace_location|
        if binding
-          binding.source_location.should == [backtrace_location.path, backtrace_location.lineno]
+          # YJIT modifies Array#each backtraces but leaves its source_location as is
+          unless defined?(RubyVM::YJIT) && klass == Array && iseq.label == "each"
+            binding.source_location.should == [backtrace_location.path, backtrace_location.lineno]
+          end
          method_name = binding.eval('__method__')
          if method_name
            method_name.should == backtrace_location.base_label.to_sym
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@ -1677,6 +1677,71 @@ class TestYJIT < Test::Unit::TestCase
    RUBY
  end

+  def test_yjit_option_uses_array_each_in_ruby
+    assert_separately(["--yjit"], <<~'RUBY')
+      # Array#each should be implemented in Ruby for YJIT
+      assert_equal "<internal:array>", Array.instance_method(:each).source_location.first
+
+      # The backtrace, however, should not be `from <internal:array>:XX:in 'Array#each'`
+      begin
+        [nil].each { raise }
+      rescue => e
+        assert_equal "-:11:in 'Array#each'", e.backtrace[1]
+      end
+    RUBY
+  end
+
+  def test_yjit_enable_replaces_array_each
+    assert_separately([*("--disable=yjit" if RubyVM::YJIT.enabled?)], <<~'RUBY')
+      # Array#each should be implemented in C for the interpreter
+      assert_nil Array.instance_method(:each).source_location
+
+      # The backtrace should not be `from <internal:array>:XX:in 'Array#each'`
+      begin
+        [nil].each { raise }
+      rescue => e
+        assert_equal "-:11:in 'Array#each'", e.backtrace[1]
+      end
+
+      RubyVM::YJIT.enable
+
+      # Array#each should be implemented in Ruby for YJIT
+      assert_equal "<internal:array>", Array.instance_method(:each).source_location.first
+
+      # However, the backtrace should still not be `from <internal:array>:XX:in 'Array#each'`
+      begin
+        [nil].each { raise }
+      rescue => e
+        assert_equal "-:23:in 'Array#each'", e.backtrace[1]
+      end
+    RUBY
+  end
+
+  def test_yjit_enable_preserves_array_each_monkey_patch
+    assert_separately([*("--disable=yjit" if RubyVM::YJIT.enabled?)], <<~'RUBY')
+      # Array#each should be implemented in C initially
+      assert_nil Array.instance_method(:each).source_location
+
+      # Override Array#each
+      $called = false
+      Array.prepend(Module.new {
+        def each
+          $called = true
+          super
+        end
+      })
+
+      RubyVM::YJIT.enable
+
+      # The monkey-patch should still be alive
+      [].each {}
+      assert_true $called
+
+      # YJIT should not replace Array#each with the "<internal:array>" one
+      assert_equal "-", Array.instance_method(:each).source_location.first
+    RUBY
+  end
+
  private

  def code_gc_helpers
--- a/tool/mk_builtin_loader.rb
+++ b/tool/mk_builtin_loader.rb
@ -6,7 +6,7 @@ require_relative 'ruby_vm/helpers/c_escape'

 SUBLIBS = {}
 REQUIRED = {}
-BUILTIN_ATTRS = %w[leaf inline_block use_block]
+BUILTIN_ATTRS = %w[leaf inline_block use_block c_trace]

 module CompileWarning
  @@warnings = 0
--- a/vm.c
+++ b/vm.c
@ -4437,6 +4437,9 @@ Init_vm_objects(void)
 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 RJIT units
 #if !USE_RJIT
 void Init_builtin_rjit(void) {}
--- a/vm_backtrace.c
+++ b/vm_backtrace.c
@ -265,10 +265,26 @@ retry:
    }
 }

+// Return true if a given location is a C method or supposed to behave like one.
+static inline bool
+location_cfunc_p(rb_backtrace_location_t *loc)
+{
+    if (!loc->cme) return false;
+
+    switch (loc->cme->def->type) {
+      case VM_METHOD_TYPE_CFUNC:
+        return true;
+      case VM_METHOD_TYPE_ISEQ:
+        return rb_iseq_attr_p(loc->cme->def->body.iseq.iseqptr, BUILTIN_ATTR_C_TRACE);
+      default:
+        return false;
+    }
+}
+
 static VALUE
 location_label(rb_backtrace_location_t *loc)
 {
-    if (loc->cme && loc->cme->def->type == VM_METHOD_TYPE_CFUNC) {
+    if (location_cfunc_p(loc)) {
        return rb_gen_method_name(loc->cme->owner, rb_id2str(loc->cme->def->original_id));
    }
    else {
@ -314,7 +330,7 @@ location_label_m(VALUE self)
 static VALUE
 location_base_label(rb_backtrace_location_t *loc)
 {
-    if (loc->cme && loc->cme->def->type == VM_METHOD_TYPE_CFUNC) {
+    if (location_cfunc_p(loc)) {
        return rb_id2str(loc->cme->def->original_id);
    }

@ -448,7 +464,7 @@ location_to_str(rb_backtrace_location_t *loc)
    VALUE file, owner = Qnil, name;
    int lineno;

-    if (loc->cme && loc->cme->def->type == VM_METHOD_TYPE_CFUNC) {
+    if (location_cfunc_p(loc)) {
        if (loc->iseq && loc->pc) {
            file = rb_iseq_path(loc->iseq);
            lineno = calc_lineno(loc->iseq, loc->pc);
@ -684,13 +700,21 @@ rb_ec_partial_backtrace_object(const rb_execution_context_t *ec, long start_fram
                        const VALUE *pc = cfp->pc;
                        loc = &bt->backtrace[bt->backtrace_size++];
                        RB_OBJ_WRITE(btobj, &loc->cme, rb_vm_frame_method_entry(cfp));
-                        RB_OBJ_WRITE(btobj, &loc->iseq, iseq);
-                        loc->pc = pc;
-                        bt_update_cfunc_loc(cfunc_counter, loc-1, iseq, pc);
-                        if (do_yield) {
-                            bt_yield_loc(loc - cfunc_counter, cfunc_counter+1, btobj);
+                        // Ruby methods with `Primitive.attr! :c_trace` should behave like C methods
+                        if (rb_iseq_attr_p(cfp->iseq, BUILTIN_ATTR_C_TRACE)) {
+                            loc->iseq = NULL;
+                            loc->pc = NULL;
+                            cfunc_counter++;
+                        }
+                        else {
+                            RB_OBJ_WRITE(btobj, &loc->iseq, iseq);
+                            loc->pc = pc;
+                            bt_update_cfunc_loc(cfunc_counter, loc-1, iseq, pc);
+                            if (do_yield) {
+                                bt_yield_loc(loc - cfunc_counter, cfunc_counter+1, btobj);
+                            }
+                            cfunc_counter = 0;
                        }
-                        cfunc_counter = 0;
                    }
                    skip_next_frame = is_rescue_or_ensure_frame(cfp);
                }
--- a/vm_core.h
+++ b/vm_core.h
@ -395,6 +395,8 @@ enum rb_builtin_attr {
    BUILTIN_ATTR_SINGLE_NOARG_LEAF = 0x02,
    // This attribute signals JIT to duplicate the iseq for each block iseq so that its `yield` will be monomorphic.
    BUILTIN_ATTR_INLINE_BLOCK = 0x04,
+    // The iseq acts like a C method in backtraces.
+    BUILTIN_ATTR_C_TRACE = 0x08,
 };

 typedef VALUE (*rb_jit_func_t)(struct rb_execution_context_struct *, struct rb_control_frame_struct *);
@ -604,6 +606,12 @@ rb_iseq_check(const rb_iseq_t *iseq)
    return iseq;
 }

+static inline bool
+rb_iseq_attr_p(const rb_iseq_t *iseq, enum rb_builtin_attr attr)
+{
+    return (ISEQ_BODY(iseq)->builtin_attrs & attr) == attr;
+}
+
 static inline const rb_iseq_t *
 def_iseq_ptr(rb_method_definition_t *def)
 {
--- a/vm_method.c
+++ b/vm_method.c
@ -641,6 +641,11 @@ 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
+                if (rb_iseq_attr_p(iseq, BUILTIN_ATTR_C_TRACE)) {
+                    METHOD_ENTRY_BASIC_SET((rb_method_entry_t *)me, TRUE);
+                }
+
                if (ISEQ_BODY(iseq)->mandatory_only_iseq) def->iseq_overload = 1;

                if (0) vm_cref_dump("rb_method_definition_create", cref);
--- a/yjit.c
+++ b/yjit.c
@ -1244,6 +1244,14 @@ 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_c_builtin_p(rb_execution_context_t *ec, VALUE self);
+
+// Allow YJIT_C_BUILTIN macro to force --yjit-c-builtin
+#ifdef YJIT_C_BUILTIN
+static VALUE yjit_c_builtin_p(rb_execution_context_t *ec, VALUE self) { return Qtrue; }
+#else
+#define yjit_c_builtin_p rb_yjit_c_builtin_p
+#endif

 // Preprocessed yjit.rb generated during build
 #include "yjit.rbinc"
--- a/yjit.rb
+++ b/yjit.rb
@ -37,7 +37,7 @@ module RubyVM::YJIT
  # whether to enable \YJIT compilation logging or not.
  #
  # `stats`:
-  # * `false`: Disable stats.
+  # * `false`: Don't enable stats.
  # * `true`: Enable stats. Print stats at exit.
  # * `:quiet`: Enable stats. Do not print stats at exit.
  #
@ -48,6 +48,7 @@ module RubyVM::YJIT
  def self.enable(stats: false, log: false)
    return false if enabled?
    at_exit { print_and_dump_stats } if stats
+    call_yjit_hooks
    Primitive.rb_yjit_enable(stats, stats != :quiet, log, log != :quiet)
  end

@ -247,10 +248,26 @@ module RubyVM::YJIT
    at_exit { print_and_dump_stats }
  end

+  # Blocks that are called when YJIT is enabled
+  @yjit_hooks = []
+
  class << self
    # :stopdoc:
    private

+    # Register a block to be called when YJIT is enabled
+    def add_yjit_hook(hook)
+      @yjit_hooks << hook
+    end
+
+    # Run YJIT hooks registered by RubyVM::YJIT.with_yjit
+    def call_yjit_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
+    end
+
    # Print stats and dump exit locations
    def print_and_dump_stats # :nodoc:
      if Primitive.rb_yjit_print_stats_p
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@ -492,6 +492,7 @@ pub type rb_iseq_type = u32;
 pub const BUILTIN_ATTR_LEAF: rb_builtin_attr = 1;
 pub const BUILTIN_ATTR_SINGLE_NOARG_LEAF: rb_builtin_attr = 2;
 pub const BUILTIN_ATTR_INLINE_BLOCK: rb_builtin_attr = 4;
+pub const BUILTIN_ATTR_C_TRACE: rb_builtin_attr = 8;
 pub type rb_builtin_attr = u32;
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
--- a/yjit/src/options.rs
+++ b/yjit/src/options.rs
@ -1,5 +1,5 @@
 use std::{ffi::{CStr, CString}, ptr::null, fs::File};
-use crate::{backend::current::TEMP_REGS, stats::Counter};
+use crate::{backend::current::TEMP_REGS, cruby::*, stats::Counter};
 use std::os::raw::{c_char, c_int, c_uint};

 // Call threshold for small deployments and command-line apps
@ -46,6 +46,9 @@ 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
+    pub c_builtin: bool,
+
    // Capture stats
    pub gen_stats: bool,

@ -94,6 +97,7 @@ pub static mut OPTIONS: Options = Options {
    no_type_prop: false,
    max_versions: 4,
    num_temp_regs: 5,
+    c_builtin: false,
    gen_stats: false,
    trace_exits: None,
    print_stats: true,
@ -270,6 +274,10 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
            }
        },

+        ("c-builtin", _) => unsafe {
+            OPTIONS.c_builtin = true;
+        },
+
        ("code-gc", _) => unsafe {
            OPTIONS.code_gc = true;
        },
@ -413,3 +421,13 @@ pub extern "C" fn rb_yjit_show_usage(help: c_int, highlight: c_int, width: c_uin
        unsafe { ruby_show_usage_line(name.as_ptr(), null(), description.as_ptr(), help, highlight, width, columns) }
    }
 }
+
+/// Return true if --yjit-c-builtin is given
+#[no_mangle]
+pub extern "C" fn rb_yjit_c_builtin_p(_ec: EcPtr, _self: VALUE) -> VALUE {
+    if get_option!(c_builtin) {
+        Qtrue
+    } else {
+        Qfalse
+    }
+}
--- a/yjit_hook.rb
+++ b/yjit_hook.rb
@ -0,0 +1,9 @@
+# If YJIT is enabled, load the YJIT-only version of builtin methods
+if defined?(RubyVM::YJIT) && RubyVM::YJIT.enabled?
+  RubyVM::YJIT.send(:call_yjit_hooks)
+end
+
+# Remove the helper defined in kernel.rb
+module Kernel
+  undef :with_yjit
+end