[Mlir-commits] [mlir] 7b8bc1b - Revert "[mlir][bufferization] implement BufferizableOpInterface for concat op (#140171)"
Kazu Hirata
llvmlistbot at llvm.org
Fri May 16 20:23:24 PDT 2025
Author: Kazu Hirata
Date: 2025-05-16T20:23:18-07:00
New Revision: 7b8bc1b3d1ae99894b4c7741e08a0b9bfb2ffb80
URL: https://github.com/llvm/llvm-project/commit/7b8bc1b3d1ae99894b4c7741e08a0b9bfb2ffb80
DIFF: https://github.com/llvm/llvm-project/commit/7b8bc1b3d1ae99894b4c7741e08a0b9bfb2ffb80.diff
LOG: Revert "[mlir][bufferization] implement BufferizableOpInterface for concat op (#140171)"
This reverts commit 6d9ce6767d259a5231ae312a19459f8fea3bd0ca.
Multiple builtbot failures have been reported:
https://github.com/llvm/llvm-project/pull/140171
Added:
Modified:
mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp
mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
mlir/test/Dialect/Tensor/bufferize.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp b/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp
index e7d8f52d309c9..8af087cbf0f61 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp
@@ -49,8 +49,8 @@ void TensorDialect::initialize() {
>();
addInterfaces<TensorInlinerInterface>();
declarePromisedInterfaces<
- bufferization::BufferizableOpInterface, CastOp, CollapseShapeOp, ConcatOp,
- DimOp, EmptyOp, ExpandShapeOp, ExtractSliceOp, ExtractOp, FromElementsOp,
+ bufferization::BufferizableOpInterface, CastOp, CollapseShapeOp, DimOp,
+ EmptyOp, ExpandShapeOp, ExtractSliceOp, ExtractOp, FromElementsOp,
GenerateOp, InsertOp, InsertSliceOp, PadOp, ParallelInsertSliceOp, RankOp,
ReshapeOp, SplatOp>();
declarePromisedInterfaces<transform::FindPayloadReplacementOpInterface,
diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index e1706b841eb31..c0e697292d2a0 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1048,134 +1048,6 @@ struct SplatOpInterface
}
};
-/// Bufferization of tensor.concat. Bufferizes to a new allocation that is
-/// filled with copy ops. Similar to tensor.from_elements, but using memref.copy
-/// on subviews instead of memref.store.
-struct ConcatOpInterface
- : public BufferizableOpInterface::ExternalModel<ConcatOpInterface,
- tensor::ConcatOp> {
-
- bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
-
- bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
- const AnalysisState &state) const {
- return false;
- }
-
- bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
- const AnalysisState &state) const {
- return true;
- }
-
- AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
- const AnalysisState &state) const {
- return {};
- }
-
- LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
- const BufferizationOptions &options) const {
- OpBuilder::InsertionGuard g(rewriter);
- auto concatOp = cast<tensor::ConcatOp>(op);
-
- // Allocate memory.
- Location loc = op->getLoc();
- FailureOr<Value> tensorAlloc = allocateTensorForShapedValue(
- rewriter, loc, concatOp.getResult(), options,
- /*copy=*/false);
- if (failed(tensorAlloc))
- return failure();
- auto tensorType = cast<RankedTensorType>(tensorAlloc->getType());
-
- // TODO: Implement memory space for this op.
- if (options.defaultMemorySpaceFn(tensorType) != Attribute())
- return op->emitError("memory space not implemented yet");
-
- MemRefLayoutAttrInterface layout;
- MemRefType memrefType =
- MemRefType::get(concatOp.getResultType().getShape(),
- concatOp.getResultType().getElementType(), layout);
- Value dstBuffer = rewriter.create<bufferization::ToMemrefOp>(
- op->getLoc(), memrefType, *tensorAlloc);
-
- // Extract the dimension for the concat op
- uint64_t concatDim = concatOp.getDim();
- bool dynamicConcatDim = false;
-
- SmallVector<OpFoldResult> offsets(tensorType.getRank(),
- rewriter.getIndexAttr(0));
- SmallVector<OpFoldResult> strides(tensorType.getRank(),
- rewriter.getIndexAttr(1));
- SmallVector<OpFoldResult> sizes;
-
- for (const auto &[dimIdx, dimSize] :
- llvm::enumerate(tensorType.getShape())) {
- if (dimSize == ShapedType::kDynamic) {
- auto dimOp = rewriter.create<memref::DimOp>(loc, dstBuffer, dimIdx);
- sizes.push_back(dimOp.getResult());
- if (dimIdx == concatDim)
- dynamicConcatDim = true;
- } else {
- sizes.push_back(rewriter.getIndexAttr(dimSize));
- }
- }
-
- int64_t concatDimOffset = 0;
- std::optional<Value> dynamicOffset;
- std::optional<Value> dynamicSize;
- if (dynamicConcatDim) {
- // One or more operands have dynamic size, so we must accumulate the
- // offset with arith ops.
- dynamicOffset = rewriter.create<arith::ConstantIndexOp>(loc, 0);
- }
-
- for (auto operand : concatOp.getInputs()) {
- // Get the buffer for the operand.
- FailureOr<Value> srcBuffer = getBuffer(rewriter, operand, options);
- if (failed(srcBuffer))
- return failure();
-
- // Each operand may have a
diff erent size along the concat dimension,
- // so the offset on that axis must accumulate through the loop, and the
- // size must change to the size of the current operand.
- auto operandTensorType = cast<RankedTensorType>(operand.getType());
- int64_t operandConcatDimSize = operandTensorType.getDimSize(concatDim);
-
- if (dynamicConcatDim) {
- offsets[concatDim] = dynamicOffset.value();
- dynamicSize = rewriter.create<memref::DimOp>(loc, *srcBuffer, concatDim)
- .getResult();
- sizes[concatDim] = dynamicSize.value();
- } else {
- sizes[concatDim] = rewriter.getIndexAttr(operandConcatDimSize);
- offsets[concatDim] = rewriter.getIndexAttr(concatDimOffset);
- }
-
- // Create a subview of the destination buffer.
- auto dstMemrefType = cast<MemRefType>(memrefType);
- MemRefType subviewMemRefType =
- memref::SubViewOp::inferRankReducedResultType(
- operandTensorType.getShape(), dstMemrefType, offsets, sizes,
- strides);
- Value subview = rewriter.create<memref::SubViewOp>(
- loc, subviewMemRefType, dstBuffer, offsets, sizes, strides);
-
- // Copy the source buffer into the destination subview.
- if (failed(options.createMemCpy(rewriter, loc, *srcBuffer, subview)))
- return failure();
-
- if (dynamicConcatDim) {
- dynamicOffset = rewriter.create<arith::AddIOp>(
- loc, dynamicOffset.value(), dynamicSize.value());
- } else {
- concatDimOffset += operandConcatDimSize;
- }
- }
-
- replaceOpWithBufferizedValues(rewriter, op, dstBuffer);
- return success();
- }
-};
-
} // namespace
} // namespace tensor
} // namespace mlir
@@ -1185,7 +1057,6 @@ void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
CastOp::attachInterface<CastOpInterface>(*ctx);
CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
- ConcatOp::attachInterface<ConcatOpInterface>(*ctx);
DimOp::attachInterface<DimOpInterface>(*ctx);
EmptyOp::attachInterface<EmptyOpInterface>(*ctx);
ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
diff --git a/mlir/test/Dialect/Tensor/bufferize.mlir b/mlir/test/Dialect/Tensor/bufferize.mlir
index 308e52a6d9b9a..567c4abea488e 100644
--- a/mlir/test/Dialect/Tensor/bufferize.mlir
+++ b/mlir/test/Dialect/Tensor/bufferize.mlir
@@ -615,97 +615,6 @@ func.func @tensor.splat(%f: f32) -> tensor<10x2x4xf32> {
// -----
-// CHECK-LABEL: func @tensor.concat(
-// CHECK-SAME: %[[F:.*]]: tensor<8xf32>)
-// CHECK: %[[F_MEMREF:.*]] = bufferization.to_memref %[[F]]
-// CHECK: %[[ALLOC:.*]] = memref.alloc() {{.*}} : memref<16xf32>
-// CHECK: %[[SUBVIEW1:.*]] = memref.subview %[[ALLOC]][0] [8] [1]
-// CHECK: memref.copy %[[F_MEMREF]], %[[SUBVIEW1]]
-// CHECK: %[[SUBVIEW2:.*]] = memref.subview %[[ALLOC]][8] [8] [1]
-// CHECK: memref.copy %[[F_MEMREF]], %[[SUBVIEW2]]
-// CHECK: %[[RET:.*]] = bufferization.to_tensor %[[ALLOC]]
-// CHECK: return %[[RET]]
-// CHECK: }
-func.func @tensor.concat(%f: tensor<8xf32>) -> tensor<16xf32> {
- %t = tensor.concat dim(0) %f, %f : (tensor<8xf32>, tensor<8xf32>) -> tensor<16xf32>
- return %t : tensor<16xf32>
-}
-
-// -----
-
-// CHECK-LABEL: func @tensor.concat_
diff erent_shapes(
-// CHECK-SAME: %[[F:.*]]: tensor<8x4xf32>
-// CHECK-SAME: %[[G:.*]]: tensor<8x5xf32>
-// CHECK-DAG: %[[F_MEMREF:.*]] = bufferization.to_memref %[[F]]
-// CHECK-DAG: %[[G_MEMREF:.*]] = bufferization.to_memref %[[G]]
-// CHECK: %[[ALLOC:.*]] = memref.alloc() {{.*}} : memref<8x9xf32>
-// CHECK: %[[SUBVIEW1:.*]] = memref.subview %[[ALLOC]][0, 0] [8, 4] [1, 1]
-// CHECK: memref.copy %[[F_MEMREF]], %[[SUBVIEW1]]
-// CHECK: %[[SUBVIEW2:.*]] = memref.subview %[[ALLOC]][0, 4] [8, 5] [1, 1]
-// CHECK: memref.copy %[[G_MEMREF]], %[[SUBVIEW2]]
-// CHECK: %[[RET:.*]] = bufferization.to_tensor %[[ALLOC]]
-// CHECK: return %[[RET]]
-// CHECK: }
-func.func @tensor.concat_
diff erent_shapes(%f: tensor<8x4xf32>, %g: tensor<8x5xf32>) -> tensor<8x9xf32> {
- %t = tensor.concat dim(1) %f, %g : (tensor<8x4xf32>, tensor<8x5xf32>) -> tensor<8x9xf32>
- return %t : tensor<8x9xf32>
-}
-
-// -----
-
-// CHECK-LABEL: func @tensor.concat_dynamic(
-// CHECK-SAME: %[[F:.*]]: tensor<8x?xf32>,
-// CHECK-SAME: %[[G:.*]]: tensor<8x?xf32>
-// CHECK-DAG: %[[F_MEMREF:.*]] = bufferization.to_memref %[[F]]
-// CHECK-DAG: %[[G_MEMREF:.*]] = bufferization.to_memref %[[G]]
-// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
-// CHECK-DAG: %[[F_DIM:.*]] = memref.dim %[[F_MEMREF]], %[[c1]]
-// CHECK-DAG: %[[G_DIM:.*]] = memref.dim %[[G_MEMREF]], %[[c1]]
-// CHECK: %[[ALLOC:.*]] = memref.alloc
-// CHECK-SAME: memref<8x?xf32>
-// CHECK-DAG: %[[OFFSET:.*]] = arith.constant 0 : index
-// CHECK: %[[SUBVIEW1:.*]] = memref.subview %[[ALLOC]][0, %[[OFFSET]]] [8, %[[F_DIM]]] [1, 1]
-// CHECK: memref.copy %[[F_MEMREF]], %[[SUBVIEW1]]
-// CHECK: %[[OFFSET_2:.*]] = arith.addi %[[OFFSET]], %[[F_DIM]] : index
-// CHECK: %[[SUBVIEW2:.*]] = memref.subview %[[ALLOC]][0, %[[OFFSET_2]]] [8, %[[G_DIM]]] [1, 1]
-// CHECK: memref.copy %[[G_MEMREF]], %[[SUBVIEW2]]
-// CHECK: %[[RET:.*]] = bufferization.to_tensor %[[ALLOC]]
-// CHECK: return %[[RET]]
-// CHECK: }
-func.func @tensor.concat_dynamic(%f: tensor<8x?xf32>, %g: tensor<8x?xf32>) -> tensor<8x?xf32> {
- %t = tensor.concat dim(1) %f, %g : (tensor<8x?xf32>, tensor<8x?xf32>) -> tensor<8x?xf32>
- return %t : tensor<8x?xf32>
-}
-
-// -----
-
-// CHECK-LABEL: func @tensor.concat_dynamic_nonconcat_dim(
-// CHECK-SAME: %[[F:.*]]: tensor<?x?xf32>,
-// CHECK-SAME: %[[G:.*]]: tensor<?x?xf32>
-// CHECK-DAG: %[[F_MEMREF:.*]] = bufferization.to_memref %[[F]]
-// CHECK-DAG: %[[G_MEMREF:.*]] = bufferization.to_memref %[[G]]
-// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
-// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
-// CHECK-DAG: %[[F_DIM:.*]] = memref.dim %[[F_MEMREF]], %[[c1]]
-// CHECK-DAG: %[[G_DIM:.*]] = memref.dim %[[G_MEMREF]], %[[c1]]
-// CHECK: %[[ALLOC:.*]] = memref.alloc
-// CHECK-SAME: memref<?x?xf32>
-// CHECK-DAG: %[[NON_CONCAT_DIM:.*]] = memref.dim %[[ALLOC]], %[[c0]]
-// CHECK: %[[SUBVIEW1:.*]] = memref.subview %[[ALLOC]][0, %[[c0]]] [%[[NON_CONCAT_DIM]], %[[F_DIM]]] [1, 1]
-// CHECK: memref.copy %[[F_MEMREF]], %[[SUBVIEW1]]
-// CHECK: %[[OFFSET_2:.*]] = arith.addi %[[c0]], %[[F_DIM]] : index
-// CHECK: %[[SUBVIEW2:.*]] = memref.subview %[[ALLOC]][0, %[[OFFSET_2]]] [%[[NON_CONCAT_DIM]], %[[G_DIM]]] [1, 1]
-// CHECK: memref.copy %[[G_MEMREF]], %[[SUBVIEW2]]
-// CHECK: %[[RET:.*]] = bufferization.to_tensor %[[ALLOC]]
-// CHECK: return %[[RET]]
-// CHECK: }
-func.func @tensor.concat_dynamic_nonconcat_dim(%f: tensor<?x?xf32>, %g: tensor<?x?xf32>) -> tensor<?x?xf32> {
- %t = tensor.concat dim(1) %f, %g : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
- return %t : tensor<?x?xf32>
-}
-
-// -----
-
// CHECK-LABEL: func @tensor.splat_dynamic(
// CHECK-SAME: %[[F:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[M:[a-zA-Z0-9_]+]]: index
More information about the Mlir-commits
mailing list