Support atomic struct accessors (#7155)

Implement support for both sequentially consistent and acquire-release
variants of `struct.atomic.get` and `struct.atomic.set`, as proposed by
shared-everything-threads. Introduce a new `MemoryOrdering` enum for
describing different levels of atomicity (or the lack thereof). This new
enum should eventually be adopted by linear memory atomic accessors as
well to support acquire-release semantics, but for now just use it in
`StructGet` and `StructSet`.

In addition to implementing parsing and emitting for the instructions,
validate that shared-everything is enabled to use them, mark them as
having synchronization side effects, and lightly optimize them by
relaxing acquire-release accesses to non-shared structs to normal,
unordered accesses. This is valid because such accesses cannot possibly
synchronize with other threads. Also update Precompute to avoid
optimizing out synchronization points.

There are probably other passes that need to be updated to avoid
incorrectly optimizing synchronizing accesses, but identifying and
fixing them is left as future work.

Author: tlively

scripts/gen-s-parser.py

@@ -617,7 +617,11 @@
     ("struct.get",           "makeStructGet()"),
     ("struct.get_s",         "makeStructGet(true)"),
     ("struct.get_u",         "makeStructGet(false)"),
+    ("struct.atomic.get",    "makeAtomicStructGet()"),
+    ("struct.atomic.get_s",  "makeAtomicStructGet(true)"),
+    ("struct.atomic.get_u",  "makeAtomicStructGet(false)"),
     ("struct.set",           "makeStructSet()"),
+    ("struct.atomic.set",    "makeAtomicStructSet()"),
     ("array.new",            "makeArrayNew(false)"),
     ("array.new_default",    "makeArrayNew(true)"),
     ("array.new_data",       "makeArrayNewData()"),

src/gen-s-parser.inc

@@ -5013,6 +5013,45 @@ switch (buf[0]) {
           }
           case 'u': {
             switch (buf[7]) {
+              case 'a': {
+                switch (buf[14]) {
+                  case 'g': {
+                    switch (buf[17]) {
+                      case '\0':
+                        if (op == "struct.atomic.get"sv) {
+                          CHECK_ERR(makeAtomicStructGet(ctx, pos, annotations));
+                          return Ok{};
+                        }
+                        goto parse_error;
+                      case '_': {
+                        switch (buf[18]) {
+                          case 's':
+                            if (op == "struct.atomic.get_s"sv) {
+                              CHECK_ERR(makeAtomicStructGet(ctx, pos, annotations, true));
+                              return Ok{};
+                            }
+                            goto parse_error;
+                          case 'u':
+                            if (op == "struct.atomic.get_u"sv) {
+                              CHECK_ERR(makeAtomicStructGet(ctx, pos, annotations, false));
+                              return Ok{};
+                            }
+                            goto parse_error;
+                          default: goto parse_error;
+                        }
+                      }
+                      default: goto parse_error;
+                    }
+                  }
+                  case 's':
+                    if (op == "struct.atomic.set"sv) {
+                      CHECK_ERR(makeAtomicStructSet(ctx, pos, annotations));
+                      return Ok{};
+                    }
+                    goto parse_error;
+                  default: goto parse_error;
+                }
+              }
               case 'g': {
                 switch (buf[10]) {
                   case '\0':

src/ir/effects.h

@@ -872,6 +872,18 @@ class EffectAnalyzer {
       if (curr->ref->type.isNullable()) {
         parent.implicitTrap = true;
       }
+      switch (curr->order) {
+        case MemoryOrder::Unordered:
+          break;
+        case MemoryOrder::SeqCst:
+          // Synchronizes with other threads.
+          parent.isAtomic = true;
+          break;
+        case MemoryOrder::AcqRel:
+          // Only synchronizes if other threads can read the field.
+          parent.isAtomic = curr->ref->type.getHeapType().isShared();
+          break;
+      }
     }
     void visitStructSet(StructSet* curr) {
       if (curr->ref->type.isNull()) {
@@ -883,6 +895,9 @@ class EffectAnalyzer {
       if (curr->ref->type.isNullable()) {
         parent.implicitTrap = true;
       }
+      if (curr->order != MemoryOrder::Unordered) {
+        parent.isAtomic = true;
+      }
     }
     void visitArrayNew(ArrayNew* curr) {}
     void visitArrayNewData(ArrayNewData* curr) {

src/parser/contexts.h

@@ -735,13 +735,20 @@ struct NullInstrParserCtx {
     return Ok{};
   }
   template<typename HeapTypeT>
-  Result<> makeStructGet(
-    Index, const std::vector<Annotation>&, HeapTypeT, FieldIdxT, bool) {
+  Result<> makeStructGet(Index,
+                         const std::vector<Annotation>&,
+                         HeapTypeT,
+                         FieldIdxT,
+                         bool,
+                         MemoryOrder = MemoryOrder::Unordered) {
     return Ok{};
   }
   template<typename HeapTypeT>
-  Result<>
-  makeStructSet(Index, const std::vector<Annotation>&, HeapTypeT, FieldIdxT) {
+  Result<> makeStructSet(Index,
+                         const std::vector<Annotation>&,
+                         HeapTypeT,
+                         FieldIdxT,
+                         MemoryOrder = MemoryOrder::Unordered) {
     return Ok{};
   }
   template<typename HeapTypeT>
@@ -2448,15 +2455,17 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
                          const std::vector<Annotation>& annotations,
                          HeapType type,
                          Index field,
-                         bool signed_) {
-    return withLoc(pos, irBuilder.makeStructGet(type, field, signed_));
+                         bool signed_,
+                         MemoryOrder order = MemoryOrder::Unordered) {
+    return withLoc(pos, irBuilder.makeStructGet(type, field, signed_, order));
   }
 
   Result<> makeStructSet(Index pos,
                          const std::vector<Annotation>& annotations,
                          HeapType type,
-                         Index field) {
-    return withLoc(pos, irBuilder.makeStructSet(type, field));
+                         Index field,
+                         MemoryOrder order = MemoryOrder::Unordered) {
+    return withLoc(pos, irBuilder.makeStructSet(type, field, order));
   }
 
   Result<> makeArrayNew(Index pos,

src/parser/parsers.h

@@ -48,6 +48,7 @@ template<typename Ctx> Result<typename Ctx::LimitsT> limits64(Ctx&);
 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> Result<typename Ctx::TableTypeT> tabletype(Ctx&);
 template<typename Ctx>
 Result<typename Ctx::TableTypeT> tabletypeContinued(Ctx&, Type addressType);
@@ -246,8 +247,15 @@ Result<> makeStructGet(Ctx&,
                        const std::vector<Annotation>&,
                        bool signed_ = false);
 template<typename Ctx>
+Result<> makeAtomicStructGet(Ctx&,
+                             Index,
+                             const std::vector<Annotation>&,
+                             bool signed_ = false);
+template<typename Ctx>
 Result<> makeStructSet(Ctx&, Index, const std::vector<Annotation>&);
 template<typename Ctx>
+Result<> makeAtomicStructSet(Ctx&, Index, const std::vector<Annotation>&);
+template<typename Ctx>
 Result<>
 makeArrayNew(Ctx&, Index, const std::vector<Annotation>&, bool default_);
 template<typename Ctx>
@@ -801,6 +809,17 @@ 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;
+}
+
 // tabletype ::= (limits32 | 'i32' limits32 | 'i64' limit64) reftype
 template<typename Ctx> Result<typename Ctx::TableTypeT> tabletype(Ctx& ctx) {
   Type addressType = Type::i32;
@@ -2224,6 +2243,20 @@ Result<> makeStructGet(Ctx& ctx,
   return ctx.makeStructGet(pos, annotations, *type, *field, signed_);
 }
 
+template<typename Ctx>
+Result<> makeAtomicStructGet(Ctx& ctx,
+                             Index pos,
+                             const std::vector<Annotation>& annotations,
+                             bool signed_) {
+  auto order = memorder(ctx);
+  CHECK_ERR(order);
+  auto type = typeidx(ctx);
+  CHECK_ERR(type);
+  auto field = fieldidx(ctx, *type);
+  CHECK_ERR(field);
+  return ctx.makeStructGet(pos, annotations, *type, *field, signed_, *order);
+}
+
 template<typename Ctx>
 Result<>
 makeStructSet(Ctx& ctx, Index pos, const std::vector<Annotation>& annotations) {
@@ -2234,6 +2267,19 @@ makeStructSet(Ctx& ctx, Index pos, const std::vector<Annotation>& annotations) {
   return ctx.makeStructSet(pos, annotations, *type, *field);
 }
 
+template<typename Ctx>
+Result<> makeAtomicStructSet(Ctx& ctx,
+                             Index pos,
+                             const std::vector<Annotation>& annotations) {
+  auto order = memorder(ctx);
+  CHECK_ERR(order);
+  auto type = typeidx(ctx);
+  CHECK_ERR(type);
+  auto field = fieldidx(ctx, *type);
+  CHECK_ERR(field);
+  return ctx.makeStructSet(pos, annotations, *type, *field, *order);
+}
+
 template<typename Ctx>
 Result<> makeArrayNew(Ctx& ctx,
                       Index pos,

src/passes/OptimizeInstructions.cpp

@@ -1831,6 +1831,12 @@ struct OptimizeInstructions
   void visitStructGet(StructGet* curr) {
     skipNonNullCast(curr->ref, curr);
     trapOnNull(curr, curr->ref);
+    // Relax acquire loads of unshared fields to unordered because they cannot
+    // synchronize with other threads.
+    if (curr->order == MemoryOrder::AcqRel && curr->ref->type.isRef() &&
+        !curr->ref->type.getHeapType().isShared()) {
+      curr->order = MemoryOrder::Unordered;
+    }
   }
 
   void visitStructSet(StructSet* curr) {
@@ -1847,6 +1853,13 @@ struct OptimizeInstructions
         optimizeStoredValue(curr->value, fields[curr->index].getByteSize());
       }
     }
+
+    // Relax release stores of unshared fields to unordered because they cannot
+    // synchronize with other threads.
+    if (curr->order == MemoryOrder::AcqRel && curr->ref->type.isRef() &&
+        !curr->ref->type.getHeapType().isShared()) {
+      curr->order = MemoryOrder::Unordered;
+    }
   }
 
   void visitArrayNew(ArrayNew* curr) {

src/passes/Precompute.cpp

@@ -134,23 +134,37 @@ class PrecomputingExpressionRunner
   }
   Flow visitStructSet(StructSet* curr) { return Flow(NONCONSTANT_FLOW); }
   Flow visitStructGet(StructGet* curr) {
-    if (curr->ref->type != Type::unreachable && !curr->ref->type.isNull()) {
-      // If this field is immutable then we may be able to precompute this, as
-      // if we also created the data in this function (or it was created in an
-      // immutable global) then we know the value in the field. If it is
-      // immutable, call the super method which will do the rest here. That
-      // includes checking for the data being properly created, as if it was
-      // not then we will not have a constant value for it, which means the
-      // local.get of that value will stop us.
-      auto& field =
-        curr->ref->type.getHeapType().getStruct().fields[curr->index];
-      if (field.mutable_ == Immutable) {
-        return Super::visitStructGet(curr);
-      }
+    if (curr->ref->type == Type::unreachable || curr->ref->type.isNull()) {
+      return Flow(NONCONSTANT_FLOW);
     }
-
-    // Otherwise, we've failed to precompute.
-    return Flow(NONCONSTANT_FLOW);
+    switch (curr->order) {
+      case MemoryOrder::Unordered:
+        // This can always be precomputed.
+        break;
+      case MemoryOrder::SeqCst:
+        // This can never be precomputed away because it synchronizes with other
+        // threads.
+        return Flow(NONCONSTANT_FLOW);
+      case MemoryOrder::AcqRel:
+        // This synchronizes only with writes to the same data, so it can still
+        // be precomputed if the data is not shared with other threads.
+        if (curr->ref->type.getHeapType().isShared()) {
+          return Flow(NONCONSTANT_FLOW);
+        }
+        break;
+    }
+    // If this field is immutable then we may be able to precompute this, as
+    // if we also created the data in this function (or it was created in an
+    // immutable global) then we know the value in the field. If it is
+    // immutable, call the super method which will do the rest here. That
+    // includes checking for the data being properly created, as if it was
+    // not then we will not have a constant value for it, which means the
+    // local.get of that value will stop us.
+    auto& field = curr->ref->type.getHeapType().getStruct().fields[curr->index];
+    if (field.mutable_ == Mutable) {
+      return Flow(NONCONSTANT_FLOW);
+    }
+    return Super::visitStructGet(curr);
   }
   Flow visitArrayNew(ArrayNew* curr) {
     auto flow = Super::visitArrayNew(curr);

src/passes/Print.cpp

@@ -2276,24 +2276,47 @@ struct PrintExpressionContents
       o << index;
     }
   }
+  void printMemoryOrder(MemoryOrder order) {
+    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
+      // for more concise output.
+      case MemoryOrder::Unordered:
+      case MemoryOrder::SeqCst:
+        break;
+      case MemoryOrder::AcqRel:
+        o << "acqrel ";
+        break;
+    }
+  }
   void visitStructGet(StructGet* curr) {
     auto heapType = curr->ref->type.getHeapType();
     const auto& field = heapType.getStruct().fields[curr->index];
+    printMedium(o, "struct");
+    if (curr->order != MemoryOrder::Unordered) {
+      printMedium(o, ".atomic");
+    }
     if (field.type == Type::i32 && field.packedType != Field::not_packed) {
       if (curr->signed_) {
-        printMedium(o, "struct.get_s ");
+        printMedium(o, ".get_s ");
       } else {
-        printMedium(o, "struct.get_u ");
+        printMedium(o, ".get_u ");
       }
     } else {
-      printMedium(o, "struct.get ");
+      printMedium(o, ".get ");
     }
+    printMemoryOrder(curr->order);
     printHeapType(heapType);
     o << ' ';
     printFieldName(heapType, curr->index);
   }
   void visitStructSet(StructSet* curr) {
-    printMedium(o, "struct.set ");
+    if (curr->order == MemoryOrder::Unordered) {
+      printMedium(o, "struct.set ");
+    } else {
+      printMedium(o, "struct.atomic.set ");
+    }
+    printMemoryOrder(curr->order);
     auto heapType = curr->ref->type.getHeapType();
     printHeapType(heapType);
     o << ' ';

src/wasm-binary.h

@@ -1125,6 +1125,15 @@ enum ASTNodes {
   I31GetU = 0x1e,
   RefI31Shared = 0x1f,
 
+  // Shared GC Opcodes
+
+  OrderSeqCst = 0x0,
+  OrderAcqRel = 0x1,
+  StructAtomicGet = 0x5c,
+  StructAtomicGetS = 0x5d,
+  StructAtomicGetU = 0x5e,
+  StructAtomicSet = 0x5f,
+
   // stringref opcodes
 
   StringConst = 0x82,
@@ -1352,6 +1361,8 @@ class WasmBinaryWriter {
 
   void writeField(const Field& field);
 
+  void writeMemoryOrder(MemoryOrder order);
+
 private:
   Module* wasm;
   BufferWithRandomAccess& o;
@@ -1587,6 +1598,7 @@ class WasmBinaryReader {
 
   Index readMemoryAccess(Address& alignment, Address& offset);
   std::tuple<Name, Address, Address> getMemarg();
+  MemoryOrder getMemoryOrder();
 
   [[noreturn]] void throwError(std::string text) {
     throw ParseException(text, 0, pos);