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.

3 files changed, 129 insertions, 7 deletions

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.