[Mlir-commits] [mlir] [mlir] Convert `expand_shape` to more static form (PR #112265)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 21 07:53:52 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tensor
Author: Ian Wood (IanWood1)
<details>
<summary>Changes</summary>
Initially, my idea was to sink `tensor.cast` op through `tensor.expand_shape` ops when it makes the expand op more static. But then I realized that the SSA `output_shape` operands are capturing shape info that can't be propagated. >From the commit's description:
>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).
My main concern here is that the generated casts are not guaranteed to fold with other casts. This is somewhat similar to what `linalg` does where it introduces casts before operands when their shapes are inferred. But, I'm not sure if this is suited for a canonicalization pattern (I could just add a check to make sure the pattern would fold >1 adjacent cast).
Also, the opposite might happen as well. Where `output_sizes` are unknown constants but there is a `tensor.cast` consumer that has the output size information.
Sidenote: I disabled CI because `drop-unit-extent-dims.mlir` will fail. There is a cast that gets converted to a static form. Just wanted to wait for review to determine if a fix is needed.
---
Full diff: https://github.com/llvm/llvm-project/pull/112265.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+79-1)
- (modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+54)
``````````diff
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 4d6c5965c4fcc3..ee0e8c2d201226 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -24,6 +24,7 @@
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallBitVector.h"
@@ -1982,6 +1983,83 @@ struct FoldDimOfCollapseShape : public OpRewritePattern<DimOp> {
return success();
}
};
+
+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();
+
+ const ArrayRef<int64_t> castSrcShape =
+ castOp.getSource().getType().getShape();
+ const SmallVector<ReassociationIndices, 4> reassoc =
+ expandOp.getReassociationIndices();
+
+ SmallVector<int64_t> newOutputShape(expandOp.getResultType().getShape());
+ SmallVector<Value> dynamicOutputShape;
+ auto outputIt = expandOp.getOutputShape().begin();
+
+ for (const auto &[inputDim, innerReassoc] : llvm::enumerate(reassoc)) {
+ for (const uint64_t outDim : innerReassoc) {
+ if (!ShapedType::isDynamic(newOutputShape[outDim]))
+ continue;
+
+ // If the cast's src type is dynamic, don't infer any of the
+ // corresponding expanded dimensions. `tensor.expand_shape` requires at
+ // least one of the expanded dimensions to be dynamic if the input is
+ // dynamic.
+ Value val = *outputIt;
+ ++outputIt;
+ if (ShapedType::isDynamic(castSrcShape[inputDim])) {
+ dynamicOutputShape.push_back(val);
+ continue;
+ }
+
+ APInt cst;
+ if (matchPattern(val, m_ConstantInt(&cst))) {
+ newOutputShape[outDim] = 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<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;
+ }
+ }
+
+ 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>(
+ expandOp.getLoc(), outputType, inputCast.getResult(),
+ expandOp.getReassociationIndices(), outputOfr);
+ rewriter.replaceOpWithNewOp<CastOp>(expandOp, expandOp.getType(),
+ newExpand.getResult());
+ return success();
+ }
+};
} // namespace
void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
@@ -1989,7 +2067,7 @@ void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
results.add<
ComposeReassociativeReshapeOps<ExpandShapeOp, ReshapeOpKind::kExpand>,
ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>,
- 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 0aa2d33ef17ed4..63f394a14d3899 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2718,3 +2718,57 @@ 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<10x1x10xf32> {
+ %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>
+ %2 = tensor.cast %1 : tensor<?x?x?xf32> to tensor<10x1x10xf32>
+ return %2 : tensor<10x1x10xf32>
+}
+// CHECK-LABEL: func.func @fold_expand_of_cast
+// CHECK: %[[RES:.+]] = tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]] output_shape [10, 1, 10]
+// CHECK: return %[[RES]]
+
+// -----
+
+func.func @sink_expand_of_cast(%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 @sink_expand_of_cast
+// CHECK-DAG: %[[C10:.*]] = arith.constant 10
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1
+// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]]
+// CHECK-SAME: output_shape [%[[C10]], %[[C1]], 10]
+// CHECK: %[[RES:.+]] = tensor.cast %[[EXPAND]]
+// CHECK: return %[[RES]]
+
+// -----
+
+func.func @partial_sink_expand_of_cast(%arg0 : tensor<10x10xf32>, %arg1 : index, %arg2 : index)
+ -> tensor<?x?x?xf32> {
+ %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 [%arg1, %arg2, %c10]
+ : tensor<?x?xf32> into tensor<?x?x?xf32>
+ return %1 : tensor<?x?x?xf32>
+}
+// CHECK-LABEL: func.func @partial_sink_expand_of_cast
+// CHECK: %[[CAST:.+]] = tensor.cast
+// CHECK-SAME: tensor<10x10xf32> to tensor<?x10xf32>
+// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]]
+// CHECK-SAME: output_shape [%{{.*}}, %{{.*}}, 10]
+// CHECK: %[[RES:.+]] = tensor.cast %[[EXPAND]]
+// CHECK-SAME: tensor<?x?x10xf32> to tensor<?x?x?xf32>
+// CHECK: return %[[RES]]
``````````
</details>
https://github.com/llvm/llvm-project/pull/112265
More information about the Mlir-commits
mailing list