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

[nicolasvasilache]

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.

[bkramer]

lgtm

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

LogicalResult LinalgCopyVTRForwardingPattern::matchAndRewrite(
   if (maybeFillOp && xferOp.padding() != maybeFillOp.value())
     return failure();
   if (maybeFillOp)
     LLVM_DEBUG(llvm::dbgs() << dbgPref << "with maybeFillOp " << *maybeFillOp);
 
   // `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);
   rewriter.eraseOp(copyOp);
   rewriter.replaceOp(xferOp, res);
 
   return success();
 }
LogicalResult LinalgCopyVTWForwardingPattern::matchAndRewrite(
   }
   if (!copyOp)
     return failure();
 
   // `out` is the subview copied into that we replace.
   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);
 
   return success();
 }

mlir/test/Dialect/Linalg/forward-vector-transfers.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -test-linalg-transform-patterns=test-vector-transfer-forwarding-patterns | FileCheck %s
 
 // CHECK-LABEL: testAllocRead
 //  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
 //   CHECK-NOT: linalg.fill
 //   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>
 }
 
 // CHECK-LABEL: testAllocFillRead
 //  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
 //   CHECK-NOT: linalg.fill
 //   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
   %alloc = alloc() : memref<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>
 }
 
 // CHECK-LABEL: testViewRead
 //  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
 //   CHECK-NOT: linalg.fill
 //   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
   %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>
   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>
 }
 
 // CHECK-LABEL: testViewFillRead
 //  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
 //   CHECK-NOT: linalg.fill
 //   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
   %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>
   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>
 }
 
 // CHECK-LABEL: testAllocWrite
 //  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector
 //  CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: memref
 //   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
 }
 
 // CHECK-LABEL: testViewWrite
 //  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector
 //  CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: memref
 //   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
 }
 
 ///===--------------------------------------------------------------------===///
 // Negative tests
 ///===--------------------------------------------------------------------===///