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

[llvmlistbot at llvm.org]

https://github.com/Pecco-314 created https://github.com/llvm/llvm-project/pull/234906

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)

>From dc3a48e582ae708fa444a9071da005f9f4f7871d Mon Sep 17 00:00:00 2001
From: Yanfeng Gao <gaoyanfeng.gyf at alibaba-inc.com>
Date: Fri, 23 Oct 2026 18:31:28 +0800
Subject: [PATCH] [mlir][vector] Fix transfer_write bufferization coverage
 checks

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)
---
 .../BufferizableOpInterfaceImpl.cpp           |   4 +-
 .../Dialect/Vector/one-shot-bufferize.mlir    | 116 ++++++++++++++++++
 2 files changed, 118 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
index 352f477a8746e2..037e9ad07572fe 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 a190c947530086..14cd6132117203 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>
+}