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

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

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index bbdb8e7b46b4c6e44ce66230e05976382444940e..d923ea1bea76e5e53c7ac5c35b8da87ea1539116 100644 (file)
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -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);

diff --git a/mlir/test/Dialect/Linalg/forward-vector-transfers.mlir b/mlir/test/Dialect/Linalg/forward-vector-transfers.mlir
index 7f56234219fe6ab62a187a110adc82222d14f557..35e9b7f13a7e2b6d5f05c3463dfcb6270cdf47c1 100644 (file)
--- a/mlir/test/Dialect/Linalg/forward-vector-transfers.mlir
+++ b/mlir/test/Dialect/Linalg/forward-vector-transfers.mlir
@@ -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