Handle StructCmpxchg::expected in Heap2Local (#8491)

We previously assumed that if we were optimizing a StructCmpxchg in
Heap2Local, then the flow must be from the `ref` operand. This was based
on an assumption that allocations are processed in order of appearance,
and because the `ref` operand appears before the `expected operand`, it
must be the one getting optimized. But this neglected the fact the array
allocations are processed before struct allocations, so if `expected` is
an array, it could be optimized first. This faulty assumption led to
assertion failures and invalid code.

Fix the problem by handling flows from `expected` explicitly. When a
non-escaping allocation flows into `expected`, we know it cannot
possibly match the value already in the accessed struct, so we can
optimize the `cmpxchg` to a `struct.atomic.get`.

The replacement pattern uses a scratch local to propagate value of the
`ref` operand past the dropped `expected` expression to the new
`struct.atomic.get`. In case the `ref` operand uses another allocation
that will be processed later, we must update the parents map to account
for the new data flow from the `ref` into the scratch local and from the
get of the scratch local through to the `struct.atomic.get`, which fully
consumes the value.

Author: tlively

src/ir/parents.h

@@ -17,7 +17,8 @@
 #ifndef wasm_ir_parents_h
 #define wasm_ir_parents_h
 
-#include "parsing.h"
+#include "wasm-traversal.h"
+#include "wasm.h"
 
 namespace wasm {
 
@@ -32,6 +33,10 @@ struct Parents {
     return nullptr;
   }
 
+  void setParent(Expression* child, Expression* parent) {
+    inner.parentMap[child] = parent;
+  }
+
 private:
   struct Inner
     : public ExpressionStackWalker<Inner, UnifiedExpressionVisitor<Inner>> {

src/passes/Heap2Local.cpp

@@ -157,11 +157,11 @@
 #include "ir/local-graph.h"
 #include "ir/parents.h"
 #include "ir/properties.h"
-#include "ir/type-updating.h"
 #include "ir/utils.h"
 #include "pass.h"
 #include "support/unique_deferring_queue.h"
 #include "wasm-builder.h"
+#include "wasm-type.h"
 #include "wasm.h"
 
 namespace wasm {
@@ -192,6 +192,10 @@ enum class ParentChildInteraction : int8_t {
   None,
 };
 
+// When we insert scratch locals, we sometimes need to record the flow between
+// their set and subsequent get.
+using ScratchInfo = std::unordered_map<LocalSet*, LocalGet*>;
+
 // Core analysis that provides an escapes() method to check if an allocation
 // escapes in a way that prevents optimizing it away as described above. It also
 // stashes information about the relevant expressions as it goes, which helps
@@ -201,21 +205,29 @@ struct EscapeAnalyzer {
   // parents, and via branches, and through locals.
   //
   // We use a lazy graph here because we only need this for reference locals,
-  // and even among them, only ones we see an allocation is stored to.
+  // and even among them, only ones we see an allocation is stored to. The
+  // LocalGraph is is augmented by ScratchInfo, since the LocalGraph does not
+  // know about scratch locals we add. We currently only record scratch locals
+  // that might possibly have another optimized allocation flowing through them.
+  // If it's not possible for another optimized allocation to flow through the
+  // scratch local, then we will never look at it again after creating it and do
+  // not need to record it here.
   const LazyLocalGraph& localGraph;
-  const Parents& parents;
+  ScratchInfo& scratchInfo;
+  Parents& parents;
   const BranchUtils::BranchTargets& branchTargets;
 
   const PassOptions& passOptions;
   Module& wasm;
 
   EscapeAnalyzer(const LazyLocalGraph& localGraph,
-                 const Parents& parents,
+                 ScratchInfo& scratchInfo,
+                 Parents& parents,
                  const BranchUtils::BranchTargets& branchTargets,
                  const PassOptions& passOptions,
                  Module& wasm)
-    : localGraph(localGraph), parents(parents), branchTargets(branchTargets),
-      passOptions(passOptions), wasm(wasm) {}
+    : localGraph(localGraph), scratchInfo(scratchInfo), parents(parents),
+      branchTargets(branchTargets), passOptions(passOptions), wasm(wasm) {}
 
   // We must track all the local.sets that write the allocation, to verify
   // exclusivity.
@@ -275,15 +287,23 @@ struct EscapeAnalyzer {
       }
 
       if (auto* set = parent->dynCast<LocalSet>()) {
-        // This is one of the sets we are written to, and so we must check for
-        // exclusive use of our allocation by all the gets that read the value.
-        // Note the set, and we will check the gets at the end once we know all
-        // of our sets.
-        sets.insert(set);
-
-        // We must also look at how the value flows from those gets.
-        for (auto* get : localGraph.getSetInfluences(set)) {
+
+        // We must also look at how the value flows from those gets. Check the
+        // scratchInfo first because it contains sets that localGraph doesn't
+        // know about.
+        if (auto it = scratchInfo.find(set); it != scratchInfo.end()) {
+          auto* get = it->second;
           flows.push({get, parents.getParent(get)});
+        } else {
+          // This is one of the sets we are written to, and so we must check for
+          // exclusive use of our allocation by all the gets that read the
+          // value. Note the set, and we will check the gets at the end once we
+          // know all of our sets. (For scratch locals above, we know all the
+          // sets are already accounted for.)
+          sets.insert(set);
+          for (auto* get : localGraph.getSetInfluences(set)) {
+            flows.push({get, parents.getParent(get)});
+          }
         }
       }
 
@@ -1110,17 +1130,42 @@ struct Struct2Local : PostWalker<Struct2Local> {
   }
 
   void visitStructCmpxchg(StructCmpxchg* curr) {
-    if (analyzer.getInteraction(curr->ref) != ParentChildInteraction::Flows) {
-      // The allocation can't flow into `replacement` if we've made it this far,
-      // but it might flow into `expected`, in which case we don't need to do
-      // anything because we would still be performing the cmpxchg on a real
-      // struct. We only need to replace the cmpxchg if the ref is being
-      // replaced with locals.
+    if (curr->type == Type::unreachable) {
+      // Leave this for DCE.
       return;
     }
 
-    if (curr->type == Type::unreachable) {
-      // As with RefGetDesc and StructGet, above.
+    // The allocation might flow into `ref` or `expected`, but not
+    // `replacement`, because then it would be considered to have escaped.
+    if (analyzer.getInteraction(curr->expected) ==
+        ParentChildInteraction::Flows) {
+      // Since the allocation does not escape, it cannot possibly match the
+      // value already in the struct. The cmpxchg will just do a read. Drop the
+      // other arguments and do the atomic read at the end, when the cmpxchg
+      // would have happened. Use a nullable scratch local in case we also
+      // optimize `ref` later and need to replace it with a null.
+      auto refType = curr->ref->type.with(Nullable);
+      auto refScratch = builder.addVar(func, refType);
+      auto* setRefScratch = builder.makeLocalSet(refScratch, curr->ref);
+      auto* getRefScratch = builder.makeLocalGet(refScratch, refType);
+      auto* structGet = builder.makeStructGet(
+        curr->index, getRefScratch, curr->order, curr->type);
+      auto* block = builder.makeBlock({setRefScratch,
+                                       builder.makeDrop(curr->expected),
+                                       builder.makeDrop(curr->replacement),
+                                       structGet});
+      replaceCurrent(block);
+      // Record the new data flow into and out of the new scratch local. This is
+      // necessary in case `ref` gets processed later so we can detect that it
+      // flows to the new struct.atomic.get, which may need to be replaced.
+      analyzer.parents.setParent(curr->ref, setRefScratch);
+      analyzer.scratchInfo.insert({setRefScratch, getRefScratch});
+      analyzer.parents.setParent(getRefScratch, structGet);
+      return;
+    }
+    if (analyzer.getInteraction(curr->ref) != ParentChildInteraction::Flows) {
+      // Since the allocation does not flow from `ref`, it must not flow through
+      // this cmpxchg at all.
       return;
     }
 
@@ -1468,6 +1513,7 @@ struct Heap2Local {
   const PassOptions& passOptions;
 
   LazyLocalGraph localGraph;
+  ScratchInfo scratchInfo;
   Parents parents;
   BranchUtils::BranchTargets branchTargets;
 
@@ -1539,7 +1585,7 @@ struct Heap2Local {
         continue;
       }
       EscapeAnalyzer analyzer(
-        localGraph, parents, branchTargets, passOptions, wasm);
+        localGraph, scratchInfo, parents, branchTargets, passOptions, wasm);
       if (!analyzer.escapes(allocation)) {
         // Convert the allocation and all its uses into a struct. Then convert
         // the struct into locals.
@@ -1559,7 +1605,7 @@ struct Heap2Local {
       // Check for escaping, noting relevant information as we go. If this does
       // not escape, optimize it into locals.
       EscapeAnalyzer analyzer(
-        localGraph, parents, branchTargets, passOptions, wasm);
+        localGraph, scratchInfo, parents, branchTargets, passOptions, wasm);
       if (!analyzer.escapes(allocation)) {
         Struct2Local(allocation, analyzer, func, wasm);
         optimized = true;