diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h index 752a4a2c6f42a..eabab5c110152 100644 --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h @@ -95,7 +95,20 @@ struct DeallocationOptions { // A pass option indicating whether private functions should be modified to // pass the ownership of MemRef values instead of adhering to the function // boundary ABI. - bool privateFuncDynamicOwnership = false; + bool privateFuncDynamicOwnership = true; + + /// Inserts `cf.assert` operations to verify the function boundary ABI at + /// runtime. Currently, it is only checked that the ownership of returned + /// MemRefs is 'true'. This also ensures that the returned memref does not + /// originate from the same allocation as a function argument. + /// Note: The function boundary ABI is disabled for non-external private + /// functions if `privateFuncDynamicOwnership` is enabled and thus this option + /// does not apply to them. + /// TODO: check that returned MemRefs don't alias each other. + /// If it can be determined statically that the ABI is not adhered + /// to, an error will already be emitted at compile time. This cannot be + /// changed with this option. + bool verifyFunctionBoundaryABI = true; }; /// This class collects all the state that we need to perform the buffer @@ -138,12 +151,12 @@ class DeallocationState { void getLiveMemrefsIn(Block *block, SmallVectorImpl &memrefs); /// Given an SSA value of MemRef type, this function queries the ownership and - /// if it is not already in the 'Unique' state, potentially inserts IR to get - /// a new SSA value, returned as the first element of the pair, which has - /// 'Unique' ownership and can be used instead of the passed Value with the - /// the ownership indicator returned as the second element of the pair. - std::pair - getMemrefWithUniqueOwnership(OpBuilder &builder, Value memref, Block *block); + /// if it is not already in the 'Unique' state, potentially inserts IR to + /// determine the ownership (which might involve expensive aliasing checks at + /// runtime). + Value materializeMemRefOwnership(const DeallocationOptions &options, + OpBuilder &builder, Value memref, + Block *block); /// Given two basic blocks and the values passed via block arguments to the /// destination block, compute the list of MemRefs that have to be retained in @@ -220,6 +233,16 @@ FailureOr insertDeallocOpForReturnLike(DeallocationState &state, Operation *op, ValueRange operands, SmallVectorImpl &updatedOperandOwnerships); + +/// Materializes IR that extracts the allocated pointers of the MemRef operands +/// of the defining operation of `memref` as indices and compares them. The +/// ownership of the first one that matches is returned and intended to be +/// assigned to `memref`. +Value defaultComputeMemRefOwnership(const DeallocationOptions &options, + DeallocationState &state, + OpBuilder &builder, Value memref, + Block *block); + } // namespace deallocation_impl } // namespace bufferization diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td index 3e11432c65c5f..9410557b22ffa 100644 --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.td @@ -55,8 +55,38 @@ def BufferDeallocationOpInterface : ownership indicator when needed, it should be implemented using this method (which is especially important if operations are created that cannot be easily canonicalized away anymore). + Ownership indicators have to be materialized when + * needed for the condition operands of a `bufferization.dealloc` op + * passing them as additional operands to nested regions (e.g., + init_args of `scf.for`) + * passing them as additional operands to a call operation when + `private-function-dynamic-ownership` is enabled + + In the following example, the deallocation pass would add an + additional block argument to `^bb1` for passing the ownership of `%0` + along and thus the ownership indicator has to be materialized before + the `cf.br` operation and added as a forwarded operand. + ```mlir + %0 = arith.select %cond, %m1, %m2 : memref + cf.br ^bb1(%0 : memref) + ^bb1(%arg0: memref) + ... + ``` + The `arith.select` operation could implement this interface method to + materialize another `arith.select` operation to select the + corresponding ownership indicator. + ```mlir + %0 = arith.select %cond, %m1, %m2 : memref + %0_ownership = arith.select %cond, %m1_ownership, %m2_ownership : i1 + cf.br ^bb1(%0, %0_ownership : memref, i1) + ^bb1(%arg0: memref, %arg1: i1) + ... + ``` + + The default implementation assumes that all MemRef operands already + have 'Unique' ownership. }], - /*retType=*/"std::pair", + /*retType=*/"Value", /*methodName=*/"materializeUniqueOwnershipForMemref", /*args=*/(ins "DeallocationState &":$state, "const DeallocationOptions &":$options, @@ -64,8 +94,8 @@ def BufferDeallocationOpInterface : "Value":$memref), /*methodBody=*/[{}], /*defaultImplementation=*/[{ - return state.getMemrefWithUniqueOwnership( - builder, memref, memref.getParentBlock()); + return deallocation_impl::defaultComputeMemRefOwnership( + options, state, builder, memref, memref.getParentBlock()); }]>, ]; } diff --git a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h index 7acacb763cd2c..7578351d2c4f5 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h +++ b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h @@ -17,6 +17,7 @@ namespace mlir { namespace bufferization { +struct DeallocationOptions; /// Options for the buffer deallocation pipeline. struct BufferDeallocationPipelineOptions @@ -27,7 +28,23 @@ struct BufferDeallocationPipelineOptions "Allows to add additional arguments to private functions to " "dynamically pass ownership of memrefs to callees. This can enable " "earlier deallocations."), - llvm::cl::init(false)}; + llvm::cl::init(true)}; + PassOptions::Option verifyFunctionBoundaryABI{ + *this, "verify-function-boundary-abi", + llvm::cl::desc( + "Inserts `cf.assert` operations to verify the function boundary ABI " + "at runtime. Currently, it is only checked that the ownership of " + "returned MemRefs is 'true'. This makes sure that ownership is " + "yielded and the returned MemRef does not originate from the same " + "allocation as a function argument. If it can be determined " + "statically that the ABI is not adhered to, an error will already be " + "emitted at compile time. This cannot be changed with this option."), + llvm::cl::init(true)}; + + /// Convert this BufferDeallocationPipelineOptions struct to a + /// DeallocationOptions struct to be passed to the + /// OwnershipBasedBufferDeallocationPass. + DeallocationOptions asDeallocationOptions() const; }; //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h index a6f668b26aa10..2bf82dd6f88c8 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h @@ -1,6 +1,7 @@ #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H +#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h" #include "mlir/Pass/Pass.h" namespace mlir { @@ -31,7 +32,7 @@ std::unique_ptr createBufferDeallocationPass(); /// Creates an instance of the OwnershipBasedBufferDeallocation pass to free all /// allocated buffers. std::unique_ptr createOwnershipBasedBufferDeallocationPass( - bool privateFuncDynamicOwnership = false); + const DeallocationOptions &options = DeallocationOptions()); /// Creates a pass that optimizes `bufferization.dealloc` operations. For /// example, it reduces the number of alias checks needed at runtime using @@ -134,8 +135,9 @@ func::FuncOp buildDeallocationLibraryFunction(OpBuilder &builder, Location loc, LogicalResult deallocateBuffers(Operation *op); /// Run ownership basedbuffer deallocation. -LogicalResult deallocateBuffersOwnershipBased(FunctionOpInterface op, - bool privateFuncDynamicOwnership); +LogicalResult deallocateBuffersOwnershipBased( + FunctionOpInterface op, + const DeallocationOptions &options = DeallocationOptions()); /// Creates a pass that moves allocations upwards to reduce the number of /// required copies that are inserted during the BufferDeallocation pass. diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td index e01f36b8daa18..5de17cf7faa7e 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td @@ -219,10 +219,20 @@ def OwnershipBasedBufferDeallocation : Pass< }]; let options = [ Option<"privateFuncDynamicOwnership", "private-function-dynamic-ownership", - "bool", /*default=*/"false", + "bool", /*default=*/"true", "Allows to add additional arguments to private functions to " "dynamically pass ownership of memrefs to callees. This can enable " "earlier deallocations.">, + Option<"verifyFunctionBoundaryABI", "verify-function-boundary-abi", + "bool", /*default=*/"true", + "Inserts `cf.assert` operations to verify the function boundary ABI " + "at runtime. Currently, it is only checked that the ownership of " + "returned MemRefs is 'true'. This makes sure that ownership is " + "yielded and the returned MemRef does not originate from the same " + "allocation as a function argument. " + "If it can be determined statically that the ABI is not adhered " + "to, an error will already be emitted at compile time. This cannot " + "be changed with this option.">, ]; let constructor = "mlir::bufferization::createOwnershipBasedBufferDeallocationPass()"; diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp index f2e7732e8ea4a..6c241f468539c 100644 --- a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp +++ b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp @@ -53,25 +53,21 @@ struct SelectOpInterface return op; // nothing to do } - std::pair - materializeUniqueOwnershipForMemref(Operation *op, DeallocationState &state, - const DeallocationOptions &options, - OpBuilder &builder, Value value) const { + Value materializeUniqueOwnershipForMemref(Operation *op, + DeallocationState &state, + const DeallocationOptions &options, + OpBuilder &builder, + Value value) const { auto selectOp = cast(op); assert(value == selectOp.getResult() && "Value not defined by this operation"); Block *block = value.getParentBlock(); - if (!state.getOwnership(selectOp.getTrueValue(), block).isUnique() || - !state.getOwnership(selectOp.getFalseValue(), block).isUnique()) - return state.getMemrefWithUniqueOwnership(builder, value, - value.getParentBlock()); - Value ownership = builder.create( op->getLoc(), selectOp.getCondition(), state.getOwnership(selectOp.getTrueValue(), block).getIndicator(), state.getOwnership(selectOp.getFalseValue(), block).getIndicator()); - return {selectOp.getResult(), ownership}; + return ownership; } }; diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp index 8d21446f1eb77..d3a12c76458ce 100644 --- a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp @@ -132,30 +132,83 @@ void DeallocationState::getLiveMemrefsIn(Block *block, memrefs.append(liveMemrefs); } -std::pair -DeallocationState::getMemrefWithUniqueOwnership(OpBuilder &builder, - Value memref, Block *block) { - auto iter = ownershipMap.find({memref, block}); - assert(iter != ownershipMap.end() && - "Value must already have been registered in the ownership map"); - - Ownership ownership = iter->second; - if (ownership.isUnique()) - return {memref, ownership.getIndicator()}; - - // Instead of inserting a clone operation we could also insert a dealloc - // operation earlier in the block and use the updated ownerships returned by - // the op for the retained values. Alternatively, we could insert code to - // check aliasing at runtime and use this information to combine two unique - // ownerships more intelligently to not end up with an 'Unknown' ownership in - // the first place. - auto cloneOp = - builder.create(memref.getLoc(), memref); - Value condition = buildBoolValue(builder, memref.getLoc(), true); - Value newMemref = cloneOp.getResult(); - updateOwnership(newMemref, condition); - memrefsToDeallocatePerBlock[newMemref.getParentBlock()].push_back(newMemref); - return {newMemref, condition}; +Value DeallocationState::materializeMemRefOwnership( + const DeallocationOptions &options, OpBuilder &builder, Value memref, + Block *block) { + // NOTE: Starts at the operation defining `memref` and performs a DFS along + // the reverse def/use chain until MemRef values with 'Unique' ownership are + // found. For the operation being currently processed: + // * if none of the operands have the same allocated pointer (i.e., originate + // from the same allocation), a new memref was allocated and thus the + // operation should have the allocate side-effect defined on that result + // value and thus the correct unique ownership is pre-populated by the + // ownership pass (unless an interface implementation is incorrect). Note + // that this is problematic for operations of unregistered dialects because + // the allocation side-effect cannot be represented in the assembly format. + // * if exactly one operand has the same allocated pointer, this returnes the + // ownership of exactly that operand + // * if multiple operands match the allocated pointer of the result, the + // ownership indicators of all of them always have to evaluate to the same + // value because no dealloc operations may be present and because of the + // rules they are passed to nested regions and successor blocks. This could + // be verified at runtime by inserting `cf.assert` operations, but would + // require O(|operands|^2) additional operations to check and is thus not + // implemented yet (would need to insert a library function to avoid + // code-size explosion which would make the deallocation pass a module pass) + auto ipSave = builder.saveInsertionPoint(); + SmallVector worklist; + worklist.push_back(memref); + + while (!worklist.empty()) { + Value curr = worklist.back(); + + // If the value already has unique ownership, we don't have to process it + // anymore. + Ownership ownership = getOwnership(curr, block); + if (ownership.isUnique()) { + worklist.pop_back(); + continue; + } + + // Check if all operands of MemRef type have unique ownership. + Operation *defOp = curr.getDefiningOp(); + assert(defOp && + "the ownership-based deallocation pass should be written in a way " + "that pre-populates ownership for block arguments"); + + bool allKnown = true; + for (Value val : llvm::make_filter_range(defOp->getOperands(), isMemref)) { + Ownership ownership = getOwnership(val, block); + if (ownership.isUnique()) + continue; + + worklist.push_back(val); + allKnown = false; + } + + // If all MemRef operands have unique ownership, we can check if the op + // implements the BufferDeallocationOpInterface and call that or, otherwise, + // we call the generic implementation manually here. + if (allKnown) { + builder.setInsertionPointAfter(defOp); + if (auto deallocInterface = + dyn_cast(defOp); + deallocInterface && curr.getParentBlock() == block) + ownership = deallocInterface.materializeUniqueOwnershipForMemref( + *this, options, builder, curr); + else + ownership = deallocation_impl::defaultComputeMemRefOwnership( + options, *this, builder, curr, block); + + // Ownership is already 'Unknown', so we need to override instead of + // joining. + resetOwnerships(curr, block); + updateOwnership(curr, ownership, block); + } + } + + builder.restoreInsertionPoint(ipSave); + return getOwnership(memref, block).getIndicator(); } void DeallocationState::getMemrefsToRetain( @@ -313,3 +366,25 @@ FailureOr deallocation_impl::insertDeallocOpForReturnLike( return op; } + +Value deallocation_impl::defaultComputeMemRefOwnership( + const DeallocationOptions &options, DeallocationState &state, + OpBuilder &builder, Value memref, Block *block) { + Operation *defOp = memref.getDefiningOp(); + SmallVector operands( + llvm::make_filter_range(defOp->getOperands(), isMemref)); + Value resultPtr = builder.create( + defOp->getLoc(), memref); + Value ownership = state.getOwnership(operands.front(), block).getIndicator(); + + for (Value val : ArrayRef(operands).drop_front()) { + Value operandPtr = builder.create( + defOp->getLoc(), val); + Value isSameBuffer = builder.create( + defOp->getLoc(), arith::CmpIPredicate::eq, resultPtr, operandPtr); + Value newOwnership = state.getOwnership(val, block).getIndicator(); + ownership = builder.create(defOp->getLoc(), isSameBuffer, + newOwnership, ownership); + } + return ownership; +} diff --git a/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp b/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp index a2878f0b80fa1..ea05fa29ea608 100644 --- a/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp +++ b/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp @@ -8,22 +8,34 @@ #include "mlir/Dialect/Bufferization/Pipelines/Passes.h" +#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h" #include "mlir/Dialect/Bufferization/Transforms/Passes.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/Transforms/Passes.h" #include "mlir/Pass/PassManager.h" #include "mlir/Transforms/Passes.h" +using namespace mlir; +using namespace bufferization; + //===----------------------------------------------------------------------===// // Pipeline implementation. //===----------------------------------------------------------------------===// +DeallocationOptions +BufferDeallocationPipelineOptions::asDeallocationOptions() const { + DeallocationOptions opts; + opts.privateFuncDynamicOwnership = privateFunctionDynamicOwnership.getValue(); + opts.verifyFunctionBoundaryABI = verifyFunctionBoundaryABI.getValue(); + return opts; +} + void mlir::bufferization::buildBufferDeallocationPipeline( OpPassManager &pm, const BufferDeallocationPipelineOptions &options) { pm.addPass(memref::createExpandReallocPass(/*emitDeallocs=*/false)); pm.addPass(createCanonicalizerPass()); pm.addPass(createOwnershipBasedBufferDeallocationPass( - options.privateFunctionDynamicOwnership.getValue())); + options.asDeallocationOptions())); pm.addPass(createCanonicalizerPass()); pm.addPass(createBufferDeallocationSimplificationPass()); pm.addPass(createLowerDeallocationsPass()); diff --git a/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt index 6e8dab64ba6b9..d67b28b308fa1 100644 --- a/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt +++ b/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt @@ -5,6 +5,7 @@ add_mlir_dialect_library(MLIRBufferizationPipelines ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization LINK_LIBS PUBLIC + MLIRBufferizationDialect MLIRBufferizationTransforms MLIRMemRefTransforms MLIRFuncDialect diff --git a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt index ed8dbd57bf40b..7cef74be6177f 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt @@ -26,6 +26,7 @@ add_mlir_dialect_library(MLIRBufferizationTransforms LINK_LIBS PUBLIC MLIRArithDialect MLIRBufferizationDialect + MLIRControlFlowDialect MLIRControlFlowInterfaces MLIRFuncDialect MLIRFunctionInterfaces diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp index fd36716163d0a..b06459fded872 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp @@ -21,6 +21,7 @@ #include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Bufferization/Transforms/Passes.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" @@ -139,10 +140,8 @@ namespace { /// program have a corresponding de-allocation. class BufferDeallocation { public: - BufferDeallocation(Operation *op, bool privateFuncDynamicOwnership) - : state(op) { - options.privateFuncDynamicOwnership = privateFuncDynamicOwnership; - } + BufferDeallocation(Operation *op, DeallocationOptions options) + : state(op), options(options) {} /// Performs the actual placement/creation of all dealloc operations. LogicalResult deallocate(FunctionOpInterface op); @@ -373,27 +372,11 @@ class BufferDeallocation { /// operations, etc.). void populateRemainingOwnerships(Operation *op); - /// Given an SSA value of MemRef type, returns the same of a new SSA value - /// which has 'Unique' ownership where the ownership indicator is guaranteed - /// to be always 'true'. - Value materializeMemrefWithGuaranteedOwnership(OpBuilder &builder, - Value memref, Block *block); - /// Returns whether the given operation implements FunctionOpInterface, has /// private visibility, and the private-function-dynamic-ownership pass option /// is enabled. bool isFunctionWithoutDynamicOwnership(Operation *op); - /// Given an SSA value of MemRef type, this function queries the - /// BufferDeallocationOpInterface of the defining operation of 'memref' for a - /// materialized ownership indicator for 'memref'. If the op does not - /// implement the interface or if the block for which the materialized value - /// is requested does not match the block in which 'memref' is defined, the - /// default implementation in - /// `DeallocationState::getMemrefWithUniqueOwnership` is queried instead. - std::pair - materializeUniqueOwnership(OpBuilder &builder, Value memref, Block *block); - /// Checks all the preconditions for operations implementing the /// FunctionOpInterface that have to hold for the deallocation to be /// applicable: @@ -430,7 +413,7 @@ class BufferDeallocation { DeallocationState state; /// Collects all pass options in a single place. - DeallocationOptions options; + const DeallocationOptions options; }; } // namespace @@ -439,28 +422,6 @@ class BufferDeallocation { // BufferDeallocation Implementation //===----------------------------------------------------------------------===// -std::pair -BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref, - Block *block) { - // The interface can only materialize ownership indicators in the same block - // as the defining op. - if (memref.getParentBlock() != block) - return state.getMemrefWithUniqueOwnership(builder, memref, block); - - Operation *owner = memref.getDefiningOp(); - if (!owner) - owner = memref.getParentBlock()->getParentOp(); - - // If the op implements the interface, query it for a materialized ownership - // value. - if (auto deallocOpInterface = dyn_cast(owner)) - return deallocOpInterface.materializeUniqueOwnershipForMemref( - state, options, builder, memref); - - // Otherwise use the default implementation. - return state.getMemrefWithUniqueOwnership(builder, memref, block); -} - static bool regionOperatesOnMemrefValues(Region ®ion) { WalkResult result = region.walk([](Block *block) { if (llvm::any_of(block->getArguments(), isMemref)) @@ -712,41 +673,6 @@ BufferDeallocation::handleInterface(RegionBranchOpInterface op) { return newOp.getOperation(); } -Value BufferDeallocation::materializeMemrefWithGuaranteedOwnership( - OpBuilder &builder, Value memref, Block *block) { - // First, make sure we at least have 'Unique' ownership already. - std::pair newMemrefAndOnwership = - materializeUniqueOwnership(builder, memref, block); - Value newMemref = newMemrefAndOnwership.first; - Value condition = newMemrefAndOnwership.second; - - // Avoid inserting additional IR if ownership is already guaranteed. In - // particular, this is already the case when we had 'Unknown' ownership - // initially and a clone was inserted to get to 'Unique' ownership. - if (matchPattern(condition, m_One())) - return newMemref; - - // Insert a runtime check and only clone if we still don't have ownership at - // runtime. - Value maybeClone = - builder - .create( - memref.getLoc(), condition, - [&](OpBuilder &builder, Location loc) { - builder.create(loc, newMemref); - }, - [&](OpBuilder &builder, Location loc) { - Value clone = - builder.create(loc, newMemref); - builder.create(loc, clone); - }) - .getResult(0); - Value trueVal = buildBoolValue(builder, memref.getLoc(), true); - state.updateOwnership(maybeClone, trueVal); - state.addMemrefToDeallocate(maybeClone, maybeClone.getParentBlock()); - return maybeClone; -} - FailureOr BufferDeallocation::handleInterface(BranchOpInterface op) { if (op->getNumSuccessors() > 1) @@ -819,10 +745,11 @@ FailureOr BufferDeallocation::handleInterface(CallOpInterface op) { newOperands.push_back(operand); continue; } - auto [memref, condition] = - materializeUniqueOwnership(builder, operand, op->getBlock()); - newOperands.push_back(memref); - ownershipIndicatorsToAdd.push_back(condition); + Value ownership = state.materializeMemRefOwnership( + options, builder, operand, op->getBlock()); + + newOperands.push_back(operand); + ownershipIndicatorsToAdd.push_back(ownership); } newOperands.append(ownershipIndicatorsToAdd.begin(), ownershipIndicatorsToAdd.end()); @@ -903,8 +830,14 @@ BufferDeallocation::handleInterface(RegionBranchTerminatorOpInterface op) { if (!isMemref(val.get())) continue; - val.set(materializeMemrefWithGuaranteedOwnership(builder, val.get(), - op->getBlock())); + if (options.verifyFunctionBoundaryABI) { + Value ownership = state.materializeMemRefOwnership( + options, builder, val.get(), op->getBlock()); + builder.create( + op->getLoc(), ownership, + builder.getStringAttr("Must have ownership of operand #" + + Twine(val.getOperandNumber()))); + } } } @@ -978,17 +911,21 @@ struct OwnershipBasedBufferDeallocationPass : public bufferization::impl::OwnershipBasedBufferDeallocationBase< OwnershipBasedBufferDeallocationPass> { OwnershipBasedBufferDeallocationPass() = default; - OwnershipBasedBufferDeallocationPass(bool privateFuncDynamicOwnership) + OwnershipBasedBufferDeallocationPass(const DeallocationOptions &options) : OwnershipBasedBufferDeallocationPass() { - this->privateFuncDynamicOwnership.setValue(privateFuncDynamicOwnership); + privateFuncDynamicOwnership.setValue(options.privateFuncDynamicOwnership); + verifyFunctionBoundaryABI.setValue(options.verifyFunctionBoundaryABI); } void runOnOperation() override { auto status = getOperation()->walk([&](func::FuncOp func) { if (func.isExternal()) return WalkResult::skip(); - if (failed(deallocateBuffersOwnershipBased(func, - privateFuncDynamicOwnership))) + DeallocationOptions options; + options.privateFuncDynamicOwnership = + privateFuncDynamicOwnership.getValue(); + options.verifyFunctionBoundaryABI = verifyFunctionBoundaryABI.getValue(); + if (failed(deallocateBuffersOwnershipBased(func, options))) return WalkResult::interrupt(); return WalkResult::advance(); @@ -1005,9 +942,9 @@ struct OwnershipBasedBufferDeallocationPass //===----------------------------------------------------------------------===// LogicalResult bufferization::deallocateBuffersOwnershipBased( - FunctionOpInterface op, bool privateFuncDynamicOwnership) { + FunctionOpInterface op, const DeallocationOptions &options) { // Gather all required allocation nodes and prepare the deallocation phase. - BufferDeallocation deallocation(op, privateFuncDynamicOwnership); + BufferDeallocation deallocation(op, options); // Place all required temporary clone and dealloc nodes. return deallocation.deallocate(op); @@ -1019,7 +956,6 @@ LogicalResult bufferization::deallocateBuffersOwnershipBased( std::unique_ptr mlir::bufferization::createOwnershipBasedBufferDeallocationPass( - bool privateFuncDynamicOwnership) { - return std::make_unique( - privateFuncDynamicOwnership); + const DeallocationOptions &options) { + return std::make_unique(options); } diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir index 3ae0529ab7d74..0e51c2380eb42 100644 --- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir @@ -1,6 +1,6 @@ // RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation \ // RUN: -buffer-deallocation-simplification -split-input-file %s | FileCheck %s -// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null +// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=false -split-input-file %s > /dev/null // RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir index 02bf2d10e9e3f..b310fcc4731bf 100644 --- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir @@ -1,6 +1,6 @@ // RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=false \ // RUN: -buffer-deallocation-simplification -split-input-file %s | FileCheck %s -// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true \ +// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation \ // RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s --check-prefix=CHECK-DYNAMIC // RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir index 13c55d0289880..306ce0e098b7a 100644 --- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir @@ -1,6 +1,6 @@ // RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=false \ // RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s -// RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true \ +// RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation \ // RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s --check-prefix=CHECK-DYNAMIC // RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null @@ -94,34 +94,26 @@ func.func private @redundantOperations(%arg0: memref<2xf32>) { func.func private @memref_in_function_results( %arg0: memref<5xf32>, - %arg1: memref<10xf32>, - %arg2: memref<5xf32>) -> (memref<10xf32>, memref<15xf32>) { + %arg2: memref<5xf32>) -> (memref<15xf32>) { %x = memref.alloc() : memref<15xf32> %y = memref.alloc() : memref<5xf32> test.buffer_based in(%arg0: memref<5xf32>) out(%y: memref<5xf32>) test.copy(%y, %arg2) : (memref<5xf32>, memref<5xf32>) - return %arg1, %x : memref<10xf32>, memref<15xf32> + return %x : memref<15xf32> } // CHECK-LABEL: func private @memref_in_function_results -// CHECK: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>, +// CHECK: (%[[ARG0:.*]]: memref<5xf32>, // CHECK-SAME: %[[RESULT:.*]]: memref<5xf32>) // CHECK: %[[X:.*]] = memref.alloc() // CHECK: %[[Y:.*]] = memref.alloc() // CHECK: test.copy -// CHECK-NEXT: %[[V0:.+]] = scf.if %false -// CHECK-NEXT: scf.yield %[[ARG1]] -// CHECK-NEXT: } else { -// CHECK-NEXT: %[[CLONE:.+]] = bufferization.clone %[[ARG1]] -// CHECK-NEXT: scf.yield %[[CLONE]] -// CHECK-NEXT: } // CHECK: bufferization.dealloc (%[[Y]] : {{.*}}) if (%true{{[0-9_]*}}) // CHECK-NOT: retain -// CHECK: return %[[V0]], %[[X]] +// CHECK: return %[[X]] // CHECK-DYNAMIC-LABEL: func private @memref_in_function_results -// CHECK-DYNAMIC: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>, -// CHECK-DYNAMIC-SAME: %[[RESULT:.*]]: memref<5xf32>, %[[ARG3:.*]]: i1, %[[ARG4:.*]]: i1, %[[ARG5:.*]]: i1) +// CHECK-DYNAMIC-SAME: (%[[ARG0:.*]]: memref<5xf32>, %[[RESULT:.*]]: memref<5xf32>, %[[ARG3:.*]]: i1, %[[ARG5:.*]]: i1) // CHECK-DYNAMIC: %[[X:.*]] = memref.alloc() // CHECK-DYNAMIC: %[[Y:.*]] = memref.alloc() // CHECK-DYNAMIC: test.copy @@ -129,6 +121,6 @@ func.func private @memref_in_function_results( // CHECK-DYNAMIC: %[[BASE1:[a-zA-Z0-9_]+]], {{.+}} = memref.extract_strided_metadata %[[RESULT]] // CHECK-DYNAMIC: bufferization.dealloc (%[[Y]] : {{.*}}) if (%true{{[0-9_]*}}) // CHECK-DYNAMIC-NOT: retain -// CHECK-DYNAMIC: [[OWN:%.+]] = bufferization.dealloc (%[[BASE0]], %[[BASE1]] : {{.*}}) if (%[[ARG3]], %[[ARG5]]) retain (%[[ARG1]] : -// CHECK-DYNAMIC: [[OR:%.+]] = arith.ori [[OWN]], %[[ARG4]] -// CHECK-DYNAMIC: return %[[ARG1]], %[[X]], [[OR]], %true +// CHECK-DYNAMIC: bufferization.dealloc (%[[BASE0]], %[[BASE1]] : {{.*}}) if (%[[ARG3]], %[[ARG5]]) +// CHECK-DYNAMIC-NOT: retain +// CHECK-DYNAMIC: return %[[X]], %true diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir index 44cf16385603e..269b8b71f7beb 100644 --- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir @@ -1,6 +1,6 @@ // RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation \ // RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s -// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null +// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=false -split-input-file %s > /dev/null // RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null @@ -100,27 +100,10 @@ func.func @dealloc_existing_clones(%arg0: memref, %arg1: memref, %[[ARG1:.*]]: memref) // CHECK: %[[RES0:.*]] = bufferization.clone %[[ARG0]] // CHECK: %[[RES1:.*]] = bufferization.clone %[[ARG1]] -// CHECK-NEXT: bufferization.dealloc (%[[RES1]] :{{.*}}) if (%true{{[0-9_]*}}) +// CHECK: bufferization.dealloc (%[[RES1]] :{{.*}}) if (%true{{[0-9_]*}}) // CHECK-NOT: retain // CHECK-NEXT: return %[[RES0]] // TODO: The retain operand could be dropped to avoid runtime aliasing checks // since We can guarantee at compile-time that it will never alias with the // dealloc operand - -// ----- - -memref.global "private" constant @__constant_4xf32 : memref<4xf32> = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> - -func.func @op_without_aliasing_and_allocation() -> memref<4xf32> { - %0 = memref.get_global @__constant_4xf32 : memref<4xf32> - return %0 : memref<4xf32> -} - -// CHECK-LABEL: func @op_without_aliasing_and_allocation -// CHECK: [[GLOBAL:%.+]] = memref.get_global @__constant_4xf32 -// CHECK: [[RES:%.+]] = scf.if %false -// CHECK: scf.yield [[GLOBAL]] : -// CHECK: [[CLONE:%.+]] = bufferization.clone [[GLOBAL]] -// CHECK: scf.yield [[CLONE]] : -// CHECK: return [[RES]] : diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir index dc372749fc074..f1b753e405531 100644 --- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir @@ -1,6 +1,6 @@ // RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation \ // RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s -// RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null +// RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=false -split-input-file %s > /dev/null // RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file --verify-diagnostics > /dev/null @@ -85,13 +85,7 @@ func.func @nested_region_control_flow( // CHECK: bufferization.dealloc ([[ALLOC1]] :{{.*}}) if (%true{{[0-9_]*}}) // CHECK-NOT: retain // CHECK: scf.yield [[ALLOC]], %false -// CHECK: [[V1:%.+]] = scf.if [[V0]]#1 -// CHECK: scf.yield [[V0]]#0 -// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0 -// CHECK: scf.yield [[CLONE]] -// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0 -// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] : {{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] : -// CHECK: return [[V1]] +// CHECK: return [[V0]]#0 // ----- @@ -120,13 +114,8 @@ func.func @nested_region_control_flow_div( // CHECK: scf.yield [[ALLOC]], %false // CHECK: [[ALLOC1:%.+]] = memref.alloc( // CHECK: scf.yield [[ALLOC1]], %true -// CHECK: [[V1:%.+]] = scf.if [[V0]]#1 -// CHECK: scf.yield [[V0]]#0 -// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0 -// CHECK: scf.yield [[CLONE]] -// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0 -// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] : -// CHECK: return [[V1]] +// CHECK: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true{{[0-9_]*}}) retain ([[V0]]#0 : +// CHECK: return [[V0]]#0 // ----- @@ -158,13 +147,8 @@ func.func @inner_region_control_flow(%arg0 : index) -> memref { // CHECK: test.region_if_yield [[ARG1]], [[ARG2]] // CHECK: ^bb0([[ARG1:%.+]]: memref, [[ARG2:%.+]]: i1): // CHECK: test.region_if_yield [[ARG1]], [[ARG2]] -// CHECK: [[V1:%.+]] = scf.if [[V0]]#1 -// CHECK: scf.yield [[V0]]#0 -// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0 -// CHECK: scf.yield [[CLONE]] -// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0 -// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] : -// CHECK: return [[V1]] +// CHECK-NOT: bufferization.dealloc +// CHECK: return [[V0]]#0 // ----- @@ -232,13 +216,8 @@ func.func @nestedRegionControlFlowAlloca( // CHECK: scf.yield [[ALLOC]], %false // CHECK: memref.alloca( // CHECK: scf.yield [[ALLOC]], %false -// CHECK: [[V1:%.+]] = scf.if [[V0]]#1 -// CHECK: scf.yield [[V0]]#0 -// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0 -// CHECK: scf.yield [[CLONE]] -// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0 -// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] : -// CHECK: return [[V1]] +// CHECK-NOT: bufferization.dealloc +// CHECK: return [[V0]]#0 // ----- @@ -364,13 +343,8 @@ func.func @loop_nested_if_alloc( // CHECK: [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V1]]#1 // CHECK: scf.yield [[V1]]#0, [[OWN_AGG]] // CHECK: } -// CHECK: [[V2:%.+]] = scf.if [[V0]]#1 -// CHECK: scf.yield [[V0]]#0 -// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0 -// CHECK: scf.yield [[CLONE]] -// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0 -// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V2]] : -// CHECK: return [[V2]] +// CHECK: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true{{[0-9_]*}}) retain ([[V0]]#0 : +// CHECK: return [[V0]]#0 // ----- @@ -626,13 +600,7 @@ func.func @test_affine_if_1(%arg0: memref<10xf32>) -> memref<10xf32> { // CHECK: [[ALLOC:%.+]] = memref.alloc() // CHECK: affine.yield [[ALLOC]], %true // CHECK: affine.yield [[ARG0]], %false -// CHECK: [[V1:%.+]] = scf.if [[V0]]#1 -// CHECK: scf.yield [[V0]]#0 -// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0 -// CHECK: scf.yield [[CLONE]] -// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0 -// CHECK: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[V0]]#1) retain ([[V1]] : -// CHECK: return [[V1]] +// CHECK: return [[V0]]#0 // TODO: the dealloc could be optimized away since the memref to be deallocated // either aliases with V1 or the condition is false @@ -652,19 +620,14 @@ func.func @test_affine_if_2() -> memref<10xf32> { } return %0 : memref<10xf32> } + // CHECK-LABEL: func @test_affine_if_2 // CHECK: [[ALLOC:%.+]] = memref.alloc() // CHECK: [[V0:%.+]]:2 = affine.if // CHECK: affine.yield [[ALLOC]], %false // CHECK: [[ALLOC1:%.+]] = memref.alloc() // CHECK: affine.yield [[ALLOC1]], %true -// CHECK: [[V1:%.+]] = scf.if [[V0]]#1 -// CHECK: scf.yield [[V0]]#0 -// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0 -// CHECK: scf.yield [[CLONE]] -// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0 -// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] : -// CHECK: return [[V1]] +// CHECK: return [[V0]]#0 // ----- @@ -688,10 +651,4 @@ func.func @test_affine_if_3() -> memref<10xf32> { // CHECK: [[ALLOC1:%.+]] = memref.alloc() // CHECK: affine.yield [[ALLOC1]], %true // CHECK: affine.yield [[ALLOC]], %false -// CHECK: [[V1:%.+]] = scf.if [[V0]]#1 -// CHECK: scf.yield [[V0]]#0 -// CHECK: [[CLONE:%.+]] = bufferization.clone [[V0]]#0 -// CHECK: scf.yield [[CLONE]] -// CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0 -// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] -// CHECK: return [[V1]] +// CHECK: return [[V0]]#0 diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-runtime-verification.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-runtime-verification.mlir new file mode 100644 index 0000000000000..9951347efa45d --- /dev/null +++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-runtime-verification.mlir @@ -0,0 +1,13 @@ +// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation -split-input-file %s | FileCheck %s + +memref.global "private" constant @__constant_4xf32 : memref<4xf32> = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> + +func.func @op_without_aliasing_and_allocation() -> memref<4xf32> { + %0 = memref.get_global @__constant_4xf32 : memref<4xf32> + return %0 : memref<4xf32> +} + +// CHECK-LABEL: func @op_without_aliasing_and_allocation +// CHECK: [[GLOBAL:%.+]] = memref.get_global @__constant_4xf32 +// CHECK: cf.assert %false{{[0-9_]*}}, "Must have ownership of operand #0" +// CHECK: return [[GLOBAL]] : diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir index 35523319de154..bf55730c8fbb3 100644 --- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir @@ -1,6 +1,6 @@ // RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation \ // RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s -// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null +// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=false -split-input-file %s > /dev/null // RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-unknown-ops.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-unknown-ops.mlir new file mode 100644 index 0000000000000..6808053500e6c --- /dev/null +++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-unknown-ops.mlir @@ -0,0 +1,29 @@ +// RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation \ +// RUN: -buffer-deallocation-simplification -split-input-file %s | FileCheck %s + +func.func private @callee(%arg0: memref) -> memref { + return %arg0 : memref +} + +func.func @generic_ownership_materialization() { + %a1 = memref.alloc() : memref + %a2 = memref.alloca() : memref + %0 = "my_dialect.select_randomly"(%a1, %a2, %a1) : (memref, memref, memref) -> memref + %1 = call @callee(%0) : (memref) -> memref + return +} + +// CHECK-LABEL: func @generic_ownership_materialization +// CHECK: [[ALLOC:%.+]] = memref.alloc( +// CHECK: [[ALLOCA:%.+]] = memref.alloca( +// CHECK: [[SELECT:%.+]] = "my_dialect.select_randomly"([[ALLOC]], [[ALLOCA]], [[ALLOC]]) +// CHECK: [[SELECT_PTR:%.+]] = memref.extract_aligned_pointer_as_index [[SELECT]] +// CHECK: [[ALLOCA_PTR:%.+]] = memref.extract_aligned_pointer_as_index [[ALLOCA]] +// CHECK: [[EQ1:%.+]] = arith.cmpi eq, [[SELECT_PTR]], [[ALLOCA_PTR]] +// CHECK: [[OWN1:%.+]] = arith.select [[EQ1]], %false{{[0-9_]*}}, %true +// CHECK: [[ALLOC_PTR:%.+]] = memref.extract_aligned_pointer_as_index [[ALLOC]] +// CHECK: [[EQ2:%.+]] = arith.cmpi eq, [[SELECT_PTR]], [[ALLOC_PTR]] +// CHECK: [[OWN2:%.+]] = arith.select [[EQ2]], %true{{[0-9_]*}}, [[OWN1]] +// CHECK: [[CALL:%.+]]:2 = call @callee([[SELECT]], [[OWN2]]) +// CHECK: [[BASE:%.+]],{{.*}} = memref.extract_strided_metadata [[CALL]]#0 +// CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[CALL]]#1) diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/test-collapse-tensor.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/test-collapse-tensor.mlir index 43e423d4c3e8e..a0a41d04e1a00 100644 --- a/mlir/test/Integration/Dialect/Linalg/CPU/test-collapse-tensor.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/test-collapse-tensor.mlir @@ -33,7 +33,7 @@ func.func @main() { func.func private @printMemrefF32(%ptr : tensor<*xf32>) -func.func @collapse_dynamic_shape(%arg0 : tensor<2x?x?x?xf32>) -> tensor<2x?x?xf32> { +func.func private @collapse_dynamic_shape(%arg0 : tensor<2x?x?x?xf32>) -> tensor<2x?x?xf32> { %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]]: tensor<2x?x?x?xf32> into tensor<2x?x?xf32> return %0 : tensor<2x?x?xf32> } diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/test-expand-tensor.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/test-expand-tensor.mlir index a101b76ef186b..0aa1b81b5a426 100644 --- a/mlir/test/Integration/Dialect/Linalg/CPU/test-expand-tensor.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/test-expand-tensor.mlir @@ -34,7 +34,7 @@ func.func @main() { func.func private @printMemrefF32(%ptr : tensor<*xf32>) -func.func @expand_dynamic_shape(%arg0 : tensor<2x?x?xf32>) -> tensor<2x2x?x1x?xf32> { +func.func private @expand_dynamic_shape(%arg0 : tensor<2x?x?xf32>) -> tensor<2x2x?x1x?xf32> { %0 = tensor.expand_shape %arg0 [[0], [1, 2, 3], [4]]: tensor<2x?x?xf32> into tensor<2x2x?x1x?xf32> return %0 : tensor<2x2x?x1x?xf32> } diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/test-one-shot-bufferize.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/test-one-shot-bufferize.mlir index 06165515d4613..d58414bb43cc3 100644 --- a/mlir/test/Integration/Dialect/Linalg/CPU/test-one-shot-bufferize.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/test-one-shot-bufferize.mlir @@ -9,7 +9,7 @@ #map0 = affine_map<(d0, d1)[s0] -> ((d1 - d0) ceildiv s0)> #map1 = affine_map<(d0, d1)[s0] -> ((d0 - d1) ceildiv s0)> -func.func @init_and_dot(%arg0: tensor<64xf32>, %arg1: tensor<64xf32>, %arg2: tensor) -> tensor { +func.func private @init_and_dot(%arg0: tensor<64xf32>, %arg1: tensor<64xf32>, %arg2: tensor) -> tensor { %c64 = arith.constant 64 : index %cst = arith.constant 0.000000e+00 : f32 %c2 = arith.constant 2 : index diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/test-tensor-e2e.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/test-tensor-e2e.mlir index 38b49cd444df3..dcaa484315af1 100644 --- a/mlir/test/Integration/Dialect/Linalg/CPU/test-tensor-e2e.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/test-tensor-e2e.mlir @@ -5,7 +5,7 @@ // RUN: -shared-libs=%mlir_runner_utils \ // RUN: | FileCheck %s -func.func @foo() -> tensor<4xf32> { +func.func private @foo() -> tensor<4xf32> { %0 = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32> return %0 : tensor<4xf32> } diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel index cd6d83db9680b..257d4311ff25f 100644 --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -12390,6 +12390,7 @@ cc_library( hdrs = ["include/mlir/Dialect/Bufferization/Pipelines/Passes.h"], includes = ["include"], deps = [ + ":BufferizationDialect", ":BufferizationToMemRef", ":BufferizationTransforms", ":FuncDialect",