[Mlir-commits] [mlir] [mlir][tosa] Improve broadcasting behaviour in elementwise folders (PR #181114)
Luke Hutton
llvmlistbot at llvm.org
Mon Feb 16 00:35:25 PST 2026
https://github.com/lhutton1 updated https://github.com/llvm/llvm-project/pull/181114
>From de4568580364bc62d7af8184a17ade479b98a1a1 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Wed, 11 Feb 2026 18:06:56 +0000
Subject: [PATCH 1/2] [mlir][tosa] Improve broadcasting behaviour in
elementwise folders
This commit aims to improve elementwise folder behaviour when broadcasting
is involved. In particular, it ensures correctness of folders when the
input operands are dynamic and it is not clear whether broadcasting is
involved.
For example, previously, the tosa.add folder could result in shape
information loss:
```
func.func @test(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
%one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
%div = tosa.add %one, %arg0 : (tensor<2x1xi32>, tensor<?x4xi32>) -> tensor<?x4xi32>
return %div : tensor<?x4xi32>
}
$ mlir-opt --canonicalize test.mlir
func.func @test(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
return %arg0 : tensor<?x4xi32>
}
```
In this example we have lost the contraint on the first dimension provided
by the lhs operand. This commit ensures these cases are not folded.
Similar logic has also been applied to tosa.sub, tosa.intdiv, tosa.mul.
tosa.select previously handled the above example correctly, but did this
by failing to fold when dynamic shapes were involved (a conservative
approach). This commit improves the folder such that it follows the
behaviour of the other elementwise folders.
Note: some tests have been moved from canonicalize.mlir to
constant_folding.mlir to help keep folding tests together.
Change-Id: If3d3209299afb6488600f1eb1374ba9c1e5ccf96
---
.../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 100 +++--
mlir/test/Dialect/Tosa/canonicalize.mlir | 166 --------
mlir/test/Dialect/Tosa/constant_folding.mlir | 384 ++++++++++++++++++
3 files changed, 445 insertions(+), 205 deletions(-)
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 42033ce8a3b02..230c7d8bdb15f 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -16,6 +16,7 @@
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
#include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
+#include "mlir/Dialect/Traits.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
@@ -1213,9 +1214,11 @@ OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
auto rhsAttr =
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
- if (lhsTy == resultTy && isSplatZero(resultETy, rhsAttr))
+ const bool isBroadcastable = OpTrait::util::staticallyKnownBroadcastable(
+ lhsTy.getShape(), rhsTy.getShape());
+ if (isBroadcastable && lhsTy == resultTy && isSplatZero(resultETy, rhsAttr))
return getInput1();
- if (rhsTy == resultTy && isSplatZero(resultETy, lhsAttr))
+ if (isBroadcastable && rhsTy == resultTy && isSplatZero(resultETy, lhsAttr))
return getInput2();
if (!lhsAttr || !rhsAttr)
@@ -1247,7 +1250,7 @@ OpFoldResult IntDivOp::fold(FoldAdaptor adaptor) {
auto resultTy = llvm::dyn_cast<RankedTensorType>(getType());
if (!lhsTy || !rhsTy || !resultTy)
return {};
- if (lhsTy != rhsTy)
+ if (lhsTy.getElementType() != rhsTy.getElementType())
return {};
// IntDivOp inputs must be integer type, no need to check for quantized type
@@ -1257,13 +1260,16 @@ OpFoldResult IntDivOp::fold(FoldAdaptor adaptor) {
auto rhsAttr =
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
if (lhsAttr && lhsAttr.isSplat()) {
- if (llvm::isa<IntegerType>(resultETy) &&
+ if (llvm::isa<IntegerType>(resultETy) && resultTy.hasStaticShape() &&
lhsAttr.getSplatValue<APInt>().isZero())
- return lhsAttr;
+ return lhsAttr.resizeSplat(resultTy);
}
if (rhsAttr && rhsAttr.isSplat()) {
- if (llvm::isa<IntegerType>(resultETy) &&
+ const bool isBroadcastable = OpTrait::util::staticallyKnownBroadcastable(
+ lhsTy.getShape(), rhsTy.getShape());
+ if (isBroadcastable && lhsTy == resultTy &&
+ llvm::isa<IntegerType>(resultETy) &&
rhsAttr.getSplatValue<APInt>().isOne())
return getInput1();
}
@@ -1369,19 +1375,22 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
}
}
- if (rhsTy == resultTy) {
- if (isSplatZero(resultETy, lhsAttr) && resultTy.hasStaticShape())
- // constant values can only be resized if resulting type is static
- return lhsAttr.resizeSplat(resultTy);
- if (isSplatOne(resultETy, lhsAttr, shift))
- return rhs;
- }
- if (lhsTy == resultTy) {
- if (isSplatZero(resultETy, rhsAttr) && resultTy.hasStaticShape())
- return rhsAttr.resizeSplat(resultTy);
- if (isSplatOne(resultETy, rhsAttr, shift))
- return lhs;
- }
+ if (rhsTy == resultTy && isSplatZero(resultETy, lhsAttr) &&
+ resultTy.hasStaticShape())
+ // constant values can only be resized if resulting type is static
+ return lhsAttr.resizeSplat(resultTy);
+ if (lhsTy == resultTy && isSplatZero(resultETy, rhsAttr) &&
+ resultTy.hasStaticShape())
+ return rhsAttr.resizeSplat(resultTy);
+
+ const bool isBroadcastable = OpTrait::util::staticallyKnownBroadcastable(
+ lhsTy.getShape(), rhsTy.getShape());
+ if (isBroadcastable && rhsTy == resultTy &&
+ isSplatOne(resultETy, lhsAttr, shift))
+ return rhs;
+ if (isBroadcastable && lhsTy == resultTy &&
+ isSplatOne(resultETy, rhsAttr, shift))
+ return lhs;
return mulBinaryFolder(lhsAttr, rhsAttr, resultTy, shift);
}
@@ -1404,7 +1413,9 @@ OpFoldResult SubOp::fold(FoldAdaptor adaptor) {
auto rhsAttr =
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
- if (lhsTy == resultTy && isSplatZero(resultETy, rhsAttr))
+ const bool isBroadcastable = OpTrait::util::staticallyKnownBroadcastable(
+ lhsTy.getShape(), rhsTy.getShape());
+ if (isBroadcastable && lhsTy == resultTy && isSplatZero(resultETy, rhsAttr))
return getInput1();
if (!lhsAttr || !rhsAttr)
@@ -1717,36 +1728,47 @@ OpFoldResult SliceOp::fold(FoldAdaptor adaptor) {
return {};
}
-static bool
-mayRequireBroadcast(ValueTypeRange<mlir::OperandRange> operandTypes) {
- const auto isDynamic = [](Type ty) {
- const auto shapedTy = llvm::dyn_cast<ShapedType>(ty);
- return !shapedTy || !shapedTy.hasStaticShape();
- };
-
- return llvm::any_of(operandTypes, isDynamic) ||
- failed(verifyCompatibleShapes(operandTypes));
-}
-
OpFoldResult tosa::SelectOp::fold(FoldAdaptor adaptor) {
- // Select allows operand shapes to be broadcast to the output shape. For
- // now, don't support folding when we cannot prove no broadcasting is
- // involved.
- if (mayRequireBroadcast(getOperandTypes()))
+ const Value pred = getPred();
+ const Value onTrue = getOnTrue();
+ const Value onFalse = getOnFalse();
+
+ const auto predTy = llvm::dyn_cast<RankedTensorType>(pred.getType());
+ const auto onTrueTy = llvm::dyn_cast<RankedTensorType>(onTrue.getType());
+ const auto onFalseTy = llvm::dyn_cast<RankedTensorType>(onFalse.getType());
+ if (!predTy || !onTrueTy || !onFalseTy)
return {};
- if (getOnTrue() == getOnFalse())
+ const Type resultTy = getType();
+
+ const ArrayRef<int64_t> predShape = predTy.getShape();
+ const ArrayRef<int64_t> onTrueShape = onTrueTy.getShape();
+
+ if (onTrue == onFalse && onTrueTy == resultTy &&
+ OpTrait::util::staticallyKnownBroadcastable(predShape, onTrueShape))
return getOnTrue();
auto predicate =
llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getInput1());
if (!predicate)
return {};
-
if (!predicate.isSplat())
return {};
- return predicate.getSplatValue<APInt>().getBoolValue() ? getOnTrue()
- : getOnFalse();
+
+ const bool predicateValue = predicate.getSplatValue<APInt>().getBoolValue();
+
+ SmallVector<SmallVector<int64_t>, 3> shapes;
+ shapes.emplace_back(predShape);
+ shapes.emplace_back(onTrueShape);
+ shapes.emplace_back(onFalseTy.getShape());
+ const bool isBroadcastable =
+ OpTrait::util::staticallyKnownBroadcastable(shapes);
+
+ if (predicateValue == true && onTrueTy == resultTy && isBroadcastable)
+ return getOnTrue();
+ if (predicateValue == false && onFalseTy == resultTy && isBroadcastable)
+ return getOnFalse();
+ return {};
}
OpFoldResult TileOp::fold(FoldAdaptor adaptor) {
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 81e537babf9ab..e85d349459aaa 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -539,130 +539,6 @@ func.func @pad_determine_val_quant(%arg0: tensor<?x?xi32>, %arg1 : tensor<2x2xi3
// -----
-// CHECK-LABEL: @mul_one_float
-func.func @mul_one_float(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
- // CHECK: return %arg0
- // CHECK-NOT: tosa.mul
- %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
- %ones = "tosa.const"() {values = dense<1.0> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
- %1 = tosa.mul %arg0, %ones, %shift : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<1xi8>) -> tensor<2x3xf32>
- return %1 : tensor<2x3xf32>
-}
-
-// -----
-
-// CHECK-LABEL: @mul_bcast_one_float
-func.func @mul_bcast_one_float(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
- // CHECK: return %arg0
- // CHECK-NOT: tosa.mul
- %ones = "tosa.const"() {values = dense<1.0> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
- %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
- %1 = tosa.mul %ones, %arg0, %shift : (tensor<1x1xf32>, tensor<2x3xf32>, tensor<1xi8>) -> tensor<2x3xf32>
- return %1 : tensor<2x3xf32>
-}
-
-// -----
-
-// CHECK-LABEL: @mul_one_int
-func.func @mul_one_int(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> {
- // CHECK: return %arg0
- // CHECK-NOT: tosa.mul
- %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
- %ones = "tosa.const"() {values = dense<1> : tensor<2x3xi32>} : () -> tensor<2x3xi32>
- %1 = tosa.mul %arg0, %ones, %shift : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<1xi8>) -> tensor<2x3xi32>
- return %1 : tensor<2x3xi32>
-}
-
-// -----
-
-// CHECK-LABEL: @mul_one_int_and_shift
-func.func @mul_one_int_and_shift(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> {
- // CHECK-DAG: %[[VAL_1:.*]] = "tosa.const"() <{values = dense<1> : tensor<2x3xi32>}>
- // CHECK-DAG: %[[VAL_2:.*]] = "tosa.const"() <{values = dense<31> : tensor<1xi8>}>
- // CHECK: %[[VAL_3:.*]] = tosa.mul %arg0, %[[VAL_1]], %[[VAL_2]] : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<1xi8>)
- %ones = "tosa.const"() {values = dense<1> : tensor<2x3xi32>} : () -> tensor<2x3xi32>
- %shift = "tosa.const"() <{values = dense<31> : tensor<1xi8>}> : () -> tensor<1xi8>
- %1 = tosa.mul %arg0, %ones, %shift : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<1xi8>) -> tensor<2x3xi32>
- return %1 : tensor<2x3xi32>
-}
-
-// -----
-
-// CHECK-LABEL: @mul_zero_broadcast
-func.func @mul_zero_broadcast(%arg0: tensor<2x3xf32>) -> (tensor<2x3xf32>, tensor<2x3xf32>) {
- // CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<2x3xf32>}
- // CHECK-NOT: tosa.mul
- %zeros = "tosa.const"() {values = dense<0.0> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
- %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
- %1 = tosa.mul %arg0, %zeros, %shift : (tensor<2x3xf32>, tensor<1x1xf32>, tensor<1xi8>) -> tensor<2x3xf32>
-
- // CHECK-NOT: tosa.mul
- // CHECK: return %[[ZERO]], %[[ZERO]]
- %2 = tosa.mul %zeros, %arg0, %shift : (tensor<1x1xf32>, tensor<2x3xf32>, tensor<1xi8>) -> tensor<2x3xf32>
- return %1, %2 : tensor<2x3xf32>, tensor<2x3xf32>
-}
-
-// -----
-
-// CHECK-LABEL: @mul_zero_dynamic_nofold
-// CHECK-SAME: %[[ARG0:.*]]: tensor<?x17xf32>) -> tensor<?x17xf32> {
-// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1x1xf32>}> : () -> tensor<1x1xf32>
-// CHECK: %[[SHIFT:.*]] = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
-// CHECK: %[[MUL:.*]] = tosa.mul %[[ARG0]], %[[ZERO]], %[[SHIFT]] : (tensor<?x17xf32>, tensor<1x1xf32>, tensor<1xi8>) -> tensor<?x17xf32>
-// CHECK: return %[[MUL]]
-func.func @mul_zero_dynamic_nofold(%arg0: tensor<?x17xf32>) -> tensor<?x17xf32> {
- %0 = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1x1xf32>}> : () -> tensor<1x1xf32>
- %1 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
- %2 = tosa.mul %arg0, %0, %1 : (tensor<?x17xf32>, tensor<1x1xf32>, tensor<1xi8>) -> tensor<?x17xf32>
- return %2 : tensor<?x17xf32>
-}
-
-// -----
-
-// CHECK-LABEL: @mul_one_dynamic_fold
-// CHECK-SAME: %[[ARG0:.*]]: tensor<?x17xf32>) -> tensor<?x17xf32> {
-// CHECK: return %[[ARG0]]
-func.func @mul_one_dynamic_fold(%arg0: tensor<?x17xf32>) -> tensor<?x17xf32> {
- %0 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x1xf32>}> : () -> tensor<1x1xf32>
- %1 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
- %2 = tosa.mul %arg0, %0, %1 : (tensor<?x17xf32>, tensor<1x1xf32>, tensor<1xi8>) -> tensor<?x17xf32>
- return %2 : tensor<?x17xf32>
-}
-
-// -----
-
-// CHECK-LABEL: @select_same_value
-func.func @select_same_value(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
- %0 = tosa.select %arg0, %arg1, %arg1 : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
- // CHECK: return %arg1
- // CHECK-NOT: tosa.select
- return %0 : tensor<2x3xi32>
-}
-
-// -----
-
-// CHECK-LABEL: @select_true_value
-func.func @select_true_value(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
- %c1 = "tosa.const"() {values = dense<1> : tensor<2x3xi1>} : () -> tensor<2x3xi1>
- %0 = tosa.select %c1, %arg0, %arg1 : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
- // CHECK: return %arg0
- // CHECK-NOT: tosa.select
- return %0 : tensor<2x3xi32>
-}
-
-// -----
-
-// CHECK-LABEL: @select_false_value
-func.func @select_false_value(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
- %c0 = "tosa.const"() {values = dense<0> : tensor<2x3xi1>} : () -> tensor<2x3xi1>
- %0 = tosa.select %c0, %arg0, %arg1 : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
- // CHECK: return %arg1
- // CHECK-NOT: tosa.select
- return %0 : tensor<2x3xi32>
-}
-
-// -----
-
// CHECK-LABEL: @select_not_pred
func.func @select_not_pred(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>, %arg2: tensor<2x3xi32>) -> tensor<2x3xi32> {
%0 = tosa.logical_not %arg0 : (tensor<2x3xi1>) -> tensor<2x3xi1>
@@ -673,48 +549,6 @@ func.func @select_not_pred(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>, %arg2:
// -----
-// CHECK-LABEL: @select_broadcast_same_value_no_fold
-func.func @select_broadcast_same_value_no_fold(%arg0: tensor<2x2xi1>, %arg1: tensor<1x1xf32>) -> tensor<2x2xf32> {
- // CHECK: tosa.select %arg0, %arg1, %arg1
- %0 = tosa.select %arg0, %arg1, %arg1 : (tensor<2x2xi1>, tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<2x2xf32>
- return %0 : tensor<2x2xf32>
-}
-
-// -----
-
-// CHECK-LABEL: @select_broadcast_true_value_no_fold
-func.func @select_broadcast_true_value_no_fold(%arg0: tensor<1x1xf32>, %arg1: tensor<2x2xf32>) -> tensor<?x?xf32> {
- // CHECK: %[[CONST:.*]] = "tosa.const"
- %0 = "tosa.const"() {values = dense<1> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
- // CHECK: tosa.select %[[CONST]], %arg0, %arg1
- %1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<1x1xf32>, tensor<2x2xf32>) -> tensor<?x?xf32>
- return %1 : tensor<?x?xf32>
-}
-
-// -----
-
-// CHECK-LABEL: @select_broadcast_false_value_no_fold
-func.func @select_broadcast_false_value_no_fold(%arg0: tensor<2x2xf32>, %arg1: tensor<1x1xf32>) -> tensor<2x2xf32> {
- // CHECK: %[[CONST:.*]] = "tosa.const"
- %0 = "tosa.const"() {values = dense<0> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
- // CHECK: tosa.select %[[CONST]], %arg0, %arg1
- %1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<2x2xf32>, tensor<1x1xf32>) -> tensor<2x2xf32>
- return %1 : tensor<2x2xf32>
-}
-
-// -----
-
-// CHECK-LABEL: @select_broadcast_false_value_dynamic_operand_no_fold
-func.func @select_broadcast_false_value_dynamic_operand_no_fold(%arg0: tensor<2x?xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {
- // CHECK: %[[CONST:.*]] = "tosa.const"
- %0 = "tosa.const"() {values = dense<0> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
- // CHECK: tosa.select %[[CONST]], %arg0, %arg1
- %1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<2x?xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
- return %1 : tensor<2x2xf32>
-}
-
-// -----
-
// CHECK-LABEL: @reduce_all_fold
func.func @reduce_all_fold(%arg0: tensor<?x1xf32>) -> tensor<?x1xf32> {
// CHECK: return %arg0
diff --git a/mlir/test/Dialect/Tosa/constant_folding.mlir b/mlir/test/Dialect/Tosa/constant_folding.mlir
index c3186279a30ae..5fa5f1d5143f3 100644
--- a/mlir/test/Dialect/Tosa/constant_folding.mlir
+++ b/mlir/test/Dialect/Tosa/constant_folding.mlir
@@ -136,6 +136,56 @@ func.func @fold_add_splat_ui8_overflow() -> tensor<10xui8> {
// -----
+// CHECK-LABEL: @no_fold_add_unknown_broadcast_zero_lhs
+func.func @no_fold_add_unknown_broadcast_zero_lhs(%arg0: tensor<1x4xi32>) -> tensor<2x4xi32> {
+ %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.add %one, %arg0 : (tensor<2x1xi32>, tensor<1x4xi32>) -> tensor<2x4xi32>
+ // CHECK: tosa.add
+ return %div : tensor<2x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_dynamic_add_unknown_broadcast_zero_lhs
+func.func @no_fold_dynamic_add_unknown_broadcast_zero_lhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
+ %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.add %one, %arg0 : (tensor<2x1xi32>, tensor<?x4xi32>) -> tensor<?x4xi32>
+ // CHECK: tosa.add
+ return %div : tensor<?x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_dynamic_add_broadcast_zero_lhs
+func.func @fold_dynamic_add_broadcast_zero_lhs(%arg0: tensor<?x17xi32>) -> tensor<?x17xi32> {
+ %one = "tosa.const"() {values = dense<0> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
+ %div = tosa.add %one, %arg0 : (tensor<1x1xi32>, tensor<?x17xi32>) -> tensor<?x17xi32>
+ // CHECK: return %arg0
+ return %div : tensor<?x17xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_dynamic_add_unknown_broadcast_zero_rhs
+func.func @no_fold_dynamic_add_unknown_broadcast_zero_rhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
+ %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.add %arg0, %one : (tensor<?x4xi32>, tensor<2x1xi32>) -> tensor<?x4xi32>
+ // CHECK: tosa.add
+ return %div : tensor<?x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_dynamic_add_broadcast_zero_rhs
+func.func @fold_dynamic_add_broadcast_zero_rhs(%arg0: tensor<?x17xi32>) -> tensor<?x17xi32> {
+ %one = "tosa.const"() {values = dense<0> : tensor<1x17xi32>} : () -> tensor<1x17xi32>
+ %div = tosa.add %arg0, %one : (tensor<?x17xi32>, tensor<1x17xi32>) -> tensor<?x17xi32>
+ // CHECK: return %arg0
+ return %div : tensor<?x17xi32>
+}
+
+// -----
+
// CHECK-LABEL: @fold_div_zero_lhs_i32
func.func @fold_div_zero_lhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
%zero = "tosa.const"() {values = dense<0> : tensor<i32>} : () -> tensor<i32>
@@ -147,6 +197,47 @@ func.func @fold_div_zero_lhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
// -----
+// CHECK-LABEL: @no_fold_dynamic_div_zero_lhs
+func.func @no_fold_dynamic_div_zero_lhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
+ %zero = "tosa.const"() {values = dense<0> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
+ // CHECK: tosa.intdiv
+ %div = tosa.intdiv %zero, %arg0 : (tensor<1x1xi32>, tensor<?x4xi32>) -> tensor<?x4xi32>
+ return %div : tensor<?x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_div_zero_lhs_broadcast
+func.func @fold_div_zero_lhs_broadcast(%arg0: tensor<2x4xi32>) -> tensor<2x4xi32> {
+ %zero = "tosa.const"() {values = dense<0> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
+ // CHECK: %[[ZERO:.+]] = "tosa.const"() <{values = dense<0> : tensor<2x4xi32>
+ %div = tosa.intdiv %zero, %arg0 : (tensor<1x1xi32>, tensor<2x4xi32>) -> tensor<2x4xi32>
+ // CHECK: return %[[ZERO]]
+ return %div : tensor<2x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_div_unknown_broadcast_one_rhs
+func.func @no_fold_div_unknown_broadcast_one_rhs(%arg0: tensor<1x4xi32>) -> tensor<2x4xi32> {
+ %one = "tosa.const"() {values = dense<1> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ // CHECK: tosa.intdiv
+ %div = tosa.intdiv %arg0, %one : (tensor<1x4xi32>, tensor<2x1xi32>) -> tensor<2x4xi32>
+ return %div : tensor<2x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_div_broadcast_one_rhs
+func.func @fold_div_broadcast_one_rhs(%arg0: tensor<?x17xi32>) -> tensor<?x17xi32> {
+ %one = "tosa.const"() {values = dense<1> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
+ %div = tosa.intdiv %arg0, %one : (tensor<?x17xi32>, tensor<1x1xi32>) -> tensor<?x17xi32>
+ // CHECK: return %arg0
+ return %div : tensor<?x17xi32>
+}
+
+// -----
+
// CHECK-LABEL: @fold_div_one_rhs_i32
func.func @fold_div_one_rhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
%one = "tosa.const"() {values = dense<1> : tensor<i32>} : () -> tensor<i32>
@@ -157,6 +248,16 @@ func.func @fold_div_one_rhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
// -----
+// CHECK-LABEL: @no_fold_dynamic_div_unknown_broadcast_one_rhs
+func.func @no_fold_dynamic_div_unknown_broadcast_one_rhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
+ %one = "tosa.const"() {values = dense<1> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ // CHECK: tosa.intdiv
+ %div = tosa.intdiv %arg0, %one : (tensor<?x4xi32>, tensor<2x1xi32>) -> tensor<?x4xi32>
+ return %div : tensor<?x4xi32>
+}
+
+// -----
+
// CHECK-LABEL: @fold_div_splat_i32
func.func @fold_div_splat_i32() -> tensor<i32> {
%lhs = "tosa.const"() {values = dense<10> : tensor<i32>} : () -> tensor<i32>
@@ -229,6 +330,28 @@ func.func @fold_mul_one_rhs_f32(%arg0: tensor<f32>) -> tensor<f32> {
// -----
+// CHECK-LABEL: @no_fold_mul_unknown_broadcast_one_rhs_f32
+func.func @no_fold_mul_unknown_broadcast_one_rhs_f32(%arg0: tensor<1x4xf32>) -> tensor<2x4xf32> {
+ %one = "tosa.const"() {values = dense<1.0> : tensor<2x1xf32>} : () -> tensor<2x1xf32>
+ %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ // CHECK: tosa.mul
+ %mul = tosa.mul %arg0, %one, %shift : (tensor<1x4xf32>, tensor<2x1xf32>, tensor<1xi8>) -> tensor<2x4xf32>
+ return %mul : tensor<2x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_dynamic_mul_unknown_broadcast_one_rhs_f32
+func.func @no_fold_dynamic_mul_unknown_broadcast_one_rhs_f32(%arg0: tensor<?x4xf32>) -> tensor<?x4xf32> {
+ %one = "tosa.const"() {values = dense<1.0> : tensor<2x1xf32>} : () -> tensor<2x1xf32>
+ %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ // CHECK: tosa.mul
+ %mul = tosa.mul %arg0, %one, %shift : (tensor<?x4xf32>, tensor<2x1xf32>, tensor<1xi8>) -> tensor<?x4xf32>
+ return %mul : tensor<?x4xf32>
+}
+
+// -----
+
// CHECK-LABEL: @fold_mul_one_lhs_f32
func.func @fold_mul_one_lhs_f32(%arg0: tensor<f32>) -> tensor<f32> {
%one = "tosa.const"() {values = dense<1.0> : tensor<f32>} : () -> tensor<f32>
@@ -262,6 +385,120 @@ func.func @fold_mul_one_lhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
// -----
+// CHECK-LABEL: @no_fold_mul_unknown_broadcast_one_lhs_f32
+func.func @no_fold_mul_unknown_broadcast_one_lhs_f32(%arg0: tensor<1x4xf32>) -> tensor<2x4xf32> {
+ %one = "tosa.const"() {values = dense<1.0> : tensor<2x1xf32>} : () -> tensor<2x1xf32>
+ %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ // CHECK: tosa.mul
+ %mul = tosa.mul %one, %arg0, %shift : (tensor<2x1xf32>, tensor<1x4xf32>, tensor<1xi8>) -> tensor<2x4xf32>
+ return %mul : tensor<2x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_dynamic_mul_unknown_broadcast_one_lhs_f32
+func.func @no_fold_dynamic_mul_unknown_broadcast_one_lhs_f32(%arg0: tensor<?x4xf32>) -> tensor<?x4xf32> {
+ %one = "tosa.const"() {values = dense<1.0> : tensor<2x1xf32>} : () -> tensor<2x1xf32>
+ %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ // CHECK: tosa.mul
+ %mul = tosa.mul %one, %arg0, %shift : (tensor<2x1xf32>, tensor<?x4xf32>, tensor<1xi8>) -> tensor<?x4xf32>
+ return %mul : tensor<?x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_mul_one_float
+func.func @fold_mul_one_float(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
+ // CHECK: return %arg0
+ // CHECK-NOT: tosa.mul
+ %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %ones = "tosa.const"() {values = dense<1.0> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
+ %1 = tosa.mul %arg0, %ones, %shift : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<1xi8>) -> tensor<2x3xf32>
+ return %1 : tensor<2x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_mul_bcast_one_float
+func.func @fold_mul_bcast_one_float(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
+ // CHECK: return %arg0
+ // CHECK-NOT: tosa.mul
+ %ones = "tosa.const"() {values = dense<1.0> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
+ %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %1 = tosa.mul %ones, %arg0, %shift : (tensor<1x1xf32>, tensor<2x3xf32>, tensor<1xi8>) -> tensor<2x3xf32>
+ return %1 : tensor<2x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_mul_one_int
+func.func @fold_mul_one_int(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> {
+ // CHECK: return %arg0
+ // CHECK-NOT: tosa.mul
+ %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %ones = "tosa.const"() {values = dense<1> : tensor<2x3xi32>} : () -> tensor<2x3xi32>
+ %1 = tosa.mul %arg0, %ones, %shift : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<1xi8>) -> tensor<2x3xi32>
+ return %1 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_mul_one_int_and_shift
+func.func @fold_mul_one_int_and_shift(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> {
+ // CHECK-DAG: %[[VAL_1:.*]] = "tosa.const"() <{values = dense<1> : tensor<2x3xi32>}>
+ // CHECK-DAG: %[[VAL_2:.*]] = "tosa.const"() <{values = dense<31> : tensor<1xi8>}>
+ // CHECK: %[[VAL_3:.*]] = tosa.mul %arg0, %[[VAL_1]], %[[VAL_2]] : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<1xi8>)
+ %ones = "tosa.const"() {values = dense<1> : tensor<2x3xi32>} : () -> tensor<2x3xi32>
+ %shift = "tosa.const"() <{values = dense<31> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %1 = tosa.mul %arg0, %ones, %shift : (tensor<2x3xi32>, tensor<2x3xi32>, tensor<1xi8>) -> tensor<2x3xi32>
+ return %1 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_mul_zero_broadcast
+func.func @fold_mul_zero_broadcast(%arg0: tensor<2x3xf32>) -> (tensor<2x3xf32>, tensor<2x3xf32>) {
+ // CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<2x3xf32>}
+ // CHECK-NOT: tosa.mul
+ %zeros = "tosa.const"() {values = dense<0.0> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
+ %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %1 = tosa.mul %arg0, %zeros, %shift : (tensor<2x3xf32>, tensor<1x1xf32>, tensor<1xi8>) -> tensor<2x3xf32>
+
+ // CHECK-NOT: tosa.mul
+ // CHECK: return %[[ZERO]], %[[ZERO]]
+ %2 = tosa.mul %zeros, %arg0, %shift : (tensor<1x1xf32>, tensor<2x3xf32>, tensor<1xi8>) -> tensor<2x3xf32>
+ return %1, %2 : tensor<2x3xf32>, tensor<2x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_mul_zero_dynamic_nofold
+// CHECK-SAME: %[[ARG0:.*]]: tensor<?x17xf32>) -> tensor<?x17xf32> {
+// CHECK: %[[ZERO:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1x1xf32>}> : () -> tensor<1x1xf32>
+// CHECK: %[[SHIFT:.*]] = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK: %[[MUL:.*]] = tosa.mul %[[ARG0]], %[[ZERO]], %[[SHIFT]] : (tensor<?x17xf32>, tensor<1x1xf32>, tensor<1xi8>) -> tensor<?x17xf32>
+// CHECK: return %[[MUL]]
+func.func @fold_mul_zero_dynamic_nofold(%arg0: tensor<?x17xf32>) -> tensor<?x17xf32> {
+ %0 = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1x1xf32>}> : () -> tensor<1x1xf32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %2 = tosa.mul %arg0, %0, %1 : (tensor<?x17xf32>, tensor<1x1xf32>, tensor<1xi8>) -> tensor<?x17xf32>
+ return %2 : tensor<?x17xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_mul_one_dynamic_fold
+// CHECK-SAME: %[[ARG0:.*]]: tensor<?x17xf32>) -> tensor<?x17xf32> {
+// CHECK: return %[[ARG0]]
+func.func @fold_mul_one_dynamic_fold(%arg0: tensor<?x17xf32>) -> tensor<?x17xf32> {
+ %0 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x1xf32>}> : () -> tensor<1x1xf32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %2 = tosa.mul %arg0, %0, %1 : (tensor<?x17xf32>, tensor<1x1xf32>, tensor<1xi8>) -> tensor<?x17xf32>
+ return %2 : tensor<?x17xf32>
+}
+
+// -----
+
// CHECK-LABEL: @fold_mul_splat_i8
func.func @fold_mul_splat_i8() -> tensor<10xi32> {
%one = "tosa.const"() {values = dense<17> : tensor<10xi8>} : () -> tensor<10xi8>
@@ -308,6 +545,36 @@ func.func @fold_sub_zero_rhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
// -----
+// CHECK-LABEL: @no_fold_sub_unknown_broadcast_zero_rhs
+func.func @no_fold_sub_unknown_broadcast_zero_rhs(%arg0: tensor<1x4xi32>) -> tensor<2x4xi32> {
+ %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.sub %arg0, %one : (tensor<1x4xi32>, tensor<2x1xi32>) -> tensor<2x4xi32>
+ // CHECK: tosa.sub
+ return %div : tensor<2x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_dynamic_sub_unknown_broadcast_zero_rhs
+func.func @no_fold_dynamic_sub_unknown_broadcast_zero_rhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
+ %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.sub %arg0, %one : (tensor<?x4xi32>, tensor<2x1xi32>) -> tensor<?x4xi32>
+ // CHECK: tosa.sub
+ return %div : tensor<?x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_dynamic_sub_broadcast_zero_rhs
+func.func @fold_dynamic_sub_broadcast_zero_rhs(%arg0: tensor<?x17xi32>) -> tensor<?x17xi32> {
+ %one = "tosa.const"() {values = dense<0> : tensor<1x17xi32>} : () -> tensor<1x17xi32>
+ %div = tosa.sub %arg0, %one : (tensor<?x17xi32>, tensor<1x17xi32>) -> tensor<?x17xi32>
+ // CHECK: return %arg0
+ return %div : tensor<?x17xi32>
+}
+
+// -----
+
// CHECK-LABEL: @fold_sub_splat_i32
func.func @fold_sub_splat_i32() -> tensor<10xi32> {
%one = "tosa.const"() {values = dense<1> : tensor<10xi32>} : () -> tensor<10xi32>
@@ -653,6 +920,123 @@ func.func @no_shift_op_reorder (%arg0 : tensor<44x1xi16>, %arg1 : tensor<1xi8>)
// -----
+// CHECK-LABEL: @fold_select_same_value
+func.func @fold_select_same_value(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
+ %0 = tosa.select %arg0, %arg1, %arg1 : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
+ // CHECK: return %arg1
+ // CHECK-NOT: tosa.select
+ return %0 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_select_same_value_dynamic_broadcast
+func.func @fold_select_same_value_dynamic_broadcast(%arg0: tensor<1x1xi1>, %arg1: tensor<?x3xi32>) -> tensor<?x3xi32> {
+ %0 = tosa.select %arg0, %arg1, %arg1 : (tensor<1x1xi1>, tensor<?x3xi32>, tensor<?x3xi32>) -> tensor<?x3xi32>
+ // CHECK: return %arg1
+ // CHECK-NOT: tosa.select
+ return %0 : tensor<?x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_select_true_value
+func.func @fold_select_true_value(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
+ %c1 = "tosa.const"() {values = dense<1> : tensor<2x3xi1>} : () -> tensor<2x3xi1>
+ %0 = tosa.select %c1, %arg0, %arg1 : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
+ // CHECK: return %arg0
+ // CHECK-NOT: tosa.select
+ return %0 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_select_false_value
+func.func @fold_select_false_value(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
+ %c0 = "tosa.const"() {values = dense<0> : tensor<2x3xi1>} : () -> tensor<2x3xi1>
+ %0 = tosa.select %c0, %arg0, %arg1 : (tensor<2x3xi1>, tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
+ // CHECK: return %arg1
+ // CHECK-NOT: tosa.select
+ return %0 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_select_broadcast_same_value
+func.func @no_fold_select_broadcast_same_value(%arg0: tensor<2x2xi1>, %arg1: tensor<1x1xf32>) -> tensor<2x2xf32> {
+ // CHECK: tosa.select %arg0, %arg1, %arg1
+ %0 = tosa.select %arg0, %arg1, %arg1 : (tensor<2x2xi1>, tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<2x2xf32>
+ return %0 : tensor<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_select_broadcast_true_value
+func.func @no_fold_select_broadcast_true_value(%arg0: tensor<1x1xf32>, %arg1: tensor<2x2xf32>) -> tensor<?x?xf32> {
+ // CHECK: %[[CONST:.*]] = "tosa.const"
+ %0 = "tosa.const"() {values = dense<1> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
+ // CHECK: tosa.select %[[CONST]], %arg0, %arg1
+ %1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<1x1xf32>, tensor<2x2xf32>) -> tensor<?x?xf32>
+ return %1 : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_select_broadcast_false_value
+func.func @no_fold_select_broadcast_false_value(%arg0: tensor<2x2xf32>, %arg1: tensor<1x1xf32>) -> tensor<2x2xf32> {
+ // CHECK: %[[CONST:.*]] = "tosa.const"
+ %0 = "tosa.const"() {values = dense<0> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
+ // CHECK: tosa.select %[[CONST]], %arg0, %arg1
+ %1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<2x2xf32>, tensor<1x1xf32>) -> tensor<2x2xf32>
+ return %1 : tensor<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_select_unknown_broadcast_true_value_dynamic_operand
+func.func @no_fold_select_unknown_broadcast_true_value_dynamic_operand(%arg0: tensor<2x?xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x?xf32> {
+ // CHECK: %[[CONST:.*]] = "tosa.const"
+ %0 = "tosa.const"() {values = dense<1> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
+ // CHECK: tosa.select %[[CONST]], %arg0, %arg1
+ %1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<2x?xf32>, tensor<2x2xf32>) -> tensor<2x?xf32>
+ return %1 : tensor<2x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_select_unknown_broadcast_false_value_dynamic_operand
+func.func @no_fold_select_unknown_broadcast_false_value_dynamic_operand(%arg0: tensor<2x?xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {
+ // CHECK: %[[CONST:.*]] = "tosa.const"
+ %0 = "tosa.const"() {values = dense<0> : tensor<2x2xi1>} : () -> tensor<2x2xi1>
+ // CHECK: tosa.select %[[CONST]], %arg0, %arg1
+ %1 = tosa.select %0, %arg0, %arg1 : (tensor<2x2xi1>, tensor<2x?xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+ return %1 : tensor<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_select_dynamic_broadcast_true_value
+func.func @fold_select_dynamic_broadcast_true_value(%arg0: tensor<2x?xf32>, %arg1: tensor<2x1xf32>) -> tensor<2x?xf32> {
+ %0 = "tosa.const"() {values = dense<1> : tensor<2x1xi1>} : () -> tensor<2x1xi1>
+ %1 = tosa.select %0, %arg0, %arg1 : (tensor<2x1xi1>, tensor<2x?xf32>, tensor<2x1xf32>) -> tensor<2x?xf32>
+ // CHECK: return %arg0
+ // CHECK-NOT: tosa.select
+ return %1 : tensor<2x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_select_dynamic_broadcast_false_value
+func.func @fold_select_dynamic_broadcast_false_value(%arg0: tensor<1x1xf32>, %arg1: tensor<?x2xf32>) -> tensor<?x2xf32> {
+ %0 = "tosa.const"() {values = dense<0> : tensor<1x1xi1>} : () -> tensor<1x1xi1>
+ %1 = tosa.select %0, %arg0, %arg1 : (tensor<1x1xi1>, tensor<1x1xf32>, tensor<?x2xf32>) -> tensor<?x2xf32>
+ // CHECK: return %arg1
+ // CHECK-NOT: tosa.select
+ return %1 : tensor<?x2xf32>
+}
+
+// -----
+
// CHECK-LABEL: @test_fold_add_shape
// CHECK: tosa.const_shape {values = dense<[2, 4, 6, 8, 10, 12]> : tensor<6xindex>} : () -> !tosa.shape<6>
func.func @test_fold_add_shape() -> !tosa.shape<6> {
>From e874634e940c149e20baf91d327525aea096b47d Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Mon, 16 Feb 2026 08:31:55 +0000
Subject: [PATCH 2/2] Fix unit tests: one -> zero
Change-Id: I1aa588a3a55cda560ca9dadae859b41229ca5dc5
---
mlir/test/Dialect/Tosa/constant_folding.mlir | 32 ++++++++++----------
1 file changed, 16 insertions(+), 16 deletions(-)
diff --git a/mlir/test/Dialect/Tosa/constant_folding.mlir b/mlir/test/Dialect/Tosa/constant_folding.mlir
index 5fa5f1d5143f3..1a209093d2578 100644
--- a/mlir/test/Dialect/Tosa/constant_folding.mlir
+++ b/mlir/test/Dialect/Tosa/constant_folding.mlir
@@ -138,8 +138,8 @@ func.func @fold_add_splat_ui8_overflow() -> tensor<10xui8> {
// CHECK-LABEL: @no_fold_add_unknown_broadcast_zero_lhs
func.func @no_fold_add_unknown_broadcast_zero_lhs(%arg0: tensor<1x4xi32>) -> tensor<2x4xi32> {
- %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
- %div = tosa.add %one, %arg0 : (tensor<2x1xi32>, tensor<1x4xi32>) -> tensor<2x4xi32>
+ %zero = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.add %zero, %arg0 : (tensor<2x1xi32>, tensor<1x4xi32>) -> tensor<2x4xi32>
// CHECK: tosa.add
return %div : tensor<2x4xi32>
}
@@ -148,8 +148,8 @@ func.func @no_fold_add_unknown_broadcast_zero_lhs(%arg0: tensor<1x4xi32>) -> ten
// CHECK-LABEL: @no_fold_dynamic_add_unknown_broadcast_zero_lhs
func.func @no_fold_dynamic_add_unknown_broadcast_zero_lhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
- %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
- %div = tosa.add %one, %arg0 : (tensor<2x1xi32>, tensor<?x4xi32>) -> tensor<?x4xi32>
+ %zero = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.add %zero, %arg0 : (tensor<2x1xi32>, tensor<?x4xi32>) -> tensor<?x4xi32>
// CHECK: tosa.add
return %div : tensor<?x4xi32>
}
@@ -158,8 +158,8 @@ func.func @no_fold_dynamic_add_unknown_broadcast_zero_lhs(%arg0: tensor<?x4xi32>
// CHECK-LABEL: @fold_dynamic_add_broadcast_zero_lhs
func.func @fold_dynamic_add_broadcast_zero_lhs(%arg0: tensor<?x17xi32>) -> tensor<?x17xi32> {
- %one = "tosa.const"() {values = dense<0> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
- %div = tosa.add %one, %arg0 : (tensor<1x1xi32>, tensor<?x17xi32>) -> tensor<?x17xi32>
+ %zero = "tosa.const"() {values = dense<0> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
+ %div = tosa.add %zero, %arg0 : (tensor<1x1xi32>, tensor<?x17xi32>) -> tensor<?x17xi32>
// CHECK: return %arg0
return %div : tensor<?x17xi32>
}
@@ -168,8 +168,8 @@ func.func @fold_dynamic_add_broadcast_zero_lhs(%arg0: tensor<?x17xi32>) -> tenso
// CHECK-LABEL: @no_fold_dynamic_add_unknown_broadcast_zero_rhs
func.func @no_fold_dynamic_add_unknown_broadcast_zero_rhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
- %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
- %div = tosa.add %arg0, %one : (tensor<?x4xi32>, tensor<2x1xi32>) -> tensor<?x4xi32>
+ %zero = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.add %arg0, %zero : (tensor<?x4xi32>, tensor<2x1xi32>) -> tensor<?x4xi32>
// CHECK: tosa.add
return %div : tensor<?x4xi32>
}
@@ -178,8 +178,8 @@ func.func @no_fold_dynamic_add_unknown_broadcast_zero_rhs(%arg0: tensor<?x4xi32>
// CHECK-LABEL: @fold_dynamic_add_broadcast_zero_rhs
func.func @fold_dynamic_add_broadcast_zero_rhs(%arg0: tensor<?x17xi32>) -> tensor<?x17xi32> {
- %one = "tosa.const"() {values = dense<0> : tensor<1x17xi32>} : () -> tensor<1x17xi32>
- %div = tosa.add %arg0, %one : (tensor<?x17xi32>, tensor<1x17xi32>) -> tensor<?x17xi32>
+ %zero = "tosa.const"() {values = dense<0> : tensor<1x17xi32>} : () -> tensor<1x17xi32>
+ %div = tosa.add %arg0, %zero : (tensor<?x17xi32>, tensor<1x17xi32>) -> tensor<?x17xi32>
// CHECK: return %arg0
return %div : tensor<?x17xi32>
}
@@ -547,8 +547,8 @@ func.func @fold_sub_zero_rhs_i32(%arg0: tensor<i32>) -> tensor<i32> {
// CHECK-LABEL: @no_fold_sub_unknown_broadcast_zero_rhs
func.func @no_fold_sub_unknown_broadcast_zero_rhs(%arg0: tensor<1x4xi32>) -> tensor<2x4xi32> {
- %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
- %div = tosa.sub %arg0, %one : (tensor<1x4xi32>, tensor<2x1xi32>) -> tensor<2x4xi32>
+ %zero = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.sub %arg0, %zero : (tensor<1x4xi32>, tensor<2x1xi32>) -> tensor<2x4xi32>
// CHECK: tosa.sub
return %div : tensor<2x4xi32>
}
@@ -557,8 +557,8 @@ func.func @no_fold_sub_unknown_broadcast_zero_rhs(%arg0: tensor<1x4xi32>) -> ten
// CHECK-LABEL: @no_fold_dynamic_sub_unknown_broadcast_zero_rhs
func.func @no_fold_dynamic_sub_unknown_broadcast_zero_rhs(%arg0: tensor<?x4xi32>) -> tensor<?x4xi32> {
- %one = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
- %div = tosa.sub %arg0, %one : (tensor<?x4xi32>, tensor<2x1xi32>) -> tensor<?x4xi32>
+ %zero = "tosa.const"() {values = dense<0> : tensor<2x1xi32>} : () -> tensor<2x1xi32>
+ %div = tosa.sub %arg0, %zero : (tensor<?x4xi32>, tensor<2x1xi32>) -> tensor<?x4xi32>
// CHECK: tosa.sub
return %div : tensor<?x4xi32>
}
@@ -567,8 +567,8 @@ func.func @no_fold_dynamic_sub_unknown_broadcast_zero_rhs(%arg0: tensor<?x4xi32>
// CHECK-LABEL: @fold_dynamic_sub_broadcast_zero_rhs
func.func @fold_dynamic_sub_broadcast_zero_rhs(%arg0: tensor<?x17xi32>) -> tensor<?x17xi32> {
- %one = "tosa.const"() {values = dense<0> : tensor<1x17xi32>} : () -> tensor<1x17xi32>
- %div = tosa.sub %arg0, %one : (tensor<?x17xi32>, tensor<1x17xi32>) -> tensor<?x17xi32>
+ %zero = "tosa.const"() {values = dense<0> : tensor<1x17xi32>} : () -> tensor<1x17xi32>
+ %div = tosa.sub %arg0, %zero : (tensor<?x17xi32>, tensor<1x17xi32>) -> tensor<?x17xi32>
// CHECK: return %arg0
return %div : tensor<?x17xi32>
}
More information about the Mlir-commits
mailing list