[Mlir-commits] [mlir] 557c3bf - [mlir][tosa] Add support for ARGMIN (#227997)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 2 02:41:47 PDT 2026
Author: Iliyan Georgiev
Date: 2026-10-02T10:41:39+01:00
New Revision: 557c3bf999fc86f15cc91d4f9ccfa34c3b784fdc
URL: https://github.com/llvm/llvm-project/commit/557c3bf999fc86f15cc91d4f9ccfa34c3b784fdc
DIFF: https://github.com/llvm/llvm-project/commit/557c3bf999fc86f15cc91d4f9ccfa34c3b784fdc.diff
LOG: [mlir][tosa] Add support for ARGMIN (#227997)
- Add ARGMIN table gen
- Add validators and folders
- Add positive and negative tests
- Add compliance info
Signed-off-by: Iliyan Georgiev <iliyan.georgiev at arm.com>
Added:
Modified:
mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp
mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
mlir/test/Dialect/Tosa/canonicalize.mlir
mlir/test/Dialect/Tosa/constant_folding.mlir
mlir/test/Dialect/Tosa/invalid.mlir
mlir/test/Dialect/Tosa/invalid_extension.mlir
mlir/test/Dialect/Tosa/level_check.mlir
mlir/test/Dialect/Tosa/ops.mlir
mlir/test/Dialect/Tosa/profile_pro_int_unsupported.mlir
mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir
mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index 5dcfe3a737e2d..e3789986735fa 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 5f34856859bfa..8b5803c3074f0 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 4cdd0b097e7b6..dd6ce622c4646 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 cb82e87a5cfd0..5d465065497ab 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -743,6 +743,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();
@@ -799,6 +814,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 16462130bb23b..65fd585954255 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> {
return %0 : tensor<14xi32>
}
+// -----
+func.func @test_argmin(%arg0: tensor<14x19xbf16>) -> tensor<14xi32> {
+ // expected-error at +1 {{'tosa.argmin' op illegal: requires any of [bf16] profiles/extensions to be specified in the target environment}}
+ %0 = tosa.argmin %arg0 axis(1) : (tensor<14x19xbf16>) -> tensor<14xi32>
+ return %0 : tensor<14xi32>
+}
+
// -----
func.func @test_avg_pool2d(%arg0: tensor<1x7x7x9xbf16>, %arg1: tensor<1xbf16>, %arg2: tensor<1xbf16>) -> tensor<1x7x7x9xbf16> {
// expected-error at +1 {{'tosa.avg_pool2d' op illegal: requires any of [bf16] profiles/extensions to be specified in the target environment}}
diff --git a/mlir/test/Dialect/Tosa/level_check.mlir b/mlir/test/Dialect/Tosa/level_check.mlir
index 5c2514057ea2d..2d1dfbf5b45b8 100644
--- a/mlir/test/Dialect/Tosa/level_check.mlir
+++ b/mlir/test/Dialect/Tosa/level_check.mlir
@@ -8,6 +8,14 @@ func.func @test_argmax_rank_invalid(%arg0: tensor<1x1x1x1x29x29x4xf32>) -> tenso
// -----
+func.func @test_argmin_rank_invalid(%arg0: tensor<1x1x1x1x29x29x4xf32>) -> tensor<1x1x1x1x29x4xi32> {
+ // expected-error at +1 {{'tosa.argmin' op failed level check: operand rank(shape) <= MAX_RANK}}
+ %0 = "tosa.argmin"(%arg0) <{axis = 4 : i32}> : (tensor<1x1x1x1x29x29x4xf32>) -> tensor<1x1x1x1x29x4xi32>
+ return %0 : tensor<1x1x1x1x29x4xi32>
+}
+
+// -----
+
func.func @test_clamp_rank_invalid(%arg0: tensor<1x1x1x1x13x21x3xf32>) -> tensor<1x1x1x1x13x21x3xf32> {
// expected-error at +1 {{'tosa.clamp' op failed level check: operand rank(shape) <= MAX_RANK}}
%0 = tosa.clamp %arg0 min_val(-3.40282347E+38 : f32) max_val(3.40282347E+38 : f32) : (tensor<1x1x1x1x13x21x3xf32>) -> tensor<1x1x1x1x13x21x3xf32>
@@ -1020,7 +1028,7 @@ func.func @test_maxpool2d_adaptive_pad_first(%arg0: tensor<1x32x8194x8xf32>) ->
%a = tosa.const_shape values(dense<[1, 1]> : tensor<2xindex>) : () -> !tosa.shape<2>
%b = tosa.const_shape values(dense<[1, 1]> : tensor<2xindex>) : () -> !tosa.shape<2>
%kernel = tosa.add_shape %a, %b : (!tosa.shape<2>, !tosa.shape<2>) -> !tosa.shape<2>
-
+
%stride = tosa.const_shape values(dense<[1, 1]> : tensor<2xindex>) : () -> !tosa.shape<2>
%pad = tosa.const_shape values(dense<[8193, 0, 0, 0]> : tensor<4xindex>) : () -> !tosa.shape<4>
// expected-error at +1 {{'tosa.max_pool2d_adaptive' op failed level check: pad <= MAX_KERNEL (8192), got 8193}}
diff --git a/mlir/test/Dialect/Tosa/ops.mlir b/mlir/test/Dialect/Tosa/ops.mlir
index eac2f58772d8d..2592b08b35d39 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -9,6 +9,13 @@ func.func @test_argmax(%arg0: tensor<14x19xf32>) -> tensor<14xi32> {
return %0 : tensor<14xi32>
}
+// -----
+// CHECK-LABEL: argmin
+func.func @test_argmin(%arg0: tensor<14x19xf32>) -> tensor<14xi32> {
+ %0 = tosa.argmin %arg0 axis(1) : (tensor<14x19xf32>) -> tensor<14xi32>
+ return %0 : tensor<14xi32>
+}
+
// -----
// CHECK-LABEL: avg_pool2d_f32
func.func @test_avg_pool2d_f32(%arg0: tensor<1x7x7x9xf32>) -> tensor<1x7x7x9xf32> {
@@ -1412,6 +1419,13 @@ func.func @test_argmax_f8E5M2(%arg0: tensor<12x8x16xf8E5M2>) -> tensor<12x16xi32
return %0 : tensor<12x16xi32>
}
+// -----
+// CHECK-LABEL: argmin_f8E5M2
+func.func @test_argmin_f8E5M2(%arg0: tensor<12x8x16xf8E5M2>) -> tensor<12x16xi32> {
+ %0 = tosa.argmin %arg0 axis(1) : (tensor<12x8x16xf8E5M2>) -> tensor<12x16xi32>
+ return %0 : tensor<12x16xi32>
+}
+
// -----
// CHECK-LABEL: avg_pool2d_f8E5M2
func.func @test_avg_pool2d_f8E5M2(%arg0: tensor<1x7x7x9xf8E5M2>) -> tensor<1x7x7x9xf8E5M2> {
@@ -1618,6 +1632,13 @@ func.func @test_argmax_f8E4M3FN(%arg0: tensor<12x8x16xf8E4M3FN>) -> tensor<12x16
return %0 : tensor<12x16xi32>
}
+// -----
+// CHECK-LABEL: argmin_f8E4M3FN
+func.func @test_argmin_f8E4M3FN(%arg0: tensor<12x8x16xf8E4M3FN>) -> tensor<12x16xi32> {
+ %0 = tosa.argmin %arg0 axis(1) : (tensor<12x8x16xf8E4M3FN>) -> tensor<12x16xi32>
+ return %0 : tensor<12x16xi32>
+}
+
// -----
// CHECK-LABEL: avg_pool2d_f8E4M3FN
func.func @test_avg_pool2d_f8E4M3FN(%arg0: tensor<1x7x7x9xf8E4M3FN>) -> tensor<1x7x7x9xf8E4M3FN> {
@@ -2068,21 +2089,21 @@ func.func @test_assert_equal_shape() {
// CHECK-LABEL: test_block_scaled_const
func.func @test_block_scaled_const() -> tensor<2x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>> {
- %0 = "tosa.const"() <{values = dense<tensor<2x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN, {1.0, 1.0}>> : [[0.0 : f8E4M3FN, 1.0 : f8E4M3FN, 0.001953125 : f8E4M3FN, 0.0078125 : f8E4M3FN,
- 2.0 : f8E4M3FN, 2.25 : f8E4M3FN, 2.5 : f8E4M3FN, 2.75 : f8E4M3FN,
- 15.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ %0 = "tosa.const"() <{values = dense<tensor<2x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN, {1.0, 1.0}>> : [[0.0 : f8E4M3FN, 1.0 : f8E4M3FN, 0.001953125 : f8E4M3FN, 0.0078125 : f8E4M3FN,
+ 2.0 : f8E4M3FN, 2.25 : f8E4M3FN, 2.5 : f8E4M3FN, 2.75 : f8E4M3FN,
+ 15.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN],
+ [0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
+ 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN],
- [0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
- 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN,
0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN, 0.0 : f8E4M3FN]]>}> : () -> tensor<2x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
return %0 : tensor<2x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>
}
diff --git a/mlir/test/Dialect/Tosa/profile_pro_int_unsupported.mlir b/mlir/test/Dialect/Tosa/profile_pro_int_unsupported.mlir
index 4243b6f449e6a..ee6ecb81ec623 100644
--- a/mlir/test/Dialect/Tosa/profile_pro_int_unsupported.mlir
+++ b/mlir/test/Dialect/Tosa/profile_pro_int_unsupported.mlir
@@ -23,6 +23,13 @@ func.func @test_argmax(%arg0: tensor<14x19xi8>) -> tensor<14xi32> {
return %0 : tensor<14xi32>
}
+// -----
+func.func @test_argmin(%arg0: tensor<14x19xi8>) -> tensor<14xi32> {
+ // expected-error at +1 {{'tosa.argmin' op illegal: requires any of [pro_int] profiles/extensions to be specified in the target environment}}
+ %0 = tosa.argmin %arg0 axis(1) : (tensor<14x19xi8>) -> tensor<14xi32>
+ return %0 : tensor<14xi32>
+}
+
// -----
func.func @test_avg_pool2d_adaptive_missing_pro_int(%arg0: tensor<1x7x7x9xi8>, %arg1: tensor<1xi8>, %arg2: tensor<1xi8>) -> tensor<1x7x7x9xi8> {
%kernel = tosa.const_shape values(dense<[2, 2]> : tensor<2xindex>) : () -> !tosa.shape<2>
diff --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
index 1c64c97708a6d..4673f22460ea7 100644
--- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
@@ -306,6 +306,30 @@ func.func @test_dynamic_argmax(%arg0 : tensor<2x?xi32>) -> () {
// -----
+// CHECK-LABEL: @test_static_argmin
+func.func @test_static_argmin(%arg0 : tensor<2x3xi32>) -> () {
+ // CHECK: tosa.argmin %arg0 axis(0) : (tensor<2x3xi32>) -> tensor<3xi32>
+ %0 = tosa.argmin %arg0 axis(0) : (tensor<2x3xi32>) -> tensor<?xi32>
+
+ // CHECK: tosa.argmin %arg0 axis(1) : (tensor<2x3xi32>) -> tensor<2xi32>
+ %1 = tosa.argmin %arg0 axis(1) : (tensor<2x3xi32>) -> tensor<?xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_dynamic_argmin
+func.func @test_dynamic_argmin(%arg0 : tensor<2x?xi32>) -> () {
+ // CHECK: tosa.argmin %arg0 axis(0) : (tensor<2x?xi32>) -> tensor<?xi32>
+ %0 = tosa.argmin %arg0 axis(0) : (tensor<2x?xi32>) -> tensor<?xi32>
+
+ // CHECK: tosa.argmin %arg0 axis(1) : (tensor<2x?xi32>) -> tensor<2xi32>
+ %1 = tosa.argmin %arg0 axis(1) : (tensor<2x?xi32>) -> tensor<?xi32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @test_static_matmul
func.func @test_static_matmul(%arg0 : tensor<2x3x4xi32>, %arg1 : tensor<2x4x5xi32>) -> () {
// CHECK tosa.matmul %arg0, %arg1, %0, %1 : (tensor<2x3x4xi32>, tensor<2x4x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2x3x5xi32>
diff --git a/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir b/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
index 1bef68e895b93..aeb509d2f21a4 100644
--- a/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
@@ -1,23 +1,6 @@
// RUN: mlir-opt --split-input-file --tosa-layerwise-constant-fold %s | FileCheck %s
// RUN: mlir-opt --split-input-file --tosa-layerwise-constant-fold="aggressive-reduce-constant=true" %s | FileCheck %s --check-prefix=AGGRESIVE
-// CHECK-LABEL: @armax_fold_dim_size_1
-func.func @armax_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: @transpose_fold
func.func @transpose_fold(%arg0: tensor<3x4xf32>) -> tensor<3x4xf32> {
diff --git a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir
index b7a7f0543d911..cc4ec84ba3a34 100644
--- a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32-aggressive.mlir
@@ -10,6 +10,15 @@ func.func @test_i64_argmax_large_axis_dim(%arg0: tensor<1x513x513x2147483650xi8>
// -----
+// CHECK-LABEL: test_i64_argmin_large_axis_dim
+func.func @test_i64_argmin_large_axis_dim(%arg0: tensor<1x513x513x2147483650xi8>) -> tensor<1x513x513xi64> {
+ // DEFAULT: tosa.argmin %arg0 axis(3) : (tensor<1x513x513x2147483650xi8>) -> tensor<1x513x513xi32>
+ %0 = tosa.argmin %arg0 axis(3) : (tensor<1x513x513x2147483650xi8>) -> tensor<1x513x513xi64>
+ return %0 : tensor<1x513x513xi64>
+}
+
+// -----
+
// CHECK-LABEL: test_convert_input_parameters
// DEFAULT: %[[IN:.*]]: tensor<1x513x513x3xi64>
// FUNCBOUND: %[[IN:.*]]: tensor<1x513x513x3xi32>
diff --git a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
index 2c787eec4f25c..128b75d0a8504 100644
--- a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
@@ -15,6 +15,18 @@ func.func @test_i64_argmax(%arg0: tensor<1x513x513x19xi8>) -> tensor<1x513x513xi
// -----
+// CHECK-LABEL: test_i64_argmin
+func.func @test_i64_argmin(%arg0: tensor<1x513x513x19xi8>) -> tensor<1x513x513xi64> {
+ // COMMON: %[[ARGMIN:.*]] = tosa.argmin %arg0 axis(3) : (tensor<1x513x513x19xi8>) -> tensor<1x513x513xi32>
+ %0 = tosa.argmin %arg0 axis(3) : (tensor<1x513x513x19xi8>) -> tensor<1x513x513xi64>
+
+ // DEFAULT: %[[CAST:.*]] = tosa.cast %[[ARGMIN]] input_unsigned(false) : (tensor<1x513x513xi32>) -> tensor<1x513x513xi64>
+ // FUNCBOUND: return %[[ARGMIN]] : tensor<1x513x513xi32>
+ return %0 : tensor<1x513x513xi64>
+}
+
+// -----
+
// CHECK-LABEL: test_i64_const
func.func @test_i64_const() -> tensor<2xi64> {
// COMMON: %[[CONST:.*]] = tosa.const values(dense<[1, 2]> : tensor<2xi32>) : () -> tensor<2xi32>
@@ -38,6 +50,17 @@ func.func @test_i64_argmax_cast(%arg0: tensor<1x513x513x19xi8>) -> tensor<1x513x
// -----
+// CHECK-LABEL: test_i64_argmin_cast
+func.func @test_i64_argmin_cast(%arg0: tensor<1x513x513x19xi8>) -> tensor<1x513x513xf32> {
+ // COMMON: %[[ARGMIN:.*]] = tosa.argmin %arg0 axis(3) : (tensor<1x513x513x19xi8>) -> tensor<1x513x513xi32>
+ %0 = tosa.argmin %arg0 axis(3) : (tensor<1x513x513x19xi8>) -> tensor<1x513x513xi64>
+ // COMMON: tosa.cast %[[ARGMIN]] {test.marker = "keep"} : (tensor<1x513x513xi32>) -> tensor<1x513x513xf32>
+ %1 = tosa.cast %0 {test.marker = "keep"} : (tensor<1x513x513xi64>) -> tensor<1x513x513xf32>
+ return %1 : tensor<1x513x513xf32>
+}
+
+// -----
+
// CHECK-LABEL: test_i64_argmax_large_axis_dim
func.func @test_i64_argmax_large_axis_dim(%arg0: tensor<1x513x513x2147483650xi8>) -> tensor<1x513x513xi64> {
// expected-error @+1 {{failed to legalize operation 'tosa.argmax'}}
@@ -47,6 +70,15 @@ func.func @test_i64_argmax_large_axis_dim(%arg0: tensor<1x513x513x2147483650xi8>
// -----
+// CHECK-LABEL: test_i64_argmin_large_axis_dim
+func.func @test_i64_argmin_large_axis_dim(%arg0: tensor<1x513x513x2147483650xi8>) -> tensor<1x513x513xi64> {
+ // expected-error @+1 {{failed to legalize operation 'tosa.argmin'}}
+ %0 = tosa.argmin %arg0 axis(3) : (tensor<1x513x513x2147483650xi8>) -> tensor<1x513x513xi64>
+ return %0 : tensor<1x513x513xi64>
+}
+
+// -----
+
// CHECK-LABEL: test_add
func.func @test_add(%arg0: tensor<13x21x1xi64>, %arg1: tensor<13x21x3xi64>) -> tensor<13x21x3xi64> {
// expected-error @+1 {{failed to legalize operation 'tosa.add'}}
@@ -75,6 +107,25 @@ func.func @test_regions(%arg0: tensor<1x2xi32>, %arg1: tensor<1xi32>, %arg2: ten
// -----
+// CHECK-LABEL: test_regions
+func.func @test_regions(%arg0: tensor<1x2xi32>, %arg1: tensor<1xi32>, %arg2: tensor<i1>) -> tensor<1xi32> {
+ // COMMON: %[[IF_RESULT:.*]] = tosa.cond_if %arg2 : tensor<i1> -> tensor<1xi32>
+ %0 = tosa.cond_if %arg2 : tensor<i1> -> tensor<1xi32> {
+ // COMMON: %[[ARGMIN:.*]] = tosa.argmin %arg0 axis(1) : (tensor<1x2xi32>) -> tensor<1xi32>
+ %1 = tosa.argmin %arg0 axis(1) : (tensor<1x2xi32>) -> tensor<1xi64>
+ // COMMON: %[[CAST:.*]] = tosa.cast %[[ARGMIN]] : (tensor<1xi32>) -> tensor<1xi32>
+ %2 = tosa.cast %1 : (tensor<1xi64>) -> tensor<1xi32>
+ // COMMON: tosa.yield %[[CAST]] : tensor<1xi32>
+ tosa.yield %2 : tensor<1xi32>
+ } else {
+ tosa.yield %arg1 : tensor<1xi32>
+ }
+ // COMMON: return %[[IF_RESULT]] : tensor<1xi32>
+ return %0 : tensor<1xi32>
+}
+
+// -----
+
// CHECK-LABEL: test_cond_if_i64_yield_with_call
module {
func.func @m0(%arg0: tensor<4xi64>) -> tensor<4xi64> {
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
index 3f020494d7c22..32719d7a0ab78 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
@@ -256,6 +256,14 @@ func.func @test_dyanmic_dims(%arg0: tensor<1x8x16xi8>) {
// -----
+func.func @test_dyanmic_dims(%arg0: tensor<1x8x16xi8>) {
+ // expected-error at +1 {{'tosa.argmin' op illegal: requires specification version compatible with 1.1.draft (got 1.0)}}
+ %1 = tosa.argmin %arg0 axis(1) : (tensor<1x8x16xi8>) -> tensor<1x16xi32>
+ return
+}
+
+// -----
+
func.func @test_block_scaled_identity(%arg0: tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>) -> tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>> {
// expected-error at +1 {{'tosa.identity' op illegal: requires specification version compatible with 1.1.draft (got 1.0) and requires all of [mx_common, mx_fp6e2m3] profiles/extensions to be specified in the target environment}}
%0 = tosa.identity %arg0 : (tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>) -> tensor<13x21x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
index c84496018dfc6..2927b9a0165d0 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -194,6 +194,14 @@ func.func @test_argmax_int64(%arg0: tensor<1x13x13x5xf32>) -> tensor<1x13x13xi64
// -----
+// CHECK-LABEL: test_argmin_int64
+func.func @test_argmin_int64(%arg0: tensor<1x13x13x5xf32>) -> tensor<1x13x13xi64> {
+ %0 = tosa.argmin %arg0 axis(3) : (tensor<1x13x13x5xf32>) -> tensor<1x13x13xi64>
+ return %0 : tensor<1x13x13xi64>
+}
+
+// -----
+
// CHECK-LABEL: test_const_i64
func.func @test_const_i64() -> tensor<4xi64> {
%0 = "tosa.const"() <{values = dense<[3, 0, 1, 2]> : tensor<4xi64>}> : () -> tensor<4xi64>
@@ -303,6 +311,22 @@ func.func @test_argmax_bf16_i64(%arg0: tensor<12x8x16xbf16>) -> tensor<12x16xi64
// -----
+// CHECK-LABEL: test_argmin_fp8_i64
+func.func @test_argmin_fp8_i64(%arg0: tensor<12x8x16xf8E5M2>) -> tensor<12x16xi64> {
+ %0 = tosa.argmin %arg0 axis(1) : (tensor<12x8x16xf8E5M2>) -> tensor<12x16xi64>
+ return %0 : tensor<12x16xi64>
+}
+
+// -----
+
+// CHECK-LABEL: test_argmin_bf16_i64
+func.func @test_argmin_bf16_i64(%arg0: tensor<12x8x16xbf16>) -> tensor<12x16xi64> {
+ %0 = tosa.argmin %arg0 axis(1) : (tensor<12x8x16xbf16>) -> tensor<12x16xi64>
+ return %0 : tensor<12x16xi64>
+}
+
+// -----
+
// CHECK-LABEL: test_scatter_const_indices_int64
func.func @test_scatter_const_indices_int64(%arg0: tensor<2x52x3xf32>, %arg2: tensor<2x12x3xf32>) -> tensor<2x52x3xf32> {
%indices = "tosa.const"() <{ values = dense<[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]]> : tensor<2x12xi64> }> : () -> tensor<2x12xi64>
More information about the Mlir-commits
mailing list