Atomic struct RMW instructions (#7194)

Add `StructRMW` and `StructCmpxchg` expression classes with binary and text printing and parsing as well as validation.
WebAssembly · Jan 11, 2025 · caefb33 · caefb33
1 parent b6eacd7
commit caefb33
Show file tree

Hide file tree

Showing 29 changed files with 962 additions and 54 deletions.
diff --git a/scripts/gen-s-parser.py b/scripts/gen-s-parser.py
@@ -622,6 +622,13 @@
     ("struct.atomic.get_u",  "makeAtomicStructGet(false)"),
     ("struct.set",           "makeStructSet()"),
     ("struct.atomic.set",    "makeAtomicStructSet()"),
+    ("struct.atomic.rmw.add", "makeStructRMW(RMWAdd)"),
+    ("struct.atomic.rmw.sub", "makeStructRMW(RMWSub)"),
+    ("struct.atomic.rmw.and", "makeStructRMW(RMWAnd)"),
+    ("struct.atomic.rmw.or", "makeStructRMW(RMWOr)"),
+    ("struct.atomic.rmw.xor", "makeStructRMW(RMWXor)"),
+    ("struct.atomic.rmw.xchg", "makeStructRMW(RMWXchg)"),
+    ("struct.atomic.rmw.cmpxchg", "makeStructCmpxchg()"),
     ("array.new",            "makeArrayNew(false)"),
     ("array.new_default",    "makeArrayNew(true)"),
     ("array.new_data",       "makeArrayNewData()"),

diff --git a/src/gen-s-parser.inc b/src/gen-s-parser.inc
@@ -5043,6 +5043,63 @@ switch (buf[0]) {
                       default: goto parse_error;
                     }
                   }
+                  case 'r': {
+                    switch (buf[18]) {
+                      case 'a': {
+                        switch (buf[19]) {
+                          case 'd':
+                            if (op == "struct.atomic.rmw.add"sv) {
+                              CHECK_ERR(makeStructRMW(ctx, pos, annotations, RMWAdd));
+                              return Ok{};
+                            }
+                            goto parse_error;
+                          case 'n':
+                            if (op == "struct.atomic.rmw.and"sv) {
+                              CHECK_ERR(makeStructRMW(ctx, pos, annotations, RMWAnd));
+                              return Ok{};
+                            }
+                            goto parse_error;
+                          default: goto parse_error;
+                        }
+                      }
+                      case 'c':
+                        if (op == "struct.atomic.rmw.cmpxchg"sv) {
+                          CHECK_ERR(makeStructCmpxchg(ctx, pos, annotations));
+                          return Ok{};
+                        }
+                        goto parse_error;
+                      case 'o':
+                        if (op == "struct.atomic.rmw.or"sv) {
+                          CHECK_ERR(makeStructRMW(ctx, pos, annotations, RMWOr));
+                          return Ok{};
+                        }
+                        goto parse_error;
+                      case 's':
+                        if (op == "struct.atomic.rmw.sub"sv) {
+                          CHECK_ERR(makeStructRMW(ctx, pos, annotations, RMWSub));
+                          return Ok{};
+                        }
+                        goto parse_error;
+                      case 'x': {
+                        switch (buf[19]) {
+                          case 'c':
+                            if (op == "struct.atomic.rmw.xchg"sv) {
+                              CHECK_ERR(makeStructRMW(ctx, pos, annotations, RMWXchg));
+                              return Ok{};
+                            }
+                            goto parse_error;
+                          case 'o':
+                            if (op == "struct.atomic.rmw.xor"sv) {
+                              CHECK_ERR(makeStructRMW(ctx, pos, annotations, RMWXor));
+                              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));

diff --git a/src/ir/ReFinalize.cpp b/src/ir/ReFinalize.cpp
@@ -158,6 +158,8 @@ void ReFinalize::visitBrOn(BrOn* curr) {
 void ReFinalize::visitStructNew(StructNew* curr) { curr->finalize(); }
 void ReFinalize::visitStructGet(StructGet* curr) { curr->finalize(); }
 void ReFinalize::visitStructSet(StructSet* curr) { curr->finalize(); }
+void ReFinalize::visitStructRMW(StructRMW* curr) { curr->finalize(); }
+void ReFinalize::visitStructCmpxchg(StructCmpxchg* curr) { curr->finalize(); }
 void ReFinalize::visitArrayNew(ArrayNew* curr) { curr->finalize(); }
 void ReFinalize::visitArrayNewData(ArrayNewData* curr) { curr->finalize(); }
 void ReFinalize::visitArrayNewElem(ArrayNewElem* curr) { curr->finalize(); }

diff --git a/src/ir/child-typer.h b/src/ir/child-typer.h
@@ -891,6 +891,29 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     note(&curr->value, fields[curr->index].type);
   }
 
+  void visitStructRMW(StructRMW* curr,
+                      std::optional<HeapType> ht = std::nullopt) {
+    if (!ht) {
+      ht = curr->ref->type.getHeapType();
+    }
+    const auto& fields = ht->getStruct().fields;
+    assert(curr->index < fields.size());
+    note(&curr->ref, Type(*ht, Nullable));
+    note(&curr->value, fields[curr->index].type);
+  }
+
+  void visitStructCmpxchg(StructCmpxchg* curr,
+                          std::optional<HeapType> ht = std::nullopt) {
+    if (!ht) {
+      ht = curr->ref->type.getHeapType();
+    }
+    const auto& fields = ht->getStruct().fields;
+    assert(curr->index < fields.size());
+    note(&curr->ref, Type(*ht, Nullable));
+    note(&curr->expected, fields[curr->index].type);
+    note(&curr->replacement, fields[curr->index].type);
+  }
+
   void visitArrayNew(ArrayNew* curr) {
     if (!curr->isWithDefault()) {
       note(&curr->init, curr->type.getHeapType().getArray().element.type);

diff --git a/src/ir/cost.h b/src/ir/cost.h
@@ -687,6 +687,14 @@ struct CostAnalyzer : public OverriddenVisitor<CostAnalyzer, CostType> {
   CostType visitStructSet(StructSet* curr) {
     return 2 + nullCheckCost(curr->ref) + visit(curr->ref) + visit(curr->value);
   }
+  CostType visitStructRMW(StructRMW* curr) {
+    return AtomicCost + nullCheckCost(curr->ref) + visit(curr->ref) +
+           visit(curr->value);
+  }
+  CostType visitStructCmpxchg(StructCmpxchg* curr) {
+    return AtomicCost + nullCheckCost(curr->ref) + visit(curr->ref) +
+           visit(curr->expected) + visit(curr->replacement);
+  }
   CostType visitArrayNew(ArrayNew* curr) {
     return 4 + visit(curr->size) + maybeVisit(curr->init);
   }

diff --git a/src/ir/effects.h b/src/ir/effects.h
@@ -899,6 +899,32 @@ class EffectAnalyzer {
         parent.isAtomic = true;
       }
     }
+    void visitStructRMW(StructRMW* curr) {
+      if (curr->ref->type.isNull()) {
+        parent.trap = true;
+        return;
+      }
+      parent.readsMutableStruct = true;
+      parent.writesStruct = true;
+      if (curr->ref->type.isNullable()) {
+        parent.implicitTrap = true;
+      }
+      assert(curr->order != MemoryOrder::Unordered);
+      parent.isAtomic = true;
+    }
+    void visitStructCmpxchg(StructCmpxchg* curr) {
+      if (curr->ref->type.isNull()) {
+        parent.trap = true;
+        return;
+      }
+      parent.readsMutableStruct = true;
+      parent.writesStruct = true;
+      if (curr->ref->type.isNullable()) {
+        parent.implicitTrap = true;
+      }
+      assert(curr->order != MemoryOrder::Unordered);
+      parent.isAtomic = true;
+    }
     void visitArrayNew(ArrayNew* curr) {}
     void visitArrayNewData(ArrayNewData* curr) {
       // Traps on out of bounds access to segments or access to dropped

diff --git a/src/ir/possible-contents.cpp b/src/ir/possible-contents.cpp
@@ -1016,6 +1016,22 @@ struct InfoCollector
     addChildParentLink(curr->ref, curr);
     addChildParentLink(curr->value, curr);
   }
+  void visitStructRMW(StructRMW* curr) {
+    if (curr->ref->type == Type::unreachable) {
+      return;
+    }
+    // TODO: Model the modification part of the RMW in addition to the read and
+    // the write.
+    addRoot(curr);
+  }
+  void visitStructCmpxchg(StructCmpxchg* curr) {
+    if (curr->ref->type == Type::unreachable) {
+      return;
+    }
+    // TODO: Model the modification part of the RMW in addition to the read and
+    // the write.
+    addRoot(curr);
+  }
   // Array operations access the array's location, parallel to how structs work.
   void visitArrayGet(ArrayGet* curr) {
     if (!isRelevant(curr->ref)) {
@@ -1552,6 +1568,10 @@ void TNHOracle::scan(Function* func,
 
     void visitStructGet(StructGet* curr) { notePossibleTrap(curr->ref); }
     void visitStructSet(StructSet* curr) { notePossibleTrap(curr->ref); }
+    void visitStructRMW(StructRMW* curr) { notePossibleTrap(curr->ref); }
+    void visitStructCmpxchg(StructCmpxchg* curr) {
+      notePossibleTrap(curr->ref);
+    }
     void visitArrayGet(ArrayGet* curr) { notePossibleTrap(curr->ref); }
     void visitArraySet(ArraySet* curr) { notePossibleTrap(curr->ref); }
     void visitArrayLen(ArrayLen* curr) { notePossibleTrap(curr->ref); }
@@ -2354,7 +2374,8 @@ bool Flower::updateContents(LocationIndex locationIndex,
       // we must only combine the filtered contents (e.g. if 0xff arrives which
       // as a signed read is truly 0xffffffff then we cannot first combine the
       // existing 0xffffffff with the new 0xff, as they are different, and the
-      // result will no longer be a constant).
+      // result will no longer be a constant). There is no need to filter atomic
+      // RMW operations here because they always do unsigned reads.
       filterPackedDataReads(newContents, *exprLoc);
 #if defined(POSSIBLE_CONTENTS_DEBUG) && POSSIBLE_CONTENTS_DEBUG >= 2
       std::cout << "  pre-filtered packed read contents:\n";

diff --git a/src/ir/subtype-exprs.h b/src/ir/subtype-exprs.h
@@ -324,6 +324,21 @@ struct SubtypingDiscoverer : public OverriddenVisitor<SubType> {
     const auto& fields = curr->ref->type.getHeapType().getStruct().fields;
     self()->noteSubtype(curr->value, fields[curr->index].type);
   }
+  void visitStructRMW(StructRMW* curr) {
+    if (!curr->ref->type.isStruct()) {
+      return;
+    }
+    const auto& fields = curr->ref->type.getHeapType().getStruct().fields;
+    self()->noteSubtype(curr->value, fields[curr->index].type);
+  }
+  void visitStructCmpxchg(StructCmpxchg* curr) {
+    if (!curr->ref->type.isStruct()) {
+      return;
+    }
+    const auto& fields = curr->ref->type.getHeapType().getStruct().fields;
+    self()->noteSubtype(curr->expected, fields[curr->index].type);
+    self()->noteSubtype(curr->replacement, fields[curr->index].type);
+  }
   void visitArrayNew(ArrayNew* curr) {
     if (!curr->type.isArray() || curr->isWithDefault()) {
       return;

diff --git a/src/parser/contexts.h b/src/parser/contexts.h
@@ -749,6 +749,20 @@ struct NullInstrParserCtx {
     return Ok{};
   }
   template<typename HeapTypeT>
+  Result<> makeStructRMW(Index,
+                         const std::vector<Annotation>&,
+                         AtomicRMWOp,
+                         HeapTypeT,
+                         FieldIdxT,
+                         MemoryOrder) {
+    return Ok{};
+  }
+  template<typename HeapTypeT>
+  Result<> makeStructCmpxchg(
+    Index, const std::vector<Annotation>&, HeapTypeT, FieldIdxT, MemoryOrder) {
+    return Ok{};
+  }
+  template<typename HeapTypeT>
   Result<> makeArrayNew(Index, const std::vector<Annotation>&, HeapTypeT) {
     return Ok{};
   }
@@ -2453,18 +2467,35 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
                          HeapType type,
                          Index field,
                          bool signed_,
-                         MemoryOrder order = MemoryOrder::Unordered) {
+                         MemoryOrder order) {
     return withLoc(pos, irBuilder.makeStructGet(type, field, signed_, order));
   }
 
   Result<> makeStructSet(Index pos,
                          const std::vector<Annotation>& annotations,
                          HeapType type,
                          Index field,
-                         MemoryOrder order = MemoryOrder::Unordered) {
+                         MemoryOrder order) {
     return withLoc(pos, irBuilder.makeStructSet(type, field, order));
   }
 
+  Result<> makeStructRMW(Index pos,
+                         const std::vector<Annotation>& annotations,
+                         AtomicRMWOp op,
+                         HeapType type,
+                         Index field,
+                         MemoryOrder order) {
+    return withLoc(pos, irBuilder.makeStructRMW(op, type, field, order));
+  }
+
+  Result<> makeStructCmpxchg(Index pos,
+                             const std::vector<Annotation>& annotations,
+                             HeapType type,
+                             Index field,
+                             MemoryOrder order) {
+    return withLoc(pos, irBuilder.makeStructCmpxchg(type, field, order));
+  }
+
   Result<> makeArrayNew(Index pos,
                         const std::vector<Annotation>& annotations,
                         HeapType type) {

diff --git a/src/parser/parsers.h b/src/parser/parsers.h
@@ -256,6 +256,10 @@ Result<> makeStructSet(Ctx&, Index, const std::vector<Annotation>&);
 template<typename Ctx>
 Result<> makeAtomicStructSet(Ctx&, Index, const std::vector<Annotation>&);
 template<typename Ctx>
+Result<> makeStructRMW(AtomicRMWOp, Index, const std::vector<Annotation>&);
+template<typename Ctx>
+Result<> makeStructCmpxchg(Index, const std::vector<Annotation>&);
+template<typename Ctx>
 Result<>
 makeArrayNew(Ctx&, Index, const std::vector<Annotation>&, bool default_);
 template<typename Ctx>
@@ -2282,6 +2286,43 @@ Result<> makeAtomicStructSet(Ctx& ctx,
   return ctx.makeStructSet(pos, annotations, *type, *field, *order);
 }
 
+template<typename Ctx>
+Result<> makeStructRMW(Ctx& ctx,
+                       Index pos,
+                       const std::vector<Annotation>& annotations,
+                       AtomicRMWOp op) {
+  auto order1 = memorder(ctx);
+  CHECK_ERR(order1);
+  auto order2 = memorder(ctx);
+  CHECK_ERR(order2);
+  if (*order1 != *order2) {
+    return ctx.in.err(pos, "struct.atomic.rmw memory orders must be identical");
+  }
+  auto type = typeidx(ctx);
+  CHECK_ERR(type);
+  auto field = fieldidx(ctx, *type);
+  CHECK_ERR(field);
+  return ctx.makeStructRMW(pos, annotations, op, *type, *field, *order1);
+}
+
+template<typename Ctx>
+Result<> makeStructCmpxchg(Ctx& ctx,
+                           Index pos,
+                           const std::vector<Annotation>& annotations) {
+  auto order1 = memorder(ctx);
+  CHECK_ERR(order1);
+  auto order2 = memorder(ctx);
+  CHECK_ERR(order2);
+  if (*order1 != *order2) {
+    return ctx.in.err(pos, "struct.atomic.rmw memory orders must be identical");
+  }
+  auto type = typeidx(ctx);
+  CHECK_ERR(type);
+  auto field = fieldidx(ctx, *type);
+  CHECK_ERR(field);
+  return ctx.makeStructCmpxchg(pos, annotations, *type, *field, *order1);
+}
+
 template<typename Ctx>
 Result<> makeArrayNew(Ctx& ctx,
                       Index pos,