[wasm-split] Remove unnecessary trampolines for ref.func initializers

src/ir/module-splitting.cpp

@@ -45,7 +45,8 @@
 //      instantiation.
 //
 //   8. Export globals, tags, tables, and memories from the primary module and
-//      import them in the secondary modules.
+//      import them in the secondary modules. If possible, move those module
+//      items instead to the secondary modules.
 //
 // Functions can be used or referenced three ways in a WebAssembly module: they
 // can be exported, called, or referenced with ref.func. The above procedure
@@ -583,6 +584,25 @@ Expression* ModuleSplitter::maybeLoadSecondary(Builder& builder,
   return builder.makeSequence(loadSecondary, callIndirect);
 }
 
+// Helper to walk expressions in segments but NOT in globals.
+template<typename Walker>
+static void walkSegments(Walker& walker, Module* module) {
+  walker.setModule(module);
+  for (auto& curr : module->elementSegments) {
+    if (curr->offset) {
+      walker.walk(curr->offset);
+    }
+    for (auto* item : curr->data) {
+      walker.walk(item);
+    }
+  }
+  for (auto& curr : module->dataSegments) {
+    if (curr->offset) {
+      walker.walk(curr->offset);
+    }
+  }
+}
+
 void ModuleSplitter::indirectReferencesToSecondaryFunctions() {
   // Turn references to secondary functions into references to thunks that
   // perform a direct call to the original referent. The direct calls in the
@@ -611,7 +631,25 @@ void ModuleSplitter::indirectReferencesToSecondaryFunctions() {
       }
     }
   } gatherer(*this);
-  gatherer.walkModule(&primary);
+  // We shouldn't use collector.walkModuleCode here, because we don't want to
+  // walk on global initializers. At this point, all globals are still in the
+  // primary module, so if we walk on 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());
   }
@@ -977,7 +1015,19 @@ void ModuleSplitter::shareImportableItems() {
     }
 
     NameCollector collector(used);
-    collector.walkModuleCode(&module);
+    // We shouldn't use collector.walkModuleCode here, because we don't want to
+    // walk on global initializers. At this point, all globals are still in the
+    // primary module, so if we walk on global initializers here, globals appear
+    // in their initialalizers will be all marked as used in the primary module,
+    // which is not we want.
+    //
+    // For example, we have (global $a i32 (global.get $b)). Because $a is at
+    // this point still in the primary module, $b will be marked as "used" in
+    // the primary module. But $a can be moved to a secondary module later if it
+    // is used exclusively by that module. Then $b can be also moved, in case it
+    // doesn't have other uses. But if it is marked as "used" in the primary
+    // module, it can't.
+    walkSegments(collector, &module);
     for (auto& segment : module.dataSegments) {
       if (segment->memory.is()) {
         used.memories.insert(segment->memory);
@@ -1009,7 +1059,6 @@ void ModuleSplitter::shareImportableItems() {
           break;
       }
     }
-
     return used;
   };
 
@@ -1019,25 +1068,33 @@ void ModuleSplitter::shareImportableItems() {
     secondaryUsed.push_back(getUsedNames(*secondaryPtr));
   }
 
-  // Compute globals referenced in other globals' initializers. Since globals
-  // can reference other globals, we must ensure that if a global is used in a
-  // module, all its dependencies are also marked as used.
-  auto computeDependentItems = [&](UsedNames& used) {
+  // Compute transitive closure of globals referenced in other globals'
+  // initializers. Since globals can reference other globals, we must ensure
+  // 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());
-    for (auto name : worklist) {
+    std::unordered_set<Name> visited(used.globals.begin(), used.globals.end());
+    while (!worklist.empty()) {
+      Name name = worklist.back();
+      worklist.pop_back();
       // 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) {
-          used.globals.insert(get->name);
+          if (visited.insert(get->name).second) {
+            worklist.push_back(get->name);
+            used.globals.insert(get->name);
+          }
         }
       }
     }
   };
 
+  computeTransitiveGlobals(primaryUsed);
   for (auto& used : secondaryUsed) {
-    computeDependentItems(used);
+    computeTransitiveGlobals(used);
   }
 
   // Given a name and module item kind, returns the list of secondary modules
@@ -1127,20 +1184,47 @@ void ModuleSplitter::shareImportableItems() {
       getUsingSecondaries(global->name, &UsedNames::globals);
     bool usedInPrimary = primaryUsed.globals.count(global->name);
     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 (allSecondaryFuncs.count(ref->func)) {
+            Index targetIndex = funcToSecondaryIndex.at(ref->func);
+            if (secondaries[targetIndex].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 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 {
+
       for (auto* secondary : usingSecondaries) {
         auto* secondaryGlobal =
           ModuleUtils::copyGlobal(global.get(), *secondary);

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

@@ -8,21 +8,11 @@
 ;; TODO Use $split in the secondary module directly in the split global
 
 (module
-  ;; PRIMARY:      (export "trampoline_split" (func $trampoline_split))
-
   ;; PRIMARY:      (func $keep
   ;; PRIMARY-NEXT: )
   (func $keep)
 
-  ;; PRIMARY:      (func $trampoline_split
-  ;; PRIMARY-NEXT:  (call_indirect (type $0)
-  ;; PRIMARY-NEXT:   (i32.const 0)
-  ;; PRIMARY-NEXT:  )
-  ;; PRIMARY-NEXT: )
-
-
-  ;; SECONDARY:      (import "primary" "trampoline_split" (func $trampoline_split (exact)))
-  ;; SECONDARY:      (global $a funcref (ref.func $trampoline_split))
+  ;; SECONDARY:      (global $a funcref (ref.func $split))
   (global $a funcref (ref.func $split))
 
   ;; SECONDARY:      (func $split

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:  )

test/lit/wasm-split/transitive-globals.wast

@@ -3,26 +3,20 @@
 ;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
 
 ;; Check that transitive dependencies in global initializers are correctly
-;; analyzed and exported from the primary module to the secondary module.
-;; TODO Move $b and $c to the secondary module
+;; analyzed and moved to the secondary module.
 
 (module
-  ;; PRIMARY:      (global $c i32 (i32.const 42))
+  ;; SECONDARY:      (global $c i32 (i32.const 42))
   (global $c i32 (i32.const 42))
 
   ;; $b depends on $c.
-  ;; PRIMARY:      (global $b i32 (global.get $c))
+  ;; SECONDARY:      (global $b i32 (global.get $c))
   (global $b i32 (global.get $c))
 
-  ;; Globals $b is exported to the secondary module
-  ;; PRIMARY:      (export "global" (global $b))
-
-  ;; Globals $b is imported from the primary module
-  ;; SECONDARY:      (import "primary" "global" (global $b i32))
-
-  ;; $a depends on $b. Since $a is exclusively used by the secondary module,
-  ;; it will be moved there. Its dependency $b should be exported from the
-  ;; primary module and imported into the secondary module.
+  ;; $a depends on $b. since $a is exclusively used by the secondary module,
+  ;; it will be moved there. The transitive dependency must ensure that $b (and
+  ;; $c) are moved to the secondary module too, because they are not used in the
+  ;; primary module.
   ;; SECONDARY:      (global $a i32 (global.get $b))
   (global $a i32 (global.get $b))