YJIT: Compile exception handlers (#8171)

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>

Author: k0kubun

bootstraptest/test_yjit.rb

@@ -4004,3 +4004,25 @@ def calling_func
     _x, _y = func.call
   end.call
 }
+
+# Catch TAG_BREAK in a non-FINISH frame with JIT code
+assert_equal '1', %q{
+  def entry
+    catch_break
+  end
+
+  def catch_break
+    while_true do
+      break
+    end
+    1
+  end
+
+  def while_true
+    while true
+      yield
+    end
+  end
+
+  entry
+}

lib/ruby_vm/rjit/compiler.rb

@@ -64,7 +64,7 @@ def compile(iseq, cfp)
       asm = Assembler.new
       compile_prologue(asm, iseq, pc)
       compile_block(asm, jit:, pc:)
-      iseq.body.jit_func = @cb.write(asm)
+      iseq.body.jit_entry = @cb.write(asm)
     rescue Exception => e
       $stderr.puts e.full_message
       exit 1
@@ -176,8 +176,8 @@ def invalidate_blocks(iseq, pc)
 
       # If they were the ISEQ's first blocks, re-compile RJIT entry as well
       if iseq.body.iseq_encoded.to_i == pc
-        iseq.body.jit_func = 0
-        iseq.body.total_calls = 0
+        iseq.body.jit_entry = 0
+        iseq.body.jit_entry_calls = 0
       end
     end
 

lib/ruby_vm/rjit/invariants.rb

@@ -143,11 +143,11 @@ def invalidate_all
 
         C.rjit_for_each_iseq do |iseq|
           # Avoid entering past code
-          iseq.body.jit_func = 0
+          iseq.body.jit_entry = 0
           # Avoid reusing past code
           iseq.body.rjit_blocks.clear if iseq.body.rjit_blocks
           # Compile this again if not converted to trace_* insns
-          iseq.body.total_calls = 0
+          iseq.body.jit_entry_calls = 0
         end
       end
     end

rjit_c.rb

@@ -1177,8 +1177,8 @@ def C.rb_iseq_constant_body
       ), Primitive.cexpr!("OFFSETOF((*((struct rb_iseq_constant_body *)NULL)), mark_bits)")],
       outer_variables: [CType::Pointer.new { self.rb_id_table }, Primitive.cexpr!("OFFSETOF((*((struct rb_iseq_constant_body *)NULL)), outer_variables)")],
       mandatory_only_iseq: [CType::Pointer.new { self.rb_iseq_t }, Primitive.cexpr!("OFFSETOF((*((struct rb_iseq_constant_body *)NULL)), mandatory_only_iseq)")],
-      jit_func: [self.rb_jit_func_t, Primitive.cexpr!("OFFSETOF((*((struct rb_iseq_constant_body *)NULL)), jit_func)")],
-      total_calls: [CType::Immediate.parse("unsigned long"), Primitive.cexpr!("OFFSETOF((*((struct rb_iseq_constant_body *)NULL)), total_calls)")],
+      jit_entry: [self.rb_jit_func_t, Primitive.cexpr!("OFFSETOF((*((struct rb_iseq_constant_body *)NULL)), jit_entry)")],
+      jit_entry_calls: [CType::Immediate.parse("unsigned long"), Primitive.cexpr!("OFFSETOF((*((struct rb_iseq_constant_body *)NULL)), jit_entry_calls)")],
       rjit_blocks: [self.VALUE, Primitive.cexpr!("OFFSETOF((*((struct rb_iseq_constant_body *)NULL)), rjit_blocks)"), true],
     )
   end

tool/rjit/bindgen.rb

@@ -638,7 +638,7 @@ def push_target(target)
   skip_fields: {
     'rb_execution_context_struct.machine': %w[regs], # differs between macOS and Linux
     rb_execution_context_struct: %w[method_missing_reason], # non-leading bit fields not supported
-    rb_iseq_constant_body: %w[yjit_payload], # conditionally defined
+    rb_iseq_constant_body: %w[jit_exception jit_exception_calls yjit_payload], # conditionally defined
     rb_thread_struct: %w[status has_dedicated_nt to_kill abort_on_exception report_on_exception pending_interrupt_queue_checked],
     :'' => %w[is_from_method is_lambda is_isolated], # rb_proc_t
   },

vm.c

@@ -370,7 +370,14 @@ extern VALUE rb_vm_invoke_bmethod(rb_execution_context_t *ec, rb_proc_t *proc, V
 static VALUE vm_invoke_proc(rb_execution_context_t *ec, rb_proc_t *proc, VALUE self, int argc, const VALUE *argv, int kw_splat, VALUE block_handler);
 
 #if USE_RJIT || USE_YJIT
-// Try to compile the current ISeq in ec. Return 0 if not compiled.
+// Generate JIT code that supports the following kinds of ISEQ entries:
+//   * The first ISEQ on vm_exec (e.g. <main>, or Ruby methods/blocks
+//     called by a C method). The current frame has VM_FRAME_FLAG_FINISH.
+//     The current vm_exec stops if JIT code returns a non-Qundef value.
+//   * ISEQs called by the interpreter on vm_sendish (e.g. Ruby methods or
+//     blocks called by a Ruby frame that isn't compiled or side-exited).
+//     The current frame doesn't have VM_FRAME_FLAG_FINISH. The current
+//     vm_exec does NOT stop whether JIT code returns Qundef or not.
 static inline rb_jit_func_t
 jit_compile(rb_execution_context_t *ec)
 {
@@ -379,35 +386,29 @@ jit_compile(rb_execution_context_t *ec)
     struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
     bool yjit_enabled = rb_yjit_compile_new_iseqs();
     if (yjit_enabled || rb_rjit_call_p) {
-        body->total_calls++;
+        body->jit_entry_calls++;
     }
     else {
-        return 0;
-    }
-
-    // Don't try to compile the function if it's already compiled
-    if (body->jit_func) {
-        return body->jit_func;
+        return NULL;
     }
 
-    // Trigger JIT compilation as needed
-    if (yjit_enabled) {
-        if (rb_yjit_threshold_hit(iseq)) {
-            rb_yjit_compile_iseq(iseq, ec);
+    // Trigger JIT compilation if not compiled
+    if (body->jit_entry == NULL) {
+        if (yjit_enabled) {
+            if (rb_yjit_threshold_hit(iseq, body->jit_entry_calls)) {
+                rb_yjit_compile_iseq(iseq, ec, false);
+            }
         }
-    }
-    else { // rb_rjit_call_p
-        if (body->total_calls == rb_rjit_call_threshold()) {
-            rb_rjit_compile(iseq);
+        else { // rb_rjit_call_p
+            if (body->jit_entry_calls == rb_rjit_call_threshold()) {
+                rb_rjit_compile(iseq);
+            }
         }
     }
-
-    return body->jit_func;
+    return body->jit_entry;
 }
 
-// Try to execute the current iseq in ec.  Use JIT code if it is ready.
-// If it is not, add ISEQ to the compilation queue and return Qundef for RJIT.
-// YJIT compiles on the thread running the iseq.
+// Execute JIT code compiled by jit_compile()
 static inline VALUE
 jit_exec(rb_execution_context_t *ec)
 {
@@ -425,6 +426,51 @@ jit_exec(rb_execution_context_t *ec)
 # define jit_exec(ec) Qundef
 #endif
 
+#if USE_YJIT
+// Generate JIT code that supports the following kind of ISEQ entry:
+//   * The first ISEQ pushed by vm_exec_handle_exception. The frame would
+//     point to a location specified by a catch table, and it doesn't have
+//     VM_FRAME_FLAG_FINISH. The current vm_exec stops if JIT code returns
+//     a non-Qundef value. So you should not return a non-Qundef value
+//     until ec->cfp is changed to a frame with VM_FRAME_FLAG_FINISH.
+static inline rb_jit_func_t
+jit_compile_exception(rb_execution_context_t *ec)
+{
+    // Increment the ISEQ's call counter
+    const rb_iseq_t *iseq = ec->cfp->iseq;
+    struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
+    if (rb_yjit_compile_new_iseqs()) {
+        body->jit_exception_calls++;
+    }
+    else {
+        return NULL;
+    }
+
+    // Trigger JIT compilation if not compiled
+    if (body->jit_exception == NULL && rb_yjit_threshold_hit(iseq, body->jit_exception_calls)) {
+        rb_yjit_compile_iseq(iseq, ec, true);
+    }
+    return body->jit_exception;
+}
+
+// Execute JIT code compiled by jit_compile_exception()
+static inline VALUE
+jit_exec_exception(rb_execution_context_t *ec)
+{
+    rb_jit_func_t func = jit_compile_exception(ec);
+    if (func) {
+        // Call the JIT code
+        return func(ec, ec->cfp);
+    }
+    else {
+        return Qundef;
+    }
+}
+#else
+# define jit_compile_exception(ec) ((rb_jit_func_t)0)
+# define jit_exec_exception(ec) Qundef
+#endif
+
 #include "vm_insnhelper.c"
 
 #include "vm_exec.c"
@@ -2381,8 +2427,11 @@ vm_exec_loop(rb_execution_context_t *ec, enum ruby_tag_type state,
 
     rb_ec_raised_reset(ec, RAISED_STACKOVERFLOW | RAISED_NOMEMORY);
     while (UNDEF_P(result = vm_exec_handle_exception(ec, state, result))) {
-        /* caught a jump, exec the handler */
-        result = vm_exec_core(ec);
+        // caught a jump, exec the handler. JIT code in jit_exec_exception()
+        // may return Qundef to run remaining frames with vm_exec_core().
+        if (UNDEF_P(result = jit_exec_exception(ec))) {
+            result = vm_exec_core(ec);
+        }
       vm_loop_start:
         VM_ASSERT(ec->tag == tag);
         /* when caught `throw`, `tag.state` is set. */

vm_core.h

@@ -503,10 +503,17 @@ struct rb_iseq_constant_body {
     const rb_iseq_t *mandatory_only_iseq;
 
 #if USE_RJIT || USE_YJIT
-    // Function pointer for JIT code
-    rb_jit_func_t jit_func;
-    // Number of total calls with jit_exec()
-    long unsigned total_calls;
+    // Function pointer for JIT code on jit_exec()
+    rb_jit_func_t jit_entry;
+    // Number of calls on jit_exec()
+    long unsigned jit_entry_calls;
+#endif
+
+#if USE_YJIT
+    // Function pointer for JIT code on jit_exec_exception()
+    rb_jit_func_t jit_exception;
+    // Number of calls on jit_exec_exception()
+    long unsigned jit_exception_calls;
 #endif
 
 #if USE_RJIT

vm_insnhelper.c

@@ -2485,6 +2485,12 @@ vm_base_ptr(const rb_control_frame_t *cfp)
     }
 }
 
+VALUE *
+rb_vm_base_ptr(const rb_control_frame_t *cfp)
+{
+    return vm_base_ptr(cfp);
+}
+
 /* method call processes with call_info */
 
 #include "vm_args.c"

yjit.c

@@ -422,10 +422,12 @@ void
 rb_iseq_reset_jit_func(const rb_iseq_t *iseq)
 {
     RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(iseq, imemo_iseq));
-    iseq->body->jit_func = NULL;
+    iseq->body->jit_entry = NULL;
+    iseq->body->jit_exception = NULL;
     // Enable re-compiling this ISEQ. Event when it's invalidated for TracePoint,
     // we'd like to re-compile ISEQs that haven't been converted to trace_* insns.
-    iseq->body->total_calls = 0;
+    iseq->body->jit_entry_calls = 0;
+    iseq->body->jit_exception_calls = 0;
 }
 
 // Get the PC for a given index in an iseq
@@ -597,12 +599,6 @@ rb_get_def_bmethod_proc(rb_method_definition_t *def)
     return def->body.bmethod.proc;
 }
 
-unsigned long
-rb_get_iseq_body_total_calls(const rb_iseq_t *iseq)
-{
-    return iseq->body->total_calls;
-}
-
 const rb_iseq_t *
 rb_get_iseq_body_local_iseq(const rb_iseq_t *iseq)
 {
@@ -832,6 +828,8 @@ rb_get_cfp_ep_level(struct rb_control_frame_struct *cfp, uint32_t lv)
     return ep;
 }
 
+extern VALUE *rb_vm_base_ptr(struct rb_control_frame_struct *cfp);
+
 VALUE
 rb_yarv_class_of(VALUE obj)
 {
@@ -1047,27 +1045,24 @@ rb_yjit_vm_unlock(unsigned int *recursive_lock_level, const char *file, int line
     rb_vm_lock_leave(recursive_lock_level, file, line);
 }
 
-bool
-rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec)
+void
+rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception)
 {
-    bool success = true;
     RB_VM_LOCK_ENTER();
     rb_vm_barrier();
 
-    // Compile a block version starting at the first instruction
-    uint8_t *rb_yjit_iseq_gen_entry_point(const rb_iseq_t *iseq, rb_execution_context_t *ec); // defined in Rust
-    uint8_t *code_ptr = rb_yjit_iseq_gen_entry_point(iseq, ec);
+    // Compile a block version starting at the current instruction
+    uint8_t *rb_yjit_iseq_gen_entry_point(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception); // defined in Rust
+    uint8_t *code_ptr = rb_yjit_iseq_gen_entry_point(iseq, ec, jit_exception);
 
-    if (code_ptr) {
-        iseq->body->jit_func = (rb_jit_func_t)code_ptr;
+    if (jit_exception) {
+        iseq->body->jit_exception = (rb_jit_func_t)code_ptr;
     }
     else {
-        iseq->body->jit_func = 0;
-        success = false;
+        iseq->body->jit_entry = (rb_jit_func_t)code_ptr;
     }
 
     RB_VM_LOCK_LEAVE();
-    return success;
 }
 
 // GC root for interacting with the GC
@@ -1143,6 +1138,35 @@ rb_yjit_invokeblock_sp_pops(const struct rb_callinfo *ci)
     return 1 - sp_inc_of_invokeblock(ci); // + 1 to ignore return value push
 }
 
+// Setup jit_return to avoid returning a non-Qundef value on a non-FINISH frame.
+// See [jit_compile_exception] for details.
+void
+rb_yjit_set_exception_return(rb_control_frame_t *cfp, void *leave_exit, void *leave_exception)
+{
+    if (VM_FRAME_FINISHED_P(cfp)) {
+        // If it's a FINISH frame, just normally exit with a non-Qundef value.
+        cfp->jit_return = leave_exit;
+    }
+    else if (cfp->jit_return) {
+        while (!VM_FRAME_FINISHED_P(cfp)) {
+            if (cfp->jit_return == leave_exit) {
+                // Unlike jit_exec(), leave_exit is not safe on a non-FINISH frame on
+                // jit_exec_exception(). See [jit_exec] and [jit_exec_exception] for
+                // details. Exit to the interpreter with Qundef to let it keep executing
+                // other Ruby frames.
+                cfp->jit_return = leave_exception;
+                return;
+            }
+            cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp);
+        }
+    }
+    else {
+        // If the caller was not JIT code, exit to the interpreter with Qundef
+        // to keep executing Ruby frames with the interpreter.
+        cfp->jit_return = leave_exception;
+    }
+}
+
 // Primitives used by yjit.rb
 VALUE rb_yjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_trace_exit_locations_enabled_p(rb_execution_context_t *ec, VALUE self);

yjit.h

@@ -27,12 +27,12 @@
 // Expose these as declarations since we are building YJIT.
 bool rb_yjit_enabled_p(void);
 bool rb_yjit_compile_new_iseqs(void);
-bool rb_yjit_threshold_hit(const rb_iseq_t *const iseq);
+bool rb_yjit_threshold_hit(const rb_iseq_t *const iseq, unsigned long total_calls);
 void rb_yjit_invalidate_all_method_lookup_assumptions(void);
 void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme);
 void rb_yjit_collect_binding_alloc(void);
 void rb_yjit_collect_binding_set(void);
-bool rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec);
+void rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception);
 void rb_yjit_init(void);
 void rb_yjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
 void rb_yjit_constant_state_changed(ID id);
@@ -49,12 +49,12 @@ void rb_yjit_tracing_invalidate_all(void);
 
 static inline bool rb_yjit_enabled_p(void) { return false; }
 static inline bool rb_yjit_compile_new_iseqs(void) { return false; }
-static inline bool rb_yjit_threshold_hit(const rb_iseq_t *const iseq) { return false; }
+static inline bool rb_yjit_threshold_hit(const rb_iseq_t *const iseq, unsigned long total_calls) { return false; }
 static inline void rb_yjit_invalidate_all_method_lookup_assumptions(void) {}
 static inline void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme) {}
 static inline void rb_yjit_collect_binding_alloc(void) {}
 static inline void rb_yjit_collect_binding_set(void) {}
-static inline bool rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec) { return false; }
+static inline void rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception) {}
 static inline void rb_yjit_init(void) {}
 static inline void rb_yjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop) {}
 static inline void rb_yjit_constant_state_changed(ID id) {}