Multi-Memories wasm-split (#4977)

Adds an --in-secondary-memory switch to the wasm-split tool that allows profile data to be stored in a separate memory from module main memory. 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.

4 files changed, 148 insertions, 10 deletions

diff --git a/test/lit/help/wasm-split.test b/test/lit/help/wasm-split.test
index 1994c681a..2e798b077 100644
--- a/test/lit/help/wasm-split.test
+++ b/test/lit/help/wasm-split.test
@@ -54,9 +54,13 @@
 ;; CHECK-NEXT:   --placeholdermap                    [split] Write a file mapping placeholder
 ;; CHECK-NEXT:                                       indices to the function names.
 ;; CHECK-NEXT:
-;; CHECK-NEXT:   --import-namespace                  [split] The namespace from which to import
-;; CHECK-NEXT:                                       objects from the primary module into the
-;; CHECK-NEXT:                                       secondary module.
+;; CHECK-NEXT:   --import-namespace                  [split, instrument] When provided as an
+;; CHECK-NEXT:                                       option for module splitting, the namespace
+;; CHECK-NEXT:                                       from which to import objects from the
+;; CHECK-NEXT:                                       primary module into the secondary module.
+;; CHECK-NEXT:                                       In instrument mode, refers to the
+;; CHECK-NEXT:                                       namespace from which to import the
+;; CHECK-NEXT:                                       secondary memory, if any.
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --placeholder-namespace             [split] The namespace from which to import
 ;; CHECK-NEXT:                                       placeholder functions into the primary
@@ -85,6 +89,18 @@
 ;; CHECK-NEXT:                                       does not use the initial memory region for
 ;; CHECK-NEXT:                                       anything else.
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --in-secondary-memory               [instrument] Store profile information in
+;; CHECK-NEXT:                                       a separate memory, rather than in module
+;; CHECK-NEXT:                                       main memory or globals (the default). With
+;; CHECK-NEXT:                                       this option, users do not need to reserve
+;; CHECK-NEXT:                                       the initial memory region for profile data
+;; CHECK-NEXT:                                       and the data can be shared between
+;; CHECK-NEXT:                                       multiple threads.
+;; CHECK-NEXT:
+;; CHECK-NEXT:   --secondary-memory-name             [instrument] The name of the secondary
+;; CHECK-NEXT:                                       memory created to store profile
+;; CHECK-NEXT:                                       information.
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --emit-module-names                 [split, instrument] Emit module names,
 ;; CHECK-NEXT:                                       even if not emitting the rest of the names
 ;; CHECK-NEXT:                                       section. Can help differentiate the
diff --git a/test/lit/wasm-split/instrument-in-secondary-memory-custom-names.wast b/test/lit/wasm-split/instrument-in-secondary-memory-custom-names.wast
new file mode 100644
index 000000000..9035a0b8b
--- /dev/null
+++ b/test/lit/wasm-split/instrument-in-secondary-memory-custom-names.wast
@@ -0,0 +1,31 @@
+;; RUN: wasm-split %s --instrument --in-secondary-memory --import-namespace=custom_env --secondary-memory-name=custom_name -all -S -o - | filecheck %s
+
+;; Check that the output round trips and validates as well
+;; RUN: wasm-split %s --instrument --in-secondary-memory -all -g -o %t.wasm
+;; RUN: wasm-opt -all %t.wasm -S -o -
+
+(module
+  (import "env" "foo" (func $foo))
+  (export "bar" (func $bar))
+  (memory $0 1 1)
+  (func $bar
+    (call $foo)
+  )
+  (func $baz (param i32) (result i32)
+    (local.get 0)
+  )
+)
+
+;; Check that a memory import has been added for secondary memory
+;; CHECK: (import "custom_env" "custom_name" (memory $custom_name (shared 1 1)))
+
+;; And the profiling function exported
+;; CHECK: (export "__write_profile" (func $__write_profile))
+
+;; And main memory has been exported
+;; CHECK: (export "profile-memory" (memory $0))
+
+;; Check that the function instrumentation uses the correct memory name
+;; CHECK:  (i32.atomic.store8 $custom_name
+;; CHECK:  (i32.atomic.store8 $custom_name offset=1
+;; CHECK:  (i32.atomic.load8_u $custom_name
diff --git a/test/lit/wasm-split/instrument-in-secondary-memory.wast b/test/lit/wasm-split/instrument-in-secondary-memory.wast
new file mode 100644
index 000000000..3e012293b
--- /dev/null
+++ b/test/lit/wasm-split/instrument-in-secondary-memory.wast
@@ -0,0 +1,92 @@
+;; RUN: wasm-split %s --instrument --in-secondary-memory -all -S -o - | filecheck %s
+
+;; Check that the output round trips and validates as well
+;; RUN: wasm-split %s --instrument --in-secondary-memory -all -g -o %t.wasm
+;; RUN: wasm-opt -all %t.wasm -S -o -
+
+(module
+  (import "env" "foo" (func $foo))
+  (export "bar" (func $bar))
+  (memory $0 1 1)
+  (func $bar
+    (call $foo)
+  )
+  (func $baz (param i32) (result i32)
+    (local.get 0)
+  )
+)
+
+;; Check that a memory import has been added for secondary memory
+;; CHECK: (import "env" "profile-data" (memory $profile-data (shared 1 1)))
+
+;; And the profiling function exported
+;; CHECK: (export "__write_profile" (func $__write_profile))
+
+;; And main memory has been exported
+;; CHECK: (export "profile-memory" (memory $0))
+
+;; Check that the function instrumentation is correct
+
+;; CHECK:      (func $bar
+;; CHECK-NEXT:  (i32.atomic.store8 $profile-data
+;; 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 $profile-data 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 $0 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 $0 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 $profile-data
+;; 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/invalid-options.wast b/test/lit/wasm-split/invalid-options.wast
index 2f9c0a148..c5a56b579 100644
--- a/test/lit/wasm-split/invalid-options.wast
+++ b/test/lit/wasm-split/invalid-options.wast
@@ -17,14 +17,9 @@
 ;; RUN: not wasm-split %s --instrument --symbolmap 2>&1 \
 ;; RUN:   | filecheck %s --check-prefix INSTRUMENT-SYMBOLMAP
 
-;; --instrument cannot be used with --import-namespace
-;; RUN: not wasm-split %s --instrument --import-namespace=foo 2>&1 \
-;; RUN:   | filecheck %s --check-prefix INSTRUMENT-IMPORT-NS
-
 ;; --instrument cannot be used with --placeholder-namespace
 ;; RUN: not wasm-split %s --instrument --placeholder-namespace=foo 2>&1 \
 ;; RUN:   | filecheck %s --check-prefix INSTRUMENT-PLACEHOLDER-NS
-
 ;; --instrument cannot be used with --export-prefix
 ;; RUN: not wasm-split %s --instrument --export-prefix=foo 2>&1 \
 ;; RUN:   | filecheck %s --check-prefix INSTRUMENT-EXPORT-PREFIX
@@ -45,6 +40,10 @@
 ;; RUN: not wasm-split %s --profile-export=foo 2>&1 \
 ;; RUN:   | filecheck %s --check-prefix SPLIT-PROFILE-EXPORT
 
+;; --secondary-memory-name cannot be used with Split mode
+;; RUN: not wasm-split %s --secondary-memory-name=foo 2>&1 \
+;; RUN:   | filecheck %s --check-prefix SPLIT-SECONDARY-MEMORY-NAME
+
 ;; -S cannot be used with --merge-profiles
 ;; RUN: not wasm-split %s --merge-profiles -S 2>&1 \
 ;; RUN:   | filecheck %s --check-prefix MERGE-EMIT-TEXT
@@ -73,8 +72,6 @@
 
 ;; INSTRUMENT-SYMBOLMAP: error: Option --symbolmap cannot be used in instrument mode.
 
-;; INSTRUMENT-IMPORT-NS: error: Option --import-namespace cannot be used in instrument mode.
-
 ;; INSTRUMENT-PLACEHOLDER-NS: error: Option --placeholder-namespace cannot be used in instrument mode.
 
 ;; INSTRUMENT-EXPORT-PREFIX: error: Option --export-prefix cannot be used in instrument mode.
@@ -87,6 +84,8 @@
 
 ;; SPLIT-PROFILE-EXPORT: error: Option --profile-export cannot be used in split mode.
 
+;; SPLIT-SECONDARY-MEMORY-NAME: error: Option --secondary-memory-name cannot be used in split mode.
+
 ;; MERGE-EMIT-TEXT: error: Option --emit-text cannot be used in merge-profiles mode.
 
 ;; MERGE-DEBUGINFO: error: Option --debuginfo cannot be used in merge-profiles mode.