diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td @@ -862,6 +862,37 @@ let assemblyFormat = "$matrix attr-dict `:` type($matrix) `into` type($res)"; } +// +// LLVM masked operations. +// + +/// Create a call to Masked Load intrinsic. +def LLVM_MaskedLoadOp + : LLVM_OneResultOp<"intr.masked.load">, + Arguments<(ins LLVM_Type:$data, LLVM_Type:$mask, + Variadic:$pass_thru, I32Attr:$alignment)> { + string llvmBuilder = [{ + $res = $pass_thru.empty() ? builder.CreateMaskedLoad( + $data, llvm::Align($alignment.getZExtValue()), $mask) : + builder.CreateMaskedLoad( + $data, llvm::Align($alignment.getZExtValue()), $mask, $pass_thru[0]); + }]; + let assemblyFormat = + "operands attr-dict `:` functional-type(operands, results)"; +} + +/// Create a call to Masked Store intrinsic. +def LLVM_MaskedStoreOp + : LLVM_ZeroResultOp<"intr.masked.store">, + Arguments<(ins LLVM_Type:$value, LLVM_Type:$data, LLVM_Type:$mask, + I32Attr:$alignment)> { + string llvmBuilder = [{ + builder.CreateMaskedStore( + $value, $data, llvm::Align($alignment.getZExtValue()), $mask); + }]; + let assemblyFormat = "$value `,` $data `,` $mask attr-dict `:` " + "type($value) `,` type($mask) `into` type($data)"; +} // // Atomic operations. // diff --git a/mlir/test/Target/llvmir-intrinsics.mlir b/mlir/test/Target/llvmir-intrinsics.mlir --- a/mlir/test/Target/llvmir-intrinsics.mlir +++ b/mlir/test/Target/llvmir-intrinsics.mlir @@ -152,6 +152,20 @@ llvm.return } +// CHECK-LABEL: @masked_intrinsics +llvm.func @masked_intrinsics(%A: !llvm<"<7 x float>*">, %mask: !llvm<"<7 x i1>">) { + // CHECK: call <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>* %{{.*}}, i32 1, <7 x i1> %{{.*}}, <7 x float> undef) + %a = llvm.intr.masked.load %A, %mask { alignment = 1: i32} : + (!llvm<"<7 x float>*">, !llvm<"<7 x i1>">) -> !llvm<"<7 x float>"> + // CHECK: call <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>* %{{.*}}, i32 1, <7 x i1> %{{.*}}, <7 x float> %{{.*}}) + %b = llvm.intr.masked.load %A, %mask, %a { alignment = 1: i32} : + (!llvm<"<7 x float>*">, !llvm<"<7 x i1>">, !llvm<"<7 x float>">) -> !llvm<"<7 x float>"> + // CHECK: call void @llvm.masked.store.v7f32.p0v7f32(<7 x float> %{{.*}}, <7 x float>* %0, i32 {{.*}}, <7 x i1> %{{.*}}) + llvm.intr.masked.store %b, %A, %mask { alignment = 1: i32} : + !llvm<"<7 x float>">, !llvm<"<7 x i1>"> into !llvm<"<7 x float>*"> + 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 @@ -179,3 +193,5 @@ // CHECK-DAG: declare <48 x float> @llvm.matrix.transpose.v48f32(<48 x float>, i32 immarg, i32 immarg) // CHECK-DAG: declare <48 x float> @llvm.matrix.columnwise.load.v48f32.p0f32(float*, i32, i32 immarg, i32 immarg) // CHECK-DAG: declare void @llvm.matrix.columnwise.store.v48f32.p0f32(<48 x float>, float* writeonly, i32, i32 immarg, i32 immarg) +// CHECK-DAG: declare <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>*, i32 immarg, <7 x i1>, <7 x float>) +// CHECK-DAG: declare void @llvm.masked.store.v7f32.p0v7f32(<7 x float>, <7 x float>*, i32 immarg, <7 x i1>)