[Shared-Everything] ArrayRMW and ArrayCmpxchg

scripts/gen-s-parser.py

@@ -653,6 +653,13 @@
     ("array.fill",           "makeArrayFill()"),
     ("array.init_data",      "makeArrayInitData()"),
     ("array.init_elem",      "makeArrayInitElem()"),
+    ("array.atomic.rmw.add", "makeArrayRMW(RMWAdd)"),
+    ("array.atomic.rmw.sub", "makeArrayRMW(RMWSub)"),
+    ("array.atomic.rmw.and", "makeArrayRMW(RMWAnd)"),
+    ("array.atomic.rmw.or",  "makeArrayRMW(RMWOr)"),
+    ("array.atomic.rmw.xor", "makeArrayRMW(RMWXor)"),
+    ("array.atomic.rmw.xchg", "makeArrayRMW(RMWXchg)"),
+    ("array.atomic.rmw.cmpxchg", "makeArrayCmpxchg()"),
     ("ref.as_non_null",      "makeRefAs(RefAsNonNull)"),
     ("extern.internalize",   "makeRefAs(AnyConvertExtern)"),  # Deprecated
     ("extern.externalize",   "makeRefAs(ExternConvertAny)"),  # Deprecated

src/gen-s-parser.inc

@@ -52,6 +52,63 @@ switch (buf[0]) {
                   default: goto parse_error;
                 }
               }
+              case 'r': {
+                switch (buf[17]) {
+                  case 'a': {
+                    switch (buf[18]) {
+                      case 'd':
+                        if (op == "array.atomic.rmw.add"sv) {
+                          CHECK_ERR(makeArrayRMW(ctx, pos, annotations, RMWAdd));
+                          return Ok{};
+                        }
+                        goto parse_error;
+                      case 'n':
+                        if (op == "array.atomic.rmw.and"sv) {
+                          CHECK_ERR(makeArrayRMW(ctx, pos, annotations, RMWAnd));
+                          return Ok{};
+                        }
+                        goto parse_error;
+                      default: goto parse_error;
+                    }
+                  }
+                  case 'c':
+                    if (op == "array.atomic.rmw.cmpxchg"sv) {
+                      CHECK_ERR(makeArrayCmpxchg(ctx, pos, annotations));
+                      return Ok{};
+                    }
+                    goto parse_error;
+                  case 'o':
+                    if (op == "array.atomic.rmw.or"sv) {
+                      CHECK_ERR(makeArrayRMW(ctx, pos, annotations, RMWOr));
+                      return Ok{};
+                    }
+                    goto parse_error;
+                  case 's':
+                    if (op == "array.atomic.rmw.sub"sv) {
+                      CHECK_ERR(makeArrayRMW(ctx, pos, annotations, RMWSub));
+                      return Ok{};
+                    }
+                    goto parse_error;
+                  case 'x': {
+                    switch (buf[18]) {
+                      case 'c':
+                        if (op == "array.atomic.rmw.xchg"sv) {
+                          CHECK_ERR(makeArrayRMW(ctx, pos, annotations, RMWXchg));
+                          return Ok{};
+                        }
+                        goto parse_error;
+                      case 'o':
+                        if (op == "array.atomic.rmw.xor"sv) {
+                          CHECK_ERR(makeArrayRMW(ctx, pos, annotations, RMWXor));
+                          return Ok{};
+                        }
+                        goto parse_error;
+                      default: goto parse_error;
+                    }
+                  }
+                  default: goto parse_error;
+                }
+              }
               case 's':
                 if (op == "array.atomic.set"sv) {
                   CHECK_ERR(makeAtomicArraySet(ctx, pos, annotations));

src/interpreter/interpreter.cpp

@@ -263,6 +263,8 @@ struct ExpressionInterpreter : OverriddenVisitor<ExpressionInterpreter, Flow> {
   Flow visitArrayFill(ArrayFill* curr) { WASM_UNREACHABLE("TODO"); }
   Flow visitArrayInitData(ArrayInitData* curr) { WASM_UNREACHABLE("TODO"); }
   Flow visitArrayInitElem(ArrayInitElem* curr) { WASM_UNREACHABLE("TODO"); }
+  Flow visitArrayRMW(ArrayRMW* curr) { WASM_UNREACHABLE("TODO"); }
+  Flow visitArrayCmpxchg(ArrayCmpxchg* curr) { WASM_UNREACHABLE("TODO"); }
   Flow visitRefAs(RefAs* curr) { WASM_UNREACHABLE("TODO"); }
   Flow visitStringNew(StringNew* curr) { WASM_UNREACHABLE("TODO"); }
   Flow visitStringConst(StringConst* curr) { WASM_UNREACHABLE("TODO"); }

src/ir/ReFinalize.cpp

@@ -172,6 +172,8 @@ void ReFinalize::visitArrayCopy(ArrayCopy* curr) { curr->finalize(); }
 void ReFinalize::visitArrayFill(ArrayFill* curr) { curr->finalize(); }
 void ReFinalize::visitArrayInitData(ArrayInitData* curr) { curr->finalize(); }
 void ReFinalize::visitArrayInitElem(ArrayInitElem* curr) { curr->finalize(); }
+void ReFinalize::visitArrayRMW(ArrayRMW* curr) { curr->finalize(); }
+void ReFinalize::visitArrayCmpxchg(ArrayCmpxchg* curr) { curr->finalize(); }
 void ReFinalize::visitRefAs(RefAs* curr) { curr->finalize(); }
 void ReFinalize::visitStringNew(StringNew* curr) { curr->finalize(); }
 void ReFinalize::visitStringConst(StringConst* curr) { curr->finalize(); }

src/ir/child-typer.h

@@ -1126,6 +1126,35 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     note(&curr->size, Type::i32);
   }
 
+  void visitArrayRMW(ArrayRMW* curr,
+                     std::optional<HeapType> ht = std::nullopt) {
+    if (!ht) {
+      if (self().skipUnreachable() && !curr->ref->type.isRef()) {
+        return;
+      }
+      ht = curr->ref->type.getHeapType();
+    }
+    auto type = ht->getArray().element.type;
+    note(&curr->ref, Type(*ht, Nullable));
+    note(&curr->index, Type::i32);
+    note(&curr->value, type);
+  }
+
+  void visitArrayCmpxchg(ArrayCmpxchg* curr,
+                         std::optional<HeapType> ht = std::nullopt) {
+    if (!ht) {
+      if (self().skipUnreachable() && !curr->ref->type.isRef()) {
+        return;
+      }
+      ht = curr->ref->type.getHeapType();
+    }
+    auto type = ht->getArray().element.type;
+    note(&curr->ref, Type(*ht, Nullable));
+    note(&curr->index, Type::i32);
+    note(&curr->expected, type);
+    note(&curr->replacement, type);
+  }
+
   void visitRefAs(RefAs* curr) {
     switch (curr->op) {
       case RefAsNonNull:

src/ir/cost.h

@@ -752,6 +752,15 @@ struct CostAnalyzer : public OverriddenVisitor<CostAnalyzer, CostType> {
     return 6 + visit(curr->ref) + visit(curr->index) + visit(curr->offset) +
            visit(curr->size);
   }
+  CostType visitArrayRMW(ArrayRMW* curr) {
+    return AtomicCost + nullCheckCost(curr->ref) + visit(curr->ref) +
+           visit(curr->index) + visit(curr->value);
+  }
+  CostType visitArrayCmpxchg(ArrayCmpxchg* curr) {
+    return AtomicCost + nullCheckCost(curr->ref) + visit(curr->ref) +
+           visit(curr->index) + visit(curr->expected) +
+           visit(curr->replacement);
+  }
   CostType visitRefAs(RefAs* curr) { return 1 + visit(curr->value); }
   CostType visitStringNew(StringNew* curr) {
     return 8 + visit(curr->ref) + maybeVisit(curr->start) +

src/ir/effects.h

@@ -1000,6 +1000,32 @@ class EffectAnalyzer {
     }
     void visitArrayInitData(ArrayInitData* curr) { visitArrayInit(curr); }
     void visitArrayInitElem(ArrayInitElem* curr) { visitArrayInit(curr); }
+    void visitArrayRMW(ArrayRMW* curr) {
+      if (curr->ref->type.isNull()) {
+        parent.trap = true;
+        return;
+      }
+      parent.readsArray = true;
+      parent.writesArray = true;
+      if (curr->ref->type.isNullable()) {
+        parent.implicitTrap = true;
+      }
+      assert(curr->order != MemoryOrder::Unordered);
+      parent.isAtomic = true;
+    }
+    void visitArrayCmpxchg(ArrayCmpxchg* curr) {
+      if (curr->ref->type.isNull()) {
+        parent.trap = true;
+        return;
+      }
+      parent.readsArray = true;
+      parent.writesArray = true;
+      if (curr->ref->type.isNullable()) {
+        parent.implicitTrap = true;
+      }
+      assert(curr->order != MemoryOrder::Unordered);
+      parent.isAtomic = true;
+    }
     void visitRefAs(RefAs* curr) {
       if (curr->op == AnyConvertExtern || curr->op == ExternConvertAny) {
         // These conversions are infallible.

src/ir/possible-contents.cpp

@@ -1100,6 +1100,22 @@ struct InfoCollector
   }
   void visitArrayInitData(ArrayInitData* curr) { visitArrayInit(curr); }
   void visitArrayInitElem(ArrayInitElem* curr) { visitArrayInit(curr); }
+  void visitArrayRMW(ArrayRMW* 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 visitArrayCmpxchg(ArrayCmpxchg* 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 visitStringNew(StringNew* curr) {
     if (curr->type == Type::unreachable) {
       return;
@@ -1618,6 +1634,8 @@ void TNHOracle::scan(Function* func,
     void visitArrayInitElem(ArrayInitElem* curr) {
       notePossibleTrap(curr->ref);
     }
+    void visitArrayRMW(ArrayRMW* curr) { notePossibleTrap(curr->ref); }
+    void visitArrayCmpxchg(ArrayCmpxchg* curr) { notePossibleTrap(curr->ref); }
 
     void visitFunction(Function* curr) {
       // In optimized TNH code, a function that always traps will be turned

src/ir/subtype-exprs.h

@@ -399,6 +399,21 @@ struct SubtypingDiscoverer : public OverriddenVisitor<SubType> {
     auto* seg = self()->getModule()->getElementSegment(curr->segment);
     self()->noteSubtype(seg->type, array.element.type);
   }
+  void visitArrayRMW(ArrayRMW* curr) {
+    if (!curr->ref->type.isArray()) {
+      return;
+    }
+    auto array = curr->ref->type.getHeapType().getArray();
+    self()->noteSubtype(curr->value, array.element.type);
+  }
+  void visitArrayCmpxchg(ArrayCmpxchg* curr) {
+    if (!curr->ref->type.isArray()) {
+      return;
+    }
+    auto array = curr->ref->type.getHeapType().getArray();
+    self()->noteSubtype(curr->expected, array.element.type);
+    self()->noteSubtype(curr->replacement, array.element.type);
+  }
   void visitRefAs(RefAs* curr) {
     if (curr->op == RefAsNonNull) {
       self()->noteCast(curr->value, curr);

src/parser/contexts.h

@@ -848,6 +848,21 @@ struct NullInstrParserCtx {
                              ElemIdxT) {
     return Ok{};
   }
+  template<typename HeapTypeT>
+  Result<> makeArrayRMW(Index,
+                        const std::vector<Annotation>&,
+                        AtomicRMWOp,
+                        HeapTypeT,
+                        MemoryOrder) {
+    return Ok{};
+  }
+  template<typename HeapTypeT>
+  Result<> makeArrayCmpxchg(Index,
+                            const std::vector<Annotation>&,
+                            HeapTypeT,
+                            MemoryOrder) {
+    return Ok{};
+  }
   Result<> makeRefAs(Index, const std::vector<Annotation>&, RefAsOp) {
     return Ok{};
   }
@@ -2738,6 +2753,21 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
     return withLoc(pos, irBuilder.makeArrayInitElem(type, elem));
   }
 
+  Result<> makeArrayRMW(Index pos,
+                        const std::vector<Annotation>& annotations,
+                        AtomicRMWOp op,
+                        HeapType type,
+                        MemoryOrder order) {
+    return withLoc(pos, irBuilder.makeArrayRMW(op, type, order));
+  }
+
+  Result<> makeArrayCmpxchg(Index pos,
+                            const std::vector<Annotation>& annotations,
+                            HeapType type,
+                            MemoryOrder order) {
+    return withLoc(pos, irBuilder.makeArrayCmpxchg(type, order));
+  }
+
   Result<>
   makeRefAs(Index pos, const std::vector<Annotation>& annotations, RefAsOp op) {
     return withLoc(pos, irBuilder.makeRefAs(op));

src/parser/parsers.h

@@ -292,6 +292,10 @@ Result<> makeArrayInitData(Ctx&, Index, const std::vector<Annotation>&);
 template<typename Ctx>
 Result<> makeArrayInitElem(Ctx&, Index, const std::vector<Annotation>&);
 template<typename Ctx>
+Result<> makeArrayRMW(AtomicRMWOp, Index, const std::vector<Annotation>&);
+template<typename Ctx>
+Result<> makeArrayCmpxchg(Index, const std::vector<Annotation>&);
+template<typename Ctx>
 Result<> makeRefAs(Ctx&, Index, const std::vector<Annotation>&, RefAsOp op);
 template<typename Ctx>
 Result<>
@@ -2499,6 +2503,39 @@ Result<> makeArrayInitElem(Ctx& ctx,
   return ctx.makeArrayInitElem(pos, annotations, *type, *elem);
 }
 
+template<typename Ctx>
+Result<> makeArrayRMW(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, "array.atomic.rmw memory orders must be identical");
+  }
+  auto type = typeidx(ctx);
+  CHECK_ERR(type);
+  return ctx.makeArrayRMW(pos, annotations, op, *type, *order1);
+}
+
+template<typename Ctx>
+Result<> makeArrayCmpxchg(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, "array.atomic.rmw memory orders must be identical");
+  }
+  auto type = typeidx(ctx);
+  CHECK_ERR(type);
+  return ctx.makeArrayCmpxchg(pos, annotations, *type, *order1);
+}
+
 template<typename Ctx>
 Result<> makeRefAs(Ctx& ctx,
                    Index pos,

src/passes/Heap2Local.cpp

@@ -448,6 +448,18 @@ struct EscapeAnalyzer {
         escapes = false;
         fullyConsumes = true;
       }
+      void visitArrayRMW(ArrayRMW* curr) {
+        if (curr->ref == child) {
+          escapes = false;
+          fullyConsumes = true;
+        }
+      }
+      void visitArrayCmpxchg(ArrayCmpxchg* curr) {
+        if (curr->ref == child || curr->expected == child) {
+          escapes = false;
+          fullyConsumes = true;
+        }
+      }
       // TODO other GC operations
     } checker;
 

src/passes/OptimizeInstructions.cpp

@@ -2240,6 +2240,18 @@ struct OptimizeInstructions
     trapOnNull(curr, curr->destRef) || trapOnNull(curr, curr->srcRef);
   }
 
+  void visitArrayRMW(ArrayRMW* curr) {
+    skipNonNullCast(curr->ref, curr);
+    trapOnNull(curr, curr->ref);
+    // TODO: more opts like StructRMW
+  }
+
+  void visitArrayCmpxchg(ArrayCmpxchg* curr) {
+    skipNonNullCast(curr->ref, curr);
+    trapOnNull(curr, curr->ref);
+    // TODO: more opts like StructCmpxchg
+  }
+
   void visitRefCast(RefCast* curr) {
     // Note we must check the ref's type here and not our own, since we only
     // refinalize at the end, which means our type may not have been updated yet

src/passes/Print.cpp

@@ -2444,6 +2444,26 @@ struct PrintExpressionContents
     o << ' ';
     curr->segment.print(o);
   }
+  void visitArrayRMW(ArrayRMW* curr) {
+    prepareColor(o);
+    o << "array.atomic.rmw.";
+    printAtomicRMWOp(curr->op);
+    restoreNormalColor(o);
+    o << ' ';
+    printMemoryOrder(curr->order);
+    printMemoryOrder(curr->order);
+    auto heapType = curr->ref->type.getHeapType();
+    printHeapTypeName(heapType);
+  }
+  void visitArrayCmpxchg(ArrayCmpxchg* curr) {
+    prepareColor(o);
+    o << "array.atomic.rmw.cmpxchg ";
+    restoreNormalColor(o);
+    printMemoryOrder(curr->order);
+    printMemoryOrder(curr->order);
+    auto heapType = curr->ref->type.getHeapType();
+    printHeapTypeName(heapType);
+  }
   void visitRefAs(RefAs* curr) {
     switch (curr->op) {
       case RefAsNonNull: