[Mlir-commits] [mlir] [mlir][tosa] Add tosa.argmin lowering to Linalg (PR #233754)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 26 03:12:58 PDT 2026
https://github.com/karmeh01 updated https://github.com/llvm/llvm-project/pull/233754
>From de207ba9d19b30f2dcd78180d91168fd8ec03915 Mon Sep 17 00:00:00 2001
From: Kareem Mehanna <kareemhisham.mehanna at arm.com>
Date: Mon, 19 Oct 2026 16:27:46 +0100
Subject: [PATCH] [mlir][tosa] Add tosa.argmin lowering to Linalg
Add ArgMinConverter following the existing ArgMaxConverter,
including reduction initialization and pattern registration.
Add conversion tests covering static and dynamic shapes, NaN modes,
f16 inputs, and i64 indices.
Change-Id: I6e378cd82c58c07b89f7b45aeeba344fcfaf54d9
---
.../Conversion/TosaToLinalg/TosaToLinalg.cpp | 126 +++---
.../TosaToLinalg/tosa-to-linalg.mlir | 382 ++++++++++++++++--
2 files changed, 420 insertions(+), 88 deletions(-)
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 867ebbd6ce5ef0..18d799e8fd72d4 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -1155,6 +1155,16 @@ static TypedAttr createInitialValueForReduceOp(Operation *op, Type elementTy,
return rewriter.getIntegerAttr(
elementTy, APInt::getSignedMinValue(elementTy.getIntOrFloatBitWidth()));
+ if (isa<tosa::ArgMinOp>(op) && isa<FloatType>(elementTy))
+ return rewriter.getFloatAttr(
+ elementTy,
+ getFloatMinMaxIdentity(cast<FloatType>(elementTy).getFloatSemantics(),
+ /*negative=*/false, allowNonFinites));
+
+ if (isa<tosa::ArgMinOp>(op) && isa<IntegerType>(elementTy))
+ return rewriter.getIntegerAttr(
+ elementTy, APInt::getSignedMaxValue(elementTy.getIntOrFloatBitWidth()));
+
return {};
}
@@ -2664,40 +2674,42 @@ struct TileConverter : public OpConversionPattern<tosa::TileOp> {
}
};
-// Tosa argmax lowering represents the ArgMax op as an linalg.indexed_generic
+// Tosa argmin/argmax lowering represents the op as a linalg.generic
// op, producing two output buffers.
//
-// The first output buffer contains the index of the found maximum value. It is
+// The first output buffer contains the index of the found min/max value. It is
// initialized to 0 and is resulting integer type.
//
-// The second output buffer contains the maximum value found. It is initialized
-// to the minimum representable value of the input element type. After being
-// populated by indexed_generic, this buffer is disgarded as only the index is
+// The second output buffer contains the min/max value found. It is initialized
+// to the max/min representable value of the input element type. After being
+// populated by generic, this buffer is discarded as only the index is
// requested.
//
-// The indexed_generic op updates both the maximum value and index if the
-// current value exceeds the running max.
-class ArgMaxConverter : public OpRewritePattern<tosa::ArgMaxOp> {
+// The generic op updates both the min/max value and index if the
+// current value is lower than/exceeds the running min/max.
+template <typename SrcOp>
+class ArgMinMaxConverter : public OpRewritePattern<SrcOp> {
public:
- ArgMaxConverter(MLIRContext *context, bool allowNonFinites)
- : OpRewritePattern<tosa::ArgMaxOp>(context),
- allowNonFinites(allowNonFinites) {}
+ ArgMinMaxConverter(MLIRContext *context, bool allowNonFinites)
+ : OpRewritePattern<SrcOp>(context), allowNonFinites(allowNonFinites) {}
- LogicalResult matchAndRewrite(tosa::ArgMaxOp argmaxOp,
+ LogicalResult matchAndRewrite(SrcOp argOp,
PatternRewriter &rewriter) const final {
- auto loc = argmaxOp.getLoc();
- Value input = argmaxOp.getInput();
+ constexpr bool isMin = std::is_same_v<SrcOp, tosa::ArgMinOp>;
+
+ auto loc = argOp.getLoc();
+ Value input = argOp.getInput();
auto inputTy = cast<ShapedType>(input.getType());
- auto resultTy = cast<ShapedType>(argmaxOp.getOutput().getType());
+ auto resultTy = cast<ShapedType>(argOp.getOutput().getType());
auto inElementTy = inputTy.getElementType();
auto outElementTy = resultTy.getElementType();
- int axis = argmaxOp.getAxis();
- auto resultMaxTy = RankedTensorType::get(resultTy.getShape(), inElementTy);
+ int axis = argOp.getAxis();
+ auto resultValueTy =
+ RankedTensorType::get(resultTy.getShape(), inElementTy);
if (!isa<IntegerType>(outElementTy))
- return rewriter.notifyMatchFailure(
- argmaxOp,
- "tosa.arg_max to linalg.* requires integer-like result type");
+ return rewriter.notifyMatchFailure(argOp,
+ "requires integer-like result type");
SmallVector<Value> dynDims;
for (int i = 0; i < inputTy.getRank(); i++) {
@@ -2718,27 +2730,25 @@ class ArgMaxConverter : public OpRewritePattern<tosa::ArgMaxOp> {
ValueRange{emptyTensorIdx})
.result();
- // Second fill the output buffer for the running max.
- auto emptyTensorMax =
+ // Second fill the output buffer for the running min/max.
+ auto emptyTensorValue =
tensor::EmptyOp::create(rewriter, loc, resultTy.getShape(), inElementTy,
dynDims)
.getResult();
- auto fillValueMaxAttr = createInitialValueForReduceOp(
- argmaxOp, inElementTy, rewriter, allowNonFinites);
+ auto fillValueAttr = createInitialValueForReduceOp(
+ argOp, inElementTy, rewriter, allowNonFinites);
- if (!fillValueMaxAttr)
- return rewriter.notifyMatchFailure(
- argmaxOp, "unsupported tosa.argmax element type");
+ if (!fillValueAttr)
+ return rewriter.notifyMatchFailure(argOp, "unsupported element type");
- auto fillValueMax =
- arith::ConstantOp::create(rewriter, loc, fillValueMaxAttr);
- auto filledTensorMax =
- linalg::FillOp::create(rewriter, loc, ValueRange{fillValueMax},
- ValueRange{emptyTensorMax})
+ auto fillValue = arith::ConstantOp::create(rewriter, loc, fillValueAttr);
+ auto filledTensorValue =
+ linalg::FillOp::create(rewriter, loc, ValueRange{fillValue},
+ ValueRange{emptyTensorValue})
.result();
- // We need to reduce along the arg-max axis, with parallel operations along
- // the rest.
+ // We need to reduce along the specified axis, with parallel operations
+ // along the rest.
SmallVector<utils::IteratorType, 4> iteratorTypes;
iteratorTypes.resize(inputTy.getRank(), utils::IteratorType::parallel);
iteratorTypes[axis] = utils::IteratorType::reduction;
@@ -2755,8 +2765,8 @@ class ArgMaxConverter : public OpRewritePattern<tosa::ArgMaxOp> {
auto maps = AffineMap::inferFromExprList({srcExprs, dstExprs, dstExprs},
rewriter.getContext());
auto linalgOp = linalg::GenericOp::create(
- rewriter, loc, ArrayRef<Type>({resultTy, resultMaxTy}), input,
- ValueRange({filledTensorIdx, filledTensorMax}), maps, iteratorTypes,
+ rewriter, loc, ArrayRef<Type>({resultTy, resultValueTy}), input,
+ ValueRange({filledTensorIdx, filledTensorValue}), maps, iteratorTypes,
[&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange blockArgs) {
auto newValue = blockArgs[0];
@@ -2769,47 +2779,50 @@ class ArgMaxConverter : public OpRewritePattern<tosa::ArgMaxOp> {
Value predicate;
if (isa<FloatType>(inElementTy)) {
- if (argmaxOp.getNanMode() == NanPropagationMode::IGNORE) {
- // Only update index & max value for non NaN values. If all
+ if (argOp.getNanMode() == NanPropagationMode::IGNORE) {
+ // Only update index & value for non NaN values. If all
// values are NaNs, the initial index will be return which is 0.
- predicate = arith::CmpFOp::create(rewriter, nestedLoc,
- arith::CmpFPredicate::OGT,
- newValue, oldValue);
+ predicate = arith::CmpFOp::create(
+ rewriter, nestedLoc,
+ isMin ? arith::CmpFPredicate::OLT : arith::CmpFPredicate::OGT,
+ newValue, oldValue);
} else {
- // Update max value if either of the following is true:
- // - new value is bigger
- // - cur max is not NaN and new value is NaN
- Value gt = arith::CmpFOp::create(rewriter, nestedLoc,
- arith::CmpFPredicate::UGT,
- newValue, oldValue);
+ // Update value if either of the following is true:
+ // - new value is smaller/bigger
+ // - cur value is not NaN and new value is NaN
+ Value comparison = arith::CmpFOp::create(
+ rewriter, nestedLoc,
+ isMin ? arith::CmpFPredicate::ULT : arith::CmpFPredicate::UGT,
+ newValue, oldValue);
Value oldNonNaN = arith::CmpFOp::create(rewriter, nestedLoc,
arith::CmpFPredicate::ORD,
oldValue, oldValue);
- predicate = arith::AndIOp::create(
- rewriter, nestedLoc, rewriter.getI1Type(), gt, oldNonNaN);
+ predicate = arith::AndIOp::create(rewriter, nestedLoc,
+ rewriter.getI1Type(),
+ comparison, oldNonNaN);
}
} else if (isa<IntegerType>(inElementTy)) {
predicate = arith::CmpIOp::create(rewriter, nestedLoc,
- arith::CmpIPredicate::sgt,
+ isMin ? arith::CmpIPredicate::slt
+ : arith::CmpIPredicate::sgt,
newValue, oldValue);
} else {
didEncounterError = true;
return;
}
- auto resultMax = arith::SelectOp::create(
+ auto resultValue = arith::SelectOp::create(
rewriter, nestedLoc, predicate, newValue, oldValue);
auto resultIndex = arith::SelectOp::create(
rewriter, nestedLoc, predicate, newIndex, oldIndex);
linalg::YieldOp::create(nestedBuilder, nestedLoc,
- ValueRange({resultIndex, resultMax}));
+ ValueRange({resultIndex, resultValue}));
});
if (didEncounterError)
- return rewriter.notifyMatchFailure(
- argmaxOp, "unsupported tosa.argmax element type");
+ return rewriter.notifyMatchFailure(argOp, "unsupported element type");
- rewriter.replaceOp(argmaxOp, linalgOp.getResult(0));
+ rewriter.replaceOp(argOp, linalgOp.getResult(0));
return success();
}
@@ -3502,6 +3515,7 @@ void mlir::tosa::populateTosaToLinalgConversionPatterns(
ReduceConverter<tosa::ReduceMaxOp>,
ReduceConverter<tosa::ReduceSumOp>,
ReduceConverter<tosa::ReduceProductOp>,
- ArgMaxConverter>(patterns->getContext(), options.allowNonFinites);
+ ArgMinMaxConverter<tosa::ArgMaxOp>,
+ ArgMinMaxConverter<tosa::ArgMinOp>>(patterns->getContext(), options.allowNonFinites);
// clang-format on
}
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index 67a51e1e2e80f9..ac4236ea26b3e7 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -1808,6 +1808,8 @@ func.func @tile_dyn_multiples(%arg0 : tensor<2x3xi8>) -> () {
// CHECK: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)>
// CHECK: #[[$MAP4:.*]] = affine_map<(d0) -> ()>
+// CHECK-LABEL: @argmax(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]+]]: tensor<3x2xi32>, [[FLOAT_INPUT:%.+]]: tensor<6xf32>
func.func @argmax(%arg0 : tensor<3x2xi32>, %arg1 : tensor<6xf32>) -> () {
// CHECK: [[IDX_INIT:%.+]] = tensor.empty()
// CHECK: [[IDX_MIN:%.+]] = arith.constant 0 : i32
@@ -1831,7 +1833,7 @@ func.func @argmax(%arg0 : tensor<3x2xi32>, %arg1 : tensor<6xf32>) -> () {
// CHECK: [[VAL_INIT:%.+]] = tensor.empty()
// CHECK: [[VAL_MIN:%.+]] = arith.constant -2147483648
// CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]]{{.*}}outs([[VAL_INIT]]
- // CHECK: linalg.generic {indexing_maps = [#map, #map2, #map2], iterator_types = ["parallel", "reduction"]} ins(%[[ARG0]] : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>)
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins(%[[ARG0]] : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>)
// CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32
// CHECK: [[IDX:%.+]] = linalg.index 1
// CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]]
@@ -1841,16 +1843,24 @@ func.func @argmax(%arg0 : tensor<3x2xi32>, %arg1 : tensor<6xf32>) -> () {
// CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]]
%1 = tosa.argmax %arg0 axis(1) : (tensor<3x2xi32>) -> tensor<3xi32>
- // FINITE: arith.constant -3.40282347E+38 : f32
- // NONFINITE: arith.constant 0xFF800000 : f32
- // CHECK: linalg.index
- // CHECK: arith.index_cast
- // CHECK: arith.cmpf ugt
- // CHECK: arith.cmpf ord
- // CHECK: andi
- // CHECK: select
- // CHECK: select
- // CHECK: linalg.yield
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<i32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<i32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<f32>
+ // FINITE: [[VAL_MIN:%.+]] = arith.constant -3.40282347E+38 : f32
+ // NONFINITE: [[VAL_MIN:%.+]] = arith.constant 0xFF800000 : f32
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]] : f32) outs([[VAL_INIT]] : tensor<f32>)
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]], #[[$MAP4]]], iterator_types = ["reduction"]}
+ // CHECK-SAME: ins([[FLOAT_INPUT]] : tensor<6xf32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<i32>, tensor<f32>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: f32, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: f32):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpf ugt, [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[OLD_NOT_NAN:%.+]] = arith.cmpf ord, [[OLD_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[UPDATE:%.+]] = arith.andi [[CMP]], [[OLD_NOT_NAN]] : i1
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[UPDATE]], [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[UPDATE]], [[CAST]], [[OLD_INDEX]] : i32
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, f32
%2 = tosa.argmax %arg1 axis(0) : (tensor<6xf32>) -> tensor<i32>
return
@@ -1861,6 +1871,8 @@ func.func @argmax(%arg0 : tensor<3x2xi32>, %arg1 : tensor<6xf32>) -> () {
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
+// CHECK-LABEL: @argmax_dyn_non_axis(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]+]]: tensor<3x?xi32>
func.func @argmax_dyn_non_axis(%arg0 : tensor<3x?xi32>) -> () {
// CHECK: %[[CST1:.+]] = arith.constant 1
// CHECK: %[[DYN:.+]] = tensor.dim %[[ARG0]], %[[CST1]]
@@ -1887,6 +1899,8 @@ func.func @argmax_dyn_non_axis(%arg0 : tensor<3x?xi32>) -> () {
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)>
+// CHECK-LABEL: @argmax_dyn_axis(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]+]]: tensor<3x?xi32>
func.func @argmax_dyn_axis(%arg0 : tensor<3x?xi32>) -> () {
// CHECK: %[[IDX_INIT:.+]] = tensor.empty()
// CHECK: %[[IDX_MIN:.+]] = arith.constant 0 : i32
@@ -1895,6 +1909,7 @@ func.func @argmax_dyn_axis(%arg0 : tensor<3x?xi32>) -> () {
// CHECK: %[[VAL_MIN:.+]] = arith.constant -2147483648
// CHECK: %[[VAL_FILL:.+]] = linalg.fill ins(%[[VAL_MIN]]{{.*}}outs(%[[VAL_INIT]]
// CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "reduction"]} ins(%[[ARG0]] : tensor<3x?xi32>) outs(%[[IDX_FILL]], %[[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>)
+ // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32
// CHECK: %[[IDX:.+]] = linalg.index 1
// CHECK: %[[CAST:.+]] = arith.index_cast %[[IDX]]
// CHECK: %[[CMP:.+]] = arith.cmpi sgt, %[[ARG1]], %[[ARG3]]
@@ -1907,6 +1922,126 @@ func.func @argmax_dyn_axis(%arg0 : tensor<3x?xi32>) -> () {
// -----
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
+// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
+// CHECK: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)>
+// CHECK: #[[$MAP4:.*]] = affine_map<(d0) -> ()>
+
+// CHECK-LABEL: @argmin(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]+]]: tensor<3x2xi32>, [[FLOAT_INPUT:%.+]]: tensor<6xf32>
+func.func @argmin(%arg0 : tensor<3x2xi32>, %arg1 : tensor<6xf32>) -> () {
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty()
+ // CHECK: [[IDX_MIN:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_MIN]]{{.*}}outs([[IDX_INIT]]
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty()
+ // CHECK: [[VAL_MAX:%.+]] = arith.constant 2147483647
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MAX]]{{.*}}outs([[VAL_INIT]]
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins(%[[ARG0]] : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<2xi32>, tensor<2xi32>)
+ // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]]
+ // CHECK: [[CMP:%.+]] = arith.cmpi slt, %[[ARG1]], %[[ARG3]]
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], %[[ARG1]], %[[ARG3]]
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], %[[ARG2]]
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]]
+ %0 = tosa.argmin %arg0 axis(0) : (tensor<3x2xi32>) -> tensor<2xi32>
+
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty()
+ // CHECK: [[IDX_MIN:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_MIN]]{{.*}}outs([[IDX_INIT]]
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty()
+ // CHECK: [[VAL_MAX:%.+]] = arith.constant 2147483647
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MAX]]{{.*}}outs([[VAL_INIT]]
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins(%[[ARG0]] : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>)
+ // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32
+ // CHECK: [[IDX:%.+]] = linalg.index 1
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]]
+ // CHECK: [[CMP:%.+]] = arith.cmpi slt, %[[ARG1]], %[[ARG3]]
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], %[[ARG1]], %[[ARG3]]
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], %[[ARG2]]
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]]
+ %1 = tosa.argmin %arg0 axis(1) : (tensor<3x2xi32>) -> tensor<3xi32>
+
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<i32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<i32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<f32>
+ // FINITE: [[VAL_MAX:%.+]] = arith.constant 3.40282347E+38 : f32
+ // NONFINITE: [[VAL_MAX:%.+]] = arith.constant 0x7F800000 : f32
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MAX]] : f32) outs([[VAL_INIT]] : tensor<f32>)
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]], #[[$MAP4]]], iterator_types = ["reduction"]}
+ // CHECK-SAME: ins([[FLOAT_INPUT]] : tensor<6xf32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<i32>, tensor<f32>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: f32, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: f32):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpf ult, [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[OLD_NOT_NAN:%.+]] = arith.cmpf ord, [[OLD_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[UPDATE:%.+]] = arith.andi [[CMP]], [[OLD_NOT_NAN]] : i1
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[UPDATE]], [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[UPDATE]], [[CAST]], [[OLD_INDEX]] : i32
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, f32
+ %2 = tosa.argmin %arg1 axis(0) : (tensor<6xf32>) -> tensor<i32>
+
+ return
+}
+
+// -----
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
+
+// CHECK-LABEL: @argmin_dyn_non_axis(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]+]]: tensor<3x?xi32>
+func.func @argmin_dyn_non_axis(%arg0 : tensor<3x?xi32>) -> () {
+ // CHECK: %[[CST1:.+]] = arith.constant 1
+ // CHECK: %[[DYN:.+]] = tensor.dim %[[ARG0]], %[[CST1]]
+ // CHECK: %[[IDX_INIT:.+]] = tensor.empty(%[[DYN]])
+ // CHECK: %[[IDX_MIN:.+]] = arith.constant 0 : i32
+ // CHECK: %[[IDX_FILL:.+]] = linalg.fill ins(%[[IDX_MIN]]{{.*}}outs(%[[IDX_INIT]]
+ // CHECK: %[[VAL_INIT:.+]] = tensor.empty(%[[DYN]])
+ // CHECK: %[[VAL_MAX:.+]] = arith.constant 2147483647
+ // CHECK: %[[VAL_FILL:.+]] = linalg.fill ins(%[[VAL_MAX]]{{.*}}outs(%[[VAL_INIT]]
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins(%[[ARG0]] : tensor<3x?xi32>) outs(%[[IDX_FILL]], %[[VAL_FILL]] : tensor<?xi32>, tensor<?xi32>)
+ // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32
+ // CHECK: %[[IDX:.+]] = linalg.index 0
+ // CHECK: %[[CAST:.+]] = arith.index_cast %[[IDX]]
+ // CHECK: %[[CMP:.+]] = arith.cmpi slt, %[[ARG1]], %[[ARG3]]
+ // CHECK: %[[SELECT_VAL:.+]] = arith.select %[[CMP]], %[[ARG1]], %[[ARG3]]
+ // CHECK: %[[SELECT_IDX:.+]] = arith.select %[[CMP]], %[[CAST]], %[[ARG2]]
+ // CHECK: linalg.yield %[[SELECT_IDX]], %[[SELECT_VAL]]
+ %0 = tosa.argmin %arg0 axis(0) : (tensor<3x?xi32>) -> tensor<?xi32>
+ return
+}
+
+// -----
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-LABEL: @argmin_dyn_axis(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]+]]: tensor<3x?xi32>
+func.func @argmin_dyn_axis(%arg0 : tensor<3x?xi32>) -> () {
+ // CHECK: %[[IDX_INIT:.+]] = tensor.empty()
+ // CHECK: %[[IDX_MIN:.+]] = arith.constant 0 : i32
+ // CHECK: %[[IDX_FILL:.+]] = linalg.fill ins(%[[IDX_MIN]]{{.*}}outs(%[[IDX_INIT]]
+ // CHECK: %[[VAL_INIT:.+]] = tensor.empty()
+ // CHECK: %[[VAL_MAX:.+]] = arith.constant 2147483647
+ // CHECK: %[[VAL_FILL:.+]] = linalg.fill ins(%[[VAL_MAX]]{{.*}}outs(%[[VAL_INIT]]
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "reduction"]} ins(%[[ARG0]] : tensor<3x?xi32>) outs(%[[IDX_FILL]], %[[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>)
+ // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32
+ // CHECK: %[[IDX:.+]] = linalg.index 1
+ // CHECK: %[[CAST:.+]] = arith.index_cast %[[IDX]]
+ // CHECK: %[[CMP:.+]] = arith.cmpi slt, %[[ARG1]], %[[ARG3]]
+ // CHECK: %[[SELECT_VAL:.+]] = arith.select %[[CMP]], %[[ARG1]], %[[ARG3]]
+ // CHECK: %[[SELECT_IDX:.+]] = arith.select %[[CMP]], %[[CAST]], %[[ARG2]]
+ // CHECK: linalg.yield %[[SELECT_IDX]], %[[SELECT_VAL]]
+ %0 = tosa.argmin %arg0 axis(1) : (tensor<3x?xi32>) -> tensor<3xi32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @gather_float
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]
@@ -2645,16 +2780,28 @@ func.func @maximum_nan_ignore(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) ->
// -----
// CHECK-LABEL: @argmax_nan_propagate
+// CHECK-SAME: ([[INPUT:%[0-9a-zA-Z_]+]]: tensor<5x4xf32>
func.func @argmax_nan_propagate(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) -> () {
- // CHECK: linalg.generic
- // CHECK: arith.cmpf ugt
- // CHECK: arith.cmpf ord
- // CHECK: andi
- // CHECK: arith.select
- // CHECK: arith.select
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<4xi32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<4xi32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<4xf32>
+ // FINITE: [[VAL_MIN:%.+]] = arith.constant -3.40282347E+38 : f32
+ // NONFINITE: [[VAL_MIN:%.+]] = arith.constant 0xFF800000 : f32
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]] : f32) outs([[VAL_INIT]] : tensor<4xf32>)
+ // CHECK: linalg.generic
+ // CHECK-SAME: ins([[INPUT]] : tensor<5x4xf32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<4xi32>, tensor<4xf32>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: f32, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: f32):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpf ugt, [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[OLD_NOT_NAN:%.+]] = arith.cmpf ord, [[OLD_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[UPDATE:%.+]] = arith.andi [[CMP]], [[OLD_NOT_NAN]] : i1
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[UPDATE]], [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[UPDATE]], [[CAST]], [[OLD_INDEX]] : i32
// CHECK-NOT: arith.cmpf uno
// CHECK-NOT: arith.select
- // CHECK: linalg.yield
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, f32
%11 = tosa.argmax %arg0 axis(0) : (tensor<5x4xf32>) -> tensor<4xi32>
return
}
@@ -2662,16 +2809,27 @@ func.func @argmax_nan_propagate(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>)
// -----
// CHECK-LABEL: @argmax_nan_ignore_int
+// CHECK-SAME: ([[INPUT:%[0-9a-zA-Z_]+]]: tensor<5x4xi8>
func.func @argmax_nan_ignore_int(%arg0: tensor<5x4xi8>, %arg1: tensor<5x4xi8>) -> () {
- // CHECK: linalg.generic
- // CHECK: arith.cmpi sgt
- // CHECK: arith.select
- // CHECK: arith.select
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<4xi32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<4xi32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<4xi8>
+ // CHECK: [[VAL_MIN:%.+]] = arith.constant -128 : i8
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]] : i8) outs([[VAL_INIT]] : tensor<4xi8>)
+ // CHECK: linalg.generic
+ // CHECK-SAME: ins([[INPUT]] : tensor<5x4xi8>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<4xi32>, tensor<4xi8>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: i8, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: i8):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpi sgt, [[NEW_VALUE]], [[OLD_VALUE]] : i8
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], [[NEW_VALUE]], [[OLD_VALUE]] : i8
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], [[OLD_INDEX]] : i32
// CHECK-NOT: arith.cmpf uno
// CHECK-NOT: arith.cmpf uno
// CHECK-NOT: arith.select
// CHECK-NOT: arith.select
- // CHECK: linalg.yield
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, i8
%12 = tosa.argmax %arg0 axis(0) nan_mode<IGNORE> : (tensor<5x4xi8>) -> tensor<4xi32>
return
}
@@ -2679,18 +2837,112 @@ func.func @argmax_nan_ignore_int(%arg0: tensor<5x4xi8>, %arg1: tensor<5x4xi8>) -
// -----
// CHECK-LABEL: @argmax_nan_ignore
+// CHECK-SAME: ([[INPUT:%[0-9a-zA-Z_]+]]: tensor<5x4xf32>
func.func @argmax_nan_ignore(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) -> () {
- // CHECK: linalg.generic
- // CHECK: arith.cmpf ogt
- // CHECK: arith.select
- // CHECK: arith.select
- // CHECK: linalg.yield
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<4xi32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<4xi32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<4xf32>
+ // FINITE: [[VAL_MIN:%.+]] = arith.constant -3.40282347E+38 : f32
+ // NONFINITE: [[VAL_MIN:%.+]] = arith.constant 0xFF800000 : f32
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]] : f32) outs([[VAL_INIT]] : tensor<4xf32>)
+ // CHECK: linalg.generic
+ // CHECK-SAME: ins([[INPUT]] : tensor<5x4xf32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<4xi32>, tensor<4xf32>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: f32, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: f32):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpf ogt, [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], [[OLD_INDEX]] : i32
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, f32
%12 = tosa.argmax %arg0 axis(0) nan_mode<IGNORE> : (tensor<5x4xf32>) -> tensor<4xi32>
return
}
// -----
+// CHECK-LABEL: @argmin_nan_propagate
+// CHECK-SAME: ([[INPUT:%[0-9a-zA-Z_]+]]: tensor<5x4xf32>
+func.func @argmin_nan_propagate(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) -> () {
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<4xi32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<4xi32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<4xf32>
+ // FINITE: [[VAL_MAX:%.+]] = arith.constant 3.40282347E+38 : f32
+ // NONFINITE: [[VAL_MAX:%.+]] = arith.constant 0x7F800000 : f32
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MAX]] : f32) outs([[VAL_INIT]] : tensor<4xf32>)
+ // CHECK: linalg.generic
+ // CHECK-SAME: ins([[INPUT]] : tensor<5x4xf32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<4xi32>, tensor<4xf32>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: f32, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: f32):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpf ult, [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[OLD_NOT_NAN:%.+]] = arith.cmpf ord, [[OLD_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[UPDATE:%.+]] = arith.andi [[CMP]], [[OLD_NOT_NAN]] : i1
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[UPDATE]], [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[UPDATE]], [[CAST]], [[OLD_INDEX]] : i32
+ // CHECK-NOT: arith.cmpf uno
+ // CHECK-NOT: arith.select
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, f32
+ %11 = tosa.argmin %arg0 axis(0) : (tensor<5x4xf32>) -> tensor<4xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @argmin_nan_ignore_int
+// CHECK-SAME: ([[INPUT:%[0-9a-zA-Z_]+]]: tensor<5x4xi8>
+func.func @argmin_nan_ignore_int(%arg0: tensor<5x4xi8>, %arg1: tensor<5x4xi8>) -> () {
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<4xi32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<4xi32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<4xi8>
+ // CHECK: [[VAL_MAX:%.+]] = arith.constant 127 : i8
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MAX]] : i8) outs([[VAL_INIT]] : tensor<4xi8>)
+ // CHECK: linalg.generic
+ // CHECK-SAME: ins([[INPUT]] : tensor<5x4xi8>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<4xi32>, tensor<4xi8>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: i8, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: i8):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpi slt, [[NEW_VALUE]], [[OLD_VALUE]] : i8
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], [[NEW_VALUE]], [[OLD_VALUE]] : i8
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], [[OLD_INDEX]] : i32
+ // CHECK-NOT: arith.cmpf uno
+ // CHECK-NOT: arith.cmpf uno
+ // CHECK-NOT: arith.select
+ // CHECK-NOT: arith.select
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, i8
+ %12 = tosa.argmin %arg0 axis(0) nan_mode<IGNORE> : (tensor<5x4xi8>) -> tensor<4xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @argmin_nan_ignore
+// CHECK-SAME: ([[INPUT:%[0-9a-zA-Z_]+]]: tensor<5x4xf32>
+func.func @argmin_nan_ignore(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) -> () {
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<4xi32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<4xi32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<4xf32>
+ // FINITE: [[VAL_MAX:%.+]] = arith.constant 3.40282347E+38 : f32
+ // NONFINITE: [[VAL_MAX:%.+]] = arith.constant 0x7F800000 : f32
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MAX]] : f32) outs([[VAL_INIT]] : tensor<4xf32>)
+ // CHECK: linalg.generic
+ // CHECK-SAME: ins([[INPUT]] : tensor<5x4xf32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<4xi32>, tensor<4xf32>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: f32, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: f32):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpf olt, [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], [[NEW_VALUE]], [[OLD_VALUE]] : f32
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], [[OLD_INDEX]] : i32
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, f32
+ %12 = tosa.argmin %arg0 axis(0) nan_mode<IGNORE> : (tensor<5x4xf32>) -> tensor<4xi32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @clamp_nan_propagate
func.func @clamp_nan_propagate(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) -> () {
// CHECK: linalg.generic
@@ -2932,14 +3184,80 @@ func.func @reduce_max_nan_ignore_seed(%arg0: tensor<5x4xf32>) -> () {
// -----
// CHECK-LABEL: @argmax_f16
+// CHECK-SAME: ([[INPUT:%[0-9a-zA-Z_]+]]: tensor<6xf16>
func.func @argmax_f16(%arg0: tensor<6xf16>) -> () {
- // FINITE: arith.constant -6.550400e+04 : f16
- // NONFINITE: arith.constant 0xFC00 : f16
- // CHECK: linalg.generic
- // CHECK: arith.cmpf ugt
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<i32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<i32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<f16>
+ // FINITE: [[VAL_MIN:%.+]] = arith.constant -6.550400e+04 : f16
+ // NONFINITE: [[VAL_MIN:%.+]] = arith.constant 0xFC00 : f16
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]] : f16) outs([[VAL_INIT]] : tensor<f16>)
+ // CHECK: linalg.generic
+ // CHECK-SAME: ins([[INPUT]] : tensor<6xf16>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<i32>, tensor<f16>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: f16, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: f16):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpf ugt, [[NEW_VALUE]], [[OLD_VALUE]] : f16
+ // CHECK: [[OLD_NOT_NAN:%.+]] = arith.cmpf ord, [[OLD_VALUE]], [[OLD_VALUE]] : f16
+ // CHECK: [[UPDATE:%.+]] = arith.andi [[CMP]], [[OLD_NOT_NAN]] : i1
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[UPDATE]], [[NEW_VALUE]], [[OLD_VALUE]] : f16
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[UPDATE]], [[CAST]], [[OLD_INDEX]] : i32
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, f16
%0 = tosa.argmax %arg0 axis(0) : (tensor<6xf16>) -> tensor<i32>
return
}
+// -----
+
+// CHECK-LABEL: @argmin_f16
+// CHECK-SAME: ([[INPUT:%[0-9a-zA-Z_]+]]: tensor<6xf16>
+func.func @argmin_f16(%arg0: tensor<6xf16>) -> () {
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<i32>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i32
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i32) outs([[IDX_INIT]] : tensor<i32>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<f16>
+ // FINITE: [[VAL_MAX:%.+]] = arith.constant 6.550400e+04 : f16
+ // NONFINITE: [[VAL_MAX:%.+]] = arith.constant 0x7C00 : f16
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MAX]] : f16) outs([[VAL_INIT]] : tensor<f16>)
+ // CHECK: linalg.generic
+ // CHECK-SAME: ins([[INPUT]] : tensor<6xf16>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<i32>, tensor<f16>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: f16, [[OLD_INDEX:%.+]]: i32, [[OLD_VALUE:%.+]]: f16):
+ // CHECK: [[IDX:%.+]] = linalg.index 0
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i32
+ // CHECK: [[CMP:%.+]] = arith.cmpf ult, [[NEW_VALUE]], [[OLD_VALUE]] : f16
+ // CHECK: [[OLD_NOT_NAN:%.+]] = arith.cmpf ord, [[OLD_VALUE]], [[OLD_VALUE]] : f16
+ // CHECK: [[UPDATE:%.+]] = arith.andi [[CMP]], [[OLD_NOT_NAN]] : i1
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[UPDATE]], [[NEW_VALUE]], [[OLD_VALUE]] : f16
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[UPDATE]], [[CAST]], [[OLD_INDEX]] : i32
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i32, f16
+ %0 = tosa.argmin %arg0 axis(0) : (tensor<6xf16>) -> tensor<i32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @argmin_i64_index
+// CHECK-SAME: [[ARGMIN_INPUT:%.*]]: tensor<3x2xi8>
+func.func @argmin_i64_index(%arg0: tensor<3x2xi8>) -> tensor<3xi64> {
+ // CHECK: [[IDX_INIT:%.+]] = tensor.empty() : tensor<3xi64>
+ // CHECK: [[IDX_ZERO:%.+]] = arith.constant 0 : i64
+ // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_ZERO]] : i64) outs([[IDX_INIT]] : tensor<3xi64>)
+ // CHECK: [[VAL_INIT:%.+]] = tensor.empty() : tensor<3xi8>
+ // CHECK: [[VAL_MAX:%.+]] = arith.constant 127 : i8
+ // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MAX]] : i8) outs([[VAL_INIT]] : tensor<3xi8>)
+ // CHECK: [[RESULT:%.+]]:2 = linalg.generic
+ // CHECK-SAME: ins([[ARGMIN_INPUT]] : tensor<3x2xi8>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<3xi64>, tensor<3xi8>)
+ // CHECK: ^bb0([[NEW_VALUE:%.+]]: i8, [[OLD_INDEX:%.+]]: i64, [[OLD_VALUE:%.+]]: i8):
+ // CHECK: [[IDX:%.+]] = linalg.index 1
+ // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] : index to i64
+ // CHECK: [[CMP:%.+]] = arith.cmpi slt, [[NEW_VALUE]], [[OLD_VALUE]] : i8
+ // CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], [[NEW_VALUE]], [[OLD_VALUE]] : i8
+ // CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], [[OLD_INDEX]] : i64
+ // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] : i64, i8
+ %0 = tosa.argmin %arg0 axis(1) : (tensor<3x2xi8>) -> tensor<3xi64>
+ // CHECK: return [[RESULT]]#0 : tensor<3xi64>
+ return %0 : tensor<3xi64>
+}
// -----
More information about the Mlir-commits
mailing list