Refactor Orc JIT to enable lazy compilation (#850)

.github/workflows/compilation_on_android_ubuntu_macos.yml

@@ -31,8 +31,8 @@ env:
   AOT_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=1 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=0 -DWAMR_BUILD_JIT=0 -DWAMR_BUILD_LAZY_JIT=0"
   CLASSIC_INTERP_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=0 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=1 -DWAMR_BUILD_JIT=0 -DWAMR_BUILD_LAZY_JIT=0"
   FAST_INTERP_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=0 -DWAMR_BUILD_FAST_INTERP=1 -DWAMR_BUILD_INTERP=1 -DWAMR_BUILD_JIT=0 -DWAMR_BUILD_LAZY_JIT=0"
-  JIT_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=1 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=0 -DWAMR_BUILD_JIT=1 -DWAMR_BUILD_LAZY_JIT=0"
   LAZY_JIT_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=1 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=0 -DWAMR_BUILD_JIT=1 -DWAMR_BUILD_LAZY_JIT=1"
+  MC_JIT_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=1 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=0 -DWAMR_BUILD_JIT=1 -DWAMR_BUILD_LAZY_JIT=0"
   LLVM_CACHE_SUFFIX: "build-llvm_libraries_ex"
 
 jobs:
@@ -193,8 +193,8 @@ jobs:
             # Running mode
             $CLASSIC_INTERP_BUILD_OPTIONS,
             $FAST_INTERP_BUILD_OPTIONS,
-            $JIT_BUILD_OPTIONS,
             $LAZY_JIT_BUILD_OPTIONS,
+            $MC_JIT_BUILD_OPTIONS,
             $AOT_BUILD_OPTIONS,
           ]
         make_options_feature: [
@@ -234,11 +234,11 @@ jobs:
             make_options_feature: "-DWAMR_BUILD_DEBUG_INTERP=1"
           # uncompatiable mode and feature
           # MULTI_MODULE only on INTERP mode
-          - make_options_run_mode: $JIT_BUILD_OPTIONS
+          - make_options_run_mode: $LAZY_JIT_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_MULTI_MODULE=1"
           - make_options_run_mode: $AOT_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_MULTI_MODULE=1"
-          - make_options_run_mode: $LAZY_JIT_BUILD_OPTIONS
+          - make_options_run_mode: $MC_JIT_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_MULTI_MODULE=1"
           # SIMD only on JIT/AOT mode
           - make_options_run_mode: $CLASSIC_INTERP_BUILD_OPTIONS
@@ -248,10 +248,10 @@ jobs:
           # DEBUG_INTERP only on CLASSIC INTERP mode
           - make_options_run_mode: $AOT_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_DEBUG_INTERP=1"
-          - make_options_run_mode: $JIT_BUILD_OPTIONS
-            make_options_feature: "-DWAMR_BUILD_DEBUG_INTERP=1"
           - make_options_run_mode: $LAZY_JIT_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_DEBUG_INTERP=1"
+          - make_options_run_mode: $MC_JIT_BUILD_OPTIONS
+            make_options_feature: "-DWAMR_BUILD_DEBUG_INTERP=1"
           - make_options_run_mode: $FAST_INTERP_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_DEBUG_INTERP=1"
           # DEBUG_AOT only on JIT/AOT mode
@@ -260,17 +260,17 @@ jobs:
           - make_options_run_mode: $FAST_INTERP_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_DEBUG_AOT=1"
           # TODO: DEBUG_AOT on JIT
-          - make_options_run_mode: $JIT_BUILD_OPTIONS
-            make_options_feature: "-DWAMR_BUILD_DEBUG_AOT=1"
           - make_options_run_mode: $LAZY_JIT_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_DEBUG_AOT=1"
+          - make_options_run_mode: $MC_JIT_BUILD_OPTIONS
+            make_options_feature: "-DWAMR_BUILD_DEBUG_AOT=1"
           # MINI_LOADER only on INTERP mode
           - make_options_run_mode: $AOT_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_MINI_LOADER=1"
-          - make_options_run_mode: $JIT_BUILD_OPTIONS
-            make_options_feature: "-DWAMR_BUILD_MINI_LOADER=1"
           - make_options_run_mode: $LAZY_JIT_BUILD_OPTIONS
             make_options_feature: "-DWAMR_BUILD_MINI_LOADER=1"
+          - make_options_run_mode: $MC_JIT_BUILD_OPTIONS
+            make_options_feature: "-DWAMR_BUILD_MINI_LOADER=1"
         include:
           - os: ubuntu-18.04
             light: ${{ needs.check_repo.outputs.traffic_light_on_ubuntu_1804 }}
@@ -320,8 +320,8 @@ jobs:
             # Running mode
             $CLASSIC_INTERP_BUILD_OPTIONS,
             $FAST_INTERP_BUILD_OPTIONS,
-            $JIT_BUILD_OPTIONS,
             $LAZY_JIT_BUILD_OPTIONS,
+            $MC_JIT_BUILD_OPTIONS,
             $AOT_BUILD_OPTIONS,
           ]
         os: [ubuntu-18.04, ubuntu-20.04, macos-latest]
@@ -341,11 +341,11 @@ jobs:
         exclude:
           # TODO: a .aot compatiable problem
           - os: macos-latest
-            make_options: $JIT_BUILD_OPTIONS
+            make_options: $LAZY_JIT_BUILD_OPTIONS
           - os: macos-latest
             make_options: $AOT_BUILD_OPTIONS
           - os: macos-latest
-            make_options: $LAZY_JIT_BUILD_OPTIONS
+            make_options: $MC_JIT_BUILD_OPTIONS
     steps:
       - name: light status
         run: echo "matrix.os=${{ matrix.os }}, light=${{ matrix.light }}"

.github/workflows/compilation_on_sgx.yml

@@ -31,8 +31,8 @@ env:
   AOT_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=1 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=0 -DWAMR_BUILD_JIT=0 -DWAMR_BUILD_LAZY_JIT=0"
   CLASSIC_INTERP_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=0 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=1 -DWAMR_BUILD_JIT=0 -DWAMR_BUILD_LAZY_JIT=0"
   FAST_INTERP_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=0 -DWAMR_BUILD_FAST_INTERP=1 -DWAMR_BUILD_INTERP=1 -DWAMR_BUILD_JIT=0 -DWAMR_BUILD_LAZY_JIT=0"
-  JIT_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=1 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=0 -DWAMR_BUILD_JIT=1 -DWAMR_BUILD_LAZY_JIT=0"
   LAZY_JIT_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=1 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=0 -DWAMR_BUILD_JIT=1 -DWAMR_BUILD_LAZY_JIT=1"
+  MC_JIT_BUILD_OPTIONS: "-DWAMR_BUILD_AOT=1 -DWAMR_BUILD_FAST_INTERP=0 -DWAMR_BUILD_INTERP=0 -DWAMR_BUILD_JIT=1 -DWAMR_BUILD_LAZY_JIT=0"
   LLVM_CACHE_SUFFIX: "build-llvm_libraries_ex"
 
 jobs:
@@ -116,8 +116,8 @@ jobs:
             $CLASSIC_INTERP_BUILD_OPTIONS,
             $FAST_INTERP_BUILD_OPTIONS,
             # doesn't support
-            # $JIT_BUILD_OPTIONS,
             # $LAZY_JIT_BUILD_OPTIONS,
+            # $MC_JIT_BUILD_OPTIONS,
             $AOT_BUILD_OPTIONS,
           ]
         make_options_feature: [
@@ -208,8 +208,8 @@ jobs:
             $CLASSIC_INTERP_BUILD_OPTIONS,
             $FAST_INTERP_BUILD_OPTIONS,
             # doesn't support
-            #$JIT_BUILD_OPTIONS,
             #$LAZY_JIT_BUILD_OPTIONS,
+            #$MC_JIT_BUILD_OPTIONS,
             #$AOT_BUILD_OPTIONS,
           ]
         os: [ubuntu-20.04]

build-scripts/config_common.cmake

@@ -83,7 +83,9 @@ endif ()
 if (WAMR_BUILD_JIT EQUAL 1)
   if (WAMR_BUILD_AOT EQUAL 1)
     add_definitions("-DWASM_ENABLE_JIT=1")
-    if (WAMR_BUILD_LAZY_JIT EQUAL 1)
+    if (NOT WAMR_BUILD_LAZY_JIT EQUAL 0)
+      # Enable Lazy JIT by default
+      set (WAMR_BUILD_LAZY_JIT 1)
       add_definitions("-DWASM_ENABLE_LAZY_JIT=1")
     endif ()
     if (NOT DEFINED LLVM_DIR)
@@ -124,11 +126,10 @@ else ()
   message ("     WAMR AOT disabled")
 endif ()
 if (WAMR_BUILD_JIT EQUAL 1)
-  message ("     WAMR JIT enabled")
   if (WAMR_BUILD_LAZY_JIT EQUAL 1)
-    message ("     WAMR LazyJIT enabled")
+    message ("     WAMR Lazy JIT enabled")
   else ()
-    message ("     WAMR LazyJIT disabled")
+    message ("     WAMR MC JIT enabled")
   endif ()
 else ()
   message ("     WAMR JIT disabled")

core/config.h

@@ -81,6 +81,10 @@
 #define WASM_ENABLE_LAZY_JIT 0
 #endif
 
+#ifndef WASM_LAZY_JIT_COMPILE_THREAD_NUM
+#define WASM_LAZY_JIT_COMPILE_THREAD_NUM 4
+#endif
+
 #if (WASM_ENABLE_AOT == 0) && (WASM_ENABLE_JIT != 0)
 /* LazyJIT or MCJIT can only be enabled when AOT is enabled */
 #undef WASM_ENABLE_JIT

core/iwasm/aot/aot_loader.c

@@ -2699,6 +2699,89 @@ aot_load_from_aot_file(const uint8 *buf, uint32 size, char *error_buf,
 }
 
 #if WASM_ENABLE_JIT != 0
+#if WASM_ENABLE_LAZY_JIT != 0
+/* Orc JIT thread arguments */
+typedef struct OrcJitThreadArg {
+    AOTCompData *comp_data;
+    AOTCompContext *comp_ctx;
+    AOTModule *module;
+    int32 group_idx;
+    int32 group_stride;
+} OrcJitThreadArg;
+
+static bool orcjit_stop_compiling = false;
+static korp_tid orcjit_threads[WASM_LAZY_JIT_COMPILE_THREAD_NUM];
+static OrcJitThreadArg orcjit_thread_args[WASM_LAZY_JIT_COMPILE_THREAD_NUM];
+
+static void *
+orcjit_thread_callback(void *arg)
+{
+    LLVMErrorRef error;
+    LLVMOrcJITTargetAddress func_addr = 0;
+    OrcJitThreadArg *thread_arg = (OrcJitThreadArg *)arg;
+    AOTCompData *comp_data = thread_arg->comp_data;
+    AOTCompContext *comp_ctx = thread_arg->comp_ctx;
+    AOTModule *module = thread_arg->module;
+    char func_name[32];
+    int32 i;
+
+    /* Compile wasm functions of this group */
+    for (i = thread_arg->group_idx; i < (int32)comp_data->func_count;
+         i += thread_arg->group_stride) {
+        if (!module->func_ptrs[i]) {
+            snprintf(func_name, sizeof(func_name), "%s%d", AOT_FUNC_PREFIX, i);
+            if ((error = LLVMOrcLLJITLookup(comp_ctx->orc_lazyjit, &func_addr,
+                                            func_name))) {
+                char *err_msg = LLVMGetErrorMessage(error);
+                os_printf("failed to compile orc jit function: %s", err_msg);
+                LLVMDisposeErrorMessage(err_msg);
+                break;
+            }
+            /**
+             * No need to lock the func_ptr[func_idx] here as it is basic
+             * data type, the load/store for it can be finished by one cpu
+             * instruction, and there can be only one cpu instruction
+             * loading/storing at the same time.
+             */
+            module->func_ptrs[i] = (void *)func_addr;
+        }
+        if (orcjit_stop_compiling) {
+            break;
+        }
+    }
+
+    /* Try to compile functions that haven't been compiled by other threads */
+    for (i = (int32)comp_data->func_count - 1; i > 0; i--) {
+        if (orcjit_stop_compiling) {
+            break;
+        }
+        if (!module->func_ptrs[i]) {
+            snprintf(func_name, sizeof(func_name), "%s%d", AOT_FUNC_PREFIX, i);
+            if ((error = LLVMOrcLLJITLookup(comp_ctx->orc_lazyjit, &func_addr,
+                                            func_name))) {
+                char *err_msg = LLVMGetErrorMessage(error);
+                os_printf("failed to compile orc jit function: %s", err_msg);
+                LLVMDisposeErrorMessage(err_msg);
+                break;
+            }
+            module->func_ptrs[i] = (void *)func_addr;
+        }
+    }
+
+    return NULL;
+}
+
+static void
+orcjit_stop_compile_threads()
+{
+    uint32 i;
+    orcjit_stop_compiling = true;
+    for (i = 0; i < WASM_LAZY_JIT_COMPILE_THREAD_NUM; i++) {
+        os_thread_join(orcjit_threads[i], NULL);
+    }
+}
+#endif
+
 static AOTModule *
 aot_load_from_comp_data(AOTCompData *comp_data, AOTCompContext *comp_ctx,
                         char *error_buf, uint32 error_buf_size)
@@ -2708,13 +2791,6 @@ aot_load_from_comp_data(AOTCompData *comp_data, AOTCompContext *comp_ctx,
     char func_name[32];
     AOTModule *module;
 
-#if WASM_ENABLE_LAZY_JIT != 0
-    LLVMOrcThreadSafeModuleRef ts_module;
-    LLVMOrcJITDylibRef main_dylib;
-    LLVMErrorRef error;
-    LLVMOrcJITTargetAddress func_addr = 0;
-#endif
-
     /* Allocate memory for module */
     if (!(module =
               loader_malloc(sizeof(AOTModule), error_buf, error_buf_size))) {
@@ -2778,44 +2854,28 @@ aot_load_from_comp_data(AOTCompData *comp_data, AOTCompContext *comp_ctx,
     }
 
 #if WASM_ENABLE_LAZY_JIT != 0
-    bh_assert(comp_ctx->lazy_orcjit);
-
-    main_dylib = LLVMOrcLLLazyJITGetMainJITDylib(comp_ctx->lazy_orcjit);
-    if (!main_dylib) {
-        set_error_buf(error_buf, error_buf_size,
-                      "failed to get dynmaic library reference");
-        goto fail3;
-    }
-
-    ts_module = LLVMOrcCreateNewThreadSafeModule(comp_ctx->module,
-                                                 comp_ctx->ts_context);
-    if (!ts_module) {
-        set_error_buf(error_buf, error_buf_size,
-                      "failed to create thread safe module");
-        goto fail3;
-    }
-
-    if ((error = LLVMOrcLLLazyJITAddLLVMIRModule(comp_ctx->lazy_orcjit,
-                                                 main_dylib, ts_module))) {
-        /*
-         * If adding the ThreadSafeModule fails then we need to clean it up
-         * ourselves. If adding it succeeds the JIT will manage the memory.
-         */
-        aot_handle_llvm_errmsg(error_buf, error_buf_size,
-                               "failed to addIRModule: ", error);
-        goto fail4;
-    }
+    /* Create threads to compile the wasm functions */
+    for (i = 0; i < WASM_LAZY_JIT_COMPILE_THREAD_NUM; i++) {
+        orcjit_thread_args[i].comp_data = comp_data;
+        orcjit_thread_args[i].comp_ctx = comp_ctx;
+        orcjit_thread_args[i].module = module;
+        orcjit_thread_args[i].group_idx = (int32)i;
+        orcjit_thread_args[i].group_stride = WASM_LAZY_JIT_COMPILE_THREAD_NUM;
+        if (os_thread_create(&orcjit_threads[i], orcjit_thread_callback,
+                             (void *)&orcjit_thread_args[i],
+                             APP_THREAD_STACK_SIZE_DEFAULT)
+            != 0) {
+            uint32 j;
 
-    for (i = 0; i < comp_data->func_count; i++) {
-        snprintf(func_name, sizeof(func_name), "%s%d", AOT_FUNC_PREFIX, i);
-        if ((error = LLVMOrcLLLazyJITLookup(comp_ctx->lazy_orcjit, &func_addr,
-                                            func_name))) {
-            aot_handle_llvm_errmsg(error_buf, error_buf_size,
-                                   "cannot lookup: ", error);
+            set_error_buf(error_buf, error_buf_size,
+                          "create orcjit compile thread failed");
+            /* Terminate the threads created */
+            orcjit_stop_compiling = true;
+            for (j = 0; j < i; j++) {
+                os_thread_join(orcjit_threads[j], NULL);
+            }
             goto fail3;
         }
-        module->func_ptrs[i] = (void *)func_addr;
-        func_addr = 0;
     }
 #else
     /* Resolve function addresses */
@@ -2850,6 +2910,29 @@ aot_load_from_comp_data(AOTCompData *comp_data, AOTCompContext *comp_ctx,
                   < module->import_func_count + module->func_count);
         /* TODO: fix issue that start func cannot be import func */
         if (comp_data->start_func_index >= module->import_func_count) {
+#if WASM_ENABLE_LAZY_JIT != 0
+            if (!module->func_ptrs[comp_data->start_func_index
+                                   - module->import_func_count]) {
+                LLVMErrorRef error;
+                LLVMOrcJITTargetAddress func_addr = 0;
+
+                snprintf(func_name, sizeof(func_name), "%s%d", AOT_FUNC_PREFIX,
+                         comp_data->start_func_index
+                             - module->import_func_count);
+                if ((error = LLVMOrcLLJITLookup(comp_ctx->orc_lazyjit,
+                                                &func_addr, func_name))) {
+                    char *err_msg = LLVMGetErrorMessage(error);
+                    set_error_buf_v(error_buf, error_buf_size,
+                                    "failed to compile orc jit function: %s",
+                                    err_msg);
+                    LLVMDisposeErrorMessage(err_msg);
+                    goto fail4;
+                }
+                module->func_ptrs[comp_data->start_func_index
+                                  - module->import_func_count] =
+                    (void *)func_addr;
+            }
+#endif
             module->start_function =
                 module->func_ptrs[comp_data->start_func_index
                                   - module->import_func_count];
@@ -2885,10 +2968,15 @@ aot_load_from_comp_data(AOTCompData *comp_data, AOTCompContext *comp_ctx,
 
 #if WASM_ENABLE_LAZY_JIT != 0
 fail4:
-    LLVMOrcDisposeThreadSafeModule(ts_module);
+    if (module->func_type_indexes)
+        wasm_runtime_free(module->func_type_indexes);
 #endif
 
 fail3:
+#if WASM_ENABLE_LAZY_JIT != 0
+    /* Terminate all threads before free module->func_ptrs */
+    orcjit_stop_compile_threads();
+#endif
     if (module->func_ptrs)
         wasm_runtime_free(module->func_ptrs);
 fail2:
@@ -2973,6 +3061,10 @@ void
 aot_unload(AOTModule *module)
 {
 #if WASM_ENABLE_JIT != 0
+#if WASM_ENABLE_LAZY_JIT != 0
+    orcjit_stop_compile_threads();
+#endif
+
     if (module->comp_data)
         aot_destroy_comp_data(module->comp_data);