[Mlir-commits] [mlir] [mlir][vector] Preserve dynamic mask bounds in shape cast folding (PR #229699)
purnima shrivastava
llvmlistbot at llvm.org
Wed Oct 14 05:54:31 PDT 2026
https://github.com/purnima-nlp updated https://github.com/llvm/llvm-project/pull/229699
>From 8145e16c369ae76583e3061facd25a1decc1508f Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Wed, 7 Oct 2026 14:24:00 +0530
Subject: [PATCH] [mlir][vector] Preserve dynamic mask bounds in shape cast
folding
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 15 +++---
mlir/test/Dialect/Vector/canonicalize.mlir | 61 ++++++++++++++++++++--
2 files changed, 63 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index d5fd5095b18b00..93180555d19e26 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -7094,16 +7094,14 @@ 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);
+ auto maskOperandsToDrop = (Side == UnitDimSide::Trailing)
+ ? 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 +7117,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 6b271ee5cce19d..f08a6e5400eb65 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3506,9 +3506,9 @@ func.func @all_true_vector_mask_no_result(%a : vector<3x4xf32>, %m : memref<3x4x
// Tests for ShapeCastCreateMaskFolderBoundaryUnitDim<UnitDimSide::Trailing>
// +---------------------------------------------------------------------------
-// CHECK-LABEL: func.func @fold_shape_cast_with_mask_trailing_unit(
+// CHECK-LABEL: func.func @fold_shape_cast_with_mask_trailing_unit_provable_true(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<1x4xi1> {
-func.func @fold_shape_cast_with_mask_trailing_unit(%arg0: tensor<1x?xf32>) -> vector<1x4xi1> {
+func.func @fold_shape_cast_with_mask_trailing_unit_provable_true(%arg0: tensor<1x?xf32>) -> vector<1x4xi1> {
// CHECK-NOT: vector.shape_cast
// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
// CHECK: %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<1x?xf32>
@@ -3523,6 +3523,23 @@ func.func @fold_shape_cast_with_mask_trailing_unit(%arg0: tensor<1x?xf32>) -> ve
// -----
+// CHECK-LABEL: func.func @no_fold_shape_cast_with_mask_trailing_unit_dynamic_unknown(
+// CHECK-SAME: %[[N:.*]]: index) -> vector<8xi1> {
+func.func @no_fold_shape_cast_with_mask_trailing_unit_dynamic_unknown(%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> {
@@ -3573,9 +3590,9 @@ func.func @fold_shape_cast_with_constant_mask_trailing_unit() -> vector<4xi1>{
// Tests for ShapeCastCreateMaskFolderBoundaryUnitDim<UnitDimSide::Leading>
// +---------------------------------------------------------------------------
-// CHECK-LABEL: func.func @fold_shape_cast_with_mask_leading_unit(
+// CHECK-LABEL: func.func @fold_shape_cast_with_mask_leading_unit_provable_true(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<4x1xi1> {
-func.func @fold_shape_cast_with_mask_leading_unit(%arg0: tensor<1x?xf32>) -> vector<4x1xi1> {
+func.func @fold_shape_cast_with_mask_leading_unit_provable_true(%arg0: tensor<1x?xf32>) -> vector<4x1xi1> {
// CHECK-NOT: vector.shape_cast
// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
// CHECK: %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<1x?xf32>
@@ -3590,6 +3607,42 @@ func.func @fold_shape_cast_with_mask_leading_unit(%arg0: tensor<1x?xf32>) -> vec
// -----
+// CHECK-LABEL: func.func @no_fold_shape_cast_with_mask_leading_unit_dynamic_unknown(
+// CHECK-SAME: %[[N:.*]]: index) -> vector<8xi1> {
+func.func @no_fold_shape_cast_with_mask_leading_unit_dynamic_unknown(%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>
+}
+
+// -----
+
+// With four operands and two dropped dims, the operand slice has the expected
+// length. Verify that any unknown removed bound still blocks the fold.
+// CHECK-LABEL: func.func @no_fold_shape_cast_with_mask_leading_unit_mixed_unknown(
+// CHECK-SAME: %[[N:.*]]: index, %[[M:.*]]: index) -> vector<8x4xi1> {
+func.func @no_fold_shape_cast_with_mask_leading_unit_mixed_unknown(%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> {
More information about the Mlir-commits
mailing list