[wasm-split] Add an option for recording profile data in memory (#4120)

To avoid requiring a static memory allocation, wasm-split's instrumentation
defaults to recording profile data in Wasm globals. This causes problems for
multithreaded applications because the globals are thread-local, but it is not
always feasible to arrange for a separate profile to be dumped on each thread.
To simplify the profiling of such multithreaded applications, add a new
instrumentation mode that stores the profiling data in shared memory instead of
in globals. This allows a single profile to be written that correctly reflects
the called functions on all threads.

This new mode is not on by default because it requires users to ensure that the
program will not trample the in-memory profiling data. The data is stored
beginning at address zero and occupies one byte per declared function in the
instrumented module. Emscripten can be told to leave this memory free using the
GLOBAL_BASE option.

5 files changed, 165 insertions, 55 deletions

diff --git a/src/tools/wasm-split/instrumenter.cpp b/src/tools/wasm-split/instrumenter.cpp
index 0c5e96b54..79d6a98a4 100644
--- a/src/tools/wasm-split/instrumenter.cpp
+++ b/src/tools/wasm-split/instrumenter.cpp
@@ -22,9 +22,8 @@
 
 namespace wasm {
 
-Instrumenter::Instrumenter(const std::string& profileExport,
-                           uint64_t moduleHash)
-  : profileExport(profileExport), moduleHash(moduleHash) {}
+Instrumenter::Instrumenter(const WasmSplitOptions& options, uint64_t moduleHash)
+  : options(options), moduleHash(moduleHash) {}
 
 void Instrumenter::run(PassRunner* runner, Module* wasm) {
   this->runner = runner;
@@ -35,6 +34,10 @@ void Instrumenter::run(PassRunner* runner, Module* wasm) {
 }
 
 void Instrumenter::addGlobals() {
+  if (options.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());
@@ -60,42 +63,65 @@ void Instrumenter::addGlobals() {
 }
 
 void Instrumenter::instrumentFuncs() {
-  // Inject the following 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.
-  //
-  //   (if (i32.eqz (global.get $timestamp))
-  //     (block
-  //       (global.set $monotonic_counter
-  //         (i32.add
-  //           (global.get $monotonic_counter)
-  //           (i32.const 1)
-  //         )
-  //       )
-  //       (global.set $timestamp
-  //         (global.get $monotonic_counter)
-  //       )
-  //     )
-  //   )
+  // 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);
-  auto globalIt = functionGlobals.begin();
-  ModuleUtils::iterDefinedFunctions(*wasm, [&](Function* func) {
-    func->body = builder.makeSequence(
-      builder.makeIf(
-        builder.makeUnary(EqZInt32,
-                          builder.makeGlobalGet(*globalIt, Type::i32)),
-        builder.makeSequence(
-          builder.makeGlobalSet(
-            counterGlobal,
-            builder.makeBinary(AddInt32,
-                               builder.makeGlobalGet(counterGlobal, Type::i32),
-                               builder.makeConst(Literal::makeOne(Type::i32)))),
-          builder.makeGlobalSet(
-            *globalIt, builder.makeGlobalGet(counterGlobal, Type::i32)))),
-      func->body,
-      func->body->type);
-    ++globalIt;
-  });
+  switch (options.storageKind) {
+    case WasmSplitOptions::StorageKind::InGlobals: {
+      // (if (i32.eqz (global.get $timestamp))
+      //   (block
+      //     (global.set $monotonic_counter
+      //       (i32.add
+      //         (global.get $monotonic_counter)
+      //         (i32.const 1)
+      //       )
+      //     )
+      //     (global.set $timestamp
+      //       (global.get $monotonic_counter)
+      //     )
+      //   )
+      // )
+      auto globalIt = functionGlobals.begin();
+      ModuleUtils::iterDefinedFunctions(*wasm, [&](Function* func) {
+        func->body = builder.makeSequence(
+          builder.makeIf(
+            builder.makeUnary(EqZInt32,
+                              builder.makeGlobalGet(*globalIt, Type::i32)),
+            builder.makeSequence(
+              builder.makeGlobalSet(
+                counterGlobal,
+                builder.makeBinary(
+                  AddInt32,
+                  builder.makeGlobalGet(counterGlobal, Type::i32),
+                  builder.makeConst(Literal::makeOne(Type::i32)))),
+              builder.makeGlobalSet(
+                *globalIt, builder.makeGlobalGet(counterGlobal, Type::i32)))),
+          func->body,
+          func->body->type);
+        ++globalIt;
+      });
+      break;
+    }
+    case WasmSplitOptions::StorageKind::InMemory: {
+      if (!wasm->features.hasAtomics()) {
+        Fatal() << "error: --in-memory requires atomics to be enabled";
+      }
+      // (i32.atomic.store8 offset=funcidx (i32.const 0) (i32.const 1))
+      Index funcIdx = 0;
+      ModuleUtils::iterDefinedFunctions(*wasm, [&](Function* func) {
+        func->body = builder.makeSequence(
+          builder.makeAtomicStore(1,
+                                  funcIdx,
+                                  builder.makeConstPtr(0),
+                                  builder.makeConst(uint32_t(1)),
+                                  Type::i32),
+          func->body,
+          func->body->type);
+        ++funcIdx;
+      });
+      break;
+    }
+  }
 }
 
 // wasm-split profile format:
@@ -118,17 +144,20 @@ void Instrumenter::addProfileExport() {
   // 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, profileExport);
+  auto name = Names::getValidFunctionName(*wasm, options.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
-  const size_t profileSize = 8 + 4 * functionGlobals.size();
+  const size_t profileSize = 8 + 4 * numFuncs;
 
   // Create the function body
   Builder builder(*wasm);
@@ -142,18 +171,76 @@ void Instrumenter::addProfileExport() {
   // Write the hash followed by all the time stamps
   Expression* writeData =
     builder.makeStore(8, 0, 1, getAddr(), hashConst(), Type::i64);
-
   uint32_t offset = 8;
-  for (const auto& global : functionGlobals) {
-    writeData = builder.blockify(
-      writeData,
-      builder.makeStore(4,
-                        offset,
-                        1,
-                        getAddr(),
-                        builder.makeGlobalGet(global, Type::i32),
-                        Type::i32));
-    offset += 4;
+
+  switch (options.storageKind) {
+    case WasmSplitOptions::StorageKind::InGlobals: {
+      for (const auto& global : functionGlobals) {
+        writeData = builder.blockify(
+          writeData,
+          builder.makeStore(4,
+                            offset,
+                            1,
+                            getAddr(),
+                            builder.makeGlobalGet(global, Type::i32),
+                            Type::i32));
+        offset += 4;
+      }
+      break;
+    }
+    case WasmSplitOptions::StorageKind::InMemory: {
+      Index funcIdxVar =
+        Builder::addVar(writeProfile.get(), "funcIdx", Type::i32);
+      auto getFuncIdx = [&]() {
+        return builder.makeLocalGet(funcIdxVar, Type::i32);
+      };
+      // (block $outer
+      //   (loop $l
+      //     (br_if $outer (i32.eq (local.get $fucIdx) (i32.const numFuncs))
+      //     (i32.store offset=8
+      //       (i32.add
+      //         (local.get $addr)
+      //         (i32.mul (local.get $funcIdx) (i32.const 4))
+      //       )
+      //       (i32.atomic.load8_u (local.get $funcIdx))
+      //     )
+      //     (local.set $funcIdx
+      //      (i32.add (local.get $fundIdx) (i32.const 1)
+      //     )
+      //     (br $l)
+      //   )
+      // )
+      writeData = builder.blockify(
+        writeData,
+        builder.makeBlock(
+          "outer",
+          builder.makeLoop(
+            "l",
+            builder.blockify(
+              builder.makeBreak(
+                "outer",
+                nullptr,
+                builder.makeBinary(EqInt32,
+                                   getFuncIdx(),
+                                   builder.makeConst(uint32_t(numFuncs)))),
+              builder.makeStore(
+                4,
+                offset,
+                4,
+                builder.makeBinary(
+                  AddInt32,
+                  getAddr(),
+                  builder.makeBinary(
+                    MulInt32, getFuncIdx(), builder.makeConst(uint32_t(4)))),
+                builder.makeAtomicLoad(1, 0, getFuncIdx(), Type::i32),
+                Type::i32),
+              builder.makeLocalSet(
+                funcIdxVar,
+                builder.makeBinary(
+                  AddInt32, getFuncIdx(), builder.makeConst(uint32_t(1)))),
+              builder.makeBreak("l")))));
+      break;
+    }
   }
 
   writeProfile->body = builder.makeSequence(
@@ -164,7 +251,7 @@ void Instrumenter::addProfileExport() {
   // Create an export for the function
   wasm->addFunction(std::move(writeProfile));
   wasm->addExport(
-    Builder::makeExport(profileExport, name, ExternalKind::Function));
+    Builder::makeExport(options.profileExport, name, ExternalKind::Function));
 
   // Also make sure there is a memory with enough pages to write into
   size_t pages = (profileSize + Memory::kPageSize - 1) / Memory::kPageSize;
diff --git a/src/tools/wasm-split/instrumenter.h b/src/tools/wasm-split/instrumenter.h
index 4f714fde9..7de5a9135 100644
--- a/src/tools/wasm-split/instrumenter.h
+++ b/src/tools/wasm-split/instrumenter.h
@@ -18,6 +18,8 @@
 #define wasm_tools_wasm_split_instrumenter_h
 
 #include "pass.h"
+#include "split-options.h"
+#include "wasm.h"
 
 namespace wasm {
 
@@ -28,15 +30,17 @@ struct Instrumenter : public Pass {
   PassRunner* runner = nullptr;
   Module* wasm = nullptr;
 
-  const std::string& profileExport;
+  const WasmSplitOptions& options;
   uint64_t moduleHash;
 
   Name counterGlobal;
   std::vector<Name> functionGlobals;
 
-  Instrumenter(const std::string& profileExport, uint64_t moduleHash);
+  Instrumenter(const WasmSplitOptions& options, uint64_t moduleHash);
 
   void run(PassRunner* runner, Module* wasm) override;
+
+private:
   void addGlobals();
   void instrumentFuncs();
   void addProfileExport();
diff --git a/src/tools/wasm-split/split-options.cpp b/src/tools/wasm-split/split-options.cpp
index 419555f45..8de08d5a1 100644
--- a/src/tools/wasm-split/split-options.cpp
+++ b/src/tools/wasm-split/split-options.cpp
@@ -197,6 +197,19 @@ WasmSplitOptions::WasmSplitOptions()
            profileExport = argument;
          })
     .add(
+      "--in-memory",
+      "",
+      "Store profile information in memory (starting at address 0 and taking "
+      "one byte per function) rather than globals (the default) so that "
+      "it can be shared between multiple threads. Users are responsible for "
+      "ensuring that the module does not use the initial memory region for "
+      "anything else.",
+      {Mode::Instrument},
+      Options::Arguments::Zero,
+      [&](Options* o, const std::string& argument) {
+        storageKind = StorageKind::InMemory;
+      })
+    .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 5c811c32c..faa6ee30f 100644
--- a/src/tools/wasm-split/split-options.h
+++ b/src/tools/wasm-split/split-options.h
@@ -33,6 +33,12 @@ struct WasmSplitOptions : ToolOptions {
   constexpr static size_t NumModes =
     static_cast<unsigned>(Mode::MergeProfiles) + 1;
 
+  enum class StorageKind : unsigned {
+    InGlobals, // Store profile data in WebAssembly Globals
+    InMemory,  // Store profile data in memory, accessible from all threads
+  };
+  StorageKind storageKind = StorageKind::InGlobals;
+
   bool verbose = false;
   bool emitBinary = true;
   bool symbolMap = false;
diff --git a/src/tools/wasm-split/wasm-split.cpp b/src/tools/wasm-split/wasm-split.cpp
index bedba8957..b1ae153c9 100644
--- a/src/tools/wasm-split/wasm-split.cpp
+++ b/src/tools/wasm-split/wasm-split.cpp
@@ -109,7 +109,7 @@ void instrumentModule(const WasmSplitOptions& options) {
 
   uint64_t moduleHash = hashFile(options.inputFiles[0]);
   PassRunner runner(&wasm, options.passOptions);
-  Instrumenter(options.profileExport, moduleHash).run(&runner, &wasm);
+  Instrumenter(options, moduleHash).run(&runner, &wasm);
 
   adjustTableSize(wasm, options.initialTableSize);