[wasm-split] Move exclusive items to secondary modules (#8441)

When splitting a module, if non-function items (memories, tables,
globals, tags) are exclusively used by a single secondary module, this
moves them directly to that secondary module rather than exporting them
from the primary module.

When a global is moved, its initializer can contain `global.get` or
`ref.func`s, creating dependences on other globals and functions. For
now, this PR just exports all the dependences from the primary module to
the secondary module. This will be improved by follow-up PRs.

This PR reduces the size of the primary module for acx_gallery by 12.5%.
Follow-up PRs will reduce it further.

This also sadly increases wasm-split's running time on acx_gallery from
16.5s -> 24.7s, by 49%, due to more computations in
`shareImportableItems`.

---

The below is `wasm-objdump -h` result of the primary modules:

- Before
```
     Type start=0x0000000c end=0x00035e09 (size=0x00035dfd) count: 11192
   Import start=0x00035e0e end=0x004bd669 (size=0x0048785b) count: 65720
 Function start=0x004bd66d end=0x004d0519 (size=0x00012eac) count: 62890
    Table start=0x004d051c end=0x004d4059 (size=0x00003b3d) count: 2921
      Tag start=0x004d405b end=0x004d405f (size=0x00000004) count: 1
   Global start=0x004d4063 end=0x00689ff8 (size=0x001b5f95) count: 80766
   Export start=0x00689ffc end=0x0071b16c (size=0x00091170) count: 60877
    Start start=0x0071b16e end=0x0071b170 (size=0x00000002) start: 828
     Elem start=0x0071b174 end=0x00784fb9 (size=0x00069e45) count: 12303
DataCount start=0x00784fbb end=0x00784fbc (size=0x00000001) count: 1
     Code start=0x00784fc1 end=0x009b4958 (size=0x0022f997) count: 62890
     Data start=0x009b495c end=0x009d44e9 (size=0x0001fb8d) count: 1
```

- After (This PR)
```
     Type start=0x0000000c end=0x00035d44 (size=0x00035d38) count: 11185
   Import start=0x00035d49 end=0x003faf6f (size=0x003c5226) count: 56805
 Function start=0x003faf73 end=0x0040de1f (size=0x00012eac) count: 62890
    Table start=0x0040de22 end=0x0041195d (size=0x00003b3b) count: 2921
      Tag start=0x0041195f end=0x00411963 (size=0x00000004) count: 1
   Global start=0x00411967 end=0x005541c5 (size=0x0014285e) count: 47771
   Export start=0x005541c9 end=0x005dfc2c (size=0x0008ba63) count: 59077
    Start start=0x005dfc2e end=0x005dfc30 (size=0x00000002) start: 828
     Elem start=0x005dfc34 end=0x00649a77 (size=0x00069e43) count: 12303
DataCount start=0x00649a79 end=0x00649a7a (size=0x00000001) count: 1
     Code start=0x00649a7f end=0x00879385 (size=0x0022f906) count: 62890
     Data start=0x00879389 end=0x00898f16 (size=0x0001fb8d) count: 1
```

Follow-ups: #8442 and #8443

Author: aheejin

src/ir/module-splitting.cpp

@@ -47,9 +47,6 @@
 //   8. Export globals, tags, tables, and memories from the primary module and
 //      import them in the secondary modules.
 //
-//   9. Run RemoveUnusedModuleElements pass on the secondary modules in order to
-//      remove unused imports.
-//
 // 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
 // introduces a layer of indirection to each of those mechanisms that removes
@@ -77,10 +74,9 @@
 #include "ir/module-splitting.h"
 #include "asmjs/shared-constants.h"
 #include "ir/export-utils.h"
+#include "ir/find_all.h"
 #include "ir/module-utils.h"
 #include "ir/names.h"
-#include "ir/utils.h"
-#include "pass.h"
 #include "support/insert_ordered.h"
 #include "wasm-builder.h"
 #include "wasm.h"
@@ -963,13 +959,11 @@ void ModuleSplitter::shareImportableItems() {
     }
   };
 
-  for (auto& secondaryPtr : secondaries) {
-    Module& secondary = *secondaryPtr;
-
-    // Collect names used in the secondary module
+  // Given a module, collect names used in the module
+  auto getUsedNames = [&](Module& module) {
     UsedNames used;
     ModuleUtils::ParallelFunctionAnalysis<UsedNames> nameCollector(
-      secondary, [&](Function* func, UsedNames& used) {
+      module, [&](Function* func, UsedNames& used) {
         if (!func->imported()) {
           NameCollector(used).walk(func->body);
         }
@@ -983,65 +977,201 @@ void ModuleSplitter::shareImportableItems() {
     }
 
     NameCollector collector(used);
-    collector.walkModuleCode(&secondary);
-    for (auto& segment : secondary.dataSegments) {
+    collector.walkModuleCode(&module);
+    for (auto& segment : module.dataSegments) {
       if (segment->memory.is()) {
         used.memories.insert(segment->memory);
       }
     }
-    for (auto& segment : secondary.elementSegments) {
+    for (auto& segment : module.elementSegments) {
       if (segment->table.is()) {
         used.tables.insert(segment->table);
       }
     }
 
-    // Export module items that are used in the secondary module
-    for (auto& memory : primary.memories) {
-      if (!used.memories.count(memory->name)) {
-        continue;
+    // If primary module has exports, they are "used" in it. Secondary modules
+    // don't have exports, so this only applies to the primary module.
+    for (auto& ex : module.exports) {
+      switch (ex->kind) {
+        case ExternalKind::Global:
+          used.globals.insert(*ex->getInternalName());
+          break;
+        case ExternalKind::Memory:
+          used.memories.insert(*ex->getInternalName());
+          break;
+        case ExternalKind::Table:
+          used.tables.insert(*ex->getInternalName());
+          break;
+        case ExternalKind::Tag:
+          used.tags.insert(*ex->getInternalName());
+          break;
+        default:
+          break;
+      }
+    }
+
+    return used;
+  };
+
+  UsedNames primaryUsed = getUsedNames(primary);
+  std::vector<UsedNames> secondaryUsed;
+  for (auto& secondaryPtr : secondaries) {
+    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) {
+    std::vector<Name> worklist(used.globals.begin(), used.globals.end());
+    for (auto name : worklist) {
+      // 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);
+        }
+      }
+    }
+  };
+
+  for (auto& used : secondaryUsed) {
+    computeDependentItems(used);
+  }
+
+  // Given a name and module item kind, returns the list of secondary modules
+  // using that name
+  auto getUsingSecondaries = [&](const Name& name, auto UsedNames::* field) {
+    std::vector<Module*> usingModules;
+    for (size_t i = 0; i < secondaries.size(); ++i) {
+      if ((secondaryUsed[i].*field).count(name)) {
+        usingModules.push_back(secondaries[i].get());
       }
-      auto secondaryMemory = ModuleUtils::copyMemory(memory.get(), secondary);
-      makeImportExport(
-        *memory, *secondaryMemory, "memory", ExternalKind::Memory);
     }
+    return usingModules;
+  };
 
-    for (auto& table : primary.tables) {
-      // 1. In case we copied this table to this secondary module in
-      //    setupTablePatching(), secondary.getTableOrNull(table->name) is not
-      //    null, and we need to export it.
-      // 2. As in the case with other module elements, if the table is used in
-      //    the secondary module's instructions, we need to export it.
-      auto secondaryTable = secondary.getTableOrNull(table->name);
-      if (!secondaryTable && !used.tables.count(table->name)) {
-        continue;
+  // Share module items with secondary modules.
+  // 1. Only share an item with the modules that use it
+  // 2. If an item is used by only a single secondary module, move the item to
+  //    that secondary module. If an item is used by multiple modules (including
+  //    the primary and secondary modules), export the item from the primary and
+  //    import it from the using secondary modules.
+
+  std::vector<Name> memoriesToRemove;
+  for (auto& memory : primary.memories) {
+    auto usingSecondaries =
+      getUsingSecondaries(memory->name, &UsedNames::memories);
+    bool usedInPrimary = primaryUsed.memories.count(memory->name);
+
+    if (!usedInPrimary && usingSecondaries.size() == 1) {
+      auto* secondary = usingSecondaries[0];
+      ModuleUtils::copyMemory(memory.get(), *secondary);
+      memoriesToRemove.push_back(memory->name);
+    } else {
+      for (auto* secondary : usingSecondaries) {
+        auto* secondaryMemory =
+          ModuleUtils::copyMemory(memory.get(), *secondary);
+        makeImportExport(
+          *memory, *secondaryMemory, "memory", ExternalKind::Memory);
       }
-      if (!secondaryTable) {
-        secondaryTable = ModuleUtils::copyTable(table.get(), secondary);
+    }
+  }
+  for (auto& name : memoriesToRemove) {
+    primary.removeMemory(name);
+  }
+
+  std::vector<Name> tablesToRemove;
+  for (auto& table : primary.tables) {
+    auto usingSecondaries =
+      getUsingSecondaries(table->name, &UsedNames::tables);
+    bool usedInPrimary = primaryUsed.tables.count(table->name);
+
+    if (!usedInPrimary && usingSecondaries.size() == 1) {
+      auto* secondary = usingSecondaries[0];
+      //  In case we copied this table to this secondary module in
+      //  setupTablePatching(), !usedInPrimary can't be satisfied, because the
+      //  primary module should have an element segment that refers to this
+      //  table.
+      assert(!secondary->getTableOrNull(table->name));
+      ModuleUtils::copyTable(table.get(), *secondary);
+      tablesToRemove.push_back(table->name);
+    } else {
+      for (auto* secondary : usingSecondaries) {
+        // 1. In case we copied this table to this secondary module in
+        //    setupTablePatching(), secondary.getTableOrNull(table->name) is not
+        //    null, and we need to import it.
+        // 2. As in the case with other module elements, if the table is used in
+        //    the secondary module's instructions, we need to export it.
+        auto secondaryTable = secondary->getTableOrNull(table->name);
+        if (!secondaryTable) {
+          secondaryTable = ModuleUtils::copyTable(table.get(), *secondary);
+        }
+        makeImportExport(*table, *secondaryTable, "table", ExternalKind::Table);
       }
-      makeImportExport(*table, *secondaryTable, "table", ExternalKind::Table);
     }
+  }
+  for (auto& name : tablesToRemove) {
+    primary.removeTable(name);
+  }
 
-    for (auto& global : primary.globals) {
-      if (!used.globals.count(global->name)) {
-        continue;
+  std::vector<Name> globalsToRemove;
+  for (auto& global : primary.globals) {
+    if (global->mutable_) {
+      assert(primary.features.hasMutableGlobals() &&
+             "TODO: add wrapper functions for disallowed mutable globals");
+    }
+
+    auto usingSecondaries =
+      getUsingSecondaries(global->name, &UsedNames::globals);
+    bool usedInPrimary = primaryUsed.globals.count(global->name);
+    if (!usedInPrimary && usingSecondaries.size() == 1) {
+      auto* secondary = usingSecondaries[0];
+      ModuleUtils::copyGlobal(global.get(), *secondary);
+      globalsToRemove.push_back(global->name);
+      // Import global initializer's ref.func dependences
+      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 (global->mutable_) {
-        assert(primary.features.hasMutableGlobals() &&
-               "TODO: add wrapper functions for disallowed mutable globals");
+    } else {
+      for (auto* secondary : usingSecondaries) {
+        auto* secondaryGlobal =
+          ModuleUtils::copyGlobal(global.get(), *secondary);
+        makeImportExport(
+          *global, *secondaryGlobal, "global", ExternalKind::Global);
       }
-      auto* secondaryGlobal = ModuleUtils::copyGlobal(global.get(), secondary);
-      makeImportExport(
-        *global, *secondaryGlobal, "global", ExternalKind::Global);
     }
+  }
+  for (auto& name : globalsToRemove) {
+    primary.removeGlobal(name);
+  }
+
+  std::vector<Name> tagsToRemove;
+  for (auto& tag : primary.tags) {
+    auto usingSecondaries = getUsingSecondaries(tag->name, &UsedNames::tags);
+    bool usedInPrimary = primaryUsed.tags.count(tag->name);
 
-    for (auto& tag : primary.tags) {
-      if (!used.tags.count(tag->name)) {
-        continue;
+    if (!usedInPrimary && usingSecondaries.size() == 1) {
+      auto* secondary = usingSecondaries[0];
+      ModuleUtils::copyTag(tag.get(), *secondary);
+      tagsToRemove.push_back(tag->name);
+    } else {
+      for (auto* secondary : usingSecondaries) {
+        auto* secondaryTag = ModuleUtils::copyTag(tag.get(), *secondary);
+        makeImportExport(*tag, *secondaryTag, "tag", ExternalKind::Tag);
       }
-      auto* secondaryTag = ModuleUtils::copyTag(tag.get(), secondary);
-      makeImportExport(*tag, *secondaryTag, "tag", ExternalKind::Tag);
     }
   }
+  for (auto& name : tagsToRemove) {
+    primary.removeTag(name);
+  }
 }
 
 } // anonymous namespace

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

@@ -0,0 +1,38 @@
+;; 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, currently we create its trampoline in the primary module and export
+;; it.
+;; 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))
+  (global $a funcref (ref.func $split))
+
+  ;; SECONDARY:      (func $split
+  ;; SECONDARY-NEXT:  (drop
+  ;; SECONDARY-NEXT:   (global.get $a)
+  ;; SECONDARY-NEXT:  )
+  ;; SECONDARY-NEXT: )
+  (func $split
+    (drop
+      (global.get $a)
+    )
+  )
+)