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.

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);

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_data_reg(JitFrame *frame, uint32 mem_idx)
+get_memory_inst_reg(JitFrame *frame, uint32 mem_idx)
 {
     JitCompContext *cc = frame->cc;
     JitReg module_inst_reg = get_module_inst_reg(frame);
-    uint32 memory_data_offset;
+    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
-    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;
+    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[0] */
-        GEN_INSN(LDPTR, memories_0_addr, memories_addr, NEW_CONST(I32, 0));
-        /* memories[0]->memory_data */
+        /* 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;
+    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
+    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
 
-    bh_assert(mem_idx == 0);
+    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;
 
-    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;
+    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);
+        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);

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);
 

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;