[Mlir-commits] [mlir] d0e0d62 - [mlir][vector] Keep a write mask that does not cover the whole vector (#228668)

[llvmlistbot at llvm.org]

Author: Dhairyashil R G
Date: 2026-10-16T22:26:28+05:30
New Revision: d0e0d6245df75b752de18af6a61593fc8484eab2

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

LOG: [mlir][vector] Keep a write mask that does not cover the whole vector (#228668)

`isMaskTriviallyFoldable` decides whether a transfer mask is all-true
and can be left out. For constant sizes and indices it checked that the
masked part stays in bounds, but not that the mask covers the whole
vector.

Since #174810 the write mask size is `dim - index`. With a non-zero
constant index that can be smaller than the vector, and the mask was
still dropped.

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

```mlir
%r = tensor.insert_slice %src into %dest[6] [%size] [1]
    : tensor<?xi32> into tensor<8xi32>
```
gave an unmasked 4-element `vector.transfer_write` at index 6, marked
`in_bounds`. It is now masked with size 2.

The function now also requires each mask size to be at least the vector
size. Its doc comment listed the two existing conditions with the
comparison reversed. They now match the code.

A write that exactly fills the destination at a non-zero index is still
unmasked. It takes a different shortcut, which a follow-up will fix.

Part of #227821.

Assisted-by: Claude

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
    mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index c38213850c6ea5..bd889a82bdb59a 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -328,8 +328,9 @@ bool vector::isLinearizableVector(VectorType type) {
 ///
 /// If the pre-conditions are met, the method checks for each destination
 /// dimension `d`:
-///   (1) destDimSize[rankDiff + d] <= maskShape[d]
-///   (2) destDimSize[rankDiff + d] <= index[d] + maskSize[d]
+///   (1) maskShape[d] <= destDimSize[rankDiff + d]
+///   (2) index[d] + maskSize[d] <= destDimSize[rankDiff + d]
+///   (3) maskShape[d] <= maskSize[d]
 ///
 /// rankDiff = rank(dest) - rank(mask).
 ///
@@ -388,7 +389,8 @@ static bool isMaskTriviallyFoldable(SmallVector<OpFoldResult> &maskSizes,
   if (cstIndices.size() != baseShape.size())
     return false;
 
-  // Go over all destination dims and check (1) and (2). Take into account that:
+  // Go over all destination dims and check (1), (2) and (3). Take into account
+  // that:
   //  * The number of mask sizes will match the rank of the vector to
   //    load/store. This could be lower than the rank of the destination tensor.
   //  * Mask sizes could be larger than the corresponding mask shape (hence
@@ -399,7 +401,8 @@ static bool isMaskTriviallyFoldable(SmallVector<OpFoldResult> &maskSizes,
     if (/*(1)*/ maskShape[i] > baseShape[rankDiff + i] ||
         /*(2)*/ baseShape[rankDiff + i] <
             (std::clamp(cstMaskSizes[i], int64_t(0), maskShape[i]) +
-             cstIndices[i]))
+             cstIndices[i]) ||
+        /*(3)*/ maskShape[i] > cstMaskSizes[i])
       return false;
   }
 

diff  --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
index 34e251c527e316..a73b4b8035adbf 100644
--- a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -292,6 +292,31 @@ func.func private @insert_slice_dynamic_source_and_dest_dim(%source: tensor<?x3x
 
 // -----
 
+// The destination is static and the insert index is a non-zero constant. The
+// vector fits in the destination, but not from offset 6: only 2 elements are
+// left, so the write needs a mask.
+
+func.func private @insert_slice_constant_non_zero_offset(
+    %source: tensor<?xi32>, %dest: tensor<8xi32>, %size: index) -> tensor<8xi32> {
+  %res = tensor.insert_slice %source into %dest[6] [%size] [1] : tensor<?xi32> into tensor<8xi32>
+  return %res : tensor<8xi32>
+}
+// CHECK-LABEL:   func.func private @insert_slice_constant_non_zero_offset(
+// CHECK:           %[[READ:.*]] = vector.mask %{{.*}} { vector.transfer_read
+// CHECK:           %[[C_2:.*]] = arith.constant 2 : index
+// CHECK:           %[[MASK_WRITE:.*]] = vector.create_mask %[[C_2]] : vector<4xi1>
+// CHECK:           vector.mask %[[MASK_WRITE]] { vector.transfer_write %[[READ]]
+
+ 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] : !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).