[Mlir-commits] [mlir] 57f2572 - [mlir][tosa] Lower ROW_GATHER to Linalg (#225417)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 6 04:09:22 PDT 2026
Author: Jack Frankland
Date: 2026-10-06T12:09:15+01:00
New Revision: 57f2572b545ae1083c663b6636bae61191ba8c02
URL: https://github.com/llvm/llvm-project/commit/57f2572b545ae1083c663b6636bae61191ba8c02
DIFF: https://github.com/llvm/llvm-project/commit/57f2572b545ae1083c663b6636bae61191ba8c02.diff
LOG: [mlir][tosa] Lower ROW_GATHER to Linalg (#225417)
Lower ROW_GATHER to a linalg.generic that maps each expanded output row
back to its source index slot and consecutive row offset. Support
dynamic output dimensions and both i32 and i64 indices.
Assisted-by: Codex
Added:
Modified:
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index b2a8360ee0c2b..555e5515cf439 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -2671,6 +2671,75 @@ class GatherConverter : public OpConversionPattern<tosa::GatherOp> {
}
};
+class RowGatherConverter : public OpConversionPattern<tosa::RowGatherOp> {
+public:
+ using OpConversionPattern<tosa::RowGatherOp>::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(tosa::RowGatherOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const final {
+ auto valuesTy = dyn_cast<RankedTensorType>(adaptor.getValues().getType());
+ auto indicesTy = dyn_cast<RankedTensorType>(adaptor.getIndices().getType());
+ auto rowCountTy =
+ dyn_cast<RankedTensorType>(adaptor.getRowCount().getType());
+ auto resultTy = dyn_cast<RankedTensorType>(op.getType());
+ if (!valuesTy || !indicesTy || !rowCountTy || !resultTy)
+ return rewriter.notifyMatchFailure(op, "unranked tensors not supported");
+
+ Location loc = op.getLoc();
+ Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ Value rowCount = tensor::ExtractOp::create(
+ rewriter, loc, adaptor.getRowCount(), ValueRange{zero});
+ Value rowCountIndex = arith::IndexCastOp::create(
+ rewriter, loc, rewriter.getIndexType(), rowCount);
+
+ SmallVector<Value> dynamicDims;
+ if (resultTy.isDynamicDim(0))
+ dynamicDims.push_back(
+ tensor::DimOp::create(rewriter, loc, adaptor.getValues(), 0));
+ if (resultTy.isDynamicDim(1)) {
+ Value indicesWidth =
+ tensor::DimOp::create(rewriter, loc, adaptor.getIndices(), 1);
+ dynamicDims.push_back(
+ arith::MulIOp::create(rewriter, loc, indicesWidth, rowCountIndex));
+ }
+ if (resultTy.isDynamicDim(2))
+ dynamicDims.push_back(
+ tensor::DimOp::create(rewriter, loc, adaptor.getValues(), 2));
+
+ Value emptyTensor =
+ tensor::EmptyOp::create(rewriter, loc, resultTy.getShape(),
+ resultTy.getElementType(), dynamicDims);
+ SmallVector<AffineMap> affineMaps = {
+ rewriter.getMultiDimIdentityMap(resultTy.getRank())};
+
+ auto genericOp = linalg::GenericOp::create(
+ rewriter, loc, ArrayRef<Type>{resultTy}, ValueRange{},
+ ValueRange{emptyTensor}, affineMaps,
+ getNParallelLoopsAttrs(resultTy.getRank()),
+ [&](OpBuilder &builder, Location nestedLoc, ValueRange) {
+ Value batch = linalg::IndexOp::create(builder, nestedLoc, 0);
+ Value outputRow = linalg::IndexOp::create(builder, nestedLoc, 1);
+ Value channel = linalg::IndexOp::create(builder, nestedLoc, 2);
+ Value indexSlot = arith::DivUIOp::create(builder, nestedLoc,
+ outputRow, rowCountIndex);
+ Value rowOffset = arith::RemUIOp::create(builder, nestedLoc,
+ outputRow, rowCountIndex);
+ Value index = tensor::ExtractOp::create(builder, nestedLoc,
+ adaptor.getIndices(),
+ ValueRange{batch, indexSlot});
+ Value row = arith::IndexCastOp::create(builder, nestedLoc,
+ builder.getIndexType(), index);
+ row = arith::AddIOp::create(builder, nestedLoc, row, rowOffset);
+ Value result =
+ tensor::ExtractOp::create(builder, nestedLoc, adaptor.getValues(),
+ ValueRange{batch, row, channel});
+ linalg::YieldOp::create(builder, nestedLoc, result);
+ });
+ rewriter.replaceOp(op, genericOp.getResult(0));
+ return success();
+ }
+};
+
// Lowerings the TableOp to a series of gathers and numerica operations. This
// includes interpolation between the high/low values. For the I8 varient, this
// simplifies to a single gather operation.
@@ -3199,6 +3268,7 @@ void mlir::tosa::populateTosaToLinalgConversionPatterns(
patterns->add<
IdentityNConverter<tosa::IdentityOp>,
GatherConverter,
+ RowGatherConverter,
RescaleConverter,
ReverseConverter,
RFFT2dConverter,
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
index d3882e1cede0a..57116f03ad34f 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
@@ -73,3 +73,12 @@ func.func @unranked_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>)
%0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xf32>, tensor<13x26xi32>) -> tensor<*xf32>
return %0 : tensor<*xf32>
}
+
+// -----
+
+// Verify ROW_GATHER rejects unranked results during lowering.
+func.func @unranked_row_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>, %row_count: tensor<1xi32>) -> tensor<*xf32> {
+ // expected-error at +1 {{failed to legalize operation 'tosa.row_gather'}}
+ %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<*xf32>
+ return %0 : tensor<*xf32>
+}
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index 9aa6d72405426..2929fa83930f4 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -1987,6 +1987,77 @@ func.func @gather_int(%arg0: tensor<2x3x2xi32>, %arg1: tensor<2x3xi32>) -> () {
// -----
+// Verify batched ROW_GATHER lowering and row_count output expansion.
+// CHECK-LABEL: @row_gather
+// CHECK-SAME: (%[[VALUES:[0-9a-zA-Z_]*]]
+// CHECK-SAME: %[[INDICES:[0-9a-zA-Z_]*]]
+// CHECK-SAME: %[[ROW_COUNT_TENSOR:[0-9a-zA-Z_]*]]
+func.func @row_gather(%values: tensor<2x5x3xf32>, %indices: tensor<2x2xi32>, %row_count: tensor<1xi32>) -> tensor<2x4x3xf32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0 : index
+ // CHECK: %[[ROW_COUNT:.+]] = tensor.extract %[[ROW_COUNT_TENSOR]][%[[C0]]] : tensor<1xi32>
+ // CHECK: %[[ROW_COUNT_INDEX:.+]] = arith.index_cast %[[ROW_COUNT]] : i32 to index
+ // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x4x3xf32>
+ // CHECK: %[[RESULT:.+]] = linalg.generic
+ // CHECK-SAME: outs(%[[INIT]] : tensor<2x4x3xf32>)
+ // CHECK: %[[BATCH:.+]] = linalg.index 0 : index
+ // CHECK: %[[OUTPUT_ROW:.+]] = linalg.index 1 : index
+ // CHECK: %[[CHANNEL:.+]] = linalg.index 2 : index
+ // CHECK: %[[INDEX_SLOT:.+]] = arith.divui %[[OUTPUT_ROW]], %[[ROW_COUNT_INDEX]] : index
+ // CHECK: %[[ROW_OFFSET:.+]] = arith.remui %[[OUTPUT_ROW]], %[[ROW_COUNT_INDEX]] : index
+ // CHECK: %[[INDEX:.+]] = tensor.extract %[[INDICES]][%[[BATCH]], %[[INDEX_SLOT]]] : tensor<2x2xi32>
+ // CHECK: %[[ROW:.+]] = arith.index_cast %[[INDEX]] : i32 to index
+ // CHECK: %[[INPUT_ROW:.+]] = arith.addi %[[ROW]], %[[ROW_OFFSET]] : index
+ // CHECK: %[[VALUE:.+]] = tensor.extract %[[VALUES]][%[[BATCH]], %[[INPUT_ROW]], %[[CHANNEL]]] : tensor<2x5x3xf32>
+ // CHECK: linalg.yield %[[VALUE]] : f32
+ // CHECK: return %[[RESULT]]
+ %0 = tosa.row_gather %values, %indices, %row_count : (tensor<2x5x3xf32>, tensor<2x2xi32>, tensor<1xi32>) -> tensor<2x4x3xf32>
+ return %0 : tensor<2x4x3xf32>
+}
+
+// -----
+
+// Verify i64 indices and the largest valid start row.
+// CHECK-LABEL: @row_gather_i64_boundary_indices
+func.func @row_gather_i64_boundary_indices(%values: tensor<2x5x2xi16>) -> tensor<2x4x2xi16> {
+ // CHECK: %[[INDICES:.+]] = "tosa.const"()
+ // CHECK: %[[ROW_COUNT:.+]] = "tosa.const"()
+ // CHECK: linalg.generic
+ // CHECK: %[[INDEX:.+]] = tensor.extract %[[INDICES]][{{.+}}] : tensor<2x2xi64>
+ // CHECK: %[[ROW:.+]] = arith.index_cast %[[INDEX]] : i64 to index
+ // CHECK: %[[VALUE:.+]] = tensor.extract %arg0[{{.+}}, %{{.+}}, %{{.+}}] : tensor<2x5x2xi16>
+ // CHECK: linalg.yield %[[VALUE]] : i16
+ %indices = "tosa.const"() <{values = dense<[[0, 3], [3, 1]]> : tensor<2x2xi64>}> : () -> tensor<2x2xi64>
+ %row_count = "tosa.const"() <{values = dense<2> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %0 = tosa.row_gather %values, %indices, %row_count : (tensor<2x5x2xi16>, tensor<2x2xi64>, tensor<1xi32>) -> tensor<2x4x2xi16>
+ return %0 : tensor<2x4x2xi16>
+}
+
+// -----
+
+// Verify dynamic output dimensions are materialized from the operands.
+// CHECK-LABEL: @row_gather_all_dynamic
+// CHECK-SAME: (%[[VALUES:[0-9a-zA-Z_]*]]
+// CHECK-SAME: %[[INDICES:[0-9a-zA-Z_]*]]
+// CHECK-SAME: %[[ROW_COUNT_TENSOR:[0-9a-zA-Z_]*]]
+func.func @row_gather_all_dynamic(%values: tensor<?x?x?xi8>, %indices: tensor<?x?xi32>, %row_count: tensor<1xi32>) -> tensor<?x?x?xi8> {
+ // CHECK: %[[C0:.+]] = arith.constant 0 : index
+ // CHECK: %[[ROW_COUNT:.+]] = tensor.extract %[[ROW_COUNT_TENSOR]][%[[C0]]] : tensor<1xi32>
+ // CHECK: %[[ROW_COUNT_INDEX:.+]] = arith.index_cast %[[ROW_COUNT]] : i32 to index
+ // CHECK: %[[BATCH:.+]] = tensor.dim %[[VALUES]], %[[C0]]
+ // CHECK: %[[C1:.+]] = arith.constant 1 : index
+ // CHECK: %[[WIDTH:.+]] = tensor.dim %[[INDICES]], %[[C1]]
+ // CHECK: %[[OUTPUT_ROWS:.+]] = arith.muli %[[WIDTH]], %[[ROW_COUNT_INDEX]] : index
+ // CHECK: %[[C2:.+]] = arith.constant 2 : index
+ // CHECK: %[[CHANNELS:.+]] = tensor.dim %[[VALUES]], %[[C2]]
+ // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]], %[[OUTPUT_ROWS]], %[[CHANNELS]]) : tensor<?x?x?xi8>
+ // CHECK: linalg.generic
+ // CHECK-SAME: outs(%[[INIT]] : tensor<?x?x?xi8>)
+ %0 = tosa.row_gather %values, %indices, %row_count : (tensor<?x?x?xi8>, tensor<?x?xi32>, tensor<1xi32>) -> tensor<?x?x?xi8>
+ return %0 : tensor<?x?x?xi8>
+}
+
+// -----
+
// CHECK-LABEL: @table8
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
More information about the Mlir-commits
mailing list