[Mlir-commits] [mlir] [mlir][tosa] Add support for ARGMIN (PR #227997)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 1 01:52:18 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tosa
Author: Iliyan Georgiev (iliyan-georgiev-arm)
<details>
<summary>Changes</summary>
- Add ARGMIN table gen
- Add validators and folders
- Add positive and negative tests
- Add compliance info
---
Patch is 41.32 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/227997.diff
20 Files Affected:
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc (+30)
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td (+36)
- (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (+17)
- (modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+73-47)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp (+12-8)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp (+1)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp (+16)
- (modified) mlir/test/Dialect/Tosa/canonicalize.mlir (+17)
- (modified) mlir/test/Dialect/Tosa/constant_folding.mlir (+36)
- (modified) mlir/test/Dialect/Tosa/invalid.mlir (+16)
- (modified) mlir/test/Dialect/Tosa/invalid_extension.mlir (+7)
- (modified) mlir/test/Dialect/Tosa/level_check.mlir (+9-1)
- (modified) mlir/test/Dialect/Tosa/ops.mlir (+35-14)
- (modified) mlir/test/Dialect/Tosa/profile_pro_int_unsupported.mlir (+7)
- (modified) mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir (+24)
- (modified) mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir (-17)
- (modified) mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir (+9)
- (modified) mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir (+51)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir (+8)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir (+24)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index 216dc65da9610..d4d1e676a171f 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -6,6 +6,13 @@ profileComplianceMap = {
{{Profile::pro_fp},
{{{fp16T, i32T}, SpecificationVersion::V_1_0},
{{fp32T, i32T}, SpecificationVersion::V_1_0}}}}},
+ {"tosa.argmin",
+ {{{Profile::pro_int}, {{{i8T, i32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Profile::pro_fp},
+ {{{fp16T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
{"tosa.avg_pool2d",
{{{Profile::pro_int},
{{{i8T, i8T, i8T, i32T, i8T}, SpecificationVersion::V_1_0}}},
@@ -595,6 +602,29 @@ extensionComplianceMap = {
{{Extension::bf16, Extension::int64},
{{{bf16T, i64T}, SpecificationVersion::V_1_1_DRAFT}},
allOf}}},
+ {"tosa.argmin",
+ {{{Extension::int16}, {{{i16T, i32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::int64},
+ {{{i8T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i16T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::fp8e4m3},
+ {{{fp8e4m3T, i32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::fp8e5m2},
+ {{{fp8e5m2T, i32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::bf16}, {{{bf16T, i32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::fp8e4m3, Extension::int64},
+ {{{fp8e4m3T, i64T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::fp8e5m2, Extension::int64},
+ {{{fp8e5m2T, i64T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::bf16, Extension::int64},
+ {{{bf16T, i64T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf}}},
{"tosa.avg_pool2d",
{{{Extension::int16},
{{{i16T, i16T, i16T, i32T, i16T}, SpecificationVersion::V_1_0}}},
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 01ef662cc4f71..6f88c6fc589bc 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -63,6 +63,42 @@ def Tosa_ArgMaxOp : Tosa_InferShapedTypeOp<"argmax", [NoMemoryEffect]> {
"functional-type(operands, results)";
}
+
+//===----------------------------------------------------------------------===//
+// Operator: argmin
+//===----------------------------------------------------------------------===//
+def Tosa_ArgMinOp : Tosa_InferShapedTypeOp<"argmin", [NoMemoryEffect]> {
+ let summary = "Perform argmin on the input.";
+
+ let description = [{
+ This returns the index with the smallest value across the given axis of the
+ input tensor. If multiple locations have equal values, returns the first
+ match along the search axis.
+
+ This operation is not pure. Undefined behaviour may occur if the min index
+ is out of bounds for the output data type.
+ }];
+
+ let arguments = (ins
+ Tosa_TensorAtLeast1D: $input,
+ I32Attr: $axis,
+ DefaultValuedAttr<Tosa_NanPropagationModeAttr, "::mlir::tosa::NanPropagationMode::PROPAGATE">:$nan_mode
+ );
+
+ let results = (outs
+ Tosa_Tensor: $output
+ );
+
+ let hasFolder = 1;
+ let hasVerifier = 1;
+
+ let assemblyFormat =
+ "operands `axis` `(` $axis `)` "
+ "oilist(`nan_mode` `` $nan_mode) attr-dict `:` "
+ "functional-type(operands, results)";
+}
+
+
//===----------------------------------------------------------------------===//
// Accumulator types.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 533d3f65095e2..21dd49875d420 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1772,6 +1772,23 @@ OpFoldResult ArgMaxOp::fold(FoldAdaptor adaptor) {
return {};
}
+OpFoldResult ArgMinOp::fold(FoldAdaptor adaptor) {
+ auto inputTy = llvm::dyn_cast<RankedTensorType>(getInput().getType());
+ auto outputTy = llvm::dyn_cast<RankedTensorType>(getType());
+ if (!inputTy || !outputTy || !inputTy.hasStaticShape() ||
+ !outputTy.hasStaticShape())
+ return {};
+
+ const Type outputElementTy = getElementTypeOrSelf(outputTy);
+ if (inputTy.getDimSize(getAxis()) == 1 && outputElementTy.isInteger()) {
+ const auto outputElemIntTy = cast<IntegerType>(outputElementTy);
+ const APInt zero = APInt::getZero(outputElemIntTy.getWidth());
+ return DenseElementsAttr::get(outputTy, zero);
+ }
+
+ return {};
+}
+
OpFoldResult IntDivOp::fold(FoldAdaptor adaptor) {
auto lhsTy = llvm::dyn_cast<RankedTensorType>(getInput1().getType());
auto rhsTy = llvm::dyn_cast<RankedTensorType>(getInput2().getType());
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index feac95c5a021a..17782b07f54b2 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -586,9 +586,71 @@ LogicalResult mlir::tosa::BlockScaledType::convertFromAttribute(
result);
}
+//===----------------------------------------------------------------------===//
+// TOSA Operator shape inference
+//===----------------------------------------------------------------------===//
+
+template <typename A, std::enable_if_t<std::is_same_v<A, ArgMaxOp::Adaptor> ||
+ std::is_same_v<A, ArgMinOp::Adaptor>,
+ int> = 0>
+LogicalResult inferArgMaxMinReturnTypeComponents(
+ MLIRContext *context, ::std::optional<Location> location, A adaptor,
+ SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+ ShapeAdaptor inputShape(adaptor.getInput().getType());
+ IntegerAttr axis = adaptor.getProperties().axis;
+ int32_t axisVal = axis.getValue().getSExtValue();
+
+ if (!inputShape.hasRank()) {
+ inferredReturnShapes.push_back(ShapedTypeComponents());
+ return success();
+ }
+
+ const auto inputRank = inputShape.getRank();
+ SmallVector<int64_t> outShape;
+ outShape.reserve(inputRank - 1);
+ for (int i = 0, s = inputRank; i < s; i++) {
+ if (i == axisVal)
+ continue;
+ outShape.push_back(inputShape.getDimSize(i));
+ }
+
+ inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// TOSA Operator Verifiers.
//===----------------------------------------------------------------------===//
+template <typename T>
+LogicalResult argMaxMinVerify(T op) {
+ const ShapedType resultType = llvm::cast<ShapedType>(op.getType());
+
+ if (const auto resultETy = resultType.getElementType();
+ !resultETy.isIntOrIndex())
+ return op.emitOpError("result tensor is not of integer type");
+
+ const auto inputType = llvm::cast<ShapedType>(op.getInput().getType());
+ if (!inputType.hasRank())
+ return success();
+
+ // Ensure axis is within the tensor rank
+ const int64_t axis = op.getAxisAttr().getInt();
+ if (((axis < 0) || axis >= inputType.getRank()))
+ return op.emitOpError("specified axis is outside the rank of the tensor");
+
+ if (!resultType.hasRank())
+ return success();
+
+ const ArrayRef<int64_t> inputShape = inputType.getShape();
+ const ArrayRef<int64_t> outputShape = resultType.getShape();
+ llvm::SmallVector<int64_t> expectedOutputShape(inputShape);
+ expectedOutputShape.erase(expectedOutputShape.begin() + axis);
+ if (failed(verifyCompatibleShape(expectedOutputShape, outputShape)))
+ return op.emitOpError("expected output shape '")
+ << expectedOutputShape << "', got '" << outputShape << "'";
+
+ return success();
+}
template <typename T>
static LogicalResult verifyConvOp(T op) {
@@ -973,36 +1035,9 @@ static LogicalResult verifySameElementTypes(Operation *op, Type aType,
return success();
}
-LogicalResult tosa::ArgMaxOp::verify() {
- const ShapedType resultType = llvm::cast<ShapedType>(getType());
-
- // Ensure output is of 32-bit integer
- if (const auto resultETy = resultType.getElementType();
- !resultETy.isIntOrIndex())
- return emitOpError("result tensor is not of integer type");
-
- const auto inputType = llvm::cast<ShapedType>(getInput().getType());
- if (!inputType.hasRank())
- return success();
+LogicalResult tosa::ArgMaxOp::verify() { return argMaxMinVerify(*this); }
- // Ensure axis is within the tensor rank
- const int64_t axis = getAxisAttr().getInt();
- if (((axis < 0) || axis >= inputType.getRank()))
- return emitOpError("specified axis is outside the rank of the tensor");
-
- if (!resultType.hasRank())
- return success();
-
- const ArrayRef<int64_t> inputShape = inputType.getShape();
- const ArrayRef<int64_t> outputShape = resultType.getShape();
- llvm::SmallVector<int64_t> expectedOutputShape(inputShape);
- expectedOutputShape.erase(expectedOutputShape.begin() + axis);
- if (failed(verifyCompatibleShape(expectedOutputShape, outputShape)))
- return emitOpError("expected output shape '")
- << expectedOutputShape << "', got '" << outputShape << "'";
-
- return success();
-}
+LogicalResult tosa::ArgMinOp::verify() { return argMaxMinVerify(*this); }
static LogicalResult verifyPoolingOpImpl(Operation *op,
ArrayRef<int64_t> kernel,
@@ -1581,25 +1616,16 @@ LogicalResult tosa::ArgMaxOp::inferReturnTypeComponents(
MLIRContext *context, ::std::optional<Location> location,
ArgMaxOp::Adaptor adaptor,
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
- ShapeAdaptor inputShape(adaptor.getInput().getType());
- IntegerAttr axis = adaptor.getProperties().axis;
- int32_t axisVal = axis.getValue().getSExtValue();
-
- if (!inputShape.hasRank()) {
- inferredReturnShapes.push_back(ShapedTypeComponents());
- return success();
- }
-
- SmallVector<int64_t> outShape;
- outShape.reserve(inputShape.getRank() - 1);
- for (int i = 0, s = inputShape.getRank(); i < s; i++) {
- if (i == axisVal)
- continue;
- outShape.push_back(inputShape.getDimSize(i));
- }
+ return inferArgMaxMinReturnTypeComponents(context, location, adaptor,
+ inferredReturnShapes);
+}
- inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
- return success();
+LogicalResult tosa::ArgMinOp::inferReturnTypeComponents(
+ MLIRContext *context, ::std::optional<Location> location,
+ ArgMinOp::Adaptor adaptor,
+ SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+ return inferArgMaxMinReturnTypeComponents(context, location, adaptor,
+ inferredReturnShapes);
}
LogicalResult tosa::RFFT2dOp::inferReturnTypeComponents(
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp
index 70fb94e386b09..7466ae5457576 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp
@@ -572,13 +572,14 @@ class ConvertCastOpWithBoundsChecking
}
};
-// ArgMax indices must fit the axis dimension, so we guard the integer rewrite.
-class ConvertArgMaxOpWithBoundsChecking
- : public OpConversionPattern<tosa::ArgMaxOp> {
- using OpConversionPattern::OpConversionPattern;
+// ArgMax/Min indices must fit the axis dimension, so we guard the integer
+// rewrite.
+template <typename OpTy>
+class ConvertArgMaxMinOpWithBoundsChecking : public OpConversionPattern<OpTy> {
+ using OpConversionPattern<OpTy>::OpConversionPattern;
LogicalResult
- matchAndRewrite(tosa::ArgMaxOp op, typename tosa::ArgMaxOp::Adaptor adaptor,
+ matchAndRewrite(OpTy op, typename OpTy::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const final {
const int32_t axis = op.getAxis();
const auto inputType = dyn_cast<ShapedType>(adaptor.getInput().getType());
@@ -593,8 +594,8 @@ class ConvertArgMaxOpWithBoundsChecking
const Type resultType = op.getOutput().getType();
const Type newResultType =
this->getTypeConverter()->convertType(resultType);
- rewriter.replaceOpWithNewOp<tosa::ArgMaxOp>(op, newResultType,
- adaptor.getInput(), axis);
+ rewriter.replaceOpWithNewOp<OpTy>(op, newResultType, adaptor.getInput(),
+ axis);
return success();
}
};
@@ -798,7 +799,10 @@ LogicalResult runTosaNarrowing(Operation *op, bool aggressiveRewrite,
typeConverter, context, allowLossyConversion, convertAccumulatorType);
} else {
if constexpr (Kind == TosaNarrowKind::Int64ToInt32) {
- patterns.add<ConvertArgMaxOpWithBoundsChecking>(typeConverter, context);
+ patterns.add<ConvertArgMaxMinOpWithBoundsChecking<tosa::ArgMaxOp>>(
+ typeConverter, context);
+ patterns.add<ConvertArgMaxMinOpWithBoundsChecking<tosa::ArgMinOp>>(
+ typeConverter, context);
patterns.add<ConvertClampOpWithBoundsChecking<Kind>>(typeConverter,
context);
}
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
index abd1aed65d910..a9dea6c86b28e 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
@@ -383,6 +383,7 @@ LogicalResult ProfileInfoDepot::populatationDispatch(Operation *op) {
POPULATE_PROFILE_INFO_COMMON(CastToBlockScaled)
POPULATE_PROFILE_INFO_COMMON(Const)
POPULATE_PROFILE_INFO_COMMON(ArgMax)
+ POPULATE_PROFILE_INFO_COMMON(ArgMin)
POPULATE_PROFILE_INFO_COMMON(Sub)
POPULATE_PROFILE_INFO_COMMON(Maximum)
POPULATE_PROFILE_INFO_COMMON(Minimum)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
index 9a37e57af0d2a..53c1fa80448b1 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -715,6 +715,21 @@ LogicalResult TosaValidation::levelCheckRanks(tosa::ArgMaxOp tosaOp) {
return success();
}
+template <>
+LogicalResult TosaValidation::levelCheckRanks(tosa::ArgMinOp tosaOp) {
+ auto *op = tosaOp.getOperation();
+ if (failed(levelCheckRank(op, tosaOp.getInput(), "operand",
+ targetEnv.getLevel().MAX_RANK)))
+ return failure();
+
+ // rank(output) = rank(input) - 1
+ if (failed(levelCheckRank(op, tosaOp.getOutput(), "result",
+ targetEnv.getLevel().MAX_RANK - 1)))
+ return failure();
+
+ return success();
+}
+
template <>
LogicalResult TosaValidation::levelCheckRanks(tosa::IfOp tosaOp) {
auto *op = tosaOp.getOperation();
@@ -771,6 +786,7 @@ LogicalResult TosaValidation::levelCheckRanksAndSizes(Operation *op) {
// Tensor Operators
CHECK_RANKS_AND_SIZES(ArgMax);
+ CHECK_RANKS_AND_SIZES(ArgMin);
// Activation Functions
CHECK_RANKS_AND_SIZES(Clamp);
CHECK_RANKS_AND_SIZES(Erf);
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 8776182bf16c5..f39a5c4379eb4 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -16,6 +16,23 @@ func.func @test_argmax_fold_i64_index(%arg0: tensor<1xi8>) -> tensor<i64> {
return %0 : tensor<i64>
}
+
+// CHECK-LABEL: @argmin_nofold
+func.func @argmin_nofold(%arg0: tensor<?x1xf32>) -> tensor<1xi32> {
+ // CHECK: tosa.argmin
+ %0 = tosa.argmin %arg0 axis(0): (tensor<?x1xf32>) -> tensor<1xi32>
+ return %0 : tensor<1xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_argmin_fold_i64_index
+func.func @test_argmin_fold_i64_index(%arg0: tensor<1xi8>) -> tensor<i64> {
+ // CHECK: tosa.const values(dense<0> : tensor<i64>) : () -> tensor<i64>
+ %0 = tosa.argmin %arg0 axis(0) : (tensor<1xi8>) -> tensor<i64>
+ return %0 : tensor<i64>
+}
+
// -----
// CHECK-LABEL: @row_gather_row_count_one_to_gather
diff --git a/mlir/test/Dialect/Tosa/constant_folding.mlir b/mlir/test/Dialect/Tosa/constant_folding.mlir
index 9d5d98df74ab0..3aefb1c40ff8e 100644
--- a/mlir/test/Dialect/Tosa/constant_folding.mlir
+++ b/mlir/test/Dialect/Tosa/constant_folding.mlir
@@ -1,5 +1,41 @@
// RUN: mlir-opt --split-input-file --test-single-fold %s | FileCheck %s
+// CHECK-LABEL: @argmax_fold_dim_size_1
+func.func @argmax_fold_dim_size_1(%arg0: tensor<2x1x3xf32>) -> tensor<2x3xi32> {
+ // CHECK: tosa.const values(dense<0> : tensor<2x3xi32>) : () -> tensor<2x3xi32>
+ %0 = tosa.argmax %arg0 axis(1): (tensor<2x1x3xf32>) -> tensor<2x3xi32>
+ return %0 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @argmax_dynamic_shape_no_fold_dim_size_1
+func.func @argmax_dynamic_shape_no_fold_dim_size_1(%arg0: tensor<?x1x3xf32>) -> tensor<?x3xi32> {
+ // CHECK: tosa.argmax
+ %0 = tosa.argmax %arg0 axis(1): (tensor<?x1x3xf32>) -> tensor<?x3xi32>
+ return %0 : tensor<?x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @argmin_fold_dim_size_1
+func.func @argmin_fold_dim_size_1(%arg0: tensor<2x1x3xf32>) -> tensor<2x3xi32> {
+ // CHECK: tosa.const values(dense<0> : tensor<2x3xi32>) : () -> tensor<2x3xi32>
+ %0 = tosa.argmin %arg0 axis(1): (tensor<2x1x3xf32>) -> tensor<2x3xi32>
+ return %0 : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @argmin_dynamic_shape_no_fold_dim_size_1
+func.func @argmin_dynamic_shape_no_fold_dim_size_1(%arg0: tensor<?x1x3xf32>) -> tensor<?x3xi32> {
+ // CHECK: tosa.argmin
+ %0 = tosa.argmin %arg0 axis(1): (tensor<?x1x3xf32>) -> tensor<?x3xi32>
+ return %0 : tensor<?x3xi32>
+}
+
+// -----
+
// CHECK-LABEL: func @test_const
func.func @test_const(%arg0 : index) -> tensor<4xi32> {
// CHECK: tosa.const
diff --git a/mlir/test/Dialect/Tosa/invalid.mlir b/mlir/test/Dialect/Tosa/invalid.mlir
index dd61969b4d1df..6d28ed3870bd3 100644
--- a/mlir/test/Dialect/Tosa/invalid.mlir
+++ b/mlir/test/Dialect/Tosa/invalid.mlir
@@ -1412,6 +1412,14 @@ func.func @test_argmax_invalid_output_shape(%arg0: tensor<1x2x3xf32>) -> tensor<
// -----
+func.func @test_argmin_invalid_output_shape(%arg0: tensor<1x2x3xf32>) -> tensor<1x2x3xi32> {
+ // expected-error at +1 {{'tosa.argmin' op expected output shape '2, 3', got '1, 2, 3'}}
+ %0 = tosa.argmin %arg0 axis(0): (tensor<1x2x3xf32>) -> tensor<1x2x3xi32>
+ return %0 : tensor<1x2x3xi32>
+}
+
+// -----
+
func.func @test_rescale_invalid_input_type(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xi32> {
%multiplier = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32> }> : () -> tensor<1xi32>
%shift = "tosa.const"() <{values = dense<30> : tensor<1xi8> }> : () -> tensor<1xi8>
@@ -2005,6 +2013,14 @@ func.func @test_scalar_argmax(%arg0: tensor<i32>) -> tensor<i32> {
// -----
+func.func @test_scalar_argmin(%arg0: tensor<i32>) -> tensor<i32> {
+ // expected-error at +1 {{'tosa.argmin' op operand #0 must be tosa-conformant tensor of at least rank 1 of number values, but got 'tensor<i32>'}}
+ %0 = tosa.argmin %arg0 axis(0) : (tensor<i32>) -> tensor<i32>
+ return %0 : tensor<i32>
+}
+
+// -----
+
func.func @test_scalar_reduce_all(%arg0: tensor<i1>) -> tensor<i1> {
// expected-error at +1 {{'tosa.reduce_all' op operand #0 must be tosa-conformant tensor of at least rank 1 of number values, but got 'tensor<i1>'}}
%0 = tosa.reduce_all %arg0 axis(0) : (tensor<i1>) -> tensor<i1>
diff --git a/mlir/test/Dialect/Tosa/invalid_extension.mlir b/mlir/test/Dialect/Tosa/invalid_extension.mlir
index f4e70710dc320..01b8a3175ccf4 100644
--- a/mlir/test/Dialect/Tosa/invalid_extension.mlir
+++ b/mlir/test/Dialect/Tosa/invalid_extension.mlir
@@ -11,6 +11,13 @@ func.func @test_argmax(%arg0: tensor<14x19xbf16>) -> tensor<14xi32> {
...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/227997
More information about the Mlir-commits
mailing list