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@+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_]*]]: _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
