[Mlir-commits] [mlir] [mlir][affine] Mask transfers in SuperVectorize instead of using in_bounds (PR #233244)

[Dhairyashil R G]

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

SuperVectorize built its transfers with `useInBoundsInsteadOfMasking=true`. That marks a transfer `in_bounds` when the memref dim is a multiple of the vector size, without looking at the index. With vectors of 8 on a `memref<64xf32>`, a read of `%src[%i + 1]` reads one element past the end and a write to `%dst[%i + 4]` writes four past the end.

Both call sites now pass `false`, so each transfer is masked by what is left of the memref from its index. `createReadOrMaskedRead` built its mask from the source sizes only. It now uses `dim - index` for each dim the permutation map selects, as `createWriteOrMaskedWrite` does, and leaves a read that only broadcasts unmasked. SuperVectorize is the only in-tree caller that passes indices or a map to it.

The masks are `vector.mask` regions, so a pipeline that lowers the output to LLVM needs `-lower-vector-mask` before `-convert-vector-to-llvm`.

Loops that divide evenly do not keep the masks: `-eliminate-vector-masks -canonicalize` removes them. `vectorize_2d_inbounds.mlir` now runs those passes and still checks for plain `llvm.load`/`llvm.store`.

In `vectorize_reduction.mlir`, the case that needs no reduction mask now checks for no `arith.select`, because the read has a mask of its own.

Part of #227821.

Assisted-by: Claude

>From f5160cd46466c0d7f49853618df1656249947963 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Sat, 17 Oct 2026 16:04:54 +0530
Subject: [PATCH] [mlir][vector] Check the written dims when folding a write
 mask

`isMaskTriviallyFoldable` compares dest dim `rankDiff + i` with index `i`, so
when the vector has fewer dims than the destination, it checks dims against the
wrong indices.
`createWriteOrMaskedWrite` therefore did not call it when given a permutation
map, and kept the mask even when it is all true. It now passes only the sizes
and indices of the dims that the map selects, and the mask shape, so the check
runs for every map.

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

  %r = tensor.insert_slice %src into %dest[2, 6, 0] [4, 1, 4] [1, 1, 1]
      : tensor<4x4xi32> into tensor<8x8x8xi32>

gave a write masked with `vector.create_mask` of sizes 6 and 8 on a
`vector<4x4xi1>`, which is all true. The write is now unmasked.

Of the in-tree callers that mask, only `tensor.insert_slice` vectorization
passes a map, for a rank-reducing slice, so nothing else changes.

Part of #227821.

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

diff --git a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
index 14a6571549897..42790ecc34d38 100644
--- a/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
+++ b/mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
@@ -610,16 +610,18 @@ Operation *vector::createWriteOrMaskedWrite(OpBuilder &builder, Location loc,
     writtenDims.push_back(cast<AffineDimExpr>(expr).getPosition());
   llvm::sort(writtenDims);
   SmallVector<OpFoldResult> maskSizes;
+  SmallVector<Value> writtenIndices;
+  SmallVector<int64_t> writtenShape;
   for (int64_t dim : writtenDims) {
     Value size = getValueOrCreateConstantIndexOp(builder, loc, destSizes[dim]);
     maskSizes.push_back(
         builder.createOrFold<arith::SubIOp>(loc, size, writeIndices[dim]));
+    writtenIndices.push_back(writeIndices[dim]);
+    writtenShape.push_back(destShape[dim]);
   }
 
-  // TODO: isMaskTriviallyFoldable checks the trailing dims against the leading
-  // indices. Make it use the dims that the map selects and their indices.
-  if (!permutationMap && isMaskTriviallyFoldable(maskSizes, writeIndices,
-                                                 destShape, vecToStoreShape))
+  if (isMaskTriviallyFoldable(maskSizes, writtenIndices, writtenShape,
+                              writeMaskType.getShape()))
     return write;
 
   Value maskForWrite =
diff --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
index 319095f97e151..3ea024b34462d 100644
--- a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -455,3 +455,41 @@ func.func private @insert_slice_leading_unit_dim_dropped_non_zero_offset(
     transform.yield
   }
  }
+
+// -----
+
+/// Dim 1 is dropped, so the vector is written to result dims 0 and 2. The
+/// vector fits in both from their indices (4 <= 8 - 2, 4 <= 8 - 0), so the
+/// write needs no mask. The index of the dropped dim (6) does not matter.
+
+func.func private @insert_slice_non_leading_unit_dim_dropped_non_zero_offsets(
+    %source: tensor<4x4xi32>, %dest: tensor<8x8x8xi32>) -> tensor<8x8x8xi32> {
+  %res = tensor.insert_slice %source into %dest[2, 6, 0] [4, 1, 4] [1, 1, 1] : tensor<4x4xi32> into tensor<8x8x8xi32>
+  return %res : tensor<8x8x8xi32>
+}
+
+// CHECK-DAG: #[[$MAP_D0_D2:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+
+// CHECK-LABEL:   func.func private @insert_slice_non_leading_unit_dim_dropped_non_zero_offsets(
+// CHECK-SAME:      %[[SRC:.*]]: tensor<4x4xi32>,
+// CHECK-SAME:      %[[DEST:.*]]: tensor<8x8x8xi32>) -> tensor<8x8x8xi32> {
+// CHECK:           %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK:           %[[C_0:.*]] = arith.constant 0 : index
+// CHECK:           %[[C_0_1:.*]] = arith.constant 0 : index
+// CHECK:           %[[READ:.*]] = vector.transfer_read %[[SRC]][%[[C_0_1]], %[[C_0_1]]], %[[PAD]] {in_bounds = [true, true]} : tensor<4x4xi32>, vector<4x4xi32>
+// CHECK:           %[[C_2:.*]] = arith.constant 2 : index
+// CHECK:           %[[C_6:.*]] = arith.constant 6 : index
+// CHECK:           %[[C_0_2:.*]] = arith.constant 0 : index
+// CHECK:           %[[RES:.*]] = vector.transfer_write %[[READ]], %[[DEST]][%[[C_2]], %[[C_6]], %[[C_0_2]]] {in_bounds = [true, true], permutation_map = #[[$MAP_D0_D2]]} : vector<4x4xi32>, tensor<8x8x8xi32>
+// CHECK:           %[[C_8:.*]] = arith.constant 8 : index
+// CHECK:           %[[C_6_1:.*]] = arith.constant 6 : index
+// CHECK:           %[[C_8_1:.*]] = arith.constant 8 : index
+// CHECK:           return %[[RES]] : tensor<8x8x8xi32>
+
+ 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, 4] : !transform.any_op
+    transform.yield
+  }
+ }