[wasm-split] Add a multi-split mode

Add a mode that splits a module into arbitrarily many parts based on a
simple manifest file. This is currently implemented by splitting out one
module at a time in a loop, but this could change in the future if
splitting out all the modules at once would improve the quality of the
output.

Author: tlively

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

@@ -61,6 +61,9 @@ std::ostream& operator<<(std::ostream& o, WasmSplitOptions::Mode& mode) {
     case WasmSplitOptions::Mode::Split:
       o << "split";
       break;
+    case WasmSplitOptions::Mode::MultiSplit:
+      o << "multi-split";
+      break;
     case WasmSplitOptions::Mode::Instrument:
       o << "instrument";
       break;
@@ -91,7 +94,14 @@ WasmSplitOptions::WasmSplitOptions()
          "Split an input module into two output modules. The default mode.",
          WasmSplitOption,
          Options::Arguments::Zero,
-         [&](Options* o, const std::string& arugment) { mode = Mode::Split; })
+         [&](Options* o, const std::string& argument) { mode = Mode::Split; })
+    .add(
+      "--multi-split",
+      "",
+      "Split an input module into an arbitrary number of output modules.",
+      WasmSplitOption,
+      Options::Arguments::Zero,
+      [&](Options* o, const std::string& argument) { mode = Mode::MultiSplit; })
     .add(
       "--instrument",
       "",
@@ -151,6 +161,25 @@ WasmSplitOptions::WasmSplitOptions()
          [&](Options* o, const std::string& argument) {
            splitFuncs = parseNameList(argument);
          })
+    .add(
+      "--manifest",
+      "",
+      "File describing the functions to be split into each module. Each "
+      "section separated by a blank line begins with the base name of an "
+      "output module, which is followed by a list of functions to place in "
+      "that module, one per line.",
+      WasmSplitOption,
+      {Mode::MultiSplit},
+      Options::Arguments::One,
+      [&](Options* o, const std::string& argument) { manifestFile = argument; })
+    .add("--out-prefix",
+         "",
+         "Prefix prepended to module names in the manifest file to create "
+         "output file names.",
+         WasmSplitOption,
+         {Mode::MultiSplit},
+         Options::Arguments::One,
+         [&](Options* o, const std::string& argument) { outPrefix = argument; })
     .add("--primary-output",
          "-o1",
          "Output file for the primary module.",
@@ -313,7 +342,7 @@ WasmSplitOptions::WasmSplitOptions()
          "-g",
          "Emit names section in wasm binary (or full debuginfo in wast)",
          WasmSplitOption,
-         {Mode::Split, Mode::Instrument},
+         {Mode::Split, Mode::MultiSplit, Mode::Instrument},
          Options::Arguments::Zero,
          [&](Options* o, const std::string& arguments) {
            passOptions.debugInfo = true;
@@ -322,7 +351,7 @@ WasmSplitOptions::WasmSplitOptions()
          "-o",
          "Output file.",
          WasmSplitOption,
-         {Mode::Instrument, Mode::MergeProfiles},
+         {Mode::Instrument, Mode::MergeProfiles, Mode::MultiSplit},
          Options::Arguments::One,
          [&](Options* o, const std::string& argument) { output = argument; })
     .add("--unescape",
@@ -407,6 +436,7 @@ bool WasmSplitOptions::validate() {
   }
   switch (mode) {
     case Mode::Split:
+    case Mode::MultiSplit:
     case Mode::Instrument:
       if (inputFiles.size() > 1) {
         fail("Cannot have more than one input file.");

src/tools/wasm-split/split-options.h

@@ -26,6 +26,7 @@ const std::string DEFAULT_PROFILE_EXPORT("__write_profile");
 struct WasmSplitOptions : ToolOptions {
   enum class Mode : unsigned {
     Split,
+    MultiSplit,
     Instrument,
     MergeProfiles,
     PrintProfile,
@@ -68,6 +69,9 @@ struct WasmSplitOptions : ToolOptions {
   std::string secondaryMemoryName;
   std::string exportPrefix;
 
+  std::string manifestFile;
+  std::string outPrefix;
+
   // A hack to ensure the split and instrumented modules have the same table
   // size when using Emscripten's SPLIT_MODULE mode with dynamic linking. TODO:
   // Figure out a more elegant solution for that use case and remove this.

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

@@ -362,6 +362,82 @@ void splitModule(const WasmSplitOptions& options) {
   writeModule(*secondary, options.secondaryOutput, options);
 }
 
+void multiSplitModule(const WasmSplitOptions& options) {
+  if (options.manifestFile.empty()) {
+    Fatal() << "--multi-split requires --manifest";
+  }
+  if (options.output.empty()) {
+    Fatal() << "--multi-split requires --output";
+  }
+
+  std::ifstream manifest(options.manifestFile);
+  if (!manifest.is_open()) {
+    Fatal() << "File not found: " << options.manifestFile;
+  }
+
+  Module wasm;
+  parseInput(wasm, options);
+
+  std::map<std::string, std::unordered_set<std::string>> moduleFuncs;
+  std::string currModule;
+  std::unordered_set<std::string>* currFuncs = nullptr;
+  std::unordered_map<std::string, std::string> funcModules;
+
+  std::string line;
+  bool newSection = true;
+  while (std::getline(manifest, line)) {
+    if (line.empty()) {
+      newSection = true;
+      continue;
+    }
+    if (newSection) {
+      currModule = line;
+      currFuncs = &moduleFuncs[line];
+      newSection = false;
+      continue;
+    }
+    assert(currFuncs);
+    currFuncs->insert(line);
+    auto [it, inserted] = funcModules.insert({line, currModule});
+    if (!inserted && it->second != currModule) {
+      Fatal() << "Function " << line << "cannot be assigned to module "
+              << currModule << "; it is already assigned to module "
+              << it->second << '\n';
+    }
+    if (inserted && !options.quiet && !wasm.getFunctionOrNull(line)) {
+      std::cerr << "warning: Function " << line << " does not exist\n";
+    }
+  }
+
+  ModuleSplitting::Config config;
+  config.usePlaceholders = false;
+  config.importNamespace = "";
+  config.minimizeNewExportNames = true;
+  for (auto& func : wasm.functions) {
+    config.primaryFuncs.insert(func->name);
+  }
+  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";
+    }
+    for (auto& func : funcs) {
+      config.primaryFuncs.erase(Name(func));
+    }
+    auto splitResults = ModuleSplitting::splitFunctions(wasm, config);
+    // TODO: symbolMap, placeholderMap, emitModuleNames
+    // TODO: Support --emit-text and use .wast in that case.
+    auto moduleName = options.outPrefix + mod + ".wasm";
+    PassRunner runner(&*splitResults.secondary);
+    runner.add("remove-unused-module-elements");
+    runner.run();
+    writeModule(*splitResults.secondary, moduleName, options);
+  }
+  writeModule(wasm, options.output, options);
+}
+
 void mergeProfiles(const WasmSplitOptions& options) {
   // Read the initial profile. We will merge other profiles into this one.
   ProfileData data = readProfile(options.inputFiles[0]);
@@ -503,6 +579,9 @@ int main(int argc, const char* argv[]) {
     case WasmSplitOptions::Mode::Split:
       splitModule(options);
       break;
+    case WasmSplitOptions::Mode::MultiSplit:
+      multiSplitModule(options);
+      break;
     case WasmSplitOptions::Mode::Instrument:
       instrumentModule(options);
       break;