diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp index 9a831e4c322ea..f74c6255c196b 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp @@ -625,46 +625,6 @@ class BufferDeallocation : public BufferPlacementTransformationBase { // BufferDeallocationPass //===----------------------------------------------------------------------===// -struct DefaultAllocationInterface - : public bufferization::AllocationOpInterface::ExternalModel< - DefaultAllocationInterface, memref::AllocOp> { - static std::optional buildDealloc(OpBuilder &builder, - Value alloc) { - return builder.create(alloc.getLoc(), alloc) - .getOperation(); - } - static std::optional buildClone(OpBuilder &builder, Value alloc) { - return builder.create(alloc.getLoc(), alloc) - .getResult(); - } - static ::mlir::HoistingKind getHoistingKind() { - return HoistingKind::Loop | HoistingKind::Block; - } - static ::std::optional<::mlir::Operation *> - buildPromotedAlloc(OpBuilder &builder, Value alloc) { - Operation *definingOp = alloc.getDefiningOp(); - return builder.create( - definingOp->getLoc(), cast(definingOp->getResultTypes()[0]), - definingOp->getOperands(), definingOp->getAttrs()); - } -}; - -struct DefaultAutomaticAllocationHoistingInterface - : public bufferization::AllocationOpInterface::ExternalModel< - DefaultAutomaticAllocationHoistingInterface, memref::AllocaOp> { - static ::mlir::HoistingKind getHoistingKind() { return HoistingKind::Loop; } -}; - -struct DefaultReallocationInterface - : public bufferization::AllocationOpInterface::ExternalModel< - DefaultAllocationInterface, memref::ReallocOp> { - static std::optional buildDealloc(OpBuilder &builder, - Value realloc) { - return builder.create(realloc.getLoc(), realloc) - .getOperation(); - } -}; - /// The actual buffer deallocation pass that inserts and moves dealloc nodes /// into the right positions. Furthermore, it inserts additional clones if /// necessary. It uses the algorithm described at the top of the file. @@ -725,16 +685,6 @@ LogicalResult bufferization::deallocateBuffers(Operation *op) { return success(); } -void bufferization::registerAllocationOpInterfaceExternalModels( - DialectRegistry ®istry) { - registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) { - memref::AllocOp::attachInterface(*ctx); - memref::AllocaOp::attachInterface< - DefaultAutomaticAllocationHoistingInterface>(*ctx); - memref::ReallocOp::attachInterface(*ctx); - }); -} - //===----------------------------------------------------------------------===// // BufferDeallocationPass construction //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp index 8fca041fe6aaa..7358d0d465d3e 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp @@ -682,3 +682,59 @@ BufferizationOptions bufferization::getPartialBufferizationOptions() { options.opFilter.allowDialect(); return options; } + +//===----------------------------------------------------------------------===// +// Default AllocationOpInterface implementation and registration +//===----------------------------------------------------------------------===// + +namespace { +struct DefaultAllocationInterface + : public bufferization::AllocationOpInterface::ExternalModel< + DefaultAllocationInterface, memref::AllocOp> { + static std::optional buildDealloc(OpBuilder &builder, + Value alloc) { + return builder.create(alloc.getLoc(), alloc) + .getOperation(); + } + static std::optional buildClone(OpBuilder &builder, Value alloc) { + return builder.create(alloc.getLoc(), alloc) + .getResult(); + } + static ::mlir::HoistingKind getHoistingKind() { + return HoistingKind::Loop | HoistingKind::Block; + } + static ::std::optional<::mlir::Operation *> + buildPromotedAlloc(OpBuilder &builder, Value alloc) { + Operation *definingOp = alloc.getDefiningOp(); + return builder.create( + definingOp->getLoc(), cast(definingOp->getResultTypes()[0]), + definingOp->getOperands(), definingOp->getAttrs()); + } +}; + +struct DefaultAutomaticAllocationHoistingInterface + : public bufferization::AllocationOpInterface::ExternalModel< + DefaultAutomaticAllocationHoistingInterface, memref::AllocaOp> { + static ::mlir::HoistingKind getHoistingKind() { return HoistingKind::Loop; } +}; + +struct DefaultReallocationInterface + : public bufferization::AllocationOpInterface::ExternalModel< + DefaultAllocationInterface, memref::ReallocOp> { + static std::optional buildDealloc(OpBuilder &builder, + Value realloc) { + return builder.create(realloc.getLoc(), realloc) + .getOperation(); + } +}; +} // namespace + +void bufferization::registerAllocationOpInterfaceExternalModels( + DialectRegistry ®istry) { + registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) { + memref::AllocOp::attachInterface(*ctx); + memref::AllocaOp::attachInterface< + DefaultAutomaticAllocationHoistingInterface>(*ctx); + memref::ReallocOp::attachInterface(*ctx); + }); +}