Optimize added constants with propagation only if we see we will remove all uses of the original add, as otherwise we may just be adding work (both an offset, and an add). Refactor local-utils.h, and make UnneededSetRemover also check for side effects, so it cleanly removes all traces of unneeded sets.

5 files changed, 201 insertions, 83 deletions

diff --git a/src/ir/count.h b/src/ir/count.h
deleted file mode 100644
index d9c89b4ce..000000000
--- a/src/ir/count.h
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Copyright 2016 WebAssembly Community Group participants
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#ifndef wasm_ir_count_h
-#define wasm_ir_count_h
-
-namespace wasm {
-
-struct GetLocalCounter : public PostWalker<GetLocalCounter> {
-  std::vector<Index> num;
-
-  GetLocalCounter() = default;
-  GetLocalCounter(Function* func) {
-    analyze(func, func->body);
-  }
-  GetLocalCounter(Function* func, Expression* ast) {
-    analyze(func, ast);
-  }
-
-  void analyze(Function* func) {
-    analyze(func, func->body);
-  }
-  void analyze(Function* func, Expression* ast) {
-    num.resize(func->getNumLocals());
-    std::fill(num.begin(), num.end(), 0);
-    walk(ast);
-  }
-
-  void visitGetLocal(GetLocal *curr) {
-    num[curr->index]++;
-  }
-};
-
-} // namespace wasm
-
-#endif // wasm_ir_count_h
-
diff --git a/src/ir/local-utils.h b/src/ir/local-utils.h
new file mode 100644
index 000000000..2feeeeb00
--- /dev/null
+++ b/src/ir/local-utils.h
@@ -0,0 +1,86 @@
+/*
+ * Copyright 2019 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_ir_local_utils_h
+#define wasm_ir_local_utils_h
+
+#include <ir/effects.h>
+
+namespace wasm {
+
+struct GetLocalCounter : public PostWalker<GetLocalCounter> {
+  std::vector<Index> num;
+
+  GetLocalCounter() = default;
+  GetLocalCounter(Function* func) {
+    analyze(func, func->body);
+  }
+  GetLocalCounter(Function* func, Expression* ast) {
+    analyze(func, ast);
+  }
+
+  void analyze(Function* func) {
+    analyze(func, func->body);
+  }
+  void analyze(Function* func, Expression* ast) {
+    num.resize(func->getNumLocals());
+    std::fill(num.begin(), num.end(), 0);
+    walk(ast);
+  }
+
+  void visitGetLocal(GetLocal *curr) {
+    num[curr->index]++;
+  }
+};
+
+struct UnneededSetRemover : public PostWalker<UnneededSetRemover> {
+  PassOptions& passOptions;
+
+  GetLocalCounter* getLocalCounter = nullptr;
+
+  UnneededSetRemover(Function* func, PassOptions& passOptions) : passOptions(passOptions) {
+    GetLocalCounter counter(func);
+    UnneededSetRemover inner(counter, func, passOptions);
+    removed = inner.removed;
+  }
+
+  UnneededSetRemover(GetLocalCounter& getLocalCounter, Function* func, PassOptions& passOptions) : passOptions(passOptions), getLocalCounter(&getLocalCounter) {
+    walk(func->body);
+  }
+
+  bool removed = false;
+
+  void visitSetLocal(SetLocal *curr) {
+    if (getLocalCounter->num[curr->index] == 0) {
+      auto* value = curr->value;
+      if (curr->isTee()) {
+        this->replaceCurrent(value);
+      } else if (EffectAnalyzer(passOptions, curr->value).hasSideEffects()) {
+        Drop* drop = ExpressionManipulator::convert<SetLocal, Drop>(curr);
+        drop->value = value;
+        drop->finalize();
+      } else {
+        ExpressionManipulator::nop(curr);
+      }
+      removed = true;
+    }
+  }
+};
+
+} // namespace wasm
+
+#endif // wasm_ir_local_utils_h
+
diff --git a/src/ir/parents.h b/src/ir/parents.h
new file mode 100644
index 000000000..71f2ae1d4
--- /dev/null
+++ b/src/ir/parents.h
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2019 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_ir_parents_h
+#define wasm_ir_parents_h
+
+namespace wasm {
+
+struct Parents {
+  Parents(Expression* expr) {
+    inner.walk(expr);
+  }
+
+  Expression* getParent(Expression* curr) {
+    return inner.parentMap[curr];
+  }
+
+private:
+  struct Inner : public ExpressionStackWalker<Inner, UnifiedExpressionVisitor<Inner>> {
+    void visitExpression(Expression* curr) {
+      parentMap[curr] = getParent();
+    }
+
+    std::map<Expression*, Expression *> parentMap;
+  } inner;
+};
+
+} // namespace wasm
+
+#endif // wasm_ir_parents_h
+
diff --git a/src/passes/OptimizeAddedConstants.cpp b/src/passes/OptimizeAddedConstants.cpp
index f6158cb2d..e2cfb1418 100644
--- a/src/passes/OptimizeAddedConstants.cpp
+++ b/src/passes/OptimizeAddedConstants.cpp
@@ -34,6 +34,8 @@
 #include <pass.h>
 #include <wasm-builder.h>
 #include <ir/local-graph.h>
+#include <ir/local-utils.h>
+#include <ir/parents.h>
 
 namespace wasm {
 
@@ -76,8 +78,9 @@ public:
         auto& sets = localGraph->getSetses[get];
         if (sets.size() == 1) {
           auto* set = *sets.begin();
-          // May be a zero-init (in which case, we can ignore it).
-          if (set) {
+          // May be a zero-init (in which case, we can ignore it). Must also be valid
+          // to propagate, as checked earlier in the parent.
+          if (set && parent->isPropagatable(set)) {
             auto* value = set->value;
             if (auto* add = value->template dynCast<Binary>()) {
               if (add->op == AddInt32) {
@@ -234,8 +237,6 @@ struct OptimizeAddedConstants : public WalkerPass<PostWalker<OptimizeAddedConsta
 
   Pass* create() override { return new OptimizeAddedConstants(propagate); }
 
-  std::unique_ptr<LocalGraph> localGraph;
-
   void visitLoad(Load* curr) {
     MemoryAccessOptimizer<OptimizeAddedConstants, Load> optimizer(this, curr, getModule(), localGraph.get());
     if (optimizer.optimize()) {
@@ -258,16 +259,23 @@ struct OptimizeAddedConstants : public WalkerPass<PostWalker<OptimizeAddedConsta
     // case (as 4 + 8 would be optimized directly if it were adjacent).
     while (1) {
       propagated = false;
+      helperIndexes.clear();
+      propagatable.clear();
       if (propagate) {
         localGraph = make_unique<LocalGraph>(func);
+        localGraph->computeInfluences();
         localGraph->computeSSAIndexes();
+        findPropagatable();
       }
       super::doWalkFunction(func);
       if (!helperIndexes.empty()) {
         createHelperIndexes();
-        helperIndexes.clear();
       }
-      if (!propagated) return;
+      if (propagated) {
+        cleanUpAfterPropagation();
+      } else {
+        return;
+      }
     }
   }
 
@@ -284,11 +292,63 @@ struct OptimizeAddedConstants : public WalkerPass<PostWalker<OptimizeAddedConsta
     return helperIndexes[set] = Builder(*getModule()).addVar(getFunction(), i32);
   }
 
-private:
-  std::map<SetLocal*, Index> helperIndexes;
+  bool isPropagatable(SetLocal* set) {
+    return propagatable.count(set);
+  }
 
+private:
   bool propagated;
 
+  std::unique_ptr<LocalGraph> localGraph;
+
+  // Whether a set is propagatable.
+  std::set<SetLocal*> propagatable;
+
+  void findPropagatable() {
+    // Conservatively, only propagate if all uses can be removed of the original. That is,
+    //  x = a + 10
+    //  f(x)
+    //  g(x)
+    // should be optimized to
+    //  f(a, offset=10)
+    //  g(a, offset=10)
+    // but if x has other uses, then avoid doing so - we'll be doing that add anyhow, so
+    // the load/store offset trick won't actually help.
+    Parents parents(getFunction()->body);
+    for (auto& pair : localGraph->locations) {
+      auto* location = pair.first;
+      if (auto* set = location->dynCast<SetLocal>()) {
+        if (auto* add = set->value->dynCast<Binary>()) {
+          if (add->op == AddInt32) {
+            if (add->left->is<Const>() || add->right->is<Const>()) {
+              // Looks like this might be relevant, check all uses.
+              bool canPropagate = true;
+              for (auto* get :localGraph->setInfluences[set]) {
+                auto* parent = parents.getParent(get);
+                assert(parent); // if this is at the top level, it's the whole body - no set can exist!
+                if (!(parent->is<Load>() || parent->is<Store>())) {
+                  canPropagate = false;
+                  break;
+                }
+              }
+              if (canPropagate) {
+                propagatable.insert(set);
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
+  void cleanUpAfterPropagation() {
+    // Remove sets that no longer have uses. This allows further propagation by letting
+    // us see the accurate amount of uses of each set.
+    UnneededSetRemover remover(getFunction(), getPassOptions());
+  }
+
+  std::map<SetLocal*, Index> helperIndexes;
+
   void createHelperIndexes() {
     struct Creator : public PostWalker<Creator> {
       std::map<SetLocal*, Index>& helperIndexes;
diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index a4931f06c..8007d92ff 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -51,7 +51,7 @@
 #include <wasm-traversal.h>
 #include <pass.h>
 #include <ir/branch-utils.h>
-#include <ir/count.h>
+#include <ir/local-utils.h>
 #include <ir/effects.h>
 #include "ir/equivalent_sets.h"
 #include <ir/find_all.h>
@@ -844,31 +844,9 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals<a
     // We may have already had a local with no uses, or we may have just
     // gotten there thanks to the EquivalentOptimizer. If there are such
     // locals, remove all their sets.
-    struct UneededSetRemover : public PostWalker<UneededSetRemover> {
-      std::vector<Index>* numGetLocals;
-
-      bool anotherCycle = false;
-
-      void visitSetLocal(SetLocal *curr) {
-        if ((*numGetLocals)[curr->index] == 0) {
-          auto* value = curr->value;
-          if (curr->isTee()) {
-            this->replaceCurrent(value);
-          } else {
-            Drop* drop = ExpressionManipulator::convert<SetLocal, Drop>(curr);
-            drop->value = value;
-            drop->finalize();
-          }
-          anotherCycle = true;
-        }
-      }
-    };
-
-    UneededSetRemover setRemover;
-    setRemover.numGetLocals = &getCounter.num;
-    setRemover.walkFunction(func);
+    UnneededSetRemover setRemover(getCounter, func, this->getPassOptions());
 
-    return eqOpter.anotherCycle || setRemover.anotherCycle;
+    return eqOpter.anotherCycle || setRemover.removed;
   }
 };