[Mlir-commits] [mlir] [mlir][tosa] Add tosa.argmin lowering to Linalg (PR #233754)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 20 02:10:10 PDT 2026
https://github.com/karmeh01 created https://github.com/llvm/llvm-project/pull/233754
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.
>From f0b6bafaf9f8fe1236cfa05106489c91b693e952 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.
---
.../Conversion/TosaToLinalg/TosaToLinalg.cpp | 168 +++++++++++++++-
.../TosaToLinalg/tosa-to-linalg.mlir | 188 ++++++++++++++++++
2 files changed, 354 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 867ebbd6ce5ef0..9e90fac4ec9fe0 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 {};
}
@@ -2672,7 +2682,7 @@ struct TileConverter : public OpConversionPattern<tosa::TileOp> {
//
// 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
+// populated by indexed_generic, this buffer is discarded as only the index is
// requested.
//
// The indexed_generic op updates both the maximum value and index if the
@@ -2817,6 +2827,159 @@ class ArgMaxConverter : public OpRewritePattern<tosa::ArgMaxOp> {
bool allowNonFinites;
};
+// Tosa argmin lowering represents the ArgMin op as an linalg.generic
+// op, producing two output buffers.
+//
+// The first output buffer contains the index of the found minimum value. It is
+// initialized to 0 and is resulting integer type.
+//
+// The second output buffer contains the minimum value found. It is initialized
+// to the maximum representable value of the input element type. After being
+// populated by linalg.generic, this buffer is discarded as only the index is
+// requested.
+//
+// The linalg.generic op updates both the minimum value and index if the
+// current value is smaller than the running min.
+class ArgMinConverter : public OpRewritePattern<tosa::ArgMinOp> {
+public:
+ ArgMinConverter(MLIRContext *context, bool allowNonFinites)
+ : OpRewritePattern<tosa::ArgMinOp>(context),
+ allowNonFinites(allowNonFinites) {}
+
+ LogicalResult matchAndRewrite(tosa::ArgMinOp argminOp,
+ PatternRewriter &rewriter) const final {
+ auto loc = argminOp.getLoc();
+ Value input = argminOp.getInput();
+ auto inputTy = cast<ShapedType>(input.getType());
+ auto resultTy = cast<ShapedType>(argminOp.getOutput().getType());
+ auto inElementTy = inputTy.getElementType();
+ auto outElementTy = resultTy.getElementType();
+ int axis = argminOp.getAxis();
+ auto resultMinTy = RankedTensorType::get(resultTy.getShape(), inElementTy);
+
+ if (!isa<IntegerType>(outElementTy))
+ return rewriter.notifyMatchFailure(
+ argminOp,
+ "tosa.argmin to linalg.* requires integer-like result type");
+
+ SmallVector<Value> dynDims;
+ for (int i = 0; i < inputTy.getRank(); i++) {
+ if (inputTy.isDynamicDim(i) && i != axis) {
+ dynDims.push_back(tensor::DimOp::create(rewriter, loc, input, i));
+ }
+ }
+
+ // First fill the output buffer for the index.
+ auto emptyTensorIdx =
+ tensor::EmptyOp::create(rewriter, loc, resultTy.getShape(),
+ outElementTy, dynDims)
+ .getResult();
+ auto fillValueIdx = arith::ConstantOp::create(
+ rewriter, loc, rewriter.getIntegerAttr(outElementTy, 0));
+ auto filledTensorIdx =
+ linalg::FillOp::create(rewriter, loc, ValueRange{fillValueIdx},
+ ValueRange{emptyTensorIdx})
+ .result();
+
+ // Second fill the output buffer for the running min.
+ auto emptyTensorMin =
+ tensor::EmptyOp::create(rewriter, loc, resultTy.getShape(), inElementTy,
+ dynDims)
+ .getResult();
+ auto fillValueMinAttr = createInitialValueForReduceOp(
+ argminOp, inElementTy, rewriter, allowNonFinites);
+
+ if (!fillValueMinAttr)
+ return rewriter.notifyMatchFailure(
+ argminOp, "unsupported tosa.argmin element type");
+
+ auto fillValueMin =
+ arith::ConstantOp::create(rewriter, loc, fillValueMinAttr);
+ auto filledTensorMin =
+ linalg::FillOp::create(rewriter, loc, ValueRange{fillValueMin},
+ ValueRange{emptyTensorMin})
+ .result();
+
+ // We need to reduce along the arg-min axis, with parallel operations along
+ // the rest.
+ SmallVector<utils::IteratorType, 4> iteratorTypes;
+ iteratorTypes.resize(inputTy.getRank(), utils::IteratorType::parallel);
+ iteratorTypes[axis] = utils::IteratorType::reduction;
+
+ SmallVector<AffineExpr, 2> srcExprs;
+ SmallVector<AffineExpr, 2> dstExprs;
+ for (int i = 0, rank = inputTy.getRank(); i != rank; ++i) {
+ srcExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext()));
+ if (axis != i)
+ dstExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext()));
+ }
+
+ bool didEncounterError = false;
+ auto maps = AffineMap::inferFromExprList({srcExprs, dstExprs, dstExprs},
+ rewriter.getContext());
+ auto linalgOp = linalg::GenericOp::create(
+ rewriter, loc, ArrayRef<Type>({resultTy, resultMinTy}), input,
+ ValueRange({filledTensorIdx, filledTensorMin}), maps, iteratorTypes,
+ [&](OpBuilder &nestedBuilder, Location nestedLoc,
+ ValueRange blockArgs) {
+ auto newValue = blockArgs[0];
+ auto oldIndex = blockArgs[1];
+ auto oldValue = blockArgs[2];
+
+ Value newIndex = arith::IndexCastOp::create(
+ rewriter, nestedLoc, oldIndex.getType(),
+ linalg::IndexOp::create(rewriter, loc, axis));
+
+ Value predicate;
+ if (isa<FloatType>(inElementTy)) {
+ if (argminOp.getNanMode() == NanPropagationMode::IGNORE) {
+ // Only update index & min 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::OLT,
+ newValue, oldValue);
+ } else {
+ // Update min value if either of the following is true:
+ // - new value is smaller
+ // - cur min is not NaN and new value is NaN
+ Value lt = arith::CmpFOp::create(rewriter, nestedLoc,
+ arith::CmpFPredicate::ULT,
+ newValue, oldValue);
+ Value oldNonNaN = arith::CmpFOp::create(rewriter, nestedLoc,
+ arith::CmpFPredicate::ORD,
+ oldValue, oldValue);
+ predicate = arith::AndIOp::create(
+ rewriter, nestedLoc, rewriter.getI1Type(), lt, oldNonNaN);
+ }
+ } else if (isa<IntegerType>(inElementTy)) {
+ predicate = arith::CmpIOp::create(rewriter, nestedLoc,
+ arith::CmpIPredicate::slt,
+ newValue, oldValue);
+ } else {
+ didEncounterError = true;
+ return;
+ }
+
+ auto resultMin = 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, resultMin}));
+ });
+
+ if (didEncounterError)
+ return rewriter.notifyMatchFailure(
+ argminOp, "unsupported tosa.argmin element type");
+
+ rewriter.replaceOp(argminOp, linalgOp.getResult(0));
+ return success();
+ }
+
+private:
+ bool allowNonFinites;
+};
+
class GatherConverter : public OpConversionPattern<tosa::GatherOp> {
public:
using OpConversionPattern<tosa::GatherOp>::OpConversionPattern;
@@ -3502,6 +3665,7 @@ void mlir::tosa::populateTosaToLinalgConversionPatterns(
ReduceConverter<tosa::ReduceMaxOp>,
ReduceConverter<tosa::ReduceSumOp>,
ReduceConverter<tosa::ReduceProductOp>,
- ArgMaxConverter>(patterns->getContext(), options.allowNonFinites);
+ ArgMaxConverter,
+ ArgMinConverter>(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..74cf0e5a070a11 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -1907,6 +1907,111 @@ 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) -> ()>
+
+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 = [#map, #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>
+
+ // FINITE: arith.constant 3.40282347E+38 : f32
+ // NONFINITE: arith.constant 0x7F800000 : f32
+ // CHECK: linalg.index
+ // CHECK: arith.index_cast
+ // CHECK: arith.cmpf ult
+ // CHECK: arith.cmpf ord
+ // CHECK: andi
+ // CHECK: select
+ // CHECK: select
+ // CHECK: linalg.yield
+ %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)>
+
+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)>
+
+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: %[[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_]*]]
@@ -2691,6 +2796,53 @@ func.func @argmax_nan_ignore(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) ->
// -----
+// CHECK-LABEL: @argmin_nan_propagate
+func.func @argmin_nan_propagate(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) -> () {
+ // CHECK: linalg.generic
+ // CHECK: arith.cmpf ult
+ // CHECK: arith.cmpf ord
+ // CHECK: andi
+ // CHECK: arith.select
+ // CHECK: arith.select
+ // CHECK-NOT: arith.cmpf uno
+ // CHECK-NOT: arith.select
+ // CHECK: linalg.yield
+ %11 = tosa.argmin %arg0 axis(0) : (tensor<5x4xf32>) -> tensor<4xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @argmin_nan_ignore_int
+func.func @argmin_nan_ignore_int(%arg0: tensor<5x4xi8>, %arg1: tensor<5x4xi8>) -> () {
+ // CHECK: linalg.generic
+ // CHECK: arith.cmpi slt
+ // CHECK: arith.select
+ // CHECK: arith.select
+ // CHECK-NOT: arith.cmpf uno
+ // CHECK-NOT: arith.cmpf uno
+ // CHECK-NOT: arith.select
+ // CHECK-NOT: arith.select
+ // CHECK: linalg.yield
+ %12 = tosa.argmin %arg0 axis(0) nan_mode<IGNORE> : (tensor<5x4xi8>) -> tensor<4xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @argmin_nan_ignore
+func.func @argmin_nan_ignore(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) -> () {
+ // CHECK: linalg.generic
+ // CHECK: arith.cmpf olt
+ // CHECK: arith.select
+ // CHECK: arith.select
+ // CHECK: linalg.yield
+ %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
@@ -2940,6 +3092,42 @@ func.func @argmax_f16(%arg0: tensor<6xf16>) -> () {
%0 = tosa.argmax %arg0 axis(0) : (tensor<6xf16>) -> tensor<i32>
return
}
+// -----
+
+// CHECK-LABEL: @argmin_f16
+func.func @argmin_f16(%arg0: tensor<6xf16>) -> () {
+ // FINITE: arith.constant 6.550400e+04 : f16
+ // NONFINITE: arith.constant 0x7C00 : f16
+ // CHECK: linalg.generic
+ // CHECK: arith.cmpf ult
+ %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