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

[llvmlistbot at llvm.org]

Author: Luke Hutton
Date: 2026-07-02T17:50:27+01:00
New Revision: 025c4c1d6bd017749e71ad3dc0b0d7b253606b4f

URL: https://github.com/llvm/llvm-project/commit/025c4c1d6bd017749e71ad3dc0b0d7b253606b4f
DIFF: https://github.com/llvm/llvm-project/commit/025c4c1d6bd017749e71ad3dc0b0d7b253606b4f.diff

LOG: [mlir][tosa] Fold unit-expanded reshapes in reduce-transposes (#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.

Added: 
    

Modified: 
    mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
    mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
index 65ef49bdc3077..0bd415040f22f 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) {
@@ -392,23 +409,18 @@ std::optional<Value> TosaReduceTransposes::buildMappedToValue(
   auto reshapeOutput = reshapeOp.getOutput();
   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 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.
-
-  // 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;
 
   // Do not insert a TransposeOp, instead we fold the reshape and its attribute.
@@ -421,8 +433,9 @@ 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(
+          applyTOSAPermutation(reshapeOutputShape, hoistedPerms),
+          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