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);
50 if (isa<LoopLikeOpInterface>(op))
55 auto regionInterface = dyn_cast<RegionBranchOpInterface>(op);
59 return regionInterface.hasLoop();
65 auto allocOp = dyn_cast<AllocationOpInterface>(op);
73 auto allocOp = dyn_cast<AllocationOpInterface>(op);
82 unsigned maxRankOfAllocatedMemRef) {
83 auto type = dyn_cast<ShapedType>(alloc.
getType());
86 if (!type.hasStaticShape()) {
92 if (type.getRank() <= maxRankOfAllocatedMemRef) {
95 return operand.getDefiningOp<memref::RankOp>();
102 return type.getNumElements() * bitwidth <= maximumSizeInBytes * 8;
109 for (
Value alias : aliases) {
110 for (
auto *use : alias.getUsers()) {
114 if (isa<RegionBranchTerminatorOpInterface>(use) &&
115 use->getParentRegion() == parentRegion)
151 struct BufferAllocationHoistingStateBase {
159 Block *placementBlock;
163 Block *placementBlock)
164 : dominators(dominators), allocValue(allocValue),
165 placementBlock(placementBlock) {}
169 template <
typename StateT>
174 postDominators(op), scopeOp(op) {}
180 allocsAndAllocas.push_back(std::get<0>(entry));
181 scopeOp->walk([&](memref::AllocaOp op) {
182 allocsAndAllocas.push_back(op.getMemref());
185 for (
auto allocValue : allocsAndAllocas) {
186 if (!StateT::shouldHoistOpType(allocValue.getDefiningOp()))
188 Operation *definingOp = allocValue.getDefiningOp();
189 assert(definingOp &&
"No defining op");
191 auto resultAliases = aliases.resolve(allocValue);
193 Block *dominatorBlock =
196 StateT state(&dominators, allocValue, allocValue.getParentBlock());
199 Block *dependencyBlock =
nullptr;
203 for (
Value depValue : operands) {
204 Block *depBlock = depValue.getParentBlock();
205 if (!dependencyBlock || dominators.dominates(dependencyBlock, depBlock))
206 dependencyBlock = depBlock;
212 Block *placementBlock = findPlacementBlock(
213 state, state.computeUpperBound(dominatorBlock, dependencyBlock));
215 allocValue, placementBlock, liveness);
218 Operation *allocOperation = allocValue.getDefiningOp();
227 Block *findPlacementBlock(StateT &state,
Block *upperBound) {
228 Block *currentBlock = state.placementBlock;
239 (parentBlock = parentOp->
getBlock()) &&
241 dominators.properlyDominates(upperBound, currentBlock))) {
248 idom = dominators.getNode(currentBlock)->getIDom();
250 if (idom && dominators.properlyDominates(parentBlock, idom->getBlock())) {
253 currentBlock = idom->getBlock();
254 state.recordMoveToDominator(currentBlock);
262 !state.isLegalPlacement(parentOp))
266 currentBlock = parentBlock;
267 state.recordMoveToParent(currentBlock);
271 return state.placementBlock;
293 struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
294 using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
297 Block *computeUpperBound(
Block *dominatorBlock,
Block *dependencyBlock) {
300 if (!dependencyBlock)
301 return dominatorBlock;
305 return dominators->properlyDominates(dominatorBlock, dependencyBlock)
314 static bool shouldHoistOpType(
Operation *op) {
319 void recordMoveToDominator(
Block *block) { placementBlock = block; }
322 void recordMoveToParent(
Block *block) { recordMoveToDominator(block); }
327 struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
328 using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
331 Block *aliasDominatorBlock =
nullptr;
334 Block *computeUpperBound(
Block *dominatorBlock,
Block *dependencyBlock) {
335 aliasDominatorBlock = dominatorBlock;
338 return dependencyBlock ? dependencyBlock :
nullptr;
347 !dominators->dominates(aliasDominatorBlock, op->
getBlock());
351 static bool shouldHoistOpType(
Operation *op) {
357 void recordMoveToDominator(
Block *block) {}
360 void recordMoveToParent(
Block *block) { placementBlock = block; }
376 Value alloc = std::get<0>(entry);
382 if (!isSmallAlloc(alloc) || dealloc ||
392 if (
auto allocInterface = dyn_cast<AllocationOpInterface>(allocOp)) {
394 allocInterface.buildPromotedAlloc(builder, alloc).value();
411 struct BufferHoistingPass
412 :
public bufferization::impl::BufferHoistingBase<BufferHoistingPass> {
414 void runOnOperation()
override {
416 BufferAllocationHoisting<BufferAllocationHoistingState> optimizer(
423 struct BufferLoopHoistingPass
424 :
public bufferization::impl::BufferLoopHoistingBase<
425 BufferLoopHoistingPass> {
427 void runOnOperation()
override {
435 class PromoteBuffersToStackPass
436 :
public bufferization::impl::PromoteBuffersToStackBase<
437 PromoteBuffersToStackPass> {
439 PromoteBuffersToStackPass(
unsigned maxAllocSizeInBytes,
440 unsigned maxRankOfAllocatedMemRef) {
441 this->maxAllocSizeInBytes = maxAllocSizeInBytes;
442 this->maxRankOfAllocatedMemRef = maxRankOfAllocatedMemRef;
445 explicit PromoteBuffersToStackPass(std::function<
bool(
Value)> isSmallAlloc)
446 : isSmallAlloc(std::move(isSmallAlloc)) {}
449 if (isSmallAlloc ==
nullptr) {
450 isSmallAlloc = [=](
Value alloc) {
452 maxRankOfAllocatedMemRef);
458 void runOnOperation()
override {
460 BufferPlacementPromotion optimizer(getOperation());
461 optimizer.promote(isSmallAlloc);
465 std::function<bool(
Value)> isSmallAlloc;
471 BufferAllocationHoisting<BufferAllocationLoopHoistingState> optimizer(op);
476 return std::make_unique<BufferHoistingPass>();
480 return std::make_unique<BufferLoopHoistingPass>();
484 unsigned maxAllocSizeInBytes,
unsigned maxRankOfAllocatedMemRef) {
485 return std::make_unique<PromoteBuffersToStackPass>(maxAllocSizeInBytes,
486 maxRankOfAllocatedMemRef);
490 std::function<
bool(
Value)> isSmallAlloc) {
491 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 isLoop(Operation *op)
Returns true if the given operation represents a loop by testing whether it implements the LoopLikeOp...
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...
Block * findCommonDominator(Value value, const BufferViewFlowAnalysis::ValueSetT &values, const DominatorT &doms)
Finds a common dominator for the given value while taking the positions of the values in the value se...
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.