SimplifyGlobals: Apply known constant values in linear traces (#2340)

This optimizes stuff like

(global.set $x (i32.const 123))
(global.get $x)

into

(global.set $x (i32.const 123))
(i32.const 123)

This doesn't help much with LLVM output as it's rare to use globals (except for the stack pointer, and that's already well optimized), but it may help on general wasm. It can also help with Asyncify that does use globals extensively.

3 files changed, 83 insertions, 14 deletions

diff --git a/src/passes/SimplifyGlobals.cpp b/src/passes/SimplifyGlobals.cpp
index fce9df964..88f27f8be 100644
--- a/src/passes/SimplifyGlobals.cpp
+++ b/src/passes/SimplifyGlobals.cpp
@@ -22,13 +22,21 @@
 //  * If an immutable global is a copy of another, use the earlier one,
 //    to allow removal of the copies later.
 //  * Apply the constant values of immutable globals.
+//  * Apply the constant values of previous global.sets, in a linear
+//    execution trace.
 //
 // Some globals may not have uses after these changes, which we leave
 // to other passes to optimize.
 //
+// This pass has a "optimize" variant (similar to inlining and DAE)
+// that also runs general function optimizations where we managed to replace
+// a constant value. That is helpful as such a replacement often opens up
+// further optimization opportunities.
+//
 
 #include <atomic>
 
+#include "ir/effects.h"
 #include "ir/utils.h"
 #include "pass.h"
 #include "wasm-builder.h"
@@ -84,26 +92,73 @@ private:
 };
 
 struct ConstantGlobalApplier
-  : public WalkerPass<PostWalker<ConstantGlobalApplier>> {
+  : public WalkerPass<
+      LinearExecutionWalker<ConstantGlobalApplier,
+                            UnifiedExpressionVisitor<ConstantGlobalApplier>>> {
   bool isFunctionParallel() override { return true; }
 
-  ConstantGlobalApplier(NameSet* constantGlobals)
-    : constantGlobals(constantGlobals) {}
+  ConstantGlobalApplier(NameSet* constantGlobals, bool optimize)
+    : constantGlobals(constantGlobals), optimize(optimize) {}
 
   ConstantGlobalApplier* create() override {
-    return new ConstantGlobalApplier(constantGlobals);
+    return new ConstantGlobalApplier(constantGlobals, optimize);
   }
 
-  void visitGlobalGet(GlobalGet* curr) {
-    if (constantGlobals->count(curr->name)) {
-      auto* global = getModule()->getGlobal(curr->name);
-      assert(global->init->is<Const>());
-      replaceCurrent(ExpressionManipulator::copy(global->init, *getModule()));
+  void visitExpression(Expression* curr) {
+    if (auto* set = curr->dynCast<GlobalSet>()) {
+      if (auto* c = set->value->dynCast<Const>()) {
+        currConstantGlobals[set->name] = c->value;
+      } else {
+        currConstantGlobals.erase(set->name);
+      }
+      return;
+    } else if (auto* get = curr->dynCast<GlobalGet>()) {
+      // Check if the global is known to be constant all the time.
+      if (constantGlobals->count(get->name)) {
+        auto* global = getModule()->getGlobal(get->name);
+        assert(global->init->is<Const>());
+        replaceCurrent(ExpressionManipulator::copy(global->init, *getModule()));
+        replaced = true;
+        return;
+      }
+      // Check if the global has a known value in this linear trace.
+      auto iter = currConstantGlobals.find(get->name);
+      if (iter != currConstantGlobals.end()) {
+        Builder builder(*getModule());
+        replaceCurrent(builder.makeConst(iter->second));
+        replaced = true;
+      }
+      return;
+    }
+    // Otherwise, invalidate if we need to.
+    EffectAnalyzer effects(getPassOptions());
+    effects.visit(curr);
+    assert(effects.globalsWritten.empty()); // handled above
+    if (effects.calls) {
+      currConstantGlobals.clear();
+    }
+  }
+
+  static void doNoteNonLinear(ConstantGlobalApplier* self, Expression** currp) {
+    self->currConstantGlobals.clear();
+  }
+
+  void visitFunction(Function* curr) {
+    if (replaced && optimize) {
+      PassRunner runner(getModule(), getPassRunner()->options);
+      runner.setIsNested(true);
+      runner.addDefaultFunctionOptimizationPasses();
+      runner.runOnFunction(curr);
     }
   }
 
 private:
   NameSet* constantGlobals;
+  bool optimize;
+  bool replaced = false;
+
+  // The globals currently constant in the linear trace.
+  std::map<Name, Literal> currConstantGlobals;
 };
 
 } // anonymous namespace
@@ -113,6 +168,9 @@ struct SimplifyGlobals : public Pass {
   Module* module;
 
   GlobalInfoMap map;
+  bool optimize;
+
+  SimplifyGlobals(bool optimize = false) : optimize(optimize) {}
 
   void run(PassRunner* runner_, Module* module_) override {
     runner = runner_;
@@ -214,12 +272,14 @@ struct SimplifyGlobals : public Pass {
         constantGlobals.insert(global->name);
       }
     }
-    if (!constantGlobals.empty()) {
-      ConstantGlobalApplier(&constantGlobals).run(runner, module);
-    }
+    ConstantGlobalApplier(&constantGlobals, optimize).run(runner, module);
   }
 };
 
-Pass* createSimplifyGlobalsPass() { return new SimplifyGlobals(); }
+Pass* createSimplifyGlobalsPass() { return new SimplifyGlobals(false); }
+
+Pass* createSimplifyGlobalsOptimizingPass() {
+  return new SimplifyGlobals(true);
+}
 
 } // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index daa575c6d..7d17510fc 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -267,6 +267,10 @@ void PassRegistry::registerPasses() {
   registerPass("simplify-globals",
                "miscellaneous globals-related optimizations",
                createSimplifyGlobalsPass);
+  registerPass("simplify-globals-optimizing",
+               "miscellaneous globals-related optimizations, and optimizes "
+               "where we replaced global.gets with constants",
+               createSimplifyGlobalsOptimizingPass);
   registerPass("simplify-locals",
                "miscellaneous locals-related optimizations",
                createSimplifyLocalsPass);
@@ -416,7 +420,11 @@ void PassRunner::addDefaultGlobalOptimizationPostPasses() {
   // optimizations show more functions as duplicate
   add("duplicate-function-elimination");
   add("duplicate-import-elimination");
-  add("simplify-globals");
+  if (options.optimizeLevel >= 2 || options.shrinkLevel >= 2) {
+    add("simplify-globals-optimizing");
+  } else {
+    add("simplify-globals");
+  }
   add("remove-unused-module-elements");
   add("memory-packing");
   // may allow more inlining/dae/etc., need --converge for that
diff --git a/src/passes/passes.h b/src/passes/passes.h
index 115e9bab6..6342ea15b 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -92,6 +92,7 @@ Pass* createRedundantSetEliminationPass();
 Pass* createSafeHeapPass();
 Pass* createSimplifyLocalsPass();
 Pass* createSimplifyGlobalsPass();
+Pass* createSimplifyGlobalsOptimizingPass();
 Pass* createSimplifyLocalsNoNestingPass();
 Pass* createSimplifyLocalsNoTeePass();
 Pass* createSimplifyLocalsNoStructurePass();