Add relaxed atomics support for stores

Author: stevenfontanella

src/binaryen-c.cpp

@@ -1368,7 +1368,8 @@ BinaryenExpressionRef BinaryenAtomicStore(BinaryenModuleRef module,
                        (Expression*)ptr,
                        (Expression*)value,
                        Type(type),
-                       getMemoryName(module, memoryName)));
+                       getMemoryName(module, memoryName),
+                       MemoryOrder::SeqCst));
 }
 BinaryenExpressionRef BinaryenAtomicRMW(BinaryenModuleRef module,
                                         BinaryenOp op,

src/parser/contexts.h

@@ -565,7 +565,8 @@ struct NullInstrParserCtx {
                      int,
                      bool,
                      MemoryIdxT*,
-                     MemargT) {
+                     MemargT,
+                     MemoryOrder) {
     return Ok{};
   }
   Result<> makeAtomicRMW(Index,
@@ -2255,12 +2256,13 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
                      int bytes,
                      bool isAtomic,
                      Name* mem,
-                     Memarg memarg) {
+                     Memarg memarg,
+                     MemoryOrder order) {
     auto m = getMemory(pos, mem);
     CHECK_ERR(m);
     if (isAtomic) {
-      return withLoc(pos,
-                     irBuilder.makeAtomicStore(bytes, memarg.offset, type, *m));
+      return withLoc(
+        pos, irBuilder.makeAtomicStore(bytes, memarg.offset, type, *m, order));
     }
     return withLoc(
       pos, irBuilder.makeStore(bytes, memarg.offset, memarg.align, type, *m));

src/parser/parsers.h

@@ -850,17 +850,6 @@ Result<typename Ctx::MemTypeT> memtypeContinued(Ctx& ctx, Type addressType) {
   return ctx.makeMemType(addressType, *limits, shared);
 }
 
-// memorder ::= '' | 'seqcst' | 'acqrel'
-template<typename Ctx> Result<MemoryOrder> memorder(Ctx& ctx) {
-  if (ctx.in.takeKeyword("seqcst"sv)) {
-    return MemoryOrder::SeqCst;
-  }
-  if (ctx.in.takeKeyword("acqrel"sv)) {
-    return MemoryOrder::AcqRel;
-  }
-  return MemoryOrder::SeqCst;
-}
-
 // memorder ::= 'seqcst' | 'acqrel'
 template<typename Ctx> MaybeResult<MemoryOrder> maybeMemOrder(Ctx& ctx) {
   if (ctx.in.takeKeyword("seqcst"sv)) {
@@ -873,6 +862,13 @@ template<typename Ctx> MaybeResult<MemoryOrder> maybeMemOrder(Ctx& ctx) {
   return {};
 }
 
+// memorder ::= '' | 'seqcst' | 'acqrel'
+template<typename Ctx> Result<MemoryOrder> memorder(Ctx& ctx) {
+  auto order = maybeMemOrder(ctx);
+  CHECK_ERR(order);
+  return order ? *order : MemoryOrder::SeqCst;
+}
+
 // tabletype ::= (limits32 | 'i32' limits32 | 'i64' limit64) reftype
 template<typename Ctx> Result<typename Ctx::TableTypeT> tabletype(Ctx& ctx) {
   Type addressType = Type::i32;
@@ -1755,7 +1751,8 @@ Result<> makeLoad(Ctx& ctx,
   CHECK_ERR(mem);
 
   // We could only parse this when `isAtomic`, but this way gives a clearer
-  // error since `memIdx` can never be mistaken for a `memOrder`.
+  // error when a memorder is given for non-atomic operations
+  // since the next token can never be mistaken for a `memOrder`.
   auto maybeOrder = maybeMemOrder(ctx);
   CHECK_ERR(maybeOrder);
 
@@ -1787,10 +1784,26 @@ Result<> makeStore(Ctx& ctx,
                    bool isAtomic) {
   auto mem = maybeMemidx(ctx);
   CHECK_ERR(mem);
+
+  auto maybeOrder = maybeMemOrder(ctx);
+  CHECK_ERR(maybeOrder);
+
+  if (maybeOrder && !isAtomic) {
+    return Err{"Memory ordering can only be provided for atomic stores."};
+  }
+
   auto arg = memarg(ctx, bytes);
   CHECK_ERR(arg);
-  return ctx.makeStore(
-    pos, annotations, type, bytes, isAtomic, mem.getPtr(), *arg);
+  return ctx.makeStore(pos,
+                       annotations,
+                       type,
+                       bytes,
+                       isAtomic,
+                       mem.getPtr(),
+                       *arg,
+                       maybeOrder ? *maybeOrder
+                       : isAtomic ? MemoryOrder::SeqCst
+                                  : MemoryOrder::Unordered);
 }
 
 template<typename Ctx>

src/passes/Print.cpp

@@ -605,6 +605,7 @@ struct PrintExpressionContents
     }
     restoreNormalColor(o);
     printMemoryName(curr->memory, o, wasm);
+    printMemoryOrder(curr->order, /*prependSpace=*/true, /*appendSpace=*/false);
     if (curr->offset) {
       o << " offset=" << curr->offset;
     }

src/tools/wasm-split/instrumenter.cpp

@@ -151,7 +151,8 @@ void Instrumenter::instrumentFuncs() {
                                   builder.makeConstPtr(0, Type::i32),
                                   builder.makeConst(uint32_t(1)),
                                   Type::i32,
-                                  memoryName),
+                                  memoryName,
+                                  MemoryOrder::SeqCst),
           func->body,
           func->body->type);
         ++funcIdx;

src/wasm-builder.h

@@ -451,9 +451,13 @@ class Builder {
                          Expression* ptr,
                          Expression* value,
                          Type type,
-                         Name memory) {
+                         Name memory,
+                         MemoryOrder order) {
+    assert(order != MemoryOrder::Unordered &&
+           "Atomic stores can't be unordered");
+
     Store* store = makeStore(bytes, offset, bytes, ptr, value, type, memory);
-    store->order = MemoryOrder::SeqCst;
+    store->order = order;
     return store;
   }
   AtomicRMW* makeAtomicRMW(AtomicRMWOp op,

src/wasm-ir-builder.h

@@ -158,7 +158,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     unsigned bytes, Address offset, unsigned align, Type type, Name mem);
   Result<> makeAtomicLoad(
     unsigned bytes, Address offset, Type type, Name mem, MemoryOrder order);
-  Result<> makeAtomicStore(unsigned bytes, Address offset, Type type, Name mem);
+  Result<> makeAtomicStore(unsigned bytes, Address offset, Type type, Name mem, MemoryOrder order);
   Result<> makeAtomicRMW(
     AtomicRMWOp op, unsigned bytes, Address offset, Type type, Name mem);
   Result<>

src/wasm/wasm-binary.cpp

@@ -3685,37 +3685,44 @@ Result<> WasmBinaryReader::readInst() {
         case BinaryConsts::I32AtomicStore8: {
           auto [mem, align, offset, memoryOrder] =
             getMemarg(/*isAtomic=*/true, /*isRMW=*/false);
-          return builder.makeAtomicStore(1, offset, Type::i32, mem);
+          return builder.makeAtomicStore(
+            1, offset, Type::i32, mem, memoryOrder);
         }
         case BinaryConsts::I32AtomicStore16: {
           auto [mem, align, offset, memoryOrder] =
             getMemarg(/*isAtomic=*/true, /*isRMW=*/false);
-          return builder.makeAtomicStore(2, offset, Type::i32, mem);
+          return builder.makeAtomicStore(
+            2, offset, Type::i32, mem, memoryOrder);
         }
         case BinaryConsts::I32AtomicStore: {
           auto [mem, align, offset, memoryOrder] =
             getMemarg(/*isAtomic=*/true, /*isRMW=*/false);
-          return builder.makeAtomicStore(4, offset, Type::i32, mem);
+          return builder.makeAtomicStore(
+            4, offset, Type::i32, mem, memoryOrder);
         }
         case BinaryConsts::I64AtomicStore8: {
           auto [mem, align, offset, memoryOrder] =
             getMemarg(/*isAtomic=*/true, /*isRMW=*/false);
-          return builder.makeAtomicStore(1, offset, Type::i64, mem);
+          return builder.makeAtomicStore(
+            1, offset, Type::i64, mem, memoryOrder);
         }
         case BinaryConsts::I64AtomicStore16: {
           auto [mem, align, offset, memoryOrder] =
             getMemarg(/*isAtomic=*/true, /*isRMW=*/false);
-          return builder.makeAtomicStore(2, offset, Type::i64, mem);
+          return builder.makeAtomicStore(
+            2, offset, Type::i64, mem, memoryOrder);
         }
         case BinaryConsts::I64AtomicStore32: {
           auto [mem, align, offset, memoryOrder] =
             getMemarg(/*isAtomic=*/true, /*isRMW=*/false);
-          return builder.makeAtomicStore(4, offset, Type::i64, mem);
+          return builder.makeAtomicStore(
+            4, offset, Type::i64, mem, memoryOrder);
         }
         case BinaryConsts::I64AtomicStore: {
           auto [mem, align, offset, memoryOrder] =
             getMemarg(/*isAtomic=*/true, /*isRMW=*/false);
-          return builder.makeAtomicStore(8, offset, Type::i64, mem);
+          return builder.makeAtomicStore(
+            8, offset, Type::i64, mem, memoryOrder);
         }
 
 #define RMW(op)                                                                \

src/wasm/wasm-ir-builder.cpp

@@ -1492,15 +1492,14 @@ Result<> IRBuilder::makeAtomicLoad(
   return Ok{};
 }
 
-Result<> IRBuilder::makeAtomicStore(unsigned bytes,
-                                    Address offset,
-                                    Type type,
-                                    Name mem) {
+Result<> IRBuilder::makeAtomicStore(
+  unsigned bytes, Address offset, Type type, Name mem, MemoryOrder order) {
   Store curr;
   curr.memory = mem;
   curr.valueType = type;
   CHECK_ERR(visitStore(&curr));
-  push(builder.makeAtomicStore(bytes, offset, curr.ptr, curr.value, type, mem));
+  push(builder.makeAtomicStore(
+    bytes, offset, curr.ptr, curr.value, type, mem, order));
   return Ok{};
 }
 

test/spec/relaxed-atomics.wast

@@ -2,12 +2,14 @@
  (memory $0 23 256 shared)
 
  (func $acqrel (result i32)
+   (i32.atomic.store acqrel (i32.const 0) (i32.const 0))
    (i32.atomic.load acqrel
     (i32.const 1)
    )
  )
  (func $seqcst (result i32)
-   ;; seqcst may be omitted for atomic loads, it's the default
+   ;; seqcst immediate may be omitted for atomic operations, it's the default
+   (i32.atomic.store seqcst (i32.const 0) (i32.const 0))
    (i32.atomic.load seqcst
     (i32.const 1)
    )
@@ -26,6 +28,38 @@
  "Memory ordering can only be provided for atomic loads."
 )
 
+;; Parses acquire-release immediate
+;; (module
+;;  (memory $0 23 256 shared)
+;;  (func $acqrel
+;;   (i32.atomic.store (i32.const 0) (i32.const 0))
+;;  )
+;; )
+(module binary
+ "\00asm\01\00\00\00"
+ "\01\04\01\60\00\00\03\02\01\00\05\05\01\03\17\80\02" ;; other sections
+ "\0a\0d\01" ;; code section
+   "\0b\00" ;; func $acqrel
+   "\41\00\41\00" ;; (i32.const 0) (i32.const 0)
+   "\fe\17" ;; i32.atomic.store
+   "\22" ;; 2 | (1<<5): Alignment of 2 (32-bit store), with bit 5 set indicating that then next byte is a memory ordering
+   "\01" ;; acqrel ordering
+   "\00" ;; offset
+   "\0b" ;; end
+)
+
+(module binary
+ "\00asm\01\00\00\00"
+ "\01\04\01\60\00\00\03\02\01\00\05\05\01\03\17\80\02" ;; other sections
+ "\0a\0d\01" ;; code section
+   "\0b\00" ;; func $acqrel
+   "\41\00\41\00" ;; (i32.const 0) (i32.const 0)
+   "\fe\17" ;; i32.atomic.store
+   "\22" ;; 2 | (1<<5): Alignment of 2 (32-bit store), with bit 5 set indicating that then next byte is a memory ordering
+   "\00" ;; seqcst ordering
+   "\00" ;; offset
+   "\0b" ;; end
+)
 
 ;; Parses acquire-release immediate
 ;; Equivalent to