Skip to content

Commit 38cf274

Browse files
authored
Optimize memory initialization handling in AOT loader (#3983)
Save memory if the file buffer is always exist before exit. Signed-off-by: Huang Qi <[email protected]>
1 parent 04f1071 commit 38cf274

File tree

3 files changed

+89
-13
lines changed

3 files changed

+89
-13
lines changed

core/iwasm/aot/aot_loader.c

+68-7
Original file line numberDiff line numberDiff line change
@@ -967,13 +967,29 @@ destroy_import_memories(AOTImportMemory *import_memories)
967967
wasm_runtime_free(import_memories);
968968
}
969969

970+
/**
971+
* Free memory initialization data segments.
972+
*
973+
* @param module the AOT module containing the data
974+
* @param data_list array of memory initialization data segments to free
975+
* @param count number of segments in the data_list array
976+
*/
977+
970978
static void
971-
destroy_mem_init_data_list(AOTMemInitData **data_list, uint32 count)
979+
destroy_mem_init_data_list(AOTModule *module, AOTMemInitData **data_list,
980+
uint32 count)
972981
{
973982
uint32 i;
983+
/* Free each memory initialization data segment */
974984
for (i = 0; i < count; i++)
975-
if (data_list[i])
985+
if (data_list[i]) {
986+
/* If the module owns the binary data, free the bytes buffer */
987+
if (module->is_binary_freeable && data_list[i]->bytes)
988+
wasm_runtime_free(data_list[i]->bytes);
989+
/* Free the data segment structure itself */
976990
wasm_runtime_free(data_list[i]);
991+
}
992+
/* Free the array of data segment pointers */
977993
wasm_runtime_free(data_list);
978994
}
979995

@@ -982,6 +998,22 @@ load_init_expr(const uint8 **p_buf, const uint8 *buf_end, AOTModule *module,
982998
InitializerExpression *expr, char *error_buf,
983999
uint32 error_buf_size);
9841000

1001+
/**
1002+
* Load memory initialization data segments from the AOT module.
1003+
*
1004+
* This function reads memory initialization data segments from the buffer and
1005+
* creates AOTMemInitData structures for each segment. The data can either be
1006+
* cloned into new memory or referenced directly from the buffer.
1007+
*
1008+
* @param p_buf pointer to buffer containing memory init data
1009+
* @param buf_end end of buffer
1010+
* @param module the AOT module being loaded
1011+
* @param error_buf buffer for error messages
1012+
* @param error_buf_size size of error buffer
1013+
*
1014+
* @return true if successful, false if error occurred
1015+
*/
1016+
9851017
static bool
9861018
load_mem_init_data_list(const uint8 **p_buf, const uint8 *buf_end,
9871019
AOTModule *module, char *error_buf,
@@ -1013,8 +1045,8 @@ load_mem_init_data_list(const uint8 **p_buf, const uint8 *buf_end,
10131045
return false;
10141046
}
10151047
read_uint32(buf, buf_end, byte_count);
1016-
size = offsetof(AOTMemInitData, bytes) + (uint64)byte_count;
1017-
if (!(data_list[i] = loader_malloc(size, error_buf, error_buf_size))) {
1048+
if (!(data_list[i] = loader_malloc(sizeof(AOTMemInitData), error_buf,
1049+
error_buf_size))) {
10181050
return false;
10191051
}
10201052

@@ -1026,8 +1058,22 @@ load_mem_init_data_list(const uint8 **p_buf, const uint8 *buf_end,
10261058
data_list[i]->offset.init_expr_type = init_value.init_expr_type;
10271059
data_list[i]->offset.u = init_value.u;
10281060
data_list[i]->byte_count = byte_count;
1029-
read_byte_array(buf, buf_end, data_list[i]->bytes,
1030-
data_list[i]->byte_count);
1061+
data_list[i]->bytes = NULL;
1062+
/* If the module owns the binary data, clone the bytes buffer */
1063+
if (module->is_binary_freeable) {
1064+
if (byte_count > 0) {
1065+
if (!(data_list[i]->bytes = loader_malloc(byte_count, error_buf,
1066+
error_buf_size))) {
1067+
return false;
1068+
}
1069+
read_byte_array(buf, buf_end, data_list[i]->bytes,
1070+
data_list[i]->byte_count);
1071+
}
1072+
}
1073+
else {
1074+
data_list[i]->bytes = (uint8 *)buf;
1075+
buf += byte_count;
1076+
}
10311077
}
10321078

10331079
*p_buf = buf;
@@ -1036,6 +1082,21 @@ load_mem_init_data_list(const uint8 **p_buf, const uint8 *buf_end,
10361082
return false;
10371083
}
10381084

1085+
/**
1086+
* Load memory information from the AOT module.
1087+
*
1088+
* This function reads memory-related data including import memory count,
1089+
* memory count, memory flags, page sizes, and memory initialization data.
1090+
*
1091+
* @param p_buf pointer to buffer containing memory info
1092+
* @param buf_end end of buffer
1093+
* @param module the AOT module being loaded
1094+
* @param error_buf buffer for error messages
1095+
* @param error_buf_size size of error buffer
1096+
*
1097+
* @return true if successful, false if error occurred
1098+
*/
1099+
10391100
static bool
10401101
load_memory_info(const uint8 **p_buf, const uint8 *buf_end, AOTModule *module,
10411102
char *error_buf, uint32 error_buf_size)
@@ -4356,7 +4417,7 @@ aot_unload(AOTModule *module)
43564417
wasm_runtime_free(module->memories);
43574418

43584419
if (module->mem_init_data_list)
4359-
destroy_mem_init_data_list(module->mem_init_data_list,
4420+
destroy_mem_init_data_list(module, module->mem_init_data_list,
43604421
module->mem_init_data_count);
43614422

43624423
if (module->native_symbol_list)

core/iwasm/compilation/aot.c

+20-5
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,11 @@ aot_destroy_mem_init_data_list(AOTMemInitData **data_list, uint32 count)
3636
{
3737
uint32 i;
3838
for (i = 0; i < count; i++)
39-
if (data_list[i])
39+
if (data_list[i]) {
40+
if (data_list[i]->bytes)
41+
wasm_runtime_free(data_list[i]->bytes);
4042
wasm_runtime_free(data_list[i]);
43+
}
4144
wasm_runtime_free(data_list);
4245
}
4346

@@ -60,27 +63,39 @@ aot_create_mem_init_data_list(const WASMModule *module)
6063

6164
/* Create each memory data segment */
6265
for (i = 0; i < module->data_seg_count; i++) {
63-
size = offsetof(AOTMemInitData, bytes)
64-
+ (uint64)module->data_segments[i]->data_length;
66+
size = sizeof(AOTMemInitData);
6567
if (size >= UINT32_MAX
6668
|| !(data_list[i] = wasm_runtime_malloc((uint32)size))) {
6769
aot_set_last_error("allocate memory failed.");
6870
goto fail;
6971
}
7072

7173
#if WASM_ENABLE_BULK_MEMORY != 0
74+
/* Set bulk memory specific properties if enabled */
7275
data_list[i]->is_passive = module->data_segments[i]->is_passive;
7376
data_list[i]->memory_index = module->data_segments[i]->memory_index;
7477
#endif
7578
data_list[i]->offset = module->data_segments[i]->base_offset;
7679
data_list[i]->byte_count = module->data_segments[i]->data_length;
77-
memcpy(data_list[i]->bytes, module->data_segments[i]->data,
78-
module->data_segments[i]->data_length);
80+
data_list[i]->bytes = NULL;
81+
/* Allocate memory for AOT compiler is OK, because the data segment
82+
* is small and the host memory is enough */
83+
if (data_list[i]->byte_count > 0) {
84+
data_list[i]->bytes = wasm_runtime_malloc(data_list[i]->byte_count);
85+
if (!data_list[i]->bytes) {
86+
aot_set_last_error("allocate memory failed.");
87+
goto fail;
88+
}
89+
/* Copy the actual data bytes from the WASM module */
90+
memcpy(data_list[i]->bytes, module->data_segments[i]->data,
91+
module->data_segments[i]->data_length);
92+
}
7993
}
8094

8195
return data_list;
8296

8397
fail:
98+
/* Clean up allocated memory in case of failure */
8499
aot_destroy_mem_init_data_list(data_list, module->data_seg_count);
85100
return NULL;
86101
}

core/iwasm/compilation/aot.h

+1-1
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ typedef struct AOTMemInitData {
103103
/* Byte count */
104104
uint32 byte_count;
105105
/* Byte array */
106-
uint8 bytes[1];
106+
uint8 *bytes;
107107
} AOTMemInitData;
108108

109109
/**

0 commit comments

Comments
 (0)