[mlir][tosa] Fold tiled input into binary elementwise operations (#203941)
Canonicalizes explicit broadcasting into implicit broadcasting by
folding the tile input into the binary element-wise operation.
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 5e5760e..e3b74f5a0 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -2485,6 +2485,7 @@
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 4c98aca..19f3575 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"
@@ -2326,6 +2328,82 @@
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 (outputType.isDynamicDim(index))
+ return failure();
+ if (inputType.getDimSize(index) != 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");
+
+ 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 newOtherType = cast<ShapedType>(otherOperand.getType());
+ const ShapedType newTileType = cast<ShapedType>(tileInput.getType());
+ SmallVector<int64_t> broadcastedShape;
+ OpTrait::util::getBroadcastedShape(
+ newOtherType.getShape(), newTileType.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 6a3d84e8d..585589a 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1996,6 +1996,166 @@
// -----
+// 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_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: @canonicalize_tile_broadcast_mul
+// CHECK-SAME: %[[ARG0:[^:]+]]: tensor<1x197x768xf32>, %[[ARG1:[^:]+]]: tensor<1x197x1xf32>, %[[SHIFT:[^:]+]]: tensor<1xi8>
+// CHECK-NOT: tosa.tile
+// CHECK: %[[MUL:.+]] = tosa.mul %[[ARG0]], %[[ARG1]], %[[SHIFT]] : (tensor<1x197x768xf32>, tensor<1x197x1xf32>, tensor<1xi8>) -> tensor<1x197x768xf32>
+// CHECK: return %[[MUL]]
+func.func @canonicalize_tile_broadcast_mul(%arg0: tensor<1x197x768xf32>, %arg1: tensor<1x197x1xf32>, %shift: tensor<1xi8>) -> tensor<1x197x768xf32> {
+ %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>
+ %mul = tosa.mul %arg0, %tile, %shift : (tensor<1x197x768xf32>, tensor<1x197x768xf32>, tensor<1xi8>) -> tensor<1x197x768xf32>
+ return %mul : tensor<1x197x768xf32>
+}
+
+
+// -----
+
+// 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_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
+// 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>
