MicroJIT: compile after ten calls

compile.c

@@ -862,46 +862,16 @@ rb_iseq_translate_threaded_code(rb_iseq_t *iseq)
 {
 #if OPT_DIRECT_THREADED_CODE || OPT_CALL_THREADED_CODE
     const void * const *table = rb_vm_get_insns_address_table();
+    unsigned int i;
     VALUE *encoded = (VALUE *)iseq->body->iseq_encoded;
-    VALUE translated_insns_buf;
 
-    unsigned int insn_idx, translated_idx;
-    unsigned int next_ujit_idx = 0;
-    unsigned int translated_len = 0;
-
-    bool ujit_enabled = rb_ujit_enabled_p();
-
-    VALUE *translated_insns = ALLOCV_N(VALUE, translated_insns_buf, iseq->body->iseq_size);
-    for (insn_idx = 0; insn_idx < iseq->body->iseq_size; /* */) {
-        int insn = (int)encoded[insn_idx];
-        int len = insn_len(insn);
-        VALUE translated;
-
-        uint8_t* native_code_ptr = NULL;
-
-        // If ujit is enabled and hasn't already compiled this instruction
-        if (ujit_enabled && insn_idx >= next_ujit_idx)
-            native_code_ptr = ujit_compile_insn(iseq, insn_idx, &next_ujit_idx);
-
-        if (native_code_ptr)
-            translated = (VALUE)native_code_ptr;
-        else
-            translated = (VALUE)table[insn];
-        translated_insns[translated_len++] = translated;
-
-    	insn_idx += len;
+    for (i = 0; i < iseq->body->iseq_size; /* */ ) {
+	int insn = (int)iseq->body->iseq_encoded[i];
+	int len = insn_len(insn);
+	encoded[i] = (VALUE)table[insn];
+	i += len;
     }
-
-    insn_idx = 0;
-    for (translated_idx = 0; translated_idx < translated_len; translated_idx++) {
-        int insn = (int)encoded[insn_idx];
-        int len = insn_len(insn);
-        encoded[insn_idx] = translated_insns[translated_idx];
-        insn_idx += len;
-    }
-
     FL_SET((VALUE)iseq, ISEQ_TRANSLATED);
-    ALLOCV_END(translated_insns_buf);
 #endif
     return COMPILE_OK;
 }

mjit.h

@@ -16,6 +16,7 @@
 
 #include "debug_counter.h"
 #include "ruby.h"
+#include "ujit_compile.h"
 
 // Special address values of a function generated from the
 // corresponding iseq by MJIT:
@@ -142,7 +143,7 @@ mjit_exec(rb_execution_context_t *ec)
     const rb_iseq_t *iseq;
     struct rb_iseq_constant_body *body;
 
-    if (!mjit_call_p)
+    if (!mjit_call_p && !rb_ujit_enabled_p())
         return Qundef;
     RB_DEBUG_COUNTER_INC(mjit_exec);
 
@@ -150,6 +151,11 @@ mjit_exec(rb_execution_context_t *ec)
     body = iseq->body;
     body->total_calls++;
 
+    const int ujit_call_threashold = 10;
+    if (body->total_calls == ujit_call_threashold) {
+        rb_ujit_compile_iseq(iseq);
+    }
+
     mjit_func_t func = body->jit_func;
     if (UNLIKELY((uintptr_t)func <= LAST_JIT_ISEQ_FUNC)) {
 #  ifdef MJIT_HEADER

ujit_compile.c

@@ -4,6 +4,7 @@
 #include "vm_core.h"
 #include "vm_callinfo.h"
 #include "builtin.h"
+#include "internal/compile.h"
 #include "insns_info.inc"
 #include "ujit_compile.h"
 #include "ujit_asm.h"
@@ -19,6 +20,8 @@
 #define PLATFORM_SUPPORTED_P 1
 #endif
 
+bool rb_ujit_enabled;
+
 // Hash table of encoded instructions
 extern st_table *rb_encoded_insn_data;
 
@@ -26,11 +29,13 @@ extern st_table *rb_encoded_insn_data;
 typedef struct ctx_struct
 {
     // Current PC
-    VALUE* pc;
+    VALUE *pc;
 
     // Difference between the current stack pointer and actual stack top
     int32_t stack_diff;
 
+    const rb_iseq_t *iseq;
+
 } ctx_t;
 
 // MicroJIT code generation function signature
@@ -63,12 +68,25 @@ addr2insn_bookkeeping(void *code_ptr, int insn)
     }
 }
 
-// Get the current instruction opcode from the context object
-int ctx_get_opcode(ctx_t* ctx)
+static int
+opcode_at_pc(const rb_iseq_t *iseq, const VALUE *pc)
 {
-    return (int)(*ctx->pc);
+    const VALUE at_pc = *pc;
+    if (FL_TEST_RAW((VALUE)iseq, ISEQ_TRANSLATED)) {
+        return rb_vm_insn_addr2insn((const void *)at_pc);
+    }
+    else {
+        return (int)at_pc;
+    }
 }
 
+// Get the current instruction opcode from the context object
+int ctx_get_opcode(ctx_t *ctx)
+{
+    return opcode_at_pc(ctx->iseq, ctx->pc);
+}
+
+
 // Get an instruction argument from the context object
 VALUE ctx_get_arg(ctx_t* ctx, size_t arg_idx)
 {
@@ -167,7 +185,7 @@ ujit_side_exit(codeblock_t* cb, ctx_t* ctx, VALUE* exit_pc)
     // Otherwise the interpreter may jump right back to the
     // JITted code we're trying to exit
     const void * const *table = rb_vm_get_insns_address_table();
-    int opcode = (int)(*exit_pc);
+    int opcode = opcode_at_pc(ctx->iseq, exit_pc);
     void* old_instr = (void*)table[opcode];
     mov(cb, RAX, const_ptr_opnd(exit_pc));
     mov(cb, RCX, const_ptr_opnd(old_instr));
@@ -192,11 +210,12 @@ System V ABI reference:
 https://wiki.osdev.org/System_V_ABI#x86-64
 */
 uint8_t *
-ujit_compile_insn(rb_iseq_t *iseq, unsigned int insn_idx, unsigned int* next_ujit_idx)
+ujit_compile_insn(const rb_iseq_t *iseq, unsigned int insn_idx, unsigned int* next_ujit_idx)
 {
     if (!cb) {
         return NULL;
     }
+    VALUE *encoded = iseq->body->iseq_encoded;
 
     // NOTE: if we are ever deployed in production, we
     // should probably just log an error and return NULL here,
@@ -218,19 +237,20 @@ ujit_compile_insn(rb_iseq_t *iseq, unsigned int insn_idx, unsigned int* next_uji
     //printf("write pos: %ld\n", cb->write_pos);
 
     // Get the first opcode in the sequence
-    int first_opcode = (int)iseq->body->iseq_encoded[insn_idx];
+    int first_opcode = opcode_at_pc(iseq, &encoded[insn_idx]);
 
     // Create codegen context
     ctx_t ctx;
     ctx.pc = NULL;
     ctx.stack_diff = 0;
+    ctx.iseq = iseq;
 
     // For each instruction to compile
     size_t num_instrs;
     for (num_instrs = 0;; ++num_instrs)
     {
         // Set the current PC
-        ctx.pc = &iseq->body->iseq_encoded[insn_idx];
+        ctx.pc = &encoded[insn_idx];
 
         // Get the current opcode
         int opcode = ctx_get_opcode(&ctx);
@@ -616,10 +636,33 @@ gen_opt_send_without_block(codeblock_t* cb, codeblock_t* ocb, ctx_t* ctx)
     */
 }
 
-bool
-rb_ujit_enabled_p(void)
+
+void
+rb_ujit_compile_iseq(const rb_iseq_t *iseq)
 {
-    return !!cb;
+#if OPT_DIRECT_THREADED_CODE || OPT_CALL_THREADED_CODE
+    VALUE *encoded = (VALUE *)iseq->body->iseq_encoded;
+
+    unsigned int insn_idx;
+    unsigned int next_ujit_idx = 0;
+
+    for (insn_idx = 0; insn_idx < iseq->body->iseq_size; /* */) {
+        int insn = opcode_at_pc(iseq, &encoded[insn_idx]);
+        int len = insn_len(insn);
+
+        uint8_t *native_code_ptr = NULL;
+
+        // If ujit hasn't already compiled this instruction
+        if (insn_idx >= next_ujit_idx) {
+            native_code_ptr = ujit_compile_insn(iseq, insn_idx, &next_ujit_idx);
+        }
+
+        if (native_code_ptr) {
+            encoded[insn_idx] = (VALUE)native_code_ptr;
+        }
+        insn_idx += len;
+    }
+#endif
 }
 
 void
@@ -630,6 +673,7 @@ rb_ujit_init(void)
         return;
     }
 
+    rb_ujit_enabled = true;
     // Initialize the code blocks
     size_t mem_size = 128 * 1024 * 1024;
     uint8_t* mem_block = alloc_exec_mem(mem_size);

ujit_compile.h

@@ -10,8 +10,16 @@ typedef struct rb_iseq_struct rb_iseq_t;
 #define rb_iseq_t rb_iseq_t
 #endif
 
+extern bool rb_ujit_enabled;
+
+static inline
+bool rb_ujit_enabled_p(void)
+{
+    return rb_ujit_enabled;
+}
+
 void rb_ujit_init(void);
-bool rb_ujit_enabled_p(void);
-uint8_t* ujit_compile_insn(rb_iseq_t *iseq, unsigned int insn_idx, unsigned int* next_ujit_idx);
+uint8_t *ujit_compile_insn(const rb_iseq_t *iseq, unsigned int insn_idx, unsigned int *next_ujit_idx);
+void rb_ujit_compile_iseq(const rb_iseq_t *iseq);
 
 #endif