Implement shared heap for AOT (#3815)

Author: wenyongh

.github/workflows/compilation_on_android_ubuntu.yml

@@ -585,6 +585,7 @@ jobs:
           cmake ..
           cmake --build . --config Debug --parallel 4
           ./shared_heap_test
+          ./shared_heap_test --aot
 
   test:
     needs:

.github/workflows/compilation_on_macos.yml

@@ -394,3 +394,4 @@ jobs:
           cmake ..
           cmake --build . --config Debug --parallel 4
           ./shared_heap_test
+          ./shared_heap_test --aot

.github/workflows/nightly_run.yml

@@ -600,6 +600,7 @@ jobs:
           cmake ..
           cmake --build . --config Debug --parallel 4
           ./shared_heap_test
+          ./shared_heap_test --aot
 
   test:
     needs:

core/config.h

@@ -396,7 +396,9 @@
 #define APP_HEAP_SIZE_DEFAULT (8 * 1024)
 #endif
 #define APP_HEAP_SIZE_MIN (256)
-#define APP_HEAP_SIZE_MAX (512 * 1024 * 1024)
+/* The ems memory allocator supports maximal heap size 1GB,
+   see ems_gc_internal.h */
+#define APP_HEAP_SIZE_MAX (1024 * 1024 * 1024)
 
 /* Default min/max gc heap size of each app */
 #ifndef GC_HEAP_SIZE_DEFAULT

core/iwasm/aot/aot_runtime.c

@@ -57,6 +57,9 @@ bh_static_assert(sizeof(AOTMemoryInstance) == 120);
 bh_static_assert(offsetof(AOTTableInstance, elems) == 24);
 
 bh_static_assert(offsetof(AOTModuleInstanceExtra, stack_sizes) == 0);
+bh_static_assert(offsetof(AOTModuleInstanceExtra, shared_heap_base_addr_adj)
+                 == 8);
+bh_static_assert(offsetof(AOTModuleInstanceExtra, shared_heap_start_off) == 16);
 
 bh_static_assert(sizeof(CApiFuncImport) == sizeof(uintptr_t) * 3);
 
@@ -1885,6 +1888,24 @@ aot_instantiate(AOTModule *module, AOTModuleInstance *parent,
     extra->stack_sizes =
         aot_get_data_section_addr(module, AOT_STACK_SIZES_SECTION_NAME, NULL);
 
+    /*
+     * The AOT code checks whether the n bytes to access are in shared heap
+     * by checking whether the beginning address meets:
+     *   addr >= start_off && addr <= end_off - n-bytes + 1
+     * where n is 1/2/4/8/16 and `end_off - n-bytes + 1` is constant, e.g.,
+     *   UINT32_MAX, UINT32_MAX-1, UINT32_MAX-3 for n = 1, 2 or 4 in 32-bit
+     * target. To simplify the check, when shared heap is disabled, we set
+     * the start off to UINT64_MAX in 64-bit target and UINT32_MAX in 32-bit
+     * target, so in the checking, the above formula will be false, we don't
+     * need to check whether the shared heap is enabled or not in the AOT
+     * code.
+     */
+#if UINTPTR_MAX == UINT64_MAX
+    extra->shared_heap_start_off.u64 = UINT64_MAX;
+#else
+    extra->shared_heap_start_off.u32[0] = UINT32_MAX;
+#endif
+
 #if WASM_ENABLE_PERF_PROFILING != 0
     total_size = sizeof(AOTFuncPerfProfInfo)
                  * ((uint64)module->import_func_count + module->func_count);

core/iwasm/aot/aot_runtime.h

@@ -111,6 +111,14 @@ typedef struct AOTFunctionInstance {
 
 typedef struct AOTModuleInstanceExtra {
     DefPointer(const uint32 *, stack_sizes);
+    /*
+     * Adjusted shared heap based addr to simple the calculation
+     * in the aot code. The value is:
+     *   shared_heap->base_addr - shared_heap->start_off
+     */
+    DefPointer(uint8 *, shared_heap_base_addr_adj);
+    MemBound shared_heap_start_off;
+
     WASMModuleInstanceExtraCommon common;
     AOTFunctionInstance **functions;
     uint32 function_count;

core/iwasm/common/wasm_memory.c

@@ -165,7 +165,7 @@ runtime_malloc(uint64 size)
 WASMSharedHeap *
 wasm_runtime_create_shared_heap(SharedHeapInitArgs *init_args)
 {
-    uint64 heap_struct_size = sizeof(WASMSharedHeap);
+    uint64 heap_struct_size = sizeof(WASMSharedHeap), map_size;
     uint32 size = init_args->size;
     WASMSharedHeap *heap;
 
@@ -192,7 +192,18 @@ wasm_runtime_create_shared_heap(SharedHeapInitArgs *init_args)
         goto fail3;
     }
 
-    if (!(heap->base_addr = wasm_mmap_linear_memory(size, size))) {
+#ifndef OS_ENABLE_HW_BOUND_CHECK
+    map_size = size;
+#else
+    /* Totally 8G is mapped, the opcode load/store address range is 0 to 8G:
+     *   ea = i + memarg.offset
+     * both i and memarg.offset are u32 in range 0 to 4G
+     * so the range of ea is 0 to 8G
+     */
+    map_size = 8 * (uint64)BH_GB;
+#endif
+
+    if (!(heap->base_addr = wasm_mmap_linear_memory(map_size, size))) {
         goto fail3;
     }
     if (!mem_allocator_create_with_struct_and_pool(
@@ -213,7 +224,7 @@ wasm_runtime_create_shared_heap(SharedHeapInitArgs *init_args)
     return heap;
 
 fail4:
-    wasm_munmap_linear_memory(heap->base_addr, size, size);
+    wasm_munmap_linear_memory(heap->base_addr, size, map_size);
 fail3:
     wasm_runtime_free(heap->heap_handle);
 fail2:
@@ -245,18 +256,52 @@ wasm_runtime_attach_shared_heap_internal(WASMModuleInstanceCommon *module_inst,
 
 #if WASM_ENABLE_INTERP != 0
     if (module_inst->module_type == Wasm_Module_Bytecode) {
-        if (((WASMModuleInstance *)module_inst)->e->shared_heap) {
+        WASMModuleInstanceExtra *e =
+            (WASMModuleInstanceExtra *)((WASMModuleInstance *)module_inst)->e;
+        if (e->shared_heap) {
             LOG_WARNING("A shared heap is already attached");
             return false;
         }
-        ((WASMModuleInstance *)module_inst)->e->shared_heap = shared_heap;
-    }
+        e->shared_heap = shared_heap;
+#if WASM_ENABLE_JIT != 0
+#if UINTPTR_MAX == UINT64_MAX
+        if (memory->is_memory64)
+            e->shared_heap_start_off.u64 = shared_heap->start_off_mem64;
+        else
+            e->shared_heap_start_off.u64 = shared_heap->start_off_mem32;
+        e->shared_heap_base_addr_adj =
+            shared_heap->base_addr - e->shared_heap_start_off.u64;
+#else
+        e->shared_heap_start_off.u32[0] = (uint32)shared_heap->start_off_mem32;
+        e->shared_heap_base_addr_adj =
+            shared_heap->base_addr - e->shared_heap_start_off.u32[0];
 #endif
+#endif /* end of WASM_ENABLE_JIT != 0 */
+    }
+#endif /* end of WASM_ENABLE_INTERP != 0 */
 #if WASM_ENABLE_AOT != 0
     if (module_inst->module_type == Wasm_Module_AoT) {
-        // TODO
-    }
+        AOTModuleInstanceExtra *e =
+            (AOTModuleInstanceExtra *)((AOTModuleInstance *)module_inst)->e;
+        if (e->shared_heap) {
+            LOG_WARNING("A shared heap is already attached");
+            return false;
+        }
+        e->shared_heap = shared_heap;
+#if UINTPTR_MAX == UINT64_MAX
+        if (memory->is_memory64)
+            e->shared_heap_start_off.u64 = shared_heap->start_off_mem64;
+        else
+            e->shared_heap_start_off.u64 = shared_heap->start_off_mem32;
+        e->shared_heap_base_addr_adj =
+            shared_heap->base_addr - e->shared_heap_start_off.u64;
+#else
+        e->shared_heap_start_off.u32[0] = (uint32)shared_heap->start_off_mem32;
+        e->shared_heap_base_addr_adj =
+            shared_heap->base_addr - e->shared_heap_start_off.u32[0];
 #endif
+    }
+#endif /* end of WASM_ENABLE_AOT != 0 */
 
     return true;
 }
@@ -277,14 +322,32 @@ wasm_runtime_detach_shared_heap_internal(WASMModuleInstanceCommon *module_inst)
 {
 #if WASM_ENABLE_INTERP != 0
     if (module_inst->module_type == Wasm_Module_Bytecode) {
-        ((WASMModuleInstance *)module_inst)->e->shared_heap = NULL;
-    }
+        WASMModuleInstanceExtra *e =
+            (WASMModuleInstanceExtra *)((WASMModuleInstance *)module_inst)->e;
+        e->shared_heap = NULL;
+#if WASM_ENABLE_JIT != 0
+#if UINTPTR_MAX == UINT64_MAX
+        e->shared_heap_start_off.u64 = UINT64_MAX;
+#else
+        e->shared_heap_start_off.u32[0] = UINT32_MAX;
+#endif
+        e->shared_heap_base_addr_adj = NULL;
 #endif
+    }
+#endif /* end of WASM_ENABLE_INTERP != 0 */
 #if WASM_ENABLE_AOT != 0
     if (module_inst->module_type == Wasm_Module_AoT) {
-        // TODO
-    }
+        AOTModuleInstanceExtra *e =
+            (AOTModuleInstanceExtra *)((AOTModuleInstance *)module_inst)->e;
+        e->shared_heap = NULL;
+#if UINTPTR_MAX == UINT64_MAX
+        e->shared_heap_start_off.u64 = UINT64_MAX;
+#else
+        e->shared_heap_start_off.u32[0] = UINT32_MAX;
 #endif
+        e->shared_heap_base_addr_adj = NULL;
+    }
+#endif /* end of WASM_ENABLE_AOT != 0 */
 }
 
 void
@@ -307,13 +370,21 @@ get_shared_heap(WASMModuleInstanceCommon *module_inst_comm)
 #endif
 #if WASM_ENABLE_AOT != 0
     if (module_inst_comm->module_type == Wasm_Module_AoT) {
-        // TODO
-        return NULL;
+        AOTModuleInstanceExtra *e =
+            (AOTModuleInstanceExtra *)((AOTModuleInstance *)module_inst_comm)
+                ->e;
+        return e->shared_heap;
     }
 #endif
     return NULL;
 }
 
+WASMSharedHeap *
+wasm_runtime_get_shared_heap(WASMModuleInstanceCommon *module_inst_comm)
+{
+    return get_shared_heap(module_inst_comm);
+}
+
 static bool
 is_app_addr_in_shared_heap(WASMModuleInstanceCommon *module_inst,
                            bool is_memory64, uint64 app_offset, uint32 bytes)
@@ -324,6 +395,10 @@ is_app_addr_in_shared_heap(WASMModuleInstanceCommon *module_inst,
         return false;
     }
 
+    if (bytes == 0) {
+        bytes = 1;
+    }
+
     if (!is_memory64) {
         if (app_offset >= heap->start_off_mem32
             && app_offset <= UINT32_MAX - bytes + 1) {
@@ -457,17 +532,23 @@ wasm_runtime_memory_init(mem_alloc_type_t mem_alloc_type,
 
 #if WASM_ENABLE_SHARED_HEAP != 0
 static void
-wasm_runtime_destroy_shared_heaps()
+destroy_shared_heaps()
 {
     WASMSharedHeap *heap = shared_heap_list;
     WASMSharedHeap *cur;
+    uint64 map_size;
 
     while (heap) {
         cur = heap;
         heap = heap->next;
         mem_allocator_destroy(cur->heap_handle);
         wasm_runtime_free(cur->heap_handle);
-        wasm_munmap_linear_memory(cur->base_addr, cur->size, cur->size);
+#ifndef OS_ENABLE_HW_BOUND_CHECK
+        map_size = cur->size;
+#else
+        map_size = 8 * (uint64)BH_GB;
+#endif
+        wasm_munmap_linear_memory(cur->base_addr, cur->size, map_size);
         wasm_runtime_free(cur);
     }
 }
@@ -477,7 +558,7 @@ void
 wasm_runtime_memory_destroy(void)
 {
 #if WASM_ENABLE_SHARED_HEAP != 0
-    wasm_runtime_destroy_shared_heaps();
+    destroy_shared_heaps();
 #endif
 
     if (memory_mode == MEMORY_MODE_POOL) {
@@ -1178,7 +1259,7 @@ wasm_mremap_linear_memory(void *mapped_mem, uint64 old_size, uint64 new_size,
 }
 
 static void *
-wasm_mmap_linear_memory(uint64_t map_size, uint64 commit_size)
+wasm_mmap_linear_memory(uint64 map_size, uint64 commit_size)
 {
     return wasm_mremap_linear_memory(NULL, 0, map_size, commit_size);
 }

core/iwasm/common/wasm_memory.h

@@ -54,9 +54,13 @@ wasm_runtime_attach_shared_heap_internal(WASMModuleInstanceCommon *module_inst,
 
 void
 wasm_runtime_detach_shared_heap(WASMModuleInstanceCommon *module_inst);
+
 void
 wasm_runtime_detach_shared_heap_internal(WASMModuleInstanceCommon *module_inst);
 
+WASMSharedHeap *
+wasm_runtime_get_shared_heap(WASMModuleInstanceCommon *module_inst_comm);
+
 uint64
 wasm_runtime_shared_heap_malloc(WASMModuleInstanceCommon *module_inst,
                                 uint64 size, void **p_native_addr);

core/iwasm/common/wasm_runtime_common.c

@@ -185,6 +185,9 @@ static bool
 is_sig_addr_in_guard_pages(void *sig_addr, WASMModuleInstance *module_inst)
 {
     WASMMemoryInstance *memory_inst;
+#if WASM_ENABLE_SHARED_HEAP != 0
+    WASMSharedHeap *shared_heap;
+#endif
     uint8 *mapped_mem_start_addr = NULL;
     uint8 *mapped_mem_end_addr = NULL;
     uint32 i;
@@ -202,6 +205,21 @@ is_sig_addr_in_guard_pages(void *sig_addr, WASMModuleInstance *module_inst)
         }
     }
 
+#if WASM_ENABLE_SHARED_HEAP != 0
+    shared_heap =
+        wasm_runtime_get_shared_heap((WASMModuleInstanceCommon *)module_inst);
+    if (shared_heap) {
+        mapped_mem_start_addr = shared_heap->base_addr;
+        mapped_mem_end_addr = shared_heap->base_addr + 8 * (uint64)BH_GB;
+        if (mapped_mem_start_addr <= (uint8 *)sig_addr
+            && (uint8 *)sig_addr < mapped_mem_end_addr) {
+            /* The address which causes segmentation fault is inside
+               the shared heap's guard regions */
+            return true;
+        }
+    }
+#endif
+
     return false;
 }
 

core/iwasm/common/wasm_shared_memory.c

@@ -8,6 +8,9 @@
 #if WASM_ENABLE_THREAD_MGR != 0
 #include "../libraries/thread-mgr/thread_manager.h"
 #endif
+#if WASM_ENABLE_AOT != 0
+#include "../aot/aot_runtime.h"
+#endif
 
 /*
  * Note: this lock can be per memory.
@@ -257,7 +260,9 @@ is_native_addr_in_shared_heap(WASMModuleInstanceCommon *module_inst,
 #endif
 #if WASM_ENABLE_AOT != 0
     if (module_inst->module_type == Wasm_Module_AoT) {
-        // TODO
+        AOTModuleInstanceExtra *e =
+            (AOTModuleInstanceExtra *)((AOTModuleInstance *)module_inst)->e;
+        shared_heap = e->shared_heap;
     }
 #endif