YJIT: Implement throw instruction (#7491)

* Break up jit_exec from vm_sendish

* YJIT: Implement throw instruction

* YJIT: Explain what rb_vm_throw does [ci skip]

insns.def

@@ -814,6 +814,12 @@ send
     VALUE bh = vm_caller_setup_arg_block(ec, GET_CFP(), cd->ci, blockiseq, false);
     val = vm_sendish(ec, GET_CFP(), cd, bh, mexp_search_method);
 
+    jit_func_t func;
+    if (val == Qundef && (func = jit_compile(ec))) {
+        val = func(ec, ec->cfp);
+        if (ec->tag->state) THROW_EXCEPTION(val);
+    }
+
     if (val == Qundef) {
         RESTORE_REGS();
         NEXT_INSN();
@@ -833,6 +839,12 @@ opt_send_without_block
     VALUE bh = VM_BLOCK_HANDLER_NONE;
     val = vm_sendish(ec, GET_CFP(), cd, bh, mexp_search_method);
 
+    jit_func_t func;
+    if (val == Qundef && (func = jit_compile(ec))) {
+        val = func(ec, ec->cfp);
+        if (ec->tag->state) THROW_EXCEPTION(val);
+    }
+
     if (val == Qundef) {
         RESTORE_REGS();
         NEXT_INSN();
@@ -935,6 +947,12 @@ invokesuper
     VALUE bh = vm_caller_setup_arg_block(ec, GET_CFP(), cd->ci, blockiseq, true);
     val = vm_sendish(ec, GET_CFP(), cd, bh, mexp_search_super);
 
+    jit_func_t func;
+    if (val == Qundef && (func = jit_compile(ec))) {
+        val = func(ec, ec->cfp);
+        if (ec->tag->state) THROW_EXCEPTION(val);
+    }
+
     if (val == Qundef) {
         RESTORE_REGS();
         NEXT_INSN();
@@ -954,6 +972,12 @@ invokeblock
     VALUE bh = VM_BLOCK_HANDLER_NONE;
     val = vm_sendish(ec, GET_CFP(), cd, bh, mexp_search_invokeblock);
 
+    jit_func_t func;
+    if (val == Qundef && (func = jit_compile(ec))) {
+        val = func(ec, ec->cfp);
+        if (ec->tag->state) THROW_EXCEPTION(val);
+    }
+
     if (val == Qundef) {
         RESTORE_REGS();
         NEXT_INSN();

vm.c

@@ -369,11 +369,9 @@ 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 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.
-static inline VALUE
-jit_exec(rb_execution_context_t *ec)
+// Try to compile the current ISeq in ec. Return 0 if not compiled.
+static inline jit_func_t
+jit_compile(rb_execution_context_t *ec)
 {
     // Increment the ISEQ's call counter
     const rb_iseq_t *iseq = ec->cfp->iseq;
@@ -383,43 +381,42 @@ jit_exec(rb_execution_context_t *ec)
         body->total_calls++;
     }
     else {
-        return Qundef;
+        return 0;
     }
 
     // Trigger JIT compilation as needed
-    jit_func_t func;
     if (yjit_enabled) {
         if (body->total_calls == rb_yjit_call_threshold())  {
-            // If we couldn't generate any code for this iseq, then return
-            // Qundef so the interpreter will handle the call.
-            if (!rb_yjit_compile_iseq(iseq, ec)) {
-                return Qundef;
-            }
-        }
-        // YJIT tried compiling this function once before and couldn't do
-        // it, so return Qundef so the interpreter handles it.
-        if ((func = body->jit_func) == 0) {
-            return Qundef;
+            rb_yjit_compile_iseq(iseq, ec);
         }
     }
     else { // rb_rjit_call_p
         if (body->total_calls == rb_rjit_call_threshold()) {
             rb_rjit_compile(iseq);
         }
-        if ((func = body->jit_func) == 0) {
-            return Qundef;
-        }
     }
 
-    // Call the JIT code
-    return func(ec, ec->cfp); // SystemV x64 calling convention: ec -> RDI, cfp -> RSI
+    return body->jit_func;
 }
-#else
+
+// 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.
 static inline VALUE
 jit_exec(rb_execution_context_t *ec)
 {
-    return Qundef;
+    jit_func_t func = jit_compile(ec);
+    if (func) {
+        // Call the JIT code
+        return func(ec, ec->cfp);
+    }
+    else {
+        return Qundef;
+    }
 }
+#else
+static inline jit_func_t jit_compile(rb_execution_context_t *ec) { return 0; }
+static inline VALUE jit_exec(rb_execution_context_t *ec) { return Qundef; }
 #endif
 
 #include "vm_insnhelper.c"

vm_insnhelper.c

@@ -1823,6 +1823,12 @@ vm_throw(const rb_execution_context_t *ec, rb_control_frame_t *reg_cfp,
     }
 }
 
+VALUE
+rb_vm_throw(const rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, rb_num_t throw_state, VALUE throwobj)
+{
+    return vm_throw(ec, reg_cfp, throw_state, throwobj);
+}
+
 static inline void
 vm_expandarray(VALUE *sp, VALUE ary, rb_num_t num, int flag)
 {
@@ -5211,15 +5217,7 @@ vm_sendish(
         val = vm_invokeblock_i(ec, GET_CFP(), &calling);
         break;
     }
-
-    if (!UNDEF_P(val)) {
-        return val;             /* CFUNC normal return */
-    }
-    else {
-        RESTORE_REGS();         /* CFP pushed in cc->call() */
-    }
-
-    return jit_exec(ec);
+    return val;
 }
 
 /* object.c */

yjit/src/codegen.rs

@@ -3699,6 +3699,40 @@ fn gen_branchnil(
     EndBlock
 }
 
+fn gen_throw(
+    jit: &mut JITState,
+    ctx: &mut Context,
+    asm: &mut Assembler,
+    _ocb: &mut OutlinedCb,
+) -> CodegenStatus {
+    let throw_state = jit.get_arg(0).as_u64();
+    let throwobj = asm.load(ctx.stack_pop(1));
+
+    // THROW_DATA_NEW allocates. Save SP for GC and PC for allocation tracing as
+    // well as handling the catch table. However, not using jit_prepare_routine_call
+    // since we don't need a patch point for this implementation.
+    jit_save_pc(jit, asm);
+    gen_save_sp(asm, ctx);
+
+    // rb_vm_throw verifies it's a valid throw, sets ec->tag->state, and returns throw
+    // data, which is throwobj or a vm_throw_data wrapping it. When ec->tag->state is
+    // set, JIT code callers will handle the throw with vm_exec_handle_exception.
+    extern "C" {
+        fn rb_vm_throw(ec: EcPtr, reg_cfp: CfpPtr, throw_state: u32, throwobj: VALUE) -> VALUE;
+    }
+    let val = asm.ccall(rb_vm_throw as *mut u8, vec![EC, CFP, throw_state.into(), throwobj]);
+
+    asm.comment("exit from throw");
+    asm.cpop_into(SP);
+    asm.cpop_into(EC);
+    asm.cpop_into(CFP);
+
+    asm.frame_teardown();
+
+    asm.cret(val);
+    EndBlock
+}
+
 fn gen_jump(
     jit: &mut JITState,
     ctx: &mut Context,
@@ -7741,6 +7775,7 @@ fn get_gen_fn(opcode: VALUE) -> Option<InsnGenFn> {
         YARVINSN_branchif => Some(gen_branchif),
         YARVINSN_branchunless => Some(gen_branchunless),
         YARVINSN_branchnil => Some(gen_branchnil),
+        YARVINSN_throw => Some(gen_throw),
         YARVINSN_jump => Some(gen_jump),
 
         YARVINSN_getblockparamproxy => Some(gen_getblockparamproxy),