[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:11:09 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: kmehanna (karmeh01)
<details>
<summary>Changes</summary>
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.
---
Full diff: https://github.com/llvm/llvm-project/pull/233754.diff
2 Files Affected:
- (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp (+166-2)
- (modified) mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir (+188)
``````````diff
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 867ebbd6ce5ef..9e90fac4ec9fe 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 67a51e1e2e80f..74cf0e5a070a1 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>
+}
// -----
``````````
</details>
https://github.com/llvm/llvm-project/pull/233754
More information about the Mlir-commits
mailing list