diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td @@ -119,7 +119,10 @@ // AllocOp //===----------------------------------------------------------------------===// -def MemRef_AllocOp : AllocLikeOp<"alloc", DefaultResource> { +def MemRef_AllocOp : AllocLikeOp<"alloc", DefaultResource, [ + DeclareOpInterfaceMethods] + > { let summary = "memory allocation operation"; let description = [{ The `alloc` operation allocates a region of memory, as specified by its @@ -415,7 +418,9 @@ def CloneOp : MemRef_Op<"clone", [ CopyOpInterface, - DeclareOpInterfaceMethods + DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods ]> { let builders = [ OpBuilder<(ins "Value":$value), [{ diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td --- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td +++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td @@ -16,6 +16,45 @@ include "mlir/Interfaces/SideEffectInterfaceBase.td" +//===----------------------------------------------------------------------===// +// AllocationOpInterface +//===----------------------------------------------------------------------===// + +def AllocationOpInterface : OpInterface<"AllocationOpInterface"> { + let description = [{ + This interface provides general allocation-related methods that are + designed for allocation operations. For example, it offers the ability to + construct associated deallocation and clone operations that are compatible + with the current allocation operation. + }]; + let cppNamespace = "::mlir"; + + let methods = [ + StaticInterfaceMethod<[{ + Builds a deallocation operation using the provided builder and the + current allocation value (which refers to the current Op implementing + this interface). The allocation value is a result of the current + operation implementing this interface. If there is no compatible + deallocation operation, this method can return ::llvm::None. + }], + "::mlir::Optional<::mlir::Operation*>", "buildDealloc", + (ins "::mlir::OpBuilder&":$opBuilder, "::mlir::Value":$alloc), [{}], + /*defaultImplementation=*/[{ return llvm::None; }] + >, + StaticInterfaceMethod<[{ + Builds a clone operation using the provided builder and the current + allocation value (which refers to the current Op implementing this + interface). The allocation value is a result of the current operation + implementing this interface. If there is no compatible clone operation, + this method can return ::llvm::None. + }], + "::mlir::Optional<::mlir::Value>", "buildClone", + (ins "::mlir::OpBuilder&":$opBuilder, "::mlir::Value":$alloc), [{}], + /*defaultImplementation=*/[{ return llvm::None; }] + > + ]; +} + //===----------------------------------------------------------------------===// // MemoryEffects //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp --- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp +++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp @@ -190,6 +190,15 @@ }; } // end anonymous namespace. +Optional AllocOp::buildDealloc(OpBuilder &builder, Value alloc) { + return builder.create(alloc.getLoc(), alloc) + .getOperation(); +} + +Optional AllocOp::buildClone(OpBuilder &builder, Value alloc) { + return builder.create(alloc.getLoc(), alloc).getResult(); +} + void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { results.add, SimplifyDeadAlloc>(context); @@ -638,6 +647,15 @@ return succeeded(foldMemRefCast(*this)) ? getResult() : Value(); } +Optional CloneOp::buildDealloc(OpBuilder &builder, Value alloc) { + return builder.create(alloc.getLoc(), alloc) + .getOperation(); +} + +Optional CloneOp::buildClone(OpBuilder &builder, Value alloc) { + return builder.create(alloc.getLoc(), alloc).getResult(); +} + //===----------------------------------------------------------------------===// // DeallocOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Transforms/BufferDeallocation.cpp b/mlir/lib/Transforms/BufferDeallocation.cpp --- a/mlir/lib/Transforms/BufferDeallocation.cpp +++ b/mlir/lib/Transforms/BufferDeallocation.cpp @@ -187,12 +187,46 @@ /// The buffer deallocation transformation which ensures that all allocs in the /// program have a corresponding de-allocation. As a side-effect, it might also /// introduce clones that in turn leads to additional deallocations. -class BufferDeallocation : BufferPlacementTransformationBase { +class BufferDeallocation : public BufferPlacementTransformationBase { public: + using AliasAllocationMapT = llvm::DenseMap; + BufferDeallocation(Operation *op) : BufferPlacementTransformationBase(op), dominators(op), postDominators(op) {} + /// Checks if all allocation operations either provide an already existing + /// deallocation operation or implement the AllocationOpInterface. In + /// addition, this method initializes the internal alias to + /// AllocationOpInterface mapping in order to get compatible + /// AllocationOpInterface implementations for aliases. + LogicalResult prepare() { + for (const BufferPlacementAllocs::AllocEntry &entry : allocs) { + // Get the defining allocation operation. + Value alloc = std::get<0>(entry); + auto allocationInterface = alloc.getDefiningOp(); + // If there is no existing deallocation operation and no implementation of + // the AllocationOpInterface, we cannot apply the BufferDeallocation pass. + if (!std::get<1>(entry) && !allocationInterface) { + return alloc.getDefiningOp()->emitError( + "Allocation is not deallocated explicitly nor does the operation " + "implement the AllocationOpInterface."); + } + + // Register the current allocation interface implementation. + aliasToAllocations[alloc] = allocationInterface; + + // Get the alias information for the current allocation node. + llvm::for_each(aliases.resolve(alloc), [&](Value alias) { + // TODO: check for incompatible implementations of the + // AllocationOpInterface. This could be realized by promoting the + // AllocationOpInterface to a DialectInterface. + aliasToAllocations[alias] = allocationInterface; + }); + } + return success(); + } + /// Performs the actual placement/creation of all temporary clone and dealloc /// nodes. void deallocate() { @@ -284,18 +318,21 @@ Value sourceValue = branchInterface.getSuccessorOperands(it.getSuccessorIndex()) .getValue()[blockArg.getArgNumber()]; - // Create a new clone at the current location of the terminator. - Value clone = introduceCloneBuffers(sourceValue, terminator); // Wire new clone and successor operand. auto mutableOperands = branchInterface.getMutableSuccessorOperands(it.getSuccessorIndex()); - if (!mutableOperands.hasValue()) + if (!mutableOperands) { terminator->emitError() << "terminators with immutable successor " "operands are not supported"; - else - mutableOperands.getValue() - .slice(blockArg.getArgNumber(), 1) - .assign(clone); + continue; + } + // Create a new clone at the current location of the terminator. + auto clone = introduceCloneBuffers(sourceValue, terminator); + if (failed(clone)) + continue; + mutableOperands.getValue() + .slice(blockArg.getArgNumber(), 1) + .assign(*clone); } // Check whether the block argument has implicitly defined predecessors via @@ -333,12 +370,14 @@ Value operand = regionInterface.getSuccessorEntryOperands(argRegion->getRegionNumber()) [llvm::find(it->getSuccessorInputs(), blockArg).getIndex()]; - Value clone = introduceCloneBuffers(operand, parentOp); + auto clone = introduceCloneBuffers(operand, parentOp); + if (failed(clone)) + return; auto op = llvm::find(parentOp->getOperands(), operand); assert(op != parentOp->getOperands().end() && "parentOp does not contain operand"); - parentOp->setOperand(op.getIndex(), clone); + parentOp->setOperand(op.getIndex(), *clone); } /// Introduces temporary clones in front of all associated nested-region @@ -399,9 +438,11 @@ OperandRange immutableTerminatorOperands = terminatorOperands; Value sourceValue = immutableTerminatorOperands[operandIndex]; // Create a new clone at the current location of the terminator. - Value clone = introduceCloneBuffers(sourceValue, terminator); + auto clone = introduceCloneBuffers(sourceValue, terminator); + if (failed(clone)) + return; // Wire clone and terminator operand. - terminatorOperands.slice(operandIndex, 1).assign(clone); + terminatorOperands.slice(operandIndex, 1).assign(*clone); }); } } @@ -409,7 +450,8 @@ /// Creates a new memory allocation for the given source value and clones /// its content into the newly allocated buffer. The terminator operation is /// used to insert the clone operation at the right place. - Value introduceCloneBuffers(Value sourceValue, Operation *terminator) { + FailureOr introduceCloneBuffers(Value sourceValue, + Operation *terminator) { // Avoid multiple clones of the same source value. This can happen in the // presence of loops when a branch acts as a backedge while also having // another successor that returns to its parent operation. Note: that @@ -422,19 +464,18 @@ return sourceValue; // Create a new clone operation that copies the contents of the old // buffer to the new one. - OpBuilder builder(terminator); - auto cloneOp = - builder.create(terminator->getLoc(), sourceValue); - - // Remember the clone of original source value. - clonedValues.insert(cloneOp); - return cloneOp; + auto clone = buildClone(terminator, sourceValue); + if (succeeded(clone)) { + // Remember the clone of original source value. + clonedValues.insert(*clone); + } + return clone; } /// Finds correct dealloc positions according to the algorithm described at /// the top of the file for all alloc nodes and block arguments that can be /// handled by this analysis. - void placeDeallocs() const { + void placeDeallocs() { // Move or insert deallocs using the previously computed information. // These deallocations will be linked to their associated allocation nodes // since they don't have any aliases that can (potentially) increase their @@ -492,12 +533,52 @@ if (!nextOp) continue; // If there is no dealloc node, insert one in the right place. - OpBuilder builder(nextOp); - builder.create(alloc.getLoc(), alloc); + buildDealloc(nextOp, alloc); } } } + /// Builds a deallocation operation compatible with the given allocation + /// value. If there is no registered AllocationOpInterface implementation for + /// the given value (e.g. in the case of a function parameter), this method + /// builds a memref::DeallocOp. + void buildDealloc(Operation *op, Value alloc) { + OpBuilder builder(op); + auto it = aliasToAllocations.find(alloc); + if (it != aliasToAllocations.end()) { + // Call the allocation op interface to build a supported and + // compatible deallocation operation. + auto dealloc = it->second.buildDealloc(builder, alloc); + if (!dealloc) + op->emitError() << "allocations without compatible deallocations are " + "not supported"; + } else { + // Build a "default" DeallocOp for unknown allocation sources. + builder.create(alloc.getLoc(), alloc); + } + } + + /// Builds a clone operation compatible with the given allocation value. If + /// there is no registered AllocationOpInterface implementation for the given + /// value (e.g. in the case of a function parameter), this method builds a + /// memref::CloneOp. + FailureOr buildClone(Operation *op, Value alloc) { + OpBuilder builder(op); + auto it = aliasToAllocations.find(alloc); + if (it != aliasToAllocations.end()) { + // Call the allocation op interface to build a supported and + // compatible clone operation. + auto clone = it->second.buildClone(builder, alloc); + if (clone) + return *clone; + return (LogicalResult)(op->emitError() + << "allocations without compatible clone ops " + "are not supported"); + } + // Build a "default" CloneOp for unknown allocation sources. + return builder.create(alloc.getLoc(), alloc).getResult(); + } + /// The dominator info to find the appropriate start operation to move the /// allocs. DominanceInfo dominators; @@ -508,6 +589,9 @@ /// Stores already cloned buffers to avoid additional clones of clones. ValueSetT clonedValues; + + /// Maps aliases to their source allocation interfaces (inverse mapping). + AliasAllocationMapT aliasToAllocations; }; //===----------------------------------------------------------------------===// @@ -529,12 +613,18 @@ } // Check that the control flow structures are supported. - if (!validateSupportedControlFlow(func.getRegion())) { + if (!validateSupportedControlFlow(func.getRegion())) return signalPassFailure(); - } - // Place all required temporary clone and dealloc nodes. + // Gather all required allocation nodes and prepare the deallocation phase. BufferDeallocation deallocation(func); + + // Check for supported AllocationOpInterface implementations and prepare the + // internal deallocation pass. + if (failed(deallocation.prepare())) + return signalPassFailure(); + + // Place all required temporary clone and dealloc nodes. deallocation.deallocate(); } };