Add relaxed atomics support for stores (#8183)

Part of #8165. The unit tests are written manually for now, but I will
compact the tests later once I have a script to generate them. C and JS
apis will be updated later with all atomic instructions together.

Author: stevenfontanella

src/binaryen-c.cpp

@@ -1371,7 +1371,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);
     if (curr->offset) {
       o << " offset=" << curr->offset;
     }

src/passes/SafeHeap.cpp

@@ -63,15 +63,22 @@ static Name getLoadName(Load* curr) {
 }
 
 static Name getStoreName(Store* curr) {
-  std::string ret = "SAFE_HEAP_STORE_";
-  ret += curr->valueType.toString();
-  ret += "_" + std::to_string(curr->bytes) + "_";
-  if (curr->isAtomic()) {
-    ret += "A";
-  } else {
-    ret += std::to_string(curr->align);
+  std::vector<std::string> parts{curr->valueType.toString(),
+                                 std::to_string(curr->bytes)};
+  switch (curr->order) {
+    case MemoryOrder::Unordered: {
+      parts.push_back(std::to_string(curr->align));
+      break;
+    }
+    case MemoryOrder::SeqCst: {
+      parts.push_back("SC");
+      break;
+    }
+    case MemoryOrder::AcqRel: {
+      parts.push_back("AR");
+    }
   }
-  return ret;
+  return "SAFE_HEAP_STORE_" + String::join(parts, "_");
 }
 
 struct AccessInstrumenter : public WalkerPass<PostWalker<AccessInstrumenter>> {
@@ -287,10 +294,9 @@ struct SafeHeap : public Pass {
           if (align > bytes) {
             continue;
           }
-          for (auto isAtomic : {true, false}) {
-            store.order =
-              isAtomic ? MemoryOrder::SeqCst : MemoryOrder::Unordered;
-            if (isAtomic &&
+          for (auto memoryOrder : memoryOrdersToGenerate) {
+            store.order = memoryOrder;
+            if (store.isAtomic() &&
                 !isPossibleAtomicOperation(
                   align, bytes, module->memories[0]->shared, valueType)) {
               continue;

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,8 @@ 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

@@ -3667,31 +3667,38 @@ Result<> WasmBinaryReader::readInst() {
         }
         case BinaryConsts::I32AtomicStore8: {
           auto [mem, align, offset, memoryOrder] = getAtomicMemarg();
-          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] = getAtomicMemarg();
-          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] = getAtomicMemarg();
-          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] = getAtomicMemarg();
-          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] = getAtomicMemarg();
-          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] = getAtomicMemarg();
-          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] = getAtomicMemarg();
-          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{};
 }