Multi-Memories wasm-split (#4977)

Adds an --in-secondary-memory switch to the wasm-split tool that allows profile data to be stored in a separate memory from module main memory. With this option, users do not need to reserve the initial memory region for profile data and the data can be shared between multiple threads.

5 files changed, 118 insertions, 28 deletions

diff --git a/src/tools/wasm-split/instrumenter.cpp b/src/tools/wasm-split/instrumenter.cpp
index 50da9d034..7f9ff9c75 100644
--- a/src/tools/wasm-split/instrumenter.cpp
+++ b/src/tools/wasm-split/instrumenter.cpp
@@ -22,25 +22,31 @@
 
 namespace wasm {
 
-Instrumenter::Instrumenter(const WasmSplitOptions& options, uint64_t moduleHash)
-  : options(options), moduleHash(moduleHash) {}
+Instrumenter::Instrumenter(const InstrumenterConfig& config,
+                           uint64_t moduleHash)
+  : config(config), moduleHash(moduleHash) {}
 
 void Instrumenter::run(PassRunner* runner, Module* wasm) {
   this->runner = runner;
   this->wasm = wasm;
-  addGlobals();
+
+  size_t numFuncs = 0;
+  ModuleUtils::iterDefinedFunctions(*wasm, [&](Function*) { ++numFuncs; });
+
+  addGlobals(numFuncs);
+  addSecondaryMemory(numFuncs);
   instrumentFuncs();
-  addProfileExport();
+  addProfileExport(numFuncs);
 }
 
-void Instrumenter::addGlobals() {
-  if (options.storageKind != WasmSplitOptions::StorageKind::InGlobals) {
+void Instrumenter::addGlobals(size_t numFuncs) {
+  if (config.storageKind != WasmSplitOptions::StorageKind::InGlobals) {
     // Don't need globals
     return;
   }
   // Create fresh global names (over-reserves, but that's ok)
   counterGlobal = Names::getValidGlobalName(*wasm, "monotonic_counter");
-  functionGlobals.reserve(wasm->functions.size());
+  functionGlobals.reserve(numFuncs);
   ModuleUtils::iterDefinedFunctions(*wasm, [&](Function* func) {
     functionGlobals.push_back(Names::getValidGlobalName(
       *wasm, std::string(func->name.c_str()) + "_timestamp"));
@@ -62,11 +68,31 @@ void Instrumenter::addGlobals() {
   }
 }
 
+void Instrumenter::addSecondaryMemory(size_t numFuncs) {
+  if (config.storageKind != WasmSplitOptions::StorageKind::InSecondaryMemory) {
+    // Don't need secondary memory
+    return;
+  }
+  if (!wasm->features.hasMultiMemories()) {
+    Fatal()
+      << "error: --in-secondary-memory requires multi-memories to be enabled";
+  }
+
+  secondaryMemory =
+    Names::getValidMemoryName(*wasm, config.secondaryMemoryName);
+  // Create a memory with enough pages to write into
+  size_t pages = (numFuncs + Memory::kPageSize - 1) / Memory::kPageSize;
+  auto mem = Builder::makeMemory(secondaryMemory, pages, pages, true);
+  mem->module = config.importNamespace;
+  mem->base = config.secondaryMemoryName;
+  wasm->addMemory(std::move(mem));
+}
+
 void Instrumenter::instrumentFuncs() {
   // Inject code at the beginning of each function to advance the monotonic
   // counter and set the function's timestamp if it hasn't already been set.
   Builder builder(*wasm);
-  switch (options.storageKind) {
+  switch (config.storageKind) {
     case WasmSplitOptions::StorageKind::InGlobals: {
       // (if (i32.eqz (global.get $timestamp))
       //   (block
@@ -102,13 +128,22 @@ void Instrumenter::instrumentFuncs() {
       });
       break;
     }
-    case WasmSplitOptions::StorageKind::InMemory: {
+    case WasmSplitOptions::StorageKind::InMemory:
+    case WasmSplitOptions::StorageKind::InSecondaryMemory: {
       if (!wasm->features.hasAtomics()) {
-        Fatal() << "error: --in-memory requires atomics to be enabled";
+        const char* command =
+          config.storageKind == WasmSplitOptions::StorageKind::InMemory
+            ? "in-memory"
+            : "in-secondary-memory";
+        Fatal() << "error: --" << command << " requires atomics to be enabled";
       }
       // (i32.atomic.store8 offset=funcidx (i32.const 0) (i32.const 1))
       Index funcIdx = 0;
       assert(!wasm->memories.empty());
+      Name memoryName =
+        config.storageKind == WasmSplitOptions::StorageKind::InMemory
+          ? wasm->memories[0]->name
+          : secondaryMemory;
       ModuleUtils::iterDefinedFunctions(*wasm, [&](Function* func) {
         func->body = builder.makeSequence(
           builder.makeAtomicStore(1,
@@ -116,7 +151,7 @@ void Instrumenter::instrumentFuncs() {
                                   builder.makeConstPtr(0, Type::i32),
                                   builder.makeConst(uint32_t(1)),
                                   Type::i32,
-                                  wasm->memories[0]->name),
+                                  memoryName),
           func->body,
           func->body->type);
         ++funcIdx;
@@ -141,21 +176,18 @@ void Instrumenter::instrumentFuncs() {
 // otherwise. Functions with smaller non-zero timestamps were called earlier in
 // the instrumented run than funtions with larger timestamps.
 
-void Instrumenter::addProfileExport() {
+void Instrumenter::addProfileExport(size_t numFuncs) {
   // Create and export a function to dump the profile into a given memory
   // buffer. The function takes the available address and buffer size as
   // arguments and returns the total size of the profile. It only actually
   // writes the profile if the given space is sufficient to hold it.
-  auto name = Names::getValidFunctionName(*wasm, options.profileExport);
+  auto name = Names::getValidFunctionName(*wasm, config.profileExport);
   auto writeProfile = Builder::makeFunction(
     name, Signature({Type::i32, Type::i32}, Type::i32), {});
   writeProfile->hasExplicitName = true;
   writeProfile->setLocalName(0, "addr");
   writeProfile->setLocalName(1, "size");
 
-  size_t numFuncs = 0;
-  ModuleUtils::iterDefinedFunctions(*wasm, [&](Function*) { ++numFuncs; });
-
   // Calculate the size of the profile:
   //   8 bytes module hash +
   //   4 bytes for the timestamp for each function
@@ -188,7 +220,7 @@ void Instrumenter::addProfileExport() {
     8, 0, 1, getAddr(), hashConst(), Type::i64, wasm->memories[0]->name);
   uint32_t offset = 8;
 
-  switch (options.storageKind) {
+  switch (config.storageKind) {
     case WasmSplitOptions::StorageKind::InGlobals: {
       for (const auto& global : functionGlobals) {
         writeData = builder.blockify(
@@ -204,12 +236,17 @@ void Instrumenter::addProfileExport() {
       }
       break;
     }
-    case WasmSplitOptions::StorageKind::InMemory: {
+    case WasmSplitOptions::StorageKind::InMemory:
+    case WasmSplitOptions::StorageKind::InSecondaryMemory: {
       Index funcIdxVar =
         Builder::addVar(writeProfile.get(), "funcIdx", Type::i32);
       auto getFuncIdx = [&]() {
         return builder.makeLocalGet(funcIdxVar, Type::i32);
       };
+      Name loadMemoryName =
+        config.storageKind == WasmSplitOptions::StorageKind::InMemory
+          ? wasm->memories[0]->name
+          : secondaryMemory;
       // (block $outer
       //   (loop $l
       //     (br_if $outer (i32.eq (local.get $fucIdx) (i32.const numFuncs))
@@ -249,7 +286,7 @@ void Instrumenter::addProfileExport() {
                   builder.makeBinary(
                     MulInt32, getFuncIdx(), builder.makeConst(uint32_t(4)))),
                 builder.makeAtomicLoad(
-                  1, 0, getFuncIdx(), Type::i32, wasm->memories[0]->name),
+                  1, 0, getFuncIdx(), Type::i32, loadMemoryName),
                 Type::i32,
                 wasm->memories[0]->name),
               builder.makeLocalSet(
@@ -269,7 +306,7 @@ void Instrumenter::addProfileExport() {
   // Create an export for the function
   wasm->addFunction(std::move(writeProfile));
   wasm->addExport(
-    Builder::makeExport(options.profileExport, name, ExternalKind::Function));
+    Builder::makeExport(config.profileExport, name, ExternalKind::Function));
 
   // Export the memory if it is not already exported or imported.
   if (!wasm->memories[0]->imported()) {
diff --git a/src/tools/wasm-split/instrumenter.h b/src/tools/wasm-split/instrumenter.h
index 7de5a9135..f4f2f2119 100644
--- a/src/tools/wasm-split/instrumenter.h
+++ b/src/tools/wasm-split/instrumenter.h
@@ -23,6 +23,20 @@
 
 namespace wasm {
 
+struct InstrumenterConfig {
+  // The namespace from which to import the secondary memory
+  Name importNamespace = "env";
+  // The name of the secondary memory created to store profile data during
+  // instrumentation
+  Name secondaryMemoryName = "profile-data";
+  // Where to store the profile data
+  WasmSplitOptions::StorageKind storageKind =
+    WasmSplitOptions::StorageKind::InGlobals;
+  // The export name of the function the embedder calls to write the profile
+  // into memory
+  std::string profileExport = DEFAULT_PROFILE_EXPORT;
+};
+
 // Add a global monotonic counter and a timestamp global for each function, code
 // at the beginning of each function to set its timestamp, and a new exported
 // function for dumping the profile data.
@@ -30,20 +44,23 @@ struct Instrumenter : public Pass {
   PassRunner* runner = nullptr;
   Module* wasm = nullptr;
 
-  const WasmSplitOptions& options;
+  const InstrumenterConfig& config;
   uint64_t moduleHash;
 
   Name counterGlobal;
   std::vector<Name> functionGlobals;
 
-  Instrumenter(const WasmSplitOptions& options, uint64_t moduleHash);
+  Name secondaryMemory;
+
+  Instrumenter(const InstrumenterConfig& config, uint64_t moduleHash);
 
   void run(PassRunner* runner, Module* wasm) override;
 
 private:
-  void addGlobals();
+  void addGlobals(size_t numFuncs);
+  void addSecondaryMemory(size_t numFuncs);
   void instrumentFuncs();
-  void addProfileExport();
+  void addProfileExport(size_t numFuncs);
 };
 
 } // namespace wasm
diff --git a/src/tools/wasm-split/split-options.cpp b/src/tools/wasm-split/split-options.cpp
index b8ccde83e..d077c70bc 100644
--- a/src/tools/wasm-split/split-options.cpp
+++ b/src/tools/wasm-split/split-options.cpp
@@ -185,10 +185,12 @@ WasmSplitOptions::WasmSplitOptions()
       [&](Options* o, const std::string& argument) { placeholderMap = true; })
     .add("--import-namespace",
          "",
-         "The namespace from which to import objects from the primary "
-         "module into the secondary module.",
+         "When provided as an option for module splitting, the namespace from "
+         "which to import objects from the primary "
+         "module into the secondary module. In instrument mode, refers to the "
+         "namespace from which to import the secondary memory, if any.",
          WasmSplitOption,
-         {Mode::Split},
+         {Mode::Split, Mode::Instrument},
          Options::Arguments::One,
          [&](Options* o, const std::string& argument) {
            importNamespace = argument;
@@ -246,6 +248,29 @@ WasmSplitOptions::WasmSplitOptions()
         storageKind = StorageKind::InMemory;
       })
     .add(
+      "--in-secondary-memory",
+      "",
+      "Store profile information in a separate memory, rather than in module "
+      "main memory or globals (the default). With this option, users do not "
+      "need to reserve the initial memory region for profile data and the "
+      "data can be shared between multiple threads.",
+      WasmSplitOption,
+      {Mode::Instrument},
+      Options::Arguments::Zero,
+      [&](Options* o, const std::string& argument) {
+        storageKind = StorageKind::InSecondaryMemory;
+      })
+    .add("--secondary-memory-name",
+         "",
+         "The name of the secondary memory created to store profile "
+         "information.",
+         WasmSplitOption,
+         {Mode::Instrument},
+         Options::Arguments::One,
+         [&](Options* o, const std::string& argument) {
+           secondaryMemoryName = argument;
+         })
+    .add(
       "--emit-module-names",
       "",
       "Emit module names, even if not emitting the rest of the names section. "
diff --git a/src/tools/wasm-split/split-options.h b/src/tools/wasm-split/split-options.h
index 9f6253626..5aba0ba39 100644
--- a/src/tools/wasm-split/split-options.h
+++ b/src/tools/wasm-split/split-options.h
@@ -37,6 +37,7 @@ struct WasmSplitOptions : ToolOptions {
   enum class StorageKind : unsigned {
     InGlobals, // Store profile data in WebAssembly Globals
     InMemory,  // Store profile data in memory, accessible from all threads
+    InSecondaryMemory, // Store profile data in memory separate from main memory
   };
   StorageKind storageKind = StorageKind::InGlobals;
 
@@ -63,6 +64,7 @@ struct WasmSplitOptions : ToolOptions {
 
   std::string importNamespace;
   std::string placeholderNamespace;
+  std::string secondaryMemoryName;
   std::string exportPrefix;
 
   // A hack to ensure the split and instrumented modules have the same table
diff --git a/src/tools/wasm-split/wasm-split.cpp b/src/tools/wasm-split/wasm-split.cpp
index 43db1b7d2..ec1f5555c 100644
--- a/src/tools/wasm-split/wasm-split.cpp
+++ b/src/tools/wasm-split/wasm-split.cpp
@@ -111,7 +111,16 @@ void instrumentModule(const WasmSplitOptions& options) {
 
   uint64_t moduleHash = hashFile(options.inputFiles[0]);
   PassRunner runner(&wasm, options.passOptions);
-  Instrumenter(options, moduleHash).run(&runner, &wasm);
+  InstrumenterConfig config;
+  if (options.importNamespace.size()) {
+    config.importNamespace = options.importNamespace;
+  }
+  if (options.secondaryMemoryName.size()) {
+    config.secondaryMemoryName = options.secondaryMemoryName;
+  }
+  config.storageKind = options.storageKind;
+  config.profileExport = options.profileExport;
+  Instrumenter(config, moduleHash).run(&runner, &wasm);
 
   adjustTableSize(wasm, options.initialTableSize);