Implement Multi-tier JIT (#1774)

Implement 2-level Multi-tier JIT engine: tier-up from Fast JIT to LLVM JIT to
get quick cold startup by Fast JIT and better performance by gradually
switching to LLVM JIT when the LLVM JIT functions are compiled by the
backend threads.

Refer to:
https://github.com/bytecodealliance/wasm-micro-runtime/issues/1302

core/iwasm/fast-jit/fe/jit_emit_function.c

@@ -409,7 +409,7 @@ jit_compile_op_call(JitCompContext *cc, uint32 func_idx, bool tail_call)
 
         res = create_first_res_reg(cc, func_type);
 
-        GEN_INSN(CALLBC, res, 0, jitted_code);
+        GEN_INSN(CALLBC, res, 0, jitted_code, NEW_CONST(I32, func_idx));
 
         if (!post_return(cc, func_type, res, true)) {
             goto fail;
@@ -700,7 +700,7 @@ jit_compile_op_call_indirect(JitCompContext *cc, uint32 type_idx,
                 goto fail;
         }
     }
-    GEN_INSN(CALLBC, res, 0, jitted_code);
+    GEN_INSN(CALLBC, res, 0, jitted_code, func_idx);
     /* Store res into current frame, so that post_return in
         block func_return can get the value */
     n = cc->jit_frame->sp - cc->jit_frame->lp;

core/iwasm/fast-jit/jit_codecache.c

@@ -58,8 +58,7 @@ jit_pass_register_jitted_code(JitCompContext *cc)
 {
     uint32 jit_func_idx =
         cc->cur_wasm_func_idx - cc->cur_wasm_module->import_function_count;
-    cc->cur_wasm_func->fast_jit_jitted_code = cc->jitted_addr_begin;
     cc->cur_wasm_module->fast_jit_func_ptrs[jit_func_idx] =
-        cc->jitted_addr_begin;
+        cc->cur_wasm_func->fast_jit_jitted_code = cc->jitted_addr_begin;
     return true;
 }

core/iwasm/fast-jit/jit_codegen.h

@@ -65,6 +65,14 @@ jit_codegen_gen_native(JitCompContext *cc);
 bool
 jit_codegen_lower(JitCompContext *cc);
 
+#if WASM_ENABLE_LAZY_JIT != 0 && WASM_ENABLE_JIT != 0
+void *
+jit_codegen_compile_call_to_llvm_jit(const WASMType *func_type);
+
+void *
+jit_codegen_compile_call_to_fast_jit(const WASMModule *module, uint32 func_idx);
+#endif
+
 /**
  * Dump native code in the given range to assembly.
  *
@@ -75,7 +83,8 @@ void
 jit_codegen_dump_native(void *begin_addr, void *end_addr);
 
 int
-jit_codegen_interp_jitted_glue(void *self, JitInterpSwitchInfo *info, void *pc);
+jit_codegen_interp_jitted_glue(void *self, JitInterpSwitchInfo *info,
+                               uint32 func_idx, void *pc);
 
 #ifdef __cplusplus
 }

core/iwasm/fast-jit/jit_compiler.c

@@ -10,9 +10,9 @@
 #include "../interpreter/wasm.h"
 
 typedef struct JitCompilerPass {
-    /* Name of the pass.  */
+    /* Name of the pass */
     const char *name;
-    /* The entry of the compiler pass.  */
+    /* The entry of the compiler pass */
     bool (*run)(JitCompContext *cc);
 } JitCompilerPass;
 
@@ -30,7 +30,7 @@ static JitCompilerPass compiler_passes[] = {
 #undef REG_PASS
 };
 
-/* Number of compiler passes.  */
+/* Number of compiler passes */
 #define COMPILER_PASS_NUM (sizeof(compiler_passes) / sizeof(compiler_passes[0]))
 
 #if WASM_ENABLE_FAST_JIT_DUMP == 0
@@ -43,14 +43,17 @@ static const uint8 compiler_passes_with_dump[] = {
 };
 #endif
 
-/* The exported global data of JIT compiler.  */
+/* The exported global data of JIT compiler */
 static JitGlobals jit_globals = {
 #if WASM_ENABLE_FAST_JIT_DUMP == 0
     .passes = compiler_passes_without_dump,
 #else
     .passes = compiler_passes_with_dump,
 #endif
-    .return_to_interp_from_jitted = NULL
+    .return_to_interp_from_jitted = NULL,
+#if WASM_ENABLE_LAZY_JIT != 0
+    .compile_fast_jit_and_then_call = NULL,
+#endif
 };
 /* clang-format on */
 
@@ -60,7 +63,7 @@ apply_compiler_passes(JitCompContext *cc)
     const uint8 *p = jit_globals.passes;
 
     for (; *p; p++) {
-        /* Set the pass NO.  */
+        /* Set the pass NO */
         cc->cur_pass_no = p - jit_globals.passes;
         bh_assert(*p < COMPILER_PASS_NUM);
 
@@ -120,37 +123,53 @@ jit_compiler_get_pass_name(unsigned i)
 bool
 jit_compiler_compile(WASMModule *module, uint32 func_idx)
 {
-    JitCompContext *cc;
+    JitCompContext *cc = NULL;
     char *last_error;
-    bool ret = true;
+    bool ret = false;
+    uint32 i = func_idx - module->import_function_count;
+    uint32 j = i % WASM_ORC_JIT_BACKEND_THREAD_NUM;
 
-    /* Initialize compilation context.  */
-    if (!(cc = jit_calloc(sizeof(*cc))))
-        return false;
+    /* Lock to avoid duplicated compilation by other threads */
+    os_mutex_lock(&module->fast_jit_thread_locks[j]);
+
+    if (jit_compiler_is_compiled(module, func_idx)) {
+        /* Function has been compiled */
+        os_mutex_unlock(&module->fast_jit_thread_locks[j]);
+        return true;
+    }
+
+    /* Initialize the compilation context */
+    if (!(cc = jit_calloc(sizeof(*cc)))) {
+        goto fail;
+    }
 
     if (!jit_cc_init(cc, 64)) {
-        jit_free(cc);
-        return false;
+        goto fail;
     }
 
     cc->cur_wasm_module = module;
-    cc->cur_wasm_func =
-        module->functions[func_idx - module->import_function_count];
+    cc->cur_wasm_func = module->functions[i];
     cc->cur_wasm_func_idx = func_idx;
     cc->mem_space_unchanged = (!cc->cur_wasm_func->has_op_memory_grow
                                && !cc->cur_wasm_func->has_op_func_call)
                               || (!module->possible_memory_grow);
 
-    /* Apply compiler passes.  */
+    /* Apply compiler passes */
     if (!apply_compiler_passes(cc) || jit_get_last_error(cc)) {
         last_error = jit_get_last_error(cc);
         os_printf("fast jit compilation failed: %s\n",
                   last_error ? last_error : "unknown error");
-        ret = false;
+        goto fail;
     }
 
-    /* Delete the compilation context.  */
-    jit_cc_delete(cc);
+    ret = true;
+
+fail:
+    /* Destroy the compilation context */
+    if (cc)
+        jit_cc_delete(cc);
+
+    os_mutex_unlock(&module->fast_jit_thread_locks[j]);
 
     return ret;
 }
@@ -169,8 +188,92 @@ jit_compiler_compile_all(WASMModule *module)
     return true;
 }
 
-int
-jit_interp_switch_to_jitted(void *exec_env, JitInterpSwitchInfo *info, void *pc)
+bool
+jit_compiler_is_compiled(const WASMModule *module, uint32 func_idx)
 {
-    return jit_codegen_interp_jitted_glue(exec_env, info, pc);
+    uint32 i = func_idx - module->import_function_count;
+
+    bh_assert(func_idx >= module->import_function_count
+              && func_idx
+                     < module->import_function_count + module->function_count);
+
+#if WASM_ENABLE_LAZY_JIT == 0
+    return module->fast_jit_func_ptrs[i] ? true : false;
+#else
+    return module->fast_jit_func_ptrs[i]
+                   != jit_globals.compile_fast_jit_and_then_call
+               ? true
+               : false;
+#endif
+}
+
+#if WASM_ENABLE_LAZY_JIT != 0 && WASM_ENABLE_JIT != 0
+bool
+jit_compiler_set_call_to_llvm_jit(WASMModule *module, uint32 func_idx)
+{
+    uint32 i = func_idx - module->import_function_count;
+    uint32 j = i % WASM_ORC_JIT_BACKEND_THREAD_NUM;
+    WASMType *func_type = module->functions[i]->func_type;
+    uint32 k =
+        ((uint32)(uintptr_t)func_type >> 3) % WASM_ORC_JIT_BACKEND_THREAD_NUM;
+    void *func_ptr = NULL;
+
+    /* Compile code block of call_to_llvm_jit_from_fast_jit of
+       this kind of function type if it hasn't been compiled */
+    if (!(func_ptr = func_type->call_to_llvm_jit_from_fast_jit)) {
+        os_mutex_lock(&module->fast_jit_thread_locks[k]);
+        if (!(func_ptr = func_type->call_to_llvm_jit_from_fast_jit)) {
+            if (!(func_ptr = func_type->call_to_llvm_jit_from_fast_jit =
+                      jit_codegen_compile_call_to_llvm_jit(func_type))) {
+                os_mutex_unlock(&module->fast_jit_thread_locks[k]);
+                return false;
+            }
+        }
+        os_mutex_unlock(&module->fast_jit_thread_locks[k]);
+    }
+
+    /* Switch current fast jit func ptr to the code block */
+    os_mutex_lock(&module->fast_jit_thread_locks[j]);
+    module->fast_jit_func_ptrs[i] = func_ptr;
+    os_mutex_unlock(&module->fast_jit_thread_locks[j]);
+    return true;
+}
+
+bool
+jit_compiler_set_call_to_fast_jit(WASMModule *module, uint32 func_idx)
+{
+    void *func_ptr = NULL;
+
+    func_ptr = jit_codegen_compile_call_to_fast_jit(module, func_idx);
+    if (func_ptr) {
+        jit_compiler_set_llvm_jit_func_ptr(module, func_idx, func_ptr);
+    }
+
+    return func_ptr ? true : false;
+}
+
+void
+jit_compiler_set_llvm_jit_func_ptr(WASMModule *module, uint32 func_idx,
+                                   void *func_ptr)
+{
+    WASMModuleInstance *instance;
+    uint32 i = func_idx - module->import_function_count;
+
+    module->functions[i]->llvm_jit_func_ptr = module->func_ptrs[i] = func_ptr;
+
+    os_mutex_lock(&module->instance_list_lock);
+    instance = module->instance_list;
+    while (instance) {
+        instance->func_ptrs[func_idx] = func_ptr;
+        instance = instance->e->next;
+    }
+    os_mutex_unlock(&module->instance_list_lock);
+}
+#endif /* end of WASM_ENABLE_LAZY_JIT != 0 && WASM_ENABLE_JIT != 0 */
+
+int
+jit_interp_switch_to_jitted(void *exec_env, JitInterpSwitchInfo *info,
+                            uint32 func_idx, void *pc)
+{
+    return jit_codegen_interp_jitted_glue(exec_env, info, func_idx, pc);
 }

core/iwasm/fast-jit/jit_compiler.h

@@ -18,6 +18,9 @@ typedef struct JitGlobals {
     /* Compiler pass sequence, the last element must be 0 */
     const uint8 *passes;
     char *return_to_interp_from_jitted;
+#if WASM_ENABLE_LAZY_JIT != 0
+    char *compile_fast_jit_and_then_call;
+#endif
 } JitGlobals;
 
 /**
@@ -87,8 +90,24 @@ jit_compiler_compile(WASMModule *module, uint32 func_idx);
 bool
 jit_compiler_compile_all(WASMModule *module);
 
+bool
+jit_compiler_is_compiled(const WASMModule *module, uint32 func_idx);
+
+#if WASM_ENABLE_LAZY_JIT != 0 && WASM_ENABLE_JIT != 0
+bool
+jit_compiler_set_call_to_llvm_jit(WASMModule *module, uint32 func_idx);
+
+bool
+jit_compiler_set_call_to_fast_jit(WASMModule *module, uint32 func_idx);
+
+void
+jit_compiler_set_llvm_jit_func_ptr(WASMModule *module, uint32 func_idx,
+                                   void *func_ptr);
+#endif
+
 int
-jit_interp_switch_to_jitted(void *self, JitInterpSwitchInfo *info, void *pc);
+jit_interp_switch_to_jitted(void *self, JitInterpSwitchInfo *info,
+                            uint32 func_idx, void *pc);
 
 /*
  * Pass declarations:

core/iwasm/fast-jit/jit_frontend.c

@@ -2255,6 +2255,12 @@ jit_frontend_translate_func(JitCompContext *cc)
     return basic_block_entry;
 }
 
+uint32
+jit_frontend_get_jitted_return_addr_offset()
+{
+    return (uint32)offsetof(WASMInterpFrame, jitted_return_addr);
+}
+
 #if 0
 #if WASM_ENABLE_THREAD_MGR != 0
 bool

core/iwasm/fast-jit/jit_frontend.h

@@ -13,6 +13,10 @@
 #include "../aot/aot_runtime.h"
 #endif
 
+#ifdef __cplusplus
+extern "C" {
+#endif
+
 #if WASM_ENABLE_AOT == 0
 typedef enum IntCond {
     INT_EQZ = 0,
@@ -143,6 +147,9 @@ jit_frontend_translate_func(JitCompContext *cc);
 bool
 jit_frontend_lower(JitCompContext *cc);
 
+uint32
+jit_frontend_get_jitted_return_addr_offset();
+
 uint32
 jit_frontend_get_global_data_offset(const WASMModule *module,
                                     uint32 global_idx);
@@ -483,4 +490,8 @@ set_local_f64(JitFrame *frame, int n, JitReg val)
 #define PUSH_FUNCREF(v) PUSH(v, VALUE_TYPE_FUNCREF)
 #define PUSH_EXTERNREF(v) PUSH(v, VALUE_TYPE_EXTERNREF)
 
+#ifdef __cplusplus
+}
+#endif
+
 #endif

core/iwasm/fast-jit/jit_ir.def

@@ -196,7 +196,7 @@ INSN(LOOKUPSWITCH, LookupSwitch, 1, 0)
 
 /* Call and return instructions */
 INSN(CALLNATIVE, VReg, 2, 1)
-INSN(CALLBC, Reg, 3, 2)
+INSN(CALLBC, Reg, 4, 2)
 INSN(RETURNBC, Reg, 3, 0)
 INSN(RETURN, Reg, 1, 0)