[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

104–121

537–547

[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

104–121

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.

115

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"?

538

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

104–121

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

207

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

546

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

210

568–572

587–591

624

Prefer using the global failed/success functions.

[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

211

?

555

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.

558–559

574–576

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

579

[dfki-mako]

Addressed reviewer comments and enhanced error handling.

mlir/lib/Transforms/BufferDeallocation.cpp

211

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

555

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

211

(Thanks, I'm blind)

555

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.

574–576

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

555

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

574–576

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

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() {
     // Add additional clones that are required.
     introduceClones();
     // Place deallocations for all allocation entries.
     placeDeallocs();
   }
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))
+        continue;
+          .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();
void introduceBlockArgCopy(BlockArgument blockArg) {
     if (it == successorRegions.end())
       return;
 
     // 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;
 
     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
   /// terminators and copies the source values into the newly allocated buffers.
   void introduceValueCopyForRegionResult(Value value) {
     // Get the actual result index in the scope of the parent terminator.
     Operation *operation = value.getDefiningOp();
     auto regionInterface = cast<RegionBranchOpInterface>(operation);
for (Region &region : regions) {
         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;
         // Wire clone and terminator operand.
-        terminatorOperands.slice(operandIndex, 1).assign(clone);
+        terminatorOperands.slice(operandIndex, 1).assign(*clone);
       });
     }
   }
 
   /// 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 {
+  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
     // 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);
-      }
-    }
+        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.
+
+  /// 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);
+
+    // 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();
   }
 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::createBufferDeallocationPass() {
   return std::make_unique<BufferDeallocationPass>();
 }