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);