YJIT: use a context-free landing pad to optimize `leave`

This allows gen_leave to always do an indirect to cfp->jit_return.

1 files changed, 36 insertions, 15 deletions

diff --git a/yjit_codegen.c b/yjit_codegen.c
index 4d4f182b6d..39a2667f5d 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -27,6 +27,9 @@ codeblock_t* cb = NULL;
 static codeblock_t outline_block;
 codeblock_t* ocb = NULL;
 
+// Code for exiting back to the interpreter
+static void *interp_exit;
+
 // Print the current source location for debugging purposes
 RBIMPL_ATTR_MAYBE_UNUSED()
 static void
@@ -137,7 +140,7 @@ yjit_load_regs(codeblock_t* cb)
 static uint8_t *
 yjit_gen_exit(jitstate_t *jit, ctx_t *ctx, codeblock_t *cb)
 {
-    uint8_t *code_ptr = cb_get_ptr(ocb, ocb->write_pos);
+    uint8_t *code_ptr = cb_get_ptr(cb, cb->write_pos);
 
     VALUE *exit_pc = jit->pc;
 
@@ -184,6 +187,26 @@ yjit_gen_exit(jitstate_t *jit, ctx_t *ctx, codeblock_t *cb)
     return code_ptr;
 }
 
+// Generate an interpreter to REG_CFP->pc.
+static uint8_t *
+yjit_gen_context_free_exit(codeblock_t *cb)
+{
+    uint8_t *code_ptr = cb_get_ptr(cb, cb->write_pos);
+
+    // Update the CFP on the EC
+    mov(cb, member_opnd(REG_EC, rb_execution_context_t, cfp), REG_CFP);
+
+    // Put PC into the return register, which the post call bytes dispatches to
+    mov(cb, RAX, member_opnd(REG_CFP, rb_control_frame_t, pc));
+
+    cb_write_post_call_bytes(cb);
+
+    // Note, not incrementing stats here since this exit is the natural end to
+    // executing output code.
+
+    return code_ptr;
+}
+
 
 // A shorthand for generating an exit in the outline block
 static uint8_t *
@@ -251,7 +274,7 @@ yjit_comment_array_t yjit_code_comments;
 Compile an interpreter entry block to be inserted into an iseq
 Returns `NULL` if compilation fails.
 */
-uint8_t*
+uint8_t *
 yjit_entry_prologue(void)
 {
     RUBY_ASSERT(cb != NULL);
@@ -271,6 +294,11 @@ yjit_entry_prologue(void)
     // Load the current SP from the CFP into REG_SP
     mov(cb, REG_SP, member_opnd(REG_CFP, rb_control_frame_t, sp));
 
+    // Setup cfp->jit_return
+    // TODO: this could use a IP relative LEA instead of an 8 byte immediate
+    mov(cb, REG0, const_ptr_opnd(interp_exit));
+    mov(cb, member_opnd(REG_CFP, rb_control_frame_t, jit_return), REG0);
+
     return code_ptr;
 }
 
@@ -2076,21 +2104,9 @@ gen_leave(jitstate_t* jit, ctx_t* ctx)
     mov(cb, REG_SP, member_opnd(REG_CFP, rb_control_frame_t, sp));
     mov(cb, mem_opnd(64, REG_SP, -SIZEOF_VALUE), REG0);
 
-    // If the return address is NULL, fall back to the interpreter
-    ADD_COMMENT(cb, "check for jit return");
-    int FALLBACK_LABEL = cb_new_label(cb, "FALLBACK");
-    test(cb, REG1, REG1);
-    jz_label(cb, FALLBACK_LABEL);
-
     // Jump to the JIT return address
     jmp_rm(cb, REG1);
 
-    // Fall back to the interpreter
-    cb_write_label(cb, FALLBACK_LABEL);
-    cb_link_labels(cb);
-    GEN_COUNTER_INC(cb, leave_interp_return);
-    cb_write_post_call_bytes(cb);
-
     return YJIT_END_BLOCK;
 }
 
@@ -2151,12 +2167,17 @@ yjit_init_codegen(void)
 {
     // Initialize the code blocks
     uint32_t mem_size = 128 * 1024 * 1024;
-    uint8_t* mem_block = alloc_exec_mem(mem_size);
+    uint8_t *mem_block = alloc_exec_mem(mem_size);
+
     cb = █
     cb_init(cb, mem_block, mem_size/2);
+
     ocb = &outline_block;
     cb_init(ocb, mem_block + mem_size/2, mem_size/2);
 
+    // Generate interp_exit
+    interp_exit = yjit_gen_context_free_exit(cb);
+
     // Map YARV opcodes to the corresponding codegen functions
     yjit_reg_op(BIN(dup), gen_dup);
     yjit_reg_op(BIN(nop), gen_nop);