[wasm-split] Read and use profiles (#3400)

Read the profiles produced by wasm-split's instrumentation to guide splitting.
In this initial implementation, all functions that the profile shows to have
been called are kept in the initial module. In the future, users may be able to
tune this so that functions that are run later will still be split out.

4 files changed, 151 insertions, 2 deletions

diff --git a/src/support/file.h b/src/support/file.h
index ce554acaa..3dc766b2d 100644
--- a/src/support/file.h
+++ b/src/support/file.h
@@ -15,7 +15,7 @@
  */
 
 //
-// FIle helpers.
+// File helpers.
 //
 
 #ifndef wasm_support_file_h
diff --git a/src/tools/wasm-split.cpp b/src/tools/wasm-split.cpp
index 61f52d5fa..a363b825b 100644
--- a/src/tools/wasm-split.cpp
+++ b/src/tools/wasm-split.cpp
@@ -20,6 +20,7 @@
 #include "ir/module-splitting.h"
 #include "ir/module-utils.h"
 #include "ir/names.h"
+#include "support/file.h"
 #include "support/name.h"
 #include "support/utilities.h"
 #include "tool-options.h"
@@ -358,6 +359,21 @@ void Instrumenter::instrumentFuncs() {
   });
 }
 
+// wasm-split profile format::
+//
+// The wasm-split profile is a binary format designed to be simple to produce
+// and consume. It is comprised of:
+//
+//   1. An 8-byte module hash
+//
+//   2. A 4-byte timestamp for each defined function
+//
+// The module hash is meant to guard against bugs where the module that was
+// instrumented and the module that is being split are different. The timestamps
+// are non-zero for functions that were called during the instrumented run and 0
+// otherwise. Functions with smaller non-zero timestamps were called earlier in
+// the instrumented run than funtions with larger timestamps.
+
 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
@@ -423,6 +439,8 @@ void Instrumenter::addProfileExport() {
       wasm->memory.max = pages;
     }
   }
+
+  // TODO: export the memory if it is not already exported.
 }
 
 void instrumentModule(Module& wasm, const WasmSplitOptions& options) {
@@ -443,12 +461,51 @@ void instrumentModule(Module& wasm, const WasmSplitOptions& options) {
   writer.write(wasm, options.output);
 }
 
+// See "wasm-split profile format" above for more information.
+std::set<Name> readProfile(Module& wasm, const WasmSplitOptions& options) {
+  auto profileData =
+    read_file<std::vector<char>>(options.profileFile, Flags::Binary);
+  size_t i = 0;
+  auto readi32 = [&]() {
+    if (i + 4 > profileData.size()) {
+      Fatal() << "Unexpected end of profile data";
+    }
+    uint32_t i32 = 0;
+    i32 |= uint32_t(profileData[i++]);
+    i32 |= uint32_t(profileData[i++]) << 8;
+    i32 |= uint32_t(profileData[i++]) << 16;
+    i32 |= uint32_t(profileData[i++]) << 24;
+    return i32;
+  };
+
+  // TODO: Read and compare the 8-byte module hash. Just skip it for now.
+  readi32();
+  readi32();
+
+  std::set<Name> keptFuncs;
+  ModuleUtils::iterDefinedFunctions(wasm, [&](Function* func) {
+    uint32_t timestamp = readi32();
+    // TODO: provide an option to set the timestamp threshold. For now, kee the
+    // function if the profile shows it being run at all.
+    if (timestamp > 0) {
+      keptFuncs.insert(func->name);
+    }
+  });
+
+  if (i != profileData.size()) {
+    // TODO: Handle concatenated profile data.
+    Fatal() << "Unexpected extra profile data";
+  }
+
+  return keptFuncs;
+}
+
 void splitModule(Module& wasm, const WasmSplitOptions& options) {
   std::set<Name> keepFuncs;
 
   if (options.profileFile.size()) {
     // Use the profile to initialize `keepFuncs`
-    Fatal() << "TODO: implement reading profiles\n";
+    keepFuncs = readProfile(wasm, options);
   }
 
   // Add in the functions specified with --keep-funcs
diff --git a/test/lit/wasm-split/call_exports.mjs b/test/lit/wasm-split/call_exports.mjs
new file mode 100644
index 000000000..546564d2e
--- /dev/null
+++ b/test/lit/wasm-split/call_exports.mjs
@@ -0,0 +1,25 @@
+// Instantiates an instrumented module, calls the given exports, then collects
+// its wasm-split profile and writes it to a given file.
+//
+// Usage:
+//
+//  node call_exports.mjs <module> <profile> <export>*
+
+import * as fs from 'fs';
+
+let wasm = process.argv[2];
+let outFile = process.argv[3];
+
+// Create the Wasm instance
+let { _, instance } = await WebAssembly.instantiate(fs.readFileSync(wasm));
+
+// Call the specified exports
+for (let i = 4; i < process.argv.length; i++) {
+  console.log('calling', process.argv[i]);
+  instance.exports[process.argv[i]]();
+}
+
+// Create and read the profile
+let profileSize = instance.exports['__write_profile'](0, 2**32 - 1);
+let profileData = Buffer.from(instance.exports.memory.buffer, 0, profileSize);
+fs.writeFileSync(outFile, profileData);
diff --git a/test/lit/wasm-split/profile-guided.wast b/test/lit/wasm-split/profile-guided.wast
new file mode 100644
index 000000000..f2d79f99c
--- /dev/null
+++ b/test/lit/wasm-split/profile-guided.wast
@@ -0,0 +1,67 @@
+;; Instrument the binary
+
+;; RUN: wasm-split --instrument %s -o %t.instrumented.wasm
+
+;; Create profiles
+
+;; RUN: node %S/call_exports.mjs %t.instrumented.wasm %t.foo.prof foo
+;; RUN: node %S/call_exports.mjs %t.instrumented.wasm %t.bar.prof bar
+;; RUN: node %S/call_exports.mjs %t.instrumented.wasm %t.both.prof foo bar
+;; RUN: node %S/call_exports.mjs %t.instrumented.wasm %t.none.prof
+
+;; Create profile-guided splits
+
+;; RUN: wasm-split %s --profile=%t.foo.prof -v -o1 %t.foo.1.wasm -o2 %t.foo.2.wasm \
+;; RUN:   | filecheck %s --check-prefix FOO
+
+;; FOO: Keeping functions: deep_foo_callee, foo, foo_callee, shared_callee
+;; FOO: Splitting out functions: bar, bar_callee, uncalled
+
+;; RUN: wasm-split %s --profile=%t.bar.prof -v -o1 %t.bar.1.wasm -o2 %t.bar.2.wasm \
+;; RUN:   | filecheck %s --check-prefix BAR
+
+;; BAR: Keeping functions: bar, bar_callee, shared_callee
+;; BAR: Splitting out functions: deep_foo_callee, foo, foo_callee, uncalled
+
+;; RUN: wasm-split %s --profile=%t.both.prof -v -o1 %t.both.1.wasm -o2 %t.both.2.wasm \
+;; RUN:   | filecheck %s --check-prefix BOTH
+
+;; BOTH: Keeping functions: bar, bar_callee, deep_foo_callee, foo, foo_callee, shared_callee
+;; BOTH: Splitting out functions: uncalled
+
+;; RUN: wasm-split %s --profile=%t.none.prof -v -o1 %t.none.1.wasm -o2 %t.none.2.wasm \
+;; RUN:   | filecheck %s --check-prefix NONE
+
+;; NONE: Keeping functions:
+;; NONE: Splitting out functions: bar, bar_callee, deep_foo_callee, foo, foo_callee, shared_callee, uncalled
+
+
+(module
+  (memory $mem 1 1)
+  (export "memory" (memory $mem))
+  (export "foo" (func $foo))
+  (export "bar" (func $bar))
+  (export "uncalled" (func $uncalled))
+
+  (func $foo
+    (call $foo_callee)
+    (call $shared_callee)
+  )
+
+  (func $bar
+    (call $shared_callee)
+    (call $bar_callee)
+  )
+
+  (func $uncalled)
+
+  (func $foo_callee
+    (call $deep_foo_callee)
+  )
+
+  (func $bar_callee)
+
+  (func $shared_callee)
+
+  (func $deep_foo_callee)
+)