[mlir][LLVM] Add masked reads and writes.

Summary: These intrinsics will be used to lower vector transfer read/write.

Reviewers: aartbik, tetuante, jsetoain

Reviewed By: aartbik

Subscribers: mehdi_amini, rriddle, jpienaar, burmako, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, liufengdb, Joonsoo, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D75986
This commit is contained in:
Nicolas Vasilache 2020-03-18 14:55:24 -04:00
parent b40ee7ff1b
commit 1e43132209
2 changed files with 47 additions and 0 deletions

View File

@ -862,6 +862,37 @@ def LLVM_MatrixTranposeOp
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<LLVM_Type>:$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.
//

View File

@ -152,6 +152,20 @@ llvm.func @matrix_intrinsics(%A: !llvm<"<64 x float>">, %B: !llvm<"<48 x float>"
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 @@ llvm.func @matrix_intrinsics(%A: !llvm<"<64 x float>">, %B: !llvm<"<48 x float>"
// 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>)