Acquire/release memory ordering for loads

Author: stevenfontanella

src/binaryen-c.cpp

@@ -1351,7 +1351,8 @@ BinaryenExpressionRef BinaryenAtomicLoad(BinaryenModuleRef module,
                       offset,
                       (Expression*)ptr,
                       Type(type),
-                      getMemoryName(module, memoryName)));
+                      getMemoryName(module, memoryName),
+                      MemoryOrder::SeqCst));
 }
 BinaryenExpressionRef BinaryenAtomicStore(BinaryenModuleRef module,
                                           uint32_t bytes,

src/parser/contexts.h

@@ -555,7 +555,8 @@ struct NullInstrParserCtx {
                     int,
                     bool,
                     MemoryIdxT*,
-                    MemargT) {
+                    MemargT,
+                    MemoryOrder) {
     return Ok{};
   }
   Result<> makeStore(Index,
@@ -2235,12 +2236,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.makeAtomicLoad(bytes, memarg.offset, type, *m));
+      return withLoc(
+        pos, irBuilder.makeAtomicLoad(bytes, memarg.offset, type, *m, order));
     }
     return withLoc(pos,
                    irBuilder.makeLoad(

src/parser/parsers.h

@@ -49,6 +49,7 @@ template<typename Ctx> Result<typename Ctx::MemTypeT> memtype(Ctx&);
 template<typename Ctx>
 Result<typename Ctx::MemTypeT> memtypeContinued(Ctx&, Type addressType);
 template<typename Ctx> Result<MemoryOrder> memorder(Ctx&);
+template<typename Ctx> MaybeResult<MemoryOrder> maybeMemOrder(Ctx&);
 template<typename Ctx> Result<typename Ctx::TableTypeT> tabletype(Ctx&);
 template<typename Ctx>
 Result<typename Ctx::TableTypeT> tabletypeContinued(Ctx&, Type addressType);
@@ -860,6 +861,18 @@ template<typename Ctx> Result<MemoryOrder> memorder(Ctx& ctx) {
   return MemoryOrder::SeqCst;
 }
 
+// memorder ::= 'seqcst' | 'acqrel'
+template<typename Ctx> MaybeResult<MemoryOrder> maybeMemOrder(Ctx& ctx) {
+  if (ctx.in.takeKeyword("seqcst"sv)) {
+    return MemoryOrder::SeqCst;
+  }
+  if (ctx.in.takeKeyword("acqrel"sv)) {
+    return MemoryOrder::AcqRel;
+  }
+
+  return {};
+}
+
 // tabletype ::= (limits32 | 'i32' limits32 | 'i64' limit64) reftype
 template<typename Ctx> Result<typename Ctx::TableTypeT> tabletype(Ctx& ctx) {
   Type addressType = Type::i32;
@@ -1737,12 +1750,32 @@ Result<> makeLoad(Ctx& ctx,
                   bool signed_,
                   int bytes,
                   bool isAtomic) {
+
   auto mem = maybeMemidx(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`.
+  auto maybeOrder = maybeMemOrder(ctx);
+  CHECK_ERR(maybeOrder);
+
+  if (maybeOrder && !isAtomic) {
+    return Err{"Memory ordering can only be provided for atomic loads."};
+  }
+
   auto arg = memarg(ctx, bytes);
   CHECK_ERR(arg);
-  return ctx.makeLoad(
-    pos, annotations, type, signed_, bytes, isAtomic, mem.getPtr(), *arg);
+  return ctx.makeLoad(pos,
+                      annotations,
+                      type,
+                      signed_,
+                      bytes,
+                      isAtomic,
+                      mem.getPtr(),
+                      *arg,
+                      maybeOrder ? *maybeOrder
+                      : isAtomic ? MemoryOrder::SeqCst
+                                 : MemoryOrder::Unordered);
 }
 
 template<typename Ctx>

src/passes/Print.cpp

@@ -574,6 +574,7 @@ struct PrintExpressionContents
     }
     restoreNormalColor(o);
     printMemoryName(curr->memory, o, wasm);
+    printMemoryOrder(curr->order, /*prependSpace=*/true, /*appendSpace=*/false);
     if (curr->offset) {
       o << " offset=" << curr->offset;
     }
@@ -2323,7 +2324,9 @@ struct PrintExpressionContents
       o << index;
     }
   }
-  void printMemoryOrder(MemoryOrder order) {
+
+  void
+  printMemoryOrder(MemoryOrder order, bool prependSpace, bool appendSpace) {
     switch (order) {
       // Unordered should have a different base instruction, so there is nothing
       // to print. We could be explicit and print seqcst, but we choose not to
@@ -2332,10 +2335,15 @@ struct PrintExpressionContents
       case MemoryOrder::SeqCst:
         break;
       case MemoryOrder::AcqRel:
-        o << "acqrel ";
+        o << (prependSpace ? " " : "") << "acqrel" << (appendSpace ? " " : "");
         break;
     }
   }
+
+  void printMemoryOrder(MemoryOrder order) {
+    printMemoryOrder(order, /*prependSpace=*/false, /*appendSpace=*/true);
+  }
+
   void visitStructGet(StructGet* curr) {
     auto heapType = curr->ref->type.getHeapType();
     const auto& field = heapType.getStruct().fields[curr->index];

src/passes/SafeHeap.cpp

@@ -26,6 +26,7 @@
 #include "ir/import-utils.h"
 #include "ir/load-utils.h"
 #include "pass.h"
+#include "support/string.h"
 #include "wasm-builder.h"
 #include "wasm.h"
 
@@ -37,18 +38,28 @@ static const Name SEGFAULT_IMPORT("segfault");
 static const Name ALIGNFAULT_IMPORT("alignfault");
 
 static Name getLoadName(Load* curr) {
-  std::string ret = "SAFE_HEAP_LOAD_";
-  ret += curr->type.toString();
-  ret += "_" + std::to_string(curr->bytes) + "_";
+  std::vector<std::string> parts{curr->type.toString(),
+                                 std::to_string(curr->bytes)};
   if (LoadUtils::isSignRelevant(curr) && !curr->signed_) {
-    ret += "U_";
+    parts.push_back("U");
   }
-  if (curr->isAtomic()) {
-    ret += "A";
-  } else {
-    ret += std::to_string(curr->align);
+
+  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");
+      break;
+    }
   }
-  return ret;
+
+  return "SAFE_HEAP_LOAD_" + String::join(parts, "_");
 }
 
 static Name getStoreName(Store* curr) {
@@ -232,10 +243,11 @@ struct SafeHeap : public Pass {
             if (align > bytes) {
               continue;
             }
-            for (auto isAtomic : {true, false}) {
-              load.order =
-                isAtomic ? MemoryOrder::SeqCst : MemoryOrder::Unordered;
-              if (isAtomic &&
+            for (auto memoryOrder : {MemoryOrder::Unordered,
+                                     MemoryOrder::AcqRel,
+                                     MemoryOrder::SeqCst}) {
+              load.order = memoryOrder;
+              if (load.isAtomic() &&
                   !isPossibleAtomicOperation(
                     align, bytes, module->memories[0]->shared, type)) {
                 continue;

src/support/string.h

@@ -58,6 +58,23 @@ class Split : public std::vector<std::string> {
   }
 };
 
+template<typename StringLike,
+         typename = std::enable_if_t<
+           std::is_convertible_v<StringLike, std::string_view>>>
+std::string join(const std::vector<StringLike>& strs, std::string_view sep) {
+  if (strs.empty()) {
+    return "";
+  }
+
+  std::string ret = std::string(strs[0]);
+  for (size_t i = 1; i < strs.size(); i++) {
+    ret.append(sep);
+    ret.append(strs[i]);
+  }
+
+  return ret;
+}
+
 // Handles bracketing in a list initially split by ",", but the list may
 // contain nested ","s. For example,
 //   void foo(int, double)

src/tools/wasm-split/instrumenter.cpp

@@ -287,8 +287,12 @@ void Instrumenter::addProfileExport(size_t numFuncs) {
                   getAddr(),
                   builder.makeBinary(
                     MulInt32, getFuncIdx(), builder.makeConst(uint32_t(4)))),
-                builder.makeAtomicLoad(
-                  1, 0, getFuncIdx(), Type::i32, loadMemoryName),
+                builder.makeAtomicLoad(1,
+                                       0,
+                                       getFuncIdx(),
+                                       Type::i32,
+                                       loadMemoryName,
+                                       MemoryOrder::SeqCst),
                 Type::i32,
                 wasm->memories[0]->name),
               builder.makeLocalSet(

src/wasm-binary.h

@@ -1728,8 +1728,10 @@ class WasmBinaryReader {
   size_t inlineHintsLen = 0;
   void readInlineHints(size_t payloadLen);
 
-  Index readMemoryAccess(Address& alignment, Address& offset);
-  std::tuple<Name, Address, Address> getMemarg();
+  std::tuple<Address, Address, Index, MemoryOrder>
+  readMemoryAccess(bool isAtomic, bool isRMW);
+  std::tuple<Name, Address, Address, MemoryOrder> getMemarg(bool isAtomic,
+                                                            bool isRMW);
   MemoryOrder getMemoryOrder(bool isRMW = false);
 
   [[noreturn]] void throwError(std::string text) {

src/wasm-builder.h

@@ -384,10 +384,17 @@ class Builder {
     ret->finalize();
     return ret;
   }
-  Load* makeAtomicLoad(
-    unsigned bytes, Address offset, Expression* ptr, Type type, Name memory) {
+  Load* makeAtomicLoad(unsigned bytes,
+                       Address offset,
+                       Expression* ptr,
+                       Type type,
+                       Name memory,
+                       MemoryOrder order) {
+    assert(order != MemoryOrder::Unordered &&
+           "Atomic loads can't be unordered");
+
     Load* load = makeLoad(bytes, false, offset, bytes, ptr, type, memory);
-    load->order = MemoryOrder::SeqCst;
+    load->order = order;
     return load;
   }
   AtomicWait* makeAtomicWait(Expression* ptr,

src/wasm-ir-builder.h

@@ -156,7 +156,8 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
                     Name mem);
   Result<> makeStore(
     unsigned bytes, Address offset, unsigned align, Type type, Name mem);
-  Result<> makeAtomicLoad(unsigned bytes, Address offset, 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<> makeAtomicRMW(
     AtomicRMWOp op, unsigned bytes, Address offset, Type type, Name mem);