[llvm-branch-commits] [mlir] df18ff3 - Revert "[mlir][tosa] Lower ROW_GATHER to Linalg (#225417)"
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Tue Oct 6 04:45:40 PDT 2026
Author: Ingo Müller
Date: 2026-10-06T13:45:32+02:00
New Revision: df18ff3d50c7c7750e5ab392fc2a6561ec653e40
URL: https://github.com/llvm/llvm-project/commit/df18ff3d50c7c7750e5ab392fc2a6561ec653e40
DIFF: https://github.com/llvm/llvm-project/commit/df18ff3d50c7c7750e5ab392fc2a6561ec653e40.diff
LOG: Revert "[mlir][tosa] Lower ROW_GATHER to Linalg (#225417)"
This reverts commit 57f2572b545ae1083c663b6636bae61191ba8c02.
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 555e5515cf439..b2a8360ee0c2b 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -2671,75 +2671,6 @@ 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.
@@ -3268,7 +3199,6 @@ 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 57116f03ad34f..d3882e1cede0a 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
@@ -73,12 +73,3 @@ 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 2929fa83930f4..9aa6d72405426 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -1987,77 +1987,6 @@ 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 llvm-branch-commits
mailing list