[Mlir-commits] [mlir] [mlir][vector] Fix transfer_write bufferization coverage checks (PR #234906)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Pecco (Pecco-314)

<details>
<summary>Changes</summary>

The overwrite check can miss unwritten elements when a transfer has a mask, writes a lower-rank vector, or permutes dimensions. Bufferization may then skip a required copy and lose the original values of those elements.

The check now includes the transfer's mask and uses `getTransferChunkAccessed()` to determine coverage in destination dimensions. Transposed writes and writes that omit unit dimensions can still avoid a copy when they cover the entire destination.

Fixes #<!-- -->233726

Assisted-by: Codex (OpenAI)

---
Full diff: https://github.com/llvm/llvm-project/pull/234906.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp (+2-2) 
- (modified) mlir/test/Dialect/Vector/one-shot-bufferize.mlir (+116) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
index 352f477a8746e..037e9ad07572f 100644
--- a/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -92,12 +92,12 @@ struct TransferWriteOpInterface
     }
 
     // There is no mask.
-    if (writeOp.isMasked())
+    if (writeOp.getMask() || writeOp.isMasked())
       return true;
 
     // Must write at least the full dimension size.
     for (auto [d0, d1] : llvm::zip(writeOp.getShapedType().getShape(),
-                                   writeOp.getVectorType().getShape())) {
+                                   writeOp.getTransferChunkAccessed())) {
       if (d0 > d1)
         return true;
     }
diff --git a/mlir/test/Dialect/Vector/one-shot-bufferize.mlir b/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
index a190c94753008..14cd613211720 100644
--- a/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
@@ -12,6 +12,8 @@ func.func @mask(%t0: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vect
   return %0 : tensor<?xf32>
 }
 
+// -----
+
 // CHECK-LABEL: func @mask_scalable(
 //  CHECK-SAME:     %[[t0:.*]]: memref<?xf32, strided<[?], offset: ?>>
 func.func @mask_scalable(%t0: tensor<?xf32>, %val: vector<[16]xf32>, %idx: index, %m0: vector<[16]xi1>) -> tensor<?xf32> {
@@ -125,3 +127,117 @@ func.func @unknown_transfer_read_write(
       : tensor<8xf32>, vector<4xf32>
   return %written, %read : tensor<8xf32>, vector<4xf32>
 }
+
+// -----
+
+// Masked-off elements retain their values from the destination.
+
+// CHECK-LABEL: func @masked_xfer_write(
+//  CHECK-SAME: %[[T:.*]]: memref<4xf32,
+//       CHECK: %[[ALLOC:.*]] = memref.alloc
+//       CHECK: memref.copy %[[T]], %[[ALLOC]]
+//       CHECK: vector.transfer_write %{{.*}}, %[[ALLOC]][
+//       CHECK: return %[[ALLOC]]
+
+// CHECK-ANALYSIS-LABEL: func @masked_xfer_write(
+// CHECK-ANALYSIS-SAME: tensor<4xf32> {bufferization.access = "read",
+func.func @masked_xfer_write(
+    %t: tensor<4xf32> {bufferization.writable = false},
+    %v: vector<4xf32>, %mask: vector<4xi1>) -> tensor<4xf32> {
+  %c0 = arith.constant 0 : index
+  %written = vector.transfer_write %v, %t[%c0], %mask
+      : vector<4xf32>, tensor<4xf32>
+  return %written : tensor<4xf32>
+}
+
+// -----
+
+// A lower-rank write updates one row, so the other rows must be copied.
+
+// CHECK-LABEL: func @lower_rank_xfer_write(
+//  CHECK-SAME: %[[T:.*]]: memref<4x4xf32,
+//       CHECK: %[[ALLOC:.*]] = memref.alloc
+//       CHECK: memref.copy %[[T]], %[[ALLOC]]
+//       CHECK: vector.transfer_write %{{.*}}, %[[ALLOC]][
+//       CHECK: return %[[ALLOC]]
+
+// CHECK-ANALYSIS-LABEL: func @lower_rank_xfer_write(
+// CHECK-ANALYSIS-SAME: tensor<4x4xf32> {bufferization.access = "read",
+func.func @lower_rank_xfer_write(
+    %t: tensor<4x4xf32> {bufferization.writable = false},
+    %v: vector<4xf32>) -> tensor<4x4xf32> {
+  %c0 = arith.constant 0 : index
+  %written = vector.transfer_write %v, %t[%c0, %c0]
+      : vector<4xf32>, tensor<4x4xf32>
+  return %written : tensor<4x4xf32>
+}
+
+// -----
+
+// The transpose makes the write cover an 8x4 slice of a 4x8 tensor.
+// The last four columns must be preserved.
+
+// CHECK-LABEL: func @partial_permuted_xfer_write(
+//  CHECK-SAME: %[[T:.*]]: memref<4x8xf32,
+//       CHECK: %[[ALLOC:.*]] = memref.alloc
+//       CHECK: memref.copy %[[T]], %[[ALLOC]]
+//       CHECK: vector.transfer_write %{{.*}}, %[[ALLOC]][
+//       CHECK: return %[[ALLOC]]
+
+// CHECK-ANALYSIS-LABEL: func @partial_permuted_xfer_write(
+// CHECK-ANALYSIS-SAME: tensor<4x8xf32> {bufferization.access = "read",
+func.func @partial_permuted_xfer_write(
+    %t: tensor<4x8xf32> {bufferization.writable = false},
+    %v: vector<4x8xf32>) -> tensor<4x8xf32> {
+  %c0 = arith.constant 0 : index
+  %written = vector.transfer_write %v, %t[%c0, %c0]
+      {permutation_map = affine_map<(d0, d1) -> (d1, d0)>}
+      : vector<4x8xf32>, tensor<4x8xf32>
+  return %written : tensor<4x8xf32>
+}
+
+// -----
+
+// The transpose makes an 8x4 vector cover the entire 4x8 destination.
+
+// CHECK-LABEL: func @full_permuted_xfer_write(
+//       CHECK: %[[ALLOC:.*]] = memref.alloc
+//   CHECK-NOT: memref.copy
+//       CHECK: vector.transfer_write %{{.*}}, %[[ALLOC]][
+//   CHECK-NOT: memref.copy
+//       CHECK: return %[[ALLOC]]
+
+// CHECK-ANALYSIS-LABEL: func @full_permuted_xfer_write(
+// CHECK-ANALYSIS-SAME: tensor<4x8xf32> {bufferization.access = "none",
+func.func @full_permuted_xfer_write(
+    %t: tensor<4x8xf32> {bufferization.writable = false},
+    %v: vector<8x4xf32>) -> tensor<4x8xf32> {
+  %c0 = arith.constant 0 : index
+  %written = vector.transfer_write %v, %t[%c0, %c0]
+      {permutation_map = affine_map<(d0, d1) -> (d1, d0)>}
+      : vector<8x4xf32>, tensor<4x8xf32>
+  return %written : tensor<4x8xf32>
+}
+
+// -----
+
+// The dimension omitted by the map has size one, so the write covers the
+// entire destination.
+
+// CHECK-LABEL: func @full_lower_rank_xfer_write(
+//       CHECK: %[[ALLOC:.*]] = memref.alloc
+//   CHECK-NOT: memref.copy
+//       CHECK: vector.transfer_write %{{.*}}, %[[ALLOC]][
+//   CHECK-NOT: memref.copy
+//       CHECK: return %[[ALLOC]]
+
+// CHECK-ANALYSIS-LABEL: func @full_lower_rank_xfer_write(
+// CHECK-ANALYSIS-SAME: tensor<1x4xf32> {bufferization.access = "none",
+func.func @full_lower_rank_xfer_write(
+    %t: tensor<1x4xf32> {bufferization.writable = false},
+    %v: vector<4xf32>) -> tensor<1x4xf32> {
+  %c0 = arith.constant 0 : index
+  %written = vector.transfer_write %v, %t[%c0, %c0]
+      : vector<4xf32>, tensor<1x4xf32>
+  return %written : tensor<1x4xf32>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/234906