[NFC] Make Precompute use a lazy LocalGraph (#6934)

To do this, add locations and getInfluences to LazyLocalGraph. Both cannot
really be computed in a fine-grained manner, so just compute them all on the
first request. That is not as efficient as our lazy computation of getSets and
setInfluences, but they are also less important, and this change makes the
pass 20% faster.

3 files changed, 76 insertions, 10 deletions

diff --git a/src/ir/LocalGraph.cpp b/src/ir/LocalGraph.cpp
index af1370731..a5495b5b3 100644
--- a/src/ir/LocalGraph.cpp
+++ b/src/ir/LocalGraph.cpp
@@ -281,6 +281,8 @@ struct LocalGraphFlower
                                       });
           if (lastSet != pred->lastSets.end()) {
             // There is a set here, apply it, and stop the flow.
+            // TODO: If we find a computed get, apply its sets and stop? That
+            //       could help but it requires more info on FlowBlock.
             for (auto* get : gets) {
               getSetsMap[get].insert(lastSet->second);
             }
@@ -512,7 +514,9 @@ void LocalGraph::computeSetInfluences() {
   }
 }
 
-void LocalGraph::computeGetInfluences() {
+static void
+doComputeGetInfluences(const LocalGraphBase::Locations& locations,
+                       LocalGraphBase::GetInfluencesMap& getInfluences) {
   for (auto& [curr, _] : locations) {
     if (auto* set = curr->dynCast<LocalSet>()) {
       FindAll<LocalGet> findAll(set->value);
@@ -523,6 +527,10 @@ void LocalGraph::computeGetInfluences() {
   }
 }
 
+void LocalGraph::computeGetInfluences() {
+  doComputeGetInfluences(locations, getInfluences);
+}
+
 void LocalGraph::computeSSAIndexes() {
   std::unordered_map<Index, std::set<LocalSet*>> indexSets;
   for (auto& [get, sets] : getSetsMap) {
@@ -555,13 +563,12 @@ LazyLocalGraph::LazyLocalGraph(Function* func, Module* module)
   : LocalGraphBase(func, module) {}
 
 void LazyLocalGraph::makeFlower() const {
-  // Lazy graphs do not provide |locations| publicly. TODO: perhaps refactor to
-  // avoid filling in this dummy data structure, but we may want to add a lazy
-  // version of it too, so see which makes sense first.
-  LocalGraph::Locations locations;
+  // |locations| is set here and filled in by |flower|.
+  assert(!locations);
+  locations.emplace();
 
   flower =
-    std::make_unique<LocalGraphFlower>(getSetsMap, locations, func, module);
+    std::make_unique<LocalGraphFlower>(getSetsMap, *locations, func, module);
 
   flower->prepareLaziness();
 
@@ -585,6 +592,9 @@ LazyLocalGraph::~LazyLocalGraph() {
 }
 
 void LazyLocalGraph::computeGetSets(LocalGet* get) const {
+  // We must never repeat work.
+  assert(!getSetsMap.count(get));
+
   if (!flower) {
     makeFlower();
   }
@@ -592,10 +602,40 @@ void LazyLocalGraph::computeGetSets(LocalGet* get) const {
 }
 
 void LazyLocalGraph::computeSetInfluences(LocalSet* set) const {
+  // We must never repeat work.
+  assert(!setInfluences.count(set));
+
   if (!flower) {
     makeFlower();
   }
   flower->computeSetInfluences(set, setInfluences);
 }
 
+void LazyLocalGraph::computeGetInfluences() const {
+  // We must never repeat work.
+  assert(!getInfluences);
+
+  // We do not need any flow for this, but we do need |locations| to be filled
+  // in.
+  getLocations();
+  assert(locations);
+
+  getInfluences.emplace();
+  doComputeGetInfluences(*locations, *getInfluences);
+}
+
+void LazyLocalGraph::computeLocations() const {
+  // We must never repeat work.
+  assert(!locations);
+
+  // |flower| fills in |locations| as it scans the function.
+  //
+  // In theory we could be even lazier here, but it is nice that flower will
+  // fill in the locations as it goes, avoiding an additional pass. And, in
+  // practice, if we ask for locations then we likely need other things anyhow.
+  if (!flower) {
+    makeFlower();
+  }
+}
+
 } // namespace wasm
diff --git a/src/ir/local-graph.h b/src/ir/local-graph.h
index 46bc135e0..6f1e621cf 100644
--- a/src/ir/local-graph.h
+++ b/src/ir/local-graph.h
@@ -202,17 +202,43 @@ struct LazyLocalGraph : public LocalGraphBase {
     }
     return iter->second;
   }
+  const GetInfluences& getGetInfluences(LocalGet* get) const {
+    if (!getInfluences) {
+      computeGetInfluences();
+      assert(getInfluences);
+    }
+    return (*getInfluences)[get];
+  }
+
+  const Locations& getLocations() const {
+    if (!locations) {
+      computeLocations();
+      assert(locations);
+    }
+    return *locations;
+  }
 
 private:
   // These data structures are mutable so that we can memoize.
   mutable GetSetsMap getSetsMap;
 
   mutable SetInfluencesMap setInfluences;
+  // The entire |getInfluences| is computed once the first request for one
+  // arrives, so the entire thing is either present or not, unlike setInfluences
+  // which is fine-grained. The difference is that the influences of a get may
+  // include sets of other indexes, so there is no simple way to lazify that
+  // computation.
+  mutable std::optional<GetInfluencesMap> getInfluences;
+  mutable std::optional<Locations> locations;
 
   // Compute the sets for a get and store them on getSetsMap.
   void computeGetSets(LocalGet* get) const;
   // Compute influences for a set and store them on setInfluences.
   void computeSetInfluences(LocalSet* set) const;
+  // Compute influences for all gets and store them on getInfluences.
+  void computeGetInfluences() const;
+  // Compute locations and store them on getInfluences.
+  void computeLocations() const;
 
   // This remains alive as long as we are, so that we can compute things lazily.
   // It is mutable as when we construct this is an internal detail, that does
diff --git a/src/passes/Precompute.cpp b/src/passes/Precompute.cpp
index 0b8ead056..709fd7d3d 100644
--- a/src/passes/Precompute.cpp
+++ b/src/passes/Precompute.cpp
@@ -749,9 +749,9 @@ private:
     // 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). First, compute all influences/dependencies.
-    LocalGraph localGraph(func, getModule());
-    localGraph.computeInfluences();
+    // constant). We do this lazily as most locals do not end up with constant
+    // values that we can propagate.
+    LazyLocalGraph localGraph(func, getModule());
 
     // A map of sets to their constant values. If a set does not appear here
     // then it is not constant, like |getValues|.
@@ -875,7 +875,7 @@ private:
 
     // 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) {
+    for (auto& [curr, _] : localGraph.getLocations()) {
       if (auto* set = curr->dynCast<LocalSet>()) {
         checkConstantSet(set);
       } else {