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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Dhairyashil R G (dhairyashilRG)

<details>
<summary>Changes</summary>

`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


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


2 Files Affected:

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


``````````diff
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).
 

``````````

</details>


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