[Mlir-commits] [mlir] [mlir][vector] Preserve dynamic mask bounds in shape cast folding (PR #229699)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-vector

Author: purnima shrivastava (purnima-nlp)

<details>
<summary>Changes</summary>

Fixes `ShapeCastCreateMaskFolderBoundaryUnitDim` so that folding `vector.shape_cast` across `vector.create_mask` only removes boundary unit dimensions when every corresponding mask bound is known to be constant one.

The fix:
- inspects exactly `numDimsToDrop` mask operands,
- rejects the rewrite if any removed `vector.create_mask` bound is not constant one,
- fixes the analogous slice calculation for `vector.constant_mask`.

Adds regression tests covering leading, trailing, and mixed dynamic mask bounds.

Fixes #<!-- -->229017.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+5-7) 
- (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+51) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index d5fd5095b18b000..941e22dd0bda3ca 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -7094,16 +7094,15 @@ class ShapeCastCreateMaskFolderBoundaryUnitDim final
 
     if (createMaskOp) {
       auto maskOperands = createMaskOp.getOperands();
-      size_t numOperands = maskOperands.size();
 
       auto maskOperandsToDrop =
           (Side == UnitDimSide::Trailing)
-              ? maskOperands.take_back(numOperands - numDimsToDrop)
-              : maskOperands.take_front(numOperands - numDimsToDrop);
+              ? maskOperands.take_back(numDimsToDrop)
+              : maskOperands.take_front(numDimsToDrop);
 
       // Check that every mask-dim-size to-be-dropped is constant and == 1. We
       // could also check for == 0, but that's left as a TODO.
-      if (llvm::all_of(maskOperandsToDrop, [](Value maskDim) {
+      if (llvm::any_of(maskOperandsToDrop, [](Value maskDim) {
             auto cst = maskDim.getDefiningOp<arith::ConstantIndexOp>();
             return !cst || (cst.value() != 1);
           }))
@@ -7119,12 +7118,11 @@ class ShapeCastCreateMaskFolderBoundaryUnitDim final
 
     if (constantMaskOp) {
       auto maskDimSizes = constantMaskOp.getMaskDimSizes();
-      size_t numDims = maskDimSizes.size();
 
       ArrayRef<int64_t> maskDimSizesToDrop =
           (Side == UnitDimSide::Trailing)
-              ? maskDimSizes.take_back(numDims - numDimsToDrop)
-              : maskDimSizes.take_front(numDims - numDimsToDrop);
+              ? maskDimSizes.take_back(numDimsToDrop)
+              : maskDimSizes.take_front(numDimsToDrop);
 
       // Check that every mask-dim-size to-be-dropped is constant and == 1. We
       // could also check for == 0, but that's left as a TODO.
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6b271ee5cce19d0..557835693883070 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3523,6 +3523,23 @@ func.func @fold_shape_cast_with_mask_trailing_unit(%arg0: tensor<1x?xf32>) -> ve
 
 // -----
 
+// CHECK-LABEL:   func.func @fold_shape_cast_with_mask_trailing_unit_dynamic(
+// CHECK-SAME:     %[[N:.*]]: index) -> vector<8xi1> {
+func.func @fold_shape_cast_with_mask_trailing_unit_dynamic(%n: index) -> vector<8xi1> {
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+// CHECK:     %[[MASK:.*]] = vector.create_mask %[[C8]], %[[N]], %[[C1]] : vector<8x1x1xi1>
+// CHECK:     %[[CAST:.*]] = vector.shape_cast %[[MASK]] : vector<8x1x1xi1> to vector<8xi1>
+// CHECK:     return %[[CAST]] : vector<8xi1>
+  %c1 = arith.constant 1 : index
+  %c8 = arith.constant 8 : index
+  %mask = vector.create_mask %c8, %n, %c1 : vector<8x1x1xi1>
+  %cast = vector.shape_cast %mask : vector<8x1x1xi1> to vector<8xi1>
+  return %cast : vector<8xi1>
+}
+
+// -----
+
 // CHECK-LABEL:   func.func @fold_shape_cast_with_mask_trailing_unit_scalable(
 // CHECK-SAME:    %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<1x[4]xi1> {
 func.func @fold_shape_cast_with_mask_trailing_unit_scalable(%arg0: tensor<1x?xf32>) -> vector<1x[4]xi1> {
@@ -3590,6 +3607,40 @@ func.func @fold_shape_cast_with_mask_leading_unit(%arg0: tensor<1x?xf32>) -> vec
 
 // -----
 
+// CHECK-LABEL:   func.func @fold_shape_cast_with_mask_leading_unit_dynamic(
+// CHECK-SAME:     %[[N:.*]]: index) -> vector<8xi1> {
+func.func @fold_shape_cast_with_mask_leading_unit_dynamic(%n: index) -> vector<8xi1> {
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+// CHECK:     %[[MASK:.*]] = vector.create_mask %[[C1]], %[[N]], %[[C8]] : vector<1x1x8xi1>
+// CHECK:     %[[CAST:.*]] = vector.shape_cast %[[MASK]] : vector<1x1x8xi1> to vector<8xi1>
+// CHECK:     return %[[CAST]] : vector<8xi1>
+  %c1 = arith.constant 1 : index
+  %c8 = arith.constant 8 : index
+  %mask = vector.create_mask %c1, %n, %c8 : vector<1x1x8xi1>
+  %cast = vector.shape_cast %mask : vector<1x1x8xi1> to vector<8xi1>
+  return %cast : vector<8xi1>
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @fold_shape_cast_with_mask_leading_unit_dynamic_mixed_bounds(
+// CHECK-SAME:     %[[N:.*]]: index, %[[M:.*]]: index) -> vector<8x4xi1> {
+func.func @fold_shape_cast_with_mask_leading_unit_dynamic_mixed_bounds(%n: index, %m: index) -> vector<8x4xi1> {
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+// CHECK:     %[[MASK:.*]] = vector.create_mask %[[C1]], %[[N]], %[[C8]], %[[M]] : vector<1x1x8x4xi1>
+// CHECK:     %[[CAST:.*]] = vector.shape_cast %[[MASK]] : vector<1x1x8x4xi1> to vector<8x4xi1>
+// CHECK:     return %[[CAST]] : vector<8x4xi1>
+  %c1 = arith.constant 1 : index
+  %c8 = arith.constant 8 : index
+  %mask = vector.create_mask %c1, %n, %c8, %m : vector<1x1x8x4xi1>
+  %cast = vector.shape_cast %mask : vector<1x1x8x4xi1> to vector<8x4xi1>
+  return %cast : vector<8x4xi1>
+}
+
+// -----
+
 // CHECK-LABEL:   func.func @fold_shape_cast_with_mask_leading_unit_scalable(
 // CHECK-SAME:    %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<[4]x1xi1> {
 func.func @fold_shape_cast_with_mask_leading_unit_scalable(%arg0: tensor<1x?xf32>) -> vector<[4]x1xi1> {

``````````

</details>


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