YJIT: Compile exception handlers

lib/ruby_vm/rjit/compiler.rb

@@ -235,14 +235,22 @@ def invalidate_block(block)
     def compile_prologue(asm, iseq, pc)
       asm.comment('RJIT entry point')
 
+      # handle_exception is not supported yet
+      not_exception = asm.new_label('not_exception')
+      asm.test(C_ARGS[2], C_ARGS[2])
+      asm.jz(not_exception)
+      asm.mov(C_RET, Qundef)
+      asm.ret
+      asm.write_label(not_exception)
+
       # Save callee-saved registers used by JITed code
       asm.push(CFP)
       asm.push(EC)
       asm.push(SP)
 
       # Move arguments EC and CFP to dedicated registers
-      asm.mov(EC, :rdi)
-      asm.mov(CFP, :rsi)
+      asm.mov(EC, C_ARGS[0])
+      asm.mov(CFP, C_ARGS[1])
 
       # Load sp to a dedicated register
       asm.mov(SP, [CFP, C.rb_control_frame_t.offsetof(:sp)]) # rbx = cfp->sp

vm.c

@@ -372,7 +372,7 @@ static VALUE vm_invoke_proc(rb_execution_context_t *ec, rb_proc_t *proc, VALUE s
 #if USE_RJIT || USE_YJIT
 // Try to compile the current ISeq in ec. Return 0 if not compiled.
 static inline rb_jit_func_t
-jit_compile(rb_execution_context_t *ec)
+jit_compile(rb_execution_context_t *ec, bool handle_exception)
 {
     // Increment the ISEQ's call counter
     const rb_iseq_t *iseq = ec->cfp->iseq;
@@ -397,7 +397,7 @@ jit_compile(rb_execution_context_t *ec)
         }
     }
     else { // rb_rjit_call_p
-        if (body->total_calls == rb_rjit_call_threshold()) {
+        if (body->total_calls == rb_rjit_call_threshold() && !handle_exception) {
             rb_rjit_compile(iseq);
         }
     }
@@ -409,20 +409,20 @@ jit_compile(rb_execution_context_t *ec)
 // If it is not, add ISEQ to the compilation queue and return Qundef for RJIT.
 // YJIT compiles on the thread running the iseq.
 static inline VALUE
-jit_exec(rb_execution_context_t *ec)
+jit_exec(rb_execution_context_t *ec, bool handle_exception)
 {
-    rb_jit_func_t func = jit_compile(ec);
+    rb_jit_func_t func = jit_compile(ec, handle_exception);
     if (func) {
         // Call the JIT code
-        return func(ec, ec->cfp);
+        return func(ec, ec->cfp, handle_exception);
     }
     else {
         return Qundef;
     }
 }
 #else
-static inline rb_jit_func_t jit_compile(rb_execution_context_t *ec) { return 0; }
-static inline VALUE jit_exec(rb_execution_context_t *ec) { return Qundef; }
+static inline rb_jit_func_t jit_compile(rb_execution_context_t *ec, bool handle_exception) { return 0; }
+static inline VALUE jit_exec(rb_execution_context_t *ec, bool handle_exception) { return Qundef; }
 #endif
 
 #include "vm_insnhelper.c"
@@ -2336,7 +2336,7 @@ vm_exec_bottom_main(void *context)
     struct rb_vm_exec_context *ctx = (struct rb_vm_exec_context *)context;
 
     ctx->state = TAG_NONE;
-    if (UNDEF_P(ctx->result = jit_exec(ctx->ec))) {
+    if (UNDEF_P(ctx->result = jit_exec(ctx->ec, false))) {
         ctx->result = vm_exec_core(ctx->ec);
     }
     vm_exec_enter_vm_loop(ctx->ec, ctx, ctx->tag, true);
@@ -2386,7 +2386,7 @@ vm_exec(rb_execution_context_t *ec)
 
     _tag.retval = Qnil;
     if ((state = EC_EXEC_TAG()) == TAG_NONE) {
-        if (UNDEF_P(result = jit_exec(ec))) {
+        if (UNDEF_P(result = jit_exec(ec, false))) {
             result = vm_exec_core(ec);
         }
         goto vm_loop_start; /* fallback to the VM */
@@ -2396,7 +2396,9 @@ vm_exec(rb_execution_context_t *ec)
         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);
+            if (UNDEF_P(result = jit_exec(ec, true))) {
+                result = vm_exec_core(ec);
+            }
           vm_loop_start:
             VM_ASSERT(ec->tag == &_tag);
             /* when caught `throw`, `tag.state` is set. */

vm_core.h

@@ -371,7 +371,7 @@ enum rb_builtin_attr {
     BUILTIN_ATTR_SINGLE_NOARG_INLINE = 0x04,
 };
 
-typedef VALUE (*rb_jit_func_t)(struct rb_execution_context_struct *, struct rb_control_frame_struct *);
+typedef VALUE (*rb_jit_func_t)(struct rb_execution_context_struct *, struct rb_control_frame_struct *, bool);
 
 struct rb_iseq_constant_body {
     enum rb_iseq_type type;

vm_exec.h

@@ -176,8 +176,8 @@ default:                        \
 // Run the JIT from the interpreter
 #define JIT_EXEC(ec, val) do { \
     rb_jit_func_t func; \
-    if (val == Qundef && (func = jit_compile(ec))) { \
-        val = func(ec, ec->cfp); \
+    if (val == Qundef && (func = jit_compile(ec, false))) { \
+        val = func(ec, ec->cfp, false); \
         RESTORE_REGS(); /* fix cfp for tailcall */ \
         if (ec->tag->state) THROW_EXCEPTION(val); \
     } \

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

@@ -615,6 +615,12 @@ rb_get_iseq_body_parent_iseq(const rb_iseq_t *iseq)
     return iseq->body->parent_iseq;
 }
 
+void *
+rb_get_cfp_jit_return(struct rb_control_frame_struct *cfp)
+{
+    return cfp->jit_return;
+}
+
 unsigned int
 rb_get_iseq_body_local_table_size(const rb_iseq_t *iseq)
 {
@@ -832,6 +838,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)
 {
@@ -1143,6 +1151,24 @@ rb_yjit_invokeblock_sp_pops(const struct rb_callinfo *ci)
     return 1 - sp_inc_of_invokeblock(ci); // + 1 to ignore return value push
 }
 
+// Set up the JIT return address to finish the most recent vm_exec call successfully.
+void
+rb_yjit_set_jit_return(rb_control_frame_t *cfp, void *leave_compile, void *leave_exit, bool handle_exception)
+{
+    // When JIT code exits a non-FINISH frame and handle_exception is true,
+    // the return value must be Qundef because the caller will not process
+    // remaining frames in that case.
+    // To prevent that situation, we make sure each frame returns Qundef or
+    // proceeds to compile the caller frame until it reaches a FINISH frame.
+    if (handle_exception) {
+        while (!VM_FRAME_FINISHED_P(cfp)) {
+            cfp->jit_return = leave_compile;
+            cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp);
+        }
+    }
+    cfp->jit_return = leave_exit;
+}
+
 // 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/bindgen/src/main.rs

@@ -326,6 +326,7 @@ fn main() {
         .allowlist_function("rb_yjit_assert_holding_vm_lock")
         .allowlist_function("rb_yjit_sendish_sp_pops")
         .allowlist_function("rb_yjit_invokeblock_sp_pops")
+        .allowlist_function("rb_yjit_set_jit_return")
         .allowlist_type("robject_offsets")
         .allowlist_type("rstring_offsets")
 
@@ -443,6 +444,8 @@ fn main() {
         .allowlist_function("rb_yjit_array_len")
         .allowlist_function("rb_obj_class")
         .allowlist_function("rb_obj_is_proc")
+        .allowlist_function("rb_vm_base_ptr")
+        .allowlist_function("rb_get_cfp_jit_return")
 
         // We define VALUE manually, don't import it
         .blocklist_type("VALUE")

yjit/src/codegen.rs

@@ -662,31 +662,33 @@ pub fn gen_entry_prologue(cb: &mut CodeBlock, ocb: &mut OutlinedCb, iseq: IseqPt
     asm.cpush(EC);
     asm.cpush(SP);
 
-    // We are passed EC and CFP as arguments
+    // We are passed EC, CFP, and handle_exception flag as arguments
     asm.mov(EC, C_ARG_OPNDS[0]);
     asm.mov(CFP, C_ARG_OPNDS[1]);
+    let handle_exception = C_ARG_OPNDS[2];
 
     // Load the current SP from the CFP into REG_SP
     asm.mov(SP, Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SP));
 
     // Setup cfp->jit_return
-    asm.mov(
-        Opnd::mem(64, CFP, RUBY_OFFSET_CFP_JIT_RETURN),
-        Opnd::const_ptr(CodegenGlobals::get_leave_exit_code().raw_ptr()),
+    asm.ccall(
+        rb_yjit_set_jit_return as _,
+        vec![
+            CFP,
+            CodegenGlobals::get_leave_compile_code().into_usize().into(),
+            CodegenGlobals::get_leave_exit_code().into_usize().into(),
+            handle_exception,
+        ],
     );
 
-    // We're compiling iseqs that we *expect* to start at `insn_idx`. But in
-    // the case of optional parameters, the interpreter can set the pc to a
-    // different location depending on the optional parameters.  If an iseq
-    // has optional parameters, we'll add a runtime check that the PC we've
+    // We're compiling iseqs that we *expect* to start at `insn_idx`.
+    // But in the case of optional parameters or when handling exceptions,
+    // the interpreter can set the pc to a different location. For
+    // such scenarios, we'll add a runtime check that the PC we've
     // compiled for is the same PC that the interpreter wants us to run with.
     // If they don't match, then we'll jump to an entry stub and generate
     // another PC check and entry there.
-    let pending_entry = if unsafe { get_iseq_flags_has_opt(iseq) } {
-        Some(gen_entry_chain_guard(&mut asm, ocb, iseq, insn_idx)?)
-    } else {
-        None
-    };
+    let pending_entry = gen_entry_chain_guard(&mut asm, ocb, iseq, insn_idx)?;
 
     asm.compile(cb, Some(ocb));
 
@@ -699,11 +701,9 @@ pub fn gen_entry_prologue(cb: &mut CodeBlock, ocb: &mut OutlinedCb, iseq: IseqPt
             iseq_payload.pages.insert(page);
         }
         // Write an entry to the heap and push it to the ISEQ
-        if let Some(pending_entry) = pending_entry {
-            let pending_entry = Rc::try_unwrap(pending_entry)
-                .ok().expect("PendingEntry should be unique");
-            iseq_payload.entries.push(pending_entry.into_entry());
-        }
+        let pending_entry = Rc::try_unwrap(pending_entry)
+            .ok().expect("PendingEntry should be unique");
+        iseq_payload.entries.push(pending_entry.into_entry());
         Some(code_ptr)
     }
 }
@@ -8293,6 +8293,9 @@ pub struct CodegenGlobals {
     /// Code for exiting back to the interpreter from the leave instruction
     leave_exit_code: CodePtr,
 
+    /// Code for compiling and jumping to the caller
+    leave_compile_code: CodePtr,
+
     // For exiting from YJIT frame from branch_stub_hit().
     // Filled by gen_code_for_exit_from_stub().
     stub_exit_code: CodePtr,
@@ -8383,6 +8386,7 @@ impl CodegenGlobals {
 
         let ocb_start_addr = ocb.unwrap().get_write_ptr();
         let leave_exit_code = gen_leave_exit(&mut ocb);
+        let leave_compile_code = gen_leave_compile(&mut ocb);
 
         let stub_exit_code = gen_code_for_exit_from_stub(&mut ocb);
 
@@ -8403,6 +8407,7 @@ impl CodegenGlobals {
             inline_cb: cb,
             outlined_cb: ocb,
             leave_exit_code,
+            leave_compile_code,
             stub_exit_code: stub_exit_code,
             outline_full_cfunc_return_pos: cfunc_exit_code,
             branch_stub_hit_trampoline,
@@ -8523,6 +8528,10 @@ impl CodegenGlobals {
         CodegenGlobals::get_instance().leave_exit_code
     }
 
+    pub fn get_leave_compile_code() -> CodePtr {
+        CodegenGlobals::get_instance().leave_compile_code
+    }
+
     pub fn get_stub_exit_code() -> CodePtr {
         CodegenGlobals::get_instance().stub_exit_code
     }