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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


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

@llvm/pr-subscribers-mlir-linalg

Author: Dhairyashil R G (dhairyashilRG)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp (+6-4) 
- (modified) mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir (+38) 


``````````diff
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
+  }
+ }

``````````

</details>


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