maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][Vector] Degrade masking information when forwarding linalg.copy to vector.transfer 

Summary:
linalg.copy + linalg.fill can be used to create a padded local buffer.
The `masked` attribute is only valid on this padded buffer.
When forwarding to vector.transfer ops, the attribute must be reset
conservatively.

Differential Revision: https://reviews.llvm.org/D83782
This commit is contained in:
Nicolas Vasilache 2020-07-14 11:45:59 -04:00
parent 694ded37b9
commit 512da70be7
2 changed files with 22 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
View File

@ -260,10 +260,13 @@ LogicalResult LinalgCopyVTRForwardingPattern::matchAndRewrite(
// `in` is the subview that linalg.copy reads. Replace it.
Value in = copyOp.getInput(0);
// linalg.copy + linalg.fill can be used to create a padded local buffer.
// The `masked` attribute is only valid on this padded buffer.
// When forwarding to vector.transfer_read, the attribute must be reset
// conservatively.
Value res = rewriter.create<vector::TransferReadOp>(
xferOp.getLoc(), xferOp.getVectorType(), in, xferOp.indices(),
xferOp.permutation_map(), xferOp.padding(),
xferOp.masked() ? *xferOp.masked() : ArrayAttr());
xferOp.permutation_map(), xferOp.padding(), ArrayAttr());
if (maybeFillOp)
rewriter.eraseOp(maybeFillOp);
@ -308,10 +311,13 @@ LogicalResult LinalgCopyVTWForwardingPattern::matchAndRewrite(
Value out = copyOp.getOutputBuffer(0);
// Forward vector.transfer into copy.
// linalg.copy + linalg.fill can be used to create a padded local buffer.
// The `masked` attribute is only valid on this padded buffer.
// When forwarding to vector.transfer_write, the attribute must be reset
// conservatively.
rewriter.create<vector::TransferWriteOp>(
xferOp.getLoc(), xferOp.vector(), out, xferOp.indices(),
xferOp.permutation_map(),
xferOp.masked() ? *xferOp.masked() : ArrayAttr());
xferOp.permutation_map(), ArrayAttr());
rewriter.eraseOp(copyOp);
rewriter.eraseOp(xferOp);

18
mlir/test/Dialect/Linalg/forward-vector-transfers.mlir
View File

@ -6,13 +6,14 @@
// CHECK-NOT: linalg.copy
// CHECK: %[[ALLOC:.*]] = alloc
// CHECK: vector.transfer_read %[[ARG0]]
// CHECK-NOT: masked
func @testAllocRead(%in: memref<? x f32>) -> vector<32 x f32> {
%c0 = constant 0: index
%f0 = constant 0.0: f32
%alloc = alloc() : memref<32 x f32>
%subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
%0 = vector.transfer_read %alloc[%c0], %f0: memref<32 x f32>, vector<32 x f32>
%0 = vector.transfer_read %alloc[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>
dealloc %alloc : memref<32 x f32>
return %0: vector<32 x f32>
}
@ -23,6 +24,7 @@ func @testAllocRead(%in: memref<? x f32>) -> vector<32 x f32> {
// CHECK-NOT: linalg.copy
// CHECK: %[[ALLOC:.*]] = alloc
// CHECK: vector.transfer_read %[[ARG0]]
// CHECK-NOT: masked
func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
%c0 = constant 0: index
%f0 = constant 0.0: f32
@ -30,7 +32,7 @@ func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
linalg.fill(%alloc, %f0): memref<32 x f32>, f32
%subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
%0 = vector.transfer_read %alloc[%c0], %f0: memref<32 x f32>, vector<32 x f32>
%0 = vector.transfer_read %alloc[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>
dealloc %alloc : memref<32 x f32>
return %0: vector<32 x f32>
}
@ -41,6 +43,7 @@ func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
// CHECK-NOT: linalg.copy
// CHECK: %[[ALLOC:.*]] = alloc
// CHECK: vector.transfer_read %[[ARG0]]
// CHECK-NOT: masked
func @testViewRead(%in: memref<? x f32>) -> vector<32 x f32> {
%c0 = constant 0: index
%f0 = constant 0.0: f32
@ -48,7 +51,7 @@ func @testViewRead(%in: memref<? x f32>) -> vector<32 x f32> {
%view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>
%subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
%0 = vector.transfer_read %view[%c0], %f0: memref<32 x f32>, vector<32 x f32>
%0 = vector.transfer_read %view[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>
dealloc %alloc : memref<128 x i8>
return %0: vector<32 x f32>
}
@ -59,6 +62,7 @@ func @testViewRead(%in: memref<? x f32>) -> vector<32 x f32> {
// CHECK-NOT: linalg.copy
// CHECK: %[[ALLOC:.*]] = alloc
// CHECK: vector.transfer_read %[[ARG0]]
// CHECK-NOT: masked
func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
%c0 = constant 0: index
%f0 = constant 0.0: f32
@ -67,7 +71,7 @@ func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
%subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
linalg.fill(%view, %f0): memref<32 x f32>, f32
linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
%0 = vector.transfer_read %view[%c0], %f0: memref<32 x f32>, vector<32 x f32>
%0 = vector.transfer_read %view[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>
dealloc %alloc : memref<128 x i8>
return %0: vector<32 x f32>
}
@ -78,12 +82,13 @@ func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
// CHECK-NOT: linalg.copy
// CHECK: %[[ALLOC:.*]] = alloc
// CHECK: vector.transfer_write %[[ARG0]], %[[ARG1]]
// CHECK-NOT: masked
func @testAllocWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {
%c0 = constant 0: index
%f0 = constant 0.0: f32
%alloc = alloc() : memref<32 x f32>
%subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
vector.transfer_write %vec, %alloc[%c0] : vector<32 x f32>, memref<32 x f32>
vector.transfer_write %vec, %alloc[%c0] {masked = [false]} : vector<32 x f32>, memref<32 x f32>
linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>
dealloc %alloc : memref<32 x f32>
return
@ -95,13 +100,14 @@ func @testAllocWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {
// CHECK-NOT: linalg.copy
// CHECK: %[[ALLOC:.*]] = alloc
// CHECK: vector.transfer_write %[[ARG0]], %[[ARG1]]
// CHECK-NOT: masked
func @testViewWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {
%c0 = constant 0: index
%f0 = constant 0.0: f32
%alloc = alloc() : memref<128 x i8>
%view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>
%subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
vector.transfer_write %vec, %view[%c0] : vector<32 x f32>, memref<32 x f32>
vector.transfer_write %vec, %view[%c0] {masked = [false]} : vector<32 x f32>, memref<32 x f32>
linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>
dealloc %alloc : memref<128 x i8>
return