Optimize memory initialization handling in AOT loader (#3983)

Save memory if the file buffer is always exist before exit.

Signed-off-by: Huang Qi <huangqi3@xiaomi.com>

core/iwasm/aot/aot_loader.c

@@ -967,13 +967,29 @@ destroy_import_memories(AOTImportMemory *import_memories)
     wasm_runtime_free(import_memories);
 }
 
+/**
+ * Free memory initialization data segments.
+ *
+ * @param module the AOT module containing the data
+ * @param data_list array of memory initialization data segments to free
+ * @param count number of segments in the data_list array
+ */
+
 static void
-destroy_mem_init_data_list(AOTMemInitData **data_list, uint32 count)
+destroy_mem_init_data_list(AOTModule *module, AOTMemInitData **data_list,
+                           uint32 count)
 {
     uint32 i;
+    /* Free each memory initialization data segment */
     for (i = 0; i < count; i++)
-        if (data_list[i])
+        if (data_list[i]) {
+            /* If the module owns the binary data, free the bytes buffer */
+            if (module->is_binary_freeable && data_list[i]->bytes)
+                wasm_runtime_free(data_list[i]->bytes);
+            /* Free the data segment structure itself */
             wasm_runtime_free(data_list[i]);
+        }
+    /* Free the array of data segment pointers */
     wasm_runtime_free(data_list);
 }
 
@@ -982,6 +998,22 @@ load_init_expr(const uint8 **p_buf, const uint8 *buf_end, AOTModule *module,
                InitializerExpression *expr, char *error_buf,
                uint32 error_buf_size);
 
+/**
+ * Load memory initialization data segments from the AOT module.
+ *
+ * This function reads memory initialization data segments from the buffer and
+ * creates AOTMemInitData structures for each segment. The data can either be
+ * cloned into new memory or referenced directly from the buffer.
+ *
+ * @param p_buf pointer to buffer containing memory init data
+ * @param buf_end end of buffer
+ * @param module the AOT module being loaded
+ * @param error_buf buffer for error messages
+ * @param error_buf_size size of error buffer
+ *
+ * @return true if successful, false if error occurred
+ */
+
 static bool
 load_mem_init_data_list(const uint8 **p_buf, const uint8 *buf_end,
                         AOTModule *module, char *error_buf,
@@ -1013,8 +1045,8 @@ load_mem_init_data_list(const uint8 **p_buf, const uint8 *buf_end,
             return false;
         }
         read_uint32(buf, buf_end, byte_count);
-        size = offsetof(AOTMemInitData, bytes) + (uint64)byte_count;
-        if (!(data_list[i] = loader_malloc(size, error_buf, error_buf_size))) {
+        if (!(data_list[i] = loader_malloc(sizeof(AOTMemInitData), error_buf,
+                                           error_buf_size))) {
             return false;
         }
 
@@ -1026,8 +1058,22 @@ load_mem_init_data_list(const uint8 **p_buf, const uint8 *buf_end,
         data_list[i]->offset.init_expr_type = init_value.init_expr_type;
         data_list[i]->offset.u = init_value.u;
         data_list[i]->byte_count = byte_count;
-        read_byte_array(buf, buf_end, data_list[i]->bytes,
-                        data_list[i]->byte_count);
+        data_list[i]->bytes = NULL;
+        /* If the module owns the binary data, clone the bytes buffer */
+        if (module->is_binary_freeable) {
+            if (byte_count > 0) {
+                if (!(data_list[i]->bytes = loader_malloc(byte_count, error_buf,
+                                                          error_buf_size))) {
+                    return false;
+                }
+                read_byte_array(buf, buf_end, data_list[i]->bytes,
+                                data_list[i]->byte_count);
+            }
+        }
+        else {
+            data_list[i]->bytes = (uint8 *)buf;
+            buf += byte_count;
+        }
     }
 
     *p_buf = buf;
@@ -1036,6 +1082,21 @@ fail:
     return false;
 }
 
+/**
+ * Load memory information from the AOT module.
+ *
+ * This function reads memory-related data including import memory count,
+ * memory count, memory flags, page sizes, and memory initialization data.
+ *
+ * @param p_buf pointer to buffer containing memory info
+ * @param buf_end end of buffer
+ * @param module the AOT module being loaded
+ * @param error_buf buffer for error messages
+ * @param error_buf_size size of error buffer
+ *
+ * @return true if successful, false if error occurred
+ */
+
 static bool
 load_memory_info(const uint8 **p_buf, const uint8 *buf_end, AOTModule *module,
                  char *error_buf, uint32 error_buf_size)
@@ -4356,7 +4417,7 @@ aot_unload(AOTModule *module)
         wasm_runtime_free(module->memories);
 
     if (module->mem_init_data_list)
-        destroy_mem_init_data_list(module->mem_init_data_list,
+        destroy_mem_init_data_list(module, module->mem_init_data_list,
                                    module->mem_init_data_count);
 
     if (module->native_symbol_list)

core/iwasm/compilation/aot.c

@@ -36,8 +36,11 @@ aot_destroy_mem_init_data_list(AOTMemInitData **data_list, uint32 count)
 {
     uint32 i;
     for (i = 0; i < count; i++)
-        if (data_list[i])
+        if (data_list[i]) {
+            if (data_list[i]->bytes)
+                wasm_runtime_free(data_list[i]->bytes);
             wasm_runtime_free(data_list[i]);
+        }
     wasm_runtime_free(data_list);
 }
 
@@ -60,8 +63,7 @@ aot_create_mem_init_data_list(const WASMModule *module)
 
     /* Create each memory data segment */
     for (i = 0; i < module->data_seg_count; i++) {
-        size = offsetof(AOTMemInitData, bytes)
-               + (uint64)module->data_segments[i]->data_length;
+        size = sizeof(AOTMemInitData);
         if (size >= UINT32_MAX
             || !(data_list[i] = wasm_runtime_malloc((uint32)size))) {
             aot_set_last_error("allocate memory failed.");
@@ -69,18 +71,31 @@ aot_create_mem_init_data_list(const WASMModule *module)
         }
 
 #if WASM_ENABLE_BULK_MEMORY != 0
+        /* Set bulk memory specific properties if enabled */
         data_list[i]->is_passive = module->data_segments[i]->is_passive;
         data_list[i]->memory_index = module->data_segments[i]->memory_index;
 #endif
         data_list[i]->offset = module->data_segments[i]->base_offset;
         data_list[i]->byte_count = module->data_segments[i]->data_length;
-        memcpy(data_list[i]->bytes, module->data_segments[i]->data,
-               module->data_segments[i]->data_length);
+        data_list[i]->bytes = NULL;
+        /* Allocate memory for AOT compiler is OK, because the data segment
+         * is small and the host memory is enough */
+        if (data_list[i]->byte_count > 0) {
+            data_list[i]->bytes = wasm_runtime_malloc(data_list[i]->byte_count);
+            if (!data_list[i]->bytes) {
+                aot_set_last_error("allocate memory failed.");
+                goto fail;
+            }
+            /* Copy the actual data bytes from the WASM module */
+            memcpy(data_list[i]->bytes, module->data_segments[i]->data,
+                   module->data_segments[i]->data_length);
+        }
     }
 
     return data_list;
 
 fail:
+    /* Clean up allocated memory in case of failure */
     aot_destroy_mem_init_data_list(data_list, module->data_seg_count);
     return NULL;
 }

core/iwasm/compilation/aot.h

@@ -103,7 +103,7 @@ typedef struct AOTMemInitData {
     /* Byte count */
     uint32 byte_count;
     /* Byte array */
-    uint8 bytes[1];
+    uint8 *bytes;
 } AOTMemInitData;
 
 /**