1 //===- BufferOptimizations.cpp - pre-pass optimizations for bufferization -===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements logic for three optimization passes. The first two
10 // passes try to move alloc nodes out of blocks to reduce the number of
11 // allocations and copies during buffer deallocation. The third pass tries to
12 // convert heap-based allocations to stack-based allocations, if possible.
13
14 #include "PassDetail.h"
15 #include "mlir/Dialect/MemRef/IR/MemRef.h"
16 #include "mlir/IR/Operation.h"
17 #include "mlir/Interfaces/LoopLikeInterface.h"
18 #include "mlir/Pass/Pass.h"
19 #include "mlir/Transforms/BufferUtils.h"
20 #include "mlir/Transforms/Passes.h"
21
22 using namespace mlir;
23
24 /// Returns true if the given operation implements a known high-level region-
25 /// based control-flow interface.
isKnownControlFlowInterface(Operation * op)26 static bool isKnownControlFlowInterface(Operation *op) {
27 return isa<LoopLikeOpInterface, RegionBranchOpInterface>(op);
28 }
29
30 /// Check if the size of the allocation is less than the given size. The
31 /// transformation is only applied to small buffers since large buffers could
32 /// exceed the stack space.
defaultIsSmallAlloc(Value alloc,unsigned maximumSizeInBytes,unsigned bitwidthOfIndexType,unsigned maxRankOfAllocatedMemRef)33 static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
34 unsigned bitwidthOfIndexType,
35 unsigned maxRankOfAllocatedMemRef) {
36 auto type = alloc.getType().dyn_cast<ShapedType>();
37 if (!type || !alloc.getDefiningOp<memref::AllocOp>())
38 return false;
39 if (!type.hasStaticShape()) {
40 // Check if the dynamic shape dimension of the alloc is produced by RankOp.
41 // If this is the case, it is likely to be small. Furthermore, the dimension
42 // is limited to the maximum rank of the allocated memref to avoid large
43 // values by multiplying several small values.
44 if (type.getRank() <= maxRankOfAllocatedMemRef) {
45 return llvm::all_of(
46 alloc.getDefiningOp()->getOperands(),
47 [&](Value operand) { return operand.getDefiningOp<RankOp>(); });
48 }
49 return false;
50 }
51 // For index types, use the provided size, as the type does not know.
52 unsigned int bitwidth = type.getElementType().isIndex()
53 ? bitwidthOfIndexType
54 : type.getElementTypeBitWidth();
55 return type.getNumElements() * bitwidth <= maximumSizeInBytes * 8;
56 }
57
58 /// Checks whether the given aliases leave the allocation scope.
59 static bool
leavesAllocationScope(Region * parentRegion,const BufferViewFlowAnalysis::ValueSetT & aliases)60 leavesAllocationScope(Region *parentRegion,
61 const BufferViewFlowAnalysis::ValueSetT &aliases) {
62 for (Value alias : aliases) {
63 for (auto *use : alias.getUsers()) {
64 // If there is at least one alias that leaves the parent region, we know
65 // that this alias escapes the whole region and hence the associated
66 // allocation leaves allocation scope.
67 if (isRegionReturnLike(use) && use->getParentRegion() == parentRegion)
68 return true;
69 }
70 }
71 return false;
72 }
73
74 /// Checks, if an automated allocation scope for a given alloc value exists.
hasAllocationScope(Value alloc,const BufferViewFlowAnalysis & aliasAnalysis)75 static bool hasAllocationScope(Value alloc,
76 const BufferViewFlowAnalysis &aliasAnalysis) {
77 Region *region = alloc.getParentRegion();
78 do {
79 if (Operation *parentOp = region->getParentOp()) {
80 // Check if the operation is an automatic allocation scope and whether an
81 // alias leaves the scope. This means, an allocation yields out of
82 // this scope and can not be transformed in a stack-based allocation.
83 if (parentOp->hasTrait<OpTrait::AutomaticAllocationScope>() &&
84 !leavesAllocationScope(region, aliasAnalysis.resolve(alloc)))
85 return true;
86 // Check if the operation is a known control flow interface and break the
87 // loop to avoid transformation in loops. Furthermore skip transformation
88 // if the operation does not implement a RegionBeanchOpInterface.
89 if (BufferPlacementTransformationBase::isLoop(parentOp) ||
90 !isKnownControlFlowInterface(parentOp))
91 break;
92 }
93 } while ((region = region->getParentRegion()));
94 return false;
95 }
96
97 namespace {
98
99 //===----------------------------------------------------------------------===//
100 // BufferAllocationHoisting
101 //===----------------------------------------------------------------------===//
102
103 /// A base implementation compatible with the `BufferAllocationHoisting` class.
104 struct BufferAllocationHoistingStateBase {
105 /// A pointer to the current dominance info.
106 DominanceInfo *dominators;
107
108 /// The current allocation value.
109 Value allocValue;
110
111 /// The current placement block (if any).
112 Block *placementBlock;
113
114 /// Initializes the state base.
BufferAllocationHoistingStateBase__anon1c2f66b40211::BufferAllocationHoistingStateBase115 BufferAllocationHoistingStateBase(DominanceInfo *dominators, Value allocValue,
116 Block *placementBlock)
117 : dominators(dominators), allocValue(allocValue),
118 placementBlock(placementBlock) {}
119 };
120
121 /// Implements the actual hoisting logic for allocation nodes.
122 template <typename StateT>
123 class BufferAllocationHoisting : public BufferPlacementTransformationBase {
124 public:
BufferAllocationHoisting(Operation * op)125 BufferAllocationHoisting(Operation *op)
126 : BufferPlacementTransformationBase(op), dominators(op),
127 postDominators(op), scopeOp(op) {}
128
129 /// Moves allocations upwards.
hoist()130 void hoist() {
131 SmallVector<Value> allocsAndAllocas;
132 for (BufferPlacementAllocs::AllocEntry &entry : allocs)
133 allocsAndAllocas.push_back(std::get<0>(entry));
134 scopeOp->walk(
135 [&](memref::AllocaOp op) { allocsAndAllocas.push_back(op.memref()); });
136
137 for (auto allocValue : allocsAndAllocas) {
138 if (!StateT::shouldHoistOpType(allocValue.getDefiningOp()))
139 continue;
140 Operation *definingOp = allocValue.getDefiningOp();
141 assert(definingOp && "No defining op");
142 auto operands = definingOp->getOperands();
143 auto resultAliases = aliases.resolve(allocValue);
144 // Determine the common dominator block of all aliases.
145 Block *dominatorBlock =
146 findCommonDominator(allocValue, resultAliases, dominators);
147 // Init the initial hoisting state.
148 StateT state(&dominators, allocValue, allocValue.getParentBlock());
149 // Check for additional allocation dependencies to compute an upper bound
150 // for hoisting.
151 Block *dependencyBlock = nullptr;
152 // If this node has dependencies, check all dependent nodes. This ensures
153 // that all dependency values have been computed before allocating the
154 // buffer.
155 for (Value depValue : operands) {
156 Block *depBlock = depValue.getParentBlock();
157 if (!dependencyBlock || dominators.dominates(dependencyBlock, depBlock))
158 dependencyBlock = depBlock;
159 }
160
161 // Find the actual placement block and determine the start operation using
162 // an upper placement-block boundary. The idea is that placement block
163 // cannot be moved any further upwards than the given upper bound.
164 Block *placementBlock = findPlacementBlock(
165 state, state.computeUpperBound(dominatorBlock, dependencyBlock));
166 Operation *startOperation = BufferPlacementAllocs::getStartOperation(
167 allocValue, placementBlock, liveness);
168
169 // Move the alloc in front of the start operation.
170 Operation *allocOperation = allocValue.getDefiningOp();
171 allocOperation->moveBefore(startOperation);
172 }
173 }
174
175 private:
176 /// Finds a valid placement block by walking upwards in the CFG until we
177 /// either cannot continue our walk due to constraints (given by the StateT
178 /// implementation) or we have reached the upper-most dominator block.
findPlacementBlock(StateT & state,Block * upperBound)179 Block *findPlacementBlock(StateT &state, Block *upperBound) {
180 Block *currentBlock = state.placementBlock;
181 // Walk from the innermost regions/loops to the outermost regions/loops and
182 // find an appropriate placement block that satisfies the constraint of the
183 // current StateT implementation. Walk until we reach the upperBound block
184 // (if any).
185
186 // If we are not able to find a valid parent operation or an associated
187 // parent block, break the walk loop.
188 Operation *parentOp;
189 Block *parentBlock;
190 while ((parentOp = currentBlock->getParentOp()) &&
191 (parentBlock = parentOp->getBlock()) &&
192 (!upperBound ||
193 dominators.properlyDominates(upperBound, currentBlock))) {
194 // Try to find an immediate dominator and check whether the parent block
195 // is above the immediate dominator (if any).
196 DominanceInfoNode *idom = nullptr;
197
198 // DominanceInfo doesn't support getNode queries for single-block regions.
199 if (!currentBlock->isEntryBlock())
200 idom = dominators.getNode(currentBlock)->getIDom();
201
202 if (idom && dominators.properlyDominates(parentBlock, idom->getBlock())) {
203 // If the current immediate dominator is below the placement block, move
204 // to the immediate dominator block.
205 currentBlock = idom->getBlock();
206 state.recordMoveToDominator(currentBlock);
207 } else {
208 // We have to move to our parent block since an immediate dominator does
209 // either not exist or is above our parent block. If we cannot move to
210 // our parent operation due to constraints given by the StateT
211 // implementation, break the walk loop. Furthermore, we should not move
212 // allocations out of unknown region-based control-flow operations.
213 if (!isKnownControlFlowInterface(parentOp) ||
214 !state.isLegalPlacement(parentOp))
215 break;
216 // Move to our parent block by notifying the current StateT
217 // implementation.
218 currentBlock = parentBlock;
219 state.recordMoveToParent(currentBlock);
220 }
221 }
222 // Return the finally determined placement block.
223 return state.placementBlock;
224 }
225
226 /// The dominator info to find the appropriate start operation to move the
227 /// allocs.
228 DominanceInfo dominators;
229
230 /// The post dominator info to move the dependent allocs in the right
231 /// position.
232 PostDominanceInfo postDominators;
233
234 /// The map storing the final placement blocks of a given alloc value.
235 llvm::DenseMap<Value, Block *> placementBlocks;
236
237 /// The operation that this transformation is working on. It is used to also
238 /// gather allocas.
239 Operation *scopeOp;
240 };
241
242 /// A state implementation compatible with the `BufferAllocationHoisting` class
243 /// that hoists allocations into dominator blocks while keeping them inside of
244 /// loops.
245 struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
246 using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
247
248 /// Computes the upper bound for the placement block search.
computeUpperBound__anon1c2f66b40211::BufferAllocationHoistingState249 Block *computeUpperBound(Block *dominatorBlock, Block *dependencyBlock) {
250 // If we do not have a dependency block, the upper bound is given by the
251 // dominator block.
252 if (!dependencyBlock)
253 return dominatorBlock;
254
255 // Find the "lower" block of the dominator and the dependency block to
256 // ensure that we do not move allocations above this block.
257 return dominators->properlyDominates(dominatorBlock, dependencyBlock)
258 ? dependencyBlock
259 : dominatorBlock;
260 }
261
262 /// Returns true if the given operation does not represent a loop.
isLegalPlacement__anon1c2f66b40211::BufferAllocationHoistingState263 bool isLegalPlacement(Operation *op) {
264 return !BufferPlacementTransformationBase::isLoop(op);
265 }
266
267 /// Returns true if the given operation should be considered for hoisting.
shouldHoistOpType__anon1c2f66b40211::BufferAllocationHoistingState268 static bool shouldHoistOpType(Operation *op) {
269 return llvm::isa<memref::AllocOp>(op);
270 }
271
272 /// Sets the current placement block to the given block.
recordMoveToDominator__anon1c2f66b40211::BufferAllocationHoistingState273 void recordMoveToDominator(Block *block) { placementBlock = block; }
274
275 /// Sets the current placement block to the given block.
recordMoveToParent__anon1c2f66b40211::BufferAllocationHoistingState276 void recordMoveToParent(Block *block) { recordMoveToDominator(block); }
277 };
278
279 /// A state implementation compatible with the `BufferAllocationHoisting` class
280 /// that hoists allocations out of loops.
281 struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
282 using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
283
284 /// Remembers the dominator block of all aliases.
285 Block *aliasDominatorBlock;
286
287 /// Computes the upper bound for the placement block search.
computeUpperBound__anon1c2f66b40211::BufferAllocationLoopHoistingState288 Block *computeUpperBound(Block *dominatorBlock, Block *dependencyBlock) {
289 aliasDominatorBlock = dominatorBlock;
290 // If there is a dependency block, we have to use this block as an upper
291 // bound to satisfy all allocation value dependencies.
292 return dependencyBlock ? dependencyBlock : nullptr;
293 }
294
295 /// Returns true if the given operation represents a loop and one of the
296 /// aliases caused the `aliasDominatorBlock` to be "above" the block of the
297 /// given loop operation. If this is the case, it indicates that the
298 /// allocation is passed via a back edge.
isLegalPlacement__anon1c2f66b40211::BufferAllocationLoopHoistingState299 bool isLegalPlacement(Operation *op) {
300 return BufferPlacementTransformationBase::isLoop(op) &&
301 !dominators->dominates(aliasDominatorBlock, op->getBlock());
302 }
303
304 /// Returns true if the given operation should be considered for hoisting.
shouldHoistOpType__anon1c2f66b40211::BufferAllocationLoopHoistingState305 static bool shouldHoistOpType(Operation *op) {
306 return llvm::isa<memref::AllocOp, memref::AllocaOp>(op);
307 }
308
309 /// Does not change the internal placement block, as we want to move
310 /// operations out of loops only.
recordMoveToDominator__anon1c2f66b40211::BufferAllocationLoopHoistingState311 void recordMoveToDominator(Block *block) {}
312
313 /// Sets the current placement block to the given block.
recordMoveToParent__anon1c2f66b40211::BufferAllocationLoopHoistingState314 void recordMoveToParent(Block *block) { placementBlock = block; }
315 };
316
317 //===----------------------------------------------------------------------===//
318 // BufferPlacementPromotion
319 //===----------------------------------------------------------------------===//
320
321 /// Promotes heap-based allocations to stack-based allocations (if possible).
322 class BufferPlacementPromotion : BufferPlacementTransformationBase {
323 public:
BufferPlacementPromotion(Operation * op)324 BufferPlacementPromotion(Operation *op)
325 : BufferPlacementTransformationBase(op) {}
326
327 /// Promote buffers to stack-based allocations.
promote(function_ref<bool (Value)> isSmallAlloc)328 void promote(function_ref<bool(Value)> isSmallAlloc) {
329 for (BufferPlacementAllocs::AllocEntry &entry : allocs) {
330 Value alloc = std::get<0>(entry);
331 Operation *dealloc = std::get<1>(entry);
332 // Checking several requirements to transform an AllocOp into an AllocaOp.
333 // The transformation is done if the allocation is limited to a given
334 // size. Furthermore, a deallocation must not be defined for this
335 // allocation entry and a parent allocation scope must exist.
336 if (!isSmallAlloc(alloc) || dealloc ||
337 !hasAllocationScope(alloc, aliases))
338 continue;
339
340 Operation *startOperation = BufferPlacementAllocs::getStartOperation(
341 alloc, alloc.getParentBlock(), liveness);
342 // Build a new alloca that is associated with its parent
343 // `AutomaticAllocationScope` determined during the initialization phase.
344 OpBuilder builder(startOperation);
345 Operation *allocOp = alloc.getDefiningOp();
346 Operation *alloca = builder.create<memref::AllocaOp>(
347 alloc.getLoc(), alloc.getType().cast<MemRefType>(),
348 allocOp->getOperands());
349
350 // Replace the original alloc by a newly created alloca.
351 allocOp->replaceAllUsesWith(alloca);
352 allocOp->erase();
353 }
354 }
355 };
356
357 //===----------------------------------------------------------------------===//
358 // BufferOptimizationPasses
359 //===----------------------------------------------------------------------===//
360
361 /// The buffer hoisting pass that hoists allocation nodes into dominating
362 /// blocks.
363 struct BufferHoistingPass : BufferHoistingBase<BufferHoistingPass> {
364
runOnFunction__anon1c2f66b40211::BufferHoistingPass365 void runOnFunction() override {
366 // Hoist all allocations into dominator blocks.
367 BufferAllocationHoisting<BufferAllocationHoistingState> optimizer(
368 getFunction());
369 optimizer.hoist();
370 }
371 };
372
373 /// The buffer loop hoisting pass that hoists allocation nodes out of loops.
374 struct BufferLoopHoistingPass : BufferLoopHoistingBase<BufferLoopHoistingPass> {
375
runOnFunction__anon1c2f66b40211::BufferLoopHoistingPass376 void runOnFunction() override {
377 // Hoist all allocations out of loops.
378 BufferAllocationHoisting<BufferAllocationLoopHoistingState> optimizer(
379 getFunction());
380 optimizer.hoist();
381 }
382 };
383
384 /// The promote buffer to stack pass that tries to convert alloc nodes into
385 /// alloca nodes.
386 class PromoteBuffersToStackPass
387 : public PromoteBuffersToStackBase<PromoteBuffersToStackPass> {
388 public:
PromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,unsigned bitwidthOfIndexType,unsigned maxRankOfAllocatedMemRef)389 PromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
390 unsigned bitwidthOfIndexType,
391 unsigned maxRankOfAllocatedMemRef) {
392 this->maxAllocSizeInBytes = maxAllocSizeInBytes;
393 this->bitwidthOfIndexType = bitwidthOfIndexType;
394 this->maxRankOfAllocatedMemRef = maxRankOfAllocatedMemRef;
395 }
396
PromoteBuffersToStackPass(std::function<bool (Value)> isSmallAlloc)397 explicit PromoteBuffersToStackPass(std::function<bool(Value)> isSmallAlloc)
398 : isSmallAlloc(std::move(isSmallAlloc)) {}
399
initialize(MLIRContext * context)400 LogicalResult initialize(MLIRContext *context) override {
401 if (isSmallAlloc == nullptr) {
402 isSmallAlloc = [=](Value alloc) {
403 return defaultIsSmallAlloc(alloc, maxAllocSizeInBytes,
404 bitwidthOfIndexType,
405 maxRankOfAllocatedMemRef);
406 };
407 }
408 return success();
409 }
410
runOnFunction()411 void runOnFunction() override {
412 // Move all allocation nodes and convert candidates into allocas.
413 BufferPlacementPromotion optimizer(getFunction());
414 optimizer.promote(isSmallAlloc);
415 }
416
417 private:
418 std::function<bool(Value)> isSmallAlloc;
419 };
420
421 } // end anonymous namespace
422
createBufferHoistingPass()423 std::unique_ptr<Pass> mlir::createBufferHoistingPass() {
424 return std::make_unique<BufferHoistingPass>();
425 }
426
createBufferLoopHoistingPass()427 std::unique_ptr<Pass> mlir::createBufferLoopHoistingPass() {
428 return std::make_unique<BufferLoopHoistingPass>();
429 }
430
431 std::unique_ptr<Pass>
createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,unsigned bitwidthOfIndexType,unsigned maxRankOfAllocatedMemRef)432 mlir::createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
433 unsigned bitwidthOfIndexType,
434 unsigned maxRankOfAllocatedMemRef) {
435 return std::make_unique<PromoteBuffersToStackPass>(
436 maxAllocSizeInBytes, bitwidthOfIndexType, maxRankOfAllocatedMemRef);
437 }
438
439 std::unique_ptr<Pass>
createPromoteBuffersToStackPass(std::function<bool (Value)> isSmallAlloc)440 mlir::createPromoteBuffersToStackPass(std::function<bool(Value)> isSmallAlloc) {
441 return std::make_unique<PromoteBuffersToStackPass>(std::move(isSmallAlloc));
442 }
443