1054 lines · cpp
1//===- OwnershipBasedBufferDeallocation.cpp - impl. for buffer dealloc. ---===//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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements logic for computing correct `bufferization.dealloc`10// positions. Furthermore, buffer deallocation also adds required new clone11// operations to ensure that memrefs returned by functions never alias an12// argument.13//14// TODO:15// The current implementation does not support explicit-control-flow loops and16// the resulting code will be invalid with respect to program semantics.17// However, structured control-flow loops are fully supported.18//19//===----------------------------------------------------------------------===//20 21#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"22#include "mlir/Dialect/Bufferization/IR/Bufferization.h"23#include "mlir/Dialect/Bufferization/Transforms/Passes.h"24#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"25#include "mlir/Dialect/Func/IR/FuncOps.h"26#include "mlir/Dialect/MemRef/IR/MemRef.h"27#include "mlir/Dialect/SCF/IR/SCF.h"28#include "mlir/IR/Iterators.h"29#include "mlir/Interfaces/ControlFlowInterfaces.h"30 31namespace mlir {32namespace bufferization {33#define GEN_PASS_DEF_OWNERSHIPBASEDBUFFERDEALLOCATIONPASS34#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"35} // namespace bufferization36} // namespace mlir37 38using namespace mlir;39using namespace mlir::bufferization;40 41//===----------------------------------------------------------------------===//42// Helpers43//===----------------------------------------------------------------------===//44 45static Value buildBoolValue(OpBuilder &builder, Location loc, bool value) {46 return arith::ConstantOp::create(builder, loc, builder.getBoolAttr(value));47}48 49static bool isMemref(Value v) { return isa<BaseMemRefType>(v.getType()); }50 51/// Return "true" if the given op is guaranteed to have neither "Allocate" nor52/// "Free" side effects.53static bool hasNeitherAllocateNorFreeSideEffect(Operation *op) {54 return !mightHaveEffect<MemoryEffects::Allocate>(op) &&55 !mightHaveEffect<MemoryEffects::Free>(op);56}57 58/// Return "true" if the given op has buffer semantics. I.e., it has buffer59/// operands, buffer results and/or buffer region entry block arguments.60static bool hasBufferSemantics(Operation *op) {61 if (llvm::any_of(op->getOperands(), isMemref) ||62 llvm::any_of(op->getResults(), isMemref))63 return true;64 for (Region ®ion : op->getRegions())65 if (!region.empty())66 if (llvm::any_of(region.front().getArguments(), isMemref))67 return true;68 return false;69}70 71//===----------------------------------------------------------------------===//72// Backedges analysis73//===----------------------------------------------------------------------===//74 75namespace {76 77/// A straight-forward program analysis which detects loop backedges induced by78/// explicit control flow.79class Backedges {80public:81 using BlockSetT = SmallPtrSet<Block *, 16>;82 using BackedgeSetT = llvm::DenseSet<std::pair<Block *, Block *>>;83 84public:85 /// Constructs a new backedges analysis using the op provided.86 Backedges(Operation *op) { recurse(op); }87 88 /// Returns the number of backedges formed by explicit control flow.89 size_t size() const { return edgeSet.size(); }90 91 /// Returns the start iterator to loop over all backedges.92 BackedgeSetT::const_iterator begin() const { return edgeSet.begin(); }93 94 /// Returns the end iterator to loop over all backedges.95 BackedgeSetT::const_iterator end() const { return edgeSet.end(); }96 97private:98 /// Enters the current block and inserts a backedge into the `edgeSet` if we99 /// have already visited the current block. The inserted edge links the given100 /// `predecessor` with the `current` block.101 bool enter(Block ¤t, Block *predecessor) {102 bool inserted = visited.insert(¤t).second;103 if (!inserted)104 edgeSet.insert(std::make_pair(predecessor, ¤t));105 return inserted;106 }107 108 /// Leaves the current block.109 void exit(Block ¤t) { visited.erase(¤t); }110 111 /// Recurses into the given operation while taking all attached regions into112 /// account.113 void recurse(Operation *op) {114 Block *current = op->getBlock();115 // If the current op implements the `BranchOpInterface`, there can be116 // cycles in the scope of all successor blocks.117 if (isa<BranchOpInterface>(op)) {118 for (Block *succ : current->getSuccessors())119 recurse(*succ, current);120 }121 // Recurse into all distinct regions and check for explicit control-flow122 // loops.123 for (Region ®ion : op->getRegions()) {124 if (!region.empty())125 recurse(region.front(), current);126 }127 }128 129 /// Recurses into explicit control-flow structures that are given by130 /// the successor relation defined on the block level.131 void recurse(Block &block, Block *predecessor) {132 // Try to enter the current block. If this is not possible, we are133 // currently processing this block and can safely return here.134 if (!enter(block, predecessor))135 return;136 137 // Recurse into all operations and successor blocks.138 for (Operation &op : block.getOperations())139 recurse(&op);140 141 // Leave the current block.142 exit(block);143 }144 145 /// Stores all blocks that are currently visited and on the processing stack.146 BlockSetT visited;147 148 /// Stores all backedges in the format (source, target).149 BackedgeSetT edgeSet;150};151 152} // namespace153 154//===----------------------------------------------------------------------===//155// BufferDeallocation156//===----------------------------------------------------------------------===//157 158namespace {159/// The buffer deallocation transformation which ensures that all allocs in the160/// program have a corresponding de-allocation.161class BufferDeallocation {162public:163 BufferDeallocation(Operation *op, DeallocationOptions options,164 SymbolTableCollection &symbolTables)165 : state(op, symbolTables), options(options) {}166 167 /// Performs the actual placement/creation of all dealloc operations.168 LogicalResult deallocate(FunctionOpInterface op);169 170private:171 /// The base case for the recursive template below.172 template <typename... T>173 typename std::enable_if<sizeof...(T) == 0, FailureOr<Operation *>>::type174 handleOp(Operation *op) {175 return op;176 }177 178 /// Applies all the handlers of the interfaces in the template list179 /// implemented by 'op'. In particular, if an operation implements more than180 /// one of the interfaces in the template list, all the associated handlers181 /// will be applied to the operation in the same order as the template list182 /// specifies. If a handler reports a failure or removes the operation without183 /// replacement (indicated by returning 'nullptr'), no further handlers are184 /// applied and the return value is propagated to the caller of 'handleOp'.185 ///186 /// The interface handlers job is to update the deallocation state, most187 /// importantly the ownership map and list of memrefs to potentially be188 /// deallocated per block, but also to insert `bufferization.dealloc`189 /// operations where needed. Obviously, no MemRefs that may be used at a later190 /// point in the control-flow may be deallocated and the ownership map has to191 /// be updated to reflect potential ownership changes caused by the dealloc192 /// operation (e.g., if two interfaces on the same op insert a dealloc193 /// operation each, the second one should query the ownership map and use them194 /// as deallocation condition such that MemRefs already deallocated in the195 /// first dealloc operation are not deallocated a second time (double-free)).196 /// Note that currently only the interfaces on terminators may insert dealloc197 /// operations and it is verified as a precondition that a terminator op must198 /// implement exactly one of the interfaces handling dealloc insertion.199 ///200 /// The return value of the 'handleInterface' functions should be a201 /// FailureOr<Operation *> indicating whether there was a failure or otherwise202 /// returning the operation itself or a replacement operation.203 ///204 /// Note: The difference compared to `TypeSwitch` is that all205 /// matching cases are applied instead of just the first match.206 template <typename InterfaceT, typename... InterfacesU>207 FailureOr<Operation *> handleOp(Operation *op) {208 Operation *next = op;209 if (auto concreteOp = dyn_cast<InterfaceT>(op)) {210 FailureOr<Operation *> result = handleInterface(concreteOp);211 if (failed(result))212 return failure();213 next = *result;214 }215 if (!next)216 return FailureOr<Operation *>(nullptr);217 return handleOp<InterfacesU...>(next);218 }219 220 /// Apply all supported interface handlers to the given op.221 FailureOr<Operation *> handleAllInterfaces(Operation *op) {222 if (auto deallocOpInterface = dyn_cast<BufferDeallocationOpInterface>(op))223 return deallocOpInterface.process(state, options);224 225 if (failed(verifyOperationPreconditions(op)))226 return failure();227 228 return handleOp<MemoryEffectOpInterface, RegionBranchOpInterface,229 CallOpInterface, BranchOpInterface,230 RegionBranchTerminatorOpInterface>(op);231 }232 233 /// Make sure that for each forwarded MemRef value, an ownership indicator234 /// `i1` value is forwarded as well such that the successor block knows235 /// whether the MemRef has to be deallocated.236 ///237 /// Example:238 /// ```239 /// ^bb1:240 /// <more ops...>241 /// cf.br ^bb2(<forward-to-bb2>)242 /// ```243 /// becomes244 /// ```245 /// // let (m, c) = getMemrefsAndConditionsToDeallocate(bb1)246 /// // let r = getMemrefsToRetain(bb1, bb2, <forward-to-bb2>)247 /// ^bb1:248 /// <more ops...>249 /// o = bufferization.dealloc m if c retain r250 /// // replace ownership(r) with o element-wise251 /// cf.br ^bb2(<forward-to-bb2>, o)252 /// ```253 FailureOr<Operation *> handleInterface(BranchOpInterface op);254 255 /// Add an ownership indicator for every forwarding MemRef operand and result.256 /// Nested regions never take ownership of MemRefs owned by a parent region257 /// (neither via forwarding operand nor when captured implicitly when the258 /// region is not isolated from above). Ownerships will only be passed to peer259 /// regions (when an operation has multiple regions, such as scf.while), or to260 /// parent regions.261 /// Note that the block arguments in the nested region are currently handled262 /// centrally in the 'dealloc' function, but better interface support could263 /// allow us to do this here for the nested region specifically to reduce the264 /// amount of assumptions we make on the structure of ops implementing this265 /// interface.266 ///267 /// Example:268 /// ```269 /// %ret = scf.for %i = %c0 to %c10 step %c1 iter_args(%m = %memref) {270 /// <more ops...>271 /// scf.yield %m : memref<2xi32>, i1272 /// }273 /// ```274 /// becomes275 /// ```276 /// %ret:2 = scf.for %i = %c0 to %c10 step %c1277 /// iter_args(%m = %memref, %own = %false) {278 /// <more ops...>279 /// // Note that the scf.yield is handled by the280 /// // RegionBranchTerminatorOpInterface (not this handler)281 /// // let o = getMemrefWithUniqueOwnership(%own)282 /// scf.yield %m, o : memref<2xi32>, i1283 /// }284 /// ```285 FailureOr<Operation *> handleInterface(RegionBranchOpInterface op);286 287 /// If the private-function-dynamic-ownership pass option is enabled and the288 /// called function is private, additional results are added for each MemRef289 /// result to pass the dynamic ownership indicator along. Otherwise, updates290 /// the ownership map and list of memrefs to be deallocated according to the291 /// function boundary ABI, i.e., assume ownership of all returned MemRefs.292 ///293 /// Example (assume `private-function-dynamic-ownership` is enabled):294 /// ```295 /// func.func @f(%arg0: memref<2xi32>) -> memref<2xi32> {...}296 /// func.func private @g(%arg0: memref<2xi32>) -> memref<2xi32> {...}297 ///298 /// %ret_f = func.call @f(%memref) : (memref<2xi32>) -> memref<2xi32>299 /// %ret_g = func.call @g(%memref) : (memref<2xi32>) -> memref<2xi32>300 /// ```301 /// becomes302 /// ```303 /// func.func @f(%arg0: memref<2xi32>) -> memref<2xi32> {...}304 /// func.func private @g(%arg0: memref<2xi32>) -> (memref<2xi32>, i1) {...}305 ///306 /// %ret_f = func.call @f(%memref) : (memref<2xi32>) -> memref<2xi32>307 /// // set ownership(%ret_f) := true308 /// // remember to deallocate %ret_f309 ///310 /// %ret_g:2 = func.call @g(%memref) : (memref<2xi32>) -> (memref<2xi32>, i1)311 /// // set ownership(%ret_g#0) := %ret_g#1312 /// // remember to deallocate %ret_g if it comes with ownership313 /// ```314 FailureOr<Operation *> handleInterface(CallOpInterface op);315 316 /// Takes care of allocation and free side-effects. It collects allocated317 /// MemRefs that we have to add to manually deallocate, but also removes318 /// values again that are already deallocated before the end of the block. It319 /// also updates the ownership map accordingly.320 ///321 /// Example:322 /// ```323 /// %alloc = memref.alloc()324 /// %alloca = memref.alloca()325 /// ```326 /// becomes327 /// ```328 /// %alloc = memref.alloc()329 /// %alloca = memref.alloca()330 /// // set ownership(alloc) := true331 /// // set ownership(alloca) := false332 /// // remember to deallocate %alloc333 /// ```334 FailureOr<Operation *> handleInterface(MemoryEffectOpInterface op);335 336 /// Takes care that the function boundary ABI is adhered to if the parent337 /// operation implements FunctionOpInterface, inserting a338 /// `bufferization.clone` if necessary, and inserts the339 /// `bufferization.dealloc` operation according to the ops operands.340 ///341 /// Example:342 /// ```343 /// ^bb1:344 /// <more ops...>345 /// func.return <return-vals>346 /// ```347 /// becomes348 /// ```349 /// // let (m, c) = getMemrefsAndConditionsToDeallocate(bb1)350 /// // let r = getMemrefsToRetain(bb1, nullptr, <return-vals>)351 /// ^bb1:352 /// <more ops...>353 /// o = bufferization.dealloc m if c retain r354 /// func.return <return-vals>355 /// (if !isFunctionWithoutDynamicOwnership: append o)356 /// ```357 FailureOr<Operation *> handleInterface(RegionBranchTerminatorOpInterface op);358 359 /// Construct a new operation which is exactly the same as the passed 'op'360 /// except that the OpResults list is appended by new results of the passed361 /// 'types'.362 /// TODO: ideally, this would be implemented using an OpInterface because it363 /// is used to append function results, loop iter_args, etc. and thus makes364 /// some assumptions that the variadic list of those is at the end of the365 /// OpResults range.366 Operation *appendOpResults(Operation *op, ArrayRef<Type> types);367 368 /// A convenience template for the generic 'appendOpResults' function above to369 /// avoid manual casting of the result.370 template <typename OpTy>371 OpTy appendOpResults(OpTy op, ArrayRef<Type> types) {372 return cast<OpTy>(appendOpResults(op.getOperation(), types));373 }374 375 /// Performs deallocation of a single basic block. This is a private function376 /// because some internal data structures have to be set up beforehand and377 /// this function has to be called on blocks in a region in dominance order.378 LogicalResult deallocate(Block *block);379 380 /// After all relevant interfaces of an operation have been processed by the381 /// 'handleInterface' functions, this function sets the ownership of operation382 /// results that have not been set yet by the 'handleInterface' functions. It383 /// generally assumes that each result can alias with every operand of the384 /// operation, if there are MemRef typed results but no MemRef operands it385 /// assigns 'false' as ownership. This happens, e.g., for the386 /// memref.get_global operation. It would also be possible to query some alias387 /// analysis to get more precise ownerships, however, the analysis would have388 /// to be updated according to the IR modifications this pass performs (e.g.,389 /// re-building operations to have more result values, inserting clone390 /// operations, etc.).391 void populateRemainingOwnerships(Operation *op);392 393 /// Given an SSA value of MemRef type, returns the same of a new SSA value394 /// which has 'Unique' ownership where the ownership indicator is guaranteed395 /// to be always 'true'.396 Value materializeMemrefWithGuaranteedOwnership(OpBuilder &builder,397 Value memref, Block *block);398 399 /// Returns whether the given operation implements FunctionOpInterface, has400 /// private visibility, and the private-function-dynamic-ownership pass option401 /// is enabled.402 bool isFunctionWithoutDynamicOwnership(Operation *op);403 404 /// Given an SSA value of MemRef type, this function queries the405 /// BufferDeallocationOpInterface of the defining operation of 'memref' for a406 /// materialized ownership indicator for 'memref'. If the op does not407 /// implement the interface or if the block for which the materialized value408 /// is requested does not match the block in which 'memref' is defined, the409 /// default implementation in410 /// `DeallocationState::getMemrefWithUniqueOwnership` is queried instead.411 std::pair<Value, Value>412 materializeUniqueOwnership(OpBuilder &builder, Value memref, Block *block);413 414 /// Checks all the preconditions for operations implementing the415 /// FunctionOpInterface that have to hold for the deallocation to be416 /// applicable:417 /// (1) Checks that there are not explicit control flow loops.418 static LogicalResult verifyFunctionPreconditions(FunctionOpInterface op);419 420 /// Checks all the preconditions for operations inside the region of421 /// operations implementing the FunctionOpInterface that have to hold for the422 /// deallocation to be applicable:423 /// (1) Checks if all operations that have at least one attached region424 /// implement the RegionBranchOpInterface. This is not required in edge cases,425 /// where we have a single attached region and the parent operation has no426 /// results.427 /// (2) Checks that no deallocations already exist. Especially deallocations428 /// in nested regions are not properly supported yet since this requires429 /// ownership of the memref to be transferred to the nested region, which does430 /// not happen by default. This constrained can be lifted in the future.431 /// (3) Checks that terminators with more than one successor except432 /// `cf.cond_br` are not present and that either BranchOpInterface or433 /// RegionBranchTerminatorOpInterface is implemented.434 static LogicalResult verifyOperationPreconditions(Operation *op);435 436 /// When the 'private-function-dynamic-ownership' pass option is enabled,437 /// additional `i1` return values are added for each MemRef result in the438 /// function signature. This function takes care of updating the439 /// `function_type` attribute of the function according to the actually440 /// returned values from the terminators.441 static LogicalResult updateFunctionSignature(FunctionOpInterface op);442 443private:444 /// Collects all analysis state and including liveness, caches, ownerships of445 /// already processed values and operations, and the MemRefs that have to be446 /// deallocated at the end of each block.447 DeallocationState state;448 449 /// Collects all pass options in a single place.450 DeallocationOptions options;451};452 453} // namespace454 455//===----------------------------------------------------------------------===//456// BufferDeallocation Implementation457//===----------------------------------------------------------------------===//458 459std::pair<Value, Value>460BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref,461 Block *block) {462 // The interface can only materialize ownership indicators in the same block463 // as the defining op.464 if (memref.getParentBlock() != block)465 return state.getMemrefWithUniqueOwnership(builder, memref, block);466 467 Operation *owner = memref.getDefiningOp();468 if (!owner)469 owner = memref.getParentBlock()->getParentOp();470 471 // If the op implements the interface, query it for a materialized ownership472 // value.473 if (auto deallocOpInterface = dyn_cast<BufferDeallocationOpInterface>(owner))474 return deallocOpInterface.materializeUniqueOwnershipForMemref(475 state, options, builder, memref);476 477 // Otherwise use the default implementation.478 return state.getMemrefWithUniqueOwnership(builder, memref, block);479}480 481LogicalResult482BufferDeallocation::verifyFunctionPreconditions(FunctionOpInterface op) {483 // (1) Ensure that there are supported loops only (no explicit control flow484 // loops).485 Backedges backedges(op);486 if (backedges.size()) {487 op->emitError("Only structured control-flow loops are supported.");488 return failure();489 }490 491 return success();492}493 494LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op) {495 // We do not care about ops that do not operate on buffers and have no496 // Allocate/Free side effect.497 if (!hasBufferSemantics(op) && hasNeitherAllocateNorFreeSideEffect(op))498 return success();499 500 // (1) The pass does not work properly when deallocations are already present.501 // Alternatively, we could also remove all deallocations as a pre-pass.502 if (isa<DeallocOp>(op))503 return op->emitError(504 "No deallocation operations must be present when running this pass!");505 506 // (2) Memory side effects of unregistered ops are unknown. In particular, we507 // do not know whether an unregistered op allocates memory or not.508 // - Ops with recursive memory effects are allowed. All nested ops in the509 // regions of `op` will be analyzed separately.510 // - Call ops are allowed even though they typically do not implement the511 // MemoryEffectOpInterface. They usually do not have side effects apart512 // from the callee, which will be analyzed separately. (This is similar to513 // "recursive memory effects".)514 if (hasUnknownEffects(op) && !isa<CallOpInterface>(op))515 return op->emitError(516 "ops with unknown memory side effects are not supported");517 518 // (3) Check that the control flow structures are supported.519 auto regions = op->getRegions();520 // Check that if the operation has at521 // least one region it implements the RegionBranchOpInterface. If there522 // is an operation that does not fulfill this condition, we cannot apply523 // the deallocation steps. Furthermore, we accept cases, where we have a524 // region that returns no results, since, in that case, the intra-region525 // control flow does not affect the transformation.526 size_t size = regions.size();527 if (((size == 1 && !op->getResults().empty()) || size > 1) &&528 !dyn_cast<RegionBranchOpInterface>(op)) {529 return op->emitError("All operations with attached regions need to "530 "implement the RegionBranchOpInterface.");531 }532 533 // (3) Check that terminators with more than one successor except `cf.cond_br`534 // are not present and that either BranchOpInterface or535 // RegionBranchTerminatorOpInterface is implemented.536 if (op->hasTrait<OpTrait::NoTerminator>())537 return op->emitError("NoTerminator trait is not supported");538 539 if (op->hasTrait<OpTrait::IsTerminator>()) {540 // Either one of those interfaces has to be implemented on terminators, but541 // not both.542 if (!isa<BranchOpInterface, RegionBranchTerminatorOpInterface>(op) ||543 (isa<BranchOpInterface>(op) &&544 isa<RegionBranchTerminatorOpInterface>(op)))545 546 return op->emitError(547 "Terminators must implement either BranchOpInterface or "548 "RegionBranchTerminatorOpInterface (but not both)!");549 550 // We only support terminators with 0 or 1 successors for now and551 // special-case the conditional branch op.552 if (op->getSuccessors().size() > 1)553 554 return op->emitError("Terminators with more than one successor "555 "are not supported!");556 }557 558 return success();559}560 561LogicalResult562BufferDeallocation::updateFunctionSignature(FunctionOpInterface op) {563 SmallVector<TypeRange> returnOperandTypes(llvm::map_range(564 op.getFunctionBody().getOps<RegionBranchTerminatorOpInterface>(),565 [&](RegionBranchTerminatorOpInterface branchOp) {566 return branchOp567 .getSuccessorOperands(RegionSuccessor(568 op.getOperation(), op.getOperation()->getResults()))569 .getTypes();570 }));571 if (!llvm::all_equal(returnOperandTypes))572 return op->emitError(573 "there are multiple return operations with different operand types");574 575 TypeRange resultTypes = op.getResultTypes();576 // Check if we found a return operation because that doesn't necessarily577 // always have to be the case, e.g., consider a function with one block that578 // has a cf.br at the end branching to itself again (i.e., an infinite loop).579 // In that case we don't want to crash but just not update the return types.580 if (!returnOperandTypes.empty())581 resultTypes = returnOperandTypes[0];582 583 op.setFunctionTypeAttr(TypeAttr::get(FunctionType::get(584 op->getContext(), op.getFunctionBody().front().getArgumentTypes(),585 resultTypes)));586 587 return success();588}589 590LogicalResult BufferDeallocation::deallocate(FunctionOpInterface op) {591 // Stop and emit a proper error message if we don't support the input IR.592 if (failed(verifyFunctionPreconditions(op)))593 return failure();594 595 // Process the function block by block.596 auto result = op->walk<WalkOrder::PostOrder, ForwardDominanceIterator<>>(597 [&](Block *block) {598 if (failed(deallocate(block)))599 return WalkResult::interrupt();600 return WalkResult::advance();601 });602 if (result.wasInterrupted())603 return failure();604 605 // Update the function signature if the function is private, dynamic ownership606 // is enabled, and the function has memrefs as arguments or results.607 return updateFunctionSignature(op);608}609 610LogicalResult BufferDeallocation::deallocate(Block *block) {611 OpBuilder builder = OpBuilder::atBlockBegin(block);612 613 // Compute liveness transfers of ownership to this block.614 SmallVector<Value> liveMemrefs;615 state.getLiveMemrefsIn(block, liveMemrefs);616 for (auto li : liveMemrefs) {617 // Ownership of implicitly captured memrefs from other regions is never618 // taken, but ownership of memrefs in the same region (but different block)619 // is taken.620 if (li.getParentRegion() == block->getParent()) {621 state.updateOwnership(li, state.getOwnership(li, li.getParentBlock()),622 block);623 state.addMemrefToDeallocate(li, block);624 continue;625 }626 627 if (li.getParentRegion()->isProperAncestor(block->getParent())) {628 Value falseVal = buildBoolValue(builder, li.getLoc(), false);629 state.updateOwnership(li, falseVal, block);630 }631 }632 633 for (unsigned i = 0, e = block->getNumArguments(); i < e; ++i) {634 BlockArgument arg = block->getArgument(i);635 if (!isMemref(arg))636 continue;637 638 // Adhere to function boundary ABI: no ownership of function argument639 // MemRefs is taken.640 if (isa<FunctionOpInterface>(block->getParentOp()) &&641 block->isEntryBlock()) {642 Value newArg = buildBoolValue(builder, arg.getLoc(), false);643 state.updateOwnership(arg, newArg);644 state.addMemrefToDeallocate(arg, block);645 continue;646 }647 648 // Pass MemRef ownerships along via `i1` values.649 Value newArg = block->addArgument(builder.getI1Type(), arg.getLoc());650 state.updateOwnership(arg, newArg);651 state.addMemrefToDeallocate(arg, block);652 }653 654 // For each operation in the block, handle the interfaces that affect aliasing655 // and ownership of memrefs.656 for (Operation &op : llvm::make_early_inc_range(*block)) {657 FailureOr<Operation *> result = handleAllInterfaces(&op);658 if (failed(result))659 return failure();660 if (!*result)661 continue;662 663 populateRemainingOwnerships(*result);664 }665 666 // TODO: if block has no terminator, handle dealloc insertion here.667 return success();668}669 670Operation *BufferDeallocation::appendOpResults(Operation *op,671 ArrayRef<Type> types) {672 SmallVector<Type> newTypes(op->getResultTypes());673 newTypes.append(types.begin(), types.end());674 auto *newOp = Operation::create(op->getLoc(), op->getName(), newTypes,675 op->getOperands(), op->getAttrDictionary(),676 op->getPropertiesStorage(),677 op->getSuccessors(), op->getNumRegions());678 for (auto [oldRegion, newRegion] :679 llvm::zip(op->getRegions(), newOp->getRegions()))680 newRegion.takeBody(oldRegion);681 682 OpBuilder(op).insert(newOp);683 op->replaceAllUsesWith(newOp->getResults().take_front(op->getNumResults()));684 op->erase();685 686 return newOp;687}688 689FailureOr<Operation *>690BufferDeallocation::handleInterface(RegionBranchOpInterface op) {691 OpBuilder builder = OpBuilder::atBlockBegin(op->getBlock());692 693 // TODO: the RegionBranchOpInterface does not provide all the necessary694 // methods to perform this transformation without additional assumptions on695 // the structure. In particular, that696 // * additional values to be passed to the next region can be added to the end697 // of the operand list, the end of the block argument list, and the end of698 // the result value list. However, it seems to be the general guideline for699 // operations implementing this interface to follow this structure.700 // * and that the block arguments and result values match the forwarded701 // operands one-to-one (i.e., that there are no other values appended to the702 // front).703 // These assumptions are satisfied by the `scf.if`, `scf.for`, and `scf.while`704 // operations.705 706 SmallVector<RegionSuccessor> regions;707 op.getSuccessorRegions(RegionBranchPoint::parent(), regions);708 assert(!regions.empty() && "Must have at least one successor region");709 SmallVector<Value> entryOperands(710 op.getEntrySuccessorOperands(regions.front()));711 unsigned numMemrefOperands = llvm::count_if(entryOperands, isMemref);712 713 // No ownership is acquired for any MemRefs that are passed to the region from714 // the outside.715 Value falseVal = buildBoolValue(builder, op.getLoc(), false);716 op->insertOperands(op->getNumOperands(),717 SmallVector<Value>(numMemrefOperands, falseVal));718 719 int counter = op->getNumResults();720 unsigned numMemrefResults = llvm::count_if(op->getResults(), isMemref);721 SmallVector<Type> ownershipResults(numMemrefResults, builder.getI1Type());722 RegionBranchOpInterface newOp = appendOpResults(op, ownershipResults);723 724 for (auto result : llvm::make_filter_range(newOp->getResults(), isMemref)) {725 state.updateOwnership(result, newOp->getResult(counter++));726 state.addMemrefToDeallocate(result, newOp->getBlock());727 }728 729 return newOp.getOperation();730}731 732Value BufferDeallocation::materializeMemrefWithGuaranteedOwnership(733 OpBuilder &builder, Value memref, Block *block) {734 // First, make sure we at least have 'Unique' ownership already.735 std::pair<Value, Value> newMemrefAndOnwership =736 materializeUniqueOwnership(builder, memref, block);737 Value newMemref = newMemrefAndOnwership.first;738 Value condition = newMemrefAndOnwership.second;739 740 // Avoid inserting additional IR if ownership is already guaranteed. In741 // particular, this is already the case when we had 'Unknown' ownership742 // initially and a clone was inserted to get to 'Unique' ownership.743 if (matchPattern(condition, m_One()))744 return newMemref;745 746 // Insert a runtime check and only clone if we still don't have ownership at747 // runtime.748 Value maybeClone = scf::IfOp::create(749 builder, memref.getLoc(), condition,750 [&](OpBuilder &builder, Location loc) {751 scf::YieldOp::create(builder, loc, newMemref);752 },753 [&](OpBuilder &builder, Location loc) {754 Value clone = bufferization::CloneOp::create(755 builder, loc, newMemref);756 scf::YieldOp::create(builder, loc, clone);757 })758 .getResult(0);759 Value trueVal = buildBoolValue(builder, memref.getLoc(), true);760 state.updateOwnership(maybeClone, trueVal);761 state.addMemrefToDeallocate(maybeClone, maybeClone.getParentBlock());762 return maybeClone;763}764 765FailureOr<Operation *>766BufferDeallocation::handleInterface(BranchOpInterface op) {767 if (op->getNumSuccessors() > 1)768 return op->emitError("BranchOpInterface operations with multiple "769 "successors are not supported yet");770 771 if (op->getNumSuccessors() != 1)772 return emitError(op.getLoc(),773 "only BranchOpInterface operations with exactly "774 "one successor are supported yet");775 776 if (op.getSuccessorOperands(0).getProducedOperandCount() > 0)777 return op.emitError("produced operands are not supported");778 779 // Collect the values to deallocate and retain and use them to create the780 // dealloc operation.781 Block *block = op->getBlock();782 OpBuilder builder(op);783 SmallVector<Value> memrefs, conditions, toRetain;784 if (failed(state.getMemrefsAndConditionsToDeallocate(785 builder, op.getLoc(), block, memrefs, conditions)))786 return failure();787 788 OperandRange forwardedOperands =789 op.getSuccessorOperands(0).getForwardedOperands();790 state.getMemrefsToRetain(block, op->getSuccessor(0), forwardedOperands,791 toRetain);792 793 auto deallocOp = bufferization::DeallocOp::create(794 builder, op.getLoc(), memrefs, conditions, toRetain);795 796 // We want to replace the current ownership of the retained values with the797 // result values of the dealloc operation as they are always unique.798 state.resetOwnerships(deallocOp.getRetained(), block);799 for (auto [retained, ownership] :800 llvm::zip(deallocOp.getRetained(), deallocOp.getUpdatedConditions())) {801 state.updateOwnership(retained, ownership, block);802 }803 804 unsigned numAdditionalReturns = llvm::count_if(forwardedOperands, isMemref);805 SmallVector<Value> newOperands(forwardedOperands);806 auto additionalConditions =807 deallocOp.getUpdatedConditions().take_front(numAdditionalReturns);808 newOperands.append(additionalConditions.begin(), additionalConditions.end());809 op.getSuccessorOperands(0).getMutableForwardedOperands().assign(newOperands);810 811 return op.getOperation();812}813 814FailureOr<Operation *> BufferDeallocation::handleInterface(CallOpInterface op) {815 OpBuilder builder(op);816 817 // Lookup the function operation and check if it has private visibility. If818 // the function is referenced by SSA value instead of a Symbol, it's assumed819 // to be public. (And we cannot easily change the type of the SSA value820 // anyway.)821 Operation *funcOp = op.resolveCallableInTable(state.getSymbolTable());822 bool isPrivate = false;823 if (auto symbol = dyn_cast_or_null<SymbolOpInterface>(funcOp))824 isPrivate = symbol.isPrivate() && !symbol.isDeclaration();825 826 // If the private-function-dynamic-ownership option is enabled and we are827 // calling a private function, we need to add an additional `i1` result for828 // each MemRef result to dynamically pass the current ownership indicator829 // rather than adhering to the function boundary ABI.830 if (options.privateFuncDynamicOwnership && isPrivate) {831 unsigned numMemrefs = llvm::count_if(op->getResults(), isMemref);832 SmallVector<Type> ownershipTypesToAppend(numMemrefs, builder.getI1Type());833 unsigned ownershipCounter = op->getNumResults();834 op = appendOpResults(op, ownershipTypesToAppend);835 836 for (auto result : llvm::make_filter_range(op->getResults(), isMemref)) {837 state.updateOwnership(result, op->getResult(ownershipCounter++));838 state.addMemrefToDeallocate(result, result.getParentBlock());839 }840 841 return op.getOperation();842 }843 844 // According to the function boundary ABI we are guaranteed to get ownership845 // of all MemRefs returned by the function. Thus we set ownership to constant846 // 'true' and remember to deallocate it.847 Value trueVal = buildBoolValue(builder, op.getLoc(), true);848 for (auto result : llvm::make_filter_range(op->getResults(), isMemref)) {849 state.updateOwnership(result, trueVal);850 state.addMemrefToDeallocate(result, result.getParentBlock());851 }852 853 return op.getOperation();854}855 856FailureOr<Operation *>857BufferDeallocation::handleInterface(MemoryEffectOpInterface op) {858 auto *block = op->getBlock();859 OpBuilder builder = OpBuilder::atBlockBegin(block);860 861 for (auto operand : llvm::make_filter_range(op->getOperands(), isMemref)) {862 if (op.getEffectOnValue<MemoryEffects::Free>(operand).has_value()) {863 // The bufferization.manual_deallocation attribute can be attached to ops864 // with an allocation and/or deallocation side effect. It indicates that865 // the op is under a "manual deallocation" scheme. Deallocation ops are866 // usually forbidden in the input IR (not supported by the buffer867 // deallocation pass). However, if they are under manual deallocation,868 // they can be safely ignored by the buffer deallocation pass.869 if (!op->hasAttr(BufferizationDialect::kManualDeallocation))870 return op->emitError(871 "memory free side-effect on MemRef value not supported!");872 873 // Buffers that were allocated under "manual deallocation" may be874 // manually deallocated. We insert a runtime assertion to cover certain875 // cases of invalid IR where an automatically managed buffer allocation876 // is manually deallocated. This is not a bulletproof check!877 OpBuilder::InsertionGuard g(builder);878 builder.setInsertionPoint(op);879 Ownership ownership = state.getOwnership(operand, block);880 if (ownership.isUnique()) {881 Value ownershipInverted = arith::XOrIOp::create(882 builder, op.getLoc(), ownership.getIndicator(),883 buildBoolValue(builder, op.getLoc(), true));884 cf::AssertOp::create(builder, op.getLoc(), ownershipInverted,885 "expected that the block does not have ownership");886 }887 }888 }889 890 for (auto res : llvm::make_filter_range(op->getResults(), isMemref)) {891 auto allocEffect = op.getEffectOnValue<MemoryEffects::Allocate>(res);892 if (allocEffect.has_value()) {893 if (isa<SideEffects::AutomaticAllocationScopeResource>(894 allocEffect->getResource())) {895 // Make sure that the ownership of auto-managed allocations is set to896 // false. This is important for operations that have at least one memref897 // typed operand. E.g., consider an operation like `bufferization.clone`898 // that lowers to a `memref.alloca + memref.copy` instead of a899 // `memref.alloc`. If we wouldn't set the ownership of the result here,900 // the default ownership population in `populateRemainingOwnerships`901 // would assume aliasing with the MemRef operand.902 state.resetOwnerships(res, block);903 state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), false));904 continue;905 }906 907 if (op->hasAttr(BufferizationDialect::kManualDeallocation)) {908 // This allocation will be deallocated manually. Assign an ownership of909 // "false", so that it will never be deallocated by the buffer910 // deallocation pass.911 state.resetOwnerships(res, block);912 state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), false));913 continue;914 }915 916 state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), true));917 state.addMemrefToDeallocate(res, block);918 }919 }920 921 return op.getOperation();922}923 924FailureOr<Operation *>925BufferDeallocation::handleInterface(RegionBranchTerminatorOpInterface op) {926 OpBuilder builder(op);927 928 // If this is a return operation of a function that is not private or the929 // dynamic function boundary ownership is disabled, we need to return memref930 // values for which we have guaranteed ownership to pass on to adhere to the931 // function boundary ABI.932 bool funcWithoutDynamicOwnership =933 isFunctionWithoutDynamicOwnership(op->getParentOp());934 if (funcWithoutDynamicOwnership) {935 for (OpOperand &val : op->getOpOperands()) {936 if (!isMemref(val.get()))937 continue;938 939 val.set(materializeMemrefWithGuaranteedOwnership(builder, val.get(),940 op->getBlock()));941 }942 }943 944 // TODO: getSuccessorRegions is not implemented by all operations we care945 // about, but we would need to check how many successors there are and under946 // which condition they are taken, etc.947 948 MutableOperandRange operands = op.getMutableSuccessorOperands(949 RegionSuccessor(op.getOperation(), op.getOperation()->getResults()));950 951 SmallVector<Value> updatedOwnerships;952 auto result = deallocation_impl::insertDeallocOpForReturnLike(953 state, op, operands.getAsOperandRange(), updatedOwnerships);954 if (failed(result) || !*result)955 return result;956 957 // Add an additional operand for every MemRef for the ownership indicator.958 if (!funcWithoutDynamicOwnership) {959 SmallVector<Value> newOperands{operands.getAsOperandRange()};960 newOperands.append(updatedOwnerships.begin(), updatedOwnerships.end());961 operands.assign(newOperands);962 }963 964 return op.getOperation();965}966 967bool BufferDeallocation::isFunctionWithoutDynamicOwnership(Operation *op) {968 auto funcOp = dyn_cast<FunctionOpInterface>(op);969 return funcOp && (!options.privateFuncDynamicOwnership ||970 !funcOp.isPrivate() || funcOp.isExternal());971}972 973void BufferDeallocation::populateRemainingOwnerships(Operation *op) {974 for (auto res : op->getResults()) {975 if (!isMemref(res))976 continue;977 if (!state.getOwnership(res, op->getBlock()).isUninitialized())978 continue;979 980 // The op does not allocate memory, otherwise, it would have been assigned981 // an ownership during `handleInterface`. Assume the result may alias with982 // any memref operand and thus combine all their ownerships.983 for (auto operand : op->getOperands()) {984 if (!isMemref(operand))985 continue;986 987 state.updateOwnership(988 res, state.getOwnership(operand, operand.getParentBlock()),989 op->getBlock());990 }991 992 // If the ownership value is still uninitialized (e.g., because the op has993 // no memref operands), assume that no ownership is taken. E.g., this is the994 // case for "memref.get_global".995 //996 // Note: This can lead to memory leaks if memory side effects are not997 // properly specified on the op.998 if (state.getOwnership(res, op->getBlock()).isUninitialized()) {999 OpBuilder builder(op);1000 state.updateOwnership(res, buildBoolValue(builder, op->getLoc(), false));1001 }1002 }1003}1004 1005//===----------------------------------------------------------------------===//1006// OwnershipBasedBufferDeallocationPass1007//===----------------------------------------------------------------------===//1008 1009namespace {1010 1011/// The actual buffer deallocation pass that inserts and moves dealloc nodes1012/// into the right positions. Furthermore, it inserts additional clones if1013/// necessary. It uses the algorithm described at the top of the file.1014struct OwnershipBasedBufferDeallocationPass1015 : public bufferization::impl::OwnershipBasedBufferDeallocationPassBase<1016 OwnershipBasedBufferDeallocationPass> {1017 using Base::Base;1018 1019 void runOnOperation() override {1020 DeallocationOptions options;1021 options.privateFuncDynamicOwnership = privateFuncDynamicOwnership;1022 1023 mlir::SymbolTableCollection symbolTables;1024 1025 auto status = getOperation()->walk([&](func::FuncOp func) {1026 if (func.isExternal())1027 return WalkResult::skip();1028 1029 if (failed(deallocateBuffersOwnershipBased(func, options, symbolTables)))1030 return WalkResult::interrupt();1031 1032 return WalkResult::advance();1033 });1034 if (status.wasInterrupted())1035 signalPassFailure();1036 }1037};1038 1039} // namespace1040 1041//===----------------------------------------------------------------------===//1042// Implement bufferization API1043//===----------------------------------------------------------------------===//1044 1045LogicalResult bufferization::deallocateBuffersOwnershipBased(1046 FunctionOpInterface op, DeallocationOptions options,1047 SymbolTableCollection &symbolTables) {1048 // Gather all required allocation nodes and prepare the deallocation phase.1049 BufferDeallocation deallocation(op, options, symbolTables);1050 1051 // Place all required temporary clone and dealloc nodes.1052 return deallocation.deallocate(op);1053}1054