Refactor and optimize binary writing type collection (#2478)

Create a new ParallelFunctionAnalysis helper, which lets us
run in parallel on all functions and collect info from them,
without manually handling locks etc.

Use that in the binary writing code's type collection logic,
avoiding a lock for each type increment.

Also add Signature printing which was useful to debug this.

Author: kripken

src/ir/module-utils.h

@@ -262,33 +262,20 @@ template<typename T> inline void iterDefinedEvents(Module& wasm, T visitor) {
   }
 }
 
-// Helper class for analyzing the call graph.
-//
-// Provides hooks for running some initial calculation on each function (which
-// is done in parallel), writing to a FunctionInfo structure for each function.
-// Then you can call propagateBack() to propagate a property of interest to the
-// calling functions, transitively.
-//
-// For example, if some functions are known to call an import "foo", then you
-// can use this to find which functions call something that might eventually
-// reach foo, by initially marking the direct callers as "calling foo" and
-// propagating that backwards.
-template<typename T> struct CallGraphPropertyAnalysis {
+// Helper class for performing an operation on all the functions in the module,
+// in parallel, with an Info object for each one that can contain results of
+// some computation that the operation performs.
+// The operation performend should not modify the wasm module in any way.
+// TODO: enforce this
+template<typename T> struct ParallelFunctionAnalysis {
   Module& wasm;
 
-  // The basic information for each function about whom it calls and who is
-  // called by it.
-  struct FunctionInfo {
-    std::set<Function*> callsTo;
-    std::set<Function*> calledBy;
-  };
-
   typedef std::map<Function*, T> Map;
   Map map;
 
   typedef std::function<void(Function*, T&)> Func;
 
-  CallGraphPropertyAnalysis(Module& wasm, Func work) : wasm(wasm) {
+  ParallelFunctionAnalysis(Module& wasm, Func work) : wasm(wasm) {
     // Fill in map, as we operate on it in parallel (each function to its own
     // entry).
     for (auto& func : wasm.functions) {
@@ -304,30 +291,78 @@ template<typename T> struct CallGraphPropertyAnalysis {
 
     struct Mapper : public WalkerPass<PostWalker<Mapper>> {
       bool isFunctionParallel() override { return true; }
+      bool modifiesBinaryenIR() override { return false; }
 
-      Mapper(Module* module, Map* map, Func work)
+      Mapper(Module& module, Map& map, Func work)
         : module(module), map(map), work(work) {}
 
       Mapper* create() override { return new Mapper(module, map, work); }
 
-      void visitCall(Call* curr) {
-        (*map)[this->getFunction()].callsTo.insert(
-          module->getFunction(curr->target));
-      }
-
-      void visitFunction(Function* curr) {
-        assert((*map).count(curr));
-        work(curr, (*map)[curr]);
+      void doWalkFunction(Function* curr) {
+        assert(map.count(curr));
+        work(curr, map[curr]);
       }
 
     private:
-      Module* module;
-      Map* map;
+      Module& module;
+      Map& map;
       Func work;
     };
 
     PassRunner runner(&wasm);
-    Mapper(&wasm, &map, work).run(&runner, &wasm);
+    Mapper(wasm, map, work).run(&runner, &wasm);
+  }
+};
+
+// Helper class for analyzing the call graph.
+//
+// Provides hooks for running some initial calculation on each function (which
+// is done in parallel), writing to a FunctionInfo structure for each function.
+// Then you can call propagateBack() to propagate a property of interest to the
+// calling functions, transitively.
+//
+// For example, if some functions are known to call an import "foo", then you
+// can use this to find which functions call something that might eventually
+// reach foo, by initially marking the direct callers as "calling foo" and
+// propagating that backwards.
+template<typename T> struct CallGraphPropertyAnalysis {
+  Module& wasm;
+
+  // The basic information for each function about whom it calls and who is
+  // called by it.
+  struct FunctionInfo {
+    std::set<Function*> callsTo;
+    std::set<Function*> calledBy;
+  };
+
+  typedef std::map<Function*, T> Map;
+  Map map;
+
+  typedef std::function<void(Function*, T&)> Func;
+
+  CallGraphPropertyAnalysis(Module& wasm, Func work) : wasm(wasm) {
+    ParallelFunctionAnalysis<T> analysis(wasm, [&](Function* func, T& info) {
+      work(func, info);
+      if (func->imported()) {
+        return;
+      }
+      struct Mapper : public PostWalker<Mapper> {
+        Mapper(Module* module, T& info, Func work)
+          : module(module), info(info), work(work) {}
+
+        void visitCall(Call* curr) {
+          info.callsTo.insert(module->getFunction(curr->target));
+        }
+
+      private:
+        Module* module;
+        T& info;
+        Func work;
+      } mapper(&wasm, info, work);
+      mapper.walk(func->body);
+    });
+
+    map.swap(analysis.map);
 
     // Find what is called by what.
     for (auto& pair : map) {

src/wasm-type.h

@@ -96,10 +96,6 @@ struct ResultType {
   std::string toString() const;
 };
 
-std::ostream& operator<<(std::ostream& os, Type t);
-std::ostream& operator<<(std::ostream& os, ParamType t);
-std::ostream& operator<<(std::ostream& os, ResultType t);
-
 struct Signature {
   Type params;
   Type results;
@@ -112,6 +108,11 @@ struct Signature {
   bool operator<(const Signature& other) const;
 };
 
+std::ostream& operator<<(std::ostream& os, Type t);
+std::ostream& operator<<(std::ostream& os, ParamType t);
+std::ostream& operator<<(std::ostream& os, ResultType t);
+std::ostream& operator<<(std::ostream& os, Signature t);
+
 constexpr Type none = Type::none;
 constexpr Type i32 = Type::i32;
 constexpr Type i64 = Type::i64;

src/wasm/wasm-binary.cpp

@@ -16,68 +16,52 @@
 
 #include <algorithm>
 #include <fstream>
-#include <shared_mutex>
 
+#include "ir/module-utils.h"
 #include "support/bits.h"
 #include "wasm-binary.h"
 #include "wasm-stack.h"
 
 namespace wasm {
 
 void WasmBinaryWriter::prepare() {
-  // Collect function types and their frequencies
-  using Counts = std::unordered_map<Signature, size_t>;
-  using AtomicCounts = std::unordered_map<Signature, std::atomic_size_t>;
+  // Collect function types and their frequencies. Collect information in each
+  // function in parallel, then merge.
+  typedef std::unordered_map<Signature, size_t> Counts;
+  ModuleUtils::ParallelFunctionAnalysis<Counts> analysis(
+    *wasm, [&](Function* func, Counts& counts) {
+      if (func->imported()) {
+        return;
+      }
+      struct TypeCounter : PostWalker<TypeCounter> {
+        Module& wasm;
+        Counts& counts;
+
+        TypeCounter(Module& wasm, Counts& counts)
+          : wasm(wasm), counts(counts) {}
+
+        void visitCallIndirect(CallIndirect* curr) {
+          auto* type = wasm.getFunctionType(curr->fullType);
+          Signature sig(Type(type->params), type->result);
+          counts[sig]++;
+        }
+      };
+      TypeCounter(*wasm, counts).walk(func->body);
+    });
+  // Collect all the counts.
   Counts counts;
   for (auto& curr : wasm->functions) {
     counts[Signature(Type(curr->params), curr->result)]++;
   }
   for (auto& curr : wasm->events) {
     counts[curr->sig]++;
   }
-
-  // Parallelize collection of call_indirect type counts
-  struct TypeCounter : WalkerPass<PostWalker<TypeCounter>> {
-    AtomicCounts& counts;
-    std::shared_timed_mutex& mutex;
-    TypeCounter(AtomicCounts& counts, std::shared_timed_mutex& mutex)
-      : counts(counts), mutex(mutex) {}
-    bool isFunctionParallel() override { return true; }
-    bool modifiesBinaryenIR() override { return false; }
-    void visitCallIndirect(CallIndirect* curr) {
-      auto* type = getModule()->getFunctionType(curr->fullType);
-      Signature sig(Type(type->params), type->result);
-      {
-        std::shared_lock<std::shared_timed_mutex> lock(mutex);
-        auto it = counts.find(sig);
-        if (it != counts.end()) {
-          it->second++;
-          return;
-        }
-      }
-      {
-        std::lock_guard<std::shared_timed_mutex> lock(mutex);
-        counts[sig]++;
-      }
+  for (auto& pair : analysis.map) {
+    Counts& functionCounts = pair.second;
+    for (auto& innerPair : functionCounts) {
+      counts[innerPair.first] += innerPair.second;
     }
-    Pass* create() override { return new TypeCounter(counts, mutex); }
-  };
-
-  std::shared_timed_mutex mutex;
-  AtomicCounts parallelCounts;
-  for (auto& kv : counts) {
-    parallelCounts[kv.first] = 0;
   }
-
-  TypeCounter counter(parallelCounts, mutex);
-  PassRunner runner(wasm);
-  runner.setIsNested(true);
-  counter.run(&runner, wasm);
-
-  for (auto& kv : parallelCounts) {
-    counts[kv.first] += kv.second;
-  }
-
   std::vector<std::pair<Signature, size_t>> sorted(counts.begin(),
                                                    counts.end());
   std::sort(sorted.begin(), sorted.end(), [&](auto a, auto b) {

src/wasm/wasm-type.cpp

@@ -222,6 +222,10 @@ std::ostream& operator<<(std::ostream& os, ResultType param) {
   return printPrefixedTypes(os, "result", param.type);
 }
 
+std::ostream& operator<<(std::ostream& os, Signature sig) {
+  return os << "Signature(" << sig.params << " => " << sig.results << ")";
+}
+
 std::string Type::toString() const { return genericToString(*this); }
 
 std::string ParamType::toString() const { return genericToString(*this); }

test/unit/input/stack_ir.wast

@@ -0,0 +1,12 @@
+(module
+  (import "env" "bar" (func $bar (param i32) (result i32)))
+  (func "foo1" (result i32)
+    (local $x i32)
+    (local.set $x (call $bar (i32.const 0)))
+    (drop
+     (call $bar (i32.const 1))
+    )
+    (local.get $x) ;; local2stack can help here
+  )
+)
+

test/unit/test_stack_ir.py

@@ -0,0 +1,19 @@
+import os
+from scripts.test import shared
+from . import utils
+
+
+class StackIRTest(utils.BinaryenTestCase):
+    # test that stack IR opts make a difference.
+    def test_stack_ir_opts(self):
+        path = self.input_path('stack_ir.wast')
+        opt = shared.run_process(shared.WASM_OPT + [path, '-O', '--generate-stack-ir', '--optimize-stack-ir', '--print-stack-ir', '-o', 'a.wasm'], capture_output=True).stdout
+        nonopt = shared.run_process(shared.WASM_OPT + [path, '-O', '--generate-stack-ir', '--print-stack-ir', '-o', 'b.wasm'], capture_output=True).stdout
+        # see a difference in the printed stack IR (the optimizations let us
+        # remove a pair of local.set/get)
+        self.assertNotEqual(opt, nonopt)
+        self.assertLess(len(opt), len(nonopt))
+        # see a difference in the actual emitted wasm binary.
+        opt_size = os.path.getsize('a.wasm')
+        nonopt_size = os.path.getsize('b.wasm')
+        self.assertLess(opt_size, nonopt_size)