Add ModAsyncify* passes (#2404)

These passes are meant to be run after Asyncify has been run, they modify the
output. We can assume that we will always unwind if we reach an import, or
that we will never unwind, etc.

This is meant to help with lazy code loading, that is, the ability for an
initially-downloaded wasm to not contain all the code, and if code not present
there is called, we download all the rest and continue with that. That could
work something like this:

* The wasm is created. It contains calls to a special import for lazy code
  loading.
*  Asyncify is run on it.
*  The initially downloaded wasm is created by running
  --mod-asyncify-always-and-only-unwind: if the special import for lazy code
  loading is called, we will definitely unwind, and we won't rewind in this binary.
*  The lazily downloaded wasm is created by running --mod-asyncify-never-unwind:
  we will rewind into this binary, but no longer need support for unwinding.
  (Optionally, there could also be a third wasm, which has not had Asyncify run
  on it, and which we'd swap to for max speed.)

These --mod-asyncify passes allow the optimizer to do a lot of work, especially
for the initially downloaded wasm if we have lots of calls to the lazy code
loading import. In that case the optimizer will see that those calls unwind,
which means the code after them is not reached, potentially making lots of code
dead and removable.

3 files changed, 154 insertions, 1 deletions

diff --git a/src/passes/Asyncify.cpp b/src/passes/Asyncify.cpp
index 56c294b21..c4e1c58bd 100644
--- a/src/passes/Asyncify.cpp
+++ b/src/passes/Asyncify.cpp
@@ -253,6 +253,7 @@
 //
 
 #include "ir/effects.h"
+#include "ir/find_all.h"
 #include "ir/literal-utils.h"
 #include "ir/memory-utils.h"
 #include "ir/module-utils.h"
@@ -1169,6 +1170,10 @@ private:
 
 } // anonymous namespace
 
+static std::string getFullImportName(Name module, Name base) {
+  return std::string(module.str) + '.' + base.str;
+}
+
 struct Asyncify : public Pass {
   void run(PassRunner* runner, Module* module) override {
     bool optimize = runner->options.optimizeLevel > 0;
@@ -1209,7 +1214,7 @@ struct Asyncify : public Pass {
       if (allImportsCanChangeState) {
         return true;
       }
-      std::string full = std::string(module.str) + '.' + base.str;
+      auto full = getFullImportName(module, base);
       for (auto& listedImport : listedImports) {
         if (String::wildcardMatch(listedImport, full)) {
           return true;
@@ -1341,4 +1346,143 @@ private:
 
 Pass* createAsyncifyPass() { return new Asyncify(); }
 
+// Helper passes that can be run after Asyncify.
+
+template<bool neverRewind, bool neverUnwind, bool importsAlwaysUnwind>
+struct ModAsyncify
+  : public WalkerPass<LinearExecutionWalker<
+      ModAsyncify<neverRewind, neverUnwind, importsAlwaysUnwind>>> {
+  bool isFunctionParallel() override { return true; }
+
+  ModAsyncify* create() override {
+    return new ModAsyncify<neverRewind, neverUnwind, importsAlwaysUnwind>();
+  }
+
+  void doWalkFunction(Function* func) {
+    // Find the asyncify state name.
+    auto* unwind = this->getModule()->getExport(ASYNCIFY_STOP_UNWIND);
+    auto* unwindFunc = this->getModule()->getFunction(unwind->value);
+    FindAll<GlobalSet> sets(unwindFunc->body);
+    assert(sets.list.size() == 1);
+    asyncifyStateName = sets.list[0]->name;
+    // Walk and optimize.
+    this->walk(func->body);
+  }
+
+  // Note that we don't just implement GetGlobal as we may know the value is
+  // *not* 0, 1, or 2, but not know the actual value. So what we can say depends
+  // on the comparison being done on it, and so we implement Binary and
+  // Select.
+
+  void visitBinary(Binary* curr) {
+    // Check if this is a comparison of the asyncify state to a specific
+    // constant, which we may know is impossible.
+    bool flip = false;
+    if (curr->op == NeInt32) {
+      flip = true;
+    } else if (curr->op != EqInt32) {
+      return;
+    }
+    auto* c = curr->right->dynCast<Const>();
+    if (!c) {
+      return;
+    }
+    auto* get = curr->left->dynCast<GlobalGet>();
+    if (!get || get->name != asyncifyStateName) {
+      return;
+    }
+    // This is a comparison of the state to a constant, check if we know the
+    // value.
+    int32_t value;
+    auto checkedValue = c->value.geti32();
+    if ((checkedValue == int(State::Unwinding) && neverUnwind) ||
+        (checkedValue == int(State::Rewinding) && neverRewind)) {
+      // We know the state is checked against an impossible value.
+      value = 0;
+    } else if (checkedValue == int(State::Unwinding) && this->unwinding) {
+      // We know we are in fact unwinding right now.
+      value = 1;
+      unsetUnwinding();
+    } else {
+      return;
+    }
+    if (flip) {
+      value = 1 - value;
+    }
+    Builder builder(*this->getModule());
+    this->replaceCurrent(builder.makeConst(Literal(int32_t(value))));
+  }
+
+  void visitSelect(Select* curr) {
+    auto* get = curr->condition->dynCast<GlobalGet>();
+    if (!get || get->name != asyncifyStateName) {
+      return;
+    }
+    // This is a comparison of the state to zero, which means we are checking
+    // "if running normally, run this code, but if rewinding, ignore it". If
+    // we know we'll never rewind, we can optimize this.
+    if (neverRewind) {
+      Builder builder(*this->getModule());
+      curr->condition = builder.makeConst(Literal(int32_t(0)));
+    }
+  }
+
+  void visitCall(Call* curr) {
+    unsetUnwinding();
+    if (!importsAlwaysUnwind) {
+      return;
+    }
+    auto* target = this->getModule()->getFunction(curr->target);
+    if (!target->imported()) {
+      return;
+    }
+    // This is an import that definitely unwinds. Await the next check of
+    // the state in this linear execution trace, which we can turn into a
+    // constant.
+    this->unwinding = true;
+  }
+
+  void visitCallIndirect(CallIndirect* curr) { unsetUnwinding(); }
+
+  static void doNoteNonLinear(
+    ModAsyncify<neverRewind, neverUnwind, importsAlwaysUnwind>* self,
+    Expression**) {
+    // When control flow branches, stop tracking an unwinding.
+    self->unsetUnwinding();
+  }
+
+  void visitGlobalSet(GlobalSet* set) {
+    // TODO: this could be more precise
+    unsetUnwinding();
+  }
+
+private:
+  Name asyncifyStateName;
+
+  // Whether we just did a call to an import that indicates we are unwinding.
+  bool unwinding = false;
+
+  void unsetUnwinding() { this->unwinding = false; }
+};
+
+//
+// Assume imports that may unwind will always unwind, and that rewinding never
+// happens.
+//
+
+Pass* createModAsyncifyAlwaysOnlyUnwindPass() {
+  return new ModAsyncify<true, false, true>();
+}
+
+//
+// Assume that we never unwind, but may still rewind.
+//
+struct ModAsyncifyNeverUnwind : public Pass {
+  void run(PassRunner* runner, Module* module) override {}
+};
+
+Pass* createModAsyncifyNeverUnwindPass() {
+  return new ModAsyncify<false, true, false>();
+}
+
 } // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 7d17510fc..0ac70aa59 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -180,6 +180,13 @@ void PassRegistry::registerPasses() {
                "minifies both import and export names, and emits a mapping to "
                "the minified ones",
                createMinifyImportsAndExportsPass);
+  registerPass("mod-asyncify-always-and-only-unwind",
+               "apply the assumption that asyncify imports always unwind, "
+               "and we never rewind",
+               createModAsyncifyAlwaysOnlyUnwindPass);
+  registerPass("mod-asyncify-never-unwind",
+               "apply the assumption that asyncify never unwinds",
+               createModAsyncifyNeverUnwindPass);
   registerPass("nm", "name list", createNameListPass);
   registerPass("no-exit-runtime",
                "removes calls to atexit(), which is valid if the C runtime "
diff --git a/src/passes/passes.h b/src/passes/passes.h
index 6342ea15b..6732f1114 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -69,6 +69,8 @@ Pass* createOptimizeAddedConstantsPropagatePass();
 Pass* createOptimizeInstructionsPass();
 Pass* createOptimizeStackIRPass();
 Pass* createPickLoadSignsPass();
+Pass* createModAsyncifyAlwaysOnlyUnwindPass();
+Pass* createModAsyncifyNeverUnwindPass();
 Pass* createPostEmscriptenPass();
 Pass* createPrecomputePass();
 Pass* createPrecomputePropagatePass();