MergeSimilarFunctions optimization pass (#4414)

Merge similar functions that only differs constant values (like immediate
operand of const and call insts) by parameterization.
Performing this pass at post-link time can merge more functions across
objects. Inspired by Swift compiler's optimization which is derived from
LLVM's one:
https://github.com/apple/swift/blob/main/lib/LLVMPasses/LLVMMergeFunctions.cpp
https://github.com/llvm/llvm-project/blob/main/llvm/docs/MergeFunctions.rst

The basic ideas here are constant value parameterization and direct callee
parameterization by indirection.

Constant value parameterization is like below:

;;  Before
(func $big-const-42 (result i32)
  [[many instr 1]]
  (i32.const 44)
  [[many instr 2]]
)
(func $big-const-43 (result i32)
  [[many instr 1]]
  (i32.const 45)
  [[many instr 2]]
)

;; After
(func $byn$mgfn-shared$big-const-42 (result i32)
  [[many instr 1]]
  (local.get $0) ;; parameterized!!
  [[many instr 2]]
)
(func $big-const-42 (result i32)
  (call $byn$mgfn-shared$big-const-42
   (i32.const 42)
  )
)
(func $big-const-43 (result i32)
  (call $byn$mgfn-shared$big-const-42
   (i32.const 43)
  )
)

Direct callee parameterization is similar to the constant value parameterization,
but it parameterizes callee function i by ref.func instead. Therefore it is enabled
only when reference-types and typed-function-references features are enabled.

I saw 1 ~ 2 % reduction for SwiftWasm binary and Ruby's wasm port
using wasi-sdk, and 3 ~ 4.5% reduction for Unity WebGL binary when -Oz.

Author: kateinoigakukun

src/ir/ExpressionAnalyzer.cpp

@@ -262,7 +262,10 @@ struct Hasher {
   std::map<Name, Index> internalNames;
   ExpressionStack stack;
 
-  Hasher(Expression* curr, bool visitChildren) : visitChildren(visitChildren) {
+  Hasher(Expression* curr,
+         bool visitChildren,
+         ExpressionAnalyzer::ExprHasher custom)
+    : visitChildren(visitChildren) {
     stack.push_back(curr);
     // DELEGATE_CALLER_TARGET is a fake target used to denote delegating to
     // the caller. Add it here to prevent the unknown name error.
@@ -287,7 +290,11 @@ struct Hasher {
       // call_imports type, etc. The simplest thing is just to hash the
       // type for all of them.
       rehash(digest, curr->type.getID());
-      // Hash the contents of the expression.
+      // If the custom hasher handled this expr, then we have nothing to do.
+      if (custom(curr, digest)) {
+        continue;
+      }
+      // Hash the contents of the expression normally.
       hashExpression(curr);
     }
   }
@@ -365,12 +372,13 @@ struct Hasher {
 
 } // anonymous namespace
 
-size_t ExpressionAnalyzer::hash(Expression* curr) {
-  return Hasher(curr, true).digest;
+size_t ExpressionAnalyzer::flexibleHash(Expression* curr,
+                                        ExpressionAnalyzer::ExprHasher custom) {
+  return Hasher(curr, true, custom).digest;
 }
 
 size_t ExpressionAnalyzer::shallowHash(Expression* curr) {
-  return Hasher(curr, false).digest;
+  return Hasher(curr, false, ExpressionAnalyzer::nothingHasher).digest;
 }
 
 } // namespace wasm

src/ir/hashed.h

@@ -20,6 +20,7 @@
 #include "ir/utils.h"
 #include "support/hash.h"
 #include "wasm.h"
+#include <functional>
 
 namespace wasm {
 
@@ -30,9 +31,14 @@ struct FunctionHasher : public WalkerPass<PostWalker<FunctionHasher>> {
 
   struct Map : public std::map<Function*, size_t> {};
 
-  FunctionHasher(Map* output) : output(output) {}
+  FunctionHasher(Map* output, ExpressionAnalyzer::ExprHasher customHasher)
+    : output(output), customHasher(customHasher) {}
+  FunctionHasher(Map* output)
+    : output(output), customHasher(ExpressionAnalyzer::nothingHasher) {}
 
-  FunctionHasher* create() override { return new FunctionHasher(output); }
+  FunctionHasher* create() override {
+    return new FunctionHasher(output, customHasher);
+  }
 
   static Map createMap(Module* module) {
     Map hashes;
@@ -44,19 +50,29 @@ struct FunctionHasher : public WalkerPass<PostWalker<FunctionHasher>> {
     return hashes;
   }
 
-  void doWalkFunction(Function* func) { output->at(func) = hashFunction(func); }
+  void doWalkFunction(Function* func) {
+    output->at(func) = flexibleHashFunction(func, customHasher);
+  }
 
-  static size_t hashFunction(Function* func) {
+  static size_t
+  flexibleHashFunction(Function* func,
+                       ExpressionAnalyzer::ExprHasher customHasher) {
     auto digest = hash(func->type);
     for (auto type : func->vars) {
       rehash(digest, type.getID());
     }
-    hash_combine(digest, ExpressionAnalyzer::hash(func->body));
+    hash_combine(digest,
+                 ExpressionAnalyzer::flexibleHash(func->body, customHasher));
     return digest;
   }
 
+  static size_t hashFunction(Function* func) {
+    return flexibleHashFunction(func, ExpressionAnalyzer::nothingHasher);
+  }
+
 private:
   Map* output;
+  ExpressionAnalyzer::ExprHasher customHasher;
 };
 
 } // namespace wasm

src/ir/utils.h

@@ -82,9 +82,17 @@ struct ExpressionAnalyzer {
     return flexibleEqual(left, right, comparer);
   }
 
+  // Returns true if the expression is handled by the hasher.
+  using ExprHasher = std::function<bool(Expression*, size_t&)>;
+  static bool nothingHasher(Expression*, size_t&) { return false; }
+
+  static size_t flexibleHash(Expression* curr, ExprHasher hasher);
+
   // hash an expression, ignoring superficial details like specific internal
   // names
-  static size_t hash(Expression* curr);
+  static size_t hash(Expression* curr) {
+    return flexibleHash(curr, nothingHasher);
+  }
 
   // hash an expression, ignoring child nodes.
   static size_t shallowHash(Expression* curr);

src/passes/CMakeLists.txt

@@ -53,6 +53,7 @@ set(passes_SOURCES
   Memory64Lowering.cpp
   MemoryPacking.cpp
   MergeBlocks.cpp
+  MergeSimilarFunctions.cpp
   MergeLocals.cpp
   Metrics.cpp
   MinifyImportsAndExports.cpp