Flatten/LowerBrOn experiment

scripts/fuzz_opt.py

@@ -2540,16 +2540,9 @@ def get_random_opts():
             if has_flatten:
                 print('avoiding multiple --flatten in a single command, due to exponential overhead')
                 continue
-            if '--enable-multivalue' in FEATURE_OPTS and '--enable-reference-types' in FEATURE_OPTS:
-                print('avoiding --flatten due to multivalue + reference types not supporting it (spilling of non-nullable tuples)')
-                print('TODO: Resolving https://github.com/WebAssembly/binaryen/issues/4824 may fix this')
-                continue
             if '--enable-exception-handling' in FEATURE_OPTS:
                 print('avoiding --flatten due to exception-handling not supporting it (requires blocks with results)')
                 continue
-            if '--gc' not in FEATURE_OPTS:
-                print('avoiding --flatten due to GC not supporting it (spilling of non-nullable locals)')
-                continue
             if INITIAL_CONTENTS and os.path.getsize(INITIAL_CONTENTS) > 2000:
                 print('avoiding --flatten due using a large amount of initial contents, which may blow up')
                 continue

src/ir/flat.h

@@ -39,7 +39,7 @@
 //    (i32.const 1)
 //  )
 //
-// Formally, this pass flattens in the precise sense of
+// Formally, the Flatten pass flattens in the precise sense of
 // making the AST have these properties:
 //
 //  1. Aside from a local.set, the operands of an instruction must be a
@@ -54,9 +54,8 @@
 //     values is prohibited already, but e.g. a block ending in unreachable,
 //     which can normally be nested, is also disallowed).
 //
-// Note: ref.as_non_null must be allowed in a nested position because we cannot
-// spill it to a local - the result is non-null, which is not allowable in a
-// local.
+// This header defines the concept of flatness, and utilities for verifying it,
+// which each pass that assumes flat form should use.
 //
 
 #ifndef wasm_ir_flat_h
@@ -76,12 +75,15 @@ inline void verifyFlatness(Function* func) {
     void visitExpression(Expression* curr) {
       if (Properties::isControlFlowStructure(curr)) {
         verify(!curr->type.isConcrete(),
-               "control flow structures must not flow values");
+               "control flow structures must not flow values",
+               curr);
       } else if (auto* set = curr->dynCast<LocalSet>()) {
         verify(!set->isTee() || set->type == Type::unreachable,
-               "tees are not allowed, only sets");
+               "tees are not allowed, only sets",
+               set);
         verify(!Properties::isControlFlowStructure(set->value),
-               "set values cannot be control flow");
+               "set values cannot be control flow",
+               set);
       } else {
         for (auto* child : ChildIterator(curr)) {
           bool isRefAsNonNull =
@@ -90,15 +92,16 @@ inline void verifyFlatness(Function* func) {
                    child->is<LocalGet>() || child->is<Unreachable>() ||
                    isRefAsNonNull,
                  "instructions must only have constant expressions, local.get, "
-                 "or unreachable as children");
+                 "or unreachable as children",
+                 curr);
         }
       }
     }
 
-    void verify(bool condition, const char* message) {
+    void verify(bool condition, const char* message, Expression* expr) {
       if (!condition) {
         Fatal() << "IR must be flat: run --flatten beforehand (" << message
-                << ", in " << getFunction()->name << ')';
+                << ", in " << getFunction()->name << ") " << *expr;
       }
     }
   };
@@ -107,7 +110,8 @@ inline void verifyFlatness(Function* func) {
   verifier.walkFunction(func);
   verifier.setFunction(func);
   verifier.verify(!func->body->type.isConcrete(),
-                  "function bodies must not flow values");
+                  "function bodies must not flow values",
+                  nullptr);
 }
 
 inline void verifyFlatness(Module* module) {

src/passes/Flatten.cpp

@@ -35,9 +35,11 @@
 //
 // The tuple has a non-nullable type, and so it cannot currently be set to a
 // local, but in principle there's no reason it couldn't be. For now, error on
-// this.
+// this. Note: running TupleOptimization can work around this, by removing such
+// tuples.
 
 #include <ir/branch-utils.h>
+#include <ir/drop.h>
 #include <ir/effects.h>
 #include <ir/eh-utils.h>
 #include <ir/flat.h>
@@ -49,6 +51,98 @@
 
 namespace wasm {
 
+// Lowers instructions that Flatten cannot directly handle. It is simpler to
+// first lower them to instructions that it can.
+struct LowerUnflattenable : public PostWalker<LowerUnflattenable> {
+  Builder builder;
+  const PassOptions& options;
+
+  LowerUnflattenable(Module& wasm, const PassOptions& options)
+    : builder(wasm), options(options) {}
+
+  void visitBrOn(BrOn* curr) {
+    if (curr->type == Type::unreachable) {
+      replaceUnreachableWithDrops();
+      return;
+    }
+
+    // All ops stash the ref to a temp var.
+    auto refType = curr->ref->type;
+    auto refTemp = builder.addVar(getFunction(), refType);
+    // All ops will use this refTee, and more gets.
+    auto* refTee = builder.makeLocalTee(refTemp, curr->ref, refType);
+    auto getRef = [&]() { return builder.makeLocalGet(refTemp, refType); };
+
+    // The overall shape we emit is to check a condition, branch if so with
+    // some value, and if not, flow out something:
+    //
+    //   if (condition) {
+    //     br(brValue);
+    //   }
+    //   flowValue
+    //
+    // Each op fills in those things.
+    Expression* condition = nullptr; // uses refTee
+    Expression* brValue = nullptr;
+    Expression* flowValue = nullptr;
+    bool flip = false; // whether to flip the condition.
+
+    switch (curr->op) {
+      case BrOnNull: {
+        condition = builder.makeRefIsNull(refTee);
+        brValue = nullptr;
+        flowValue = builder.makeRefAs(RefAsNonNull, getRef());
+        break;
+      }
+      case BrOnNonNull: {
+        condition = builder.makeRefIsNull(refTee);
+        flip = true;
+        brValue = builder.makeRefAs(RefAsNonNull, getRef());
+        // No value flos out.
+        break;
+      }
+      case BrOnCast: {
+        condition = builder.makeRefTest(refTee, curr->castType);
+        brValue = builder.makeRefCast(getRef(), curr->castType);
+        flowValue = getRef();
+        break;
+      }
+      case BrOnCastFail: {
+        condition = builder.makeRefTest(refTee, curr->castType);
+        flip = true;
+        brValue = getRef();
+        flowValue = builder.makeRefCast(getRef(), curr->castType);
+        break;
+      }
+      case BrOnCastDescEq:
+      case BrOnCastDescEqFail: {
+        WASM_UNREACHABLE("TODO");
+      }
+    }
+
+    if (flip) {
+      condition = builder.makeUnary(EqZInt32, condition);
+    }
+
+    auto* br = builder.makeBreak(curr->name, brValue);
+    Expression* result = builder.makeIf(condition, br);
+    if (flowValue) {
+      result = builder.makeSequence(result, flowValue);
+    }
+    replaceCurrent(result);
+  }
+
+  void replaceUnreachableWithDrops() {
+    Builder builder(*getModule());
+
+    getDroppedChildrenAndAppend(getCurrent(),
+                                *getModule(),
+                                options,
+                                builder.makeUnreachable(),
+                                DropMode::IgnoreParentEffects);
+  }
+};
+
 // We use the following algorithm: we maintain a list of "preludes", code
 // that runs right before an expression. When we visit an expression we
 // must handle it and its preludes. If the expression has side effects,
@@ -333,7 +427,7 @@ struct Flatten
       }
     }
 
-    if (curr->is<BrOn>() || curr->is<TryTable>()) {
+    if (curr->is<TryTable>()) {
       Fatal() << "Unsupported instruction for Flatten: "
               << getExpressionName(curr);
     }
@@ -368,7 +462,15 @@ struct Flatten
     }
   }
 
-  void visitFunction(Function* curr) {
+  void doWalkFunction(Function* curr) {
+    // Lower things before the main walk.
+    LowerUnflattenable(*getModule(), getPassOptions()).walkFunction(curr);
+
+    WalkerPass<
+      ExpressionStackWalker<Flatten, UnifiedExpressionVisitor<Flatten>>>::
+      doWalkFunction(curr);
+
+    // Finish up.
     auto* originalBody = curr->body;
     // if the body is a block with a result, turn that into a return
     if (curr->body->type.isConcrete()) {

src/passes/ReReloop.cpp

@@ -24,6 +24,7 @@
 #include <memory>
 
 #include "cfg/Relooper.h"
+#include "ir/find_all.h"
 #include "ir/flat.h"
 #include "ir/utils.h"
 #include "pass.h"
@@ -284,7 +285,7 @@ struct ReReloop final : public Pass {
       ReturnTask::handle(*this, ret);
     } else if (auto* un = curr->dynCast<Unreachable>()) {
       UnreachableTask::handle(*this, un);
-    } else if (curr->is<Try>() || curr->is<Throw>() || curr->is<Rethrow>()) {
+    } else if (curr->is<Try>()) {
       Fatal() << "ReReloop does not support EH instructions yet";
     } else {
       // not control flow, so just a simple element
@@ -298,6 +299,11 @@ struct ReReloop final : public Pass {
   }
 
   void runOnFunction(Module* module, Function* function) override {
+    if (FindAll<Try>(function->body).list.size()) {
+      std::cout << "skip " << function->name << '\n';
+      return;
+    }
+
     Flat::verifyFlatness(function);
 
     // since control flow is flattened, this is pretty simple