Refine lock/unlock shared memory (#2682)

Split memory instance's field `uint32 ref_count` into `bool is_shared_memory`
and `uint16 ref_count`, and lock the memory only when `is_shared_memory`
flag is true, no need to acquire a lock for non-shared memory when shared
memory feature is enabled.

core/iwasm/aot/aot_runtime.c

@@ -595,6 +595,7 @@ memory_instantiate(AOTModuleInstance *module_inst, AOTModuleInstance *parent,
 
 #if WASM_ENABLE_SHARED_MEMORY != 0
     if (is_shared_memory) {
+        memory_inst->is_shared_memory = true;
         memory_inst->ref_count = 1;
     }
 #endif

core/iwasm/common/wasm_memory.c

@@ -547,8 +547,6 @@ wasm_check_app_addr_and_convert(WASMModuleInstance *module_inst, bool is_str,
         return false;
     }
 
-    SHARED_MEMORY_LOCK(memory_inst);
-
     native_addr = memory_inst->memory_data + app_buf_addr;
 
     bounds_checks = is_bounds_checks_enabled((wasm_module_inst_t)module_inst);
@@ -563,6 +561,8 @@ wasm_check_app_addr_and_convert(WASMModuleInstance *module_inst, bool is_str,
     /* No need to check the app_offset and buf_size if memory access
        boundary check with hardware trap is enabled */
 #ifndef OS_ENABLE_HW_BOUND_CHECK
+    SHARED_MEMORY_LOCK(memory_inst);
+
     if (app_buf_addr >= memory_inst->memory_data_size) {
         goto fail;
     }
@@ -583,9 +583,9 @@ wasm_check_app_addr_and_convert(WASMModuleInstance *module_inst, bool is_str,
         if (str == str_end)
             goto fail;
     }
-#endif
 
     SHARED_MEMORY_UNLOCK(memory_inst);
+#endif
 
 success:
     *p_native_addr = (void *)native_addr;

core/iwasm/common/wasm_shared_memory.c

@@ -18,7 +18,7 @@
  * - If you care performance, it's better to make the interpreters
  *   use atomic ops.
  */
-static korp_mutex _shared_memory_lock;
+korp_mutex g_shared_memory_lock;
 
 /* clang-format off */
 enum {
@@ -55,13 +55,13 @@ destroy_wait_info(void *wait_info);
 bool
 wasm_shared_memory_init()
 {
-    if (os_mutex_init(&_shared_memory_lock) != 0)
+    if (os_mutex_init(&g_shared_memory_lock) != 0)
         return false;
     /* wait map not exists, create new map */
     if (!(wait_map = bh_hash_map_create(32, true, (HashFunc)wait_address_hash,
                                         (KeyEqualFunc)wait_address_equal, NULL,
                                         destroy_wait_info))) {
-        os_mutex_destroy(&_shared_memory_lock);
+        os_mutex_destroy(&g_shared_memory_lock);
         return false;
     }
     return true;
@@ -71,79 +71,47 @@ void
 wasm_shared_memory_destroy()
 {
     bh_hash_map_destroy(wait_map);
-    os_mutex_destroy(&_shared_memory_lock);
+    os_mutex_destroy(&g_shared_memory_lock);
 }
 
-uint32
+uint16
 shared_memory_inc_reference(WASMMemoryInstance *memory)
 {
     bh_assert(shared_memory_is_shared(memory));
-    uint32 old;
-#if BH_ATOMIC_32_IS_ATOMIC == 0
-    os_mutex_lock(&_shared_memory_lock);
+    uint16 old;
+#if BH_ATOMIC_16_IS_ATOMIC == 0
+    os_mutex_lock(&g_shared_memory_lock);
 #endif
-    old = BH_ATOMIC_32_FETCH_ADD(memory->ref_count, 1);
-#if BH_ATOMIC_32_IS_ATOMIC == 0
-    os_mutex_unlock(&_shared_memory_lock);
+    old = BH_ATOMIC_16_FETCH_ADD(memory->ref_count, 1);
+#if BH_ATOMIC_16_IS_ATOMIC == 0
+    os_mutex_unlock(&g_shared_memory_lock);
 #endif
     bh_assert(old >= 1);
-    bh_assert(old < UINT32_MAX);
+    bh_assert(old < UINT16_MAX);
     return old + 1;
 }
 
-uint32
+uint16
 shared_memory_dec_reference(WASMMemoryInstance *memory)
 {
     bh_assert(shared_memory_is_shared(memory));
-    uint32 old;
-#if BH_ATOMIC_32_IS_ATOMIC == 0
-    os_mutex_lock(&_shared_memory_lock);
+    uint16 old;
+#if BH_ATOMIC_16_IS_ATOMIC == 0
+    os_mutex_lock(&g_shared_memory_lock);
 #endif
-    old = BH_ATOMIC_32_FETCH_SUB(memory->ref_count, 1);
-#if BH_ATOMIC_32_IS_ATOMIC == 0
-    os_mutex_unlock(&_shared_memory_lock);
+    old = BH_ATOMIC_16_FETCH_SUB(memory->ref_count, 1);
+#if BH_ATOMIC_16_IS_ATOMIC == 0
+    os_mutex_unlock(&g_shared_memory_lock);
 #endif
     bh_assert(old > 0);
     return old - 1;
 }
 
-bool
-shared_memory_is_shared(WASMMemoryInstance *memory)
-{
-    uint32 old;
-#if BH_ATOMIC_32_IS_ATOMIC == 0
-    os_mutex_lock(&_shared_memory_lock);
-#endif
-    old = BH_ATOMIC_32_LOAD(memory->ref_count);
-#if BH_ATOMIC_32_IS_ATOMIC == 0
-    os_mutex_unlock(&_shared_memory_lock);
-#endif
-    return old > 0;
-}
-
 static korp_mutex *
 shared_memory_get_lock_pointer(WASMMemoryInstance *memory)
 {
     bh_assert(memory != NULL);
-    return &_shared_memory_lock;
-}
-
-void
-shared_memory_lock(WASMMemoryInstance *memory)
-{
-    /*
-     * Note: exception logic is currently abusing this lock.
-     * cf. https://github.com/bytecodealliance/wasm-micro-runtime/issues/2407
-     */
-    bh_assert(memory != NULL);
-    os_mutex_lock(&_shared_memory_lock);
-}
-
-void
-shared_memory_unlock(WASMMemoryInstance *memory)
-{
-    bh_assert(memory != NULL);
-    os_mutex_unlock(&_shared_memory_lock);
+    return &g_shared_memory_lock;
 }
 
 /* Atomics wait && notify APIs */

core/iwasm/common/wasm_shared_memory.h

@@ -14,26 +14,39 @@
 extern "C" {
 #endif
 
+extern korp_mutex g_shared_memory_lock;
+
 bool
 wasm_shared_memory_init();
 
 void
 wasm_shared_memory_destroy();
 
-uint32
+uint16
 shared_memory_inc_reference(WASMMemoryInstance *memory);
 
-uint32
+uint16
 shared_memory_dec_reference(WASMMemoryInstance *memory);
 
-bool
-shared_memory_is_shared(WASMMemoryInstance *memory);
-
-void
-shared_memory_lock(WASMMemoryInstance *memory);
-
-void
-shared_memory_unlock(WASMMemoryInstance *memory);
+#define shared_memory_is_shared(memory) memory->is_shared_memory
+
+#define shared_memory_lock(memory)                                            \
+    do {                                                                      \
+        /*                                                                    \
+         * Note: exception logic is currently abusing this lock.              \
+         * cf.                                                                \
+         * https://github.com/bytecodealliance/wasm-micro-runtime/issues/2407 \
+         */                                                                   \
+        bh_assert(memory != NULL);                                            \
+        if (memory->is_shared_memory)                                         \
+            os_mutex_lock(&g_shared_memory_lock);                             \
+    } while (0)
+
+#define shared_memory_unlock(memory)                \
+    do {                                            \
+        if (memory->is_shared_memory)               \
+            os_mutex_unlock(&g_shared_memory_lock); \
+    } while (0)
 
 uint32
 wasm_runtime_atomic_wait(WASMModuleInstanceCommon *module, void *address,

core/iwasm/interpreter/wasm_runtime.c

@@ -374,6 +374,7 @@ memory_instantiate(WASMModuleInstance *module_inst, WASMModuleInstance *parent,
 
 #if WASM_ENABLE_SHARED_MEMORY != 0
     if (is_shared_memory) {
+        memory->is_shared_memory = true;
         memory->ref_count = 1;
     }
 #endif

core/iwasm/interpreter/wasm_runtime.h

@@ -79,8 +79,11 @@ typedef union {
 struct WASMMemoryInstance {
     /* Module type */
     uint32 module_type;
+
+    bool is_shared_memory;
+
     /* Shared memory flag */
-    bh_atomic_32_t ref_count; /* 0: non-shared, > 0: reference count */
+    bh_atomic_16_t ref_count; /* 0: non-shared, > 0: reference count */
 
     /* Number bytes per page */
     uint32 num_bytes_per_page;

core/shared/utils/bh_atomic.h

@@ -45,6 +45,7 @@ extern "C" {
  */
 
 typedef uint32 bh_atomic_32_t;
+typedef uint16 bh_atomic_16_t;
 
 #if defined(__GNUC_PREREQ)
 #if __GNUC_PREREQ(4, 7)
@@ -59,6 +60,7 @@ typedef uint32 bh_atomic_32_t;
 #if defined(CLANG_GCC_HAS_ATOMIC_BUILTIN)
 #define BH_ATOMIC_32_IS_ATOMIC 1
 #define BH_ATOMIC_32_LOAD(v) __atomic_load_n(&(v), __ATOMIC_SEQ_CST)
+#define BH_ATOMIC_32_STORE(v, val) __atomic_store_n(&(v), val, __ATOMIC_SEQ_CST)
 #define BH_ATOMIC_32_FETCH_OR(v, val) \
     __atomic_fetch_or(&(v), (val), __ATOMIC_SEQ_CST)
 #define BH_ATOMIC_32_FETCH_AND(v, val) \
@@ -67,13 +69,33 @@ typedef uint32 bh_atomic_32_t;
     __atomic_fetch_add(&(v), (val), __ATOMIC_SEQ_CST)
 #define BH_ATOMIC_32_FETCH_SUB(v, val) \
     __atomic_fetch_sub(&(v), (val), __ATOMIC_SEQ_CST)
+
+#define BH_ATOMIC_16_IS_ATOMIC 1
+#define BH_ATOMIC_16_LOAD(v) __atomic_load_n(&(v), __ATOMIC_SEQ_CST)
+#define BH_ATOMIC_16_STORE(v, val) __atomic_store_n(&(v), val, __ATOMIC_SEQ_CST)
+#define BH_ATOMIC_16_FETCH_OR(v, val) \
+    __atomic_fetch_or(&(v), (val), __ATOMIC_SEQ_CST)
+#define BH_ATOMIC_16_FETCH_AND(v, val) \
+    __atomic_fetch_and(&(v), (val), __ATOMIC_SEQ_CST)
+#define BH_ATOMIC_16_FETCH_ADD(v, val) \
+    __atomic_fetch_add(&(v), (val), __ATOMIC_SEQ_CST)
+#define BH_ATOMIC_16_FETCH_SUB(v, val) \
+    __atomic_fetch_sub(&(v), (val), __ATOMIC_SEQ_CST)
 #else /* else of defined(CLANG_GCC_HAS_ATOMIC_BUILTIN) */
 #define BH_ATOMIC_32_LOAD(v) (v)
+#define BH_ATOMIC_32_STORE(v, val) (v) = val
 #define BH_ATOMIC_32_FETCH_OR(v, val) nonatomic_32_fetch_or(&(v), val)
 #define BH_ATOMIC_32_FETCH_AND(v, val) nonatomic_32_fetch_and(&(v), val)
 #define BH_ATOMIC_32_FETCH_ADD(v, val) nonatomic_32_fetch_add(&(v), val)
 #define BH_ATOMIC_32_FETCH_SUB(v, val) nonatomic_32_fetch_sub(&(v), val)
 
+#define BH_ATOMIC_16_LOAD(v) (v)
+#define BH_ATOMIC_16_STORE(v) (v) = val
+#define BH_ATOMIC_16_FETCH_OR(v, val) nonatomic_16_fetch_or(&(v), val)
+#define BH_ATOMIC_16_FETCH_AND(v, val) nonatomic_16_fetch_and(&(v), val)
+#define BH_ATOMIC_16_FETCH_ADD(v, val) nonatomic_16_fetch_add(&(v), val)
+#define BH_ATOMIC_16_FETCH_SUB(v, val) nonatomic_16_fetch_sub(&(v), val)
+
 static inline uint32
 nonatomic_32_fetch_or(bh_atomic_32_t *p, uint32 val)
 {
@@ -106,6 +128,38 @@ nonatomic_32_fetch_sub(bh_atomic_32_t *p, uint32 val)
     return old;
 }
 
+static inline uint16
+nonatomic_16_fetch_or(bh_atomic_16_t *p, uint16 val)
+{
+    uint16 old = *p;
+    *p |= val;
+    return old;
+}
+
+static inline uint16
+nonatomic_16_fetch_and(bh_atomic_16_t *p, uint16 val)
+{
+    uint16 old = *p;
+    *p &= val;
+    return old;
+}
+
+static inline uint16
+nonatomic_16_fetch_add(bh_atomic_16_t *p, uint16 val)
+{
+    uint16 old = *p;
+    *p += val;
+    return old;
+}
+
+static inline uint16
+nonatomic_16_fetch_sub(bh_atomic_16_t *p, uint16 val)
+{
+    uint16 old = *p;
+    *p -= val;
+    return old;
+}
+
 /* The flag can be defined by the user if the platform
    supports atomic access to uint32 aligned memory. */
 #ifdef WASM_UINT32_IS_ATOMIC
@@ -114,6 +168,12 @@ nonatomic_32_fetch_sub(bh_atomic_32_t *p, uint32 val)
 #define BH_ATOMIC_32_IS_ATOMIC 0
 #endif /* WASM_UINT32_IS_ATOMIC */
 
+#ifdef WASM_UINT16_IS_ATOMIC
+#define BH_ATOMIC_16_IS_ATOMIC 1
+#else /* else of WASM_UINT16_IS_ATOMIC */
+#define BH_ATOMIC_16_IS_ATOMIC 0
+#endif /* WASM_UINT16_IS_ATOMIC */
+
 #endif
 
 #ifdef __cplusplus