StackIR: Run StackIR during binary writing

CHANGELOG.md

@@ -15,6 +15,19 @@ full changeset diff at the end of each section.
 Current Trunk
 -------------
 
+ - StackIR is now handled entirely during binary writing. This is mostly not
+   noticeable, except that:
+   - Text output no longer notes `(; has Stack IR ;)` (as Stack IR only exists
+     during binary writing).
+   - `--generate-stack-ir`, `--optimize-stack-ir`, and `--print-stack-ir` are
+     now flags and not passes. That means the order of operations may seem
+     different, as they apply during binary writing (or, if no binary is written
+     but we were still asked to print StackIR, `wasm-opt` does it at the very
+     end).
+   - Whether to generate, optimize, and print StackIR is now noted as part of
+     the PassOptions. As a result `BinaryenModulePrintStackIR` and similar APIs
+     do not receive an `optimize` flag, as they read the PassOption
+     `optimizeStackIR` instead.
  - The new, standards-compliant text parser is now the default. `wasm-opt` has a
    `--deprecated-wat-parser` flag that will switch back to using the old text
    parser, but that option will go away soon.

src/binaryen-c.cpp

@@ -5590,8 +5590,8 @@ void BinaryenModulePrint(BinaryenModuleRef module) {
   std::cout << *(Module*)module;
 }
 
-void BinaryenModulePrintStackIR(BinaryenModuleRef module, bool optimize) {
-  wasm::printStackIR(std::cout, (Module*)module, optimize);
+void BinaryenModulePrintStackIR(BinaryenModuleRef module) {
+  wasm::printStackIR(std::cout, (Module*)module, globalPassOptions);
 }
 
 void BinaryenModulePrintAsmjs(BinaryenModuleRef module) {
@@ -5737,7 +5737,7 @@ static BinaryenBufferSizes writeModule(BinaryenModuleRef module,
                                        char* sourceMap,
                                        size_t sourceMapSize) {
   BufferWithRandomAccess buffer;
-  WasmBinaryWriter writer((Module*)module, buffer);
+  WasmBinaryWriter writer((Module*)module, buffer, globalPassOptions);
   writer.setNamesSection(globalPassOptions.debugInfo);
   std::ostringstream os;
   if (sourceMapUrl) {
@@ -5778,12 +5778,11 @@ size_t BinaryenModuleWriteText(BinaryenModuleRef module,
 
 size_t BinaryenModuleWriteStackIR(BinaryenModuleRef module,
                                   char* output,
-                                  size_t outputSize,
-                                  bool optimize) {
+                                  size_t outputSize) {
   // use a stringstream as an std::ostream. Extract the std::string
   // representation, and then store in the output.
   std::stringstream ss;
-  wasm::printStackIR(ss, (Module*)module, optimize);
+  wasm::printStackIR(ss, (Module*)module, globalPassOptions);
 
   const auto temp = ss.str();
   const auto ctemp = temp.c_str();
@@ -5808,7 +5807,7 @@ BinaryenModuleAllocateAndWriteResult
 BinaryenModuleAllocateAndWrite(BinaryenModuleRef module,
                                const char* sourceMapUrl) {
   BufferWithRandomAccess buffer;
-  WasmBinaryWriter writer((Module*)module, buffer);
+  WasmBinaryWriter writer((Module*)module, buffer, globalPassOptions);
   writer.setNamesSection(globalPassOptions.debugInfo);
   std::ostringstream os;
   if (sourceMapUrl) {
@@ -5842,13 +5841,12 @@ char* BinaryenModuleAllocateAndWriteText(BinaryenModuleRef module) {
   return output;
 }
 
-char* BinaryenModuleAllocateAndWriteStackIR(BinaryenModuleRef module,
-                                            bool optimize) {
+char* BinaryenModuleAllocateAndWriteStackIR(BinaryenModuleRef module) {
   std::ostringstream os;
   bool colors = Colors::isEnabled();
 
   Colors::setEnabled(false); // do not use colors for writing
-  wasm::printStackIR(os, (Module*)module, optimize);
+  wasm::printStackIR(os, (Module*)module, globalPassOptions);
   Colors::setEnabled(colors); // restore colors state
 
   auto str = os.str();

src/binaryen-c.h

@@ -2984,8 +2984,7 @@ BINARYEN_API BinaryenModuleRef BinaryenModuleParse(const char* text);
 BINARYEN_API void BinaryenModulePrint(BinaryenModuleRef module);
 
 // Print a module to stdout in stack IR text format. Useful for debugging.
-BINARYEN_API void BinaryenModulePrintStackIR(BinaryenModuleRef module,
-                                             bool optimize);
+BINARYEN_API void BinaryenModulePrintStackIR(BinaryenModuleRef module);
 
 // Print a module to stdout in asm.js syntax.
 BINARYEN_API void BinaryenModulePrintAsmjs(BinaryenModuleRef module);
@@ -3126,8 +3125,7 @@ BINARYEN_API size_t BinaryenModuleWriteText(BinaryenModuleRef module,
 //         outputSize
 BINARYEN_API size_t BinaryenModuleWriteStackIR(BinaryenModuleRef module,
                                                char* output,
-                                               size_t outputSize,
-                                               bool optimize);
+                                               size_t outputSize);
 
 typedef struct BinaryenBufferSizes {
   size_t outputBytes;
@@ -3173,7 +3171,7 @@ BINARYEN_API char* BinaryenModuleAllocateAndWriteText(BinaryenModuleRef module);
 // char* with malloc(), and expects the user to free() them manually
 // once not needed anymore.
 BINARYEN_API char*
-BinaryenModuleAllocateAndWriteStackIR(BinaryenModuleRef module, bool optimize);
+BinaryenModuleAllocateAndWriteStackIR(BinaryenModuleRef module);
 
 // Deserialize a module from binary form, assuming the MVP feature set.
 BINARYEN_API BinaryenModuleRef BinaryenModuleRead(char* input,

src/ir/module-utils.cpp

@@ -69,9 +69,6 @@ Function* copyFunction(Function* func,
   ret->base = func->base;
   ret->noFullInline = func->noFullInline;
   ret->noPartialInline = func->noPartialInline;
-
-  // TODO: copy Stack IR
-  assert(!func->stackIR);
   return out.addFunction(std::move(ret));
 }
 

src/js/binaryen.js-post.js

@@ -2688,8 +2688,8 @@ function wrapModule(module, self = {}) {
     if (textPtr) _free(textPtr);
     return text;
   };
-  self['emitStackIR'] = function(optimize) {
-    let textPtr = Module['_BinaryenModuleAllocateAndWriteStackIR'](module, optimize);
+  self['emitStackIR'] = function() {
+    let textPtr = Module['_BinaryenModuleAllocateAndWriteStackIR'](module);
     let text = UTF8ToString(textPtr);
     if (textPtr) _free(textPtr);
     return text;

src/pass.h

@@ -219,6 +219,16 @@ struct PassOptions {
   bool closedWorld = false;
   // Whether to try to preserve debug info through, which are special calls.
   bool debugInfo = false;
+  // Whether to generate StackIR during binary writing. This is on by default
+  // in -O2 and above.
+  bool generateStackIR = false;
+  // Whether to optimize StackIR during binary writing. How we optimize depends
+  // on other optimization flags like optimizeLevel. This is on by default in
+  // -O2 and above.
+  bool optimizeStackIR = false;
+  // Whether to print StackIR during binary writing, and if so to what stream.
+  // This is mainly useful for debugging.
+  std::optional<std::ostream*> printStackIR;
   // Whether we are targeting JS. In that case we want to avoid emitting things
   // in the optimizer that do not translate well to JS, or that could cause us
   // to need extra lowering work or even a loop (where we optimize to something

src/passes/CMakeLists.txt

@@ -87,7 +87,6 @@ set(passes_SOURCES
   PrintFunctionMap.cpp
   RoundTrip.cpp
   SetGlobals.cpp
-  StackIR.cpp
   SignaturePruning.cpp
   SignatureRefining.cpp
   SignExtLowering.cpp

src/passes/Metrics.cpp

@@ -94,7 +94,7 @@ struct Metrics
       printCounts("global");
       // compute binary info, so we know function sizes
       BufferWithRandomAccess buffer;
-      WasmBinaryWriter writer(module, buffer);
+      WasmBinaryWriter writer(module, buffer, getPassOptions());
       writer.write();
       // print for each function
       Index binaryIndex = 0;
@@ -108,14 +108,14 @@ struct Metrics
       });
       // print for each export how much code size is due to it, i.e.,
       // how much the module could shrink without it.
-      auto sizeAfterGlobalCleanup = [](Module* module) {
+      auto sizeAfterGlobalCleanup = [&](Module* module) {
         PassRunner runner(module,
                           PassOptions::getWithDefaultOptimizationOptions());
         runner.setIsNested(true);
         runner.addDefaultGlobalOptimizationPostPasses(); // remove stuff
         runner.run();
         BufferWithRandomAccess buffer;
-        WasmBinaryWriter writer(module, buffer);
+        WasmBinaryWriter writer(module, buffer, getPassOptions());
         writer.write();
         return buffer.size();
       };

src/passes/Monomorphize.cpp

@@ -151,11 +151,6 @@ struct Monomorphize : public Pass {
     // monomorphizing.
 
     // Create a new function with refined parameters as a copy of the original.
-    // (Note we must clear stack IR on the original: atm we do not have the
-    // ability to copy stack IR, so we'd hit an internal error. But as we will
-    // be optimizing the function anyhow, we'd be throwing away stack IR later
-    // so this isn't a problem.)
-    func->stackIR.reset();
     auto refinedTarget = Names::getValidFunctionName(*module, target);
     auto* refinedFunc = ModuleUtils::copyFunction(func, *module, refinedTarget);
     TypeUpdating::updateParamTypes(refinedFunc, refinedTypes, *module);

src/passes/Print.cpp

@@ -109,11 +109,11 @@ struct PrintSExpression : public UnifiedExpressionVisitor<PrintSExpression> {
   const char* maybeSpace;
   const char* maybeNewLine;
 
-  bool full = false;    // whether to not elide nodes in output when possible
-                        // (like implicit blocks) and to emit types
-  bool stackIR = false; // whether to print stack IR if it is present
-                        // (if false, and Stack IR is there, we just
-                        // note it exists)
+  // Whether to not elide nodes in output when possible (like implicit blocks)
+  // and to emit types.
+  bool full = false;
+  // If present, it contains StackIR that we will print.
+  std::optional<ModuleStackIR> moduleStackIR;
 
   Module* currModule = nullptr;
   Function* currFunction = nullptr;
@@ -268,7 +268,9 @@ struct PrintSExpression : public UnifiedExpressionVisitor<PrintSExpression> {
 
   void setFull(bool full_) { full = full_; }
 
-  void setStackIR(bool stackIR_) { stackIR = stackIR_; }
+  void generateStackIR(const PassOptions& options) {
+    moduleStackIR.emplace(*currModule, options);
+  }
 
   void setDebugInfo(bool debugInfo_) { debugInfo = debugInfo_; }
 
@@ -2978,9 +2980,6 @@ void PrintSExpression::visitDefinedFunction(Function* curr) {
     o << " (type ";
     printHeapType(curr->type) << ')';
   }
-  if (!stackIR && curr->stackIR && !minify) {
-    o << " (; has Stack IR ;)";
-  }
   if (curr->getParams().size() > 0) {
     Index i = 0;
     for (const auto& param : curr->getParams()) {
@@ -3007,7 +3006,13 @@ void PrintSExpression::visitDefinedFunction(Function* curr) {
     o << maybeNewLine;
   }
   // Print the body.
-  if (!stackIR || !curr->stackIR) {
+  StackIR* stackIR = nullptr;
+  if (moduleStackIR) {
+    stackIR = moduleStackIR->getStackIROrNull(curr);
+  }
+  if (stackIR) {
+    printStackIR(stackIR, *this);
+  } else {
     // It is ok to emit a block here, as a function can directly contain a
     // list, even if our ast avoids that for simplicity. We can just do that
     // optimization here..
@@ -3021,9 +3026,6 @@ void PrintSExpression::visitDefinedFunction(Function* curr) {
       printFullLine(curr->body);
     }
     assert(controlFlowDepth == 0);
-  } else {
-    // Print the stack IR.
-    printStackIR(curr->stackIR.get(), *this);
   }
   if (currFunction->epilogLocation.size()) {
     // Print last debug location: mix of decIndent and printDebugLocation
@@ -3430,14 +3432,12 @@ class PrintStackIR : public Printer {
   void run(Module* module) override {
     PrintSExpression print(o);
     print.setDebugInfo(getPassOptions().debugInfo);
-    print.setStackIR(true);
     print.currModule = module;
+    print.generateStackIR(getPassOptions());
     print.visitModule(module);
   }
 };
 
-Pass* createPrintStackIRPass() { return new PrintStackIR(); }
-
 static std::ostream& printExpression(Expression* expression,
                                      std::ostream& o,
                                      bool minify,
@@ -3602,12 +3602,9 @@ static std::ostream& printStackIR(StackIR* ir, PrintSExpression& printer) {
   return o;
 }
 
-std::ostream& printStackIR(std::ostream& o, Module* module, bool optimize) {
-  wasm::PassRunner runner(module);
-  runner.add("generate-stack-ir");
-  if (optimize) {
-    runner.add("optimize-stack-ir");
-  }
+std::ostream&
+printStackIR(std::ostream& o, Module* module, const PassOptions& options) {
+  wasm::PassRunner runner(module, options);
   runner.add(std::make_unique<PrintStackIR>(&o));
   runner.run();
   return o;
@@ -3659,9 +3656,4 @@ std::ostream& operator<<(std::ostream& o, wasm::StackInst& inst) {
   return wasm::printStackInst(&inst, o);
 }
 
-std::ostream& operator<<(std::ostream& o, wasm::StackIR& ir) {
-  wasm::PrintSExpression printer(o);
-  return wasm::printStackIR(&ir, printer);
-}
-
 } // namespace std

src/passes/RoundTrip.cpp

@@ -41,7 +41,7 @@ struct RoundTrip : public Pass {
     // to tell the builder which features to build with.
     auto features = module->features;
     // Write, clear, and read the module
-    WasmBinaryWriter(module, buffer).write();
+    WasmBinaryWriter(module, buffer, getPassOptions()).write();
     ModuleUtils::clearModule(*module);
     auto input = buffer.getAsChars();
     WasmBinaryReader parser(*module, features, input);