[Mlir-commits] [mlir] [mlir][vector] Mask a write that fills its dims but starts at a non-zero index (PR #233243)

[Dhairyashil R G]

https://github.com/dhairyashilRG created https://github.com/llvm/llvm-project/pull/233243

`createWriteOrMaskedWrite` skipped the mask when the vector shape matched the destination dims it writes to. It did not look at the write indices, so a write that starts at a non-zero index was left unmasked and marked `in_bounds`, although it runs past the end of the destination.

For example, vectorizing this with `vector_sizes [8]`:

```mlir
%r = tensor.insert_slice %src into %dest[%offset] [%size] [1]
    : tensor<?xi32> into tensor<8xi32>
```

gave an unmasked `vector.transfer_write` of 8 elements at `%offset`. The mask is now skipped only when the write also starts at 0, so this write is masked with `8 - %offset`.

Of the in-tree callers that mask, only `tensor.insert_slice` vectorization passes write indices, so nothing else changes.

Part of #227821.

Assisted-by: Claude

>From b342db9a363917fe36e3c4051f8b05d82eebb4d6 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Sat, 3 Oct 2026 11:51:40 +0530
Subject: [PATCH] [mlir][vector] Mask a write that fills its dims but starts at
 a non-zero index

`createWriteOrMaskedWrite` left a write unmasked when the vector shape matched
the destination dims it writes to, without looking at the write indices. A
write that starts at a non-zero index then runs past the end of the
destination.

For example, vectorizing this with `vector_sizes [8]`:

  %r = tensor.insert_slice %src into %dest[2] [6] [1]
      : tensor<6xi32> into tensor<8xi32>

gave an unmasked `vector.transfer_write` of 8 elements at index 2, which also
writes `%dest[8]` and `%dest[9]`. The mask is now skipped only when the write
starts at 0 in those dims, so this write is masked to its first 6 elements.

Of the in-tree callers that mask, only `tensor.insert_slice` vectorization
passes write indices, so nothing else changes.

Part of #227821.

Assisted-by: Claude
---
 mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp |  8 +++-
 .../Linalg/vectorization/insert-slice.mlir    | 38 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index 14a6571549897..c69edef936ad4 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -591,8 +591,12 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
                       : AffineMap::getMinorIdentityMap(destRank, vecToStoreRank,
                                                        builder.getContext());
 
-  // Check if masking is needed. If not, exit.
-  if (llvm::equal(vecToStoreShape, applyPermutationMap(map, destShape)))
+  // A write that starts at 0 and fills the dims it writes to needs no mask.
+  bool startsAtZero = llvm::all_of(map.getResults(), [&](AffineExpr expr) {
+    return isZeroInteger(writeIndices[cast<AffineDimExpr>(expr).getPosition()]);
+  });
+  if (startsAtZero &&
+      llvm::equal(vecToStoreShape, applyPermutationMap(map, destShape)))
     return write;
 
   // Compute the mask and mask the write Op.
diff --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
index 319095f97e151..639068f588124 100644
--- a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -328,6 +328,44 @@ func.func private @insert_slice_constant_non_zero_offset(
 
 // -----
 
+// The vector shape equals the destination shape, but the insert index for dim 1
+// is 2: only 6 elements are left in that dim, so the write needs a mask.
+
+func.func private @insert_slice_vector_matches_dest_non_zero_offset(
+    %source: tensor<4x6xi32>, %dest: tensor<4x8xi32>) -> tensor<4x8xi32> {
+  %res = tensor.insert_slice %source into %dest[0, 2] [4, 6] [1, 1] : tensor<4x6xi32> into tensor<4x8xi32>
+  return %res : tensor<4x8xi32>
+}
+
+// CHECK-LABEL:   func.func private @insert_slice_vector_matches_dest_non_zero_offset(
+// CHECK-SAME:      %[[SRC:.*]]: tensor<4x6xi32>,
+// CHECK-SAME:      %[[DEST:.*]]: tensor<4x8xi32>) -> tensor<4x8xi32> {
+// CHECK:           %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK:           %[[C_0:.*]] = arith.constant 0 : index
+// CHECK:           %[[C_0_1:.*]] = arith.constant 0 : index
+// CHECK:           %[[C_4:.*]] = arith.constant 4 : index
+// CHECK:           %[[C_6:.*]] = arith.constant 6 : index
+// CHECK:           %[[MASK_READ:.*]] = vector.create_mask %[[C_4]], %[[C_6]] : vector<4x8xi1>
+// CHECK:           %[[READ:.*]] = vector.mask %[[MASK_READ]] { vector.transfer_read %[[SRC]][%[[C_0_1]], %[[C_0_1]]], %[[PAD]] {in_bounds = [true, true]} : tensor<4x6xi32>, vector<4x8xi32> } : vector<4x8xi1> -> vector<4x8xi32>
+// CHECK:           %[[C_0_2:.*]] = arith.constant 0 : index
+// CHECK:           %[[C_2:.*]] = arith.constant 2 : index
+// CHECK:           %[[C_4_1:.*]] = arith.constant 4 : index
+// CHECK:           %[[C_8:.*]] = arith.constant 8 : index
+// CHECK:           %[[C_6_1:.*]] = arith.constant 6 : index
+// CHECK:           %[[MASK_WRITE:.*]] = vector.create_mask %[[C_4_1]], %[[C_6_1]] : vector<4x8xi1>
+// CHECK:           %[[RES:.*]] = vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]], %[[DEST]][%[[C_0_2]], %[[C_2]]] {in_bounds = [true, true]} : vector<4x8xi32>, tensor<4x8xi32> } : vector<4x8xi1> -> tensor<4x8xi32>
+// CHECK:           return %[[RES]] : tensor<4x8xi32>
+
+ module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    transform.structured.vectorize %0 vector_sizes [4, 8] : !transform.any_op
+    transform.yield
+  }
+ }
+
+// -----
+
 // One of the destination dimensions is dynamic and the the corresponding
 // insert index is != 0. Make sure that the mask is computed correctly (note arith.subi in the output).