Improve LocalGraph (#1382)

This simplifies the logic there into a more standard flow operation. This is not always faster, but it is much faster on the worst cases we saw before like sqlite, and it is simpler.

The rewrite also fixes a fuzz bug.

1 files changed, 164 insertions, 211 deletions

diff --git a/src/ir/LocalGraph.cpp b/src/ir/LocalGraph.cpp
index 53c6a52e2..77f5906f0 100644
--- a/src/ir/LocalGraph.cpp
+++ b/src/ir/LocalGraph.cpp
@@ -20,250 +20,203 @@
 #include <wasm-printing.h>
 #include <ir/find_all.h>
 #include <ir/local-graph.h>
+#include <cfg/cfg-traversal.h>
 
 namespace wasm {
 
-LocalGraph::LocalGraph(Function* func, Module* module) {
-  walkFunctionInModule(func, module);
+namespace LocalGraphInternal {
+
+// A relevant action. Supports a get, a set, or an
+// "other" which can be used for other purposes, to mark
+// their position in a block
+struct Action {
+  enum What {
+    Get = 0,
+    Set = 1
+  };
+  What what;
+  Index index; // the local index read or written
+  Expression* expr; // the expression itself
+
+  Action(What what, Index index, Expression* expr) : what(what), index(index), expr(expr) {
+    if (what == Get) assert(expr->is<GetLocal>());
+    if (what == Set) assert(expr->is<SetLocal>());
+  }
 
-#ifdef LOCAL_GRAPH_DEBUG
-  std::cout << "LocalGraph::dump\n";
-  for (auto& pair : getSetses) {
-    auto* get = pair.first;
-    auto& sets = pair.second;
-    std::cout << "GET\n" << get << " is influenced by\n";
-    for (auto* set : sets) {
-      std::cout << set << '\n';
+  bool isGet() { return what == Get; }
+  bool isSet() { return what == Set; }
+};
+
+// information about a basic block
+struct Info {
+  std::vector<Action> actions; // actions occurring in this block
+  std::vector<SetLocal*> lastSets; // for each index, the last set_local for it
+
+  void dump(Function* func) {
+    if (actions.empty()) return;
+    std::cout << "    actions:\n";
+    for (auto& action : actions) {
+      std::cout << "      " << (action.isGet() ? "get" : "set") << " " << func->getLocalName(action.index) << "\n";
     }
   }
-#endif
-}
+};
 
-void LocalGraph::computeInfluences() {
-  for (auto& pair : locations) {
-    auto* curr = pair.first;
-    if (auto* set = curr->dynCast<SetLocal>()) {
-      FindAll<GetLocal> findAll(set->value);
-      for (auto* get : findAll.list) {
-        getInfluences[get].insert(set);
-      }
-    } else {
-      auto* get = curr->cast<GetLocal>();
-      for (auto* set : getSetses[get]) {
-        setInfluences[set].insert(get);
-      }
-    }
+// flow helper class. flows the gets to their sets
+
+struct Flower : public CFGWalker<Flower, Visitor<Flower>, Info> {
+  LocalGraph::GetSetses& getSetses;
+  LocalGraph::Locations& locations;
+
+  Flower(LocalGraph::GetSetses& getSetses, LocalGraph::Locations& locations, Function* func) : getSetses(getSetses), locations(locations) {
+    setFunction(func);
+    // create the CFG by walking the IR
+    CFGWalker<Flower, Visitor<Flower>, Info>::doWalkFunction(func);
+    // flow gets across blocks
+    flow(func);
   }
-}
 
-void LocalGraph::doWalkFunction(Function* func) {
-  numLocals = func->getNumLocals();
-  if (numLocals == 0) return; // nothing to do
-  // We begin with each param being assigned from the incoming value, and the zero-init for the locals,
-  // so the initial state is the identity permutation
-  currMapping.resize(numLocals);
-  for (auto& set : currMapping) {
-    set = { nullptr };
+  BasicBlock* makeBasicBlock() {
+    auto* ret = new BasicBlock();
+    auto& lastSets = ret->contents.lastSets;
+    lastSets.resize(getFunction()->getNumLocals());
+    std::fill(lastSets.begin(), lastSets.end(), nullptr);
+    return ret;
   }
-  PostWalker<LocalGraph>::walk(func->body);
-}
 
-// control flow
+  // cfg traversal work
 
-void LocalGraph::visitBlock(Block* curr) {
-  if (curr->name.is() && breakMappings.find(curr->name) != breakMappings.end()) {
-    auto& infos = breakMappings[curr->name];
-    infos.emplace_back(std::move(currMapping));
-    currMapping = std::move(merge(infos));
-    breakMappings.erase(curr->name);
+  static void doVisitGetLocal(Flower* self, Expression** currp) {
+    auto* curr = (*currp)->cast<GetLocal>();
+     // if in unreachable code, skip
+    if (!self->currBasicBlock) return;
+    self->currBasicBlock->contents.actions.emplace_back(Action::Get, curr->index, curr);
+    self->locations[curr] = currp;
   }
-}
 
-void LocalGraph::finishIf() {
-  // that's it for this if, merge
-  std::vector<Mapping> breaks;
-  breaks.emplace_back(std::move(currMapping));
-  breaks.emplace_back(std::move(mappingStack.back()));
-  mappingStack.pop_back();
-  currMapping = std::move(merge(breaks));
-}
-
-void LocalGraph::afterIfCondition(LocalGraph* self, Expression** currp) {
-  self->mappingStack.push_back(self->currMapping);
-}
-void LocalGraph::afterIfTrue(LocalGraph* self, Expression** currp) {
-  auto* curr = (*currp)->cast<If>();
-  if (curr->ifFalse) {
-    auto afterCondition = std::move(self->mappingStack.back());
-    self->mappingStack.back() = std::move(self->currMapping);
-    self->currMapping = std::move(afterCondition);
-  } else {
-    self->finishIf();
+  static void doVisitSetLocal(Flower* self, Expression** currp) {
+    auto* curr = (*currp)->cast<SetLocal>();
+    // if in unreachable code, skip
+    if (!self->currBasicBlock) return;
+    self->currBasicBlock->contents.actions.emplace_back(Action::Set, curr->index, curr);
+    self->currBasicBlock->contents.lastSets[curr->index] = curr;
+    self->locations[curr] = currp;
   }
-}
-void LocalGraph::afterIfFalse(LocalGraph* self, Expression** currp) {
-  self->finishIf();
-}
-void LocalGraph::beforeLoop(LocalGraph* self, Expression** currp) {
-  // save the state before entering the loop, for calculation later of the merge at the loop top
-  self->mappingStack.push_back(self->currMapping);
-  self->loopGetStack.push_back({});
-}
-void LocalGraph::visitLoop(Loop* curr) {
-  if (curr->name.is() && breakMappings.find(curr->name) != breakMappings.end()) {
-    auto& infos = breakMappings[curr->name];
-    infos.emplace_back(std::move(mappingStack.back()));
-    auto before = infos.back();
-    auto& merged = merge(infos);
-    // every local we created a phi for requires us to update get_local operations in
-    // the loop - the branch back has means that gets in the loop have potentially
-    // more sets reaching them.
-    // we can detect this as follows: if a get of oldIndex has the same sets
-    // as the sets at the entrance to the loop, then it is affected by the loop
-    // header sets, and we can add to there sets that looped back
-    auto linkLoopTop = [&](Index i, Sets& getSets) {
-      auto& beforeSets = before[i];
-      if (getSets.size() < beforeSets.size()) {
-        // the get trivially has fewer sets, so it overrode the loop entry sets
-        return;
+
+  void flow(Function* func) {
+    auto numLocals = func->getNumLocals();
+    std::vector<std::vector<GetLocal*>> allGets;
+    allGets.resize(numLocals);
+    std::unordered_set<BasicBlock*> seen;
+    std::vector<BasicBlock*> work;
+    for (auto& block : basicBlocks) {
+#ifdef LOCAL_GRAPH_DEBUG
+      std::cout << "basic block " << block.get() << " :\n";
+      for (auto& action : block->contents.actions) {
+        std::cout << "  action: " << action.expr << '\n';
       }
-      if (!std::includes(getSets.begin(), getSets.end(),
-                         beforeSets.begin(), beforeSets.end())) {
-        // the get has not the same sets as in the loop entry
-        return;
+      for (auto* lastSet : block->contents.lastSets) {
+        std::cout << "  last set " << lastSet << '\n';
       }
-      // the get has the entry sets, so add any new ones
-      for (auto* set : merged[i]) {
-        getSets.insert(set);
+#endif
+      // go through the block, finding each get and adding it to its index,
+      // and seeing how sets affect that
+      auto& actions = block->contents.actions;
+      // move towards the front, handling things as we go
+      for (int i = int(actions.size()) - 1; i >= 0; i--) {
+        auto& action = actions[i];
+        auto index = action.index;
+        if (action.isGet()) {
+          allGets[index].push_back(action.expr->cast<GetLocal>());
+        } else {
+          // this set is the only set for all those gets
+          auto* set = action.expr->cast<SetLocal>();
+          auto& gets = allGets[index];
+          for (auto* get : gets) {
+            getSetses[get].insert(set);
+          }
+          gets.clear();
+        }
       }
-    };
-    auto& gets = loopGetStack.back();
-    for (auto* get : gets) {
-      linkLoopTop(get->index, getSetses[get]);
-    }
-    // and the same for the loop fallthrough: any local that still has the
-    // entry sets should also have the loop-back sets as well
-    for (Index i = 0; i < numLocals; i++) {
-      linkLoopTop(i, currMapping[i]);
-    }
-    // finally, breaks still in flight must be updated too
-    for (auto& iter : breakMappings) {
-      auto name = iter.first;
-      if (name == curr->name) continue; // skip our own (which is still in use)
-      auto& mappings = iter.second;
-      for (auto& mapping : mappings) {
-        for (Index i = 0; i < numLocals; i++) {
-          linkLoopTop(i, mapping[i]);
+      // if anything is left, we must flow it back through other blocks. we
+      // can do that for all gets as a whole, they will get the same results
+      for (Index index = 0; index < numLocals; index++) {
+        auto& gets = allGets[index];
+        if (gets.empty()) continue;
+        work.push_back(block.get());
+        seen.clear();
+        // note that we may need to revisit the later parts of this initial
+        // block, if we are in a loop, so don't mark it as seen
+        while (!work.empty()) {
+          auto* curr = work.back();
+          work.pop_back();
+          // we have gone through this block; now we must handle flowing to
+          // the inputs
+          if (curr->in.empty()) {
+            if (curr == entry) {
+              // these receive a param or zero init value
+              for (auto* get : gets) {
+                getSetses[get].insert(nullptr);
+              }
+            }
+          } else {
+            for (auto* pred : curr->in) {
+              if (seen.count(pred)) continue;
+              seen.insert(pred);
+              auto* lastSet = pred->contents.lastSets[index];
+              if (lastSet) {
+                // there is a set here, apply it, and stop the flow
+                for (auto* get : gets) {
+                  getSetses[get].insert(lastSet);
+                }
+              } else {
+                // keep on flowing
+                work.push_back(pred);
+              }
+            }
+          }
         }
+        gets.clear();
       }
     }
-    // now that we are done with using the mappings, erase our own
-    breakMappings.erase(curr->name);
-  }
-  mappingStack.pop_back();
-  loopGetStack.pop_back();
-}
-void LocalGraph::visitBreak(Break* curr) {
-  if (curr->condition) {
-    breakMappings[curr->name].emplace_back(currMapping);
-  } else {
-    breakMappings[curr->name].emplace_back(std::move(currMapping));
-    setUnreachable(currMapping);
-  }
-}
-void LocalGraph::visitSwitch(Switch* curr) {
-  std::set<Name> all;
-  for (auto target : curr->targets) {
-    all.insert(target);
   }
-  all.insert(curr->default_);
-  for (auto target : all) {
-    breakMappings[target].emplace_back(currMapping);
-  }
-  setUnreachable(currMapping);
-}
-void LocalGraph::visitReturn(Return *curr) {
-  setUnreachable(currMapping);
-}
-void LocalGraph::visitUnreachable(Unreachable *curr) {
-  setUnreachable(currMapping);
-}
+};
 
-// local usage
+} // namespace LocalGraphInternal
 
-void LocalGraph::visitGetLocal(GetLocal* curr) {
-  assert(currMapping.size() == numLocals);
-  assert(curr->index < numLocals);
-  for (auto& loopGets : loopGetStack) {
-    loopGets.push_back(curr);
-  }
-  // current sets are our sets
-  getSetses[curr] = currMapping[curr->index];
-  locations[curr] = getCurrentPointer();
-}
-void LocalGraph::visitSetLocal(SetLocal* curr) {
-  assert(currMapping.size() == numLocals);
-  assert(curr->index < numLocals);
-  // current sets are just this set
-  currMapping[curr->index] = { curr }; // TODO optimize?
-  locations[curr] = getCurrentPointer();
-}
+// LocalGraph implementation
 
-// traversal
+LocalGraph::LocalGraph(Function* func) {
+  LocalGraphInternal::Flower flower(getSetses, locations, func);
 
-void LocalGraph::scan(LocalGraph* self, Expression** currp) {
-  if (auto* iff = (*currp)->dynCast<If>()) {
-    // if needs special handling
-    if (iff->ifFalse) {
-      self->pushTask(LocalGraph::afterIfFalse, currp);
-      self->pushTask(LocalGraph::scan, &iff->ifFalse);
+#ifdef LOCAL_GRAPH_DEBUG
+  std::cout << "LocalGraph::dump\n";
+  for (auto& pair : getSetses) {
+    auto* get = pair.first;
+    auto& sets = pair.second;
+    std::cout << "GET\n" << get << " is influenced by\n";
+    for (auto* set : sets) {
+      std::cout << set << '\n';
     }
-    self->pushTask(LocalGraph::afterIfTrue, currp);
-    self->pushTask(LocalGraph::scan, &iff->ifTrue);
-    self->pushTask(LocalGraph::afterIfCondition, currp);
-    self->pushTask(LocalGraph::scan, &iff->condition);
-  } else {
-    PostWalker<LocalGraph>::scan(self, currp);
   }
-
-  // loops need pre-order visiting too
-  if ((*currp)->is<Loop>()) {
-    self->pushTask(LocalGraph::beforeLoop, currp);
-  }
-}
-
-// helpers
-
-void LocalGraph::setUnreachable(Mapping& mapping) {
-  mapping.resize(numLocals); // may have been emptied by a move
-  mapping[0].clear();
-}
-
-bool LocalGraph::isUnreachable(Mapping& mapping) {
-  // we must have some set for each index, if only the zero init, so empty means we emptied it for unreachable code
-  return mapping[0].empty();
+  std::cout << "total locations: " << locations.size() << '\n';
+#endif
 }
 
-// merges a bunch of infos into one.
-// if we need phis, writes them into the provided vector. the caller should
-// ensure those are placed in the right location
-LocalGraph::Mapping& LocalGraph::merge(std::vector<Mapping>& mappings) {
-  assert(mappings.size() > 0);
-  auto& out = mappings[0];
-  if (mappings.size() == 1) {
-    return out;
-  }
-  // merge into the first
-  for (Index j = 1; j < mappings.size(); j++) {
-    auto& other = mappings[j];
-    for (Index i = 0; i < numLocals; i++) {
-      auto& outSets = out[i];
-      for (auto* set : other[i]) {
-        outSets.insert(set);
+void LocalGraph::computeInfluences() {
+  for (auto& pair : locations) {
+    auto* curr = pair.first;
+    if (auto* set = curr->dynCast<SetLocal>()) {
+      FindAll<GetLocal> findAll(set->value);
+      for (auto* get : findAll.list) {
+        getInfluences[get].insert(set);
+      }
+    } else {
+      auto* get = curr->cast<GetLocal>();
+      for (auto* set : getSetses[get]) {
+        setInfluences[set].insert(get);
       }
     }
   }
-  return out;
 }
 
 } // namespace wasm