[mlir] Added new AllocationOpInterface and adapted BufferDeallocation pass to use the new interface.

[dfki-mako]

Summary

The motivation behind this CL is the fact that the current BufferDeallocation pass supports building of tightly coupled DellocOp/CloneOp operations only. There have been many requests from the community to relax these restrictions in order to apply the BufferDeallocation pass to more domains and dialects. For this reason, this CL includes a newly introduced AllocationOpInterface to give developers the opportunity to implement their own "compatible" clone and deallocation operations.

• Introduced AllocationOpInterface to create deallocation operations on-the-fly that are compatible with the allocation operation implementing this interface.
  • Added interface implementations for AllocOp and CloneOp defined in the MemRef dialect.
  • Adapted the BufferDeallocation pass to be compatible with the interface introduced in this CL.

[rriddle]

Can you also include the "why" in your commit description? The motivation isn't clear from the patch description.

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

23–40

Why an OpInterface instead of a dialect interface? This kind of feels off for an op level granularity, and seems to be more related to transformation.

31

Did you intend for this method to be static?

36–37

Is this failable?

mlir/lib/Transforms/BufferDeallocation.cpp

108–125

549–559

[rriddle]

Thanks for the patch!

[herhut]

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

36

Looking at this change, I realized we likely also need a buildClone method here.

mlir/lib/Transforms/BufferDeallocation.cpp

108–125

With this only the first operation that is lacking the interface is reported. Is that generally preferred in MLIR, abort on first failure? Otherwise I find the current version more helpful.

119

As this iterates over all operations, maybe phrase this in terms of the operation? Something like `allocated value is not deallocated nor does the operation implement the AllocOpInterface"?

550

I think it would be better to find the alloc that was cloned here. Do we have the corresponding aliasing information available here?

While currently this can only be a clone, at some point we will need to support different clone operations for the different allocs, for the same reason that this supports different free operations.

[dfki-mako]

Added new buildClone operation to the newly introduced AllocationOpInterface.
Made Op interface methods static.
Added support for automatic building of "compatible" deallocation/clone operations based on inverse alias information.
Added support for building "default" DeallocOp/CloneOp values in the case of unknown value sources (e.g. function parameters).
Refined verification part to work on logical result values.

[dfki-mako]

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

23–40

We have also been thinking about this question before. Implementing a dialect interface seems to be more suitable on first sight. However, similar (if not even the same...) questions arise in this scope: should deallocation/clone operations be created on the type level rather than the Op/Dialect level? Should we support cases in which the operations can decide on different dealloc/clone ops depending on the context?

As you have also mentioned, an OP interface feels more related to transformations in general. That's why we have decided to go for the Op interface first. However, this is TBD :-)

36–37

An excellent questions. Currently, we assume that this operation cannot "fail" in any case, as we are implementing this interface for allocation operations that support these functions only. However, you can argue that these methods should have the ability to fail in general. In this case, we only need to adapt the BufferDeallocation pass to report proper errors :-) What do you think?

mlir/lib/Transforms/BufferDeallocation.cpp

108–125

We assumed that reporting all errors (instead of the first one) should be the preferred way of error reporting. However, we addressed @rriddle's comment by reporting the first error only for now.

[herhut]

Looks good to me. Please also wait for @rriddle.

@aartbik would this allow using the deallocation pass with sparse or would more be needed?

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

34

nit: missing ). Also, do you mean that the value is a result of the current Op?

42

nit: missing )

mlir/lib/Transforms/BufferDeallocation.cpp

211

You could use AllocationOpInterface definingOp = alloc.getDefiningOp<AllocationOpInterface>() here to simplify the code below.

603

Is this abstraction worth the reuse you get?

[rriddle]

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

36–37

If the method can't fail, it feels a tad weird for an interface that is named as generally as AllocationOpInterface. I wouldn't expect an alloca operation to have a dealloc operation, but I would still consider it an "AllocationOp". Is there a better name for this interface? Or, how do you see the functionality of this interface growing/changing? Right now it seems very limited in focus. If we see this interface growing as a general AllocOp interface, I would make this method fail-able.

mlir/lib/Transforms/BufferDeallocation.cpp

214

625–629

640

Prefer using the global failed/success functions.

644–648

[rriddle]

Sorry for the wait, coming back from OOO.

[dfki-mako]

Addressed reviewer comments and added support for failing interface functions.

[dfki-mako]

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

36–37

After rethinking the design of our interface proposal, we also agree that these methods should be fail-able. Therefore, we changed their returns types to Optional<T> values.

[rriddle]

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

25–26

Can you reword this to better capture what this interface is intended to represent in general? If the deallocation operations are optional, I imagine we want a more general description here.

36

39

Do you want to add a default implementation to these methods that returns None?

46

mlir/lib/Transforms/BufferDeallocation.cpp

215

?

612

This will result in a crash/segfault. Please use an appropriate error handling here, e.g. assert, llvm::report_fatal_error, propagate the error up, etc.

615–616

631–633

Same as above, please add proper error handling after the error.

636

[dfki-mako]

Addressed reviewer comments and enhanced error handling.

mlir/lib/Transforms/BufferDeallocation.cpp

215

Hmm... I am not sure, as alloc is a Value and not an Operation*..?

612

We enhanced our error handling logic to avoid crashes. However, the currently implemented method simply skips certain parts of the program without actually signaling a pass failure on the function level. Do you suggest that the pass should also signal a pass failure in these cases?

[rriddle]

mlir/lib/Transforms/BufferDeallocation.cpp

215

(Thanks, I'm blind)

612

It's up to you how to handle the error. If you do a runtime recoverable error, you'll need to propagate failure all the way up to the top of the pass and signalPassFailure(in general I wouldn't expect the compiler to keep running after an error). If you don't expect this to be recoverable at runtime (i.e. you don't expect this type of input), I'd just assert/report_fatal_error/etc right after the error.

631–633

Is the cast to logical result really necessary here?

[dfki-mako]

Updated error handling to propagate errors upwards in order to signal pass failures.

mlir/lib/Transforms/BufferDeallocation.cpp

612

All right. Consequently, we adjusted the error propagation to abort on the first failure during the transformation by signaling a pass failure.

631–633

The reason for this cast operation is the fact the we need to explicitly convert an mlir::InFlightDiagnostics result to an mlir::FailureOr<Value> instance.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

def AssumeAlignmentOp : MemRef_Op<"assume_alignment"> {
 
   let assemblyFormat = "$memref `,` $alignment attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//
 
-def MemRef_AllocOp : AllocLikeOp<"alloc", DefaultResource> {
+def MemRef_AllocOp : AllocLikeOp<"alloc", DefaultResource, [
+    DeclareOpInterfaceMethods<AllocationOpInterface,
+      ["buildDealloc", "buildClone"]>]
+  > {
   let summary = "memory allocation operation";
   let description = [{
     The `alloc` operation allocates a region of memory, as specified by its
     memref type.
 
     Example:
 
     ```mlir

 }
 
 //===----------------------------------------------------------------------===//
 // CloneOp
 //===----------------------------------------------------------------------===//
 
 def CloneOp : MemRef_Op<"clone", [
     CopyOpInterface,
-    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
+    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+    DeclareOpInterfaceMethods<AllocationOpInterface,
+      ["buildDealloc", "buildClone"]>
   ]> {
   let builders = [
     OpBuilder<(ins "Value":$value), [{
       return build($_builder, $_state, value.getType(), value);
     }]>];
 
   let description = [{
     Clones the data in the input view into an implicitly defined output view.

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_SIDEEFFECTS
 #define MLIR_INTERFACES_SIDEEFFECTS
 
 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
 //===----------------------------------------------------------------------===//
 
 // This def represents the definition for the memory effects interface. Users
 // should generally not use this directly, and should instead use
 // `MemoryEffects`.
 def MemoryEffectsOpInterface
     : EffectOpInterfaceBase<"MemoryEffectOpInterface",

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

for (Operation *user : llvm::make_early_inc_range(alloc->getUsers()))
       rewriter.eraseOp(user);
 
     rewriter.eraseOp(alloc);
     return success();
   }
 };
 } // end anonymous namespace.
 
+Optional<Operation *> AllocOp::buildDealloc(OpBuilder &builder, Value alloc) {
+  return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
+      .getOperation();
+}
+
+Optional<Value> AllocOp::buildClone(OpBuilder &builder, Value alloc) {
+  return builder.create<memref::CloneOp>(alloc.getLoc(), alloc).getResult();
+}
+
 void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
   results.add<SimplifyAllocConst<AllocOp>, SimplifyDeadAlloc<AllocOp>>(context);
 }
 
 void AllocaOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
   results.add<SimplifyAllocConst<AllocaOp>, SimplifyDeadAlloc<AllocaOp>>(
void CloneOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                           MLIRContext *context) {
   results.insert<SimplifyClones>(context);
 }
 
 OpFoldResult CloneOp::fold(ArrayRef<Attribute> operands) {
   return succeeded(foldMemRefCast(*this)) ? getResult() : Value();
 }
 
+Optional<Operation *> CloneOp::buildDealloc(OpBuilder &builder, Value alloc) {
+  return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
+      .getOperation();
+}
+
+Optional<Value> CloneOp::buildClone(OpBuilder &builder, Value alloc) {
+  return builder.create<memref::CloneOp>(alloc.getLoc(), alloc).getResult();
+}
+
 //===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult DeallocOp::fold(ArrayRef<Attribute> cstOperands,
                               SmallVectorImpl<OpFoldResult> &results) {
   /// dealloc(memrefcast) -> dealloc
   return foldMemRefCast(*this);

mlir/lib/Transforms/BufferDeallocation.cpp

 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/BufferUtils.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SetOperations.h"
 
 using namespace mlir;
 
 /// Walks over all immediate return-like terminators in the given region.
-static void walkReturnOperations(Region *region,
-                                 std::function<void(Operation *)> func) {
+static LogicalResult
+walkReturnOperations(Region *region,
+                     std::function<LogicalResult(Operation *)> func) {
   for (Block &block : *region) {
     Operation *terminator = block.getTerminator();
     // Skip non region-return-like terminators.
-    if (isRegionReturnLike(terminator))
-      func(terminator);
+    if (isRegionReturnLike(terminator)) {
+      if (failed(func(terminator)))
+        return failure();
     }
   }
+  return success();
+}
 
 /// Checks if all operations in a given region that have at least one attached
 /// region implement the RegionBranchOpInterface. This is not required in edge
 /// cases, where we have a single attached region and the parent operation has
 /// no results.
 static bool validateSupportedControlFlow(Region &region) {
   bool success = true;
   region.walk([&success](Operation *operation) {
if (((size == 1 && !operation->getResults().empty()) || size > 1) &&
       success = false;
     }
   });
   return success;
 }
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Backedges analysis
 //===----------------------------------------------------------------------===//
 
 /// A straight-forward program analysis which detects loop backedges induced by
 /// explicit control flow.
 class Backedges {
 public:
   using BlockSetT = SmallPtrSet<Block *, 16>;
   using BackedgeSetT = llvm::DenseSet<std::pair<Block *, Block *>>;
 
 public:
   /// Constructs a new backedges analysis using the op provided.
   Backedges(Operation *op) { recurse(op); }
 
   /// Returns the number of backedges formed by explicit control flow.
   size_t size() const { return edgeSet.size(); }
 
   /// Returns the start iterator to loop over all backedges.
   BackedgeSetT::const_iterator begin() const { return edgeSet.begin(); }
 
   /// Returns the end iterator to loop over all backedges.
   BackedgeSetT::const_iterator end() const { return edgeSet.end(); }
 
 private:
   /// Enters the current block and inserts a backedge into the `edgeSet` if we

 
 //===----------------------------------------------------------------------===//
 // BufferDeallocation
 //===----------------------------------------------------------------------===//
 
 /// 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<Value, AllocationOpInterface>;
+
   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<AllocationOpInterface>();
+      // 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() {
+  LogicalResult deallocate() {
     // Add additional clones that are required.
-    introduceClones();
+    if (failed(introduceClones()))
+      return failure();
+
     // Place deallocations for all allocation entries.
-    placeDeallocs();
+    return placeDeallocs();
   }
 
 private:
   /// Introduces required clone operations to avoid memory leaks.
-  void introduceClones() {
+  LogicalResult introduceClones() {
     // Initialize the set of values that require a dedicated memory free
     // operation since their operands cannot be safely deallocated in a post
     // dominator.
     SmallPtrSet<Value, 8> valuesToFree;
     llvm::SmallDenseSet<std::tuple<Value, Block *>> visitedValues;
     SmallVector<std::tuple<Value, Block *>, 8> toProcess;
 
     // Check dominance relation for proper dominance properties. If the given
LogicalResult introduceClones() {
     }
 
     // Update buffer aliases to ensure that we free all buffers and block
     // arguments at the correct locations.
     aliases.remove(valuesToFree);
 
     // Add new allocs and additional clone operations.
     for (Value value : valuesToFree) {
-      if (auto blockArg = value.dyn_cast<BlockArgument>())
-        introduceBlockArgCopy(blockArg);
-      else
-        introduceValueCopyForRegionResult(value);
+      if (failed(value.isa<BlockArgument>()
+                     ? introduceBlockArgCopy(value.cast<BlockArgument>())
+                     : introduceValueCopyForRegionResult(value)))
+        return failure();
 
       // Register the value to require a final dealloc. Note that we do not have
       // to assign a block here since we do not want to move the allocation node
       // to another location.
       allocs.registerAlloc(std::make_tuple(value, nullptr));
     }
+    return success();
   }
 
   /// Introduces temporary clones in all predecessors and copies the source
   /// values into the newly allocated buffers.
-  void introduceBlockArgCopy(BlockArgument blockArg) {
+  LogicalResult introduceBlockArgCopy(BlockArgument blockArg) {
     // Allocate a buffer for the current block argument in the block of
     // the associated value (which will be a predecessor block by
     // definition).
     Block *block = blockArg.getOwner();
     for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
       // Get the terminator and the value that will be passed to our
       // argument.
       Operation *terminator = (*it)->getTerminator();
       auto branchInterface = cast<BranchOpInterface>(terminator);
       // Query the associated source value.
       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))
+        return failure();
+          .assign(*clone);
     }
 
     // Check whether the block argument has implicitly defined predecessors via
     // the RegionBranchOpInterface. This can be the case if the current block
     // argument belongs to the first block in a region and the parent operation
     // implements the RegionBranchOpInterface.
     Region *argRegion = block->getParent();
     Operation *parentOp = argRegion->getParentOp();
     RegionBranchOpInterface regionInterface;
     if (!argRegion || &argRegion->front() != block ||
         !(regionInterface = dyn_cast<RegionBranchOpInterface>(parentOp)))
-      return;
+      return success();
 
-    introduceClonesForRegionSuccessors(
+    if (failed(introduceClonesForRegionSuccessors(
             regionInterface, argRegion->getParentOp()->getRegions(), blockArg,
             [&](RegionSuccessor &successorRegion) {
               // Find a predecessor of our argRegion.
               return successorRegion.getSuccessor() == argRegion;
-        });
+            })))
+      return failure();
 
     // Check whether the block argument belongs to an entry region of the
     // parent operation. In this case, we have to introduce an additional clone
     // for buffer that is passed to the argument.
     SmallVector<RegionSuccessor, 2> successorRegions;
     regionInterface.getSuccessorRegions(/*index=*/llvm::None, successorRegions);
     auto *it =
         llvm::find_if(successorRegions, [&](RegionSuccessor &successorRegion) {
           return successorRegion.getSuccessor() == argRegion;
         });
     if (it == successorRegions.end())
-      return;
+      return success();
 
     // Determine the actual operand to introduce a clone for and rewire the
     // operand to point to the clone instead.
     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 failure();
 
     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);
+    return success();
   }
 
   /// Introduces temporary clones in front of all associated nested-region
   /// terminators and copies the source values into the newly allocated buffers.
-  void introduceValueCopyForRegionResult(Value value) {
+  LogicalResult introduceValueCopyForRegionResult(Value value) {
     // Get the actual result index in the scope of the parent terminator.
     Operation *operation = value.getDefiningOp();
     auto regionInterface = cast<RegionBranchOpInterface>(operation);
     // Filter successors that return to the parent operation.
     auto regionPredicate = [&](RegionSuccessor &successorRegion) {
       // If the RegionSuccessor has no associated successor, it will return to
       // its parent operation.
       return !successorRegion.getSuccessor();
     };
     // Introduce a clone for all region "results" that are returned to the
     // parent operation. This is required since the parent's result value has
     // been considered critical. Therefore, the algorithm assumes that a clone
     // of a previously allocated buffer is returned by the operation (like in
     // the case of a block argument).
-    introduceClonesForRegionSuccessors(regionInterface, operation->getRegions(),
-                                       value, regionPredicate);
+    return introduceClonesForRegionSuccessors(
+        regionInterface, operation->getRegions(), value, regionPredicate);
   }
 
   /// Introduces buffer clones for all terminators in the given regions. The
   /// regionPredicate is applied to every successor region in order to restrict
   /// the clones to specific regions.
   template <typename TPredicate>
-  void introduceClonesForRegionSuccessors(
+  LogicalResult introduceClonesForRegionSuccessors(
       RegionBranchOpInterface regionInterface, MutableArrayRef<Region> regions,
       Value argValue, const TPredicate &regionPredicate) {
     for (Region &region : regions) {
       // Query the regionInterface to get all successor regions of the current
       // one.
       SmallVector<RegionSuccessor, 2> successorRegions;
       regionInterface.getSuccessorRegions(region.getRegionNumber(),
                                           successorRegions);
       // Try to find a matching region successor.
       RegionSuccessor *regionSuccessor =
           llvm::find_if(successorRegions, regionPredicate);
       if (regionSuccessor == successorRegions.end())
         continue;
       // Get the operand index in the context of the current successor input
       // bindings.
       size_t operandIndex =
           llvm::find(regionSuccessor->getSuccessorInputs(), argValue)
               .getIndex();
 
       // Iterate over all immediate terminator operations to introduce
       // new buffer allocations. Thereby, the appropriate terminator operand
       // will be adjusted to point to the newly allocated buffer instead.
-      walkReturnOperations(&region, [&](Operation *terminator) {
+      if (failed(walkReturnOperations(&region, [&](Operation *terminator) {
             // Get the actual mutable operands for this terminator op.
             auto terminatorOperands = *getMutableRegionBranchSuccessorOperands(
                 terminator, region.getRegionNumber());
             // Extract the source value from the current terminator.
             // This conversion needs to exist on a separate line due to a bug in
             // GCC conversion analysis.
             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 failure();
             // Wire clone and terminator operand.
-        terminatorOperands.slice(operandIndex, 1).assign(clone);
-      });
+            terminatorOperands.slice(operandIndex, 1).assign(*clone);
+            return success();
+          })))
+        return failure();
     }
+    return success();
   }
 
   /// 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<Value> 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
     // copying copied buffers can introduce memory leaks since the invariant of
     // BufferDeallocation assumes that a buffer will be only cloned once into a
     // temporary buffer. Hence, the construction of clone chains introduces
     // additional allocations that are not tracked automatically by the
     // algorithm.
     if (clonedValues.contains(sourceValue))
       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<memref::CloneOp>(terminator->getLoc(), sourceValue);
-
+    auto clone = buildClone(terminator, sourceValue);
+    if (succeeded(clone)) {
       // Remember the clone of original source value.
-    clonedValues.insert(cloneOp);
-    return cloneOp;
+      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 {
+  LogicalResult 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
     // liveness.
     for (const BufferPlacementAllocs::AllocEntry &entry : allocs) {
       Value alloc = std::get<0>(entry);
       auto aliasesSet = aliases.resolve(alloc);
       assert(!aliasesSet.empty() && "must contain at least one alias");
for (const BufferPlacementAllocs::AllocEntry &entry : allocs) {
         deallocOperation->moveAfter(endOperation);
       } else {
         // If the Dealloc position is at the terminator operation of the
         // block, then the value should escape from a deallocation.
         Operation *nextOp = endOperation->getNextNode();
         if (!nextOp)
           continue;
         // If there is no dealloc node, insert one in the right place.
-        OpBuilder builder(nextOp);
-        builder.create<memref::DeallocOp>(alloc.getLoc(), alloc);
-      }
+        if (failed(buildDealloc(nextOp, alloc)))
+          return failure();
+      }
+    }
+    return success();
+  }
+
+  /// 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.
+  LogicalResult 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)
+        return op->emitError()
+               << "allocations without compatible deallocations are "
+                  "not supported";
+    } else {
+      // Build a "default" DeallocOp for unknown allocation sources.
+    return success();
+  }
+
+  /// 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<Value> 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<memref::CloneOp>(alloc.getLoc(), alloc).getResult();
   }
 
   /// The dominator info to find the appropriate start operation to move the
   /// allocs.
   DominanceInfo dominators;
 
   /// The post dominator info to move the dependent allocs in the right
   /// position.
   PostDominanceInfo postDominators;
 
   /// Stores already cloned buffers to avoid additional clones of clones.
   ValueSetT clonedValues;
+
+  /// Maps aliases to their source allocation interfaces (inverse mapping).
+  AliasAllocationMapT aliasToAllocations;
 };
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationPass
 //===----------------------------------------------------------------------===//
 
 /// The actual buffer deallocation pass that inserts and moves dealloc nodes
 /// into the right positions. Furthermore, it inserts additional clones if
 /// necessary. It uses the algorithm described at the top of the file.
 struct BufferDeallocationPass : BufferDeallocationBase<BufferDeallocationPass> {
 
   void runOnFunction() override {
     // Ensure that there are supported loops only.
     FuncOp func = getFunction();
     Backedges backedges(func);
     if (backedges.size()) {
       func.emitError("Only structured control-flow loops are supported.");
       return signalPassFailure();
     }
 
     // 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);
-    deallocation.deallocate();
+
+    // 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.
+    if (failed(deallocation.deallocate()))
+      return signalPassFailure();
   }
 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::createBufferDeallocationPass() {
   return std::make_unique<BufferDeallocationPass>();
 }