[NFC] Optimize propagateLocals() (#6928)

That is the slowest part of --precompute-propagate, which is one of our slowest
passes. This makes that pass 25% faster.

The main change is to make the maps consider missing elements as non-constant,
rather than storing a None element in them. That saves allocating entries for the
common case of a non-constant set/get.

Also switch to a simple vector for the work queue, which is possible if we only
add work when a get/set becomes constant (which can only happen once). Another
benefit to adding work in that manner is that we don't start by adding every single
get/set as "work" at the start.

1 files changed, 156 insertions, 116 deletions

diff --git a/src/passes/Precompute.cpp b/src/passes/Precompute.cpp
index 780945b46..379f6ee49 100644
--- a/src/passes/Precompute.cpp
+++ b/src/passes/Precompute.cpp
@@ -35,13 +35,16 @@
 #include "ir/utils.h"
 #include "pass.h"
 #include "support/insert_ordered.h"
-#include "support/unique_deferring_queue.h"
+#include "support/small_vector.h"
 #include "wasm-builder.h"
 #include "wasm-interpreter.h"
 #include "wasm.h"
 
 namespace wasm {
 
+// A map of gets to their constant values. If a get does not have a constant
+// value then it does not appear in the map (that avoids allocating for the
+// majority of gets).
 using GetValues = std::unordered_map<LocalGet*, Literals>;
 
 // A map of values on the heap. This maps the expressions that create the
@@ -81,7 +84,7 @@ class PrecomputingExpressionRunner
 
   // Concrete values of gets computed during the pass, which the runner does not
   // know about since it only records values of sets it visits.
-  GetValues& getValues;
+  const GetValues& getValues;
 
   HeapValues& heapValues;
 
@@ -99,7 +102,7 @@ class PrecomputingExpressionRunner
 
 public:
   PrecomputingExpressionRunner(Module* module,
-                               GetValues& getValues,
+                               const GetValues& getValues,
                                HeapValues& heapValues,
                                bool replaceExpression)
     : ConstantExpressionRunner<PrecomputingExpressionRunner>(
@@ -114,9 +117,8 @@ public:
     auto iter = getValues.find(curr);
     if (iter != getValues.end()) {
       auto values = iter->second;
-      if (values.isConcrete()) {
-        return Flow(values);
-      }
+      assert(values.isConcrete());
+      return Flow(values);
     }
     return ConstantExpressionRunner<
       PrecomputingExpressionRunner>::visitLocalGet(curr);
@@ -743,131 +745,169 @@ private:
   // Propagates values around. Returns whether we propagated.
   bool propagateLocals(Function* func) {
     bool propagated = false;
-    // using the graph of get-set interactions, do a constant-propagation type
+
+    // Using the graph of get-set interactions, do a constant-propagation type
     // operation: note which sets are assigned locals, then see if that lets us
     // compute other sets as locals (since some of the gets they read may be
-    // constant).
-    // compute all dependencies
+    // constant). First, compute all influences/dependencies.
     LocalGraph localGraph(func, getModule());
     localGraph.computeInfluences();
-    // prepare the work list. we add things here that might change to a constant
-    // initially, that means everything
-    UniqueDeferredQueue<Expression*> work;
-    for (auto& [curr, _] : localGraph.locations) {
-      work.push(curr);
-    }
-    // the constant value, or none if not a constant
+
+    // A map of sets to their constant values. If a set does not appear here
+    // then it is not constant, like |getValues|.
     std::unordered_map<LocalSet*, Literals> setValues;
-    // propagate constant values
-    while (!work.empty()) {
-      auto* curr = work.pop();
-      // see if this set or get is actually a constant value, and if so,
-      // mark it as such and add everything it influences to the work list,
-      // as they may be constant too.
-      if (auto* set = curr->dynCast<LocalSet>()) {
-        if (setValues[set].isConcrete()) {
-          continue; // already known constant
-        }
-        // Precompute the value. Note that this executes the code from scratch
-        // each time we reach this point, and so we need to be careful about
-        // repeating side effects if those side effects are expressed *in the
-        // value*. A case where that can happen is GC data (each struct.new
-        // creates a new, unique struct, even if the data is equal), and so
-        // PrecomputingExpressionRunner has special logic to make sure that
-        // reference identity is preserved properly.
-        //
-        // (Other side effects are fine; if an expression does a call and we
-        // somehow know the entire expression precomputes to a 42, then we can
-        // propagate that 42 along to the users, regardless of whatever the call
-        // did globally.)
-        auto values = precomputeValue(Properties::getFallthrough(
-          set->value, getPassOptions(), *getModule()));
-        // Fix up the value. The computation we just did was to look at the
-        // fallthrough, then precompute that; that looks through expressions
-        // that pass through the value. Normally that does not matter here,
-        // for example, (block .. (value)) returns the value unmodified.
-        // However, some things change the type, for example RefAsNonNull has
-        // a non-null type, while its input may be nullable. That does not
-        // matter either, as if we managed to precompute it then the value had
-        // the more specific (in this example, non-nullable) type. But there
-        // is a situation where this can cause an issue: RefCast. An attempt to
-        // perform a "bad" cast, say of a function to a struct, is a case where
-        // the fallthrough value's type is very different than the actually
-        // returned value's type. To handle that, if we precomputed a value and
-        // if it has the wrong type then precompute it again without looking
-        // through to the fallthrough.
-        if (values.isConcrete() &&
-            !Type::isSubType(values.getType(), set->value->type)) {
-          values = precomputeValue(set->value);
-        }
+
+    // The work list, which will contain sets and gets that have just been
+    // found to have a constant value. As we only add them to the list when they
+    // are found to be constant, each can only be added once, and a simple
+    // vector is enough here (which we can make a small vector to avoid any
+    // allocation in small-enough functions).
+    SmallVector<Expression*, 10> work;
+
+    // Given a set, see if it has a constant value. If so, note that on
+    // setValues and add to the work list.
+    auto checkConstantSet = [&](LocalSet* set) {
+      if (setValues.count(set)) {
+        // Already known to be constant.
+        return;
+      }
+
+      // Precompute the value. Note that this executes the code from scratch
+      // each time we reach this point, and so we need to be careful about
+      // repeating side effects if those side effects are expressed *in the
+      // value*. A case where that can happen is GC data (each struct.new
+      // creates a new, unique struct, even if the data is equal), and so
+      // PrecomputingExpressionRunner has special logic to make sure that
+      // reference identity is preserved properly.
+      //
+      // (Other side effects are fine; if an expression does a call and we
+      // somehow know the entire expression precomputes to a 42, then we can
+      // propagate that 42 along to the users, regardless of whatever the call
+      // did globally.)
+      auto values = precomputeValue(
+        Properties::getFallthrough(set->value, getPassOptions(), *getModule()));
+
+      // Fix up the value. The computation we just did was to look at the
+      // fallthrough, then precompute that; that looks through expressions
+      // that pass through the value. Normally that does not matter here,
+      // for example, (block .. (value)) returns the value unmodified.
+      // However, some things change the type, for example RefAsNonNull has
+      // a non-null type, while its input may be nullable. That does not
+      // matter either, as if we managed to precompute it then the value had
+      // the more specific (in this example, non-nullable) type. But there
+      // is a situation where this can cause an issue: RefCast. An attempt to
+      // perform a "bad" cast, say of a null to non-null, is a tricky case where
+      // the fallthrough value's type is very different than the actually
+      // returned value's type. To handle that, if we precomputed a value and
+      // if it has the wrong type then precompute it again without looking
+      // through to the fallthrough.
+      if (values.isConcrete() &&
+          !Type::isSubType(values.getType(), set->value->type)) {
+        values = precomputeValue(set->value);
+      }
+
+      if (values.isConcrete()) {
         setValues[set] = values;
-        if (values.isConcrete()) {
-          for (auto* get : localGraph.getSetInfluences(set)) {
-            work.push(get);
-          }
-        }
-      } else {
-        auto* get = curr->cast<LocalGet>();
-        if (getValues[get].size() >= 1) {
-          continue; // already known constant
-        }
-        // for this get to have constant value, all sets must agree
-        Literals values;
-        bool first = true;
-        for (auto* set : localGraph.getSets(get)) {
-          Literals curr;
-          if (set == nullptr) {
-            if (getFunction()->isVar(get->index)) {
-              auto localType = getFunction()->getLocalType(get->index);
-              if (!localType.isDefaultable()) {
-                // This is a nondefaultable local that seems to read the default
-                // value at the function entry. This is either an internal error
-                // or a case of unreachable code; the latter is possible as
-                // LocalGraph is not precise in unreachable code.
-                //
-                // We cannot set zeros here (as applying them, even in
-                // unreachable code, would not validate), so just mark this as
-                // a hopeless case to ignore.
-                values = {};
-              } else {
-                curr = Literal::makeZeros(localType);
-              }
+        work.push_back(set);
+      }
+    };
+
+    // The same, for a get.
+    auto checkConstantGet = [&](LocalGet* get) {
+      if (getValues.count(get)) {
+        // Already known to be constant.
+        return;
+      }
+
+      // For this get to have constant value, all sets must agree on a constant.
+      Literals values;
+      bool first = true;
+      for (auto* set : localGraph.getSets(get)) {
+        Literals curr;
+        if (set == nullptr) {
+          if (getFunction()->isVar(get->index)) {
+            auto localType = getFunction()->getLocalType(get->index);
+            if (!localType.isDefaultable()) {
+              // This is a nondefaultable local that seems to read the default
+              // value at the function entry. This is either an internal error
+              // or a case of unreachable code; the latter is possible as
+              // LocalGraph is not precise in unreachable code. Give up.
+              return;
             } else {
-              // it's a param, so it's hopeless
-              values = {};
-              break;
+              curr = Literal::makeZeros(localType);
             }
           } else {
-            curr = setValues[set];
-          }
-          if (curr.isNone()) {
-            // not a constant, give up
-            values = {};
-            break;
+            // It's a param, so the value is non-constant. Give up.
+            return;
           }
-          // we found a concrete value. compare with the current one
-          if (first) {
-            values = curr; // this is the first
-            first = false;
-          } else {
-            if (values != curr) {
-              // not the same, give up
-              values = {};
-              break;
-            }
+        } else {
+          // If there is an entry for the set, use that constant. Otherwise, the
+          // set is not constant, and we give up.
+          auto iter = setValues.find(set);
+          if (iter == setValues.end()) {
+            return;
           }
+          curr = iter->second;
         }
-        // we may have found a value
-        if (values.isConcrete()) {
-          // we did!
-          getValues[get] = values;
-          for (auto* set : localGraph.getGetInfluences(get)) {
-            work.push(set);
-          }
-          propagated = true;
+
+        // We found a concrete value, so there is a chance, if it matches all
+        // the rest.
+        assert(curr.isConcrete());
+        if (first) {
+          // This is the first ever value we see. All later ones must match it.
+          values = curr;
+          first = false;
+        } else if (values != curr) {
+          // This later value is not the same as before, give up.
+          return;
+        }
+      }
+
+      if (values.isConcrete()) {
+        // We found a constant value!
+        getValues[get] = values;
+        work.push_back(get);
+        propagated = true;
+      } else {
+        // If it is not concrete then, since we early-exited before on any
+        // possible problem, there must be no sets for this get, which means it
+        // is in unreachable code. In that case, we never switched |first| from
+        // true to false.
+        assert(first == true);
+        // We could optimize using unreachability here, but we leave that for
+        // other passes.
+      }
+    };
+
+    // Check all gets and sets to find which are constant, mark them as such,
+    // and add propagation work based on that.
+    for (auto& [curr, _] : localGraph.locations) {
+      if (auto* set = curr->dynCast<LocalSet>()) {
+        checkConstantSet(set);
+      } else {
+        checkConstantGet(curr->cast<LocalGet>());
+      }
+    }
+
+    // Propagate constant values while work remains.
+    while (!work.empty()) {
+      auto* curr = work.back();
+      work.pop_back();
+
+      // This get or set is a constant value. Check if the things it influences
+      // become constant.
+      if (auto* set = curr->dynCast<LocalSet>()) {
+        for (auto* get : localGraph.getSetInfluences(set)) {
+          checkConstantGet(get);
+        }
+      } else {
+        auto* get = curr->cast<LocalGet>();
+        for (auto* set : localGraph.getGetInfluences(get)) {
+          checkConstantSet(set);
         }
       }
     }
+
     return propagated;
   }