[Mlir-commits] [mlir] [mlir][vector] Extend `ShapeCastCreateMaskFolderTrailingOneDim` (PR #219494)
Andrzej Warzyński
llvmlistbot at llvm.org
Thu Sep 3 01:49:54 PDT 2026
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/219494
>From 8ca60826da56dd96592cf34b08264781ee98e9ec Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 28 Aug 2026 14:41:13 +0000
Subject: [PATCH 1/2] [mlir][vector] Add
`ShapeCastCreateMaskFolderLeadingOneDim`
This pattern complements `ShapeCastCreateMaskFolderTrailingOneDim` and
will basically fold:
```mlir
%1 = vector.create_mask %c1, %c1, %dim, %c1 : vector<1x1x[4]x1xi1>
%2 = vector.shape_cast %1 : vector<1x1x[4]x1xi1> to vector<[4]x1xi1>
```
as:
```mlir
%0 = vector.create_mask %c1, %dim : vector<[4]xi1>
```
Implementation and tests mirror
`ShapeCastCreateMaskFolderTrailingOneDim`.
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 117 +++++++++++++++++++--
mlir/test/Dialect/Vector/canonicalize.mlir | 87 +++++++++++++--
2 files changed, 190 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index f8f3deb2e4789..dce252962e9de 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6952,7 +6952,7 @@ namespace {
///
/// vector<4x1x1xi1> --> vector<4x1xi1>
///
-static VectorType trimTrailingOneDims(VectorType oldType) {
+static VectorType trimTrailingUnitDims(VectorType oldType) {
ArrayRef<int64_t> oldShape = oldType.getShape();
ArrayRef<int64_t> newShape = oldShape;
@@ -6974,9 +6974,114 @@ static VectorType trimTrailingOneDims(VectorType oldType) {
return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
}
+/// Helper function that computes a new vector type based on the input vector
+/// type by removing the trailing one dims:
+///
+/// vector<4x1x1xi1> --> vector<4x1xi1>
+///
+static VectorType trimLeadingUnitDims(VectorType oldType) {
+ ArrayRef<int64_t> oldShape = oldType.getShape();
+ ArrayRef<int64_t> newShape = oldShape;
+
+ ArrayRef<bool> oldScalableDims = oldType.getScalableDims();
+ ArrayRef<bool> newScalableDims = oldScalableDims;
+
+ while (!newShape.empty() && newShape.front() == 1 &&
+ !newScalableDims.front()) {
+ newShape = newShape.drop_front(1);
+ newScalableDims = newScalableDims.drop_front(1);
+ }
+
+ // Make sure we have at least 1 dimension.
+ // TODO: Add support for 0-D vectors.
+ if (newShape.empty()) {
+ newShape = oldShape.take_back();
+ newScalableDims = oldScalableDims.take_back();
+ }
+
+ return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
+}
+
+/// Folds qualifying shape_cast(create_mask) into a new create_mask
+///
+/// Looks at `vector.shape_cast` Ops that simply "drop" the _leading_unit
+/// dimension. If the input vector comes from `vector.create_mask` for which
+/// the corresponding mask input value is 1 (e.g. `%c1` below), then it is safe
+/// to fold shape_cast into create_mask.
+///
+/// BEFORE:
+/// %1 = vector.create_mask %c1, %c1, %dim, %c1 : vector<1x1x[4]x1xi1>
+/// %2 = vector.shape_cast %1 : vector<1x1x[4]x1xi1> to vector<[4]x1xi1>
+/// AFTER:
+/// %0 = vector.create_mask %c1, %dim : vector<[4]xi1>
+class ShapeCastCreateMaskFolderLeadingOneDim final
+ : public OpRewritePattern<ShapeCastOp> {
+public:
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(ShapeCastOp shapeOp,
+ PatternRewriter &rewriter) const override {
+ Value shapeOpSrc = shapeOp->getOperand(0);
+ auto createMaskOp = shapeOpSrc.getDefiningOp<vector::CreateMaskOp>();
+ auto constantMaskOp = shapeOpSrc.getDefiningOp<vector::ConstantMaskOp>();
+ if (!createMaskOp && !constantMaskOp)
+ return failure();
+
+ VectorType shapeOpResTy = shapeOp.getResultVectorType();
+ VectorType shapeOpSrcTy = shapeOp.getSourceVectorType();
+
+ VectorType newVecType = trimLeadingUnitDims(shapeOpSrcTy);
+ if (newVecType != shapeOpResTy)
+ return rewriter.notifyMatchFailure(
+ shapeOp, "Non-leading-unit-dim dropping shape_cast Op");
+
+ auto numDimsToDrop = shapeOpSrcTy.getRank() - shapeOpResTy.getRank();
+
+ // No unit dims to drop
+ if (numDimsToDrop == 0)
+ return rewriter.notifyMatchFailure(
+ shapeOp, "Non-leading-unit-dim dropping shape_cast Op");
+
+ if (createMaskOp) {
+ auto maskOperands = createMaskOp.getOperands();
+
+ for (int64_t dimIdx = 0; dimIdx < numDimsToDrop; dimIdx++) {
+ auto constantOp =
+ maskOperands[dimIdx].getDefiningOp<arith::ConstantIndexOp>();
+ if (!constantOp || constantOp.value() != 1) {
+ return failure();
+ }
+ }
+
+ rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(
+ shapeOp, shapeOpResTy, maskOperands.drop_front(numDimsToDrop));
+
+ return success();
+ }
+
+ if (constantMaskOp) {
+ auto maskDimSizes = constantMaskOp.getMaskDimSizes();
+ auto numMaskOperands = maskDimSizes.size();
+
+ // Check every mask dim size to see whether it can be dropped
+ for (uint64_t i = 0; i < numMaskOperands; ++i) {
+ if (maskDimSizes[i] != 1)
+ return failure();
+ }
+
+ auto newMaskOperands = maskDimSizes.drop_front(numDimsToDrop);
+ rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(shapeOp, shapeOpResTy,
+ newMaskOperands);
+ return success();
+ }
+
+ return failure();
+ }
+};
+
/// Folds qualifying shape_cast(create_mask) into a new create_mask
///
-/// Looks at `vector.shape_cast` Ops that simply "drop" the trailing unit
+/// Looks at `vector.shape_cast` Ops that simply "drop" the _trailing_ unit
/// dimension. If the input vector comes from `vector.create_mask` for which
/// the corresponding mask input value is 1 (e.g. `%c1` below), then it is safe
/// to fold shape_cast into create_mask.
@@ -7002,12 +7107,11 @@ class ShapeCastCreateMaskFolderTrailingOneDim final
VectorType shapeOpResTy = shapeOp.getResultVectorType();
VectorType shapeOpSrcTy = shapeOp.getSourceVectorType();
- VectorType newVecType = trimTrailingOneDims(shapeOpSrcTy);
+ VectorType newVecType = trimTrailingUnitDims(shapeOpSrcTy);
if (newVecType != shapeOpResTy)
return failure();
- auto numDimsToDrop =
- shapeOpSrcTy.getShape().size() - shapeOpResTy.getShape().size();
+ auto numDimsToDrop = shapeOpSrcTy.getRank() - shapeOpResTy.getRank();
// No unit dims to drop
if (!numDimsToDrop)
@@ -7156,7 +7260,8 @@ class FoldShapeCastOfFromElements final : public OpRewritePattern<ShapeCastOp> {
void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<ShapeCastCreateMaskFolderTrailingOneDim, ShapeCastBroadcastFolder,
+ results.add<ShapeCastCreateMaskFolderTrailingOneDim,
+ ShapeCastCreateMaskFolderLeadingOneDim, ShapeCastBroadcastFolder,
FoldShapeCastOfFromElements>(context);
}
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 25edb72cc5eb2..d09a721e38692 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3463,9 +3463,13 @@ func.func @all_true_vector_mask_no_result(%a : vector<3x4xf32>, %m : memref<3x4x
// -----
-// CHECK-LABEL: func.func @fold_shape_cast_with_mask(
+// +---------------------------------------------------------------------------
+// Tests for ShapeCastCreateMaskFolderTrailingOneDim
+// +---------------------------------------------------------------------------
+
+// CHECK-LABEL: func.func @fold_shape_cast_with_mask_trailing_unit(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<1x4xi1> {
-func.func @fold_shape_cast_with_mask(%arg0: tensor<1x?xf32>) -> vector<1x4xi1> {
+func.func @fold_shape_cast_with_mask_trailing_unit(%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>
@@ -3480,9 +3484,9 @@ func.func @fold_shape_cast_with_mask(%arg0: tensor<1x?xf32>) -> vector<1x4xi1> {
// -----
-// CHECK-LABEL: func.func @fold_shape_cast_with_mask_scalable(
+// 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_scalable(%arg0: 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> {
// CHECK-NOT: vector.shape_cast
// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
// CHECK: %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<1x?xf32>
@@ -3498,9 +3502,9 @@ func.func @fold_shape_cast_with_mask_scalable(%arg0: tensor<1x?xf32>) -> vector<
// -----
// Check that scalable "1" (i.e. [1]) is not folded
-// CHECK-LABEL: func.func @fold_shape_cast_with_mask_scalable_one(
+// CHECK-LABEL: func.func @fold_shape_cast_with_mask_trailing_unit_scalable_one(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<1x[1]xi1> {
-func.func @fold_shape_cast_with_mask_scalable_one(%arg0: tensor<1x?xf32>) -> vector<1x[1]xi1>{
+func.func @fold_shape_cast_with_mask_trailing_unit_scalable_one(%arg0: tensor<1x?xf32>) -> vector<1x[1]xi1>{
// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
// CHECK: %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<1x?xf32>
// CHECK: %[[VAL_3:.*]] = vector.create_mask %[[VAL_1]], %[[VAL_2]] : vector<1x[1]xi1>
@@ -3514,8 +3518,8 @@ func.func @fold_shape_cast_with_mask_scalable_one(%arg0: tensor<1x?xf32>) -> vec
// -----
-// CHECK-LABEL: func.func @fold_shape_cast_with_constant_mask() -> vector<4xi1> {
-func.func @fold_shape_cast_with_constant_mask() -> vector<4xi1>{
+// CHECK-LABEL: func.func @fold_shape_cast_with_constant_mask_trailing_unit() -> vector<4xi1> {
+func.func @fold_shape_cast_with_constant_mask_trailing_unit() -> vector<4xi1>{
// CHECK-NOT: vector.shape_cast
// CHECK: %[[VAL_0:.*]] = vector.constant_mask [1] : vector<4xi1>
// CHECK: return %[[VAL_0]] : vector<4xi1>
@@ -3526,6 +3530,73 @@ func.func @fold_shape_cast_with_constant_mask() -> vector<4xi1>{
// -----
+// +---------------------------------------------------------------------------
+// Tests for CastAwayTransferWriteLeadingOneDim
+// +---------------------------------------------------------------------------
+
+// CHECK-LABEL: func.func @fold_shape_cast_with_mask_leading_unit(
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<4x1xi1> {
+func.func @fold_shape_cast_with_mask_leading_unit(%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>
+// CHECK: %[[VAL_3:.*]] = vector.create_mask %[[VAL_2]], %[[VAL_1]] : vector<4x1xi1>
+// CHECK: return %[[VAL_3]] : vector<4x1xi1>
+ %c1 = arith.constant 1 : index
+ %dim = tensor.dim %arg0, %c1 : tensor<1x?xf32>
+ %1 = vector.create_mask %c1, %c1, %dim, %c1 : vector<1x1x4x1xi1>
+ %2 = vector.shape_cast %1 : vector<1x1x4x1xi1> to vector<4x1xi1>
+ return %2 : vector<4x1xi1>
+}
+
+// -----
+
+// 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> {
+// CHECK-NOT: vector.shape_cast
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<1x?xf32>
+// CHECK: %[[VAL_3:.*]] = vector.create_mask %[[VAL_2]], %[[VAL_1]] : vector<[4]x1xi1>
+// CHECK: return %[[VAL_3]] : vector<[4]x1xi1>
+ %c1 = arith.constant 1 : index
+ %dim = tensor.dim %arg0, %c1 : tensor<1x?xf32>
+ %1 = vector.create_mask %c1, %c1, %dim, %c1 : vector<1x1x[4]x1xi1>
+ %2 = vector.shape_cast %1 : vector<1x1x[4]x1xi1> to vector<[4]x1xi1>
+ return %2 : vector<[4]x1xi1>
+}
+
+// -----
+
+// Check that scalable "1" (i.e. [1]) is not folded
+// CHECK-LABEL: func.func @fold_shape_cast_with_mask_leading_unit_scalable_one(
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x?xf32>) -> vector<[1]x1xi1> {
+func.func @fold_shape_cast_with_mask_leading_unit_scalable_one(%arg0: tensor<1x?xf32>) -> vector<[1]x1xi1>{
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<1x?xf32>
+// CHECK: %[[VAL_3:.*]] = vector.create_mask %[[VAL_2]], %[[VAL_1]] : vector<[1]x1xi1>
+// CHECK: return %[[VAL_3]] : vector<[1]x1xi1>
+ %c1 = arith.constant 1 : index
+ %dim = tensor.dim %arg0, %c1 : tensor<1x?xf32>
+ %1 = vector.create_mask %c1, %dim, %c1 : vector<1x[1]x1xi1>
+ %2 = vector.shape_cast %1 : vector<1x[1]x1xi1> to vector<[1]x1xi1>
+ return %2 : vector<[1]x1xi1>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @fold_shape_cast_with_constant_mask_leading_unit() -> vector<4xi1> {
+func.func @fold_shape_cast_with_constant_mask_leading_unit() -> vector<4xi1>{
+// CHECK-NOT: vector.shape_cast
+// CHECK: %[[VAL_0:.*]] = vector.constant_mask [1] : vector<4xi1>
+// CHECK: return %[[VAL_0]] : vector<4xi1>
+ %1 = vector.constant_mask [1, 1, 1] : vector<1x1x4xi1>
+ %2 = vector.shape_cast %1 : vector<1x1x4xi1> to vector<4xi1>
+ return %2 : vector<4xi1>
+}
+
+// -----
+
// TODO: This IR could be canonicalized but the canonicalization pattern is not
// smart enough. For now, just make sure that we do not crash.
>From e07a8d8bb87b441c634a439bc37206b72f00e6b8 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Wed, 2 Sep 2026 16:08:43 +0000
Subject: [PATCH 2/2] Address comments from Federico
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 183 ++++++++-------------
mlir/test/Dialect/Vector/canonicalize.mlir | 12 +-
2 files changed, 78 insertions(+), 117 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index dce252962e9de..ad5e21a20efb8 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6948,7 +6948,7 @@ OpFoldResult ShapeCastOp::fold(FoldAdaptor adaptor) {
namespace {
/// Helper function that computes a new vector type based on the input vector
-/// type by removing the trailing one dims:
+/// type by removing the _trailing_ unit dims:
///
/// vector<4x1x1xi1> --> vector<4x1xi1>
///
@@ -6975,9 +6975,9 @@ static VectorType trimTrailingUnitDims(VectorType oldType) {
}
/// Helper function that computes a new vector type based on the input vector
-/// type by removing the trailing one dims:
+/// type by removing the _leading_ unit dims:
///
-/// vector<4x1x1xi1> --> vector<4x1xi1>
+/// vector<1x1x4xi1> --> vector<4x1xi1>
///
static VectorType trimLeadingUnitDims(VectorType oldType) {
ArrayRef<int64_t> oldShape = oldType.getShape();
@@ -7002,19 +7002,31 @@ static VectorType trimLeadingUnitDims(VectorType oldType) {
return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
}
+enum class UnitDimSide { Leading, Trailing };
+
/// Folds qualifying shape_cast(create_mask) into a new create_mask
///
-/// Looks at `vector.shape_cast` Ops that simply "drop" the _leading_unit
-/// dimension. If the input vector comes from `vector.create_mask` for which
-/// the corresponding mask input value is 1 (e.g. `%c1` below), then it is safe
-/// to fold shape_cast into create_mask.
+/// Looks at `vector.shape_cast` Ops that simply "drop" either the _trailing_ or
+/// the _leading_unit dimension (configured with the template parameter). If the
+/// input vector comes from `vector.create_mask` for which the corresponding
+/// mask input value is 1 (e.g. `%c1` below), then it is safe to fold shape_cast
+/// into create_mask.
+///
+/// EX 1 (trailing unit dims)
+/// BEFORE:
+/// %1 = vector.create_mask %c1, %dim, %c1, %c1 : vector<1x[4]x1x1xi1>
+/// %2 = vector.shape_cast %1 : vector<1x[4]x1x1xi1> to vector<1x[4]xi1>
+/// AFTER:
+/// %0 = vector.create_mask %c1, %dim : vector<1x[4]xi1>
///
+/// EX 2 (leading unit dims)
/// BEFORE:
/// %1 = vector.create_mask %c1, %c1, %dim, %c1 : vector<1x1x[4]x1xi1>
/// %2 = vector.shape_cast %1 : vector<1x1x[4]x1xi1> to vector<[4]x1xi1>
/// AFTER:
/// %0 = vector.create_mask %c1, %dim : vector<[4]xi1>
-class ShapeCastCreateMaskFolderLeadingOneDim final
+template <UnitDimSide Side>
+class ShapeCastCreateMaskFolderBoundaryUnitDim final
: public OpRewritePattern<ShapeCastOp> {
public:
using Base::Base;
@@ -7030,107 +7042,49 @@ class ShapeCastCreateMaskFolderLeadingOneDim final
VectorType shapeOpResTy = shapeOp.getResultVectorType();
VectorType shapeOpSrcTy = shapeOp.getSourceVectorType();
- VectorType newVecType = trimLeadingUnitDims(shapeOpSrcTy);
- if (newVecType != shapeOpResTy)
- return rewriter.notifyMatchFailure(
- shapeOp, "Non-leading-unit-dim dropping shape_cast Op");
+ VectorType newVecType = (Side == UnitDimSide::Trailing)
+ ? trimTrailingUnitDims(shapeOpSrcTy)
+ : trimLeadingUnitDims(shapeOpSrcTy);
- auto numDimsToDrop = shapeOpSrcTy.getRank() - shapeOpResTy.getRank();
-
- // No unit dims to drop
- if (numDimsToDrop == 0)
- return rewriter.notifyMatchFailure(
- shapeOp, "Non-leading-unit-dim dropping shape_cast Op");
-
- if (createMaskOp) {
- auto maskOperands = createMaskOp.getOperands();
-
- for (int64_t dimIdx = 0; dimIdx < numDimsToDrop; dimIdx++) {
- auto constantOp =
- maskOperands[dimIdx].getDefiningOp<arith::ConstantIndexOp>();
- if (!constantOp || constantOp.value() != 1) {
- return failure();
- }
- }
-
- rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(
- shapeOp, shapeOpResTy, maskOperands.drop_front(numDimsToDrop));
-
- return success();
- }
-
- if (constantMaskOp) {
- auto maskDimSizes = constantMaskOp.getMaskDimSizes();
- auto numMaskOperands = maskDimSizes.size();
-
- // Check every mask dim size to see whether it can be dropped
- for (uint64_t i = 0; i < numMaskOperands; ++i) {
- if (maskDimSizes[i] != 1)
- return failure();
- }
-
- auto newMaskOperands = maskDimSizes.drop_front(numDimsToDrop);
- rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(shapeOp, shapeOpResTy,
- newMaskOperands);
- return success();
+ if (newVecType != shapeOpResTy) {
+ return (Side == UnitDimSide::Trailing)
+ ? rewriter.notifyMatchFailure(
+ shapeOp, "Non-trailing-unit-dim dropping shape_cast Op")
+ : rewriter.notifyMatchFailure(
+ shapeOp, "Non-leading-unit-dim dropping shape_cast Op");
}
- return failure();
- }
-};
-
-/// Folds qualifying shape_cast(create_mask) into a new create_mask
-///
-/// Looks at `vector.shape_cast` Ops that simply "drop" the _trailing_ unit
-/// dimension. If the input vector comes from `vector.create_mask` for which
-/// the corresponding mask input value is 1 (e.g. `%c1` below), then it is safe
-/// to fold shape_cast into create_mask.
-///
-/// BEFORE:
-/// %1 = vector.create_mask %c1, %dim, %c1, %c1 : vector<1x[4]x1x1xi1>
-/// %2 = vector.shape_cast %1 : vector<1x[4]x1x1xi1> to vector<1x[4]xi1>
-/// AFTER:
-/// %0 = vector.create_mask %c1, %dim : vector<1x[4]xi1>
-class ShapeCastCreateMaskFolderTrailingOneDim final
- : public OpRewritePattern<ShapeCastOp> {
-public:
- using Base::Base;
-
- LogicalResult matchAndRewrite(ShapeCastOp shapeOp,
- PatternRewriter &rewriter) const override {
- Value shapeOpSrc = shapeOp->getOperand(0);
- auto createMaskOp = shapeOpSrc.getDefiningOp<vector::CreateMaskOp>();
- auto constantMaskOp = shapeOpSrc.getDefiningOp<vector::ConstantMaskOp>();
- if (!createMaskOp && !constantMaskOp)
- return failure();
-
- VectorType shapeOpResTy = shapeOp.getResultVectorType();
- VectorType shapeOpSrcTy = shapeOp.getSourceVectorType();
-
- VectorType newVecType = trimTrailingUnitDims(shapeOpSrcTy);
- if (newVecType != shapeOpResTy)
- return failure();
-
auto numDimsToDrop = shapeOpSrcTy.getRank() - shapeOpResTy.getRank();
// No unit dims to drop
- if (!numDimsToDrop)
- return failure();
+ if (!numDimsToDrop) {
+ return (Side == UnitDimSide::Trailing)
+ ? rewriter.notifyMatchFailure(
+ shapeOp, "Non-trailing-unit-dim dropping shape_cast Op")
+ : rewriter.notifyMatchFailure(
+ shapeOp, "Non-leading-unit-dim dropping shape_cast Op");
+ }
if (createMaskOp) {
auto maskOperands = createMaskOp.getOperands();
- auto numMaskOperands = maskOperands.size();
-
- // Check every mask dim size to see whether it can be dropped
- for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
- --i) {
- auto constant = maskOperands[i].getDefiningOp<arith::ConstantIndexOp>();
- if (!constant || (constant.value() != 1))
- return failure();
- }
- SmallVector<Value> newMaskOperands =
- maskOperands.drop_back(numDimsToDrop);
+ size_t numOperands = maskOperands.size();
+
+ auto maskOperandsToDrop =
+ (Side == UnitDimSide::Trailing)
+ ? maskOperands.take_back(numOperands - numDimsToDrop)
+ : maskOperands.take_front(numOperands - 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) {
+ auto cst = maskDim.getDefiningOp<arith::ConstantIndexOp>();
+ return !cst || (cst.value() != 1);
+ }))
+ return failure();
+ auto newMaskOperands = (Side == UnitDimSide::Trailing)
+ ? maskOperands.drop_back(numDimsToDrop)
+ : maskOperands.drop_front(numDimsToDrop);
rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(shapeOp, shapeOpResTy,
newMaskOperands);
return success();
@@ -7138,18 +7092,25 @@ class ShapeCastCreateMaskFolderTrailingOneDim final
if (constantMaskOp) {
auto maskDimSizes = constantMaskOp.getMaskDimSizes();
- auto numMaskOperands = maskDimSizes.size();
+ size_t numDims = maskDimSizes.size();
- // Check every mask dim size to see whether it can be dropped
- for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
- --i) {
- if (maskDimSizes[i] != 1)
- return failure();
- }
+ ArrayRef<int64_t> maskDimSizesToDrop =
+ (Side == UnitDimSide::Trailing)
+ ? maskDimSizes.take_back(numDims - numDimsToDrop)
+ : maskDimSizes.take_front(numDims - 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::any_of(maskDimSizesToDrop,
+ [](int64_t dim) { return dim != 1; }))
+ return failure();
- auto newMaskOperands = maskDimSizes.drop_back(numDimsToDrop);
+ ArrayRef<int64_t> newMaskDimSizes =
+ (Side == UnitDimSide::Trailing)
+ ? maskDimSizes.drop_back(numDimsToDrop)
+ : maskDimSizes.drop_front(numDimsToDrop);
rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(shapeOp, shapeOpResTy,
- newMaskOperands);
+ newMaskDimSizes);
return success();
}
@@ -7260,9 +7221,9 @@ class FoldShapeCastOfFromElements final : public OpRewritePattern<ShapeCastOp> {
void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<ShapeCastCreateMaskFolderTrailingOneDim,
- ShapeCastCreateMaskFolderLeadingOneDim, ShapeCastBroadcastFolder,
- FoldShapeCastOfFromElements>(context);
+ results.add<ShapeCastCreateMaskFolderBoundaryUnitDim<UnitDimSide::Leading>,
+ ShapeCastCreateMaskFolderBoundaryUnitDim<UnitDimSide::Trailing>,
+ ShapeCastBroadcastFolder, FoldShapeCastOfFromElements>(context);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index d09a721e38692..8861ef3822b04 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3464,7 +3464,7 @@ func.func @all_true_vector_mask_no_result(%a : vector<3x4xf32>, %m : memref<3x4x
// -----
// +---------------------------------------------------------------------------
-// Tests for ShapeCastCreateMaskFolderTrailingOneDim
+// Tests for ShapeCastCreateMaskFolderBoundaryUnitDim<UnitDimSide::Trailing>
// +---------------------------------------------------------------------------
// CHECK-LABEL: func.func @fold_shape_cast_with_mask_trailing_unit(
@@ -3521,9 +3521,9 @@ func.func @fold_shape_cast_with_mask_trailing_unit_scalable_one(%arg0: tensor<1x
// CHECK-LABEL: func.func @fold_shape_cast_with_constant_mask_trailing_unit() -> vector<4xi1> {
func.func @fold_shape_cast_with_constant_mask_trailing_unit() -> vector<4xi1>{
// CHECK-NOT: vector.shape_cast
-// CHECK: %[[VAL_0:.*]] = vector.constant_mask [1] : vector<4xi1>
+// CHECK: %[[VAL_0:.*]] = vector.constant_mask [3] : vector<4xi1>
// CHECK: return %[[VAL_0]] : vector<4xi1>
- %1 = vector.constant_mask [1, 1, 1] : vector<4x1x1xi1>
+ %1 = vector.constant_mask [3, 1, 1] : vector<4x1x1xi1>
%2 = vector.shape_cast %1 : vector<4x1x1xi1> to vector<4xi1>
return %2 : vector<4xi1>
}
@@ -3531,7 +3531,7 @@ func.func @fold_shape_cast_with_constant_mask_trailing_unit() -> vector<4xi1>{
// -----
// +---------------------------------------------------------------------------
-// Tests for CastAwayTransferWriteLeadingOneDim
+// Tests for ShapeCastCreateMaskFolderBoundaryUnitDim<UnitDimSide::Leading>
// +---------------------------------------------------------------------------
// CHECK-LABEL: func.func @fold_shape_cast_with_mask_leading_unit(
@@ -3588,9 +3588,9 @@ func.func @fold_shape_cast_with_mask_leading_unit_scalable_one(%arg0: tensor<1x?
// CHECK-LABEL: func.func @fold_shape_cast_with_constant_mask_leading_unit() -> vector<4xi1> {
func.func @fold_shape_cast_with_constant_mask_leading_unit() -> vector<4xi1>{
// CHECK-NOT: vector.shape_cast
-// CHECK: %[[VAL_0:.*]] = vector.constant_mask [1] : vector<4xi1>
+// CHECK: %[[VAL_0:.*]] = vector.constant_mask [3] : vector<4xi1>
// CHECK: return %[[VAL_0]] : vector<4xi1>
- %1 = vector.constant_mask [1, 1, 1] : vector<1x1x4xi1>
+ %1 = vector.constant_mask [1, 1, 3] : vector<1x1x4xi1>
%2 = vector.shape_cast %1 : vector<1x1x4xi1> to vector<4xi1>
return %2 : vector<4xi1>
}
More information about the Mlir-commits
mailing list