[Mlir-commits] [mlir] [mlir][tosa] Fold tiled input into binary elementwise operations (PR #203941)
Luke Hutton
llvmlistbot at llvm.org
Mon Jun 29 06:56:02 PDT 2026
https://github.com/lhutton1 updated https://github.com/llvm/llvm-project/pull/203941
>From 427a19cd11be43f09af5283a641ef5515431f3af Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Thu, 11 Jun 2026 14:03:50 +0100
Subject: [PATCH 1/3] [mlir][tosa] Fold tiled input into binary elementwise
operations
Canonicalizes explicit broadcasting into implicit broadcasting by
folding the tile input into the binary elementwise operation.
Change-Id: I7a6d57248df9214978c15e9c4884629856b83b5c
---
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td | 1 +
.../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 86 ++++++++++
mlir/test/Dialect/Tosa/canonicalize.mlir | 147 ++++++++++++++++++
3 files changed, 234 insertions(+)
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index a99fb2fcae547..116c87f0fcfb5 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -2446,6 +2446,7 @@ def Tosa_TileOp : Tosa_InferShapedTypeOp<"tile", [Pure]> {
LogicalResult getConstantMultiples(llvm::SmallVector<int64_t> &multiples);
}];
+ 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 63d40ed4a95a5..f52cbe4f93829 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -20,8 +20,10 @@
#include "mlir/Dialect/Traits.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/InliningUtils.h"
#include "llvm/ADT/APFloat.h"
@@ -2216,6 +2218,90 @@ OpFoldResult tosa::SelectOp::fold(FoldAdaptor adaptor) {
return {};
}
+static LogicalResult verifyTileIsBroadcast(tosa::TileOp tileOp) {
+ const auto inputType = dyn_cast<RankedTensorType>(tileOp.getInput1().getType());
+ const auto outputType = dyn_cast<RankedTensorType>(tileOp.getType());
+ if (!inputType || !outputType)
+ return failure();
+
+ SmallVector<int64_t> multiples;
+ if (failed(tileOp.getConstantMultiples(multiples)))
+ return failure();
+
+ for (const auto [index, multiple] : llvm::enumerate(multiples)) {
+ if (multiple == 1)
+ continue;
+
+ if (inputType.isDynamicDim(index) || outputType.isDynamicDim(index))
+ return failure();
+ const int64_t inputDim = inputType.getDimSize(index);
+ if (inputDim * multiple != outputType.getDimSize(index))
+ return failure();
+ if (inputDim != 1)
+ return failure();
+ }
+
+ return success();
+}
+
+struct RemoveBroadcastTileFromBinaryElementwise
+ : public OpRewritePattern<tosa::TileOp> {
+ using OpRewritePattern<tosa::TileOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::TileOp tileOp,
+ PatternRewriter &rewriter) const override {
+ Value tileOutput = tileOp.getOutput();
+ if (!tileOutput.hasOneUse())
+ return rewriter.notifyMatchFailure(tileOp,
+ "tile output must have one use");
+
+ Operation *user = *tileOutput.user_begin();
+ const bool isBinaryElementwise =
+ user->getNumOperands() == 2 &&
+ user->hasTrait<OpTrait::tosa::TosaElementwiseOperator>();
+ if (!isBinaryElementwise && !isa<tosa::MulOp>(user))
+ return rewriter.notifyMatchFailure(
+ tileOp, "consumer must be binary broadcastable");
+
+ // Don't optimize a tile that feeds the shift operand of a MulOp
+ if (isa<tosa::MulOp>(user) && tileOutput == user->getOperand(2))
+ return rewriter.notifyMatchFailure(tileOp,
+ "tile feeds shift operand of MulOp");
+
+ if (failed(verifyTileIsBroadcast(tileOp)))
+ return rewriter.notifyMatchFailure(
+ tileOp, "tile must only expand statically-known singleton dims");
+
+ Value lhsOperand = user->getOperand(0);
+ Value rhsOperand = user->getOperand(1);
+ Value otherOperand = lhsOperand == tileOutput ? rhsOperand : lhsOperand;
+ Value tileInput = tileOp.getInput1();
+
+ const ShapedType newLhsType = cast<ShapedType>(otherOperand.getType());
+ const ShapedType newRhsType = cast<ShapedType>(tileInput.getType());
+ SmallVector<int64_t> broadcastedShape;
+ OpTrait::util::getBroadcastedShape(newLhsType.getShape(),
+ newRhsType.getShape(), broadcastedShape);
+
+ const ShapedType outputType = cast<ShapedType>(user->getResultTypes()[0]);
+ if (!llvm::equal(broadcastedShape, outputType.getShape()))
+ return rewriter.notifyMatchFailure(
+ tileOp, "tile output must be broadcastable to consumer operands");
+
+ rewriter.setInsertionPoint(user);
+ IRMapping mapper;
+ mapper.map(tileOutput, tileOp.getInput1());
+ Operation *newUser = rewriter.clone(*user, mapper);
+ rewriter.replaceOp(user, newUser->getResults());
+ return success();
+ }
+};
+
+void TileOp::getCanonicalizationPatterns(RewritePatternSet &results,
+ MLIRContext *context) {
+ results.add<RemoveBroadcastTileFromBinaryElementwise>(context);
+}
+
OpFoldResult TileOp::fold(FoldAdaptor adaptor) {
if (getInput1().getType() == getType()) {
if (auto multiples = llvm::dyn_cast_if_present<DenseElementsAttr>(
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 2cd040f056db8..7432767d2b57b 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1765,6 +1765,153 @@ func.func @dont_canonicalize_non_const_avg_pool2d_adaptive(%arg0: tensor<1x?x?x8
// -----
+// CHECK-LABEL: @canonicalize_tile_broadcast_sub
+// CHECK-SAME: %[[ARG0:[^:]+]]: tensor<96x56x56x96xf32>, %[[ARG1:[^:]+]]: tensor<1x56x56x1xf32>
+// CHECK-NOT: tosa.tile
+// CHECK: %[[SUB:.+]] = tosa.sub %[[ARG0]], %[[ARG1]] : (tensor<96x56x56x96xf32>, tensor<1x56x56x1xf32>) -> tensor<96x56x56x96xf32>
+// CHECK: return %[[SUB]]
+func.func @canonicalize_tile_broadcast_sub(%arg0: tensor<96x56x56x96xf32>, %arg1: tensor<1x56x56x1xf32>) -> tensor<96x56x56x96xf32> {
+ %shape = tosa.const_shape {values = dense<[96, 1, 1, 96]> : tensor<4xindex>} : () -> !tosa.shape<4>
+ %tile = tosa.tile %arg1, %shape : (tensor<1x56x56x1xf32>, !tosa.shape<4>) -> tensor<96x56x56x96xf32>
+ %sub = tosa.sub %arg0, %tile : (tensor<96x56x56x96xf32>, tensor<96x56x56x96xf32>) -> tensor<96x56x56x96xf32>
+ return %sub : tensor<96x56x56x96xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @canonicalize_tile_broadcast_mul_preserves_shift
+// CHECK-SAME: %[[ARG0:[^:]+]]: tensor<1x56x56x96xf32>, %[[ARG1:[^:]+]]: tensor<1x56x56x1xf32>, %[[SHIFT:[^:]+]]: tensor<1xi8>
+// CHECK-NOT: tosa.tile
+// CHECK: %[[MUL:.+]] = tosa.mul %[[ARG0]], %[[ARG1]], %[[SHIFT]] : (tensor<1x56x56x96xf32>, tensor<1x56x56x1xf32>, tensor<1xi8>) -> tensor<1x56x56x96xf32>
+// CHECK: return %[[MUL]]
+func.func @canonicalize_tile_broadcast_mul_preserves_shift(%arg0: tensor<1x56x56x96xf32>, %arg1: tensor<1x56x56x1xf32>, %shift: tensor<1xi8>) -> tensor<1x56x56x96xf32> {
+ %shape = tosa.const_shape {values = dense<[1, 1, 1, 96]> : tensor<4xindex>} : () -> !tosa.shape<4>
+ %tile = tosa.tile %arg1, %shape : (tensor<1x56x56x1xf32>, !tosa.shape<4>) -> tensor<1x56x56x96xf32>
+ %mul = tosa.mul %arg0, %tile, %shift : (tensor<1x56x56x96xf32>, tensor<1x56x56x96xf32>, tensor<1xi8>) -> tensor<1x56x56x96xf32>
+ return %mul : tensor<1x56x56x96xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @canonicalize_tile_broadcast_greater
+// CHECK-SAME: %[[ARG0:[^:]+]]: tensor<1x197x768xf32>, %[[ARG1:[^:]+]]: tensor<1x197x1xf32>
+// CHECK-NOT: tosa.tile
+// CHECK: %[[GT:.+]] = tosa.greater %[[ARG0]], %[[ARG1]] : (tensor<1x197x768xf32>, tensor<1x197x1xf32>) -> tensor<1x197x768xi1>
+// CHECK: return %[[GT]]
+func.func @canonicalize_tile_broadcast_greater(%arg0: tensor<1x197x768xf32>, %arg1: tensor<1x197x1xf32>) -> tensor<1x197x768xi1> {
+ %shape = tosa.const_shape {values = dense<[1, 1, 768]> : tensor<3xindex>} : () -> !tosa.shape<3>
+ %tile = tosa.tile %arg1, %shape : (tensor<1x197x1xf32>, !tosa.shape<3>) -> tensor<1x197x768xf32>
+ %gt = tosa.greater %arg0, %tile : (tensor<1x197x768xf32>, tensor<1x197x768xf32>) -> tensor<1x197x768xi1>
+ return %gt : tensor<1x197x768xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_tile_when_result_no_longer_broadcastable
+// CHECK: tosa.tile
+// CHECK: tosa.sub
+func.func @dont_canonicalize_tile_when_result_no_longer_broadcastable(%arg0: tensor<2x1xf32>) -> tensor<2x4xf32> {
+ %shape = tosa.const_shape {values = dense<[1, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %tile = tosa.tile %arg0, %shape : (tensor<2x1xf32>, !tosa.shape<2>) -> tensor<2x4xf32>
+ %sub = tosa.sub %tile, %arg0 : (tensor<2x4xf32>, tensor<2x1xf32>) -> tensor<2x4xf32>
+ return %sub : tensor<2x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_second_tile_when_result_no_longer_broadcastable
+// CHECK-SAME: %[[ARG0:[^:]+]]: tensor<2x1xf32>
+// CHECK: %[[SHAPE:.+]] = tosa.const_shape
+// CHECK: %[[TILE:.+]] = tosa.tile %[[ARG0]], %[[SHAPE]]
+// CHECK: %[[ADD:.+]] = tosa.add %[[ARG0]], %[[TILE]] : (tensor<2x1xf32>, tensor<2x4xf32>) -> tensor<2x4xf32>
+// CHECK: return %[[ADD]]
+func.func @dont_canonicalize_second_tile_when_result_no_longer_broadcastable(%arg0: tensor<2x1xf32>) -> tensor<2x4xf32> {
+ %shape = tosa.const_shape {values = dense<[1, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %tile0 = tosa.tile %arg0, %shape : (tensor<2x1xf32>, !tosa.shape<2>) -> tensor<2x4xf32>
+ %tile1 = tosa.tile %arg0, %shape : (tensor<2x1xf32>, !tosa.shape<2>) -> tensor<2x4xf32>
+ %add = tosa.add %tile0, %tile1 : (tensor<2x4xf32>, tensor<2x4xf32>) -> tensor<2x4xf32>
+ return %add : tensor<2x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_tile_with_unranked_other_operand
+// CHECK-SAME: %[[ARG0:[^:]+]]: tensor<2x1xf32>, %[[ARG1:[^:]+]]: tensor<*xf32>
+// CHECK: %[[SHAPE:.+]] = tosa.const_shape
+// CHECK: %[[TILE:.+]] = tosa.tile %[[ARG0]], %[[SHAPE]]
+// CHECK: %[[SUB:.+]] = tosa.sub %[[ARG1]], %[[TILE]] : (tensor<*xf32>, tensor<2x4xf32>) -> tensor<2x4xf32>
+// CHECK: return %[[SUB]]
+func.func @dont_canonicalize_tile_with_unranked_other_operand(%arg0: tensor<2x1xf32>, %arg1: tensor<*xf32>) -> tensor<2x4xf32> {
+ %shape = tosa.const_shape {values = dense<[1, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %tile = tosa.tile %arg0, %shape : (tensor<2x1xf32>, !tosa.shape<2>) -> tensor<2x4xf32>
+ %sub = tosa.sub %arg1, %tile : (tensor<*xf32>, tensor<2x4xf32>) -> tensor<2x4xf32>
+ return %sub : tensor<2x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_tile_non_singleton_expansion
+// CHECK: tosa.tile
+// CHECK: tosa.sub
+func.func @dont_canonicalize_tile_non_singleton_expansion(%arg0: tensor<1x56x56x96xf32>, %arg1: tensor<1x56x56x2xf32>) -> tensor<1x56x56x96xf32> {
+ %shape = tosa.const_shape {values = dense<[1, 1, 1, 48]> : tensor<4xindex>} : () -> !tosa.shape<4>
+ %tile = tosa.tile %arg1, %shape : (tensor<1x56x56x2xf32>, !tosa.shape<4>) -> tensor<1x56x56x96xf32>
+ %sub = tosa.sub %arg0, %tile : (tensor<1x56x56x96xf32>, tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32>
+ return %sub : tensor<1x56x56x96xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_tile_dynamic_expanded_dim
+// CHECK: tosa.tile
+// CHECK: tosa.sub
+func.func @dont_canonicalize_tile_dynamic_expanded_dim(%arg0: tensor<2x?xf32>, %arg1: tensor<2x4xf32>) -> tensor<2x4xf32> {
+ %shape = tosa.const_shape {values = dense<[1, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %tile = tosa.tile %arg0, %shape : (tensor<2x?xf32>, !tosa.shape<2>) -> tensor<2x4xf32>
+ %sub = tosa.sub %arg1, %tile : (tensor<2x4xf32>, tensor<2x4xf32>) -> tensor<2x4xf32>
+ return %sub : tensor<2x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_tile_multi_use
+// CHECK: tosa.tile
+// CHECK: tosa.sub
+// CHECK: tosa.add
+func.func @dont_canonicalize_tile_multi_use(%arg0: tensor<1x56x56x96xf32>, %arg1: tensor<1x56x56x1xf32>) -> (tensor<1x56x56x96xf32>, tensor<1x56x56x96xf32>) {
+ %shape = tosa.const_shape {values = dense<[1, 1, 1, 96]> : tensor<4xindex>} : () -> !tosa.shape<4>
+ %tile = tosa.tile %arg1, %shape : (tensor<1x56x56x1xf32>, !tosa.shape<4>) -> tensor<1x56x56x96xf32>
+ %sub = tosa.sub %arg0, %tile : (tensor<1x56x56x96xf32>, tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32>
+ %add = tosa.add %arg0, %tile : (tensor<1x56x56x96xf32>, tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32>
+ return %sub, %add : tensor<1x56x56x96xf32>, tensor<1x56x56x96xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_tile_unary_elementwise
+// CHECK: tosa.tile
+// CHECK: tosa.abs
+func.func @dont_canonicalize_tile_unary_elementwise(%arg0: tensor<2x1xf32>) -> tensor<2x4xf32> {
+ %shape = tosa.const_shape {values = dense<[1, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %tile = tosa.tile %arg0, %shape : (tensor<2x1xf32>, !tosa.shape<2>) -> tensor<2x4xf32>
+ %abs = tosa.abs %tile : (tensor<2x4xf32>) -> tensor<2x4xf32>
+ return %abs : tensor<2x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_tile_used_as_mul_shift
+// CHECK: tosa.tile
+// CHECK: tosa.mul
+func.func @dont_canonicalize_tile_used_as_mul_shift(%lhs: tensor<1xf32>, %rhs: tensor<1xf32>, %shift: tensor<?xi8>) -> tensor<1xf32> {
+ %multiples = tosa.const_shape {values = dense<[1]> : tensor<1xindex>} : () -> !tosa.shape<1>
+ %shift_static = tosa.tile %shift, %multiples : (tensor<?xi8>, !tosa.shape<1>) -> tensor<1xi8>
+ %mul = tosa.mul %lhs, %rhs, %shift_static : (tensor<1xf32>, tensor<1xf32>, tensor<1xi8>) -> tensor<1xf32>
+ return %mul : tensor<1xf32>
+}
+
+// -----
+
// CHECK-LABEL: test_single_concat
// CHECK: %[[VAL_1:.*]] = tosa.concat %arg0, %arg0 {axis = 1 : i32} : (tensor<1x1x7x7xf32>, tensor<1x1x7x7xf32>) -> tensor<1x2x7x7xf32>
// CHECK: return %[[VAL_1]] : tensor<1x2x7x7xf32>
>From cb9642c27b31d4f3e2c8784be3e753483058762c Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Mon, 15 Jun 2026 17:44:18 +0100
Subject: [PATCH 2/3] fix formatting
Change-Id: Ib3a2e94f618d02ac7acc6c9cb98a86a5334cd0d7
---
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index f52cbe4f93829..13e6cdd3bede7 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -2219,7 +2219,8 @@ OpFoldResult tosa::SelectOp::fold(FoldAdaptor adaptor) {
}
static LogicalResult verifyTileIsBroadcast(tosa::TileOp tileOp) {
- const auto inputType = dyn_cast<RankedTensorType>(tileOp.getInput1().getType());
+ const auto inputType =
+ dyn_cast<RankedTensorType>(tileOp.getInput1().getType());
const auto outputType = dyn_cast<RankedTensorType>(tileOp.getType());
if (!inputType || !outputType)
return failure();
>From 19714a38d1b90a7fcf1647360a044ebe14cb72fa Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Mon, 29 Jun 2026 14:55:27 +0100
Subject: [PATCH 3/3] address comments
Change-Id: If58c7e33255303e3cfb1d2da6b83533fff665201
---
.../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 23 +++++-----------
mlir/test/Dialect/Tosa/canonicalize.mlir | 26 +++++++++----------
2 files changed, 19 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 13e6cdd3bede7..b76bcba1f766a 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -2232,13 +2232,9 @@ static LogicalResult verifyTileIsBroadcast(tosa::TileOp tileOp) {
for (const auto [index, multiple] : llvm::enumerate(multiples)) {
if (multiple == 1)
continue;
-
- if (inputType.isDynamicDim(index) || outputType.isDynamicDim(index))
- return failure();
- const int64_t inputDim = inputType.getDimSize(index);
- if (inputDim * multiple != outputType.getDimSize(index))
- return failure();
- if (inputDim != 1)
+ if (outputType.isDynamicDim(index))
+ return failure();
+ if (inputType.getDimSize(index) != 1)
return failure();
}
@@ -2264,11 +2260,6 @@ struct RemoveBroadcastTileFromBinaryElementwise
return rewriter.notifyMatchFailure(
tileOp, "consumer must be binary broadcastable");
- // Don't optimize a tile that feeds the shift operand of a MulOp
- if (isa<tosa::MulOp>(user) && tileOutput == user->getOperand(2))
- return rewriter.notifyMatchFailure(tileOp,
- "tile feeds shift operand of MulOp");
-
if (failed(verifyTileIsBroadcast(tileOp)))
return rewriter.notifyMatchFailure(
tileOp, "tile must only expand statically-known singleton dims");
@@ -2278,11 +2269,11 @@ struct RemoveBroadcastTileFromBinaryElementwise
Value otherOperand = lhsOperand == tileOutput ? rhsOperand : lhsOperand;
Value tileInput = tileOp.getInput1();
- const ShapedType newLhsType = cast<ShapedType>(otherOperand.getType());
- const ShapedType newRhsType = cast<ShapedType>(tileInput.getType());
+ const ShapedType newOtherType = cast<ShapedType>(otherOperand.getType());
+ const ShapedType newTileType = cast<ShapedType>(tileInput.getType());
SmallVector<int64_t> broadcastedShape;
- OpTrait::util::getBroadcastedShape(newLhsType.getShape(),
- newRhsType.getShape(), broadcastedShape);
+ OpTrait::util::getBroadcastedShape(newOtherType.getShape(),
+ newTileType.getShape(), broadcastedShape);
const ShapedType outputType = cast<ShapedType>(user->getResultTypes()[0]);
if (!llvm::equal(broadcastedShape, outputType.getShape()))
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 7432767d2b57b..f99ce90d09c43 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1779,20 +1779,6 @@ func.func @canonicalize_tile_broadcast_sub(%arg0: tensor<96x56x56x96xf32>, %arg1
// -----
-// CHECK-LABEL: @canonicalize_tile_broadcast_mul_preserves_shift
-// CHECK-SAME: %[[ARG0:[^:]+]]: tensor<1x56x56x96xf32>, %[[ARG1:[^:]+]]: tensor<1x56x56x1xf32>, %[[SHIFT:[^:]+]]: tensor<1xi8>
-// CHECK-NOT: tosa.tile
-// CHECK: %[[MUL:.+]] = tosa.mul %[[ARG0]], %[[ARG1]], %[[SHIFT]] : (tensor<1x56x56x96xf32>, tensor<1x56x56x1xf32>, tensor<1xi8>) -> tensor<1x56x56x96xf32>
-// CHECK: return %[[MUL]]
-func.func @canonicalize_tile_broadcast_mul_preserves_shift(%arg0: tensor<1x56x56x96xf32>, %arg1: tensor<1x56x56x1xf32>, %shift: tensor<1xi8>) -> tensor<1x56x56x96xf32> {
- %shape = tosa.const_shape {values = dense<[1, 1, 1, 96]> : tensor<4xindex>} : () -> !tosa.shape<4>
- %tile = tosa.tile %arg1, %shape : (tensor<1x56x56x1xf32>, !tosa.shape<4>) -> tensor<1x56x56x96xf32>
- %mul = tosa.mul %arg0, %tile, %shift : (tensor<1x56x56x96xf32>, tensor<1x56x56x96xf32>, tensor<1xi8>) -> tensor<1x56x56x96xf32>
- return %mul : tensor<1x56x56x96xf32>
-}
-
-// -----
-
// CHECK-LABEL: @canonicalize_tile_broadcast_greater
// CHECK-SAME: %[[ARG0:[^:]+]]: tensor<1x197x768xf32>, %[[ARG1:[^:]+]]: tensor<1x197x1xf32>
// CHECK-NOT: tosa.tile
@@ -1874,6 +1860,18 @@ func.func @dont_canonicalize_tile_dynamic_expanded_dim(%arg0: tensor<2x?xf32>, %
// -----
+// CHECK-LABEL: @dont_canonicalize_tile_dynamic_output
+// CHECK: tosa.tile
+// CHECK: tosa.sub
+func.func @dont_canonicalize_tile_dynamic_output(%arg0: tensor<2x?xf32>, %arg1: tensor<2x1xf32>) -> tensor<2x?xf32> {
+ %shape = tosa.const_shape {values = dense<[1, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %tile = tosa.tile %arg1, %shape : (tensor<2x1xf32>, !tosa.shape<2>) -> tensor<2x?xf32>
+ %sub = tosa.sub %arg0, %tile : (tensor<2x?xf32>, tensor<2x?xf32>) -> tensor<2x?xf32>
+ return %sub : tensor<2x?xf32>
+}
+
+// -----
+
// CHECK-LABEL: @dont_canonicalize_tile_multi_use
// CHECK: tosa.tile
// CHECK: tosa.sub
More information about the Mlir-commits
mailing list