[wasm-split] Remove unnecessary trampolines for ref.func initializers (WebAssembly#8443)

When a global is exclusively used by a secondary module and thus moved
to that module, and its initializer has a `(ref.func $func)`, we used to
create a trampoline and export it from the primary module in all cases,
even in the case that the function is in the same secondary module. This
now only avoids creating a trampoline when the function is already in
the same secondary module.

To do this, we now skip scanning global initializers in
`indirectReferencesToSecondaryFunctions`, and selectively create
trampolines only when needed in `shareImportableItems`.

The running time of `wasm-split` hasn't really changed with this PR,
compared to the previous PR WebAssembly#8442 (~25s range in acx_gallery).

WebAssembly#8441, WebAssembly#8442, and this PR combined reduce the size of the primary module
of acx_gallery by 45.4%.

---

`wasm-objdump -h` result:

- Before (WebAssembly#8442)
```
     Type start=0x0000000c end=0x00035d44 (size=0x00035d38) count: 11185
   Import start=0x00035d48 end=0x00132efc (size=0x000fd1b4) count: 32642
 Function start=0x00132f00 end=0x00145dac (size=0x00012eac) count: 62890
    Table start=0x00145daf end=0x001498ea (size=0x00003b3b) count: 2921
      Tag start=0x001498ec end=0x001498f0 (size=0x00000004) count: 1
   Global start=0x001498f4 end=0x00289e60 (size=0x0014056c) count: 47728
   Export start=0x00289e64 end=0x002e99c1 (size=0x0005fb5d) count: 35861
    Start start=0x002e99c3 end=0x002e99c5 (size=0x00000002) start: 828
     Elem start=0x002e99c9 end=0x0035380c (size=0x00069e43) count: 12303
DataCount start=0x0035380e end=0x0035380f (size=0x00000001) count: 1
     Code start=0x00353814 end=0x005830e5 (size=0x0022f8d1) count: 62890
     Data start=0x005830e9 end=0x005a2c76 (size=0x0001fb8d) count: 1
```

- After (This PR)
```
     Type start=0x0000000c end=0x00035d38 (size=0x00035d2c) count: 11185
   Import start=0x00035d3c end=0x00132ef0 (size=0x000fd1b4) count: 32642
 Function start=0x00132ef4 end=0x001436cc (size=0x000107d8) count: 53001
    Table start=0x001436cf end=0x0014720a (size=0x00003b3b) count: 2921
      Tag start=0x0014720c end=0x00147210 (size=0x00000004) count: 1
   Global start=0x00147214 end=0x00287b75 (size=0x00140961) count: 47728
   Export start=0x00287b79 end=0x002d41ce (size=0x0004c655) count: 25972
    Start start=0x002d41d0 end=0x002d41d2 (size=0x00000002) start: 828
     Elem start=0x002d41d6 end=0x00336c36 (size=0x00062a60) count: 12303
DataCount start=0x00336c38 end=0x00336c39 (size=0x00000001) count: 1
     Code start=0x00336c3e end=0x0053dbdd (size=0x00206f9f) count: 53001
     Data start=0x0053dbe1 end=0x0055d76e (size=0x0001fb8d) count: 1
```

Author: aheejin

src/ir/module-splitting.cpp

@@ -79,6 +79,7 @@
 #include "ir/module-utils.h"
 #include "ir/names.h"
 #include "support/insert_ordered.h"
+#include "support/unique_deferring_queue.h"
 #include "wasm-builder.h"
 #include "wasm.h"
 
@@ -631,7 +632,25 @@ void ModuleSplitter::indirectReferencesToSecondaryFunctions() {
       }
     }
   } gatherer(*this);
-  gatherer.walkModule(&primary);
+  // We shouldn't use collector.walkModuleCode here, because we don't want to
+  // walk global initializers. At this point, all globals are still in the
+  // primary module, so if we walk global initializers here, it will create
+  // unnecessary trampolines.
+  //
+  // For example, we have (global $a funcref (ref.func $foo)), and $foo was
+  // split into a secondary module. Because $a is at this point still in the
+  // primary module, $foo will be considered to exist in a different module, so
+  // this will create a trampoline for $foo. But it is possible that later we
+  // find out $a is exclusively used by that secondary module and move $a there.
+  // In that case, $a can just reference $foo locally, but if we scan global
+  // initializers here, we would have created an unnecessary trampoline for
+  // $foo.
+  walkSegments(gatherer, &primary);
+  for (auto& curr : primary.functions) {
+    if (!curr->imported()) {
+      gatherer.walkFunction(curr.get());
+    }
+  }
   for (auto& secondaryPtr : secondaries) {
     gatherer.walkModule(secondaryPtr.get());
   }
@@ -1055,20 +1074,19 @@ void ModuleSplitter::shareImportableItems() {
   // that if a global is used in a module, all its dependencies are also marked
   // as used.
   auto computeTransitiveGlobals = [&](UsedNames& used) {
-    std::vector<Name> worklist(used.globals.begin(), used.globals.end());
-    std::unordered_set<Name> visited(used.globals.begin(), used.globals.end());
+    UniqueNonrepeatingDeferredQueue<Name> worklist;
+    for (auto global : used.globals) {
+      worklist.push(global);
+    }
     while (!worklist.empty()) {
-      Name name = worklist.back();
-      worklist.pop_back();
+      Name name = worklist.pop();
       // At this point all globals are still in the primary module, so this
       // exists
       auto* global = primary.getGlobal(name);
       if (!global->imported() && global->init) {
         for (auto* get : FindAll<GlobalGet>(global->init).list) {
-          if (visited.insert(get->name).second) {
-            worklist.push_back(get->name);
-            used.globals.insert(get->name);
-          }
+          worklist.push(get->name);
+          used.globals.insert(get->name);
         }
       }
     }
@@ -1168,19 +1186,45 @@ void ModuleSplitter::shareImportableItems() {
     if (!usedInPrimary && usingSecondaries.size() == 1) {
       // We are moving this global to this secondary module
       auto* secondary = usingSecondaries[0];
-      ModuleUtils::copyGlobal(global.get(), *secondary);
+      auto* secondaryGlobal = ModuleUtils::copyGlobal(global.get(), *secondary);
       globalsToRemove.push_back(global->name);
-      // Import global initializer's ref.func dependences
+
+      if (secondaryGlobal->init) {
+        // When a global's initializer contains ref.func
+        for (auto* ref : FindAll<RefFunc>(secondaryGlobal->init).list) {
+          // If ref.func's function is in a different secondary module, we
+          // create a trampoline here.
+          if (auto targetIndexIt = funcToSecondaryIndex.find(ref->func);
+              targetIndexIt != funcToSecondaryIndex.end()) {
+            if (secondaries[targetIndexIt->second].get() != secondary) {
+              ref->func = getTrampoline(ref->func);
+            }
+          }
+          // 1. If ref.func's function is in the primary module, we export it
+          //    here.
+          // 2. If ref.func's function is in a different secondary module and we
+          //    just created a trampoline for it in the primary module above, we
+          //    export the trampoline here.
+          if (primary.getFunctionOrNull(ref->func)) {
+            exportImportFunction(ref->func, {secondary});
+          }
+          // If ref.func's function is in the same secondary module, we don't
+          // need to do anything. The ref.func can directly reference the
+          // function.
+        }
+      }
+    } else { // We are NOT moving this global to the secondary module
       if (global->init) {
         for (auto* ref : FindAll<RefFunc>(global->init).list) {
-          // Here, ref->func is either a function in the primary module, or a
-          // trampoline created in indirectReferencesToSecondaryFunctions in
-          // case the original function is in one of the secondaries.
-          assert(primary.getFunctionOrNull(ref->func));
-          exportImportFunction(ref->func, {secondary});
+          // If we are exporting this global from the primary module, we should
+          // create a trampoline here, because we skipped doing it for global
+          // initializers in indirectReferencesToSecondaryFunctions.
+          if (allSecondaryFuncs.count(ref->func)) {
+            ref->func = getTrampoline(ref->func);
+          }
         }
       }
-    } else { // We are NOT moving this global to the secondary module
+
       for (auto* secondary : usingSecondaries) {
         auto* secondaryGlobal =
           ModuleUtils::copyGlobal(global.get(), *secondary);

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

@@ -3,36 +3,41 @@
 ;; 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, currently we create its trampoline in the primary module and export
-;; it.
-;; TODO Use $split in the secondary module directly in the split global
+;; function, we should NOT create any trampolines, and the split global should
+;; direclty refer to the function.
 
 (module
-  ;; PRIMARY:      (export "trampoline_split" (func $trampoline_split))
+  (global $a funcref (ref.func $split))
+  (global $b funcref (ref.func $keep))
 
-  ;; PRIMARY:      (func $keep
-  ;; PRIMARY-NEXT: )
-  (func $keep)
+  ;; PRIMARY:      (export "keep" (func $keep))
 
-  ;; PRIMARY:      (func $trampoline_split
-  ;; PRIMARY-NEXT:  (call_indirect (type $0)
-  ;; PRIMARY-NEXT:   (i32.const 0)
-  ;; PRIMARY-NEXT:  )
-  ;; PRIMARY-NEXT: )
+  ;; PRIMARY-NOT:  (export "trampoline_split"
+  ;; PRIMARY-NOT:  (func $trampoline_split
 
+  ;; SECONDARY:      (import "primary" "keep" (func $keep (exact)))
 
-  ;; SECONDARY:      (import "primary" "trampoline_split" (func $trampoline_split (exact)))
-  ;; SECONDARY:      (global $a funcref (ref.func $trampoline_split))
-  (global $a funcref (ref.func $split))
+  ;; SECONDARY:      (global $a funcref (ref.func $split))
+  ;; SECONDARY:      (global $b funcref (ref.func $keep))
+
+  ;; PRIMARY:      (func $keep
+  ;; PRIMARY-NEXT: )
+  (func $keep)
 
   ;; SECONDARY:      (func $split
   ;; SECONDARY-NEXT:  (drop
   ;; SECONDARY-NEXT:   (global.get $a)
   ;; SECONDARY-NEXT:  )
+  ;; SECONDARY-NEXT:  (drop
+  ;; SECONDARY-NEXT:   (global.get $b)
+  ;; SECONDARY-NEXT:  )
   ;; SECONDARY-NEXT: )
   (func $split
     (drop
       (global.get $a)
     )
+    (drop
+      (global.get $b)
+    )
   )
 )

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

@@ -61,7 +61,7 @@
 
  ;; SECONDARY:      (import "primary" "prime" (func $prime (exact (type $0))))
 
- ;; SECONDARY:      (elem $0 (i32.const 0) $second $second-in-table)
+ ;; SECONDARY:      (elem $0 (i32.const 0) $second-in-table $second)
 
  ;; SECONDARY:      (elem declare func $prime)
 
@@ -97,13 +97,13 @@
   ;; (but we will get a placeholder, as all split-out functions do).
  )
 )
-;; 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 0)
 ;; PRIMARY-NEXT:  )
 ;; PRIMARY-NEXT: )
 
-;; 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 1)
 ;; PRIMARY-NEXT:  )