[wasm-split] Support --placeholdermap for --multi-split (#7792)

This adds support for `--placeholdermap` for `--multi-split`. Because
there can be multiple tables where placeholders are placed, this
modifies the placeholder file format slightly by adding table indices
for each table.

Author: aheejin

CHANGELOG.md

@@ -17,10 +17,11 @@ Current Trunk
 
  - `wasm-split`'s `--multi-split` mode now supports more options:
    `--no-placeholders`, `--import-namespace`, `--emit-module-names`,
-   `--emit-text`, and `--symbolmap`. Because `--no-placeholders` is false by
-   default and until now `--multi-split` didn't use placeholders at all, this is
-   a breaking change. If you want to continue to do multi-split without
-   placeholders, you need to explicitly specify `--no-placeholders`.
+   `--emit-text`, `--symbolmap`, and `--placeholdermap`. Because
+   `--no-placeholders` is false by default and until now `--multi-split` didn't
+   use placeholders at all, this is a breaking change. If you want to continue
+   to do multi-split without placeholders, you need to explicitly specify
+   `--no-placeholders`.
  - Add a `--string-lifting` pass that raises imported string operations and
    constants into stringref in Binaryen IR (which can then be fully optimized,
    and typically lowered back down with `--string-lowering`).

src/ir/module-splitting.cpp

@@ -302,8 +302,9 @@ struct ModuleSplitter {
   // names.
   std::map<Name, Name> exportedPrimaryFuncs;
 
-  // Map placeholder indices to the names of the functions they replace.
-  std::map<size_t, Name> placeholderMap;
+  // For each table, map placeholder indices to the names of the functions they
+  // replace.
+  std::unordered_map<Name, std::map<size_t, Name>> placeholderMap;
 
   // Internal name of the LOAD_SECONDARY_MODULE function.
   Name internalLoadSecondaryModule;
@@ -723,7 +724,7 @@ void ModuleSplitter::setupTablePatching() {
       }
       assert(table == tableManager.activeTable->name);
 
-      placeholderMap[index] = ref->func;
+      placeholderMap[table][index] = ref->func;
       auto* secondaryFunc = secondary.getFunction(ref->func);
       replacedElems[index] = secondaryFunc;
       if (!config.usePlaceholders) {

src/ir/module-splitting.h

@@ -77,7 +77,7 @@ struct Config {
 
 struct Results {
   std::unique_ptr<Module> secondary;
-  std::map<size_t, Name> placeholderMap;
+  std::unordered_map<Name, std::map<size_t, Name>> placeholderMap;
 };
 
 // Returns the new secondary module and modifies the `primary` module in place.

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

@@ -221,7 +221,7 @@ WasmSplitOptions::WasmSplitOptions()
       "",
       "Write a file mapping placeholder indices to the function names.",
       WasmSplitOption,
-      {Mode::Split},
+      {Mode::Split, Mode::MultiSplit},
       Options::Arguments::Zero,
       [&](Options* o, const std::string& argument) { placeholderMap = true; })
     .add("--import-namespace",

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

@@ -20,13 +20,11 @@
 #include <fstream>
 
 #include "ir/module-splitting.h"
-#include "ir/names.h"
 #include "support/file.h"
 #include "support/name.h"
 #include "support/path.h"
 #include "support/utilities.h"
 #include "wasm-binary.h"
-#include "wasm-builder.h"
 #include "wasm-io.h"
 #include "wasm-validator.h"
 
@@ -203,12 +201,24 @@ void writeSymbolMap(Module& wasm, std::string filename) {
   runner.run();
 }
 
-void writePlaceholderMap(const std::map<size_t, Name> placeholderMap,
-                         std::string filename) {
+void writePlaceholderMap(
+  Module& wasm,
+  const std::unordered_map<Name, std::map<size_t, Name>>& placeholderMap,
+  std::string filename) {
   Output output(filename, Flags::Text);
   auto& o = output.getStream();
-  for (auto& [index, func] : placeholderMap) {
-    o << index << ':' << func << '\n';
+  for (Index i = 0; i < wasm.tables.size(); i++) {
+    const auto& table = wasm.tables[i];
+    auto it = placeholderMap.find(table->name);
+    if (it != placeholderMap.end()) {
+      o << "table " << i << "\n";
+      for (auto& [index, func] : it->second) {
+        o << index << ':' << func << '\n';
+      }
+      if (i < wasm.tables.size() - 1) {
+        o << "\n";
+      }
+    }
   }
 }
 
@@ -344,7 +354,8 @@ void splitModule(const WasmSplitOptions& options) {
   }
 
   if (options.placeholderMap) {
-    writePlaceholderMap(splitResults.placeholderMap,
+    writePlaceholderMap(wasm,
+                        splitResults.placeholderMap,
                         options.primaryOutput + ".placeholders");
   }
 
@@ -422,6 +433,7 @@ void multiSplitModule(const WasmSplitOptions& options) {
     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';
@@ -431,12 +443,14 @@ void multiSplitModule(const WasmSplitOptions& options) {
     }
     config.secondaryFuncs = std::set<Name>(funcs.begin(), funcs.end());
     auto splitResults = ModuleSplitting::splitFunctions(wasm, config);
-    // TODO: placeholderMap
     auto moduleName =
       options.outPrefix + mod + (options.emitBinary ? ".wasm" : ".wast");
     if (options.symbolMap) {
       writeSymbolMap(*splitResults.secondary, moduleName + ".symbols");
     }
+    if (options.placeholderMap) {
+      placeholderMap.merge(splitResults.placeholderMap);
+    }
     if (options.emitModuleNames) {
       splitResults.secondary->name = Path::getBaseName(moduleName);
     }
@@ -445,7 +459,9 @@ void multiSplitModule(const WasmSplitOptions& options) {
   if (options.symbolMap) {
     writeSymbolMap(wasm, options.output + ".symbols");
   }
-
+  if (options.placeholderMap) {
+    writePlaceholderMap(wasm, placeholderMap, options.output + ".placeholders");
+  }
   writeModule(wasm, options.output, options);
 }
 

test/lit/help/wasm-split.test

@@ -72,8 +72,9 @@
 ;; CHECK-NEXT:                                        functions will fail before the secondary
 ;; CHECK-NEXT:                                        module has been instantiated.
 ;; CHECK-NEXT:
-;; CHECK-NEXT:   --placeholdermap                     [split] Write a file mapping placeholder
-;; CHECK-NEXT:                                        indices to the function names.
+;; CHECK-NEXT:   --placeholdermap                     [split, multi-split] Write a file mapping
+;; CHECK-NEXT:                                        placeholder indices to the function
+;; CHECK-NEXT:                                        names.
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --import-namespace                   [split, instrument, multi-split] When
 ;; CHECK-NEXT:                                        provided as an option for module

test/lit/wasm-split/placeholdermap-multi-split.wast

@@ -0,0 +1,74 @@
+;; RUN: wasm-split -all --multi-split %s --manifest %S/multi-split.wast.manifest --out-prefix=%t --placeholdermap -o %t.wasm
+;; RUN: filecheck %s --check-prefix MAP < %t.wasm.placeholders
+
+;; MAP:      table 0
+;; MAP-NEXT: 0:A
+;; MAP-NEXT:
+;; MAP-NEXT: table 1
+;; MAP-NEXT: 0:B
+;; MAP-NEXT:
+;; MAP-NEXT: table 2
+;; MAP-NEXT: 0:C
+
+(module
+ (type $ret-i32 (func (result i32)))
+ (type $ret-i64 (func (result i64)))
+ (type $ret-f32 (func (result f32)))
+
+ (func $A (type $ret-i32) (result i32)
+  (drop
+   (call_ref $ret-i32
+    (ref.func $A)
+   )
+  )
+  (drop
+   (call_ref $ret-i64
+    (ref.func $B)
+   )
+  )
+  (drop
+   (call_ref $ret-f32
+    (ref.func $C)
+   )
+  )
+  (i32.const 0)
+ )
+
+ (func $B (type $ret-i64) (result i64)
+  (drop
+   (call_ref $ret-i32
+    (ref.func $A)
+   )
+  )
+  (drop
+   (call_ref $ret-i64
+    (ref.func $B)
+   )
+  )
+  (drop
+   (call_ref $ret-f32
+    (ref.func $C)
+   )
+  )
+  (i64.const 0)
+ )
+
+ (func $C (type $ret-f32) (result f32)
+  (drop
+   (call_ref $ret-i32
+    (ref.func $A)
+   )
+  )
+  (drop
+   (call_ref $ret-i64
+    (ref.func $B)
+   )
+  )
+  (drop
+   (call_ref $ret-f32
+    (ref.func $C)
+   )
+  )
+  (f32.const 0)
+ )
+)

test/lit/wasm-split/placeholdermap.wast

@@ -2,7 +2,8 @@
 ;; RUN: filecheck %s --check-prefix MAP < %t.1.wasm.placeholders
 ;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
 
-;; MAP:      0:foo
+;; MAP:      table 0
+;; MAP-NEXT: 0:foo
 ;; MAP-NEXT: 2:baz
 ;; MAP-NOT: bar