Skip to content

[mlir][bufferization] Move AllocationOpInterface implementations #65578

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 7, 2023
Merged

[mlir][bufferization] Move AllocationOpInterface implementations #65578

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 0 additions & 50 deletions mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -625,46 +625,6 @@ class BufferDeallocation : public BufferPlacementTransformationBase {
// BufferDeallocationPass
//===----------------------------------------------------------------------===//

struct DefaultAllocationInterface
: public bufferization::AllocationOpInterface::ExternalModel<
DefaultAllocationInterface, memref::AllocOp> {
static std::optional<Operation *> buildDealloc(OpBuilder &builder,
Value alloc) {
return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
.getOperation();
}
static std::optional<Value> buildClone(OpBuilder &builder, Value alloc) {
return builder.create<bufferization::CloneOp>(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<memref::AllocaOp>(
definingOp->getLoc(), cast<MemRefType>(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<Operation *> buildDealloc(OpBuilder &builder,
Value realloc) {
return builder.create<memref::DeallocOp>(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.
Expand Down Expand Up @@ -725,16 +685,6 @@ LogicalResult bufferization::deallocateBuffers(Operation *op) {
return success();
}

void bufferization::registerAllocationOpInterfaceExternalModels(
DialectRegistry &registry) {
registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
memref::AllocaOp::attachInterface<
DefaultAutomaticAllocationHoistingInterface>(*ctx);
memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
});
}

//===----------------------------------------------------------------------===//
// BufferDeallocationPass construction
//===----------------------------------------------------------------------===//
Expand Down
56 changes: 56 additions & 0 deletions mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -682,3 +682,59 @@ BufferizationOptions bufferization::getPartialBufferizationOptions() {
options.opFilter.allowDialect<BufferizationDialect>();
return options;
}

//===----------------------------------------------------------------------===//
// Default AllocationOpInterface implementation and registration
//===----------------------------------------------------------------------===//

namespace {
struct DefaultAllocationInterface
: public bufferization::AllocationOpInterface::ExternalModel<
DefaultAllocationInterface, memref::AllocOp> {
static std::optional<Operation *> buildDealloc(OpBuilder &builder,
Value alloc) {
return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
.getOperation();
}
static std::optional<Value> buildClone(OpBuilder &builder, Value alloc) {
return builder.create<bufferization::CloneOp>(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<memref::AllocaOp>(
definingOp->getLoc(), cast<MemRefType>(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<Operation *> buildDealloc(OpBuilder &builder,
Value realloc) {
return builder.create<memref::DeallocOp>(realloc.getLoc(), realloc)
.getOperation();
}
};
} // namespace

void bufferization::registerAllocationOpInterfaceExternalModels(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The name of the API really does not reflect that this is about the Memref dialect.

It's also not clear to me why lib/Dialect/Bufferization/Bufferize.cpp is the place to register an implementation of an interface for the Memref dialect?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wouldn't necessarily say that this is just an interface for the memref dialect. I think it's more an interface to customize various bufferization passes (e.g., buffer hoisting). In the extreme case to make them work on other dialects than just memref.

In the long-term I think it would be better to split this interface into one per bufferization pass (currently that would be one for hoisting and one for buffer promotion, the parts used for deallocation would be removed) and name them accordingly. The implementations would live in the dialect directory of the ops for which it is implemented.

In the short term I can offer to move the interface implementations to BufferOptimization.cpp where it is actually used (and not just registered). Ideally, we would move the implementation for alloc, alloca, and realloc in MemRef/Transforms/AllocationOpInterfaceImpl.cpp and the implementation of clone in Bufferization/Transforms/AllocationOpInterface.cpp, however, this is not possible because it would introduce a dependency cycle (bufferization <-> memref) unless the interface declaration is moved out of the bufferization dialect build target. But where should it go then? A new BufferizationInterfaces target?

What do you think?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wouldn't necessarily say that this is just an interface for the memref dialect. I think it's more an interface to customize various bufferization passes (e.g., buffer hoisting). In the extreme case to make them work on other dialects than just memref.

I agree that the interface isn't specific to Memref, but this is one particular implementation of the interface which is specific to the Memref dialect.

Ideally, we would move the implementation for alloc, alloca, and realloc in MemRef/Transforms/AllocationOpInterfaceImpl.cpp

Right this is what I was pointing out actually :)

however, this is not possible because it would introduce a dependency cycle (bufferization <-> memref)

Actually if it is a library of its own in the Memref dialect, there wouldn't be any cycle.
More importantly, you don't need an actually dependency to implement an interface: that is implementing AllocationOpInterfaceImpl.cpp does not need to depend on Bufferization, even if the interface is declared by Bufferization.
(see for example how mlir/IR implements a Bytecode interface but the ByteCode library depends on IR)

DialectRegistry &registry) {
registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
memref::AllocaOp::attachInterface<
DefaultAutomaticAllocationHoistingInterface>(*ctx);
memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
});
}