Add stacky mode to BinaryenIR

scripts/fuzz_opt.py

@@ -68,9 +68,11 @@ def random_size():
     return random.randint(INPUT_SIZE_MIN, 2 * INPUT_SIZE_MEAN - INPUT_SIZE_MIN)
 
 
-def run(cmd):
+def run(cmd, env_extension={}):
+    env = dict(os.environ)
+    env.update(env_extension)
     print(' '.join(cmd))
-    return subprocess.check_output(cmd, text=True)
+    return subprocess.check_output(cmd, text=True, env=env)
 
 
 def run_unchecked(cmd):
@@ -690,13 +692,40 @@ def can_run_on_feature_opts(self, feature_opts):
         return all([x in feature_opts for x in ['--disable-exception-handling', '--disable-simd', '--disable-tail-call', '--disable-reference-types', '--disable-multivalue']])
 
 
+class StackifyCodeSizeChecker(TestCaseHandler):
+    frequency = 1.0
+
+    def handle(self, wasm):
+        stack_ir = 'stackir.wasm'
+        stackify = 'stackify.wasm'
+        stackify_future = 'stackify-future.wasm'
+
+        command_base = [in_bin('wasm-opt'), wasm] + FEATURE_OPTS
+        stack_ir_command = command_base + ['--generate-stack-ir', '--optimize-stack-ir', '--optimize-level=3', '-o', stack_ir]
+        stackify_command = command_base + ['--stackify', '--stack-dce', '--stackify-locals', '--stack-remove-blocks', '--stack-dce', '--coalesce-locals', '--optimize-level=3', '-o', stackify]
+        stackify_future_command = command_base + ['--stackify', '--stack-dce', '--stackify-locals', '--stack-remove-blocks', '--stack-dce', '--coalesce-locals', '--optimize-level=3', '-o', stackify_future]
+
+        run(stack_ir_command, {'BINARYEN_USE_STACKIFY': '0'})
+        run(stackify_command, {'BINARYEN_USE_STACKIFY': '1'})
+        run(stackify_future_command, {'BINARYEN_USE_STACKIFY': '2'})
+
+        stack_ir_size = os.stat(stack_ir).st_size
+        stackify_size = os.stat(stackify).st_size
+        stackify_future_size = os.stat(stackify_future).st_size
+
+        print("sizes:", stack_ir_size, stackify_size, stackify_future_size)
+        assert stackify_size <= stack_ir_size, "Stackify not as good as Stack IR"
+        assert stackify_future_size <= stackify_size, "New pipeline is not as good"
+
+
 # The global list of all test case handlers
 testcase_handlers = [
     FuzzExec(),
     CompareVMs(),
-    CheckDeterminism(),
-    Wasm2JS(),
-    Asyncify(),
+    # CheckDeterminism(),
+    # Wasm2JS(),
+    # Asyncify(),
+    StackifyCodeSizeChecker(),
 ]
 
 
@@ -728,7 +757,7 @@ def test_one(random_input, opts, given_wasm):
     print('pre wasm size:', wasm_size)
 
     # create a second wasm for handlers that want to look at pairs.
-    generate_command = [in_bin('wasm-opt'), 'a.wasm', '-o', 'b.wasm'] + opts + FUZZ_OPTS + FEATURE_OPTS
+    generate_command = [in_bin('wasm-opt'), 'a.wasm', '-o', 'b.wasm'] + FUZZ_OPTS + opts + FEATURE_OPTS
     if PRINT_WATS:
         printed = run(generate_command + ['--print'])
         with open('b.printed.wast', 'w') as f:
@@ -828,10 +857,20 @@ def write_commands(commands, filename):
     ["--simplify-locals-notee"],
     ["--simplify-locals-notee-nostructure"],
     ["--ssa"],
+    # ["--stackify", "--unstackify"],
     ["--vacuum"],
 ]
 
 
+stack_opt_choices = [
+    ["--stack-dce"],
+    ["--stack-remove-blocks"],
+    ["--stackify-locals"],
+    ["--stackify-drops"],
+    ["--coalesce-locals"],
+]
+
+
 def randomize_opt_flags():
     flag_groups = []
     has_flatten = False
@@ -847,10 +886,22 @@ def randomize_opt_flags():
         flag_groups.append(choice)
         if len(flag_groups) > 20 or random.random() < 0.3:
             break
+
+    # stacky opts
+    if random.random() < 0.5:
+        stack_opt_group = ['--stackify']
+        while len(flag_groups) <= 20 and random.random() < 0.3:
+            stack_opt_group.extend(random.choice(stack_opt_choices))
+        flag_groups.append(stack_opt_group)
+
     # maybe add an extra round trip
     if random.random() < 0.5:
         pos = random.randint(0, len(flag_groups))
         flag_groups = flag_groups[:pos] + [['--roundtrip']] + flag_groups[pos:]
+    # maybe add an extra stackification round trip
+    # if True:
+    #     pos = random.randint(0, len(flag_groups))
+    #     flag_groups = flag_groups[:pos] + [['--stackify', '--unstackify']] + flag_groups[pos:]
     ret = [flag for group in flag_groups for flag in group]
     # modifiers (if not already implied by a -O? option)
     if '-O' not in str(ret):

src/binaryen-c.cpp

@@ -3454,6 +3454,7 @@ void BinaryenModuleOptimize(BinaryenModuleRef module) {
   PassRunner passRunner((Module*)module);
   passRunner.options = globalPassOptions;
   passRunner.addDefaultOptimizationPasses();
+  passRunner.addDefaultPreWritingPasses();
   passRunner.run();
 }
 
@@ -3706,6 +3707,7 @@ void BinaryenFunctionOptimize(BinaryenFunctionRef func,
   PassRunner passRunner((Module*)module);
   passRunner.options = globalPassOptions;
   passRunner.addDefaultOptimizationPasses();
+  passRunner.addDefaultPreWritingPasses();
   passRunner.runOnFunction((Function*)func);
 }
 void BinaryenFunctionRunPasses(BinaryenFunctionRef func,
@@ -3995,6 +3997,11 @@ class CExpressionRunner final
                     Index maxLoopIterations)
     : ConstantExpressionRunner<CExpressionRunner>(
         module, flags, maxDepth, maxLoopIterations) {}
+
+  bool isStacky() {
+    WASM_UNREACHABLE("TODO: support stackiness by tracking current function in "
+                     "traversing expression runner");
+  }
 };
 
 } // namespace wasm

src/cfg/liveness-traversal.h

@@ -104,6 +104,7 @@ struct LivenessWalker : public CFGWalker<SubType, VisitorType, Liveness> {
     typename CFGWalker<SubType, VisitorType, Liveness>::BasicBlock BasicBlock;
 
   Index numLocals;
+  IRProfile profile;
   std::unordered_set<BasicBlock*> liveBlocks;
   // canonicalized - accesses should check (low, high)
   // TODO: use a map for high N, as this tends to be sparse? or don't look at
@@ -118,6 +119,7 @@ struct LivenessWalker : public CFGWalker<SubType, VisitorType, Liveness> {
     auto* curr = (*currp)->cast<LocalGet>();
     // if in unreachable code, ignore
     if (!self->currBasicBlock) {
+      // TODO: This probably needs to be adjusted for Stacky code as well
       *currp = Builder(*self->getModule()).replaceWithIdenticalType(curr);
       return;
     }
@@ -131,9 +133,14 @@ struct LivenessWalker : public CFGWalker<SubType, VisitorType, Liveness> {
     // if it has side effects)
     if (!self->currBasicBlock) {
       if (curr->isTee()) {
-        *currp = curr->value;
+        if (self->profile == IRProfile::Normal) {
+          *currp = curr->value;
+        } else {
+          assert(self->profile == IRProfile::Stacky);
+          ExpressionManipulator::nop(curr);
+        }
       } else {
-        *currp = Builder(*self->getModule()).makeDrop(curr->value);
+        ExpressionManipulator::drop(curr, curr->value);
       }
       return;
     }
@@ -190,6 +197,7 @@ struct LivenessWalker : public CFGWalker<SubType, VisitorType, Liveness> {
 
   void doWalkFunction(Function* func) {
     numLocals = func->getNumLocals();
+    profile = func->profile;
     assert(canRun(func));
     copies.resize(numLocals * numLocals);
     std::fill(copies.begin(), copies.end(), 0);

src/ir/CMakeLists.txt

@@ -4,6 +4,7 @@ set(ir_SOURCES
   ExpressionManipulator.cpp
   LocalGraph.cpp
   ReFinalize.cpp
+  stack-utils.cpp
   ${ir_HEADERS}
 )
 add_library(ir OBJECT ${ir_SOURCES})

src/ir/ReFinalize.cpp

@@ -16,6 +16,7 @@
 
 #include "ir/branch-utils.h"
 #include "ir/find_all.h"
+#include "ir/stack-utils.h"
 #include "ir/utils.h"
 
 namespace wasm {
@@ -44,21 +45,38 @@ void ReFinalize::visitBlock(Block* curr) {
     curr->type = Type::none;
     return;
   }
+
+  // Stacky blocks can be unreachable even if they yield concrete values as long
+  // as they are not branch targets as well.
+  bool unreachableIfNotBranchTarget = false;
+
   // Get the least upper bound type of the last element and all branch return
   // values
-  curr->type = curr->list.back()->type;
+  if (getProfile() == IRProfile::Normal) {
+    curr->type = curr->list.back()->type;
+  } else {
+    assert(getProfile() == IRProfile::Stacky);
+    StackSignature sig(curr->list.begin(), curr->list.end());
+    curr->type = (sig.results == Type::none && sig.unreachable)
+                   ? Type::unreachable
+                   : sig.results;
+    unreachableIfNotBranchTarget = sig.unreachable;
+  }
   if (curr->name.is()) {
     auto iter = breakValues.find(curr->name);
     if (iter != breakValues.end()) {
       curr->type = Type::getLeastUpperBound(curr->type, iter->second);
       return;
     }
   }
+  if (unreachableIfNotBranchTarget) {
+    curr->type = Type::unreachable;
+  }
   if (curr->type == Type::unreachable) {
     return;
   }
   // type is none, but we might be unreachable
-  if (curr->type == Type::none) {
+  if (curr->type == Type::none && getProfile() == IRProfile::Normal) {
     for (auto* child : curr->list) {
       if (child->type == Type::unreachable) {
         curr->type = Type::unreachable;

src/ir/iteration.h

@@ -80,6 +80,7 @@ class ChildIterator {
 
   Iterator begin() const { return Iterator(*this, 0); }
   Iterator end() const { return Iterator(*this, children.size()); }
+  Index size() const { return children.size(); }
 };
 
 // Returns true if the current expression contains a certain kind of expression,

src/ir/manipulation.h

@@ -40,6 +40,14 @@ template<typename InputType> inline Nop* nop(InputType* target) {
   return ret;
 }
 
+template<typename InputType>
+inline Drop* drop(InputType* target, Expression* value) {
+  auto* ret = convert<InputType, Drop>(target);
+  ret->value = value;
+  ret->finalize();
+  return ret;
+}
+
 template<typename InputType> inline RefNull* refNull(InputType* target) {
   auto* ret = convert<InputType, RefNull>(target);
   ret->finalize();

src/ir/module-utils.h

@@ -402,8 +402,10 @@ collectSignatures(Module& wasm,
     struct TypeCounter
       : PostWalker<TypeCounter, UnifiedExpressionVisitor<TypeCounter>> {
       Counts& counts;
+      Function* func;
 
-      TypeCounter(Counts& counts) : counts(counts) {}
+      TypeCounter(Counts& counts, Function* func)
+        : counts(counts), func(func) {}
       void visitExpression(Expression* curr) {
         if (auto* call = curr->dynCast<CallIndirect>()) {
           counts[call->sig]++;
@@ -415,7 +417,7 @@ collectSignatures(Module& wasm,
         }
       }
     };
-    TypeCounter(counts).walk(func->body);
+    TypeCounter(counts, func).walk(func->body);
   };
 
   ModuleUtils::ParallelFunctionAnalysis<Counts> analysis(wasm, updateCounts);