[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.

4 files changed, 150 insertions, 21 deletions

diff --git a/test/lit/help/wasm-split.test b/test/lit/help/wasm-split.test
index 362c20d84..dd2b5f7e1 100644
--- a/test/lit/help/wasm-split.test
+++ b/test/lit/help/wasm-split.test
@@ -39,7 +39,7 @@
 ;; CHECK-NEXT:                                        of any profile. You can also pass a file
 ;; CHECK-NEXT:                                        with a list of functions separated by new
 ;; CHECK-NEXT:                                        lines. To do this, prepend @ before
-;; CHECK-NEXT:                                        filename (--keep-funcs @myfile) 
+;; CHECK-NEXT:                                        filename (--keep-funcs @myfile)
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --split-funcs                        [split] Comma-separated list of functions
 ;; CHECK-NEXT:                                        to split into the secondary module,
@@ -49,7 +49,7 @@
 ;; CHECK-NEXT:                                        also pass a file with a list of functions
 ;; CHECK-NEXT:                                        separated by new lines. To do this,
 ;; CHECK-NEXT:                                        prepend @ before filename (--split-funcs
-;; CHECK-NEXT:                                        @myfile) 
+;; CHECK-NEXT:                                        @myfile)
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --primary-output,-o1                 [split] Output file for the primary
 ;; CHECK-NEXT:                                        module.
@@ -80,6 +80,15 @@
 ;; CHECK-NEXT:                                        profile into memory. Defaults to
 ;; CHECK-NEXT:                                        `__write_profile`.
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --in-memory                          [instrument] Store profile information in
+;; CHECK-NEXT:                                        memory (starting at address 0 and taking
+;; CHECK-NEXT:                                        one byte per function) rather than
+;; CHECK-NEXT:                                        globals (the default) so that it can be
+;; CHECK-NEXT:                                        shared between multiple threads. Users
+;; CHECK-NEXT:                                        are responsible for ensuring that the
+;; CHECK-NEXT:                                        module does not use the initial memory
+;; CHECK-NEXT:                                        region for anything else.
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --emit-module-names                  [split, instrument] Emit module names,
 ;; CHECK-NEXT:                                        even if not emitting the rest of the
 ;; CHECK-NEXT:                                        names section. Can help differentiate the
diff --git a/test/lit/wasm-split/call_exports.mjs b/test/lit/wasm-split/call_exports.mjs
index 546564d2e..592fbad48 100644
--- a/test/lit/wasm-split/call_exports.mjs
+++ b/test/lit/wasm-split/call_exports.mjs
@@ -20,6 +20,6 @@ for (let i = 4; i < process.argv.length; 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);
+let profileSize = instance.exports['__write_profile'](1024, 2**32 - 1024);
+let profileData = Buffer.from(instance.exports.memory.buffer, 1024, profileSize);
 fs.writeFileSync(outFile, profileData);
diff --git a/test/lit/wasm-split/instrument-in-memory.wast b/test/lit/wasm-split/instrument-in-memory.wast
new file mode 100644
index 000000000..568ce2207
--- /dev/null
+++ b/test/lit/wasm-split/instrument-in-memory.wast
@@ -0,0 +1,88 @@
+;; RUN: wasm-split %s --instrument --in-memory -all -S -o - | filecheck %s
+
+;; Check that the output round trips and validates as well
+;; RUN: wasm-split %s --instrument --in-memory -all -g -o %t.wasm
+;; RUN: wasm-opt -all %t.wasm -S -o -
+
+(module
+  (import "env" "foo" (func $foo))
+  (export "bar" (func $bar))
+  (func $bar
+    (call $foo)
+  )
+  (func $baz (param i32) (result i32)
+    (local.get 0)
+  )
+)
+
+;; Check that a memory has been added
+;; CHECK: (memory $0 1 1)
+
+;; And the profiling function exported
+;; CHECK: (export "__write_profile" (func $__write_profile))
+
+;; Check that the function instrumentation is correct
+
+;; CHECK:      (func $bar
+;; CHECK-NEXT:  (i32.atomic.store8
+;; CHECK-NEXT:   (i32.const 0)
+;; CHECK-NEXT:   (i32.const 1)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (call $foo)
+;; CHECK-NEXT: )
+
+;; CHECK-NEXT: (func $baz (param $0 i32) (result i32)
+;; CHECK-NEXT:  (i32.atomic.store8 offset=1
+;; CHECK-NEXT:   (i32.const 0)
+;; CHECK-NEXT:   (i32.const 1)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (local.get $0)
+;; CHECK-NEXT: )
+
+;; Check that the profiling function is correct.
+
+;; CHECK:      (func $__write_profile (param $addr i32) (param $size i32) (result i32)
+;; CHECK-NEXT:  (local $funcIdx i32)
+;; CHECK-NEXT:  (if
+;; CHECK-NEXT:   (i32.ge_u
+;; CHECK-NEXT:    (local.get $size)
+;; CHECK-NEXT:    (i32.const 16)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:   (block
+;; CHECK-NEXT:    (i64.store align=1
+;; CHECK-NEXT:     (local.get $addr)
+;; CHECK-NEXT:     (i64.const {{.*}})
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:    (block $outer
+;; CHECK-NEXT:     (loop $l
+;; CHECK-NEXT:      (br_if $outer
+;; CHECK-NEXT:       (i32.eq
+;; CHECK-NEXT:        (local.get $funcIdx)
+;; CHECK-NEXT:        (i32.const 2)
+;; CHECK-NEXT:       )
+;; CHECK-NEXT:      )
+;; CHECK-NEXT:      (i32.store offset=8
+;; CHECK-NEXT:       (i32.add
+;; CHECK-NEXT:        (local.get $addr)
+;; CHECK-NEXT:        (i32.mul
+;; CHECK-NEXT:         (local.get $funcIdx)
+;; CHECK-NEXT:         (i32.const 4)
+;; CHECK-NEXT:        )
+;; CHECK-NEXT:       )
+;; CHECK-NEXT:       (i32.atomic.load8_u
+;; CHECK-NEXT:        (local.get $funcIdx)
+;; CHECK-NEXT:       )
+;; CHECK-NEXT:      )
+;; CHECK-NEXT:      (local.set $funcIdx
+;; CHECK-NEXT:       (i32.add
+;; CHECK-NEXT:        (local.get $funcIdx)
+;; CHECK-NEXT:        (i32.const 1)
+;; CHECK-NEXT:       )
+;; CHECK-NEXT:      )
+;; CHECK-NEXT:      (br $l)
+;; CHECK-NEXT:     )
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (i32.const 16)
+;; CHECK-NEXT: )
diff --git a/test/lit/wasm-split/profile-guided.wast b/test/lit/wasm-split/profile-guided.wast
index f2d79f99c..4a96ef55e 100644
--- a/test/lit/wasm-split/profile-guided.wast
+++ b/test/lit/wasm-split/profile-guided.wast
@@ -1,43 +1,75 @@
+;; =====================
 ;; Instrument the binary
+;; =====================
 
-;; RUN: wasm-split --instrument %s -o %t.instrumented.wasm
+;; RUN: wasm-split -all --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
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.foo.prof foo
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.bar.prof bar
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.both.prof foo bar
+;; RUN: node --experimental-wasm-threads %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: wasm-split -all %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 -all %s --profile=%t.bar.prof -v -o1 %t.bar.1.wasm -o2 %t.bar.2.wasm \
+;; RUN:   | filecheck %s --check-prefix BAR
+
+;; RUN: wasm-split -all %s --profile=%t.both.prof -v -o1 %t.both.1.wasm -o2 %t.both.2.wasm \
+;; RUN:   | filecheck %s --check-prefix BOTH
+
+;; RUN: wasm-split -all %s --profile=%t.none.prof -v -o1 %t.none.1.wasm -o2 %t.none.2.wasm \
+;; RUN:   | filecheck %s --check-prefix NONE
+
+;; =================================
+;; Do it all again using --in-memory
+;; =================================
 
-;; RUN: wasm-split %s --profile=%t.bar.prof -v -o1 %t.bar.1.wasm -o2 %t.bar.2.wasm \
+;; RUN: wasm-split -all --instrument --in-memory %s -o %t.instrumented.wasm
+
+;; Create profiles
+
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.foo.prof foo
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.bar.prof bar
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.both.prof foo bar
+;; RUN: node --experimental-wasm-threads %S/call_exports.mjs %t.instrumented.wasm %t.none.prof
+
+;; Create profile-guided splits
+
+;; RUN: wasm-split -all %s --profile=%t.foo.prof -v -o1 %t.foo.1.wasm -o2 %t.foo.2.wasm \
+;; RUN:   | filecheck %s --check-prefix FOO
+
+;; RUN: wasm-split -all %s --profile=%t.bar.prof -v -o1 %t.bar.1.wasm -o2 %t.bar.2.wasm \
 ;; RUN:   | filecheck %s --check-prefix BAR
 
+;; RUN: wasm-split -all %s --profile=%t.both.prof -v -o1 %t.both.1.wasm -o2 %t.both.2.wasm \
+;; RUN:   | filecheck %s --check-prefix BOTH
+
+;; RUN: wasm-split -all %s --profile=%t.none.prof -v -o1 %t.none.1.wasm -o2 %t.none.2.wasm \
+;; RUN:   | filecheck %s --check-prefix NONE
+
+;; =======
+;; Results
+;; =======
+
+;; FOO: Keeping functions: deep_foo_callee, foo, foo_callee, shared_callee
+;; FOO: Splitting out functions: bar, bar_callee, uncalled
+
 ;; 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)
+  (memory $mem (shared 1 1))
   (export "memory" (memory $mem))
   (export "foo" (func $foo))
   (export "bar" (func $bar))