[Mlir-commits] [mlir] [mlir][tosa] Fold unit-expanded reshapes in reduce-transposes (PR #203529)
Luke Hutton
llvmlistbot at llvm.org
Wed Jul 1 06:42:37 PDT 2026
https://github.com/lhutton1 updated https://github.com/llvm/llvm-project/pull/203529
>From 7dfb40eb9154c3c75843e6fbeb68caa8743025e8 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Wed, 10 Jun 2026 21:33:36 +0100
Subject: [PATCH 1/2] [mlir][tosa] Fold unit-expanded reshapes in
reduce-transposes
Extend tosa-reduce-transposes so transpose hosting can fold through
unit-expanded vector reshapes such as 1x...x1xC. In particular,
this improves transpose folding when hosting through broadcast binary
elementwise operations.
Change-Id: If2eb010ac322d5b9ef73f1fdda9de6049f2d2b2d
---
.../Tosa/Transforms/TosaReduceTransposes.cpp | 38 +++++++++++++------
.../Dialect/Tosa/tosa-reduce-transposes.mlir | 16 ++++++++
2 files changed, 43 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
index 65ef49bdc3077..0b4c0f04085e4 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
@@ -174,6 +174,23 @@ struct TosaReduceTransposes final
transposeDenseAttribute(DenseElementsAttr input, ArrayRef<int32_t> perms);
};
+// Check if shape is of the form 1x1x...x1xNx1x...x1x1 -> 1x1x...x1xNx1x...x1x1
+// Valid examples include:
+// - N -> 1x1xNx1
+// - Nx1x1x1 -> 1x1xNx1
+// - 1x1xNx1 -> 1x1xNx1
+static LogicalResult verifyUnitExpandedVectorShape(ArrayRef<int64_t> shape) {
+ bool nonUnitDimDetected = false;
+ for (const int64_t dim : shape) {
+ if (dim != 1) {
+ if (nonUnitDimDetected)
+ return failure();
+ nonUnitDimDetected = true;
+ }
+ }
+ return success();
+}
+
std::optional<DenseElementsAttr>
TosaReduceTransposes::transposeDenseAttribute(DenseElementsAttr input,
ArrayRef<int32_t> perms) {
@@ -393,24 +410,24 @@ std::optional<Value> TosaReduceTransposes::buildMappedToValue(
auto reshapeInputType =
llvm::dyn_cast<RankedTensorType>(reshapeOp.getInput1().getType());
auto reshapeInputShape = reshapeInputType.getShape();
- // want reshape N -> 1x1x...x1xNx1x...x1x1
- if (!reshapeInputType || reshapeInputShape.size() != 1)
+ if (!reshapeInputType)
return std::nullopt;
auto reshapeOutputType =
llvm::cast<RankedTensorType>(reshapeOutput.getType());
-
+ const ArrayRef<int64_t> reshapeOutputShape = reshapeOutputType.getShape();
// Instead of inserting a TransposeOp here, we check if we can fold it into
// the ReshapeOp. There is more complex cases where this is possible, and
// this check can be extended.
- // Checking if reshape is N -> 1x1x...x1xNx1x...x1x1
- auto shape = reshapeOutputType.getShape();
- size_t ones = llvm::count(shape, 1);
- // N == 1 and N != 1
- if (ones != shape.size() - 1 &&
- (ones != shape.size() || reshapeInputShape[0] != 1))
+ if (failed(verifyUnitExpandedVectorShape(reshapeInputShape)) ||
+ failed(verifyUnitExpandedVectorShape(reshapeOutputShape)))
return std::nullopt;
+ SmallVector<int64_t> mappedShape =
+ applyTOSAPermutation(reshapeOutputShape, hoistedPerms);
+ if (llvm::equal(mappedShape, reshapeInputShape))
+ return reshapeOp.getInput1();
+
// Do not insert a TransposeOp, instead we fold the reshape and its attribute.
llvm::SmallVector<int64_t> newShape;
if (!tosa::getConstShapeValues(reshapeOp.getShape().getDefiningOp(),
@@ -421,8 +438,7 @@ std::optional<Value> TosaReduceTransposes::buildMappedToValue(
ImplicitLocOpBuilder builder(reshapeOp.getLoc(), rewriter);
auto foldedReshape = ReshapeOp::create(
rewriter, reshapeOp.getLoc(),
- RankedTensorType::get(applyTOSAPermutation(shape, hoistedPerms),
- reshapeOutputType.getElementType()),
+ RankedTensorType::get(mappedShape, reshapeOutputType.getElementType()),
reshapeOp.getInput1(),
getTosaConstShape(builder, applyTOSAPermutation(llvm::ArrayRef(newShape),
hoistedPerms)));
diff --git a/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir b/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir
index b3f4260ede2f5..0ad3068011f7f 100644
--- a/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir
@@ -174,6 +174,22 @@ func.func @test_reshape_for_broadcast(%arg0: tensor<4x3x2xi32>) -> tensor<4x3x2x
// -----
+// CHECK-LABEL: @test_multi_dim_reshape_for_broadcast
+// CHECK-DAG: %[[SHAPE:.*]] = tosa.const_shape {values = dense<[4, 1, 1]> : tensor<3xindex>}
+// CHECK: %[[RESHAPE:.*]] = tosa.reshape %arg0, %[[SHAPE]] : (tensor<1x4x1xi32>, !tosa.shape<3>) -> tensor<4x1x1xi32>
+// CHECK: %[[ADD:.*]] = tosa.add %arg1, %[[RESHAPE]]
+// CHECK: return %[[ADD]]
+func.func @test_multi_dim_reshape_for_broadcast(%arg0: tensor<1x4x1xi32>, %arg1: tensor<4x3x2xi32>) -> tensor<4x3x2xi32> {
+ %shape = tosa.const_shape {values = dense<[1, 1, 4]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %reshape = tosa.reshape %arg0, %shape : (tensor<1x4x1xi32>, !tosa.shape<3>) -> tensor<1x1x4xi32>
+ %transpose0 = tosa.transpose %arg1 {perms = array<i32: 2, 1, 0>}: (tensor<4x3x2xi32>) -> tensor<2x3x4xi32>
+ %add = tosa.add %transpose0, %reshape : (tensor<2x3x4xi32>, tensor<1x1x4xi32>) -> tensor<2x3x4xi32>
+ %transpose1 = tosa.transpose %add {perms = array<i32: 2, 1, 0>}: (tensor<2x3x4xi32>) -> tensor<4x3x2xi32>
+ return %transpose1 : tensor<4x3x2xi32>
+}
+
+// -----
+
// COM: taken directly from ResNet18 translation.
// COM: changes: %74 as argument instead of result of conv2d
>From 87f51a82f73c4460670b267e3085c57a635c1073 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Wed, 1 Jul 2026 14:41:07 +0100
Subject: [PATCH 2/2] Address review comments
Change-Id: I571c79239a48824b5b816203465182d0f8fc3c3b
---
.../Dialect/Tosa/Transforms/TosaReduceTransposes.cpp | 11 +++--------
1 file changed, 3 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
index 0b4c0f04085e4..f025526d52d67 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
@@ -409,25 +409,20 @@ std::optional<Value> TosaReduceTransposes::buildMappedToValue(
auto reshapeOutput = reshapeOp.getOutput();
auto reshapeInputType =
llvm::dyn_cast<RankedTensorType>(reshapeOp.getInput1().getType());
- auto reshapeInputShape = reshapeInputType.getShape();
if (!reshapeInputType)
return std::nullopt;
+ auto reshapeInputShape = reshapeInputType.getShape();
auto reshapeOutputType =
llvm::cast<RankedTensorType>(reshapeOutput.getType());
const ArrayRef<int64_t> reshapeOutputShape = reshapeOutputType.getShape();
+
// Instead of inserting a TransposeOp here, we check if we can fold it into
// the ReshapeOp. There is more complex cases where this is possible, and
// this check can be extended.
-
if (failed(verifyUnitExpandedVectorShape(reshapeInputShape)) ||
failed(verifyUnitExpandedVectorShape(reshapeOutputShape)))
return std::nullopt;
- SmallVector<int64_t> mappedShape =
- applyTOSAPermutation(reshapeOutputShape, hoistedPerms);
- if (llvm::equal(mappedShape, reshapeInputShape))
- return reshapeOp.getInput1();
-
// Do not insert a TransposeOp, instead we fold the reshape and its attribute.
llvm::SmallVector<int64_t> newShape;
if (!tosa::getConstShapeValues(reshapeOp.getShape().getDefiningOp(),
@@ -438,7 +433,7 @@ std::optional<Value> TosaReduceTransposes::buildMappedToValue(
ImplicitLocOpBuilder builder(reshapeOp.getLoc(), rewriter);
auto foldedReshape = ReshapeOp::create(
rewriter, reshapeOp.getLoc(),
- RankedTensorType::get(mappedShape, reshapeOutputType.getElementType()),
+ RankedTensorType::get(applyTOSAPermutation(reshapeOutputShape, hoistedPerms), reshapeOutputType.getElementType()),
reshapeOp.getInput1(),
getTosaConstShape(builder, applyTOSAPermutation(llvm::ArrayRef(newShape),
hoistedPerms)));
More information about the Mlir-commits
mailing list