[Flow] Fold static flow.tensor.empty shapes into dims (#24622)
This PR Adds a canonicalization pattern for `flow.tensor.empty` that
folds constant dynamic dimension operands into the static result shape.
**Example :**
%0 = flow.tensor.empty : tensor<?x4xf32>{%c8}
**canonicalizes to:**
%0 = flow.tensor.empty : tensor<8x4xf32>
%cast = tensor.cast %0 : tensor<8x4xf32> to tensor<?x4xf32>
The cast preserves the original result type for existing users while
allowing downstream canonicalization to use the refined static shape.
---------
Signed-off-by: LekkalaSravya3 <lekkala.sravya@multicorewareinc.com>diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
index 76b0315..9e46d78 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp
@@ -743,9 +743,62 @@
// flow.tensor.empty
//===----------------------------------------------------------------------===//
+namespace {
+
+// Folds constant dynamic dimension operands of a flow.tensor.empty into the
+// static result shape, e.g. `flow.tensor.empty : tensor<?xi32>{%c4}` becomes
+// `flow.tensor.empty : tensor<4xi32>`. A cast is inserted to preserve the
+// original (dynamic) result type for existing users; it folds away once
+// those users can accept the refined static type.
+struct FoldTensorEmptyConstantDims : public OpRewritePattern<TensorEmptyOp> {
+ using Base::Base;
+ LogicalResult matchAndRewrite(TensorEmptyOp op,
+ PatternRewriter &rewriter) const override {
+ auto resultType = cast<RankedTensorType>(op.getResult().getType());
+ ValueRange dynamicDims = op.getResultDims();
+
+ SmallVector<int64_t> newShape(resultType.getShape());
+ // Collect only dimensions that remain dynamic after folding constants.
+ SmallVector<Value> newDynamicDims;
+ unsigned dimIdx = 0;
+ bool didFold = false;
+ for (unsigned i = 0, e = resultType.getRank(); i < e; ++i) {
+ // Skip dimensions that are already static.
+ if (!resultType.isDynamicDim(i)) {
+ continue;
+ }
+ // Get the SSA operand corresponding to dynamic dimension.
+ Value dim = dynamicDims[dimIdx++];
+ std::optional<int64_t> constantDim = getConstantIntValue(dim);
+ if (constantDim && *constantDim >= 0) {
+ // Fold the constant value into the tensor's static shape.
+ newShape[i] = *constantDim;
+ didFold = true;
+ } else {
+ newDynamicDims.push_back(
+ dim); // This dim stays dynamic, so preserve its SSA operand.
+ }
+ }
+ if (!didFold) {
+ return rewriter.notifyMatchFailure(op, "no constant dynamic dims");
+ }
+
+ auto newType = RankedTensorType::get(newShape, resultType.getElementType(),
+ resultType.getEncoding());
+ Value newEmpty =
+ TensorEmptyOp::create(rewriter, op.getLoc(), newType, newDynamicDims);
+ // Preserve the original (dynamic) result type for existing users via a
+ // cast; it folds away once those users accept the refined static type.
+ rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, newEmpty);
+ return success();
+ }
+};
+
+} // namespace
+
void TensorEmptyOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- // TODO(benvanik): fold static shapes into dims.
+ results.insert<FoldTensorEmptyConstantDims>(context);
}
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir b/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir
index a04dea6..10a470b 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir
@@ -593,11 +593,55 @@
util.func public @sliceFromZeroElements(%arg0: tensor<0xi32>) -> tensor<?xi32> {
%c0 = arith.constant 0 : index
// CHECK-NOT: flow.tensor.slice
- // CHECK: %[[RET:.+]] = flow.tensor.empty : tensor<?xi32>{%c0}
+ // CHECK: %[[RET:.+]] = flow.tensor.empty : tensor<0xi32>
+ // CHECK: %[[CAST:.+]] = tensor.cast %[[RET]] : tensor<0xi32> to tensor<?xi32>
%0 = flow.tensor.slice %arg0[%c0 for %c0] : tensor<0xi32> -> tensor<?xi32>{%c0}
- // CHECK: util.return %[[RET]]
+ // CHECK: util.return %[[CAST]]
util.return %0 : tensor<?xi32>
}
+// -----
+
+// Folds a constant dynamic dimension operand into the static result shape.
+
+// CHECK-LABEL: @emptyFoldConstantDim
+util.func public @emptyFoldConstantDim() -> tensor<?x4xf32> {
+ %c8 = arith.constant 8 : index
+ // CHECK: %[[EMPTY:.+]] = flow.tensor.empty : tensor<8x4xf32>
+ // CHECK: %[[CAST:.+]] = tensor.cast %[[EMPTY]] : tensor<8x4xf32> to tensor<?x4xf32>
+ %0 = flow.tensor.empty : tensor<?x4xf32>{%c8}
+ // CHECK: util.return %[[CAST]]
+ util.return %0 : tensor<?x4xf32>
+}
+
+
+// -----
+
+// Folds only the constant dynamic dimension; the non-constant one stays dynamic.
+
+// CHECK-LABEL: @emptyFoldMixedDims
+// CHECK-SAME: (%[[DIM:.+]]: index)
+util.func public @emptyFoldMixedDims(%dim: index) -> tensor<?x8x?xf32> {
+ %c4 = arith.constant 4 : index
+ // CHECK: %[[EMPTY:.+]] = flow.tensor.empty : tensor<?x8x4xf32>{%[[DIM]]}
+ // CHECK: %[[CAST:.+]] = tensor.cast %[[EMPTY]] : tensor<?x8x4xf32> to tensor<?x8x?xf32>
+ %0 = flow.tensor.empty : tensor<?x8x?xf32>{%dim, %c4}
+ // CHECK: util.return %[[CAST]]
+ util.return %0 : tensor<?x8x?xf32>
+}
+
+// -----
+
+// Leaves non-constant dynamic dimensions untouched.
+
+// CHECK-LABEL: @emptyKeepDynamicDim
+// CHECK-SAME: (%[[DIM:.+]]: index)
+util.func public @emptyKeepDynamicDim(%dim: index) -> tensor<?x4xf32> {
+ // CHECK: %[[RET:.+]] = flow.tensor.empty : tensor<?x4xf32>{%[[DIM]]}
+ %0 = flow.tensor.empty : tensor<?x4xf32>{%dim}
+ // CHECK: util.return %[[RET]]
+ util.return %0 : tensor<?x4xf32>
+}
+
// -----
