26 namespace bufferization {
27 #define GEN_PASS_DEF_BUFFERHOISTING
28 #define GEN_PASS_DEF_BUFFERLOOPHOISTING
29 #define GEN_PASS_DEF_PROMOTEBUFFERSTOSTACK
30 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
40 return isa<LoopLikeOpInterface, RegionBranchOpInterface>(op);
46 auto allocOp = dyn_cast<AllocationOpInterface>(op);
54 auto allocOp = dyn_cast<AllocationOpInterface>(op);
63 unsigned maxRankOfAllocatedMemRef) {
64 auto type = dyn_cast<ShapedType>(alloc.
getType());
67 if (!type.hasStaticShape()) {
73 if (type.getRank() <= maxRankOfAllocatedMemRef) {
76 return operand.getDefiningOp<memref::RankOp>();
83 return type.getNumElements() * bitwidth <= maximumSizeInBytes * 8;
90 for (
Value alias : aliases) {
91 for (
auto *use : alias.getUsers()) {
95 if (isa<RegionBranchTerminatorOpInterface>(use) &&
96 use->getParentRegion() == parentRegion)
133 struct BufferAllocationHoistingStateBase {
141 Block *placementBlock;
145 Block *placementBlock)
146 : dominators(dominators), allocValue(allocValue),
147 placementBlock(placementBlock) {}
151 template <
typename StateT>
156 postDominators(op), scopeOp(op) {}
162 allocsAndAllocas.push_back(std::get<0>(entry));
163 scopeOp->walk([&](memref::AllocaOp op) {
164 allocsAndAllocas.push_back(op.getMemref());
167 for (
auto allocValue : allocsAndAllocas) {
168 if (!StateT::shouldHoistOpType(allocValue.getDefiningOp()))
170 Operation *definingOp = allocValue.getDefiningOp();
171 assert(definingOp &&
"No defining op");
173 auto resultAliases = aliases.resolve(allocValue);
175 Block *dominatorBlock =
176 findCommonDominator(allocValue, resultAliases, dominators);
178 StateT state(&dominators, allocValue, allocValue.getParentBlock());
181 Block *dependencyBlock =
nullptr;
185 for (
Value depValue : operands) {
186 Block *depBlock = depValue.getParentBlock();
187 if (!dependencyBlock || dominators.dominates(dependencyBlock, depBlock))
188 dependencyBlock = depBlock;
194 Block *placementBlock = findPlacementBlock(
195 state, state.computeUpperBound(dominatorBlock, dependencyBlock));
197 allocValue, placementBlock, liveness);
200 Operation *allocOperation = allocValue.getDefiningOp();
209 Block *findPlacementBlock(StateT &state,
Block *upperBound) {
210 Block *currentBlock = state.placementBlock;
221 (parentBlock = parentOp->
getBlock()) &&
223 dominators.properlyDominates(upperBound, currentBlock))) {
230 idom = dominators.getNode(currentBlock)->getIDom();
232 if (idom && dominators.properlyDominates(parentBlock, idom->getBlock())) {
235 currentBlock = idom->getBlock();
236 state.recordMoveToDominator(currentBlock);
244 !state.isLegalPlacement(parentOp))
248 currentBlock = parentBlock;
249 state.recordMoveToParent(currentBlock);
253 return state.placementBlock;
275 struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
276 using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
279 Block *computeUpperBound(
Block *dominatorBlock,
Block *dependencyBlock) {
282 if (!dependencyBlock)
283 return dominatorBlock;
287 return dominators->properlyDominates(dominatorBlock, dependencyBlock)
298 static bool shouldHoistOpType(
Operation *op) {
303 void recordMoveToDominator(
Block *block) { placementBlock = block; }
306 void recordMoveToParent(
Block *block) { recordMoveToDominator(block); }
311 struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
312 using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
315 Block *aliasDominatorBlock =
nullptr;
318 Block *computeUpperBound(
Block *dominatorBlock,
Block *dependencyBlock) {
319 aliasDominatorBlock = dominatorBlock;
322 return dependencyBlock ? dependencyBlock :
nullptr;
331 !dominators->dominates(aliasDominatorBlock, op->
getBlock());
335 static bool shouldHoistOpType(
Operation *op) {
341 void recordMoveToDominator(
Block *block) {}
344 void recordMoveToParent(
Block *block) { placementBlock = block; }
360 Value alloc = std::get<0>(entry);
366 if (!isSmallAlloc(alloc) || dealloc ||
376 if (
auto allocInterface = dyn_cast<AllocationOpInterface>(allocOp)) {
378 allocInterface.buildPromotedAlloc(builder, alloc).value();
395 struct BufferHoistingPass
396 :
public bufferization::impl::BufferHoistingBase<BufferHoistingPass> {
398 void runOnOperation()
override {
400 BufferAllocationHoisting<BufferAllocationHoistingState> optimizer(
407 struct BufferLoopHoistingPass
408 :
public bufferization::impl::BufferLoopHoistingBase<
409 BufferLoopHoistingPass> {
411 void runOnOperation()
override {
419 class PromoteBuffersToStackPass
420 :
public bufferization::impl::PromoteBuffersToStackBase<
421 PromoteBuffersToStackPass> {
423 PromoteBuffersToStackPass(
unsigned maxAllocSizeInBytes,
424 unsigned maxRankOfAllocatedMemRef) {
425 this->maxAllocSizeInBytes = maxAllocSizeInBytes;
426 this->maxRankOfAllocatedMemRef = maxRankOfAllocatedMemRef;
429 explicit PromoteBuffersToStackPass(std::function<
bool(
Value)> isSmallAlloc)
430 : isSmallAlloc(std::move(isSmallAlloc)) {}
433 if (isSmallAlloc ==
nullptr) {
434 isSmallAlloc = [=](
Value alloc) {
436 maxRankOfAllocatedMemRef);
442 void runOnOperation()
override {
444 BufferPlacementPromotion optimizer(getOperation());
445 optimizer.promote(isSmallAlloc);
449 std::function<bool(
Value)> isSmallAlloc;
455 BufferAllocationHoisting<BufferAllocationLoopHoistingState> optimizer(op);
460 return std::make_unique<BufferHoistingPass>();
464 return std::make_unique<BufferLoopHoistingPass>();
468 unsigned maxAllocSizeInBytes,
unsigned maxRankOfAllocatedMemRef) {
469 return std::make_unique<PromoteBuffersToStackPass>(maxAllocSizeInBytes,
470 maxRankOfAllocatedMemRef);
474 std::function<
bool(
Value)> isSmallAlloc) {
475 return std::make_unique<PromoteBuffersToStackPass>(std::move(isSmallAlloc));
static bool leavesAllocationScope(Region *parentRegion, const BufferViewFlowAnalysis::ValueSetT &aliases)
Checks whether the given aliases leave the allocation scope.
static bool isKnownControlFlowInterface(Operation *op)
Returns true if the given operation implements a known high-level region- based control-flow interfac...
static bool hasAllocationScope(Value alloc, const BufferViewFlowAnalysis &aliasAnalysis)
Checks, if an automated allocation scope for a given alloc value exists.
static bool allowAllocDominateBlockHoisting(Operation *op)
Returns true if the given operation implements the AllocationOpInterface and it supports the dominate...
static bool allowAllocLoopHoisting(Operation *op)
Returns true if the given operation implements the AllocationOpInterface and it supports the loop hoi...
static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes, unsigned maxRankOfAllocatedMemRef)
Check if the size of the allocation is less than the given size.
Block represents an ordered list of Operations.
bool isEntryBlock()
Return if this block is the entry block in the parent region.
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
A straight-forward alias analysis which ensures that all dependencies of all values will be determine...
ValueSetT resolve(Value value) const
Find all immediate and indirect views upon this value.
static DataLayout closest(Operation *op)
Returns the layout of the closest parent operation carrying layout info.
llvm::TypeSize getTypeSizeInBits(Type t) const
Returns the size in bits of the given type in the current scope.
A class for computing basic dominance information.
MLIRContext is the top-level object for a collection of MLIR operations.
This class helps build Operations.
A trait of region holding operations that define a new scope for automatic allocations,...
Operation is the basic unit of execution within MLIR.
Block * getBlock()
Returns the operation block that contains this operation.
operand_range getOperands()
Returns an iterator on the underlying Value's.
void moveBefore(Operation *existingOp)
Unlink this operation from its current block and insert it right before existingOp which may be in th...
void replaceAllUsesWith(ValuesT &&values)
Replace all uses of results of this operation with the provided 'values'.
void erase()
Remove this operation from its parent block and delete it.
A class for computing basic postdominance information.
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Region * getParentRegion()
Return the region containing this region or nullptr if the region is attached to a top-level operatio...
Operation * getParentOp()
Return the parent operation this region is attached to.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
Block * getParentBlock()
Return the Block in which this Value is defined.
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Region * getParentRegion()
Return the Region in which this Value is defined.
static Operation * getStartOperation(Value allocValue, Block *placementBlock, const Liveness &liveness)
Get the start operation to place the given alloc value within the specified placement block.
std::tuple< Value, Operation * > AllocEntry
Represents a tuple of allocValue and deallocOperation.
void hoistBuffersFromLoops(Operation *op)
Within the given operation, hoist buffers from loops where possible.
std::unique_ptr< Pass > createBufferHoistingPass()
Creates a pass that moves allocations upwards to reduce the number of required copies that are insert...
std::unique_ptr< Pass > createBufferLoopHoistingPass()
Creates a pass that moves allocations upwards out of loops.
std::unique_ptr< Pass > createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes=1024, unsigned maxRankOfAllocatedMemRef=1)
Creates a pass that promotes heap-based allocations to stack-based ones.
void promote(RewriterBase &rewriter, scf::ForallOp forallOp)
Promotes the loop body of a scf::ForallOp to its containing block.
Include the generated interface declarations.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
llvm::DomTreeNodeBase< Block > DominanceInfoNode
This class represents an efficient way to signal success or failure.