[Stack Switching] Add basic support for resume/suspend in the interpreter (#7771)

It turns out we can implement suspend/resume on the current
interpreter with only a small amount of changes, basically using
the same idea as Asyncify: pause and resume structured control
flow by rewinding the stack in a "resume" mode that skips normal
execution. That is, this does not use a program counter (no goto),
nor continuation-passing style, and instead adds interpreter
support to unwind and rewind the stack. This is not the fastest
way to do things, but it is the simplest by far.

The basic idea is the same as in this 2019 blogpost:

https://kripken.github.io/blog/wasm/2019/07/16/asyncify.html

This is quite efficient in the things we care about: suspend/
resume is slow, but normal execution hardly regresses, which is
important to keep the Precompute pass from getting slower.
While the more intrusive experiment #7762 made that pass 2x
slower, this only adds 10-15% overhead. The main reason is that
this pass keeps us down to a single indirect call per instruction,
while e.g. separating decoding and execution - to maintain a
normal value stack - ends up doing two indirect calls. The "trick"
in this PR is that, yes, we do need a value stack (that is the only
practical way to stash the values on the stack when suspending),
but we can populate that stack only when inside a coroutine
(when we might suspend). So we still use the normal way of
getting child instruction values - just calling `visit(curr->child)`
from the parent's `visitFoo()` method - and do not lose that
speed, but still have a stack of values when we need it.

(10-15% is still significant, but it is just on a single pass, so it
seems acceptable, and there might be ways to optimize this further.)

Notes:

* Flow now has a "suspendTag" property, which denotes the tag we are
  suspending, when we suspend.
* As part of this change, `callFunction` in the interpreter returns a
  Flow, so that we can propagate suspensions out of functions.
* For testing, this adds `assert_suspension` support in `wasm-shell`.
* For testing, add part of `cont.wast` from the test suite, and minor
  fixes for it in `wasm.cpp`.
* This only adds basic support: `resume_throw` and various other parts
  of the stack switching proposal are TODOs, but I think this PR does the
  hard parts.

Author: kripken

scripts/test/fuzzing.py

@@ -111,6 +111,8 @@
     'precompute-stack-switching.wast',
     'unsubtyping-stack-switching.wast',
     'vacuum-stack-switching.wast',
+    'cont.wast',
+    'cont_simple.wast',
     # TODO: fuzzer support for custom descriptors
     'remove-unused-module-elements-refs-descriptors.wast',
     'custom-descriptors.wast',

src/ir/iteration.h

@@ -32,7 +32,7 @@ namespace wasm {
 // In general, it is preferable not to use this class and to directly access the
 // children (using e.g. iff->ifTrue etc.), as that is faster. However, in cases
 // where speed does not matter, this can be convenient. TODO: reimplement these
-// to avoid materializing all the chilren at once.
+// to avoid materializing all the children at once.
 //
 //   ChildIterator - Iterates over all children
 //

src/literal.h

@@ -33,6 +33,7 @@ namespace wasm {
 class Literals;
 struct GCData;
 struct ExnData;
+struct ContData;
 
 class Literal {
   // store only integers, whose bits are deterministic. floats
@@ -63,6 +64,8 @@ class Literal {
     std::shared_ptr<GCData> gcData;
     // A reference to Exn data.
     std::shared_ptr<ExnData> exnData;
+    // A reference to a Continuation.
+    std::shared_ptr<ContData> contData;
   };
 
 public:
@@ -93,6 +96,7 @@ class Literal {
   }
   explicit Literal(std::shared_ptr<GCData> gcData, HeapType type);
   explicit Literal(std::shared_ptr<ExnData> exnData);
+  explicit Literal(std::shared_ptr<ContData> contData);
   explicit Literal(std::string_view string);
   Literal(const Literal& other);
   Literal& operator=(const Literal& other);
@@ -105,6 +109,7 @@ class Literal {
   // a null or i31). This includes structs, arrays, and also strings.
   bool isData() const { return type.isData(); }
   bool isExn() const { return type.isExn(); }
+  bool isContinuation() const { return type.isContinuation(); }
   bool isString() const { return type.isString(); }
 
   bool isNull() const { return type.isNull(); }
@@ -312,6 +317,7 @@ class Literal {
   }
   std::shared_ptr<GCData> getGCData() const;
   std::shared_ptr<ExnData> getExnData() const;
+  std::shared_ptr<ContData> getContData() const;
 
   // careful!
   int32_t* geti32Ptr() {

src/parser/wast-parser.cpp

@@ -381,6 +381,25 @@ MaybeResult<Assertion> assertTrap(Lexer& in) {
   return Assertion{AssertModule{ModuleAssertionType::Trap, *mod}};
 }
 
+// (assert_suspension action msg)
+MaybeResult<Assertion> assertSuspension(Lexer& in) {
+  if (!in.takeSExprStart("assert_suspension"sv)) {
+    return {};
+  }
+  if (auto a = maybeAction(in)) {
+    CHECK_ERR(a);
+    auto msg = in.takeString();
+    if (!msg) {
+      return in.err("expected error message");
+    }
+    if (!in.takeRParen()) {
+      return in.err("expected end of assertion");
+    }
+    return Assertion{AssertAction{ActionAssertionType::Suspension, *a}};
+  }
+  return in.err("invalid assert_suspension");
+}
+
 MaybeResult<Assertion> assertion(Lexer& in) {
   if (auto a = assertReturn(in)) {
     CHECK_ERR(a);
@@ -402,6 +421,10 @@ MaybeResult<Assertion> assertion(Lexer& in) {
     CHECK_ERR(a);
     return *a;
   }
+  if (auto a = assertSuspension(in)) {
+    CHECK_ERR(a);
+    return *a;
+  }
   return {};
 }
 

src/parser/wat-parser.h

@@ -78,7 +78,7 @@ struct AssertReturn {
   ExpectedResults expected;
 };
 
-enum class ActionAssertionType { Trap, Exhaustion, Exception };
+enum class ActionAssertionType { Trap, Exhaustion, Exception, Suspension };
 
 struct AssertAction {
   ActionAssertionType type;

src/shell-interface.h

@@ -144,7 +144,9 @@ struct ShellExternalInterface : ModuleRunner::ExternalInterface {
       std::cout << "exit()\n";
       throw ExitException();
     } else if (auto* inst = getImportInstance(import)) {
-      return inst->callExport(import->base, arguments);
+      auto flow = inst->callExport(import->base, arguments);
+      assert(!flow.suspendTag); // TODO: support stack switching on calls
+      return flow.values;
     }
     Fatal() << "callImport: unknown import: " << import->module.str << "."
             << import->name.str;
@@ -191,7 +193,9 @@ struct ShellExternalInterface : ModuleRunner::ExternalInterface {
     if (func->imported()) {
       return callImport(func, arguments);
     } else {
-      return instance.callFunction(func->name, arguments);
+      auto flow = instance.callFunction(func->name, arguments);
+      assert(!flow.suspendTag); // TODO: support stack switching on calls
+      return flow.values;
     }
   }
 

src/tools/execution-results.h

@@ -254,7 +254,9 @@ struct LoggingExternalInterface : public ShellExternalInterface {
     }
 
     // Call the function.
-    return instance->callFunction(func->name, arguments);
+    auto flow = instance->callFunction(func->name, arguments);
+    assert(!flow.suspendTag);
+    return flow.values;
   }
 
   void setModuleRunner(ModuleRunner* instance_) { instance = instance_; }
@@ -471,7 +473,12 @@ struct ExecutionResults {
         }
         arguments.push_back(Literal::makeZero(param));
       }
-      return instance.callFunction(func->name, arguments);
+      auto flow = instance.callFunction(func->name, arguments);
+      if (flow.suspendTag) { // TODO: support stack switching here
+        std::cout << "[exception thrown: unhandled suspend]" << std::endl;
+        return Exception{};
+      }
+      return flow.values;
     } catch (const TrapException&) {
       return Trap{};
     } catch (const WasmException& e) {

src/tools/wasm-ctor-eval.cpp

@@ -350,7 +350,11 @@ struct CtorEvalExternalInterface : EvallingModuleRunner::ExternalInterface {
                                 targetFunc.toString());
     }
     if (!func->imported()) {
-      return instance.callFunction(targetFunc, arguments);
+      auto flow = instance.callFunction(targetFunc, arguments);
+      if (flow.suspendTag) {
+        throw FailToEvalException("unhandled suspend");
+      }
+      return flow.values;
     } else {
       throw FailToEvalException(
         std::string("callTable on imported function: ") +

src/tools/wasm-shell.cpp

@@ -186,8 +186,12 @@ struct Shell {
   struct TrapResult {};
   struct HostLimitResult {};
   struct ExceptionResult {};
-  using ActionResult =
-    std::variant<Literals, TrapResult, HostLimitResult, ExceptionResult>;
+  struct SuspensionResult {};
+  using ActionResult = std::variant<Literals,
+                                    TrapResult,
+                                    HostLimitResult,
+                                    ExceptionResult,
+                                    SuspensionResult>;
 
   std::string resultToString(ActionResult& result) {
     if (std::get_if<TrapResult>(&result)) {
@@ -196,6 +200,8 @@ struct Shell {
       return "exceeded host limit";
     } else if (std::get_if<ExceptionResult>(&result)) {
       return "exception";
+    } else if (std::get_if<SuspensionResult>(&result)) {
+      return "suspension";
     } else if (auto* vals = std::get_if<Literals>(&result)) {
       std::stringstream ss;
       ss << *vals;
@@ -213,8 +219,9 @@ struct Shell {
         return TrapResult{};
       }
       auto& instance = it->second;
+      Flow flow;
       try {
-        return instance->callExport(invoke->name, invoke->args);
+        flow = instance->callExport(invoke->name, invoke->args);
       } catch (TrapException&) {
         return TrapResult{};
       } catch (HostLimitException&) {
@@ -224,6 +231,10 @@ struct Shell {
       } catch (...) {
         WASM_UNREACHABLE("unexpected error");
       }
+      if (flow.suspendTag) {
+        return SuspensionResult{};
+      }
+      return flow.values;
     } else if (auto* get = std::get_if<GetAction>(&act)) {
       auto it = instances.find(get->base ? *get->base : lastModule);
       if (it == instances.end()) {
@@ -390,6 +401,12 @@ struct Shell {
         }
         err << "expected exception";
         break;
+      case ActionAssertionType::Suspension:
+        if (std::get_if<SuspensionResult>(&result)) {
+          return Ok{};
+        }
+        err << "expected suspension";
+        break;
     }
     err << ", got " << resultToString(result);
     return Err{err.str()};