[Mlir-commits] [mlir] 7f153e8 - [mlir][linalg][bufferize] Add `isAllocationHoistingBarrier` to op interface
Matthias Springer
llvmlistbot at llvm.org
Wed Nov 10 18:00:58 PST 2021
Author: Matthias Springer
Date: 2021-11-11T11:00:47+09:00
New Revision: 7f153e8ba103bc9489ecdf09b5739ff262e2c54c
URL: https://github.com/llvm/llvm-project/commit/7f153e8ba103bc9489ecdf09b5739ff262e2c54c
DIFF: https://github.com/llvm/llvm-project/commit/7f153e8ba103bc9489ecdf09b5739ff262e2c54c.diff
LOG: [mlir][linalg][bufferize] Add `isAllocationHoistingBarrier` to op interface
This make `getResultBuffer` in ComprehensiveBufferize independent of the SCF, Affine and Linalg dialects. This commit is in preparating of decoupling op interface implementations from ComprehensiveBufferize.
Differential Revision: https://reviews.llvm.org/D113380
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h
mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h
index 2889040f7593..318819cf49a7 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h
@@ -166,4 +166,57 @@ BufferRelation bufferRelation(OpOperand &opOperand);
#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc"
+namespace mlir {
+namespace linalg {
+namespace comprehensive_bufferize {
+
+/// AllocationHoistingBarrierOnly is an external implementation of
+/// BufferizableOpInterface for ops that are (not yet) bufferizable, but are
+/// known to be allocation hoisting barriers. All interface methods (except for
+/// `isAllocationHoistingBarrier`) are implemented conservatively.
+template <typename OpTy>
+struct AllocationHoistingBarrierOnly
+ : public BufferizableOpInterface::ExternalModel<
+ AllocationHoistingBarrierOnly<OpTy>, OpTy> {
+ bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
+ return true;
+ }
+
+ bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
+ return false;
+ }
+
+ SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
+ OpResult opResult) const {
+ return {};
+ }
+
+ OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
+ return OpResult();
+ }
+
+ BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
+ return BufferRelation::None;
+ }
+
+ bool isWritable(Operation *op, Value value) const { return false; }
+
+ LogicalResult bufferize(Operation *op, OpBuilder &b,
+ BlockAndValueMapping &bvm,
+ BufferizationAliasInfo &aliasInfo,
+ AllocationCallbacks &allocationFn) const {
+ auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
+ if (any_of(op->getOperandTypes(), isaTensor) ||
+ any_of(op->getResultTypes(), isaTensor))
+ return op->emitError() << "unsupported op with tensors";
+ return success();
+ }
+
+ bool isAllocationHoistingBarrier(Operation *op) const { return true; }
+};
+
+} // namespace comprehensive_bufferize
+} // namespace linalg
+} // namespace mlir
+
#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_
diff --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
index fa97e2a06b81..8c8846753304 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td
@@ -203,6 +203,22 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
/*defaultImplementation=*/[{
return value.isa<OpResult>();
}]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return `true` if the op is an allocation hoisting barrier. Buffer
+ allocations will never be beyond such ops. E.g., ops with certain
+ parallel semantics may be allocation hoisting barriers. The majority
+ of ops, however, is not a barrier. Therefore, this method returns
+ `false` by default.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"isAllocationHoistingBarrier",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
>
];
diff --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index 433edf1b94f1..6e059c3a28b3 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
@@ -873,21 +873,6 @@ static FunctionType getOrCreateBufferizedFunctionType(
// Bufferization-specific scoped alloc/dealloc insertion support.
//===----------------------------------------------------------------------===//
-template <typename... Args>
-Operation *getFirstParentOfType(Value v) {
- Operation *parent;
- if (auto bbArg = v.dyn_cast<BlockArgument>())
- parent = bbArg.getOwner()->getParentOp();
- else
- parent = v.getDefiningOp()->getParentOp();
- while (parent) {
- if (isa<Args...>(parent))
- return parent;
- parent = parent->getParentOp();
- }
- return nullptr;
-}
-
/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
/// the type of `source`.
static Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
@@ -939,17 +924,31 @@ static MemRefType getAllocationTypeAndShape(OpBuilder &b, Location loc,
// If the buffer is statically shaped, try to hoist it to the first enclosing
// parallel region.
- // TODO: this concept of parallel region and threadlocal needs interfaces.
// TODO: also hoist in the dynamic case. For now this relies on subsequent
// calls to LICM and buffer hoisting which will most likely not succeed.
// TODO: when packing, allocate a static bounding box which will enable more
// hoisting.
if (dynShape.empty()) {
- Operation *parent =
- getFirstParentOfType<FuncOp, TiledLoopOp, scf::ParallelOp,
- AffineParallelOp>(shapedValue);
- if (parent)
- b.setInsertionPointToStart(&(parent->getRegion(0).front()));
+ Operation *parent;
+ if (auto bbArg = shapedValue.dyn_cast<BlockArgument>())
+ parent = bbArg.getOwner()->getParentOp();
+ else
+ parent = shapedValue.getDefiningOp()->getParentOp();
+ while (parent) {
+ if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(parent))
+ if (bufferizableOp.isAllocationHoistingBarrier())
+ break;
+ parent = parent->getParentOp();
+ }
+
+ // FuncOp is an allocation hoisting barrier, so the above loop should never
+ // run out of parents.
+ assert(
+ (parent &&
+ cast<BufferizableOpInterface>(parent).isAllocationHoistingBarrier()) &&
+ "expected traversal to end at allocation hoisting barrier");
+
+ b.setInsertionPointToStart(&(parent->getRegion(0).front()));
}
return allocMemRefType;
}
@@ -2294,6 +2293,8 @@ struct TiledLoopOpInterface
return true;
}
+ bool isAllocationHoistingBarrier(Operation *op) const { return true; }
+
LogicalResult bufferize(Operation *op, OpBuilder &b,
BlockAndValueMapping &bvm,
BufferizationAliasInfo &aliasInfo,
@@ -3278,6 +3279,13 @@ void registerBufferizableOpInterfaceExternalModels(DialectRegistry ®istry) {
registry.addOpInterface<vector::TransferWriteOp,
vector_ext::TransferWriteOpInterface>();
+ // Ops that are not bufferizable but are allocation hoisting barriers.
+ registry.addOpInterface<FuncOp, AllocationHoistingBarrierOnly<FuncOp>>();
+ registry.addOpInterface<scf::ParallelOp,
+ AllocationHoistingBarrierOnly<scf::ParallelOp>>();
+ registry.addOpInterface<AffineParallelOp,
+ AllocationHoistingBarrierOnly<AffineParallelOp>>();
+
// Register all Linalg structured ops. `LinalgOp` is an interface and it is
// not possible to attach an external interface to an existing interface.
// Therefore, attach the `BufferizableOpInterface` to all ops one-by-one.
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp
index e965cdf89609..d95ba926f81e 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp
@@ -33,7 +33,7 @@ struct LinalgComprehensiveModuleBufferize
registry
.insert<linalg::LinalgDialect, memref::MemRefDialect,
tensor::TensorDialect, vector::VectorDialect, scf::SCFDialect,
- arith::ArithmeticDialect, StandardOpsDialect>();
+ arith::ArithmeticDialect, StandardOpsDialect, AffineDialect>();
registerBufferizableOpInterfaceExternalModels(registry);
}
};
More information about the Mlir-commits
mailing list