[Mlir-commits] [mlir] a28ce1a - [mlir][vector][bufferize] Fix transfer_write dropping mask operand

[Matthias Springer]

Author: Matthias Springer
Date: 2022-07-07T10:02:13+02:00
New Revision: a28ce1a42b39e174a346ace3359851b393bf3158

URL: https://github.com/llvm/llvm-project/commit/a28ce1a42b39e174a346ace3359851b393bf3158
DIFF: https://github.com/llvm/llvm-project/commit/a28ce1a42b39e174a346ace3359851b393bf3158.diff

LOG: [mlir][vector][bufferize] Fix transfer_write dropping mask operand

Differential Revision: https://reviews.llvm.org/D129253

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
    mlir/test/Dialect/Vector/bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
index 77f895a929c33..78c5c3bcd2b82 100644
--- a/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -106,7 +106,7 @@ struct TransferWriteOpInterface
     rewriter.create<vector::TransferWriteOp>(
         writeOp.getLoc(), writeOp.getVector(), *resultBuffer,
         writeOp.getIndices(), writeOp.getPermutationMapAttr(),
-        writeOp.getInBoundsAttr());
+        writeOp.getMask(), writeOp.getInBoundsAttr());
     replaceOpWithBufferizedValues(rewriter, op, *resultBuffer);
 
     return success();

diff  --git a/mlir/test/Dialect/Vector/bufferize.mlir b/mlir/test/Dialect/Vector/bufferize.mlir
index 170aa02fcacb1..ab271a642cbcc 100644
--- a/mlir/test/Dialect/Vector/bufferize.mlir
+++ b/mlir/test/Dialect/Vector/bufferize.mlir
@@ -15,16 +15,17 @@ func.func @transfer_read(%t: tensor<?x?xf32>, %o1: index,
 // -----
 
 // CHECK-LABEL: func @transfer_write(
-//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>, %[[o1:.*]]: index, %[[o2:.*]]: index, %[[vec:.*]]: vector<5x6xf32>)
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>, %[[o1:.*]]: index, %[[o2:.*]]: index, %[[vec:.*]]: vector<5x6xf32>, %[[mask:.*]]: vector<5x6xi1>)
 //       CHECK:   %[[m:.*]] = bufferization.to_memref %[[t]] : memref<?x?xf32>
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}, %{{.*}}) {{.*}} : memref<?x?xf32>
 //       CHECK:   memref.copy %[[m]], %[[alloc]]
-//       CHECK:   vector.transfer_write %[[vec]], %[[alloc]][%[[o1]], %[[o2]]] {in_bounds = [true, false]} : vector<5x6xf32>, memref<?x?xf32>
+//       CHECK:   vector.transfer_write %[[vec]], %[[alloc]][%[[o1]], %[[o2]]], %[[mask]] {in_bounds = [true, false]} : vector<5x6xf32>, memref<?x?xf32>
 //       CHECK:   %[[r:.*]] = bufferization.to_tensor %[[alloc]] : memref<?x?xf32>
 //       CHECK:   return %[[r]]
 func.func @transfer_write(%t: tensor<?x?xf32>, %o1: index,
-                     %o2: index, %vec: vector<5x6xf32>) -> tensor<?x?xf32> {
-  %0 = vector.transfer_write %vec, %t[%o1, %o2] {in_bounds = [true, false]}
+                          %o2: index, %vec: vector<5x6xf32>,
+                          %mask: vector<5x6xi1>) -> tensor<?x?xf32> {
+  %0 = vector.transfer_write %vec, %t[%o1, %o2], %mask {in_bounds = [true, false]}
       : vector<5x6xf32>, tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }