Add SIMD support to Multi-Memory Lowering Pass (#5336)

This PR adds support for SIMD instructions in the multi-memory lowering pass. Also includes optional bounds checks per the wasm spec guidelines, (SIMDLoad, SIMDLoadSplat, SIMDLoadExtend, SIMDLoadZero, SIMDLoadStoreLane load | store).

2 files changed, 204 insertions, 8 deletions

diff --git a/src/passes/MultiMemoryLowering.cpp b/src/passes/MultiMemoryLowering.cpp
index 57be529fa..452aa21af 100644
--- a/src/passes/MultiMemoryLowering.cpp
+++ b/src/passes/MultiMemoryLowering.cpp
@@ -107,7 +107,8 @@ struct MultiMemoryLowering : public Pass {
       replaceCurrent(builder.makeCall(funcName, {}, curr->type));
     }
 
-    template<typename T> Expression* getPtr(T* curr, Function* func) {
+    template<typename T>
+    Expression* getPtr(T* curr, Function* func, Index bytes) {
       auto memoryIdx = parent.memoryIdxMap.at(curr->memory);
       auto offsetGlobal = parent.getOffsetGlobal(memoryIdx);
       Expression* ptrValue;
@@ -123,7 +124,8 @@ struct MultiMemoryLowering : public Pass {
       if (parent.checkBounds) {
         Index ptrIdx = Builder::addVar(getFunction(), parent.pointerType);
         Expression* ptrSet = builder.makeLocalSet(ptrIdx, ptrValue);
-        Expression* boundsCheck = makeBoundsCheck(curr, ptrIdx, memoryIdx);
+        Expression* boundsCheck =
+          makeBoundsCheck(curr, ptrIdx, memoryIdx, bytes);
         Expression* ptrGet = builder.makeLocalGet(ptrIdx, parent.pointerType);
         return builder.makeBlock({ptrSet, boundsCheck, ptrGet});
       }
@@ -132,7 +134,8 @@ struct MultiMemoryLowering : public Pass {
     }
 
     template<typename T>
-    Expression* makeBoundsCheck(T* curr, Index ptrIdx, Index memoryIdx) {
+    Expression*
+    makeBoundsCheck(T* curr, Index ptrIdx, Index memoryIdx, Index bytes) {
       Name memorySizeFunc = parent.memorySizeNames[memoryIdx];
       Expression* boundsCheck = builder.makeIf(
         builder.makeBinary(
@@ -146,7 +149,7 @@ struct MultiMemoryLowering : public Pass {
               Abstract::getBinary(parent.pointerType, Abstract::Add),
               builder.makeLocalGet(ptrIdx, parent.pointerType),
               builder.makeConstPtr(curr->offset, parent.pointerType)),
-            builder.makeConstPtr(curr->bytes, parent.pointerType)),
+            builder.makeConstPtr(bytes, parent.pointerType)),
           builder.makeCall(memorySizeFunc, {}, parent.pointerType)),
         builder.makeUnreachable());
       return boundsCheck;
@@ -157,12 +160,37 @@ struct MultiMemoryLowering : public Pass {
     }
 
     void visitLoad(Load* curr) {
-      curr->ptr = getPtr(curr, getFunction());
+      curr->ptr = getPtr(curr, getFunction(), curr->bytes);
       setMemory(curr);
     }
 
     void visitStore(Store* curr) {
-      curr->ptr = getPtr(curr, getFunction());
+      curr->ptr = getPtr(curr, getFunction(), curr->bytes);
+      setMemory(curr);
+    }
+
+    void visitSIMDLoad(SIMDLoad* curr) {
+      curr->ptr = getPtr(curr, getFunction(), curr->getMemBytes());
+      setMemory(curr);
+    }
+
+    void visitSIMDLoadSplat(SIMDLoad* curr) {
+      curr->ptr = getPtr(curr, getFunction(), curr->getMemBytes());
+      setMemory(curr);
+    }
+
+    void visitSIMDLoadExtend(SIMDLoad* curr) {
+      curr->ptr = getPtr(curr, getFunction(), curr->getMemBytes());
+      setMemory(curr);
+    }
+
+    void visitSIMDLoadZero(SIMDLoad* curr) {
+      curr->ptr = getPtr(curr, getFunction(), curr->getMemBytes());
+      setMemory(curr);
+    }
+
+    void visitSIMDLoadStoreLane(SIMDLoadStoreLane* curr) {
+      curr->ptr = getPtr(curr, getFunction(), curr->getMemBytes());
       setMemory(curr);
     }
   };
diff --git a/test/lit/passes/multi-memory-lowering.wast b/test/lit/passes/multi-memory-lowering.wast
index d66db1b02..dba7260a0 100644
--- a/test/lit/passes/multi-memory-lowering.wast
+++ b/test/lit/passes/multi-memory-lowering.wast
@@ -1,6 +1,6 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
-;; RUN: wasm-opt %s --enable-multi-memories --multi-memory-lowering --enable-bulk-memory --enable-extended-const -S -o - | filecheck %s
-;; RUN: wasm-opt %s --enable-multi-memories --multi-memory-lowering-with-bounds-checks --enable-bulk-memory --enable-extended-const -S -o - | filecheck %s --check-prefix BOUNDS
+;; RUN: wasm-opt %s --enable-multi-memories --multi-memory-lowering --enable-bulk-memory --enable-extended-const --enable-simd -S -o - | filecheck %s
+;; RUN: wasm-opt %s --enable-multi-memories --multi-memory-lowering-with-bounds-checks --enable-bulk-memory --enable-extended-const --enable-simd -S -o - | filecheck %s --check-prefix BOUNDS
 
 (module
   (memory $memory1 1)
@@ -8,12 +8,16 @@
   (memory $memory3 3)
   (data (memory $memory1) (i32.const 0) "a")
   (data (memory $memory3) (i32.const 1) "123")
+  ;; CHECK:      (type $i32_=>_v128 (func (param i32) (result v128)))
+
   ;; CHECK:      (type $none_=>_i32 (func (result i32)))
 
   ;; CHECK:      (type $i32_=>_i32 (func (param i32) (result i32)))
 
   ;; CHECK:      (type $none_=>_none (func))
 
+  ;; CHECK:      (type $i32_v128_=>_v128 (func (param i32 v128) (result v128)))
+
   ;; CHECK:      (global $memory2_byte_offset (mut i32) (i32.const 65536))
 
   ;; CHECK:      (global $memory3_byte_offset (mut i32) (i32.const 196608))
@@ -50,12 +54,16 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
+  ;; BOUNDS:      (type $i32_=>_v128 (func (param i32) (result v128)))
+
   ;; BOUNDS:      (type $none_=>_i32 (func (result i32)))
 
   ;; BOUNDS:      (type $i32_=>_i32 (func (param i32) (result i32)))
 
   ;; BOUNDS:      (type $none_=>_none (func))
 
+  ;; BOUNDS:      (type $i32_v128_=>_v128 (func (param i32 v128) (result v128)))
+
   ;; BOUNDS:      (global $memory2_byte_offset (mut i32) (i32.const 65536))
 
   ;; BOUNDS:      (global $memory3_byte_offset (mut i32) (i32.const 196608))
@@ -277,6 +285,166 @@
    (i32.const 115)
   )
   )
+
+  ;; CHECK:      (func $v128.load8_splat (param $0 i32) (result v128)
+  ;; CHECK-NEXT:  (v128.load8_splat
+  ;; CHECK-NEXT:   (local.get $0)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; BOUNDS:      (func $v128.load8_splat (param $0 i32) (result v128)
+  ;; BOUNDS-NEXT:  (local $1 i32)
+  ;; BOUNDS-NEXT:  (v128.load8_splat
+  ;; BOUNDS-NEXT:   (block (result i32)
+  ;; BOUNDS-NEXT:    (local.set $1
+  ;; BOUNDS-NEXT:     (local.get $0)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (if
+  ;; BOUNDS-NEXT:     (i32.gt_u
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (i32.add
+  ;; BOUNDS-NEXT:        (local.get $1)
+  ;; BOUNDS-NEXT:        (i32.const 0)
+  ;; BOUNDS-NEXT:       )
+  ;; BOUNDS-NEXT:       (i32.const 1)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (call $memory1_size)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (unreachable)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $1)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT: )
+  (func $v128.load8_splat (param $0 i32) (result v128)
+  (v128.load8_splat $memory1
+   (local.get $0)
+  )
+ )
+
+  ;; CHECK:      (func $v128.load16_lane (param $0 i32) (param $1 v128) (result v128)
+  ;; CHECK-NEXT:  (v128.load16_lane offset=32 align=1 0
+  ;; CHECK-NEXT:   (i32.add
+  ;; CHECK-NEXT:    (global.get $memory2_byte_offset)
+  ;; CHECK-NEXT:    (local.get $0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.get $1)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; BOUNDS:      (func $v128.load16_lane (param $0 i32) (param $1 v128) (result v128)
+  ;; BOUNDS-NEXT:  (local $2 i32)
+  ;; BOUNDS-NEXT:  (v128.load16_lane offset=32 align=1 0
+  ;; BOUNDS-NEXT:   (block (result i32)
+  ;; BOUNDS-NEXT:    (local.set $2
+  ;; BOUNDS-NEXT:     (i32.add
+  ;; BOUNDS-NEXT:      (global.get $memory2_byte_offset)
+  ;; BOUNDS-NEXT:      (local.get $0)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (if
+  ;; BOUNDS-NEXT:     (i32.gt_u
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (i32.add
+  ;; BOUNDS-NEXT:        (local.get $2)
+  ;; BOUNDS-NEXT:        (i32.const 32)
+  ;; BOUNDS-NEXT:       )
+  ;; BOUNDS-NEXT:       (i32.const 2)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (call $memory2_size)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (unreachable)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $2)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:   (local.get $1)
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT: )
+  (func $v128.load16_lane (param $0 i32) (param $1 v128) (result v128)
+  (v128.load16_lane $memory2 align=1 offset=32 0
+   (local.get $0)
+   (local.get $1)
+  )
+ )
+
+  ;; CHECK:      (func $v128.load32_zero (param $0 i32) (result v128)
+  ;; CHECK-NEXT:  (v128.load32_zero offset=16 align=1
+  ;; CHECK-NEXT:   (i32.add
+  ;; CHECK-NEXT:    (global.get $memory3_byte_offset)
+  ;; CHECK-NEXT:    (local.get $0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; BOUNDS:      (func $v128.load32_zero (param $0 i32) (result v128)
+  ;; BOUNDS-NEXT:  (local $1 i32)
+  ;; BOUNDS-NEXT:  (v128.load32_zero offset=16 align=1
+  ;; BOUNDS-NEXT:   (block (result i32)
+  ;; BOUNDS-NEXT:    (local.set $1
+  ;; BOUNDS-NEXT:     (i32.add
+  ;; BOUNDS-NEXT:      (global.get $memory3_byte_offset)
+  ;; BOUNDS-NEXT:      (local.get $0)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (if
+  ;; BOUNDS-NEXT:     (i32.gt_u
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (i32.add
+  ;; BOUNDS-NEXT:        (local.get $1)
+  ;; BOUNDS-NEXT:        (i32.const 16)
+  ;; BOUNDS-NEXT:       )
+  ;; BOUNDS-NEXT:       (i32.const 4)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (call $memory3_size)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (unreachable)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $1)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT: )
+  (func $v128.load32_zero (param $0 i32) (result v128)
+  (v128.load32_zero $memory3 align=1 offset=16
+   (local.get $0)
+  )
+ )
+  ;; CHECK:      (func $v128.load32x2_s (param $0 i32) (result v128)
+  ;; CHECK-NEXT:  (v128.load32x2_s
+  ;; CHECK-NEXT:   (i32.add
+  ;; CHECK-NEXT:    (global.get $memory2_byte_offset)
+  ;; CHECK-NEXT:    (local.get $0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; BOUNDS:      (func $v128.load32x2_s (param $0 i32) (result v128)
+  ;; BOUNDS-NEXT:  (local $1 i32)
+  ;; BOUNDS-NEXT:  (v128.load32x2_s
+  ;; BOUNDS-NEXT:   (block (result i32)
+  ;; BOUNDS-NEXT:    (local.set $1
+  ;; BOUNDS-NEXT:     (i32.add
+  ;; BOUNDS-NEXT:      (global.get $memory2_byte_offset)
+  ;; BOUNDS-NEXT:      (local.get $0)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (if
+  ;; BOUNDS-NEXT:     (i32.gt_u
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (i32.add
+  ;; BOUNDS-NEXT:        (local.get $1)
+  ;; BOUNDS-NEXT:        (i32.const 0)
+  ;; BOUNDS-NEXT:       )
+  ;; BOUNDS-NEXT:       (i32.const 8)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (call $memory2_size)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (unreachable)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $1)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT: )
+  (func $v128.load32x2_s (param $0 i32) (result v128)
+  (v128.load32x2_s $memory2
+   (local.get $0)
+  )
+ )
 )
 
 ;; CHECK:      (func $memory1_size (result i32)