From 867dbd8912f8fd489ea9b24497abb9d3c81f1b40 Mon Sep 17 00:00:00 2001
From: YAMAMOTO Takashi <yamamoto@midokura.com>
Date: Tue, 25 Jun 2024 11:57:49 +0900
Subject: [PATCH] aot compiler: Propagate const-ness by ourselves (#3567)

aot_load_const_from_table() hides the const-ness of the
value and prevents optimizations like
https://github.com/bytecodealliance/wasm-micro-runtime/pull/3552.

This commit makes the aot compiler tracks the const-ness
of the value directly in the AOTValue and enables the above
mentioned optimization for XIP.
---
 core/iwasm/compilation/aot_compiler.h    |  8 ++++++++
 core/iwasm/compilation/aot_emit_const.c  |  2 ++
 core/iwasm/compilation/aot_emit_memory.c | 16 +++++++++++++---
 core/iwasm/compilation/aot_llvm.h        |  4 +++-
 4 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/core/iwasm/compilation/aot_compiler.h b/core/iwasm/compilation/aot_compiler.h
index 4172b46f2..ab74b7cb6 100644
--- a/core/iwasm/compilation/aot_compiler.h
+++ b/core/iwasm/compilation/aot_compiler.h
@@ -605,6 +605,14 @@ set_local_gc_ref(AOTCompFrame *frame, int n, LLVMValueRef value, uint8 ref_type)
 #define PUSH_PAGE_COUNT(v) \
     PUSH(v, MEMORY64_COND_VALUE(VALUE_TYPE_I64, VALUE_TYPE_I32))
 
+#define SET_CONST(v)                                                          \
+    do {                                                                      \
+        AOTValue *aot_value =                                                 \
+            func_ctx->block_stack.block_list_end->value_stack.value_list_end; \
+        aot_value->is_const = true;                                           \
+        aot_value->const_value = (v);                                         \
+    } while (0)
+
 #define TO_LLVM_TYPE(wasm_type) \
     wasm_type_to_llvm_type(comp_ctx, &comp_ctx->basic_types, wasm_type)
 
diff --git a/core/iwasm/compilation/aot_emit_const.c b/core/iwasm/compilation/aot_emit_const.c
index 4b38aa962..5665b480b 100644
--- a/core/iwasm/compilation/aot_emit_const.c
+++ b/core/iwasm/compilation/aot_emit_const.c
@@ -28,6 +28,7 @@ aot_compile_op_i32_const(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
     }
 
     PUSH_I32(value);
+    SET_CONST((uint64)(uint32)i32_const);
     return true;
 fail:
     return false;
@@ -55,6 +56,7 @@ aot_compile_op_i64_const(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
     }
 
     PUSH_I64(value);
+    SET_CONST((uint64)i64_const);
     return true;
 fail:
     return false;
diff --git a/core/iwasm/compilation/aot_emit_memory.c b/core/iwasm/compilation/aot_emit_memory.c
index a74d32681..ff4923966 100644
--- a/core/iwasm/compilation/aot_emit_memory.c
+++ b/core/iwasm/compilation/aot_emit_memory.c
@@ -107,7 +107,9 @@ aot_check_memory_overflow(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
     LLVMBasicBlockRef check_succ;
     AOTValue *aot_value_top;
     uint32 local_idx_of_aot_value = 0;
+    uint64 const_value;
     bool is_target_64bit, is_local_of_aot_value = false;
+    bool is_const = false;
 #if WASM_ENABLE_SHARED_MEMORY != 0
     bool is_shared_memory =
         comp_ctx->comp_data->memories[0].flags & SHARED_MEMORY_FLAG;
@@ -162,7 +164,9 @@ aot_check_memory_overflow(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
         /* aot_value_top is freed in the following POP_I32(addr),
            so save its fields here for further use */
         is_local_of_aot_value = aot_value_top->is_local;
+        is_const = aot_value_top->is_const;
         local_idx_of_aot_value = aot_value_top->local_idx;
+        const_value = aot_value_top->const_value;
     }
 
     POP_MEM_OFFSET(addr);
@@ -172,9 +176,15 @@ aot_check_memory_overflow(AOTCompContext *comp_ctx, AOTFuncContext *func_ctx,
      * have been thrown when converting float to integer before
      */
     /* return address directly if constant offset and inside memory space */
-    if (LLVMIsEfficientConstInt(addr)) {
-        uint64 mem_offset =
-            (uint64)LLVMConstIntGetZExtValue(addr) + (uint64)offset;
+    if (LLVMIsEfficientConstInt(addr) || is_const) {
+        uint64 value;
+        if (LLVMIsEfficientConstInt(addr)) {
+            value = (uint64)LLVMConstIntGetZExtValue(addr);
+        }
+        else {
+            value = const_value;
+        }
+        uint64 mem_offset = value + (uint64)offset;
         uint32 num_bytes_per_page =
             comp_ctx->comp_data->memories[0].num_bytes_per_page;
         uint32 init_page_count =
diff --git a/core/iwasm/compilation/aot_llvm.h b/core/iwasm/compilation/aot_llvm.h
index 289384a16..ab5242173 100644
--- a/core/iwasm/compilation/aot_llvm.h
+++ b/core/iwasm/compilation/aot_llvm.h
@@ -75,10 +75,12 @@ typedef struct AOTValue {
     struct AOTValue *next;
     struct AOTValue *prev;
     LLVMValueRef value;
+    uint64 const_value; /* valid if is_const is true */
+    uint32 local_idx;
     /* VALUE_TYPE_I32/I64/F32/F64/VOID */
     uint8 type;
     bool is_local;
-    uint32 local_idx;
+    bool is_const;
 } AOTValue;
 
 /**