[IRBuilder] Use multiple scratch locals for inputs (#9336)

Instead of using a single tuple-typed scratch local for multivalue
control flow inputs, allocate one scalar scratch local per input value.
This avoids eagerly packaging scalar inputs into a tuple on scope entry.
It also avoids allocating extra tuple scratch locals to extract elements
when inputs are consumed inside the scope.

Scalar input locals do have some downsides. Scopes that pass all of
their inputs through untouched must now repackage them with tuple.make.
Multivalue loop trampolines in fixLoopWithInput also still need a tuple
scratch local to unpack branch values into the scalar input locals.
However, real parameterized scopes usually consume their inputs as
individual values. Scalar input locals can also be coalesced directly by
standard local passes without tuple optimization, and they will allow
individual parameters to be tunneled across nested scopes with different
arities in a future commit.

Author: tlively

src/passes/Outlining.cpp

@@ -34,12 +34,6 @@
 #define ODBG(statement)
 #endif
 
-// Check a Result or MaybeResult for error and call Fatal() if the error exists.
-#define ASSERT_OK(val)                                                         \
-  if (auto _val = (val); auto err = _val.getErr()) {                           \
-    Fatal() << err->msg;                                                       \
-  }
-
 namespace wasm {
 
 // This custom hasher conforms to std::hash<Key>. Its purpose is to provide

src/support/result.h

@@ -17,6 +17,7 @@
 #ifndef wasm_support_result_h
 #define wasm_support_result_h
 
+#include <cassert>
 #include <string>
 #include <variant>
 
@@ -39,6 +40,13 @@ struct Err {
     return (typename decltype(_val)::ErrorType)(*err);                         \
   }
 
+// Assert that a Result or MaybeResult is not an error.
+#define ASSERT_OK(val)                                                         \
+  do {                                                                         \
+    [[maybe_unused]] auto _val = (val);                                        \
+    assert(!_val.getErr());                                                    \
+  } while (false)
+
 // Represent a result of type T or an error message.
 template<typename T = Ok, typename E = Err> struct [[nodiscard]] Result {
   using ErrorType = E;

src/wasm-ir-builder.h

@@ -433,10 +433,10 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     // function's `explicitLabelNames`.
     bool labelExplicit = false;
 
-    // If the control flow scope has an input type, we need to lower it using a
-    // scratch local because we cannot represent control flow input in the IR.
+    // If the control flow scope has an input type, we need to lower it using
+    // scratch locals because we cannot represent control flow input in the IR.
     Type inputType;
-    Index inputLocal = -1;
+    std::vector<Index> inputLocals;
 
     // If there are br_on_*, try_table, or resume branches that target this
     // scope and carry additional values, we need to use a scratch local to
@@ -461,13 +461,6 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
     ScopeCtx() : scope(NoScope{}) {}
     ScopeCtx(Scope scope, Type inputType)
       : scope(scope), inputType(inputType) {}
-    ScopeCtx(
-      Scope scope, Name label, bool labelUsed, Type inputType, Index inputLocal)
-      : scope(scope), label(label), labelUsed(labelUsed), inputType(inputType),
-        inputLocal(inputLocal) {}
-    ScopeCtx(Scope scope, Name label, bool labelUsed, Name branchLabel)
-      : scope(scope), label(label), branchLabel(branchLabel),
-        labelUsed(labelUsed) {}
 
     static ScopeCtx makeFunc(Function* func) {
       return ScopeCtx(FuncScope{func}, Type::none);
@@ -514,7 +507,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
       unreachable = false;
       if (!keepInput) {
         inputType = Type::none;
-        inputLocal = -1;
+        inputLocals.clear();
       }
     }
     bool isNone() { return std::get_if<NoScope>(&scope); }
@@ -715,35 +708,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Index blockHint = 0;
   Index labelHint = 0;
 
-  Result<> pushScope(ScopeCtx&& scope) {
-    if (auto label = scope.getOriginalLabel()) {
-      // Assign a fresh label to the scope, if necessary.
-      if (!scope.label) {
-        scope.label = makeFresh(label);
-      }
-      // Record the original label to handle references to it correctly.
-      labelDepths[label].push_back(scopeStack.size() + 1);
-    }
-    if (binaryPos) {
-      scope.startPos = lastBinaryPos;
-      lastBinaryPos = *binaryPos;
-    }
-    bool hasInput = scope.inputType != Type::none;
-    Index inputLocal = scope.inputLocal;
-    if (hasInput && !scope.isDelimiter()) {
-      if (inputLocal == Index(-1)) {
-        auto scratch = addScratchLocal(scope.inputType);
-        CHECK_ERR(scratch);
-        inputLocal = scope.inputLocal = *scratch;
-      }
-      CHECK_ERR(makeLocalSet(inputLocal));
-    }
-    scopeStack.emplace_back(std::move(scope));
-    if (hasInput) {
-      CHECK_ERR(makeLocalGet(inputLocal));
-    }
-    return Ok{};
-  }
+  Result<> pushScope(ScopeCtx&& scope);
 
   ScopeCtx& getScope() {
     if (scopeStack.empty()) {
@@ -806,7 +771,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Result<std::pair<Index, Name>> getExtraOutputLocalAndLabel(Index label,
                                                              size_t extraArity);
   Expression* fixExtraOutput(ScopeCtx& scope, Name label, Expression* expr);
-  void fixLoopWithInput(Loop* loop, Type inputType, Index scratch);
+  void fixLoopWithInput(Loop* loop,
+                        Type inputType,
+                        const std::vector<Index>& scratch);
 
   void applyAnnotations(Expression* expr, const CodeAnnotation& annotation);
 

src/wasm/wasm-ir-builder.cpp

@@ -778,6 +778,17 @@ struct IRBuilder::ChildPopper
     ConstraintCollector{builder, children}.visitStackSwitch(curr, ct);
     return popConstrainedChildren(children);
   }
+
+  Result<std::vector<Expression*>> popInputs(Type inputType) {
+    std::vector<Expression*> inputs(inputType.size());
+    std::vector<Child> children;
+    children.reserve(inputType.size());
+    for (size_t i = 0; i < inputType.size(); ++i) {
+      children.push_back({&inputs[i], {inputType[i]}});
+    }
+    CHECK_ERR(popConstrainedChildren(children));
+    return inputs;
+  }
 };
 
 Result<> IRBuilder::visit(Expression* curr) {
@@ -877,6 +888,42 @@ IRBuilder::visitTryTableStart(TryTable* trytable, Name label, Type inputType) {
   return pushScope(ScopeCtx::makeTryTable(trytable, label, inputType));
 }
 
+Result<> IRBuilder::pushScope(ScopeCtx&& scope) {
+  if (auto label = scope.getOriginalLabel()) {
+    // Assign a fresh label to the scope, if necessary.
+    if (!scope.label) {
+      scope.label = makeFresh(label);
+    }
+    // Record the original label to handle references to it correctly.
+    labelDepths[label].push_back(scopeStack.size() + 1);
+  }
+  if (binaryPos) {
+    scope.startPos = lastBinaryPos;
+    lastBinaryPos = *binaryPos;
+  }
+  bool hasInput = scope.inputType != Type::none;
+  if (hasInput && !scope.isDelimiter()) {
+    auto inputs = ChildPopper{*this}.popInputs(scope.inputType);
+    CHECK_ERR(inputs);
+    scope.inputLocals.reserve(inputs->size());
+    for (size_t i = 0; i < inputs->size(); ++i) {
+      auto scratch = addScratchLocal(scope.inputType[i]);
+      CHECK_ERR(scratch);
+      scope.inputLocals.push_back(*scratch);
+      pushSynthetic(builder.makeLocalSet(*scratch, (*inputs)[i]));
+    }
+  }
+  scopeStack.emplace_back(std::move(scope));
+  if (hasInput) {
+    auto& newScope = getScope();
+    for (size_t i = 0; i < newScope.inputLocals.size(); ++i) {
+      pushSynthetic(
+        builder.makeLocalGet(newScope.inputLocals[i], newScope.inputType[i]));
+    }
+  }
+  return Ok{};
+}
+
 Result<Expression*> IRBuilder::finishScope(Block* block) {
 #if IR_BUILDER_DEBUG
   if (auto* loc = std::get_if<Function::DebugLocation>(&debugLoc)) {
@@ -1212,8 +1259,8 @@ Result<> IRBuilder::visitEnd() {
     if (scope.inputType != Type::none && scope.labelUsed) {
       // Branches to this loop carry values, but Binaryen IR does not support
       // that. Fix this by trampolining the branches through new code that sets
-      // the branch value to the appropriate scratch local.
-      fixLoopWithInput(loop, scope.inputType, scope.inputLocal);
+      // the branch value to the appropriate scratch locals.
+      fixLoopWithInput(loop, scope.inputType, scope.inputLocals);
     }
     loop->finalize(loop->type);
     push(loop);
@@ -1222,8 +1269,19 @@ Result<> IRBuilder::visitEnd() {
     if (scope.inputType != Type::none) {
       // Normally an if without an else must have type none, but if there is an
       // input parameter, the empty else arm must propagate its value.
-      // Synthesize an else arm that loads the value from the scratch local.
-      iff->ifFalse = builder.makeLocalGet(scope.inputLocal, scope.inputType);
+      // Synthesize an else arm that loads the value from the scratch locals.
+      if (scope.inputType.isSingle()) {
+        iff->ifFalse =
+          builder.makeLocalGet(scope.inputLocals[0], scope.inputType);
+      } else {
+        std::vector<Expression*> gets;
+        gets.reserve(scope.inputLocals.size());
+        for (size_t i = 0; i < scope.inputLocals.size(); ++i) {
+          gets.push_back(
+            builder.makeLocalGet(scope.inputLocals[i], scope.inputType[i]));
+        }
+        iff->ifFalse = builder.makeTupleMake(gets);
+      }
     } else {
       iff->ifFalse = nullptr;
     }
@@ -1368,17 +1426,18 @@ IRBuilder::fixExtraOutput(ScopeCtx& scope, Name label, Expression* curr) {
   return curr;
 }
 
-// Branches to this loop need to be trampolined through code that sets the value
-// carried by the branch to the appropriate scratch local before branching to
-// the loop. Transform this:
+// Branches to this loop need to be trampolined through code that sets the
+// values carried by the branch to the appropriate scratch locals before
+// branching to the loop. Transform this:
 //
 //   (loop $l (param t1) (result t2) ...)
 //
-// to this:
+// to this (or, when t1 is a tuple, using a tuple scratch local to extract each
+// element into its corresponding input scratch local):
 //
 //  (loop $l0 (result t2)
 //    (block $l1 (result t2)
-//      (local.set $scratch ;; set the branch values to the scratch local
+//      (local.set $scratch ;; set the branch value to the scratch local
 //        (block $l (result t1)
 //          (br $l1 ;; exit the loop with the fallthrough value, if any.
 //            ...   ;; contains branches to $l
@@ -1388,7 +1447,9 @@ IRBuilder::fixExtraOutput(ScopeCtx& scope, Name label, Expression* curr) {
 //      (br $l0) ;; continue the loop
 //    )
 //  )
-void IRBuilder::fixLoopWithInput(Loop* loop, Type inputType, Index scratch) {
+void IRBuilder::fixLoopWithInput(Loop* loop,
+                                 Type inputType,
+                                 const std::vector<Index>& scratch) {
   auto l = loop->name;
   auto l0 = makeFresh(l, 0);
   auto l1 = makeFresh(l, 1);
@@ -1399,10 +1460,22 @@ void IRBuilder::fixLoopWithInput(Loop* loop, Type inputType, Index scratch) {
           loop->body, l, builder.makeBreak(l1), inputType)
       : builder.makeBlock(l, {builder.makeBreak(l1, loop->body)}, inputType);
 
-  Block* outer = builder.makeBlock(
-    l1,
-    {builder.makeLocalSet(scratch, inner), builder.makeBreak(l0)},
-    loop->type);
+  std::vector<Expression*> list;
+  if (inputType.isSingle()) {
+    list.push_back(builder.makeLocalSet(scratch[0], inner));
+  } else {
+    auto tupleScratch = addScratchLocal(inputType);
+    ASSERT_OK(tupleScratch);
+    list.push_back(builder.makeLocalSet(*tupleScratch, inner));
+    for (Index i = 0; i < scratch.size(); ++i) {
+      list.push_back(builder.makeLocalSet(
+        scratch[i],
+        builder.makeTupleExtract(builder.makeLocalGet(*tupleScratch, inputType),
+                                 i)));
+    }
+  }
+  list.push_back(builder.makeBreak(l0));
+  Block* outer = builder.makeBlock(l1, list, loop->type);
 
   loop->body = outer;
   loop->name = l0;