Add Atomics support to Multi-Memory Lowering Pass (#5339)

This PR adds support for Atomic instructions in the multi-memory lowering pass. Also includes optional bounds checks per the wasm spec guidelines, (visitAtomicRMW, visitAtomicCmpxchg, visitAtomicWait, visitAtomicNotify).

Note: The latter two instructions, memory.atomic.wait and memory.atomic.notify, have browser engine implementations that predate the still-in-progress threads spec. And whether or not atomic.notify should trap for out-of-bounds addresses remains an open issue. For now, this PR is using the same semantics as v8, which is to bounds check all Atomic instructions the same way and trap for out-of-bounds.

2 files changed, 226 insertions, 2 deletions

diff --git a/src/passes/MultiMemoryLowering.cpp b/src/passes/MultiMemoryLowering.cpp
index 452aa21af..0b066c39b 100644
--- a/src/passes/MultiMemoryLowering.cpp
+++ b/src/passes/MultiMemoryLowering.cpp
@@ -31,6 +31,14 @@
 // would not trap. In theory we could compute like the spec, by expanding the
 // i32s to i64s and adding there (where we won't overflow), but we don't have
 // i128s to handle i64 overflow.
+//
+// The Atomic instructions memory.atomic.wait and memory.atomic.notify, have
+// browser engine implementations that predate the still-in-progress threads
+// spec (https://github.com/WebAssembly/threads). And whether or not
+// atomic.notify should trap for out-of-bounds addresses remains an open issue
+// (https://github.com/WebAssembly/threads/issues/105). For now, we are using
+// the same semantics as v8, which is to bounds check all Atomic instructions
+// the same way and trap for out-of-bounds.
 
 #include "ir/module-utils.h"
 #include "ir/names.h"
@@ -193,6 +201,39 @@ struct MultiMemoryLowering : public Pass {
       curr->ptr = getPtr(curr, getFunction(), curr->getMemBytes());
       setMemory(curr);
     }
+
+    void visitAtomicRMW(AtomicRMW* curr) {
+      curr->ptr = getPtr(curr, getFunction(), curr->bytes);
+      setMemory(curr);
+    }
+
+    void visitAtomicCmpxchg(AtomicCmpxchg* curr) {
+      curr->ptr = getPtr(curr, getFunction(), curr->bytes);
+      setMemory(curr);
+    }
+
+    void visitAtomicWait(AtomicWait* curr) {
+      Index bytes;
+      switch (curr->expectedType.getBasic()) {
+        case Type::i32: {
+          bytes = 4;
+          break;
+        }
+        case Type::i64: {
+          bytes = 8;
+          break;
+        }
+        default:
+          WASM_UNREACHABLE("unexpected type");
+      }
+      curr->ptr = getPtr(curr, getFunction(), bytes);
+      setMemory(curr);
+    }
+
+    void visitAtomicNotify(AtomicNotify* curr) {
+      curr->ptr = getPtr(curr, getFunction(), Index(4));
+      setMemory(curr);
+    }
   };
 
   void run(Module* module) override {
diff --git a/test/lit/passes/multi-memory-lowering.wast b/test/lit/passes/multi-memory-lowering.wast
index dba7260a0..6cac24ee0 100644
--- a/test/lit/passes/multi-memory-lowering.wast
+++ b/test/lit/passes/multi-memory-lowering.wast
@@ -1,6 +1,6 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
-;; RUN: wasm-opt %s --enable-multi-memories --multi-memory-lowering --enable-bulk-memory --enable-extended-const --enable-simd -S -o - | filecheck %s
-;; RUN: wasm-opt %s --enable-multi-memories --multi-memory-lowering-with-bounds-checks --enable-bulk-memory --enable-extended-const --enable-simd -S -o - | filecheck %s --check-prefix BOUNDS
+;; RUN: wasm-opt %s --enable-multi-memories --multi-memory-lowering --enable-bulk-memory --enable-extended-const --enable-simd --enable-threads -S -o - | filecheck %s
+;; RUN: wasm-opt %s --enable-multi-memories --multi-memory-lowering-with-bounds-checks --enable-bulk-memory --enable-extended-const --enable-simd --enable-threads -S -o - | filecheck %s --check-prefix BOUNDS
 
 (module
   (memory $memory1 1)
@@ -18,6 +18,8 @@
 
   ;; CHECK:      (type $i32_v128_=>_v128 (func (param i32 v128) (result v128)))
 
+  ;; CHECK:      (type $i32_i64_=>_none (func (param i32 i64)))
+
   ;; CHECK:      (global $memory2_byte_offset (mut i32) (i32.const 65536))
 
   ;; CHECK:      (global $memory3_byte_offset (mut i32) (i32.const 196608))
@@ -64,6 +66,8 @@
 
   ;; BOUNDS:      (type $i32_v128_=>_v128 (func (param i32 v128) (result v128)))
 
+  ;; BOUNDS:      (type $i32_i64_=>_none (func (param i32 i64)))
+
   ;; BOUNDS:      (global $memory2_byte_offset (mut i32) (i32.const 65536))
 
   ;; BOUNDS:      (global $memory3_byte_offset (mut i32) (i32.const 196608))
@@ -445,6 +449,185 @@
    (local.get $0)
   )
  )
+
+  ;; CHECK:      (func $atomics (param $0 i32) (param $1 i64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.atomic.rmw.add offset=4
+  ;; CHECK-NEXT:    (i32.add
+  ;; CHECK-NEXT:     (global.get $memory3_byte_offset)
+  ;; CHECK-NEXT:     (local.get $0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.get $0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i64.atomic.rmw32.cmpxchg_u offset=48
+  ;; CHECK-NEXT:    (i32.add
+  ;; CHECK-NEXT:     (global.get $memory2_byte_offset)
+  ;; CHECK-NEXT:     (local.get $0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.get $1)
+  ;; CHECK-NEXT:    (local.get $1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (memory.atomic.wait32 offset=16
+  ;; CHECK-NEXT:    (local.get $0)
+  ;; CHECK-NEXT:    (local.get $0)
+  ;; CHECK-NEXT:    (local.get $1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (memory.atomic.notify offset=24
+  ;; CHECK-NEXT:    (i32.add
+  ;; CHECK-NEXT:     (global.get $memory2_byte_offset)
+  ;; CHECK-NEXT:     (local.get $0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.get $0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; BOUNDS:      (func $atomics (param $0 i32) (param $1 i64)
+  ;; BOUNDS-NEXT:  (local $2 i32)
+  ;; BOUNDS-NEXT:  (local $3 i32)
+  ;; BOUNDS-NEXT:  (local $4 i32)
+  ;; BOUNDS-NEXT:  (local $5 i32)
+  ;; BOUNDS-NEXT:  (drop
+  ;; BOUNDS-NEXT:   (i32.atomic.rmw.add offset=4
+  ;; BOUNDS-NEXT:    (block (result i32)
+  ;; BOUNDS-NEXT:     (local.set $2
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (global.get $memory3_byte_offset)
+  ;; BOUNDS-NEXT:       (local.get $0)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (if
+  ;; BOUNDS-NEXT:      (i32.gt_u
+  ;; BOUNDS-NEXT:       (i32.add
+  ;; BOUNDS-NEXT:        (i32.add
+  ;; BOUNDS-NEXT:         (local.get $2)
+  ;; BOUNDS-NEXT:         (i32.const 4)
+  ;; BOUNDS-NEXT:        )
+  ;; BOUNDS-NEXT:        (i32.const 4)
+  ;; BOUNDS-NEXT:       )
+  ;; BOUNDS-NEXT:       (call $memory3_size)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (unreachable)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (local.get $2)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $0)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT:  (drop
+  ;; BOUNDS-NEXT:   (i64.atomic.rmw32.cmpxchg_u offset=48
+  ;; BOUNDS-NEXT:    (block (result i32)
+  ;; BOUNDS-NEXT:     (local.set $3
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (global.get $memory2_byte_offset)
+  ;; BOUNDS-NEXT:       (local.get $0)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (if
+  ;; BOUNDS-NEXT:      (i32.gt_u
+  ;; BOUNDS-NEXT:       (i32.add
+  ;; BOUNDS-NEXT:        (i32.add
+  ;; BOUNDS-NEXT:         (local.get $3)
+  ;; BOUNDS-NEXT:         (i32.const 48)
+  ;; BOUNDS-NEXT:        )
+  ;; BOUNDS-NEXT:        (i32.const 4)
+  ;; BOUNDS-NEXT:       )
+  ;; BOUNDS-NEXT:       (call $memory2_size)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (unreachable)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (local.get $3)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $1)
+  ;; BOUNDS-NEXT:    (local.get $1)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT:  (drop
+  ;; BOUNDS-NEXT:   (memory.atomic.wait32 offset=16
+  ;; BOUNDS-NEXT:    (block (result i32)
+  ;; BOUNDS-NEXT:     (local.set $4
+  ;; BOUNDS-NEXT:      (local.get $0)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (if
+  ;; BOUNDS-NEXT:      (i32.gt_u
+  ;; BOUNDS-NEXT:       (i32.add
+  ;; BOUNDS-NEXT:        (i32.add
+  ;; BOUNDS-NEXT:         (local.get $4)
+  ;; BOUNDS-NEXT:         (i32.const 16)
+  ;; BOUNDS-NEXT:        )
+  ;; BOUNDS-NEXT:        (i32.const 4)
+  ;; BOUNDS-NEXT:       )
+  ;; BOUNDS-NEXT:       (call $memory1_size)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (unreachable)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (local.get $4)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $0)
+  ;; BOUNDS-NEXT:    (local.get $1)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT:  (drop
+  ;; BOUNDS-NEXT:   (memory.atomic.notify offset=24
+  ;; BOUNDS-NEXT:    (block (result i32)
+  ;; BOUNDS-NEXT:     (local.set $5
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (global.get $memory2_byte_offset)
+  ;; BOUNDS-NEXT:       (local.get $0)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (if
+  ;; BOUNDS-NEXT:      (i32.gt_u
+  ;; BOUNDS-NEXT:       (i32.add
+  ;; BOUNDS-NEXT:        (i32.add
+  ;; BOUNDS-NEXT:         (local.get $5)
+  ;; BOUNDS-NEXT:         (i32.const 24)
+  ;; BOUNDS-NEXT:        )
+  ;; BOUNDS-NEXT:        (i32.const 4)
+  ;; BOUNDS-NEXT:       )
+  ;; BOUNDS-NEXT:       (call $memory2_size)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (unreachable)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (local.get $5)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $0)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT: )
+  (func $atomics (param $0 i32) (param $1 i64)
+	(drop
+   (i32.atomic.rmw.add $memory3 offset=4
+    (local.get $0)
+    (local.get $0)
+   )
+  )
+(drop
+   (i64.atomic.rmw32.cmpxchg_u $memory2 offset=48
+    (local.get $0)
+    (local.get $1)
+    (local.get $1)
+   )
+  )
+(drop
+   (memory.atomic.wait32 $memory1 offset=16
+    (local.get $0)
+    (local.get $0)
+    (local.get $1)
+   )
+  )
+(drop
+   (memory.atomic.notify $memory2 offset=24
+    (local.get $0)
+    (local.get $0)
+   )
+  )
+  )
 )
 
 ;; CHECK:      (func $memory1_size (result i32)