Fix fast-jit accessing shared memory's fields issue (#2841)

For shared memory, runtime should get the memories pointer from module_inst first, then get memory instance from memories array, and then get the fields of the memory instance.
2025-02-11 09:25:20 +00:00 · 2023-12-01 17:41:24 +08:00 · 2023-12-01 17:41:24 +08:00 · 3d0342fbc8
commit 3d0342fbc8
parent 162a977006
4 changed files with 358 additions and 92 deletions
--- a/core/iwasm/fast-jit/fe/jit_emit_memory.c
+++ b/core/iwasm/fast-jit/fe/jit_emit_memory.c
@ -137,6 +137,7 @@ check_and_seek(JitCompContext *cc, JitReg addr, uint32 offset, uint32 bytes)
 {
    JitReg memory_boundary = 0, offset1;
 #ifndef OS_ENABLE_HW_BOUND_CHECK
+    JitReg cur_page_count;
    /* the default memory */
    uint32 mem_idx = 0;
 #endif
@ -146,16 +147,10 @@ check_and_seek(JitCompContext *cc, JitReg addr, uint32 offset, uint32 bytes)
    /* 1. shortcut if the memory size is 0 */
    if (cc->cur_wasm_module->memories != NULL
        && 0 == cc->cur_wasm_module->memories[mem_idx].init_page_count) {
-        JitReg module_inst, cur_page_count;
-        uint32 cur_page_count_offset =
-            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-            + (uint32)offsetof(WASMMemoryInstance, cur_page_count);
+
+        cur_page_count = get_cur_page_count_reg(cc->jit_frame, mem_idx);

        /* if (cur_mem_page_count == 0) goto EXCEPTION */
-        module_inst = get_module_inst_reg(cc->jit_frame);
-        cur_page_count = jit_cc_new_reg_I32(cc);
-        GEN_INSN(LDI32, cur_page_count, module_inst,
-                 NEW_CONST(I32, cur_page_count_offset));
        GEN_INSN(CMP, cc->cmp_reg, cur_page_count, NEW_CONST(I32, 0));
        if (!jit_emit_exception(cc, EXCE_OUT_OF_BOUNDS_MEMORY_ACCESS,
                                JIT_OP_BEQ, cc->cmp_reg, NULL)) {
@ -580,15 +575,9 @@ fail:
 bool
 jit_compile_op_memory_size(JitCompContext *cc, uint32 mem_idx)
 {
-    JitReg module_inst, cur_page_count;
-    uint32 cur_page_count_offset =
-        (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-        + (uint32)offsetof(WASMMemoryInstance, cur_page_count);
+    JitReg cur_page_count;

-    module_inst = get_module_inst_reg(cc->jit_frame);
-    cur_page_count = jit_cc_new_reg_I32(cc);
-    GEN_INSN(LDI32, cur_page_count, module_inst,
-             NEW_CONST(I32, cur_page_count_offset));
+    cur_page_count = get_cur_page_count_reg(cc->jit_frame, mem_idx);

    PUSH_I32(cur_page_count);

@ -600,18 +589,11 @@ fail:
 bool
 jit_compile_op_memory_grow(JitCompContext *cc, uint32 mem_idx)
 {
-    JitReg module_inst, grow_res, res;
+    JitReg grow_res, res;
    JitReg prev_page_count, inc_page_count, args[2];

-    /* Get current page count */
-    uint32 cur_page_count_offset =
-        (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-        + (uint32)offsetof(WASMMemoryInstance, cur_page_count);
-
-    module_inst = get_module_inst_reg(cc->jit_frame);
-    prev_page_count = jit_cc_new_reg_I32(cc);
-    GEN_INSN(LDI32, prev_page_count, module_inst,
-             NEW_CONST(I32, cur_page_count_offset));
+    /* Get current page count as prev_page_count */
+    prev_page_count = get_cur_page_count_reg(cc->jit_frame, mem_idx);

    /* Call wasm_enlarge_memory */
    POP_I32(inc_page_count);
--- a/core/iwasm/fast-jit/jit_frontend.c
+++ b/core/iwasm/fast-jit/jit_frontend.c
@ -218,42 +218,149 @@ get_aux_stack_bottom_reg(JitFrame *frame)
    return frame->aux_stack_bottom_reg;
 }

+#if WASM_ENABLE_SHARED_MEMORY != 0
+static bool
+is_shared_memory(WASMModule *module, uint32 mem_idx)
+{
+    WASMMemory *memory;
+    WASMMemoryImport *memory_import;
+    bool is_shared;
+
+    if (mem_idx < module->import_memory_count) {
+        memory_import = &(module->import_memories[mem_idx].u.memory);
+        is_shared = memory_import->flags & 0x02 ? true : false;
+    }
+    else {
+        memory = &module->memories[mem_idx - module->import_memory_count];
+        is_shared = memory->flags & 0x02 ? true : false;
+    }
+    return is_shared;
+}
+#endif
+
+JitReg
+get_memory_inst_reg(JitFrame *frame, uint32 mem_idx)
+{
+    JitCompContext *cc = frame->cc;
+    JitReg module_inst_reg = get_module_inst_reg(frame);
+    uint32 memory_inst_offset;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    JitReg memories_addr;
+    uint32 memories_offset;
+    bool is_shared;
+#endif
+
+    if (frame->memory_regs[mem_idx].memory_inst)
+        return frame->memory_regs[mem_idx].memory_inst;
+
+    frame->memory_regs[mem_idx].memory_inst =
+        cc->memory_regs[mem_idx].memory_inst;
+
+    bh_assert(mem_idx == 0);
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    is_shared = is_shared_memory(cc->cur_wasm_module, mem_idx);
+    if (is_shared) {
+        memories_addr = jit_cc_new_reg_ptr(cc);
+        memories_offset = (uint32)offsetof(WASMModuleInstance, memories);
+        /* module_inst->memories */
+        GEN_INSN(LDPTR, memories_addr, module_inst_reg,
+                 NEW_CONST(I32, memories_offset));
+        /* module_inst->memories[mem_idx], mem_idx can only be 0 now */
+        GEN_INSN(LDPTR, frame->memory_regs[mem_idx].memory_inst, memories_addr,
+                 NEW_CONST(I32, mem_idx));
+    }
+    else
+#endif
+    {
+        memory_inst_offset =
+            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes);
+        GEN_INSN(LDPTR, frame->memory_regs[mem_idx].memory_inst,
+                 module_inst_reg, NEW_CONST(I32, memory_inst_offset));
+    }
+
+    return frame->memory_regs[mem_idx].memory_inst;
+}
+
+JitReg
+get_cur_page_count_reg(JitFrame *frame, uint32 mem_idx)
+{
+    JitCompContext *cc = frame->cc;
+    JitReg module_inst_reg;
+    uint32 cur_page_count_offset;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    JitReg memory_inst_reg;
+    bool is_shared;
+#endif
+
+    if (frame->memory_regs[mem_idx].cur_page_count)
+        return frame->memory_regs[mem_idx].cur_page_count;
+
+    frame->memory_regs[mem_idx].cur_page_count =
+        cc->memory_regs[mem_idx].cur_page_count;
+
+    /* Get current page count */
+    bh_assert(mem_idx == 0);
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    is_shared = is_shared_memory(cc->cur_wasm_module, mem_idx);
+    if (is_shared) {
+        memory_inst_reg = get_memory_inst_reg(frame, mem_idx);
+        cur_page_count_offset =
+            (uint32)offsetof(WASMMemoryInstance, cur_page_count);
+        /* memories[mem_idx]->cur_page_count_offset */
+        GEN_INSN(LDI32, frame->memory_regs[mem_idx].cur_page_count,
+                 memory_inst_reg, NEW_CONST(I32, cur_page_count_offset));
+    }
+    else
+#endif
+    {
+        module_inst_reg = get_module_inst_reg(frame);
+        cur_page_count_offset =
+            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
+            + (uint32)offsetof(WASMMemoryInstance, cur_page_count);
+        GEN_INSN(LDI32, frame->memory_regs[mem_idx].cur_page_count,
+                 module_inst_reg, NEW_CONST(I32, cur_page_count_offset));
+    }
+
+    return frame->memory_regs[mem_idx].cur_page_count;
+}
+
 JitReg
 get_memory_data_reg(JitFrame *frame, uint32 mem_idx)
 {
    JitCompContext *cc = frame->cc;
-    JitReg module_inst_reg = get_module_inst_reg(frame);
+    JitReg module_inst_reg;
    uint32 memory_data_offset;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    JitReg memory_inst_reg;
+    bool is_shared;
+#endif
+
+    if (frame->memory_regs[mem_idx].memory_data)
+        return frame->memory_regs[mem_idx].memory_data;
+
+    frame->memory_regs[mem_idx].memory_data =
+        cc->memory_regs[mem_idx].memory_data;

    bh_assert(mem_idx == 0);
 #if WASM_ENABLE_SHARED_MEMORY != 0
-    uint32 memories_offset = (uint32)offsetof(WASMModuleInstance, memories);
-    JitReg memories_addr = jit_cc_new_reg_ptr(cc);
-    JitReg memories_0_addr = jit_cc_new_reg_ptr(cc);
-    memory_data_offset = (uint32)offsetof(WASMMemoryInstance, memory_data);
-    if (!frame->memory_regs[mem_idx].memory_data) {
-        frame->memory_regs[mem_idx].memory_data =
-            cc->memory_regs[mem_idx].memory_data;
-        /* module_inst->memories */
-        GEN_INSN(LDPTR, memories_addr, module_inst_reg,
-                 NEW_CONST(I32, memories_offset));
-        /* module_inst->memories[0] */
-        GEN_INSN(LDPTR, memories_0_addr, memories_addr, NEW_CONST(I32, 0));
-        /* memories[0]->memory_data */
+    is_shared = is_shared_memory(cc->cur_wasm_module, mem_idx);
+    if (is_shared) {
+        memory_inst_reg = get_memory_inst_reg(frame, mem_idx);
+        memory_data_offset = (uint32)offsetof(WASMMemoryInstance, memory_data);
+        /* memories[mem_idx]->memory_data */
        GEN_INSN(LDPTR, frame->memory_regs[mem_idx].memory_data,
-                 memories_0_addr, NEW_CONST(I32, memory_data_offset));
+                 memory_inst_reg, NEW_CONST(I32, memory_data_offset));
    }
-#else
-    memory_data_offset =
-        (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-        + (uint32)offsetof(WASMMemoryInstance, memory_data);
-    if (!frame->memory_regs[mem_idx].memory_data) {
-        frame->memory_regs[mem_idx].memory_data =
-            cc->memory_regs[mem_idx].memory_data;
+    else
+#endif
+    {
+        module_inst_reg = get_module_inst_reg(frame);
+        memory_data_offset =
+            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
+            + (uint32)offsetof(WASMMemoryInstance, memory_data);
        GEN_INSN(LDPTR, frame->memory_regs[mem_idx].memory_data,
                 module_inst_reg, NEW_CONST(I32, memory_data_offset));
    }
-#endif
    return frame->memory_regs[mem_idx].memory_data;
 }

@ -261,16 +368,37 @@ JitReg
 get_memory_data_end_reg(JitFrame *frame, uint32 mem_idx)
 {
    JitCompContext *cc = frame->cc;
-    JitReg module_inst_reg = get_module_inst_reg(frame);
-    uint32 memory_data_end_offset =
-        (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-        + (uint32)offsetof(WASMMemoryInstance, memory_data_end);
+    JitReg module_inst_reg;
+    uint32 memory_data_end_offset;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    JitReg memory_inst_reg;
+    bool is_shared;
+#endif
+
+    if (frame->memory_regs[mem_idx].memory_data_end)
+        return frame->memory_regs[mem_idx].memory_data_end;
+
+    frame->memory_regs[mem_idx].memory_data_end =
+        cc->memory_regs[mem_idx].memory_data_end;

    bh_assert(mem_idx == 0);
-
-    if (!frame->memory_regs[mem_idx].memory_data_end) {
-        frame->memory_regs[mem_idx].memory_data_end =
-            cc->memory_regs[mem_idx].memory_data_end;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    is_shared = is_shared_memory(cc->cur_wasm_module, mem_idx);
+    if (is_shared) {
+        memory_inst_reg = get_memory_inst_reg(frame, mem_idx);
+        memory_data_end_offset =
+            (uint32)offsetof(WASMMemoryInstance, memory_data_end);
+        /* memories[mem_idx]->memory_data_end */
+        GEN_INSN(LDPTR, frame->memory_regs[mem_idx].memory_data_end,
+                 memory_inst_reg, NEW_CONST(I32, memory_data_end_offset));
+    }
+    else
+#endif
+    {
+        module_inst_reg = get_module_inst_reg(frame);
+        memory_data_end_offset =
+            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
+            + (uint32)offsetof(WASMMemoryInstance, memory_data_end);
        GEN_INSN(LDPTR, frame->memory_regs[mem_idx].memory_data_end,
                 module_inst_reg, NEW_CONST(I32, memory_data_end_offset));
    }
@ -281,16 +409,43 @@ JitReg
 get_mem_bound_check_1byte_reg(JitFrame *frame, uint32 mem_idx)
 {
    JitCompContext *cc = frame->cc;
-    JitReg module_inst_reg = get_module_inst_reg(frame);
-    uint32 mem_bound_check_1byte_offset =
-        (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-        + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_1byte);
+    JitReg module_inst_reg;
+    uint32 mem_bound_check_1byte_offset;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    JitReg memory_inst_reg;
+    bool is_shared;
+#endif
+
+    if (frame->memory_regs[mem_idx].mem_bound_check_1byte)
+        return frame->memory_regs[mem_idx].mem_bound_check_1byte;
+
+    frame->memory_regs[mem_idx].mem_bound_check_1byte =
+        cc->memory_regs[mem_idx].mem_bound_check_1byte;

    bh_assert(mem_idx == 0);

-    if (!frame->memory_regs[mem_idx].mem_bound_check_1byte) {
-        frame->memory_regs[mem_idx].mem_bound_check_1byte =
-            cc->memory_regs[mem_idx].mem_bound_check_1byte;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    is_shared = is_shared_memory(cc->cur_wasm_module, mem_idx);
+    if (is_shared) {
+        memory_inst_reg = get_memory_inst_reg(frame, mem_idx);
+        mem_bound_check_1byte_offset =
+            (uint32)offsetof(WASMMemoryInstance, mem_bound_check_1byte);
+        /* memories[mem_idx]->mem_bound_check_1byte */
+#if UINTPTR_MAX == UINT64_MAX
+        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_1byte,
+                 memory_inst_reg, NEW_CONST(I32, mem_bound_check_1byte_offset));
+#else
+        GEN_INSN(LDI32, frame->memory_regs[mem_idx].mem_bound_check_1byte,
+                 memory_inst_reg, NEW_CONST(I32, mem_bound_check_1byte_offset));
+#endif
+    }
+    else
+#endif
+    {
+        module_inst_reg = get_module_inst_reg(frame);
+        mem_bound_check_1byte_offset =
+            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
+            + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_1byte);
 #if UINTPTR_MAX == UINT64_MAX
        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_1byte,
                 module_inst_reg, NEW_CONST(I32, mem_bound_check_1byte_offset));
@ -306,16 +461,45 @@ JitReg
 get_mem_bound_check_2bytes_reg(JitFrame *frame, uint32 mem_idx)
 {
    JitCompContext *cc = frame->cc;
-    JitReg module_inst_reg = get_module_inst_reg(frame);
-    uint32 mem_bound_check_2bytes_offset =
-        (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-        + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_2bytes);
+    JitReg module_inst_reg;
+    uint32 mem_bound_check_2bytes_offset;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    JitReg memory_inst_reg;
+    bool is_shared;
+#endif
+
+    if (frame->memory_regs[mem_idx].mem_bound_check_2bytes)
+        return frame->memory_regs[mem_idx].mem_bound_check_2bytes;
+
+    frame->memory_regs[mem_idx].mem_bound_check_2bytes =
+        cc->memory_regs[mem_idx].mem_bound_check_2bytes;

    bh_assert(mem_idx == 0);

-    if (!frame->memory_regs[mem_idx].mem_bound_check_2bytes) {
-        frame->memory_regs[mem_idx].mem_bound_check_2bytes =
-            cc->memory_regs[mem_idx].mem_bound_check_2bytes;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    is_shared = is_shared_memory(cc->cur_wasm_module, mem_idx);
+    if (is_shared) {
+        memory_inst_reg = get_memory_inst_reg(frame, mem_idx);
+        mem_bound_check_2bytes_offset =
+            (uint32)offsetof(WASMMemoryInstance, mem_bound_check_2bytes);
+        /* memories[mem_idx]->mem_bound_check_2bytes */
+#if UINTPTR_MAX == UINT64_MAX
+        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_2bytes,
+                 memory_inst_reg,
+                 NEW_CONST(I32, mem_bound_check_2bytes_offset));
+#else
+        GEN_INSN(LDI32, frame->memory_regs[mem_idx].mem_bound_check_2bytes,
+                 memory_inst_reg,
+                 NEW_CONST(I32, mem_bound_check_2bytes_offset));
+#endif
+    }
+    else
+#endif
+    {
+        module_inst_reg = get_module_inst_reg(frame);
+        mem_bound_check_2bytes_offset =
+            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
+            + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_2bytes);
 #if UINTPTR_MAX == UINT64_MAX
        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_2bytes,
                 module_inst_reg,
@ -333,16 +517,45 @@ JitReg
 get_mem_bound_check_4bytes_reg(JitFrame *frame, uint32 mem_idx)
 {
    JitCompContext *cc = frame->cc;
-    JitReg module_inst_reg = get_module_inst_reg(frame);
-    uint32 mem_bound_check_4bytes_offset =
-        (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-        + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_4bytes);
+    JitReg module_inst_reg;
+    uint32 mem_bound_check_4bytes_offset;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    JitReg memory_inst_reg;
+    bool is_shared;
+#endif
+
+    if (frame->memory_regs[mem_idx].mem_bound_check_4bytes)
+        return frame->memory_regs[mem_idx].mem_bound_check_4bytes;
+
+    frame->memory_regs[mem_idx].mem_bound_check_4bytes =
+        cc->memory_regs[mem_idx].mem_bound_check_4bytes;

    bh_assert(mem_idx == 0);

-    if (!frame->memory_regs[mem_idx].mem_bound_check_4bytes) {
-        frame->memory_regs[mem_idx].mem_bound_check_4bytes =
-            cc->memory_regs[mem_idx].mem_bound_check_4bytes;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    is_shared = is_shared_memory(cc->cur_wasm_module, mem_idx);
+    if (is_shared) {
+        memory_inst_reg = get_memory_inst_reg(frame, mem_idx);
+        mem_bound_check_4bytes_offset =
+            (uint32)offsetof(WASMMemoryInstance, mem_bound_check_4bytes);
+        /* memories[mem_idx]->mem_bound_check_4bytes */
+#if UINTPTR_MAX == UINT64_MAX
+        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_4bytes,
+                 memory_inst_reg,
+                 NEW_CONST(I32, mem_bound_check_4bytes_offset));
+#else
+        GEN_INSN(LDI32, frame->memory_regs[mem_idx].mem_bound_check_4bytes,
+                 memory_inst_reg,
+                 NEW_CONST(I32, mem_bound_check_4bytes_offset));
+#endif
+    }
+    else
+#endif
+    {
+        module_inst_reg = get_module_inst_reg(frame);
+        mem_bound_check_4bytes_offset =
+            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
+            + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_4bytes);
 #if UINTPTR_MAX == UINT64_MAX
        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_4bytes,
                 module_inst_reg,
@ -360,16 +573,45 @@ JitReg
 get_mem_bound_check_8bytes_reg(JitFrame *frame, uint32 mem_idx)
 {
    JitCompContext *cc = frame->cc;
-    JitReg module_inst_reg = get_module_inst_reg(frame);
-    uint32 mem_bound_check_8bytes_offset =
-        (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-        + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_8bytes);
+    JitReg module_inst_reg;
+    uint32 mem_bound_check_8bytes_offset;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    JitReg memory_inst_reg;
+    bool is_shared;
+#endif
+
+    if (frame->memory_regs[mem_idx].mem_bound_check_8bytes)
+        return frame->memory_regs[mem_idx].mem_bound_check_8bytes;
+
+    frame->memory_regs[mem_idx].mem_bound_check_8bytes =
+        cc->memory_regs[mem_idx].mem_bound_check_8bytes;

    bh_assert(mem_idx == 0);

-    if (!frame->memory_regs[mem_idx].mem_bound_check_8bytes) {
-        frame->memory_regs[mem_idx].mem_bound_check_8bytes =
-            cc->memory_regs[mem_idx].mem_bound_check_8bytes;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    is_shared = is_shared_memory(cc->cur_wasm_module, mem_idx);
+    if (is_shared) {
+        memory_inst_reg = get_memory_inst_reg(frame, mem_idx);
+        mem_bound_check_8bytes_offset =
+            (uint32)offsetof(WASMMemoryInstance, mem_bound_check_8bytes);
+        /* memories[mem_idx]->mem_bound_check_8bytes */
+#if UINTPTR_MAX == UINT64_MAX
+        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_8bytes,
+                 memory_inst_reg,
+                 NEW_CONST(I32, mem_bound_check_8bytes_offset));
+#else
+        GEN_INSN(LDI32, frame->memory_regs[mem_idx].mem_bound_check_8bytes,
+                 memory_inst_reg,
+                 NEW_CONST(I32, mem_bound_check_8bytes_offset));
+#endif
+    }
+    else
+#endif
+    {
+        module_inst_reg = get_module_inst_reg(frame);
+        mem_bound_check_8bytes_offset =
+            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
+            + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_8bytes);
 #if UINTPTR_MAX == UINT64_MAX
        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_8bytes,
                 module_inst_reg,
@ -387,16 +629,45 @@ JitReg
 get_mem_bound_check_16bytes_reg(JitFrame *frame, uint32 mem_idx)
 {
    JitCompContext *cc = frame->cc;
-    JitReg module_inst_reg = get_module_inst_reg(frame);
-    uint32 mem_bound_check_16bytes_offset =
-        (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
-        + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_16bytes);
+    JitReg module_inst_reg;
+    uint32 mem_bound_check_16bytes_offset;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    JitReg memory_inst_reg;
+    bool is_shared;
+#endif
+
+    if (frame->memory_regs[mem_idx].mem_bound_check_16bytes)
+        return frame->memory_regs[mem_idx].mem_bound_check_16bytes;
+
+    frame->memory_regs[mem_idx].mem_bound_check_16bytes =
+        cc->memory_regs[mem_idx].mem_bound_check_16bytes;

    bh_assert(mem_idx == 0);

-    if (!frame->memory_regs[mem_idx].mem_bound_check_16bytes) {
-        frame->memory_regs[mem_idx].mem_bound_check_16bytes =
-            cc->memory_regs[mem_idx].mem_bound_check_16bytes;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    is_shared = is_shared_memory(cc->cur_wasm_module, mem_idx);
+    if (is_shared) {
+        memory_inst_reg = get_memory_inst_reg(frame, mem_idx);
+        mem_bound_check_16bytes_offset =
+            (uint32)offsetof(WASMMemoryInstance, mem_bound_check_16bytes);
+        /* memories[mem_idx]->mem_bound_check_16bytes */
+#if UINTPTR_MAX == UINT64_MAX
+        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_16bytes,
+                 memory_inst_reg,
+                 NEW_CONST(I32, mem_bound_check_16bytes_offset));
+#else
+        GEN_INSN(LDI32, frame->memory_regs[mem_idx].mem_bound_check_16bytes,
+                 memory_inst_reg,
+                 NEW_CONST(I32, mem_bound_check_16bytes_offset));
+#endif
+    }
+    else
+#endif
+    {
+        module_inst_reg = get_module_inst_reg(frame);
+        mem_bound_check_16bytes_offset =
+            (uint32)offsetof(WASMModuleInstance, global_table_data.bytes)
+            + (uint32)offsetof(WASMMemoryInstance, mem_bound_check_16bytes);
 #if UINTPTR_MAX == UINT64_MAX
        GEN_INSN(LDI64, frame->memory_regs[mem_idx].mem_bound_check_16bytes,
                 module_inst_reg,
@ -462,6 +733,8 @@ clear_fixed_virtual_regs(JitFrame *frame)

    count = module->import_memory_count + module->memory_count;
    for (i = 0; i < count; i++) {
+        frame->memory_regs[i].memory_inst = 0;
+        frame->memory_regs[i].cur_page_count = 0;
        frame->memory_regs[i].memory_data = 0;
        frame->memory_regs[i].memory_data_end = 0;
        frame->memory_regs[i].mem_bound_check_1byte = 0;
@ -486,6 +759,7 @@ clear_memory_regs(JitFrame *frame)

    count = module->import_memory_count + module->memory_count;
    for (i = 0; i < count; i++) {
+        frame->memory_regs[i].cur_page_count = 0;
        frame->memory_regs[i].memory_data = 0;
        frame->memory_regs[i].memory_data_end = 0;
        frame->memory_regs[i].mem_bound_check_1byte = 0;
@ -654,6 +928,8 @@ create_fixed_virtual_regs(JitCompContext *cc)
        }

        for (i = 0; i < count; i++) {
+            cc->memory_regs[i].memory_inst = jit_cc_new_reg_ptr(cc);
+            cc->memory_regs[i].cur_page_count = jit_cc_new_reg_I32(cc);
            cc->memory_regs[i].memory_data = jit_cc_new_reg_ptr(cc);
            cc->memory_regs[i].memory_data_end = jit_cc_new_reg_ptr(cc);
            cc->memory_regs[i].mem_bound_check_1byte = jit_cc_new_reg_ptr(cc);
--- a/core/iwasm/fast-jit/jit_frontend.h
+++ b/core/iwasm/fast-jit/jit_frontend.h
@ -192,6 +192,12 @@ get_aux_stack_bound_reg(JitFrame *frame);
 JitReg
 get_aux_stack_bottom_reg(JitFrame *frame);

+JitReg
+get_memory_inst_reg(JitFrame *frame, uint32 mem_idx);
+
+JitReg
+get_cur_page_count_reg(JitFrame *frame, uint32 mem_idx);
+
 JitReg
 get_memory_data_reg(JitFrame *frame, uint32 mem_idx);

--- a/core/iwasm/fast-jit/jit_ir.h
+++ b/core/iwasm/fast-jit/jit_ir.h
@ -866,6 +866,8 @@ typedef struct JitValueSlot {
 typedef struct JitMemRegs {
    /* The following registers should be re-loaded after
       memory.grow, callbc and callnative */
+    JitReg memory_inst;
+    JitReg cur_page_count;
    JitReg memory_data;
    JitReg memory_data_end;
    JitReg mem_bound_check_1byte;