[wasm-split] Split globals' transitive global dependencies (#8442)

When a global is exclusively used by a secondary module, we can move it
to the secondary module. If its initializer contains a `global.get` of
another global, we exported it from the primary module to the secondary
module, even if it may not be used anywhere else.

When we split a global out to a secondary module, this PR computes the
transitive dependency of the split global, and if those globals in the
dependency are not used anywhere else in other modules, we move them to
the secondary module as well.

#8441 and this PR combined reduce the size of the primary module of
acx_gallery by 42.6%.

The running time of `wasm-split` hasn't really changed with this PR,
compared to #8441.

---

`wasm-objdump -h` result:

- Before (#8441)
```
     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
```

- After (This PR)
```
     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
```

Note that while the decrease in the global section is small, we have a
significant size decrease in the import and the export sections, because
we used to import globals and export them just to relay those globals to
the secondary modules.

Follow-up: #8443

Fixes #7724.

Author: aheejin

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
@@ -977,7 +997,19 @@ void ModuleSplitter::shareImportableItems() {
     }
 
     NameCollector collector(used);
-    collector.walkModuleCode(&module);
+    // 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, other globals
+    // appearing in their initializers will all be marked as used in the primary
+    // module, which is not what 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 +1041,6 @@ void ModuleSplitter::shareImportableItems() {
           break;
       }
     }
-
     return used;
   };
 
@@ -1019,25 +1050,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 the 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 +1166,21 @@ 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);
       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
+          // 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});
         }
       }
-    } else {
+    } 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/transitive-globals.wast

@@ -3,41 +3,63 @@
 ;; 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))
-  (global $c i32 (i32.const 42))
+  ;; There are two dependency chains: $a->$b->$c and $d->$e->$f. While all of
+  ;; $a, $b, and $c can be moved to the secondary module because all f them are
+  ;; used only there, $e is used in the primary module, preventing $e and $f
+  ;; from being moved to the secondary module.
 
-  ;; $b depends on $c.
-  ;; PRIMARY:      (global $b i32 (global.get $c))
+  (global $c i32 (i32.const 42))
   (global $b i32 (global.get $c))
+  (global $a i32 (global.get $b))
+
+  (global $f i32 (i32.const 42))
+  (global $e i32 (global.get $f))
+  (global $d i32 (global.get $e))
 
-  ;; Globals $b is exported to the secondary module
-  ;; PRIMARY:      (export "global" (global $b))
+  ;; PRIMARY:      (global $f i32 (i32.const 42))
+  ;; PRIMARY:      (global $e i32 (global.get $f))
 
-  ;; Globals $b is imported from the primary module
-  ;; SECONDARY:      (import "primary" "global" (global $b i32))
+  ;; PRIMARY:      (export "global" (global $f))
+  ;; PRIMARY:      (export "global_1" (global $e))
 
-  ;; $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.
+  ;; SECONDARY:      (import "primary" "global" (global $f i32))
+  ;; SECONDARY:      (import "primary" "global_1" (global $e i32))
+
+  ;; SECONDARY:      (global $c i32 (i32.const 42))
+  ;; SECONDARY:      (global $b i32 (global.get $c))
   ;; SECONDARY:      (global $a i32 (global.get $b))
-  (global $a i32 (global.get $b))
 
-  ;; PRIMARY:      (func $keep (result i32)
-  ;; PRIMARY-NEXT:  (i32.const 0)
+  ;; SECONDARY:      (global $d i32 (global.get $e))
+
+  ;; PRIMARY:      (func $keep
+  ;; PRIMARY-NEXT:  (drop
+  ;; PRIMARY-NEXT:   (global.get $e)
+  ;; PRIMARY-NEXT:  )
   ;; PRIMARY-NEXT: )
-  (func $keep (result i32)
-    (i32.const 0)
+  (func $keep
+    (drop
+      (global.get $e)
+    )
   )
 
-  ;; Exclusively uses $a, causing $a to move to the secondary module
-  ;; SECONDARY:      (func $split (result i32)
-  ;; SECONDARY-NEXT:  (global.get $a)
+  ;; Exclusively uses $a and $d, causing them to move to the secondary module
+  ;; SECONDARY:      (func $split
+  ;; SECONDARY-NEXT:  (drop
+  ;; SECONDARY-NEXT:   (global.get $a)
+  ;; SECONDARY-NEXT:  )
+  ;; SECONDARY-NEXT:  (drop
+  ;; SECONDARY-NEXT:   (global.get $d)
+  ;; SECONDARY-NEXT:  )
   ;; SECONDARY-NEXT: )
-  (func $split (result i32)
-    (global.get $a)
+  (func $split
+    (drop
+      (global.get $a)
+    )
+    (drop
+      (global.get $d)
+    )
   )
 )