[Mlir-commits] [mlir] [mlir][tosa] Fold inverse reshapes around elementwise casts (PR #236386)
Cathal Corbett
llvmlistbot at llvm.org
Wed Oct 28 04:47:48 PDT 2026
https://github.com/catcor01 updated https://github.com/llvm/llvm-project/pull/236386
>From 85527209aa16a85a174dd3b5429e7342d1c52349 Mon Sep 17 00:00:00 2001
From: Cathal Corbett <cathal.corbett at arm.com>
Date: Mon, 26 Oct 2026 09:11:57 +0100
Subject: [PATCH] [mlir][tosa] Fold inverse reshapes around elementwise casts
Canonicalize reshape -> cast -> reshape to a single cast when the final
shape matches the original input shape. The cast retains its attributes,
location and element conversion.
Restrict the pattern to statically shaped tensors with ordinary integer
or floating-point elements. Block-scaled casts depend on the last
dimension's block grouping and are excluded. Require the intermediate
cast to have a single use to avoid duplicating the conversion.
Add regression coverage for floating-point and integer conversions,
attribute preservation, different output shapes, shared casts, dynamic
shapes and block-scaled casts.
---
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td | 1 +
.../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 43 +++++++++
mlir/test/Dialect/Tosa/reshape-cast.mlir | 92 +++++++++++++++++++
3 files changed, 136 insertions(+)
create mode 100644 mlir/test/Dialect/Tosa/reshape-cast.mlir
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 4458dae5f06e8e..a315e2fc3ec79e 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -2157,6 +2157,7 @@ def Tosa_ReshapeOp : Tosa_InferTensorTypeOp<"reshape", [Pure]> {
No data conversion happens during a reshape operation.
}];
+ let hasCanonicalizer = 1;
let hasFolder = 1;
let hasVerifier = 1;
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 9853c9d1465ebd..1e57e0db55a0ac 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -2128,6 +2128,49 @@ REDUCE_FOLDER(ReduceProductOp)
REDUCE_FOLDER(ReduceSumOp)
#undef REDUCE_FOLDER
+namespace {
+// Cancel inverse reshapes around an elementwise cast, retaining its conversion.
+struct FoldReshapesAroundCast : public OpRewritePattern<tosa::ReshapeOp> {
+ using OpRewritePattern<tosa::ReshapeOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::ReshapeOp op,
+ PatternRewriter &rewriter) const override {
+ auto castOp = op.getInput1().getDefiningOp<tosa::CastOp>();
+ if (!castOp || !castOp->hasOneUse())
+ return failure();
+ auto innerReshape = castOp.getInput().getDefiningOp<tosa::ReshapeOp>();
+ if (!innerReshape)
+ return failure();
+
+ auto inputType =
+ dyn_cast<RankedTensorType>(innerReshape.getInput1().getType());
+ auto intermediateType = dyn_cast<RankedTensorType>(innerReshape.getType());
+ auto outputType = dyn_cast<RankedTensorType>(op.getType());
+ if (!inputType || !intermediateType || !outputType ||
+ !inputType.hasStaticShape() || !intermediateType.hasStaticShape() ||
+ !outputType.hasStaticShape() ||
+ inputType.getShape() != outputType.getShape())
+ return failure();
+
+ // Block-scaled casts depend on the last dimension's block grouping.
+ if (!isa<IntegerType, FloatType>(inputType.getElementType()) ||
+ !isa<IntegerType, FloatType>(outputType.getElementType()))
+ return failure();
+
+ auto replacement = tosa::CastOp::create(
+ rewriter, castOp.getLoc(), outputType, innerReshape.getInput1());
+ replacement->setAttrs(castOp->getAttrs());
+ rewriter.replaceOp(op, replacement.getOutput());
+ return success();
+ }
+};
+} // namespace
+
+void ReshapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
+ MLIRContext *context) {
+ results.add<FoldReshapesAroundCast>(context);
+}
+
OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
auto inputTy = llvm::dyn_cast<RankedTensorType>(getInput1().getType());
auto outputTy = llvm::dyn_cast<RankedTensorType>(getType());
diff --git a/mlir/test/Dialect/Tosa/reshape-cast.mlir b/mlir/test/Dialect/Tosa/reshape-cast.mlir
new file mode 100644
index 00000000000000..58c3d1f99ad7a5
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/reshape-cast.mlir
@@ -0,0 +1,92 @@
+// RUN: mlir-opt --split-input-file --canonicalize="test-convergence" %s | FileCheck %s
+
+// CHECK-LABEL: @inverse_reshapes_bf16
+func.func @inverse_reshapes_bf16(%arg0: tensor<1x4x8xf32>) -> tensor<1x4x8xbf16> {
+ // CHECK-NEXT: %[[CAST:.*]] = tosa.cast %arg0 {test.tag = "keep"} : (tensor<1x4x8xf32>) -> tensor<1x4x8xbf16>
+ // CHECK-NEXT: return %[[CAST]] : tensor<1x4x8xbf16>
+ %s2 = tosa.const_shape {values = dense<[4, 8]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %s3 = tosa.const_shape {values = dense<[1, 4, 8]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %r = tosa.reshape %arg0, %s2 : (tensor<1x4x8xf32>, !tosa.shape<2>) -> tensor<4x8xf32>
+ %c = tosa.cast %r {test.tag = "keep"} : (tensor<4x8xf32>) -> tensor<4x8xbf16>
+ %o = tosa.reshape %c, %s3 : (tensor<4x8xbf16>, !tosa.shape<3>) -> tensor<1x4x8xbf16>
+ return %o : tensor<1x4x8xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @inverse_reshapes_integer
+func.func @inverse_reshapes_integer(%arg0: tensor<2x4xi8>) -> tensor<2x4xi32> {
+ // CHECK-NEXT: %[[CAST:.*]] = tosa.cast %arg0 : (tensor<2x4xi8>) -> tensor<2x4xi32>
+ // CHECK-NEXT: return %[[CAST]] : tensor<2x4xi32>
+ %s1 = tosa.const_shape {values = dense<8> : tensor<1xindex>} : () -> !tosa.shape<1>
+ %s2 = tosa.const_shape {values = dense<[2, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %r = tosa.reshape %arg0, %s1 : (tensor<2x4xi8>, !tosa.shape<1>) -> tensor<8xi8>
+ %c = tosa.cast %r : (tensor<8xi8>) -> tensor<8xi32>
+ %o = tosa.reshape %c, %s2 : (tensor<8xi32>, !tosa.shape<2>) -> tensor<2x4xi32>
+ return %o : tensor<2x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @different_output_shape
+func.func @different_output_shape(%arg0: tensor<1x4x8xf32>) -> tensor<2x16xbf16> {
+ // CHECK: %[[RESHAPE:.*]] = tosa.reshape %arg0
+ // CHECK: %[[CAST:.*]] = tosa.cast %[[RESHAPE]]
+ // CHECK: %[[OUTPUT:.*]] = tosa.reshape %[[CAST]]
+ // CHECK: return %[[OUTPUT]]
+ %s1 = tosa.const_shape {values = dense<32> : tensor<1xindex>} : () -> !tosa.shape<1>
+ %s2 = tosa.const_shape {values = dense<[2, 16]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %r = tosa.reshape %arg0, %s1 : (tensor<1x4x8xf32>, !tosa.shape<1>) -> tensor<32xf32>
+ %c = tosa.cast %r : (tensor<32xf32>) -> tensor<32xbf16>
+ %o = tosa.reshape %c, %s2 : (tensor<32xbf16>, !tosa.shape<2>) -> tensor<2x16xbf16>
+ return %o : tensor<2x16xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @shared_cast
+func.func @shared_cast(%arg0: tensor<1x4x8xf32>) -> (tensor<4x8xbf16>, tensor<1x4x8xbf16>) {
+ // CHECK: %[[RESHAPE:.*]] = tosa.reshape %arg0
+ // CHECK: %[[CAST:.*]] = tosa.cast %[[RESHAPE]]
+ // CHECK-NOT: tosa.cast
+ // CHECK: %[[OUTPUT:.*]] = tosa.reshape %[[CAST]]
+ // CHECK-NEXT: return %[[CAST]], %[[OUTPUT]]
+ %s2 = tosa.const_shape {values = dense<[4, 8]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %s3 = tosa.const_shape {values = dense<[1, 4, 8]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %r = tosa.reshape %arg0, %s2 : (tensor<1x4x8xf32>, !tosa.shape<2>) -> tensor<4x8xf32>
+ %c = tosa.cast %r : (tensor<4x8xf32>) -> tensor<4x8xbf16>
+ %o = tosa.reshape %c, %s3 : (tensor<4x8xbf16>, !tosa.shape<3>) -> tensor<1x4x8xbf16>
+ return %c, %o : tensor<4x8xbf16>, tensor<1x4x8xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @dynamic_shapes
+func.func @dynamic_shapes(%arg0: tensor<1x?x4xf32>) -> tensor<1x?x4xbf16> {
+ // CHECK: %[[RESHAPE:.*]] = tosa.reshape %arg0
+ // CHECK: %[[CAST:.*]] = tosa.cast %[[RESHAPE]]
+ // CHECK: %[[OUTPUT:.*]] = tosa.reshape %[[CAST]]
+ // CHECK: return %[[OUTPUT]]
+ %s2 = tosa.const_shape {values = dense<[-1, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %s3 = tosa.const_shape {values = dense<[1, -1, 4]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %r = tosa.reshape %arg0, %s2 : (tensor<1x?x4xf32>, !tosa.shape<2>) -> tensor<?x4xf32>
+ %c = tosa.cast %r : (tensor<?x4xf32>) -> tensor<?x4xbf16>
+ %o = tosa.reshape %c, %s3 : (tensor<?x4xbf16>, !tosa.shape<3>) -> tensor<1x?x4xbf16>
+ return %o : tensor<1x?x4xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @block_scaled_cast
+func.func @block_scaled_cast(%arg0: tensor<64x32xf32>) -> tensor<64x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>> {
+ // CHECK: %[[RESHAPE:.*]] = tosa.reshape %arg0
+ // CHECK: %[[CAST:.*]] = tosa.cast %[[RESHAPE]]
+ // CHECK: %[[OUTPUT:.*]] = tosa.reshape %[[CAST]]
+ // CHECK: return %[[OUTPUT]]
+ %s = tosa.const_shape {values = dense<[32, 64]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %original = tosa.const_shape {values = dense<[64, 32]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %r = tosa.reshape %arg0, %s : (tensor<64x32xf32>, !tosa.shape<2>) -> tensor<32x64xf32>
+ %c = tosa.cast %r : (tensor<32x64xf32>) -> tensor<32x64x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
+ %o = tosa.reshape %c, %original : (tensor<32x64x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>, !tosa.shape<2>) -> tensor<64x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
+ return %o : tensor<64x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
+}
More information about the Mlir-commits
mailing list