diff --git a/src/ir/module-splitting.cpp b/src/ir/module-splitting.cpp index 3d19d8b02c6..89804678cfc 100644 --- a/src/ir/module-splitting.cpp +++ b/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 +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 worklist(used.globals.begin(), used.globals.end()); - for (auto name : worklist) { + std::unordered_set 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(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(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); diff --git a/test/lit/wasm-split/transitive-globals.wast b/test/lit/wasm-split/transitive-globals.wast index 90740adc3a3..6bf5ce6636f 100644 --- a/test/lit/wasm-split/transitive-globals.wast +++ b/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) + ) ) )