[Mlir-commits] [mlir] [mlir][vector] Add `ShapeCastCreateMaskFolderLeadingOneDim` (PR #219494)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 28 07:47:07 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Andrzej Warzyński (banach-space)
<details>
<summary>Changes</summary>
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`.
---
Full diff: https://github.com/llvm/llvm-project/pull/219494.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+111-6)
- (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+79-8)
``````````diff
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.
``````````
</details>
https://github.com/llvm/llvm-project/pull/219494
More information about the Mlir-commits
mailing list