[Mlir-commits] [mlir] e8f3b2d - [mlir][tosa] Lower concat to tensor.concat (#231726)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 21 02:40:47 PDT 2026
Author: Akun
Date: 2026-10-21T10:40:41+01:00
New Revision: e8f3b2d8d90ca0ed1594615188cf34a503aa53b3
URL: https://github.com/llvm/llvm-project/commit/e8f3b2d8d90ca0ed1594615188cf34a503aa53b3
DIFF: https://github.com/llvm/llvm-project/commit/e8f3b2d8d90ca0ed1594615188cf34a503aa53b3.diff
LOG: [mlir][tosa] Lower concat to tensor.concat (#231726)
Lower `tosa.concat` directly to `tensor.concat` instead of expanding it
into `tensor.empty` and `tensor.insert_slice` operations.
Update the TosaToTensor conversion tests to check the resulting concat
operations.
Added:
Modified:
mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp b/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
index d95fa14ff8bd01..622f8a6f13a198 100644
--- a/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
+++ b/mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
@@ -12,7 +12,6 @@
#include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
@@ -397,59 +396,14 @@ struct ConcatConverter : public OpConversionPattern<tosa::ConcatOp> {
LogicalResult
matchAndRewrite(tosa::ConcatOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- auto resultType = dyn_cast<RankedTensorType>(op.getType());
-
- Location loc = op.getLoc();
- int axis = op.getAxis();
- Value axisValue =
- arith::ConstantOp::create(rewriter, loc, rewriter.getIndexAttr(axis));
- int64_t rank = resultType.getRank();
-
- SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
- SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
- SmallVector<OpFoldResult> sizes =
- tensor::getMixedSizes(rewriter, op.getLoc(), adaptor.getOperands()[0]);
-
- // Pre-compute the offsets along the axis dimension.
- // The axisOffsets will be of size rank + 1, where the last value
- // will hold the total size of the tensor along the 'axis' dimension.
- SmallVector<OpFoldResult> axisOffsets;
- axisOffsets.push_back(rewriter.getIndexAttr(0));
- axisOffsets.push_back(sizes[axis]);
-
- for (auto arg : adaptor.getOperands().drop_front()) {
- auto size = rewriter.createOrFold<tensor::DimOp>(loc, arg, axisValue);
- auto currentOffset =
- getValueOrCreateConstantIndexOp(rewriter, loc, axisOffsets.back());
- auto total =
- rewriter.createOrFold<arith::AddIOp>(loc, currentOffset, size);
- axisOffsets.push_back(getAsOpFoldResult(total));
- }
- sizes[axis] = axisOffsets.back();
-
- // Compute the dynamic sizes of the tensor.empty operation.
- // This is based off of the specified result type of the tosa.concat
- // operation, since we don't want to change the result type of the operation
- // during the conversion.
- SmallVector<Value> dynDims;
- for (int64_t i = 0; i < rank; ++i) {
- if (resultType.isDynamicDim(i)) {
- dynDims.push_back(
- getValueOrCreateConstantIndexOp(rewriter, loc, sizes[i]));
- }
- }
-
- Value result =
- tensor::EmptyOp::create(rewriter, loc, resultType.getShape(),
- resultType.getElementType(), dynDims);
+ auto resultType =
+ getTypeConverter()->convertType<RankedTensorType>(op.getType());
+ if (!resultType)
+ return failure();
- for (auto [arg, offset] : llvm::zip(adaptor.getOperands(), axisOffsets)) {
- auto sizes = tensor::getMixedSizes(rewriter, op.getLoc(), arg);
- offsets[axis] = offset;
- result = rewriter.createOrFold<tensor::InsertSliceOp>(
- loc, arg, result, offsets, sizes, strides);
- }
- rewriter.replaceOp(op, result);
+ rewriter.replaceOpWithNewOp<tensor::ConcatOp>(
+ op, resultType, rewriter.getI64IntegerAttr(op.getAxis()),
+ adaptor.getInput1());
return success();
}
};
diff --git a/mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir b/mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir
index c4e98bf7bece9d..c60d340e8a7da1 100644
--- a/mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir
+++ b/mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir
@@ -589,14 +589,10 @@ func.func @pad_dyn_padding(%arg0 : tensor<1x2xf32>) -> (tensor<?x9xf32>) {
// CHECK-SAME: %[[ARG0:.+]]: tensor<5x1xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<6x1xf32>
func.func @concat(%arg0: tensor<5x1xf32>, %arg1: tensor<6x1xf32>) -> () {
- // CHECK-DAG: [[INIT:%.+]] = tensor.empty() : tensor<11x1xf32>
- // CHECK-DAG: [[INSERT0:%.+]] = tensor.insert_slice %[[ARG0]] into [[INIT]][0, 0] [5, 1] [1, 1]
- // CHECK-DAG: [[INSERT1:%.+]] = tensor.insert_slice %[[ARG1]] into [[INSERT0]][5, 0] [6, 1] [1, 1]
+ // CHECK: tensor.concat dim(0) %[[ARG0]], %[[ARG1]] : (tensor<5x1xf32>, tensor<6x1xf32>) -> tensor<11x1xf32>
%0 = "tosa.concat"(%arg0, %arg1) <{ axis = 0 : i32}> : (tensor<5x1xf32>, tensor<6x1xf32>) -> (tensor<11x1xf32>)
- // CHECK-DAG: [[INIT:%.+]] = tensor.empty() : tensor<5x2xf32>
- // CHECK-DAG: [[INSERT0:%.+]] = tensor.insert_slice %[[ARG0]] into [[INIT]][0, 0] [5, 1] [1, 1]
- // CHECK: [[INSERT1:%.+]] = tensor.insert_slice %[[ARG0]] into [[INSERT0]][0, 1] [5, 1] [1, 1]
+ // CHECK: tensor.concat dim(1) %[[ARG0]], %[[ARG0]] : (tensor<5x1xf32>, tensor<5x1xf32>) -> tensor<5x2xf32>
%1 = "tosa.concat"(%arg0, %arg0) <{ axis = 1 : i32}> : (tensor<5x1xf32>, tensor<5x1xf32>) -> (tensor<5x2xf32>)
return
}
@@ -607,16 +603,7 @@ func.func @concat(%arg0: tensor<5x1xf32>, %arg1: tensor<6x1xf32>) -> () {
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]
func.func @concat_non_axis_dyn(%arg0: tensor<5x?xf32>, %arg1: tensor<6x?xf32>) -> () {
- // CHECK-DAG: %[[AXIS:.+]] = arith.constant 0
- // CHECK-DAG: %[[IDX1:.+]] = arith.constant 1
- // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[IDX1]]
- // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[DIM0]]) : tensor<11x?xf32>
- // CHECK-DAG: %[[IDX1_1:.+]] = arith.constant 1 : index
- // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[IDX1_1]]
- // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [5, %[[DIM1]]] [1, 1]
- // CHECK-DAG: %[[IDX1_2:.+]] = arith.constant 1 : index
- // CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG1]], %[[IDX1_2]] : tensor<6x?xf32>
- // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][5, 0] [6, %[[DIM2]]] [1, 1]
+ // CHECK: tensor.concat dim(0) %[[ARG0]], %[[ARG1]] : (tensor<5x?xf32>, tensor<6x?xf32>) -> tensor<11x?xf32>
%0 = "tosa.concat"(%arg0, %arg1) <{ axis = 0 : i32}> : (tensor<5x?xf32>, tensor<6x?xf32>) -> (tensor<11x?xf32>)
return
}
@@ -627,19 +614,7 @@ func.func @concat_non_axis_dyn(%arg0: tensor<5x?xf32>, %arg1: tensor<6x?xf32>) -
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
func.func @concat_axis_dyn(%arg0: tensor<?x3xf32>, %arg1: tensor<?x3xf32>) -> () {
- // CHECK-DAG: %[[AXIS:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[IDX0:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[IDX0]] : tensor<?x3xf32>
- // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG1]], %[[AXIS]] : tensor<?x3xf32>
- // CHECK-DAG: %[[SUM:.+]] = arith.addi %[[DIM0]], %[[DIM1]] : index
- // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[SUM]]) : tensor<?x3xf32>
- // CHECK-DAG: %[[IDX0_1:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG0]], %[[IDX0_1]] : tensor<?x3xf32>
- // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM2]], 3] [1, 1] : tensor<?x3xf32> into tensor<?x3xf32>
- // CHECK-DAG: %[[IDX0_2:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[DIM3:.+]] = tensor.dim %[[ARG1]], %[[IDX0_2]] : tensor<?x3xf32>
- // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][%[[DIM0]], 0] [%[[DIM3]], 3] [1, 1] : tensor<?x3xf32> into tensor<?x3xf32>
-
+ // CHECK: tensor.concat dim(0) %[[ARG0]], %[[ARG1]] : (tensor<?x3xf32>, tensor<?x3xf32>) -> tensor<?x3xf32>
%0 = "tosa.concat"(%arg0, %arg1) <{ axis = 0 : i32}> : (tensor<?x3xf32>, tensor<?x3xf32>) -> (tensor<?x3xf32>)
return
}
@@ -651,26 +626,7 @@ func.func @concat_axis_dyn(%arg0: tensor<?x3xf32>, %arg1: tensor<?x3xf32>) -> ()
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z_]*]]:
func.func @concat_axis_dyn_mixed(%arg0: tensor<?x1xf32>, %arg1: tensor<?x1xf32>, %arg2: tensor<?x1xf32>) -> () {
- // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[C0_0:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[OFFSET0:.+]] = tensor.dim %[[ARG0]], %[[C0_0]] : tensor<?x1xf32>
- // CHECK-DAG: %[[DIM1_0:.+]] = tensor.dim %[[ARG1]], %[[C0]] : tensor<?x1xf32>
- // CHECK-DAG: %[[OFFSET1:.+]] = arith.addi %[[OFFSET0]], %[[DIM1_0]] : index
- // CHECK-DAG: %[[DIM2_2:.+]] = tensor.dim %[[ARG2]], %[[C0]] : tensor<?x1xf32>
- // CHECK-DAG: %[[OFFSET2:.+]] = arith.addi %[[OFFSET1]], %[[DIM2_2]] : index
- // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<5x1xf32>
- // CHECK-DAG: %[[C0_3:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[DIM_4:.+]] = tensor.dim %[[ARG0]], %[[C0_3]] : tensor<?x1xf32>
- // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM_4]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x1xf32>
- // CHECK-DAG: %[[C0_4:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[DIM_6:.+]] = tensor.dim %[[ARG1]], %[[C0_4]] : tensor<?x1xf32>
- // CHECK-DAG: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][%[[OFFSET0]], 0] [%[[DIM_6]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x1xf32>
- // CHECK-DAG: %[[C0_8:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[DIM_9:.+]] = tensor.dim %[[ARG2]], %[[C0_8]] : tensor<?x1xf32>
- // CHECK-DAG: %[[INSERT3:.+]] = tensor.insert_slice %[[ARG2]] into %[[INSERT1]][%[[OFFSET1]], 0] [%[[DIM_9]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x1xf32>
-
- // CHECK: return
-
+ // CHECK: tensor.concat dim(0) %[[ARG0]], %[[ARG1]], %[[ARG2]] : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>) -> tensor<5x1xf32>
%0 = "tosa.concat"(%arg0, %arg1, %arg2) <{axis = 0 : i32}> : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>) -> tensor<5x1xf32>
return
}
@@ -682,21 +638,7 @@ func.func @concat_axis_dyn_mixed(%arg0: tensor<?x1xf32>, %arg1: tensor<?x1xf32>,
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z_]*]]:
func.func @concat_non_axis_dyn_mixed(%arg0: tensor<?x1xf32>, %arg1: tensor<?x1xf32>, %arg2: tensor<?x1xf32>) -> () {
- // CHECK-DAG: %[[UNUSED0:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[UNUSED1:.+]] = tensor.dim %[[ARG0]], %[[UNUSED0]] : tensor<?x1xf32>
-
- // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<5x3xf32>
- // CHECK-DAG: %[[C0_0:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[DIM0_0:.+]] = tensor.dim %[[ARG0]], %[[C0_0]] : tensor<?x1xf32>
- // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM0_0]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x3xf32>
- // CHECK-DAG: %[[C0_1:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[DIM1_0:.+]] = tensor.dim %[[ARG1]], %[[C0_1]] : tensor<?x1xf32>
- // CHECK-DAG: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][0, 1] [%[[DIM1_0]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x3xf32>
- // CHECK-DAG: %[[C0_2:.+]] = arith.constant 0 : index
- // CHECK-DAG: %[[DIM2_0:.+]] = tensor.dim %[[ARG2]], %[[C0_2]] : tensor<?x1xf32>
- // CHECK-DAG: %[[INSERT2:.+]] = tensor.insert_slice %[[ARG2]] into %[[INSERT1]][0, 2] [%[[DIM2_0]], 1] [1, 1] : tensor<?x1xf32> into tensor<5x3xf32>
- // CHECK: return
-
+ // CHECK: tensor.concat dim(1) %[[ARG0]], %[[ARG1]], %[[ARG2]] : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>) -> tensor<5x3xf32>
%0 = "tosa.concat"(%arg0, %arg1, %arg2) <{axis = 1 : i32}> : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>) -> tensor<5x3xf32>
return
}
More information about the Mlir-commits
mailing list