[wasm-split] Initial instrumentation (#3389)

Implement an instrumentation pass that records the timestamp at which each
defined function is first called. Timestamps are not actual time, but rather
snapshots of a monotonically increasing counter. The instrumentation exports a
function that the embedder can call to dump the profile data into a memory
buffer at a given offset and size. The function returns the total size of the
profile data so the embedder can know how much to read out of the buffer or how
much it needs to grow the buffer.

Parsing and using the profile is left as future work, as is recording a hash of
the input file that will be used to guard against accidentally instrumenting one
module and trying to use the resulting profile to split a different module.

1 files changed, 197 insertions, 2 deletions

diff --git a/src/tools/wasm-split.cpp b/src/tools/wasm-split.cpp
index 5f6207e9a..61f52d5fa 100644
--- a/src/tools/wasm-split.cpp
+++ b/src/tools/wasm-split.cpp
@@ -19,10 +19,13 @@
 
 #include "ir/module-splitting.h"
 #include "ir/module-utils.h"
+#include "ir/names.h"
 #include "support/name.h"
 #include "support/utilities.h"
 #include "tool-options.h"
+#include "wasm-builder.h"
 #include "wasm-io.h"
+#include "wasm-type.h"
 #include "wasm-validator.h"
 #include <sstream>
 
@@ -43,6 +46,7 @@ std::set<Name> parseNameList(const std::string& list) {
 
 struct WasmSplitOptions : ToolOptions {
   bool verbose = false;
+  bool emitBinary = true;
 
   bool instrument = false;
 
@@ -162,6 +166,18 @@ WasmSplitOptions::WasmSplitOptions()
            verbose = true;
            quiet = false;
          })
+    .add("--emit-text",
+         "-S",
+         "Emit text instead of binary for the output file or files.",
+         Options::Arguments::Zero,
+         [&](Options* o, const std::string& argument) { emitBinary = false; })
+    .add("--debuginfo",
+         "-g",
+         "Emit names section in wasm binary (or full debuginfo in wast)",
+         Options::Arguments::Zero,
+         [&](Options* o, const std::string& arguments) {
+           passOptions.debugInfo = true;
+         })
     .add_positional(
       "INFILE",
       Options::Arguments::One,
@@ -245,8 +261,186 @@ void parseInput(Module& wasm, const WasmSplitOptions& options) {
   options.applyFeatures(wasm);
 }
 
+// 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.
+struct Instrumenter : public Pass {
+  PassRunner* runner = nullptr;
+  Module* wasm = nullptr;
+
+  const std::string& profileExport;
+  uint64_t moduleHash;
+
+  Name counterGlobal;
+  std::vector<Name> functionGlobals;
+
+  Instrumenter(const std::string& profileExport, uint64_t moduleHash);
+
+  void run(PassRunner* runner, Module* wasm) override;
+  void addGlobals();
+  void instrumentFuncs();
+  void addProfileExport();
+};
+
+Instrumenter::Instrumenter(const std::string& profileExport,
+                           uint64_t moduleHash)
+  : profileExport(profileExport), moduleHash(moduleHash) {}
+
+void Instrumenter::run(PassRunner* runner, Module* wasm) {
+  this->runner = runner;
+  this->wasm = wasm;
+  addGlobals();
+  instrumentFuncs();
+  addProfileExport();
+}
+
+void Instrumenter::addGlobals() {
+  // Create fresh global names (over-reserves, but that's ok)
+  counterGlobal = Names::getValidGlobalName(*wasm, "monotonic_counter");
+  functionGlobals.reserve(wasm->functions.size());
+  ModuleUtils::iterDefinedFunctions(*wasm, [&](Function* func) {
+    functionGlobals.push_back(Names::getValidGlobalName(
+      *wasm, std::string(func->name.c_str()) + "_timestamp"));
+  });
+
+  // Create and add new globals
+  auto addGlobal = [&](Name name) {
+    auto global = Builder::makeGlobal(
+      name,
+      Type::i32,
+      Builder(*wasm).makeConst(Literal::makeZero(Type::i32)),
+      Builder::Mutable);
+    global->hasExplicitName = true;
+    wasm->addGlobal(std::move(global));
+  };
+  addGlobal(counterGlobal);
+  for (auto& name : functionGlobals) {
+    addGlobal(name);
+  }
+}
+
+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)
+  //       )
+  //     )
+  //   )
+  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;
+  });
+}
+
+void Instrumenter::addProfileExport() {
+  // 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, profileExport);
+  auto writeProfile = Builder::makeFunction(
+    name, Signature({Type::i32, Type::i32}, Type::i32), {});
+  writeProfile->hasExplicitName = true;
+  writeProfile->setLocalName(0, "addr");
+  writeProfile->setLocalName(1, "size");
+
+  // 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();
+
+  // Create the function body
+  Builder builder(*wasm);
+  auto getAddr = [&]() { return builder.makeLocalGet(0, Type::i32); };
+  auto getSize = [&]() { return builder.makeLocalGet(1, Type::i32); };
+  auto hashConst = [&]() { return builder.makeConst(int64_t(moduleHash)); };
+  auto profileSizeConst = [&]() {
+    return builder.makeConst(int32_t(profileSize));
+  };
+
+  // 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;
+  }
+
+  writeProfile->body = builder.makeSequence(
+    builder.makeIf(builder.makeBinary(GeUInt32, getSize(), profileSizeConst()),
+                   writeData),
+    profileSizeConst());
+
+  // Create an export for the function
+  wasm->addFunction(std::move(writeProfile));
+  wasm->addExport(
+    Builder::makeExport(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;
+  if (!wasm->memory.exists) {
+    wasm->memory.exists = true;
+    wasm->memory.initial = pages;
+    wasm->memory.max = pages;
+  } else if (wasm->memory.initial < pages) {
+    wasm->memory.initial = pages;
+    if (wasm->memory.max < pages) {
+      wasm->memory.max = pages;
+    }
+  }
+}
+
 void instrumentModule(Module& wasm, const WasmSplitOptions& options) {
-  Fatal() << "TODO: implement instrumentation\n";
+  // Check that the profile export name is not already taken
+  if (wasm.getExportOrNull(options.profileExport) != nullptr) {
+    Fatal() << "error: Export " << options.profileExport << " already exists.";
+  }
+
+  // TODO: calculate module hash.
+  uint64_t moduleHash = 0;
+  PassRunner runner(&wasm, options.passOptions);
+  Instrumenter(options.profileExport, moduleHash).run(&runner, &wasm);
+
+  // Write the output modules
+  ModuleWriter writer;
+  writer.setBinary(options.emitBinary);
+  writer.setDebugInfo(options.passOptions.debugInfo);
+  writer.write(wasm, options.output);
 }
 
 void splitModule(Module& wasm, const WasmSplitOptions& options) {
@@ -337,7 +531,8 @@ void splitModule(Module& wasm, const WasmSplitOptions& options) {
 
   // Write the output modules
   ModuleWriter writer;
-  writer.setBinary(true);
+  writer.setBinary(options.emitBinary);
+  writer.setDebugInfo(options.passOptions.debugInfo);
   writer.write(wasm, options.primaryOutput);
   writer.write(*secondary, options.secondaryOutput);
 }