[Mlir-commits] [mlir] [mlir] Fold expand of cast (PR #112265)
Ian Wood
llvmlistbot at llvm.org
Mon Oct 14 19:22:47 PDT 2024
https://github.com/IanWood1 updated https://github.com/llvm/llvm-project/pull/112265
>From 67649194b893a9a017082964d285056f4c6656fa Mon Sep 17 00:00:00 2001
From: Ian Wood <ianwood2024 at u.northwestern.edu>
Date: Mon, 14 Oct 2024 13:29:42 -0500
Subject: [PATCH 1/2] [mlir] Fold expand of cast
Sink tensor.cast op through tensor.expand_shape ops when it makes the
expand op more static. This allows for other ops further down infer
their shapes.
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 31 +++++++++++++++++++++-
mlir/test/Dialect/Tensor/canonicalize.mlir | 14 ++++++++++
2 files changed, 44 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 4d6c5965c4fcc3..9be647f687e600 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -1982,6 +1982,35 @@ struct FoldDimOfCollapseShape : public OpRewritePattern<DimOp> {
return success();
}
};
+
+struct FoldExpandOfCast : public OpRewritePattern<ExpandShapeOp> {
+ using OpRewritePattern<ExpandShapeOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ExpandShapeOp expandOp,
+ PatternRewriter &rewriter) const override {
+ auto castOp = expandOp.getSrc().getDefiningOp<CastOp>();
+ if (!canFoldIntoConsumerOp(castOp))
+ return failure();
+
+ SmallVector<OpFoldResult> outputOfr =
+ getMixedValues(expandOp.getResultType().getShape(),
+ expandOp.getOutputShape(), rewriter);
+ std::optional<SmallVector<int64_t>> constantOutputShape =
+ getConstantIntValues(outputOfr);
+ if (!constantOutputShape.has_value()) {
+ return failure();
+ }
+ auto newType = RankedTensorType::get(
+ constantOutputShape.value(), expandOp.getSrcType().getElementType());
+
+ auto newExpand = rewriter.create<ExpandShapeOp>(
+ castOp.getLoc(), newType, castOp.getSource(),
+ expandOp.getReassociationIndices());
+ rewriter.replaceOpWithNewOp<CastOp>(expandOp, expandOp.getType(),
+ newExpand.getResult());
+ return success();
+ }
+};
} // namespace
void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
@@ -1989,7 +2018,7 @@ void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
results.add<
ComposeReassociativeReshapeOps<ExpandShapeOp, ReshapeOpKind::kExpand>,
ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>,
- FoldReshapeWithConstant<ExpandShapeOp>,
+ FoldExpandOfCast, FoldReshapeWithConstant<ExpandShapeOp>,
FoldReshapeWithSplat<ExpandShapeOp>,
FoldReshapeWithFromElements<ExpandShapeOp>, FoldDimOfExpandShape,
FoldDimOfCollapseShape>(context);
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 0aa2d33ef17ed4..1509d26151119d 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2718,3 +2718,17 @@ func.func @pack_dont_drop_attributes(%arg0: tensor<?x?x?xf16>, %arg1: tensor<128
%pack = tensor.pack %arg0 padding_value(%cst : f16) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [16, 1] into %arg1 {test_attr} : tensor<?x?x?xf16> -> tensor<128x?x100x16x1xf16>
return %pack : tensor<128x?x100x16x1xf16>
}
+
+// -----
+
+func.func @fold_expand_of_cast(%arg0 : tensor<10x10xf32>)
+ -> tensor<?x?x?xf32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %0 = tensor.cast %arg0 : tensor<10x10xf32> to tensor<?x?xf32>
+ %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [%c10, %c1, %c10]
+ : tensor<?x?xf32> into tensor<?x?x?xf32>
+ return %1 : tensor<?x?x?xf32>
+}
+// CHECK-LABEL: func.func @fold_expand_of_cast
+// CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]] output_shape [10, 1, 10]
>From 3f4c7bb63fc16dcfa809ae917d039c25782c7cb9 Mon Sep 17 00:00:00 2001
From: Ian Wood <ianwood2024 at u.northwestern.edu>
Date: Tue, 15 Oct 2024 02:19:37 +0000
Subject: [PATCH 2/2] Convert to static expand_shape
When output_sizes can be determined, convert to a static expand_shape
op and insert cast ops. The top cast will be (dynamic -> static) allowing
it to be propagated upwards and the bottom will be (static -> dynamic)
allowing it to propagate down (or cancel with adjacent tensor.cast ops).
[skip ci]
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 61 ++++++++++++++++------
mlir/test/Dialect/Tensor/canonicalize.mlir | 14 +++++
2 files changed, 60 insertions(+), 15 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 9be647f687e600..96384385b6a060 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -1983,29 +1983,60 @@ struct FoldDimOfCollapseShape : public OpRewritePattern<DimOp> {
}
};
-struct FoldExpandOfCast : public OpRewritePattern<ExpandShapeOp> {
+struct ConvertToStaticExpandShape : public OpRewritePattern<ExpandShapeOp> {
using OpRewritePattern<ExpandShapeOp>::OpRewritePattern;
LogicalResult matchAndRewrite(ExpandShapeOp expandOp,
PatternRewriter &rewriter) const override {
- auto castOp = expandOp.getSrc().getDefiningOp<CastOp>();
- if (!canFoldIntoConsumerOp(castOp))
- return failure();
+ SmallVector<int64_t> newOutputShape(expandOp.getResultType().getShape());
+ SmallVector<Value> dynamicOutputShape;
+ auto outputIt = expandOp.getOutputShape().begin();
+ for (auto [i, staticShape] : llvm::enumerate(newOutputShape)) {
+ if (!ShapedType::isDynamic(staticShape))
+ continue;
- SmallVector<OpFoldResult> outputOfr =
- getMixedValues(expandOp.getResultType().getShape(),
- expandOp.getOutputShape(), rewriter);
- std::optional<SmallVector<int64_t>> constantOutputShape =
- getConstantIntValues(outputOfr);
- if (!constantOutputShape.has_value()) {
+ APInt cst;
+ Value val = *outputIt;
+ ++outputIt;
+ if (matchPattern(val, m_ConstantInt(&cst))) {
+ newOutputShape[i] = cst.getSExtValue();
+ } else {
+ dynamicOutputShape.push_back(val);
+ }
+ }
+
+ // Couldn't match any values, nothing to change
+ if (expandOp.getOutputShape().size() == dynamicOutputShape.size())
return failure();
+
+ // Calculate the input shape from the output
+ SmallVector<ReassociationIndices, 4> reassoc =
+ expandOp.getReassociationIndices();
+ SmallVector<int64_t> newInputShape(expandOp.getSrcType().getRank(), 1l);
+ for (uint64_t inDim = 0; inDim < newInputShape.size(); inDim++) {
+ for (auto outDim : reassoc[inDim]) {
+ auto ofr = newOutputShape[outDim];
+ if (ShapedType::isDynamic(ofr)) {
+ newInputShape[inDim] = ShapedType::kDynamic;
+ break;
+ }
+ newInputShape[inDim] *= ofr;
+ }
}
- auto newType = RankedTensorType::get(
- constantOutputShape.value(), expandOp.getSrcType().getElementType());
+ // `inputCast` can be propagated up and the final cast can be propagated
+ // down.
+ SmallVector<OpFoldResult> outputOfr =
+ getMixedValues(newOutputShape, dynamicOutputShape, rewriter);
+ auto inputType = RankedTensorType::get(
+ newInputShape, expandOp.getSrcType().getElementType());
+ auto outputType = RankedTensorType::get(
+ newOutputShape, expandOp.getSrcType().getElementType());
+ auto inputCast = rewriter.create<CastOp>(expandOp.getLoc(), inputType,
+ expandOp.getSrc());
auto newExpand = rewriter.create<ExpandShapeOp>(
- castOp.getLoc(), newType, castOp.getSource(),
- expandOp.getReassociationIndices());
+ expandOp.getLoc(), outputType, inputCast.getResult(),
+ expandOp.getReassociationIndices(), outputOfr);
rewriter.replaceOpWithNewOp<CastOp>(expandOp, expandOp.getType(),
newExpand.getResult());
return success();
@@ -2018,7 +2049,7 @@ void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
results.add<
ComposeReassociativeReshapeOps<ExpandShapeOp, ReshapeOpKind::kExpand>,
ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>,
- FoldExpandOfCast, FoldReshapeWithConstant<ExpandShapeOp>,
+ ConvertToStaticExpandShape, FoldReshapeWithConstant<ExpandShapeOp>,
FoldReshapeWithSplat<ExpandShapeOp>,
FoldReshapeWithFromElements<ExpandShapeOp>, FoldDimOfExpandShape,
FoldDimOfCollapseShape>(context);
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 1509d26151119d..52dcfd1d427d93 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2732,3 +2732,17 @@ func.func @fold_expand_of_cast(%arg0 : tensor<10x10xf32>)
}
// CHECK-LABEL: func.func @fold_expand_of_cast
// CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]] output_shape [10, 1, 10]
+
+// -----
+
+func.func @fold_expand_of_cast_dynamic(%arg0 : tensor<?x10xf32>)
+ -> tensor<?x?x?xf32> {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %0 = tensor.cast %arg0 : tensor<?x10xf32> to tensor<?x?xf32>
+ %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [%c10, %c1, %c10]
+ : tensor<?x?xf32> into tensor<?x?x?xf32>
+ return %1 : tensor<?x?x?xf32>
+}
+// CHECK-LABEL: func.func @fold_expand_of_cast_dynamic
+// CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]] output_shape [10, 1, 10]
More information about the Mlir-commits
mailing list