[Mlir-commits] [mlir] [mlir][tosa] Fold unit-expanded reshapes in reduce-transposes (PR #203529)

[Luke Hutton]

https://github.com/lhutton1 created https://github.com/llvm/llvm-project/pull/203529

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.

>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] [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