[wasm-split] Do multi-split at once (#7956)

This does multi-splitting of modules at once, rather than splitting them
one by one by doing 2-way split n times. Previously when we did
multi-splitting we split the 1st module as the "secondary" module
assuming all other functions belonging to 2nd-nth modules as "primary"
module. And then we repeat the same task for the 2nd module, assuming
3rd-nth module functions belong to the "primary" module. This
unnecessarily repeated some tasks that could have been done once.

This reduces the running time on a reproducer provided by @biggs0125
before (to fix #7725) from 232s to 80s on my machine, reducing it by
around 66%.

Some side-products of this PR are:
- Now we only create a single table to host placeholders (or `ref.null`s
in case of `--no-placeholders`) even when reference-types is enabled.
Previously we created a table per secondary module, resulting in n
tables.
- The names of trampoline functions have been changed in the tests, but
semantically they are the same. (e.g. in
`test/lit/wasm-split/multi-split.wast`) The reason for the change is,
previously we split modules one by one, by the time we split the first
module, it assumed functions belonging to other secondary modules were
primary functions, but they later changed to trampolines as well. Now
they are all named as trampolines, arguably enhancing readability.

Author: aheejin

src/ir/module-splitting.cpp

@@ -47,11 +47,12 @@ namespace wasm::ModuleSplitting {
 static const Name LOAD_SECONDARY_MODULE("__load_secondary_module");
 
 struct Config {
-  // The set of functions to split into the secondary module. All others are
-  // kept in the primary module. Must not include the start function if it
-  // exists. May or may not include imported functions, which are always kept in
-  // the primary module regardless.
-  std::set<Name> secondaryFuncs;
+  // A vector of set of functions to split into that secondary. Each function
+  // set belongs to a single secondary module. All others are kept in the
+  // primary module. Must not include the start function if it exists. May or
+  // may not include imported functions, which are always kept in the primary
+  // module regardless.
+  std::vector<std::set<Name>> secondaryFuncs;
   // Whether to import placeholder functions into the primary module that will
   // be called when a secondary function is called before the secondary module
   // has been loaded.
@@ -76,7 +77,7 @@ struct Config {
 };
 
 struct Results {
-  std::unique_ptr<Module> secondary;
+  std::vector<std::unique_ptr<Module>> secondaries;
   std::unordered_map<Name, std::map<size_t, Name>> placeholderMap;
 };
 

src/ir/module-splitting.h

@@ -329,7 +329,7 @@ void splitModule(const WasmSplitOptions& options) {
 
   // Actually perform the splitting
   ModuleSplitting::Config config;
-  config.secondaryFuncs = std::move(splitFuncs);
+  config.secondaryFuncs.push_back(std::move(splitFuncs));
   if (options.importNamespace.size()) {
     config.importNamespace = options.importNamespace;
   }
@@ -343,7 +343,7 @@ void splitModule(const WasmSplitOptions& options) {
   config.minimizeNewExportNames = !options.passOptions.debugInfo;
   config.jspi = options.jspi;
   auto splitResults = ModuleSplitting::splitFunctions(wasm, config);
-  auto& secondary = splitResults.secondary;
+  auto& secondary = *splitResults.secondaries.begin();
 
   adjustTableSize(wasm, options.initialTableSize);
   adjustTableSize(*secondary, options.initialTableSize, /*secondary=*/true);
@@ -389,27 +389,37 @@ void multiSplitModule(const WasmSplitOptions& options) {
   Module wasm;
   parseInput(wasm, options);
 
-  // Map module names to the functions that should be in the modules.
-  std::map<Name, std::unordered_set<Name>> moduleFuncs;
   // The module for which we are currently parsing a set of functions.
   Name currModule;
   // The set of functions we are currently inserting into.
-  std::unordered_set<Name>* currFuncs = nullptr;
+  std::set<Name>* currFuncs = nullptr;
   // Map functions to their modules to ensure no function is assigned to
   // multiple modules.
   std::unordered_map<Name, Name> funcModules;
 
+  ModuleSplitting::Config config;
   std::string line;
   bool newSection = true;
+  std::vector<Name> moduleNames;
+  std::unordered_set<Name> moduleNameSet;
   while (std::getline(manifest, line)) {
     if (line.empty()) {
       newSection = true;
+      if (currFuncs->empty() && !options.quiet) {
+        std::cerr << "warning: Module " << currModule << " will be empty\n";
+      }
       continue;
     }
     Name name = WasmBinaryReader::escape(line);
     if (newSection) {
+      if (moduleNameSet.count(name)) {
+        Fatal() << "Module name " << name << " is listed more than once\n";
+      }
       currModule = name;
-      currFuncs = &moduleFuncs[name];
+      moduleNameSet.insert(currModule);
+      moduleNames.push_back(currModule);
+      config.secondaryFuncs.emplace_back(std::set<Name>());
+      currFuncs = &config.secondaryFuncs.back();
       newSection = false;
       continue;
     }
@@ -426,42 +436,33 @@ void multiSplitModule(const WasmSplitOptions& options) {
     }
   }
 
-  ModuleSplitting::Config config;
   config.usePlaceholders = options.usePlaceholders;
   config.importNamespace = options.importNamespace;
   config.minimizeNewExportNames = !options.passOptions.debugInfo;
   if (options.emitModuleNames && !wasm.name) {
     wasm.name = Path::getBaseName(options.output);
   }
 
-  std::unordered_map<Name, std::map<size_t, Name>> placeholderMap;
-  for (auto& [mod, funcs] : moduleFuncs) {
-    if (options.verbose) {
-      std::cerr << "Splitting module " << mod << '\n';
-    }
-    if (!options.quiet && funcs.empty()) {
-      std::cerr << "warning: Module " << mod << " will be empty\n";
-    }
-    config.secondaryFuncs = std::set<Name>(funcs.begin(), funcs.end());
-    auto splitResults = ModuleSplitting::splitFunctions(wasm, config);
-    auto moduleName = options.outPrefix + mod.toString() +
+  auto splitResults = ModuleSplitting::splitFunctions(wasm, config);
+  assert(moduleNames.size() == splitResults.secondaries.size());
+  for (Index i = 0, n = moduleNames.size(); i < n; i++) {
+    auto& secondary = *splitResults.secondaries[i];
+    auto moduleName = options.outPrefix + moduleNames[i].toString() +
                       (options.emitBinary ? ".wasm" : ".wast");
     if (options.symbolMap) {
-      writeSymbolMap(*splitResults.secondary, moduleName + ".symbols");
-    }
-    if (options.placeholderMap) {
-      placeholderMap.merge(splitResults.placeholderMap);
+      writeSymbolMap(secondary, moduleName + ".symbols");
     }
     if (options.emitModuleNames) {
-      splitResults.secondary->name = Path::getBaseName(moduleName);
+      secondary.name = Path::getBaseName(moduleName);
     }
-    writeModule(*splitResults.secondary, moduleName, options);
+    writeModule(secondary, moduleName, options);
   }
   if (options.symbolMap) {
     writeSymbolMap(wasm, options.output + ".symbols");
   }
   if (options.placeholderMap) {
-    writePlaceholderMap(wasm, placeholderMap, options.output + ".placeholders");
+    writePlaceholderMap(
+      wasm, splitResults.placeholderMap, options.output + ".placeholders");
   }
   writeModule(wasm, options.output, options);
 }

src/tools/wasm-split/wasm-split.cpp

@@ -51,9 +51,10 @@ void do_test(const std::set<Name>& keptFuncs, std::string&& module) {
   std::cout << "\n";
 
   ModuleSplitting::Config config;
-  config.secondaryFuncs = std::move(splitFuncs);
+  config.secondaryFuncs.push_back(std::move(splitFuncs));
   config.newExportPrefix = "%";
-  auto secondary = splitFunctions(*primary, config).secondary;
+  auto results = splitFunctions(*primary, config);
+  auto& secondary = *results.secondaries.begin();
 
   std::cout << "After:\n";
   std::cout << *primary.get();
@@ -475,11 +476,12 @@ void test_minimized_exports() {
   primary.addFunction(Builder::makeFunction("call", funcType, {}, callBody));
 
   ModuleSplitting::Config config;
-  config.secondaryFuncs = {"call"};
+  config.secondaryFuncs.push_back({"call"});
   config.newExportPrefix = "%";
   config.minimizeNewExportNames = true;
 
-  auto secondary = splitFunctions(primary, config).secondary;
+  auto results = splitFunctions(primary, config);
+  auto& secondary = *results.secondaries.begin();
   std::cout << "Minimized names primary:\n";
   std::cout << primary << "\n";
   std::cout << "Minimized names secondary:\n";

test/example/module-splitting.cpp

@@ -7,11 +7,13 @@
 ;; RUN: wasm-dis %t3.wasm | filecheck %s --check-prefix=MOD3
 
 (module
+ ;; PRIMARY:      (type $ret-i64 (func (result i64)))
+
+ ;; PRIMARY:      (type $ret-f32 (func (result f32)))
+
  ;; PRIMARY:      (type $ret-i32 (func (result i32)))
  (type $ret-i32 (func (result i32)))
- ;; PRIMARY:      (type $ret-i64 (func (result i64)))
  (type $ret-i64 (func (result i64)))
- ;; PRIMARY:      (type $ret-f32 (func (result f32)))
  (type $ret-f32 (func (result f32)))
 
  ;; MOD1:      (type $0 (func (result f32)))
@@ -20,13 +22,13 @@
 
  ;; MOD1:      (type $2 (func (result i32)))
 
- ;; MOD1:      (import "" "table" (table $timport$0 1 funcref))
+ ;; MOD1:      (import "" "table" (table $timport$0 3 funcref))
 
- ;; MOD1:      (import "" "std::operator<<\\28std::__2::basic_ostream<char\\2c\\20std::__2::char_traits<char>>&\\2c\\20wasm::Module&\\29" (func $std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29 (result f32)))
+ ;; MOD1:      (import "" "trampoline_std::operator<<\\28std::__2::basic_ostream<char\\2c\\20std::__2::char_traits<char>>&\\2c\\20wasm::Module&\\29" (func $trampoline_std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29 (result f32)))
 
- ;; MOD1:      (import "" "wasm::Literal::Literal\\28std::__2::array<wasm::Literal\\2c\\204ul>\\20const&\\29" (func $wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29 (result i64)))
+ ;; MOD1:      (import "" "trampoline_wasm::Literal::Literal\\28std::__2::array<wasm::Literal\\2c\\204ul>\\20const&\\29" (func $trampoline_wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29 (result i64)))
 
- ;; MOD1:      (elem $0 (i32.const 0) $wasm::Type::getFeatures\28\29\20const)
+ ;; MOD1:      (elem $0 (i32.const 2) $wasm::Type::getFeatures\28\29\20const)
 
  ;; MOD1:      (func $wasm::Type::getFeatures\28\29\20const (result i32)
  ;; MOD1-NEXT:  (drop
@@ -36,12 +38,12 @@
  ;; MOD1-NEXT:  )
  ;; MOD1-NEXT:  (drop
  ;; MOD1-NEXT:   (call_ref $1
- ;; MOD1-NEXT:    (ref.func $wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29)
+ ;; MOD1-NEXT:    (ref.func $trampoline_wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29)
  ;; MOD1-NEXT:   )
  ;; MOD1-NEXT:  )
  ;; MOD1-NEXT:  (drop
  ;; MOD1-NEXT:   (call_ref $0
- ;; MOD1-NEXT:    (ref.func $std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29)
+ ;; MOD1-NEXT:    (ref.func $trampoline_std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29)
  ;; MOD1-NEXT:   )
  ;; MOD1-NEXT:  )
  ;; MOD1-NEXT:  (i32.const 0)
@@ -71,9 +73,9 @@
 
  ;; MOD2:      (type $2 (func (result i64)))
 
- ;; MOD2:      (import "" "table_4" (table $timport$0 1 funcref))
+ ;; MOD2:      (import "" "table" (table $timport$0 3 funcref))
 
- ;; MOD2:      (import "" "std::operator<<\\28std::__2::basic_ostream<char\\2c\\20std::__2::char_traits<char>>&\\2c\\20wasm::Module&\\29" (func $std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29 (result f32)))
+ ;; MOD2:      (import "" "trampoline_std::operator<<\\28std::__2::basic_ostream<char\\2c\\20std::__2::char_traits<char>>&\\2c\\20wasm::Module&\\29" (func $trampoline_std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29 (result f32)))
 
  ;; MOD2:      (import "" "trampoline_wasm::Type::getFeatures\\28\\29\\20const" (func $trampoline_wasm::Type::getFeatures\28\29\20const (result i32)))
 
@@ -92,7 +94,7 @@
  ;; MOD2-NEXT:  )
  ;; MOD2-NEXT:  (drop
  ;; MOD2-NEXT:   (call_ref $0
- ;; MOD2-NEXT:    (ref.func $std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29)
+ ;; MOD2-NEXT:    (ref.func $trampoline_std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29)
  ;; MOD2-NEXT:   )
  ;; MOD2-NEXT:  )
  ;; MOD2-NEXT:  (i64.const 0)
@@ -122,13 +124,13 @@
 
  ;; MOD3:      (type $2 (func (result f32)))
 
- ;; MOD3:      (import "" "table_5" (table $timport$0 1 funcref))
+ ;; MOD3:      (import "" "table" (table $timport$0 3 funcref))
 
- ;; MOD3:      (import "" "wasm::Literal::Literal\\28std::__2::array<wasm::Literal\\2c\\204ul>\\20const&\\29" (func $trampoline_wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29 (result i64)))
+ ;; MOD3:      (import "" "trampoline_wasm::Literal::Literal\\28std::__2::array<wasm::Literal\\2c\\204ul>\\20const&\\29" (func $trampoline_wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29 (result i64)))
 
  ;; MOD3:      (import "" "trampoline_wasm::Type::getFeatures\\28\\29\\20const" (func $trampoline_wasm::Type::getFeatures\28\29\20const (result i32)))
 
- ;; MOD3:      (elem $0 (i32.const 0) $std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29)
+ ;; MOD3:      (elem $0 (i32.const 1) $std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29)
 
  ;; MOD3:      (func $std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29 (result f32)
  ;; MOD3-NEXT:  (drop
@@ -167,50 +169,38 @@
   (f32.const 0)
  )
 )
-;; PRIMARY:      (import "placeholder" "0" (func $placeholder_0 (result i32)))
+;; PRIMARY:      (import "placeholder" "0" (func $placeholder_0 (result i64)))
 
-;; PRIMARY:      (import "placeholder" "0" (func $placeholder_0_4 (result i64)))
+;; PRIMARY:      (import "placeholder" "1" (func $placeholder_1 (result f32)))
 
-;; PRIMARY:      (import "placeholder" "0" (func $placeholder_0_5 (result f32)))
+;; PRIMARY:      (import "placeholder" "2" (func $placeholder_2 (result i32)))
 
-;; PRIMARY:      (table $0 1 funcref)
+;; PRIMARY:      (table $0 3 funcref)
 
-;; PRIMARY:      (table $1 1 funcref)
+;; PRIMARY:      (elem $0 (i32.const 0) $placeholder_0 $placeholder_1 $placeholder_2)
 
-;; PRIMARY:      (table $2 1 funcref)
+;; PRIMARY:      (export "trampoline_std::operator<<\\28std::__2::basic_ostream<char\\2c\\20std::__2::char_traits<char>>&\\2c\\20wasm::Module&\\29" (func $trampoline_std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29))
 
-;; PRIMARY:      (elem $0 (table $0) (i32.const 0) func $placeholder_0)
-
-;; PRIMARY:      (elem $1 (table $1) (i32.const 0) func $placeholder_0_4)
-
-;; PRIMARY:      (elem $2 (table $2) (i32.const 0) func $placeholder_0_5)
-
-;; PRIMARY:      (export "std::operator<<\\28std::__2::basic_ostream<char\\2c\\20std::__2::char_traits<char>>&\\2c\\20wasm::Module&\\29" (func $trampoline_std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29))
-
-;; PRIMARY:      (export "wasm::Literal::Literal\\28std::__2::array<wasm::Literal\\2c\\204ul>\\20const&\\29" (func $trampoline_wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29))
-
-;; PRIMARY:      (export "table" (table $0))
+;; PRIMARY:      (export "trampoline_wasm::Literal::Literal\\28std::__2::array<wasm::Literal\\2c\\204ul>\\20const&\\29" (func $trampoline_wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29))
 
 ;; PRIMARY:      (export "trampoline_wasm::Type::getFeatures\\28\\29\\20const" (func $trampoline_wasm::Type::getFeatures\28\29\20const))
 
-;; PRIMARY:      (export "table_4" (table $1))
-
-;; PRIMARY:      (export "table_5" (table $2))
+;; PRIMARY:      (export "table" (table $0))
 
-;; PRIMARY:      (func $trampoline_wasm::Type::getFeatures\28\29\20const (result i32)
-;; PRIMARY-NEXT:  (call_indirect (type $ret-i32)
+;; PRIMARY:      (func $trampoline_wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29 (result i64)
+;; PRIMARY-NEXT:  (call_indirect (type $ret-i64)
 ;; PRIMARY-NEXT:   (i32.const 0)
 ;; PRIMARY-NEXT:  )
 ;; PRIMARY-NEXT: )
 
-;; PRIMARY:      (func $trampoline_wasm::Literal::Literal\28std::__2::array<wasm::Literal\2c\204ul>\20const&\29 (result i64)
-;; PRIMARY-NEXT:  (call_indirect $1 (type $ret-i64)
-;; PRIMARY-NEXT:   (i32.const 0)
+;; PRIMARY:      (func $trampoline_std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29 (result f32)
+;; PRIMARY-NEXT:  (call_indirect (type $ret-f32)
+;; PRIMARY-NEXT:   (i32.const 1)
 ;; PRIMARY-NEXT:  )
 ;; PRIMARY-NEXT: )
 
-;; PRIMARY:      (func $trampoline_std::operator<<\28std::__2::basic_ostream<char\2c\20std::__2::char_traits<char>>&\2c\20wasm::Module&\29 (result f32)
-;; PRIMARY-NEXT:  (call_indirect $2 (type $ret-f32)
-;; PRIMARY-NEXT:   (i32.const 0)
+;; PRIMARY:      (func $trampoline_wasm::Type::getFeatures\28\29\20const (result i32)
+;; PRIMARY-NEXT:  (call_indirect (type $ret-i32)
+;; PRIMARY-NEXT:   (i32.const 2)
 ;; PRIMARY-NEXT:  )
 ;; PRIMARY-NEXT: )