Improve local-cse (#1594)

This makes it much more effective, by rewriting it to depend on flatten. In flattened IR, it is very simple to check if an expression is equivalent to one already available for use in a local, and use that one instead, basically we just track values in locals.

Helps with #1521 

1 files changed, 77 insertions, 34 deletions

diff --git a/src/passes/LocalCSE.cpp b/src/passes/LocalCSE.cpp
index 7524ad0fa..a458937db 100644
--- a/src/passes/LocalCSE.cpp
+++ b/src/passes/LocalCSE.cpp
@@ -17,6 +17,15 @@
 //
 // Local CSE
 //
+// This requires --flatten to be run before in order to be effective,
+// and preserves flatness. The reason flatness is required is that
+// this pass assumes everything is stored in a local, and all it does
+// is alter set_locals to do get_locals of an existing value when
+// possible, replacing a recomputing of that value. That design means that
+// if there are block and if return values, nested expressions not stored
+// to a local, etc., then it can't operate on them (and will just not
+// do anything for them).
+//
 // In each linear area of execution,
 //  * track each relevant (big enough) expression
 //  * if already seen, write to a local if not already, and reuse
@@ -25,17 +34,19 @@
 // TODO: global, inter-block gvn etc.
 //
 
+#include <algorithm>
+#include <memory>
+
 #include <wasm.h>
 #include <wasm-builder.h>
 #include <wasm-traversal.h>
 #include <pass.h>
 #include <ir/effects.h>
+#include <ir/equivalent_sets.h>
 #include <ir/hashed.h>
 
 namespace wasm {
 
-const Index UNUSED = -1;
-
 struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
   bool isFunctionParallel() override { return true; }
 
@@ -43,11 +54,11 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
 
   // information for an expression we can reuse
   struct UsableInfo {
-    Expression** item;
-    Index index; // if not UNUSED, then the local we are assigned to, use that to reuse us
+    Expression* value; // the value we can reuse
+    Index index; // the local we are assigned to, get_local that to reuse us
     EffectAnalyzer effects;
 
-    UsableInfo(Expression** item, PassOptions& passOptions) : item(item), index(UNUSED), effects(passOptions, *item) {}
+    UsableInfo(Expression* value, Index index, PassOptions& passOptions) : value(value), index(index), effects(passOptions, value) {}
   };
 
   // a list of usables in a linear execution trace
@@ -56,8 +67,29 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
   // locals in current linear execution trace, which we try to sink
   Usables usables;
 
+  // We track locals containing the same value - the value is what matters, not
+  // the index.
+  EquivalentSets equivalences;
+
+  bool anotherPass;
+
+  void doWalkFunction(Function* func) {
+    anotherPass = true;
+    // we may need multiple rounds
+    while (anotherPass) {
+      anotherPass = false;
+      clear();
+      super::doWalkFunction(func);
+    }
+  }
+
   static void doNoteNonLinear(LocalCSE* self, Expression** currp) {
-    self->usables.clear();
+    self->clear();
+  }
+
+  void clear() {
+    usables.clear();
+    equivalences.clear();
   }
 
   void checkInvalidations(EffectAnalyzer& effects) {
@@ -91,9 +123,7 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
     auto* curr = *currp;
 
     // main operations
-    if (self->isRelevant(curr)) {
-      self->handle(currp, curr);
-    }
+    self->handle(curr);
 
     // post operations
 
@@ -114,38 +144,51 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
     self->pushTask(visitPre, currp);
   }
 
-  bool isRelevant(Expression* curr) {
-    if (curr->is<GetLocal>()) {
+  void handle(Expression* curr) {
+    if (auto* set = curr->dynCast<SetLocal>()) {
+      // Calculate equivalences
+      equivalences.reset(set->index);
+      if (auto* get = set->value->dynCast<GetLocal>()) {
+        equivalences.add(set->index, get->index);
+      }
+      // consider the value
+      auto* value = set->value;
+      if (isRelevant(value)) {
+        HashedExpression hashed(value);
+        auto iter = usables.find(hashed);
+        if (iter != usables.end()) {
+          // already exists in the table, this is good to reuse
+          auto& info = iter->second;
+          set->value = Builder(*getModule()).makeGetLocal(info.index, value->type);
+          anotherPass = true;
+        } else {
+          // not in table, add this, maybe we can help others later
+          usables.emplace(std::make_pair(hashed, UsableInfo(value, set->index, getPassOptions())));
+        }
+      }
+    } else if (auto* get = curr->dynCast<GetLocal>()) {
+      if (auto* set = equivalences.getEquivalents(get->index)) {
+        // Canonicalize to the lowest index. This lets hashing and comparisons
+        // "just work".
+        get->index = *std::min_element(set->begin(), set->end());
+      }
+    }
+  }
+
+  // A relevant value is a non-trivial one, something we may want to reuse
+  // and are able to.
+  bool isRelevant(Expression* value) {
+    if (value->is<GetLocal>()) {
       return false; // trivial, this is what we optimize to!
     }
-    if (!isConcreteType(curr->type)) {
+    if (!isConcreteType(value->type)) {
       return false; // don't bother with unreachable etc.
     }
-    if (EffectAnalyzer(getPassOptions(), curr).hasSideEffects()) {
+    if (EffectAnalyzer(getPassOptions(), value).hasSideEffects()) {
       return false; // we can't combine things with side effects
     }
     // check what we care about TODO: use optimize/shrink levels?
-    return Measurer::measure(curr) > 1;
-  }
-
-  void handle(Expression** currp, Expression* curr) {
-    HashedExpression hashed(curr);
-    auto iter = usables.find(hashed);
-    if (iter != usables.end()) {
-      // already exists in the table, this is good to reuse
-      auto& info = iter->second;
-      if (info.index == UNUSED) {
-        // we need to assign to a local. create a new one
-        auto index = info.index = Builder::addVar(getFunction(), curr->type);
-        (*info.item) = Builder(*getModule()).makeTeeLocal(index, *info.item);
-      }
-      replaceCurrent(
-        Builder(*getModule()).makeGetLocal(info.index, curr->type)
-      );
-    } else {
-      // not in table, add this, maybe we can help others later
-      usables.emplace(std::make_pair(hashed, UsableInfo(currp, getPassOptions())));
-    }
+    return Measurer::measure(value) > 1;
   }
 };