[Mlir-commits] [mlir] [mlir][tosa] Add row_gather operator (PR #202895)
Luke Hutton
llvmlistbot at llvm.org
Mon Jun 15 07:12:23 PDT 2026
https://github.com/lhutton1 updated https://github.com/llvm/llvm-project/pull/202895
>From 44903457bdd0b34fee3a242c835f1800ed864b2b Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Mon, 8 Jun 2026 17:59:08 +0100
Subject: [PATCH 1/3] [mlir][tosa] Add row_gather operator
Adds support for the row_gather operator defined by the TOSA
specification, see https://github.com/arm/tosa-specification/pull/60.
This includes:
- Operator definition
- Verification logic for the operator
- Output shape inference for the operator
- Validation checks to ensure compliance with the TOSA specification
including profile compliance and level checks.
- Canonicalization to replace row_gather with gather when row_count is
statically known to be 1.
It does not yet cover support for MXFP types. This will be added once
block scaled types are supported.
Change-Id: I9037f813a6c31b7dc03cfe3b5a191f5e3fce02d6
---
.../Dialect/Tosa/IR/TosaComplianceData.h.inc | 65 +++++++++++
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td | 36 ++++++
.../mlir/Dialect/Tosa/Utils/ConversionUtils.h | 5 +
.../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 21 ++++
mlir/lib/Dialect/Tosa/IR/TosaOps.cpp | 106 +++++++++++++++---
.../Tosa/Transforms/TosaProfileCompliance.cpp | 9 ++
.../Tosa/Transforms/TosaValidation.cpp | 11 ++
.../Dialect/Tosa/Utils/ConversionUtils.cpp | 16 +++
mlir/test/Dialect/Tosa/availability.mlir | 10 ++
mlir/test/Dialect/Tosa/canonicalize.mlir | 30 +++++
mlir/test/Dialect/Tosa/invalid_extension.mlir | 8 ++
mlir/test/Dialect/Tosa/level_check.mlir | 9 ++
mlir/test/Dialect/Tosa/ops.mlir | 47 ++++++++
mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir | 39 +++++++
.../tosa-validation-version-1p0-invalid.mlir | 9 ++
...a-validation-version-1p1-pro-fp-valid.mlir | 9 ++
.../tosa-validation-version-1p1-valid.mlir | 27 +++++
mlir/test/Dialect/Tosa/verifier.mlir | 45 ++++++++
18 files changed, 486 insertions(+), 16 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index 50bb9f69c6242..f17475fae563d 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -396,6 +396,33 @@ profileComplianceMap = {
{{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT},
{{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT}},
anyOf}}},
+ {"tosa.row_gather",
+ {{{Profile::pro_int},
+ {{{boolT, i32T, boolT}, SpecificationVersion::V_1_1_DRAFT},
+ {{i8T, i32T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i32T, i32T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{boolT, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT},
+ {{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Profile::pro_fp},
+ {{{boolT, i32T, boolT}, SpecificationVersion::V_1_1_DRAFT},
+ {{i8T, i32T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i32T, i32T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, i32T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, i32T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{boolT, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT},
+ {{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, i64T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, i64T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}}}},
{"tosa.row_gather_block_scaled",
{{{Profile::pro_int},
{{{i8T, i32T, i32T, i8T}, SpecificationVersion::V_1_1_DRAFT},
@@ -909,6 +936,44 @@ extensionComplianceMap = {
{{Extension::bf16, Extension::int64},
{{{bf16T, i64T, bf16T}, SpecificationVersion::V_1_1_DRAFT}},
allOf}}},
+ {"tosa.row_gather",
+ {
+ {{Extension::int16},
+ {{{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::int64},
+ {{{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{boolT, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT},
+ {{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{boolT, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT},
+ {{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, i64T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, i64T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::int16, Extension::int64},
+ {{{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::fp8e4m3},
+ {{{fp8e4m3T, i32T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::fp8e5m2},
+ {{{fp8e5m2T, i32T, fp8e5m2T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::bf16},
+ {{{bf16T, i32T, bf16T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::fp8e4m3, Extension::int64},
+ {{{fp8e4m3T, i64T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::fp8e5m2, Extension::int64},
+ {{{fp8e5m2T, i64T, fp8e5m2T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::bf16, Extension::int64},
+ {{{bf16T, i64T, bf16T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ }},
{"tosa.row_gather_block_scaled",
{{{Extension::fp8e4m3},
{{{fp8e4m3T, i32T, i32T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}}},
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index e135265b99881..75fe9c4fc54b2 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -2532,6 +2532,42 @@ def Tosa_GatherOp : Tosa_InferShapedTypeOp<"gather", [NoMemoryEffect]> {
"operands attr-dict `:` functional-type(operands, results)";
}
+//===----------------------------------------------------------------------===//
+// Operator: row_gather
+//===----------------------------------------------------------------------===//
+def Tosa_RowGatherOp : Tosa_InferShapedTypeOp<"row_gather", [NoMemoryEffect]> {
+ let summary = "Row gather operation.";
+
+ let description = [{
+ Generate a tensor based on the input indices and row_count. The number of
+ consecutive rows gathered for each index is specified in row_count. N is
+ the number of batches, W is the number of indices in each batch, K is
+ the range of each index, and C is the number of data channels for each
+ index. The values tensor has shape [N, K, C] and the output tensor has shape
+ [N, W * row_count, C].
+
+ This operation is not pure. Undefined behaviour may occur if the specified
+ indices are out of range.
+ }];
+
+ let arguments = (ins Tosa_Tensor3D:$values, Tosa_IndexTensor2D:$indices,
+ Tosa_ScalarInt32Tensor:$row_count);
+
+ let results = (outs Tosa_Tensor3D:$output);
+
+ list<Availability> availability = [Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
+ Extension<[Tosa_EXT_FP8E4M3,
+ Tosa_EXT_FP8E5M2, Tosa_EXT_BF16,
+ Tosa_EXT_INT64]>,
+ ];
+
+ let hasCanonicalizer = 1;
+ let hasVerifier = 1;
+
+ let assemblyFormat =
+ "operands attr-dict `:` functional-type(operands, results)";
+}
+
//===----------------------------------------------------------------------===//
// Operator: row_gather_block_scaled
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h b/mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h
index 3a48608e7fd94..df4dabc2afda8 100644
--- a/mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h
+++ b/mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h
@@ -271,6 +271,11 @@ std::optional<ArrayRef<T>> tryGetDenseResourceValues(ElementsAttr attr) {
return std::nullopt;
}
+// returns the value of a constant scalar int tensor, or failure if
+// the value cannot be extracted
+template <typename T>
+FailureOr<T> getConstantScalarIntValue(Value val);
+
} // namespace tosa
} // namespace mlir
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 4af185a6e534b..4d61441ff037b 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1185,6 +1185,27 @@ void CastToBlockScaledOp::getCanonicalizationPatterns(
results.add<CancellingBlockScaledCastsOptimization>(context);
}
+struct RowGatherToGather : public OpRewritePattern<tosa::RowGatherOp> {
+ using OpRewritePattern<tosa::RowGatherOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::RowGatherOp op,
+ PatternRewriter &rewriter) const override {
+ const FailureOr<int32_t> rowCount =
+ mlir::tosa::getConstantScalarIntValue<int32_t>(op.getRowCount());
+ if (failed(rowCount) || rowCount.value() != 1)
+ return failure();
+
+ rewriter.replaceOpWithNewOp<tosa::GatherOp>(
+ op, op.getOutput().getType(), op.getValues(), op.getIndices());
+ return success();
+ }
+};
+
+void RowGatherOp::getCanonicalizationPatterns(RewritePatternSet &results,
+ MLIRContext *context) {
+ results.add<RowGatherToGather>(context);
+}
+
//===----------------------------------------------------------------------===//
// Operator Folders.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 351dc1ae0dcd5..f0a44a2a22ed2 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -3104,18 +3104,6 @@ static FailureOr<int64_t> getZeroPoint(Value val, bool signExtend) {
return -1;
}
-static FailureOr<int64_t> getConstantScalarIntValue(Value val) {
- ElementsAttr attr;
- if (!matchPattern(val, m_Constant(&attr)))
- return failure();
-
- if (!llvm::isa<IntegerType>(attr.getElementType()) ||
- attr.getNumElements() != 1)
- return failure();
-
- return attr.getValues<APInt>()[0].getSExtValue();
-}
-
template <typename T>
static LogicalResult verifyZeroPoint(T op, Value val, const int64_t &zp,
const std::string &operand) {
@@ -3364,6 +3352,37 @@ LogicalResult tosa::GatherOp::inferReturnTypeComponents(
return success();
}
+LogicalResult tosa::RowGatherOp::inferReturnTypeComponents(
+ MLIRContext *context, ::std::optional<Location> location,
+ RowGatherOp::Adaptor adaptor,
+ SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+ llvm::SmallVector<int64_t> outputShape;
+ outputShape.resize(3, ShapedType::kDynamic);
+
+ const ShapeAdaptor valuesShape(adaptor.getValues().getType());
+ if (valuesShape.hasRank()) {
+ outputShape[0] = valuesShape.getDimSize(0);
+ outputShape[2] = valuesShape.getDimSize(2);
+ }
+
+ const ShapeAdaptor indicesShape(adaptor.getIndices().getType());
+ if (indicesShape.hasRank()) {
+ if (outputShape[0] == ShapedType::kDynamic)
+ outputShape[0] = indicesShape.getDimSize(0);
+
+ const FailureOr<int32_t> maybeRowCount =
+ getConstantScalarIntValue<int32_t>(adaptor.getRowCount());
+ if (succeeded(maybeRowCount) && maybeRowCount.value() > 0) {
+ const int64_t indicesW = indicesShape.getDimSize(1);
+ if (ShapedType::isStatic(indicesW))
+ outputShape[1] = indicesW * maybeRowCount.value();
+ }
+ }
+
+ inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
+ return success();
+}
+
LogicalResult tosa::RowGatherBlockScaledOp::inferReturnTypeComponents(
MLIRContext *context, ::std::optional<Location> location,
RowGatherBlockScaledOp::Adaptor adaptor,
@@ -3384,7 +3403,8 @@ LogicalResult tosa::RowGatherBlockScaledOp::inferReturnTypeComponents(
if (dataShape[0] == ShapedType::kDynamic)
dataShape[0] = indicesShape.getDimSize(0);
- if (auto rowCount = getConstantScalarIntValue(adaptor.getRowCount());
+ if (auto rowCount =
+ getConstantScalarIntValue<int32_t>(adaptor.getRowCount());
succeeded(rowCount) && rowCount.value() > 0) {
const int64_t indicesW = indicesShape.getDimSize(1);
if (ShapedType::isStatic(indicesW))
@@ -3455,6 +3475,60 @@ LogicalResult tosa::GatherOp::verify() {
return success();
}
+LogicalResult tosa::RowGatherOp::verify() {
+ if (failed(verifySameElementTypes(*this, /* inType = */ getValues().getType(),
+ /* outType = */ getOutput().getType())))
+ return failure();
+
+ const FailureOr<int32_t> maybeRowCount =
+ getConstantScalarIntValue<int32_t>(getRowCount());
+ if (succeeded(maybeRowCount) && maybeRowCount.value() <= 0)
+ return emitOpError() << "requires row_count to be > 0, got "
+ << maybeRowCount.value();
+
+ int64_t n = ShapedType::kDynamic;
+ int64_t c = ShapedType::kDynamic;
+ int64_t w = ShapedType::kDynamic;
+
+ const ShapeAdaptor valuesShape(getValues().getType());
+ if (valuesShape.hasRank()) {
+ n = valuesShape.getDimSize(0);
+ c = valuesShape.getDimSize(2);
+ }
+
+ const ShapeAdaptor indicesShape(getIndices().getType());
+ if (indicesShape.hasRank()) {
+ if (failed(tryUpdateDimOrFailure(*this, n, indicesShape.getDimSize(0),
+ "indices", "batch")))
+ return failure();
+ w = indicesShape.getDimSize(1);
+ }
+
+ const ShapeAdaptor outputShape(getOutput().getType());
+ if (outputShape.hasRank()) {
+ if (failed(tryUpdateDimOrFailure(*this, n, outputShape.getDimSize(0),
+ "output", "batch")) ||
+ failed(tryUpdateDimOrFailure(*this, c, outputShape.getDimSize(2),
+ "output", "channels")))
+ return failure();
+
+ const FailureOr<int32_t> maybeRowCount =
+ getConstantScalarIntValue<int32_t>(getRowCount());
+ if (succeeded(maybeRowCount) && maybeRowCount.value() > 0 &&
+ ShapedType::isStatic(w)) {
+ const int64_t expectedOutputRows = w * maybeRowCount.value();
+ if (ShapedType::isStatic(outputShape.getDimSize(1)) &&
+ outputShape.getDimSize(1) != expectedOutputRows)
+ return emitOpError()
+ << "requires output dimension to be equal to "
+ "indices[1]*row_count ("
+ << expectedOutputRows << "), got " << outputShape.getDimSize(1);
+ }
+ }
+
+ return success();
+}
+
LogicalResult tosa::RowGatherBlockScaledOp::verify() {
const OperandRange values = getValues();
const ResultRange output = getOutput();
@@ -3488,7 +3562,7 @@ LogicalResult tosa::RowGatherBlockScaledOp::verify() {
"values[1]", "output[1]")))
return failure();
- if (auto rowCount = getConstantScalarIntValue(getRowCount());
+ if (auto rowCount = getConstantScalarIntValue<int32_t>(getRowCount());
succeeded(rowCount) && rowCount.value() <= 0)
return emitOpError() << "requires row_count to be > 0, got "
<< rowCount.value();
@@ -3527,7 +3601,7 @@ LogicalResult tosa::RowGatherBlockScaledOp::verify() {
"output[0]", "channels")))
return failure();
- if (auto rowCount = getConstantScalarIntValue(getRowCount());
+ if (auto rowCount = getConstantScalarIntValue<int32_t>(getRowCount());
succeeded(rowCount) && rowCount.value() > 0 &&
ShapedType::isStatic(w)) {
const int64_t expectedOutputRows = w * rowCount.value();
@@ -3556,7 +3630,7 @@ LogicalResult tosa::RowGatherBlockScaledOp::verify() {
"output[1]", "batch")))
return failure();
- if (auto rowCount = getConstantScalarIntValue(getRowCount());
+ if (auto rowCount = getConstantScalarIntValue<int32_t>(getRowCount());
succeeded(rowCount) && rowCount.value() > 0 &&
ShapedType::isStatic(w)) {
const int64_t expectedOutputRows = w * rowCount.value();
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
index 9511d4da89dbd..1bf92ff562a84 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
@@ -185,6 +185,14 @@ LogicalResult ProfileInfoDepot::populateProfileInfo(tosa::GatherOp op) {
return success();
}
+template <>
+LogicalResult ProfileInfoDepot::populateProfileInfo(tosa::RowGatherOp op) {
+ addValue(op.getValues());
+ addValue(op.getIndices());
+ addValue(op.getOutput());
+ return success();
+}
+
template <>
LogicalResult
ProfileInfoDepot::populateProfileInfo(tosa::RowGatherBlockScaledOp op) {
@@ -300,6 +308,7 @@ LogicalResult ProfileInfoDepot::populatationDispatch(Operation *op) {
POPULATE_PROFILE_INFO_CUSTOM(Tile)
POPULATE_PROFILE_INFO_CUSTOM(Transpose)
POPULATE_PROFILE_INFO_CUSTOM(Gather)
+ POPULATE_PROFILE_INFO_CUSTOM(RowGather)
POPULATE_PROFILE_INFO_CUSTOM(RowGatherBlockScaled)
POPULATE_PROFILE_INFO_CUSTOM(Scatter)
POPULATE_PROFILE_INFO_CUSTOM(Resize)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
index 34ac84d237f45..4a29e5327cc55 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -125,6 +125,15 @@ checkConstantOperandRowGatherBlockScaled(Operation *op, const TargetEnv &env) {
return success();
}
+static LogicalResult checkConstantOperandRowGather(Operation *op,
+ const TargetEnv &env) {
+ if (!env.allows(Extension::dynamic) && isa<tosa::RowGatherOp>(op)) {
+ // Check 'row_count'
+ return checkConstantOperands(op, {2});
+ }
+ return success();
+}
+
static LogicalResult checkConstantOperandAvgPool2d(Operation *op,
const TargetEnv &env) {
if (!env.allows(Extension::dynamic) && isa<tosa::AvgPool2dOp>(op)) {
@@ -210,6 +219,7 @@ struct TosaValidation : public tosa::impl::TosaValidationBase<TosaValidation> {
constCheckers.emplace_back(
checkConstantOperandConvOps<tosa::TransposeConv2DOp>);
constCheckers.emplace_back(checkConstantOperandMatMul);
+ constCheckers.emplace_back(checkConstantOperandRowGather);
constCheckers.emplace_back(checkConstantOperandRowGatherBlockScaled);
constCheckers.emplace_back(checkConstantOperandAvgPool2d);
constCheckers.emplace_back(checkConstantOperandAvgPool2dAdaptive);
@@ -834,6 +844,7 @@ LogicalResult TosaValidation::levelCheckRanksAndSizes(Operation *op) {
CHECK_SIZES(RFFT2d);
// Scatter/Gather Operators
CHECK_SIZES(Gather);
+ CHECK_SIZES(RowGather);
CHECK_SIZES(Scatter);
// Image Operators
CHECK_SIZES(Resize);
diff --git a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
index dd36083e910b4..e0d4b3470c981 100644
--- a/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
+++ b/mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp
@@ -248,3 +248,19 @@ bool mlir::tosa::hasUniqueConstantScatterIndices(
return true;
}
+
+template <typename T>
+FailureOr<T> mlir::tosa::getConstantScalarIntValue(Value val) {
+ ElementsAttr attr;
+ if (!matchPattern(val, m_Constant(&attr)))
+ return failure();
+
+ if (!llvm::isa<IntegerType>(attr.getElementType()) ||
+ attr.getNumElements() != 1)
+ return failure();
+
+ return attr.getValues<T>()[0];
+}
+
+template FailureOr<int32_t>
+mlir::tosa::getConstantScalarIntValue<int32_t>(Value val);
diff --git a/mlir/test/Dialect/Tosa/availability.mlir b/mlir/test/Dialect/Tosa/availability.mlir
index 81294f2c0c308..d2a156e028fc8 100644
--- a/mlir/test/Dialect/Tosa/availability.mlir
+++ b/mlir/test/Dialect/Tosa/availability.mlir
@@ -595,6 +595,16 @@ func.func @test_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) ->
return %0 : tensor<13x26x3xf32>
}
+// -----
+// CHECK-LABEL: test_row_gather
+func.func @test_row_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // CHECK: profiles: [ [pro_int, pro_fp] ]
+ // CHECK: extensions: [ [fp8e4m3, fp8e5m2, bf16, int64] ]
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xf32>
+ return %0 : tensor<13x52x3xf32>
+}
+
// -----
// CHECK-LABEL: row_gather_block_scaled
func.func @test_row_gather_block_scaled(%arg0: tensor<13x21x32xf4E2M1FN>, %arg1: tensor<13x21x1xf8E8M0FNU>, %arg2: tensor<13x26xi32>) -> (tensor<13x52x32xf4E2M1FN>, tensor<13x52x1xf8E8M0FNU>) {
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index d4f3d23fd761e..99b2102a0a754 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -18,6 +18,36 @@ func.func @test_argmax_fold_i64_index(%arg0: tensor<1xi8>) -> tensor<i64> {
// -----
+// CHECK-LABEL: @row_gather_row_count_one_to_gather
+func.func @row_gather_row_count_one_to_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x26x3xf32> {
+ %row_count = "tosa.const"() {values = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
+ // CHECK-NOT: tosa.row_gather
+ // CHECK: tosa.gather %arg0, %arg1 : (tensor<13x21x3xf32>, tensor<13x26xi32>) -> tensor<13x26x3xf32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x26x3xf32>
+ return %0 : tensor<13x26x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @row_gather_row_count_two_no_fold
+func.func @row_gather_row_count_two_no_fold(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // CHECK: tosa.row_gather %arg0, %arg1, %{{.*}} : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xf32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xf32>
+ return %0 : tensor<13x52x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @row_gather_non_const_row_count_no_fold
+func.func @row_gather_non_const_row_count_no_fold(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>, %row_count: tensor<1xi32>) -> tensor<13x?x3xf32> {
+ // CHECK: tosa.row_gather %arg0, %arg1, %arg2 : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x?x3xf32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x?x3xf32>
+ return %0 : tensor<13x?x3xf32>
+}
+
+// -----
+
// CHECK-LABEL: @pad_wh_avg_pool2d_nofold
func.func @pad_wh_avg_pool2d_nofold(%input: tensor<1x10x8x3xf32>) -> tensor<1x6x5x3xf32> {
// CHECK: tosa.pad
diff --git a/mlir/test/Dialect/Tosa/invalid_extension.mlir b/mlir/test/Dialect/Tosa/invalid_extension.mlir
index 295d7172bc2c4..cd82e314e3b6e 100644
--- a/mlir/test/Dialect/Tosa/invalid_extension.mlir
+++ b/mlir/test/Dialect/Tosa/invalid_extension.mlir
@@ -249,6 +249,14 @@ func.func @test_gather(%arg0: tensor<13x21x3xbf16>, %arg1: tensor<13x26xi32>) ->
return %0 : tensor<13x26x3xbf16>
}
+// -----
+func.func @test_row_gather(%arg0: tensor<13x21x3xbf16>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xbf16> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op illegal: requires any of [bf16] profiles/extensions to be specified in the target environment}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xbf16>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xbf16>
+ return %0 : tensor<13x52x3xbf16>
+}
+
// -----
func.func @test_row_gather_block_scaled(%arg0: tensor<13x21x32xf4E2M1FN>, %arg1: tensor<13x21x1xf8E8M0FNU>, %arg2: tensor<13x26xi32>) -> (tensor<13x52x32xf4E2M1FN>, tensor<13x52x1xf8E8M0FNU>) {
%row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
diff --git a/mlir/test/Dialect/Tosa/level_check.mlir b/mlir/test/Dialect/Tosa/level_check.mlir
index 3042f5a23e9e6..85b5d002434bc 100644
--- a/mlir/test/Dialect/Tosa/level_check.mlir
+++ b/mlir/test/Dialect/Tosa/level_check.mlir
@@ -1286,6 +1286,15 @@ func.func @test_gather_tensor_size_invalid(%arg0: tensor<536870912x21x3xf32>, %a
// -----
+func.func @test_row_gather_tensor_size_invalid(%arg0: tensor<536870912x21x3xf32>, %arg1: tensor<536870912x26xi32>) -> tensor<536870912x52x3xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op failed level check: operand tensor size (in bytes) <= (1 << MAX_LOG2_SIZE - 1)}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<536870912x21x3xf32>, tensor<536870912x26xi32>, tensor<1xi32>) -> tensor<536870912x52x3xf32>
+ return %0 : tensor<536870912x52x3xf32>
+}
+
+// -----
+
func.func @test_custom_tensor_size_invalid(%arg0: tensor<536870912xi32>) -> tensor<536870912xi32> {
// expected-error at +1 {{'tosa.custom' op failed level check: operand tensor size (in bytes) <= (1 << MAX_LOG2_SIZE - 1)}}
%0 = tosa.custom %arg0 {operator_name="custom_test", domain_name="tosa.mlir_test", implementation_attrs="" } : (tensor<536870912xi32>) -> (tensor<536870912xi32>)
diff --git a/mlir/test/Dialect/Tosa/ops.mlir b/mlir/test/Dialect/Tosa/ops.mlir
index 5c368b3da4ff5..9c96c38e2a0d5 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -961,6 +961,45 @@ func.func @test_gather_unranked_indices(%arg0: tensor<13x21x3xf32>, %arg1: tenso
return %0 : tensor<13x26x3xf32>
}
+// -----
+// CHECK-LABEL: test_row_gather
+func.func @test_row_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xf32>
+ return %0 : tensor<13x52x3xf32>
+}
+
+// -----
+// CHECK-LABEL: test_row_gather_int64
+func.func @test_row_gather_int64(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi64>) -> tensor<13x52x3xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi64>, tensor<1xi32>) -> tensor<13x52x3xf32>
+ return %0 : tensor<13x52x3xf32>
+}
+
+// -----
+// CHECK-LABEL: test_row_gather_non_const_row_count
+func.func @test_row_gather_non_const_row_count(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>, %row_count: tensor<1xi32>) -> tensor<13x?x3xf32> {
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x?x3xf32>
+ return %0 : tensor<13x?x3xf32>
+}
+
+// -----
+// CHECK-LABEL: test_row_gather_unranked_indices
+func.func @test_row_gather_unranked_indices(%arg0: tensor<13x21x3xf32>, %arg1: tensor<*xi32>) -> tensor<13x?x3xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<*xi32>, tensor<1xi32>) -> tensor<13x?x3xf32>
+ return %0 : tensor<13x?x3xf32>
+}
+
+// -----
+// CHECK-LABEL: test_row_gather_fully_unranked
+func.func @test_row_gather_fully_unranked(%arg0: tensor<*xf32>, %arg1: tensor<*xi32>) -> tensor<*xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<*xf32>, tensor<*xi32>, tensor<1xi32>) -> tensor<*xf32>
+ return %0 : tensor<*xf32>
+}
+
// -----
// CHECK-LABEL: test_row_gather_block_scaled
func.func @test_row_gather_block_scaled(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xf32> {
@@ -1284,6 +1323,14 @@ func.func @test_concat_f8E5M2(%arg0: tensor<13x21x3xf8E5M2>, %arg1: tensor<13x21
return %0 : tensor<26x21x3xf8E5M2>
}
+// -----
+// CHECK-LABEL: row_gather_f8E5M2
+func.func @test_row_gather_f8E5M2(%arg0: tensor<13x21x3xf8E5M2>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xf8E5M2> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf8E5M2>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xf8E5M2>
+ return %0 : tensor<13x52x3xf8E5M2>
+}
+
// -----
// CHECK-LABEL: pad_f8E5M2
func.func @test_pad_f8E5M2(%arg0: tensor<13x21x3xf8E5M2>) -> tensor<13x21x3xf8E5M2> {
diff --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
index db488baef3571..5499da8c6e135 100644
--- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
@@ -771,6 +771,45 @@ func.func @gather_minimum_info(%arg0 : tensor<3x?x5xi32>, %arg1 : tensor<?x6xi32
// -----
+// CHECK-LABEL: @row_gather_static
+func.func @row_gather_static(%arg0 : tensor<3x4x5xi32>, %arg1 : tensor<3x6xi32>) {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // CHECK: tosa.row_gather %arg0, %arg1, %[[ROW_COUNT:.+]] : (tensor<3x4x5xi32>, tensor<3x6xi32>, tensor<1xi32>) -> tensor<3x12x5xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<3x4x5xi32>, tensor<3x6xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @row_gather_minimum_info
+func.func @row_gather_minimum_info(%arg0 : tensor<3x?x5xi32>, %arg1 : tensor<?x6xi32>) {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // CHECK: tosa.row_gather %arg0, %arg1, %[[ROW_COUNT:.+]] : (tensor<3x?x5xi32>, tensor<?x6xi32>, tensor<1xi32>) -> tensor<3x12x5xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<3x?x5xi32>, tensor<?x6xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @row_gather_nonconstant_row_count
+func.func @row_gather_nonconstant_row_count(%arg0 : tensor<3x4x5xi32>, %arg1 : tensor<3x6xi32>, %row_count : tensor<1xi32>) {
+ // CHECK: tosa.row_gather %arg0, %arg1, %arg2 : (tensor<3x4x5xi32>, tensor<3x6xi32>, tensor<1xi32>) -> tensor<3x?x5xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<3x4x5xi32>, tensor<3x6xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @row_gather_row_count_one
+func.func @row_gather_row_count_one(%arg0 : tensor<3x4x5xi32>, %arg1 : tensor<3x6xi32>) {
+ %row_count = "tosa.const"() {values = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
+ // CHECK: tosa.row_gather %arg0, %arg1, %[[ROW_COUNT:.+]] : (tensor<3x4x5xi32>, tensor<3x6xi32>, tensor<1xi32>) -> tensor<3x6x5xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<3x4x5xi32>, tensor<3x6xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @row_gather_block_scaled_static
func.func @row_gather_block_scaled_static(%arg0 : tensor<3x4x5xi32>, %arg1 : tensor<3x6xi32>) {
%row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
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 7ff883e8e5431..3dc10eff15a62 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
@@ -78,6 +78,15 @@ func.func @test_gather_bool_i32(%arg0: tensor<13x21x3xi1>, %arg1: tensor<13x26xi
// -----
+func.func @test_row_gather_i8_i32(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xi8> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op illegal: requires specification version compatible with 1.1.draft (got 1.0) OR requires specification version compatible with 1.1.draft (got 1.0) to be specified in the target environment}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xi8>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xi8>
+ return %0 : tensor<13x52x3xi8>
+}
+
+// -----
+
func.func @test_row_gather_block_scaled_i8_i32(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xi8> {
%row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
// expected-error at +1 {{'tosa.row_gather_block_scaled' op illegal: requires specification version compatible with 1.1.draft (got 1.0) OR requires specification version compatible with 1.1.draft (got 1.0) to be specified in the target environment}}
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-pro-fp-valid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-pro-fp-valid.mlir
index 57802fd147f23..408ee86690154 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-pro-fp-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-pro-fp-valid.mlir
@@ -16,6 +16,15 @@ func.func @test_scatter_i8_i32(%input: tensor<13x27x3xi8>, %indices: tensor<13x2
// -----
+// CHECK-LABEL: test_row_gather_i8_i32
+func.func @test_row_gather_i8_i32(%input: tensor<13x27x3xi8>, %indices: tensor<13x26xi32>) -> tensor<13x52x3xi8> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %gather = tosa.row_gather %input, %indices, %row_count : (tensor<13x27x3xi8>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xi8>
+ return %gather : tensor<13x52x3xi8>
+}
+
+// -----
+
// CHECK-LABEL: test_row_gather_block_scaled_i8_i32
func.func @test_row_gather_block_scaled_i8_i32(%input: tensor<13x27x3xi8>, %indices: tensor<13x26xi32>) -> tensor<13x52x3xi8> {
%row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
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 1bac8bacbaf40..584dc4e9aa9ed 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -223,6 +223,33 @@ func.func @test_gather_i8_i32_indices(%arg0: tensor<13x21x3xi8>, %arg1: tensor<1
// -----
+// CHECK-LABEL: test_row_gather_i8_i32_indices
+func.func @test_row_gather_i8_i32_indices(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xi8> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xi8>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xi8>
+ return %0 : tensor<13x52x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: test_row_gather_i8_i64_indices
+func.func @test_row_gather_i8_i64_indices(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi64>) -> tensor<13x52x3xi8> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xi8>, tensor<13x26xi64>, tensor<1xi32>) -> tensor<13x52x3xi8>
+ return %0 : tensor<13x52x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: test_row_gather_f8e5m2_i32_indices
+func.func @test_row_gather_f8e5m2_i32_indices(%arg0: tensor<13x21x3xf8E5M2>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xf8E5M2> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf8E5M2>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xf8E5M2>
+ return %0 : tensor<13x52x3xf8E5M2>
+}
+
+// -----
+
// CHECK-LABEL: test_row_gather_block_scaled_i8_i32_indices
func.func @test_row_gather_block_scaled_i8_i32_indices(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xi8> {
%row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
diff --git a/mlir/test/Dialect/Tosa/verifier.mlir b/mlir/test/Dialect/Tosa/verifier.mlir
index b3ae5b7a5d5f3..11e5409b4d521 100644
--- a/mlir/test/Dialect/Tosa/verifier.mlir
+++ b/mlir/test/Dialect/Tosa/verifier.mlir
@@ -499,6 +499,51 @@ func.func @test_gather_invalid_out_C(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1
return %0 : tensor<13x26x8xf32>
}
+// -----
+// CHECK-LABEL: @test_row_gather_invalid_indices_N
+func.func @test_row_gather_invalid_indices_N(%arg0: tensor<13x21x3xf32>, %arg1: tensor<12x26xi32>) -> tensor<13x52x3xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op expected batch of indices to match size 13, got 12}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<12x26xi32>, tensor<1xi32>) -> tensor<13x52x3xf32>
+ return %0 : tensor<13x52x3xf32>
+}
+
+// -----
+// CHECK-LABEL: test_row_gather_invalid_out_N
+func.func @test_row_gather_invalid_out_N(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<12x52x3xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op expected batch of output to match size 13, got 12}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<12x52x3xf32>
+ return %0 : tensor<12x52x3xf32>
+}
+
+// -----
+// CHECK-LABEL: test_row_gather_invalid_out_W
+func.func @test_row_gather_invalid_out_W(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x53x3xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op requires output dimension to be equal to indices[1]*row_count (52), got 53}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x53x3xf32>
+ return %0 : tensor<13x53x3xf32>
+}
+
+// -----
+// CHECK-LABEL: test_row_gather_invalid_out_C
+func.func @test_row_gather_invalid_out_C(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x52x8xf32> {
+ %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op expected channels of output to match size 3, got 8}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x8xf32>
+ return %0 : tensor<13x52x8xf32>
+}
+
+// -----
+// CHECK-LABEL: test_row_gather_invalid_row_count
+func.func @test_row_gather_invalid_row_count(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x26x3xf32> {
+ %row_count = "tosa.const"() {values = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op requires row_count to be > 0, got 0}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x26x3xf32>
+ return %0 : tensor<13x26x3xf32>
+}
+
// -----
// CHECK-LABEL: test_row_gather_block_scaled_output_count_mismatch
func.func @test_row_gather_block_scaled_output_count_mismatch(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> (tensor<13x52x3xf32>, tensor<13x52x3xf32>) {
>From a6c4e5da6e942e2a244864433aa081ff474c85ed Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Wed, 10 Jun 2026 21:33:36 +0100
Subject: [PATCH 2/3] [mlir][tosa] Fold unit-expanded reshapes in
reduce-transposes
Extend tosa-reduce-transposes so transpose hoisting can fold through
unit-expanded vector reshapes, not only rank-1 vector reshapes. This
covers broadcast parameter layouts such as 1x1x1xC -> 1xCx1x1, where
applying the hoisted transpose either becomes a reshape to the permuted
unit-expanded shape or folds away entirely when it maps back to the
original input shape.
Change-Id: If2eb010ac322d5b9ef73f1fdda9de6049f2d2b2d
---
.../Tosa/Transforms/TosaReduceTransposes.cpp | 38 +++++++++++++------
.../Dialect/Tosa/tosa-reduce-transposes.mlir | 28 ++++++++++++++
2 files changed, 55 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
index 65ef49bdc3077..c905219e6b5e4 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaReduceTransposes.cpp
@@ -174,6 +174,23 @@ struct TosaReduceTransposes final
transposeDenseAttribute(DenseElementsAttr input, ArrayRef<int32_t> perms);
};
+// Check if shape is of the form 1x1x...x1xNx1x...x1x1 -> 1x1x...x1xNx1x...x1x1
+// Valid examples include:
+// - N -> 1x1xNx1
+// - Nx1x1x1 -> 1x1xNx1
+// - 1x1xNx1 -> 1x1xNx1
+static LogicalResult verifyUnitExpandedVectorShape(ArrayRef<int64_t> shape) {
+ bool nonUnitDimDetected = false;
+ for (int64_t dim : shape) {
+ if (dim != 1) {
+ if (nonUnitDimDetected)
+ return failure();
+ nonUnitDimDetected = true;
+ }
+ }
+ return success();
+}
+
std::optional<DenseElementsAttr>
TosaReduceTransposes::transposeDenseAttribute(DenseElementsAttr input,
ArrayRef<int32_t> perms) {
@@ -393,24 +410,24 @@ std::optional<Value> TosaReduceTransposes::buildMappedToValue(
auto reshapeInputType =
llvm::dyn_cast<RankedTensorType>(reshapeOp.getInput1().getType());
auto reshapeInputShape = reshapeInputType.getShape();
- // want reshape N -> 1x1x...x1xNx1x...x1x1
- if (!reshapeInputType || reshapeInputShape.size() != 1)
+ if (!reshapeInputType)
return std::nullopt;
auto reshapeOutputType =
llvm::cast<RankedTensorType>(reshapeOutput.getType());
-
+ const ArrayRef<int64_t> reshapeOutputShape = reshapeOutputType.getShape();
// Instead of inserting a TransposeOp here, we check if we can fold it into
// the ReshapeOp. There is more complex cases where this is possible, and
// this check can be extended.
- // Checking if reshape is N -> 1x1x...x1xNx1x...x1x1
- auto shape = reshapeOutputType.getShape();
- size_t ones = llvm::count(shape, 1);
- // N == 1 and N != 1
- if (ones != shape.size() - 1 &&
- (ones != shape.size() || reshapeInputShape[0] != 1))
+ if (failed(verifyUnitExpandedVectorShape(reshapeInputShape)) ||
+ failed(verifyUnitExpandedVectorShape(reshapeOutputShape)))
return std::nullopt;
+ SmallVector<int64_t> mappedShape =
+ applyTOSAPermutation(reshapeOutputShape, hoistedPerms);
+ if (llvm::equal(mappedShape, reshapeInputShape))
+ return reshapeOp.getInput1();
+
// Do not insert a TransposeOp, instead we fold the reshape and its attribute.
llvm::SmallVector<int64_t> newShape;
if (!tosa::getConstShapeValues(reshapeOp.getShape().getDefiningOp(),
@@ -421,8 +438,7 @@ std::optional<Value> TosaReduceTransposes::buildMappedToValue(
ImplicitLocOpBuilder builder(reshapeOp.getLoc(), rewriter);
auto foldedReshape = ReshapeOp::create(
rewriter, reshapeOp.getLoc(),
- RankedTensorType::get(applyTOSAPermutation(shape, hoistedPerms),
- reshapeOutputType.getElementType()),
+ RankedTensorType::get(mappedShape, reshapeOutputType.getElementType()),
reshapeOp.getInput1(),
getTosaConstShape(builder, applyTOSAPermutation(llvm::ArrayRef(newShape),
hoistedPerms)));
diff --git a/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir b/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir
index b3f4260ede2f5..2988b7f0ef9d9 100644
--- a/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-reduce-transposes.mlir
@@ -174,6 +174,34 @@ func.func @test_reshape_for_broadcast(%arg0: tensor<4x3x2xi32>) -> tensor<4x3x2x
// -----
+// CHECK-LABEL: @test_multi_dim_for_broadcast
+// CHECK-NOT: tosa.transpose
+// CHECK: %[[SHAPE:.*]] = tosa.const_shape {values = dense<[1, 1, 1, 8]> : tensor<4xindex>}
+// CHECK: %[[BIAS:.*]] = tosa.reshape %arg1, %[[SHAPE]] : (tensor<8xf32>, !tosa.shape<4>) -> tensor<1x1x1x8xf32>
+// CHECK: %[[ACT:.*]] = tosa.add %arg0, %[[BIAS]] : (tensor<1x4x4x8xf32>, tensor<1x1x1x8xf32>) -> tensor<1x4x4x8xf32>
+// CHECK: %[[POOL:.*]] = tosa.avg_pool2d %[[ACT]], %arg2, %arg2
+// CHECK: %[[GATED:.*]] = tosa.mul %[[POOL]], %[[ACT]], %arg3 : (tensor<1x1x1x8xf32>, tensor<1x4x4x8xf32>, tensor<1xi8>) -> tensor<1x4x4x8xf32>
+// CHECK-NOT: tosa.transpose
+// CHECK: return %[[GATED]]
+func.func @test_multi_dim_for_broadcast(
+ %input: tensor<1x4x4x8xf32>,
+ %bias: tensor<8xf32>,
+ %zero_point: tensor<1xf32>,
+ %shift: tensor<1xi8>) -> tensor<1x4x4x8xf32> {
+ %nchw_shape = tosa.const_shape {values = dense<[1, 8, 1, 1]> : tensor<4xindex>} : () -> !tosa.shape<4>
+ %nchw = tosa.transpose %input {perms = array<i32: 0, 3, 1, 2>} : (tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32>
+ %bias4 = tosa.reshape %bias, %nchw_shape : (tensor<8xf32>, !tosa.shape<4>) -> tensor<1x8x1x1xf32>
+ %act = tosa.add %nchw, %bias4 : (tensor<1x8x4x4xf32>, tensor<1x8x1x1xf32>) -> tensor<1x8x4x4xf32>
+ %pool_input = tosa.transpose %act {perms = array<i32: 0, 2, 3, 1>} : (tensor<1x8x4x4xf32>) -> tensor<1x4x4x8xf32>
+ %pooled = tosa.avg_pool2d %pool_input, %zero_point, %zero_point {acc_type = f32, kernel = array<i64: 4, 4>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 4, 4>} : (tensor<1x4x4x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x1x1x8xf32>
+ %gate = tosa.reshape %pooled, %nchw_shape : (tensor<1x1x1x8xf32>, !tosa.shape<4>) -> tensor<1x8x1x1xf32>
+ %gated = tosa.mul %gate, %act, %shift : (tensor<1x8x1x1xf32>, tensor<1x8x4x4xf32>, tensor<1xi8>) -> tensor<1x8x4x4xf32>
+ %out = tosa.transpose %gated {perms = array<i32: 0, 2, 3, 1>} : (tensor<1x8x4x4xf32>) -> tensor<1x4x4x8xf32>
+ return %out : tensor<1x4x4x8xf32>
+}
+
+// -----
+
// COM: taken directly from ResNet18 translation.
// COM: changes: %74 as argument instead of result of conv2d
>From 2b613cc6addd973a856a7f673a865849ca1d838e Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Mon, 15 Jun 2026 15:04:53 +0100
Subject: [PATCH 3/3] address review comments
Change-Id: I882fdb7536698c8e42656a5d285fd20ff8cd779a
---
.../Dialect/Tosa/IR/TosaComplianceData.h.inc | 10 ++--------
mlir/lib/Dialect/Tosa/IR/TosaOps.cpp | 4 +---
mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir | 10 ++++++++++
mlir/test/Dialect/Tosa/verifier.mlir | 18 ++++++++++++++++++
4 files changed, 31 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index f17475fae563d..7db696cc935ea 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -939,8 +939,7 @@ extensionComplianceMap = {
{"tosa.row_gather",
{
{{Extension::int16},
- {{{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT},
- {{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT}}},
{{Extension::int64},
{{{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
{{boolT, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT},
@@ -948,15 +947,10 @@ extensionComplianceMap = {
{{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
{{i64T, i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
- {{boolT, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT},
- {{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
- {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
- {{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
{{fp16T, i64T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
{{fp32T, i64T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
{{Extension::int16, Extension::int64},
- {{{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT},
- {{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT}},
+ {{{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT}},
allOf},
{{Extension::fp8e4m3},
{{{fp8e4m3T, i32T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}}},
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index f0a44a2a22ed2..d07867067ae2c 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -3372,7 +3372,7 @@ LogicalResult tosa::RowGatherOp::inferReturnTypeComponents(
const FailureOr<int32_t> maybeRowCount =
getConstantScalarIntValue<int32_t>(adaptor.getRowCount());
- if (succeeded(maybeRowCount) && maybeRowCount.value() > 0) {
+ if (succeeded(maybeRowCount)) {
const int64_t indicesW = indicesShape.getDimSize(1);
if (ShapedType::isStatic(indicesW))
outputShape[1] = indicesW * maybeRowCount.value();
@@ -3512,8 +3512,6 @@ LogicalResult tosa::RowGatherOp::verify() {
"output", "channels")))
return failure();
- const FailureOr<int32_t> maybeRowCount =
- getConstantScalarIntValue<int32_t>(getRowCount());
if (succeeded(maybeRowCount) && maybeRowCount.value() > 0 &&
ShapedType::isStatic(w)) {
const int64_t expectedOutputRows = w * maybeRowCount.value();
diff --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
index 5499da8c6e135..cfe4ea4782037 100644
--- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
@@ -810,6 +810,16 @@ func.func @row_gather_row_count_one(%arg0 : tensor<3x4x5xi32>, %arg1 : tensor<3x
// -----
+// CHECK-LABEL: @row_gather_unranked
+func.func @row_gather_unranked(%arg0 : tensor<*xi32>, %arg1 : tensor<*xi32>) {
+ %row_count = "tosa.const"() {values = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
+ // CHECK: tosa.row_gather %arg0, %arg1, %[[ROW_COUNT:.+]] : (tensor<*xi32>, tensor<*xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<*xi32>, tensor<*xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @row_gather_block_scaled_static
func.func @row_gather_block_scaled_static(%arg0 : tensor<3x4x5xi32>, %arg1 : tensor<3x6xi32>) {
%row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
diff --git a/mlir/test/Dialect/Tosa/verifier.mlir b/mlir/test/Dialect/Tosa/verifier.mlir
index 11e5409b4d521..aa87980a5df78 100644
--- a/mlir/test/Dialect/Tosa/verifier.mlir
+++ b/mlir/test/Dialect/Tosa/verifier.mlir
@@ -544,6 +544,24 @@ func.func @test_row_gather_invalid_row_count(%arg0: tensor<13x21x3xf32>, %arg1:
return %0 : tensor<13x26x3xf32>
}
+// -----
+// CHECK-LABEL: test_row_gather_invalid_row_count_negative
+func.func @test_row_gather_invalid_row_count_negative(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x26x3xf32> {
+ %row_count = "tosa.const"() {values = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op requires row_count to be > 0, got -1}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x26x3xf32>
+ return %0 : tensor<13x26x3xf32>
+}
+
+// -----
+// CHECK-LABEL: test_row_gather_input_output_mismatch
+func.func @test_row_gather_input_output_mismatch(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x26x3xf16> {
+ %row_count = "tosa.const"() {values = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
+ // expected-error at +1 {{'tosa.row_gather' op expect input and output to have same element type, got 'f32' and 'f16'}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x26x3xf16>
+ return %0 : tensor<13x26x3xf16>
+}
+
// -----
// CHECK-LABEL: test_row_gather_block_scaled_output_count_mismatch
func.func @test_row_gather_block_scaled_output_count_mismatch(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> (tensor<13x52x3xf32>, tensor<13x52x3xf32>) {
More information about the Mlir-commits
mailing list