[Mlir-commits] [mlir] [mlir][tosa] Add row_gather operator (PR #202895)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 10 01:50:50 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Luke Hutton (lhutton1)
<details>
<summary>Changes</summary>
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.
---
Patch is 37.93 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/202895.diff
18 Files Affected:
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc (+65)
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td (+36)
- (modified) mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h (+5)
- (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (+20)
- (modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+89-16)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp (+9)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp (+11)
- (modified) mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp (+16)
- (modified) mlir/test/Dialect/Tosa/availability.mlir (+10)
- (modified) mlir/test/Dialect/Tosa/canonicalize.mlir (+30)
- (modified) mlir/test/Dialect/Tosa/invalid_extension.mlir (+8)
- (modified) mlir/test/Dialect/Tosa/level_check.mlir (+9)
- (modified) mlir/test/Dialect/Tosa/ops.mlir (+47)
- (modified) mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir (+39)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir (+9)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p1-pro-fp-valid.mlir (+9)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir (+27)
- (modified) mlir/test/Dialect/Tosa/verifier.mlir (+45)
``````````diff
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..ec5bd53951324 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1185,6 +1185,26 @@ 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..32e7d7bb47531 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,7 @@ 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 +3474,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 +3561,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 +3600,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 +3629,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,...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/202895
More information about the Mlir-commits
mailing list