[mlir][LLVMIR] Add more vector predication intrinsic ops.

[jacquesguan]

Summary

This revision adds vector predication select, merge and load/store intrinsic ops.

[jacquesguan]

rebase

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

class LLVM_VPReductionBase<string mnem, Type element>
     : LLVM_OneResultIntrOp<"vp.reduce." # mnem, [], [1], [NoSideEffect]>,
       Arguments<(ins element:$satrt_value, LLVM_VectorOf<element>:$val,
                      LLVM_VectorOf<I1>:$mask, I32:$evl)>;
 
 class LLVM_VPReductionI<string mnem> : LLVM_VPReductionBase<mnem, AnyInteger>;
 
 class LLVM_VPReductionF<string mnem> : LLVM_VPReductionBase<mnem, AnyFloat>;
 
+class LLVM_VPSelectBase<string mnem>
+    : LLVM_OneResultIntrOp<"vp." # mnem, [], [1], [NoSideEffect]>,
+      Arguments<(ins LLVM_VectorOf<I1>:$cond, LLVM_AnyVector:$true_val,
+                     LLVM_AnyVector:$false_val, I32:$evl)>;
+
 // Integer Binary
 def LLVM_VPAddOp  : LLVM_VPBinaryI<"add">;
 def LLVM_VPSubOp  : LLVM_VPBinaryI<"sub">;
 def LLVM_VPMulOp  : LLVM_VPBinaryI<"mul">;
 def LLVM_VPSDivOp : LLVM_VPBinaryI<"sdiv">;
 def LLVM_VPUDivOp : LLVM_VPBinaryI<"udiv">;
 def LLVM_VPSRemOp : LLVM_VPBinaryI<"srem">;
 def LLVM_VPURemOp : LLVM_VPBinaryI<"urem">;

 def LLVM_VPReduceUMinOp : LLVM_VPReductionI<"umin">;
 
 // Float Reduction
 def LLVM_VPReduceFAddOp : LLVM_VPReductionF<"fadd">;
 def LLVM_VPReduceFMulOp : LLVM_VPReductionF<"fmul">;
 def LLVM_VPReduceFMaxOp : LLVM_VPReductionF<"fmax">;
 def LLVM_VPReduceFMinOp : LLVM_VPReductionF<"fmin">;
 
+def LLVM_VPSelectMinOp : LLVM_VPSelectBase<"select">;
+def LLVM_VPMergeMinOp  : LLVM_VPSelectBase<"merge">;
+
+// Load/store
+def LLVM_VPLoadOp
+    : LLVM_OneResultIntrOp<"vp.load", [0], [0], []>,
+      Arguments<(ins LLVM_AnyPointer:$ptr,
+                     LLVM_VectorOf<I1>:$mask, I32:$evl)>;
+
+def LLVM_VPStoreOp
+    : LLVM_ZeroResultIntrOp<"vp.store", [0, 1], []>,
+      Arguments<(ins LLVM_AnyVector:$val,
+                     LLVM_AnyPointer:$ptr,
+                     LLVM_VectorOf<I1>:$mask, I32:$evl)>;
+
+// Strided load/store
+def LLVM_VPStridedLoadOp
+    : LLVM_OneResultIntrOp<"experimental.vp.strided.load", [0], [0, 1], []>,
+      Arguments<(ins LLVM_AnyPointer:$ptr, AnyInteger:$stride,
+                     LLVM_VectorOf<I1>:$mask, I32:$evl)>;
+
+def LLVM_VPStridedStoreOp
+    : LLVM_ZeroResultIntrOp<"experimental.vp.strided.store",[0, 1, 2], []>,
+      Arguments<(ins LLVM_AnyVector:$val, LLVM_AnyPointer:$ptr,
+                     AnyInteger:$stride, LLVM_VectorOf<I1>:$mask, I32:$evl)>;
+
 #endif // LLVM_INTRINSIC_OPS
 No newline at end of file

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

llvm.func @stack_restore(%arg0: !llvm.ptr<i8>) {
   llvm.intr.stackrestore %arg0
   llvm.return
 }
 
 // CHECK-LABEL: @vector_predication_intrinsics
 llvm.func @vector_predication_intrinsics(%A: vector<8xi32>, %B: vector<8xi32>,
                                          %C: vector<8xf32>, %D: vector<8xf32>,
                                          %i: i32, %f: f32,
+                                         %iptr : !llvm.ptr<i32>,
+                                         %fptr : !llvm.ptr<f32>,
                                          %mask: vector<8xi1>, %evl: i32) {
   // CHECK: call <8 x i32> @llvm.vp.add.v8i32
   "llvm.intr.vp.add" (%A, %B, %mask, %evl) :
          (vector<8xi32>, vector<8xi32>, vector<8xi1>, i32) -> vector<8xi32>
   // CHECK: call <8 x i32> @llvm.vp.sub.v8i32
   "llvm.intr.vp.sub" (%A, %B, %mask, %evl) :
          (vector<8xi32>, vector<8xi32>, vector<8xi1>, i32) -> vector<8xi32>
   // CHECK: call <8 x i32> @llvm.vp.mul.v8i32
"llvm.intr.vp.reduce.fmul" (%f, %C, %mask, %evl) :
          (f32, vector<8xf32>, vector<8xi1>, i32) -> f32
   // CHECK: call float @llvm.vp.reduce.fmax.v8f32
   "llvm.intr.vp.reduce.fmax" (%f, %C, %mask, %evl) :
          (f32, vector<8xf32>, vector<8xi1>, i32) -> f32
   // CHECK: call float @llvm.vp.reduce.fmin.v8f32
   "llvm.intr.vp.reduce.fmin" (%f, %C, %mask, %evl) :
          (f32, vector<8xf32>, vector<8xi1>, i32) -> f32
 
+  // CHECK: call <8 x i32> @llvm.vp.select.v8i32
+  "llvm.intr.vp.select" (%mask, %A, %B, %evl) :
+         (vector<8xi1>, vector<8xi32>, vector<8xi32>, i32) -> vector<8xi32>
+  // CHECK: call <8 x i32> @llvm.vp.merge.v8i32
+  "llvm.intr.vp.merge" (%mask, %A, %B, %evl) :
+         (vector<8xi1>, vector<8xi32>, vector<8xi32>, i32) -> vector<8xi32>
+
+  // CHECK: call void @llvm.vp.store.v8i32.p0i32
+  "llvm.intr.vp.store" (%A, %iptr, %mask, %evl) : 
+         (vector<8xi32>, !llvm.ptr<i32>, vector<8xi1>, i32) -> ()
+  // CHECK: call <8 x i32> @llvm.vp.load.v8i32.p0i32
+  "llvm.intr.vp.load" (%iptr, %mask, %evl) :
+         (!llvm.ptr<i32>, vector<8xi1>, i32) -> vector<8xi32>
+  // CHECK: call void @llvm.experimental.vp.strided.store.v8i32.p0i32.i32
+  "llvm.intr.experimental.vp.strided.store" (%A, %iptr, %i, %mask, %evl) : 
+         (vector<8xi32>, !llvm.ptr<i32>, i32, vector<8xi1>, i32) -> ()
+  // CHECK: call <8 x i32> @llvm.experimental.vp.strided.load.v8i32.p0i32.i32
+  "llvm.intr.experimental.vp.strided.load" (%iptr, %i, %mask, %evl) :
+         (!llvm.ptr<i32>, i32, vector<8xi1>, i32) -> vector<8xi32>
 
   llvm.return
 }
 
 // Check that intrinsics are declared with appropriate types.
 // CHECK-DAG: declare float @llvm.fma.f32(float, float, float)
 // CHECK-DAG: declare <8 x float> @llvm.fma.v8f32(<8 x float>, <8 x float>, <8 x float>) #0
 // CHECK-DAG: declare float @llvm.fmuladd.f32(float, float, float)

 // CHECK-DAG: declare i32 @llvm.vp.reduce.smax.v8i32(i32, <8 x i32>, <8 x i1>, i32) #0
 // CHECK-DAG: declare i32 @llvm.vp.reduce.smin.v8i32(i32, <8 x i32>, <8 x i1>, i32) #0
 // CHECK-DAG: declare i32 @llvm.vp.reduce.umax.v8i32(i32, <8 x i32>, <8 x i1>, i32) #0
 // CHECK-DAG: declare i32 @llvm.vp.reduce.umin.v8i32(i32, <8 x i32>, <8 x i1>, i32) #0
 // CHECK-DAG: declare float @llvm.vp.reduce.fadd.v8f32(float, <8 x float>, <8 x i1>, i32) #0
 // CHECK-DAG: declare float @llvm.vp.reduce.fmul.v8f32(float, <8 x float>, <8 x i1>, i32) #0
 // CHECK-DAG: declare float @llvm.vp.reduce.fmax.v8f32(float, <8 x float>, <8 x i1>, i32) #0
 // CHECK-DAG: declare float @llvm.vp.reduce.fmin.v8f32(float, <8 x float>, <8 x i1>, i32) #0
+// CHECK-DAG: declare <8 x i32> @llvm.vp.select.v8i32(<8 x i1>, <8 x i32>, <8 x i32>, i32) #12
+// CHECK-DAG: declare <8 x i32> @llvm.vp.merge.v8i32(<8 x i1>, <8 x i32>, <8 x i32>, i32) #12
+// CHECK-DAG: declare void @llvm.experimental.vp.strided.store.v8i32.p0i32.i32(<8 x i32>, i32* nocapture, i32, <8 x i1>, i32) #4
+// CHECK-DAG: declare <8 x i32> @llvm.experimental.vp.strided.load.v8i32.p0i32.i32(i32* nocapture, i32, <8 x i1>, i32) #3