User defined memory allocator for different purposes (#3316)

Some issues are related with memory fragmentation, which may cause
the linear memory cannot be allocated. In WAMR, the memory managed
by the system is often trivial, but linear memory usually directly allocates
a large block and often remains unchanged for a long time. Their sensitivity
and contribution to fragmentation are different, which is suitable for
different allocation strategies. If we can control the linear memory's allocation,
do not make it from system heap, the overhead of heap management might
be avoided.

Add `mem_alloc_usage_t usage` as the first argument for user defined
malloc/realloc/free functions when `WAMR_BUILD_ALLOC_WITH_USAGE` cmake
variable is set as 1, and make passing `Alloc_For_LinearMemory` to the
argument when allocating the linear memory.

build-scripts/config_common.cmake

@@ -556,6 +556,9 @@ else ()
   # Disable aot intrinsics for interp, fast-jit and llvm-jit
   add_definitions (-DWASM_ENABLE_AOT_INTRINSICS=0)
 endif ()
+if (WAMR_BUILD_ALLOC_WITH_USAGE EQUAL 1)
+  add_definitions(-DWASM_MEM_ALLOC_WITH_USAGE=1)
+endif()
 if (NOT WAMR_BUILD_SANITIZER STREQUAL "")
   message ("     Sanitizer ${WAMR_BUILD_SANITIZER} enabled")
 endif ()

core/config.h

@@ -587,4 +587,8 @@
 #define WASM_TABLE_MAX_SIZE 1024
 #endif
 
+#ifndef WASM_MEM_ALLOC_WITH_USAGE
+#define WASM_MEM_ALLOC_WITH_USAGE 0
+#endif
+
 #endif /* end of _CONFIG_H_ */

core/iwasm/common/wasm_memory.c

@@ -29,16 +29,35 @@ static void *enlarge_memory_error_user_data;
 
 #if WASM_MEM_ALLOC_WITH_USER_DATA != 0
 static void *allocator_user_data = NULL;
-static void *(*malloc_func)(void *user_data, unsigned int size) = NULL;
-static void *(*realloc_func)(void *user_data, void *ptr,
-                             unsigned int size) = NULL;
-static void (*free_func)(void *user_data, void *ptr) = NULL;
-#else
-static void *(*malloc_func)(unsigned int size) = NULL;
-static void *(*realloc_func)(void *ptr, unsigned int size) = NULL;
-static void (*free_func)(void *ptr) = NULL;
 #endif
 
+static void *(*malloc_func)(
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+    mem_alloc_usage_t usage,
+#endif
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+    void *user_data,
+#endif
+    unsigned int size) = NULL;
+
+static void *(*realloc_func)(
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+    mem_alloc_usage_t usage, bool full_size_mmaped,
+#endif
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+    void *user_data,
+#endif
+    void *ptr, unsigned int size) = NULL;
+
+static void (*free_func)(
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+    mem_alloc_usage_t usage,
+#endif
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+    void *user_data,
+#endif
+    void *ptr) = NULL;
+
 static unsigned int global_pool_size;
 
 static uint64
@@ -177,11 +196,14 @@ wasm_runtime_malloc_internal(unsigned int size)
         return mem_allocator_malloc(pool_allocator, size);
     }
     else if (memory_mode == MEMORY_MODE_ALLOCATOR) {
+        return malloc_func(
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+            Alloc_For_Runtime,
+#endif
 #if WASM_MEM_ALLOC_WITH_USER_DATA != 0
-        return malloc_func(allocator_user_data, size);
-#else
-        return malloc_func(size);
+            allocator_user_data,
 #endif
+            size);
     }
     else {
         return os_malloc(size);
@@ -201,11 +223,14 @@ wasm_runtime_realloc_internal(void *ptr, unsigned int size)
     }
     else if (memory_mode == MEMORY_MODE_ALLOCATOR) {
         if (realloc_func)
+            return realloc_func(
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+                Alloc_For_Runtime, false,
+#endif
 #if WASM_MEM_ALLOC_WITH_USER_DATA != 0
-            return realloc_func(allocator_user_data, ptr, size);
-#else
-            return realloc_func(ptr, size);
+                allocator_user_data,
 #endif
+                ptr, size);
         else
             return NULL;
     }
@@ -233,11 +258,14 @@ wasm_runtime_free_internal(void *ptr)
         mem_allocator_free(pool_allocator, ptr);
     }
     else if (memory_mode == MEMORY_MODE_ALLOCATOR) {
+        free_func(
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+            Alloc_For_Runtime,
+#endif
 #if WASM_MEM_ALLOC_WITH_USER_DATA != 0
-        free_func(allocator_user_data, ptr);
-#else
-        free_func(ptr);
+            allocator_user_data,
 #endif
+            ptr);
     }
     else {
         os_free(ptr);
@@ -765,6 +793,29 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count)
     bh_assert(total_size_new
               <= GET_MAX_LINEAR_MEMORY_SIZE(memory->is_memory64));
 
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+    if (!(memory_data_new =
+              realloc_func(Alloc_For_LinearMemory, full_size_mmaped,
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+                           NULL,
+#endif
+                           memory_data_old, total_size_new))) {
+        ret = false;
+        goto return_func;
+    }
+    if (heap_size > 0) {
+        if (mem_allocator_migrate(memory->heap_handle,
+                                  (char *)heap_data_old
+                                      + (memory_data_new - memory_data_old),
+                                  heap_size)
+            != 0) {
+            ret = false;
+        }
+    }
+    memory->heap_data = memory_data_new + (heap_data_old - memory_data_old);
+    memory->heap_data_end = memory->heap_data + heap_size;
+    memory->memory_data = memory_data_new;
+#else
     if (full_size_mmaped) {
 #ifdef BH_PLATFORM_WINDOWS
         if (!os_mem_commit(memory->memory_data_end,
@@ -823,6 +874,7 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count)
         os_writegsbase(memory_data_new);
 #endif
     }
+#endif /* end of WASM_MEM_ALLOC_WITH_USAGE */
 
     memory->num_bytes_per_page = num_bytes_per_page;
     memory->cur_page_count = total_page_count;
@@ -903,8 +955,19 @@ wasm_deallocate_linear_memory(WASMMemoryInstance *memory_inst)
 #else
     map_size = 8 * (uint64)BH_GB;
 #endif
+
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+    (void)map_size;
+    free_func(Alloc_For_LinearMemory,
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+              NULL,
+#endif
+              memory_inst->memory_data);
+#else
     wasm_munmap_linear_memory(memory_inst->memory_data,
                               memory_inst->memory_data_size, map_size);
+#endif
+
     memory_inst->memory_data = NULL;
 }
 
@@ -954,9 +1017,20 @@ wasm_allocate_linear_memory(uint8 **data, bool is_shared_memory,
     *memory_data_size = align_as_and_cast(*memory_data_size, page_size);
 
     if (map_size > 0) {
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+        (void)wasm_mmap_linear_memory;
+        if (!(*data = malloc_func(Alloc_For_LinearMemory,
+#if WASM_MEM_ALLOC_WITH_USER_DATA != 0
+                                  NULL,
+#endif
+                                  *memory_data_size))) {
+            return BHT_ERROR;
+        }
+#else
         if (!(*data = wasm_mmap_linear_memory(map_size, *memory_data_size))) {
             return BHT_ERROR;
         }
+#endif
     }
 
     return BHT_OK;

core/iwasm/include/wasm_export.h

@@ -113,6 +113,11 @@ typedef enum {
     Alloc_With_System_Allocator,
 } mem_alloc_type_t;
 
+typedef enum {
+    Alloc_For_Runtime,
+    Alloc_For_LinearMemory
+} mem_alloc_usage_t;
+
 /* Memory allocator option */
 typedef union MemAllocOption {
     struct {
@@ -120,6 +125,9 @@ typedef union MemAllocOption {
         uint32_t heap_size;
     } pool;
     struct {
+        /* the function signature is varied when
+        WASM_MEM_ALLOC_WITH_USER_DATA and
+        WASM_MEM_ALLOC_WITH_USAGE are defined */
         void *malloc_func;
         void *realloc_func;
         void *free_func;

doc/build_wamr.md

@@ -254,6 +254,10 @@ Currently we only profile the memory consumption of module, module_instance and
 
 > See [Enable segue optimization for wamrc when generating the aot file](./perf_tune.md#3-enable-segue-optimization-for-wamrc-when-generating-the-aot-file) for more details.
 
+#### **User defined linear memory allocator**
+- **WAMR_BUILD_ALLOC_WITH_USAGE**=1/0, default to disable if not set
+> Notes: by default, the linear memory is allocated by system. when it's set to 1 and Alloc_With_Allocator is selected, it will be allocated by customer.
+
 #### **Enable running PGO(Profile-Guided Optimization) instrumented AOT file**
 - **WAMR_BUILD_STATIC_PGO**=1/0, default to disable if not set
 > Note: See [Use the AOT static PGO method](./perf_tune.md#5-use-the-aot-static-pgo-method) for more details.

product-mini/platforms/nuttx/CMakeLists.txt

@@ -129,6 +129,10 @@ if(CONFIG_INTERPRETERS_WAMR_GLOBAL_HEAP_POOL)
   set(WAMR_BUILD_GLOBAL_HEAP_SIZE ${_HEAP_SIZE_})
 endif()
 
+if (CONFIG_INTERPRETERS_WAMR_MEM_ALLOC_WITH_USAGE)
+  set(WAMR_BUILD_MEM_ALLOC_WITH_USAGE 1)
+endif()
+
 if(CONFIG_INTERPRETERS_WAMR_ENABLE_SPEC_TEST)
   set(WAMR_BUILD_SPEC_TEST 1)
 endif()

product-mini/platforms/nuttx/wamr.mk

@@ -373,6 +373,11 @@ CFLAGS += -DWASM_ENABLE_GLOBAL_HEAP_POOL=1
 CFLAGS += -DWASM_GLOBAL_HEAP_SIZE="$(CONFIG_INTERPRETERS_WAMR_GLOBAL_HEAP_POOL_SIZE) * 1024"
 else
 CFLAGS += -DWASM_ENABLE_GLOBAL_HEAP_POOL=0
+ifeq ($(CONFIG_INTERPRETERS_WAMR_MEM_ALLOC_WITH_USAGE),y)
+CFLAGS += -DWASM_MEM_ALLOC_WITH_USAGE=1
+else
+CFLAGS += -DWASM_MEM_ALLOC_WITH_USAGE=0
+endif
 endif
 
 ifeq ($(CONFIG_INTERPRETERS_WAMR_ENABLE_SPEC_TEST),y)

product-mini/platforms/posix/main.c

@@ -445,6 +445,9 @@ static char global_heap_buf[WASM_GLOBAL_HEAP_SIZE] = { 0 };
 #else
 static void *
 malloc_func(
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+    mem_alloc_usage_t usage,
+#endif
 #if WASM_MEM_ALLOC_WITH_USER_DATA != 0
     void *user_data,
 #endif
@@ -455,6 +458,9 @@ malloc_func(
 
 static void *
 realloc_func(
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+    mem_alloc_usage_t usage, bool full_size_mmaped,
+#endif
 #if WASM_MEM_ALLOC_WITH_USER_DATA != 0
     void *user_data,
 #endif
@@ -465,6 +471,9 @@ realloc_func(
 
 static void
 free_func(
+#if WASM_MEM_ALLOC_WITH_USAGE != 0
+    mem_alloc_usage_t usage,
+#endif
 #if WASM_MEM_ALLOC_WITH_USER_DATA != 0
     void *user_data,
 #endif