[wasm-split] Pre-create trampolines for global initializers

src/ir/module-splitting.cpp

@@ -73,7 +73,6 @@
 //      from the IR before splitting.
 //
 #include "ir/module-splitting.h"
-#include "ir/export-utils.h"
 #include "ir/find_all.h"
 #include "ir/module-utils.h"
 #include "ir/names.h"
@@ -336,6 +335,7 @@ struct ModuleSplitter {
   void exportImportCalledPrimaryFunctions();
   void setupTablePatching();
   void shareImportableItems();
+  void cleanupUnusedTrampolines();
 
   ModuleSplitter(Module& primary, const Config& config)
     : config(config), primary(primary), tableManager(primary),
@@ -348,6 +348,7 @@ struct ModuleSplitter {
     exportImportCalledPrimaryFunctions();
     setupTablePatching();
     shareImportableItems();
+    cleanupUnusedTrampolines();
   }
 };
 
@@ -508,7 +509,9 @@ Name ModuleSplitter::getTrampoline(Name funcName) {
     primary, std::string("trampoline_") + funcName.toString());
   it->second = trampoline;
 
-  // Generate the call and the function.
+  // Generate the call and the function. We generate a direct call here, but
+  // this will be converted to a call_indirect in
+  // indirectCallsToSecondaryFunctions.
   std::vector<Expression*> args;
   for (Index i = 0; i < oldFunc->getNumParams(); i++) {
     args.push_back(builder.makeLocalGet(i, oldFunc->getLocalType(i)));
@@ -559,6 +562,21 @@ static void walkSegments(Walker& walker, Module* module) {
 }
 
 void ModuleSplitter::indirectReferencesToSecondaryFunctions() {
+  // Pre-create trampolines for secondary functions referenced in global
+  // initializers. We don't mutate the globals yet (that happens later in
+  // shareImportableItems), but we must create the trampolines now so they
+  // are processed correctly by indirectCallsToSecondaryFunctions and
+  // setupTablePatching.
+  for (auto& global : primary.globals) {
+    if (global->init) {
+      for (auto* ref : FindAll<RefFunc>(global->init).list) {
+        if (allSecondaryFuncs.count(ref->func)) {
+          getTrampoline(ref->func);
+        }
+      }
+    }
+  }
+
   // Turn references to secondary functions into references to thunks that
   // perform a direct call to the original referent. The direct calls in the
   // thunks will be handled like all other cross-module calls later, in
@@ -1222,6 +1240,55 @@ void ModuleSplitter::shareImportableItems() {
   }
 }
 
+void ModuleSplitter::cleanupUnusedTrampolines() {
+  // We pre-created trampolines for all functions referenced in global
+  // initializers' ref.funcs in indirectReferencesToSecondaryFunctions.
+  // But after splitting module items, if a trampoline is not used in the
+  // primary module, it means it was created for a global initializer but that
+  // global ended up moving to the same secondary module as the function
+  // referenced in the global initializer's ref.func. We can remove them here.
+  // TODO Remove placeholders and table elements created by unused trampolines
+  std::unordered_set<Name> usedFuncs;
+
+  struct FuncrefCollector : public PostWalker<FuncrefCollector> {
+    std::vector<Name>& used;
+    FuncrefCollector(std::vector<Name>& used) : used(used) {}
+    void visitRefFunc(RefFunc* curr) { used.push_back(curr->func); }
+  };
+
+  ModuleUtils::ParallelFunctionAnalysis<std::vector<Name>> analysis(
+    primary, [&](Function* func, std::vector<Name>& used) {
+      if (!func->imported()) {
+        FuncrefCollector(used).walkFunction(func);
+      }
+    });
+
+  for (auto& [_, usedList] : analysis.map) {
+    usedFuncs.insert(usedList.begin(), usedList.end());
+  }
+
+  std::vector<Name> moduleCodeUsed;
+  FuncrefCollector(moduleCodeUsed).walkModuleCode(&primary);
+  usedFuncs.insert(moduleCodeUsed.begin(), moduleCodeUsed.end());
+
+  for (auto& ex : primary.exports) {
+    if (ex->kind == ExternalKind::Function) {
+      usedFuncs.insert(*ex->getInternalName());
+    }
+  }
+
+  std::vector<Name> trampolinesToRemove;
+  for (auto& [_, trampoline] : trampolineMap) {
+    if (!usedFuncs.count(trampoline)) {
+      trampolinesToRemove.push_back(trampoline);
+    }
+  }
+  for (auto& name : trampolinesToRemove) {
+    primary.removeFunction(name);
+    primaryFuncs.erase(name);
+  }
+}
+
 } // anonymous namespace
 
 Results splitFunctions(Module& primary, const Config& config) {

test/lit/wasm-split/global-reffunc.wast

@@ -0,0 +1,51 @@
+;; RUN: wasm-split %s -all -g -o1 %t.1.wasm -o2 %t.2.wasm --keep-funcs=keep
+;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
+;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
+
+;; When a split global ($a here)'s initializer contains a ref.func of a split
+;; function, we should NOT create any trampolines, and the split global should
+;; direclty refer to the function.
+
+(module
+  (global $a funcref (ref.func $split))
+  (global $b funcref (ref.func $keep))
+
+  (func $keep)
+
+  (func $split
+    (drop
+      (global.get $a)
+    )
+    (drop
+      (global.get $b)
+    )
+  )
+)
+
+;; PRIMARY:      (module
+;; PRIMARY-NEXT:  (type $0 (func))
+;; PRIMARY-NEXT:  (import "placeholder.deferred" "0" (func $placeholder_0))
+;; PRIMARY-NEXT:  (table $0 1 funcref)
+;; PRIMARY-NEXT:  (elem $0 (i32.const 0) $placeholder_0)
+;; PRIMARY-NEXT:  (export "table" (table $0))
+;; PRIMARY-NEXT:  (export "keep" (func $keep))
+;; PRIMARY-NEXT:  (func $keep
+;; PRIMARY-NEXT:  )
+;; PRIMARY-NEXT: )
+
+;; SECONDARY:      (module
+;; SECONDARY-NEXT:  (type $0 (func))
+;; SECONDARY-NEXT:  (import "primary" "table" (table $timport$0 1 funcref))
+;; SECONDARY-NEXT:  (import "primary" "keep" (func $keep (exact)))
+;; SECONDARY-NEXT:  (global $a funcref (ref.func $split))
+;; SECONDARY-NEXT:  (global $b funcref (ref.func $keep))
+;; SECONDARY-NEXT:  (elem $0 (i32.const 0) $split)
+;; SECONDARY-NEXT:  (func $split
+;; SECONDARY-NEXT:   (drop
+;; SECONDARY-NEXT:    (global.get $a)
+;; SECONDARY-NEXT:   )
+;; SECONDARY-NEXT:   (drop
+;; SECONDARY-NEXT:    (global.get $b)
+;; SECONDARY-NEXT:   )
+;; SECONDARY-NEXT:  )
+;; SECONDARY-NEXT: )

test/lit/wasm-split/global-reffunc2.wast

@@ -0,0 +1,53 @@
+;; RUN: wasm-split %s -all -g -o1 %t.1.wasm -o2 %t.2.wasm --split-funcs=split1,split2
+;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
+;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
+
+;; Global $g1 is used (exported) in the primary module so it can't move, and
+;; global $g2 is only used in the secondary module so it will move there.
+
+(module
+  (global $g1 funcref (ref.func $split1))
+  (global $g2 funcref (ref.func $split2))
+  (export "g1" (global $g1))
+
+  (func $split1
+    (unreachable)
+  )
+
+  (func $split2
+    (drop
+      (global.get $g2)
+    )
+  )
+)
+
+;; PRIMARY-NEXT: (module
+;; PRIMARY-NEXT:  (type $0 (func))
+;; PRIMARY-NEXT:  (import "placeholder.deferred" "0" (func $placeholder_0))
+;; PRIMARY-NEXT:  (import "placeholder.deferred" "1" (func $placeholder_1))
+;; PRIMARY-NEXT:  (global $g1 funcref (ref.func $trampoline_split1))
+;; PRIMARY-NEXT:  (table $0 2 funcref)
+;; PRIMARY-NEXT:  (elem $0 (i32.const 0) $placeholder_0 $placeholder_1)
+;; PRIMARY-NEXT:  (export "g1" (global $g1))
+;; PRIMARY-NEXT:  (export "table" (table $0))
+;; PRIMARY-NEXT:  (func $trampoline_split1
+;; PRIMARY-NEXT:   (call_indirect (type $0)
+;; PRIMARY-NEXT:    (i32.const 0)
+;; PRIMARY-NEXT:   )
+;; PRIMARY-NEXT:  )
+;; PRIMARY-NEXT: )
+
+;; SECONDARY-NEXT: (module
+;; SECONDARY-NEXT:  (type $0 (func))
+;; SECONDARY-NEXT:  (import "primary" "table" (table $timport$0 2 funcref))
+;; SECONDARY-NEXT:  (global $g2 funcref (ref.func $split2))
+;; SECONDARY-NEXT:  (elem $0 (i32.const 0) $split1 $split2)
+;; SECONDARY-NEXT:  (func $split1
+;; SECONDARY-NEXT:   (unreachable)
+;; SECONDARY-NEXT:  )
+;; SECONDARY-NEXT:  (func $split2
+;; SECONDARY-NEXT:   (drop
+;; SECONDARY-NEXT:    (global.get $g2)
+;; SECONDARY-NEXT:   )
+;; SECONDARY-NEXT:  )
+;; SECONDARY-NEXT: )

test/lit/wasm-split/ref.func.wast

@@ -73,7 +73,7 @@
 
  ;; SECONDARY:      (import "primary" "prime" (func $prime (exact (type $0))))
 
- ;; SECONDARY:      (elem $0 (i32.const 0) $second-in-table $second)
+ ;; SECONDARY:      (elem $0 (i32.const 0) $second $second-in-table)
 
  ;; SECONDARY:      (elem declare func $prime)
 
@@ -109,13 +109,13 @@
   ;; (but we will get a placeholder, as all split-out functions do).
  )
 )
-;; PRIMARY:      (func $trampoline_second-in-table (type $0)
+;; PRIMARY:      (func $trampoline_second (type $0)
 ;; PRIMARY-NEXT:  (call_indirect $1 (type $0)
 ;; PRIMARY-NEXT:   (i32.const 0)
 ;; PRIMARY-NEXT:  )
 ;; PRIMARY-NEXT: )
 
-;; PRIMARY:      (func $trampoline_second (type $0)
+;; PRIMARY:      (func $trampoline_second-in-table (type $0)
 ;; PRIMARY-NEXT:  (call_indirect $1 (type $0)
 ;; PRIMARY-NEXT:   (i32.const 1)
 ;; PRIMARY-NEXT:  )