[mlir] Extended BufferPlacement to support more sophisticated scenarios in which allocations cannot be moved freely and can remain in divergent control flow.

[dfki-mako]

Summary

The current BufferPlacement pass does not support allocation nodes that carry additional dependencies (like in the case of dynamic shaped types). These allocations can often not be moved freely and in turn might remain in divergent control-flow branches. This requires a different strategy with respect to block arguments and aliases. This CL adds additional functionality to support allocation nodes in divergent control flow while avoiding memory leaks.

[herhut]

Can you add tests for dynamic allocation cases, as well?

mlir/lib/Transforms/BufferPlacement.cpp

11

Incomplete sentence to ensure that all buffers

37

In other words, the analysis uses the post-dominator to free the allocated memory. If no such post-dominator exists, aliases are removed by inserting allocs and copies. Is that the idea?

45

Nit: Remove highly.

So alloc works the same way but this time with a common dominator? If no such dominator exists, copies are inserted?

167

Optimizations should be independent of BufferPlacement, so this TODO does not belong here.

186

Would it make sense to compute this once and share it between the different phases? The moving of allocs does not influence the aliasing.

203

This should not be an assert. Instead either ignore the alloc (which AFAIK was the mode before) or fail the transformation and report an error.

Generally, the reason this was not supported before was that multiple results creates situations where the alloc cannot be moved. Can the new system not tolerate this with copies?

I am OK with this not being supported, as a multi-result alloc seems a very boutique use.

217

Allocation node with know shape? Another way to look at it would be to check whether the allocation has any operands, which structurally is the reason it cannot be moved.

227

Same here, this could just use the operands of the alloc.

229

It is a clever reuse of findPlacementBlock. Maybe it should no longer talk about aliases. Instead, it is a helper that finds a dominator/postdominator for a range of values.

233

Is this dependency information that you are lacking?

281

Would it make sense to store BlockArgument to start with? Aliases are always BlockArguments, right?

284

This breaks the aliasing between source and value. If there are further aliases of source that alias via value, they would no longer be an alias. But we would still potentially insert a copy to break the aliasing?

293

This does not account for nested blocks in nested regions. Also, why is this needed?

313

Somewhat sad that MutableOperandRange does not support to get values out...

388

Why would we want to insert before the end operation? That would free it before the last use?

446

a an -> an

mlir/test/Transforms/buffer-placement.mlir

87

The numbering is strange of ALLOC02

211

I cannot follow this without the block numbers in the checks. Why do we get copies here?

309

Same here.

[dfki-mako]

Implemented a fix-point iteration to reduce the number of required copies.
Fixed minor issues (see review comments).

[dfki-mako]

mlir/lib/Transforms/BufferPlacement.cpp

37

Yep, that's the general idea.

203

In principal yes; however, we currently remove all Dealloc nodes in the beginning of the transformation. If we keep them for unsupported Alloc ops, this should not be an issue.

293

Obsolete; we have added a new fix-point iteration.

313

Unfortunately, it seems to be the only way at the moment...

mlir/test/Transforms/buffer-placement.mlir

87

The name ALLOC02 should indicate that the allocations ALLOC0 and ALLOC2 will be freed at this location.

211

The updated algorithm does not insert any copies in these cases any more.

[mehdi_amini]

mlir/lib/Transforms/BufferPlacement.cpp

203

Can we be *safe* in this case instead of failing entirely?

[dfki-mako]

mlir/lib/Transforms/BufferPlacement.cpp

203

@mehdi_amini If we do not remove Dealloc nodes (or nodes with free semantics) that are related to unsupported alloc nodes, we can be sure that the program will not be "worse" than before with respect to these buffers.

[herhut]

mlir/lib/Transforms/BufferPlacement.cpp

242

Nit: Remove the { }.

255

Aren't there two cases here? If this immediate alias is not dominated, then a copy is inserted, the alias turns effectively into an allocation and also needs to be processed. If we do not insert a copy, then the alias still has to be processed, using the original allocation as the block that needs to dominate.

I think the second case is missing here.

[mehdi_amini]

mlir/lib/Transforms/BufferPlacement.cpp

203

Right, so we shouldn't need to emit an error here then? We can just skip here?

[dfki-mako]

Fixed corner cases in which no copies have been created due to aliasing (see reviewer comments).
Added additional test cases to simulate such special control-flow scenarios.
Changed behavior of the BufferPlacement pass to ignore Alloc operations with multiple allocation results.

[dfki-mako]

mlir/lib/Transforms/BufferPlacement.cpp

203

The latest version skips these allocation nodes without aborting the transformation process.

255

We have added a separate test case to simulate such a case and changed the CL accordingly.

[herhut]

mlir/lib/Transforms/BufferPlacement.cpp

177

Nit: Finds -> Find?

250

With this change, you consider all transitive aliases of a value. However, as soon as you insert a copy (by adding a blockarg to the list of blocks to free) that set would need updating, as you now have fewer transitive aliases.

In you example, consider the case where there is another aliasing block (say b7) after b6. Then we would get a copy from b5 to b6, as b2 does not dominate b6. This would also take care of the block b7 after b6, as long as b6 dominates it. However, with the current approach, you would still check whether b2 dominates this additional block b7 and like insert a copy from b6 to b7.

I think that, when adding a blockarg to the free list, you have to use that block as the dominating block (as it now becomes the allocating place) when processing aliases of the new freed blockarg.

305

I do not understand what the block argument refers to here.

314

Nit: delloc -> dealloc

mlir/test/Transforms/buffer-placement.mlir

146

Could you add all block numbers here. That makes it easier to follow where the br operations belong.

[pifon2a]

mlir/test/Transforms/buffer-placement.mlir

54

can this (and all other comments and also IR below) fit 80 chars?

[dfki-mako]

Updated processing of aliases to avoid the aggressive insertion of copy operations in special cases.
Reformatted test cases.

[herhut]

Thanks. This looks correct to me now.

You mentioned offline that you added a test for the case that was not covered before. I did not find it in this CL though.

mlir/lib/Transforms/BufferPlacement.cpp

241

You could also populate toProcess here and then just have the loop below. This works for me too, though.

[dfki-mako]

Resolved merge conflicts and integrated an additional test case into an existing one in order to capture an additional special case.

[dfki-mako]

mlir/lib/Transforms/BufferPlacement.cpp

241

This would also work. However, we would prefer the more explicit version to help others understand the code more easily.

[herhut]

Thanks for expanding the test.

mlir/lib/Transforms/BufferPlacement.cpp

 //===- BufferPlacement.cpp - the impl for buffer placement ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for computing correct alloc and dealloc positions.
-// The main class is the BufferPlacementPass class that implements the
-// underlying algorithm. In order to put allocations and deallocations at safe
-// positions, it is significantly important to put them into the correct blocks.
-// However, the liveness analysis does not pay attention to aliases, which can
-// occur due to branches (and their associated block arguments) in general. For
-// this purpose, BufferPlacement firstly finds all possible aliases for a single
-// value (using the BufferPlacementAliasAnalysis class). Consider the following
-// example:
+// Furthermore, buffer placement also adds required new alloc and copy
+// operations to ensure that all buffers The main class is the
+// BufferPlacementPass class that implements the underlying algorithm. In order
+// to put allocations and deallocations at safe positions, it is significantly
+// important to put them into the correct blocks. However, the liveness analysis
+// does not pay attention to aliases, which can occur due to branches (and their
+// associated block arguments) in general. For this purpose, BufferPlacement
+// firstly finds all possible aliases for a single value (using the
+// BufferPlacementAliasAnalysis class). Consider the following example:
 //
 // ^bb0(%arg0):
 //   cond_br %cond, ^bb1, ^bb2
 // ^bb1:
 //   br ^exit(%arg0)
 // ^bb2:
 //   %new_value = ...
 //   br ^exit(%new_value)
 // ^exit(%arg1):
 //   return %arg1;
 //
 // Using liveness information on its own would cause us to place the allocs and
 // deallocs in the wrong block. This is due to the fact that %new_value will not
-// be liveOut of its block. Instead, we have to place the alloc for %new_value
-// in bb0 and its associated dealloc in exit. Using the class
-// BufferPlacementAliasAnalysis, we will find out that %new_value has a
-// potential alias %arg1. In order to find the dealloc position we have to find
-// all potential aliases, iterate over their uses and find the common
-// post-dominator block. In this block we can safely be sure that %new_value
-// will die and can use liveness information to determine the exact operation
-// after which we have to insert the dealloc. Finding the alloc position is
-// highly similar and non- obvious. Again, we have to consider all potential
-// aliases and find the common dominator block to place the alloc.
+// be liveOut of its block. Instead, we can place the alloc for %new_value
+// in bb0 and its associated dealloc in exit. Alternatively, the alloc can stay
+// (or even has to stay due to additional dependencies) at this location and we
+// have to free the buffer in the same block, because it cannot be freed in the
+// post dominator. However, this requires a new copy buffer for %arg1 that will
+// contain the actual contents. Using the class BufferPlacementAliasAnalysis, we
+// will find out that %new_value has a potential alias %arg1. In order to find
+// the dealloc position we have to find all potential aliases, iterate over
+// their uses and find the common post-dominator block (note that additional
+// copies and buffers remove potential aliases and will influence the placement
+// of the deallocs). In all cases, the computed block can be safely used to free
+// the %new_value buffer (may be exit or bb2) as it will die and we can use
+// liveness information to determine the exact operation after which we have to
+// insert the dealloc. Finding the alloc position is highly similar and non-
+// obvious. However, the algorithm supports moving allocs to other places and
+// introducing copy buffers and placing deallocs in safe places to ensure that
+// all buffers will be freed in the end.
 //
 // TODO:
 // The current implementation does not support loops and the resulting code will
 // be invalid with respect to program semantics. The only thing that is
 // currently missing is a high-level loop analysis that allows us to move allocs
 // and deallocs outside of the loop blocks.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/BufferPlacement.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // BufferPlacementAliasAnalysis
 //===----------------------------------------------------------------------===//
 
 /// A straight-forward alias analysis which ensures that all aliases of all
 /// values will be determined. This is a requirement for the BufferPlacement
 /// class since you need to determine safe positions to place alloc and
 /// deallocs.
 class BufferPlacementAliasAnalysis {
 public:
   using ValueSetT = SmallPtrSet<Value, 16>;
+  using ValueMapT = llvm::DenseMap<Value, ValueSetT>;
 
 public:
   /// Constructs a new alias analysis using the op provided.
   BufferPlacementAliasAnalysis(Operation *op) { build(op->getRegions()); }
 
+  /// Finds all immediate aliases this value could potentially have.
+  ValueMapT::const_iterator find(Value value) const {
+    return aliases.find(value);
+  }
+
+  /// Returns the end iterator that can be used in combination with find.
+  ValueMapT::const_iterator end() const { return aliases.end(); }
+
   /// Finds all immediate and indirect aliases this value could potentially
   /// have. Note that the resulting set will also contain the value provided as
   /// it is an alias of itself.
   ValueSetT resolve(Value value) const {
     ValueSetT result;
     resolveRecursive(value, result);
     return result;
   }
for (Region &region : regions) {
             }
           }
         }
       }
     }
   }
 
   /// Maps values to all immediate aliases this value can have.
-  llvm::DenseMap<Value, ValueSetT> aliases;
+  ValueMapT aliases;
 };
 
 //===----------------------------------------------------------------------===//
-// BufferPlacementPositions
+// BufferPlacement
 //===----------------------------------------------------------------------===//
 
-/// Stores correct alloc and dealloc positions to place dialect-specific alloc
-/// and dealloc operations.
-struct BufferPlacementPositions {
+// The main buffer placement analysis used to place allocs, copies and deallocs.
+class BufferPlacement {
 public:
-  BufferPlacementPositions()
-      : allocPosition(nullptr), deallocPosition(nullptr) {}
+  using AllocEntry = std::pair<Value, Block *>;
+  using AllocEntryList = SmallVector<AllocEntry, 8>;
 
-  /// Creates a new positions tuple including alloc and dealloc positions.
-  BufferPlacementPositions(Operation *allocPosition, Operation *deallocPosition)
-      : allocPosition(allocPosition), deallocPosition(deallocPosition) {}
-
-  /// Returns the alloc position before which the alloc operation has to be
-  /// inserted.
-  Operation *getAllocPosition() const { return allocPosition; }
-
-  /// Returns the dealloc position after which the dealloc operation has to be
-  /// inserted.
-  Operation *getDeallocPosition() const { return deallocPosition; }
+public:
+  BufferPlacement(Operation *op)
+      : operation(op), liveness(op), dominators(op), postDominators(op) {
+    // Gather all allocation nodes
+    initBlockMapping();
+    // TODO: optimize allocations (fuse, move and reuse buffers).
+  }
+
+  /// Performs the actual placement/creation of all alloc, copy and dealloc
+  /// nodes.
+  void place() {
+    // Place all allocations.
+    placeAllocs();
+    // Add additional allocations and copies that are required.
+    introduceCopies();
+    // Place deallocations for all allocs.
+    placeDeallocs();
+  }
 
 private:
-  Operation *allocPosition;
-  Operation *deallocPosition;
-};
-
-//===----------------------------------------------------------------------===//
-// BufferPlacementAnalysis
-//===----------------------------------------------------------------------===//
+  /// Initializes the internal block mapping by discovering allocation nodes. It
+  /// maps all allocation nodes to their initial block in which they can be
+  /// safely allocated.
+  void initBlockMapping() {
+    BufferPlacementAliasAnalysis aliases(operation);
+    operation->walk([&](MemoryEffectOpInterface opInterface) {
+      // Try to find a single allocation result.
+      SmallVector<MemoryEffects::EffectInstance, 2> effects;
+      opInterface.getEffects(effects);
 
-// The main buffer placement analysis used to place allocs and deallocs.
-class BufferPlacementAnalysis {
-public:
-  using DeallocSetT = SmallPtrSet<Operation *, 2>;
+      SmallVector<MemoryEffects::EffectInstance, 2> allocateResultEffects;
+      llvm::copy_if(effects, std::back_inserter(allocateResultEffects),
+                    [=](MemoryEffects::EffectInstance &it) {
+                      Value value = it.getValue();
+                      return isa<MemoryEffects::Allocate>(it.getEffect()) &&
+                             value && value.isa<OpResult>();
+                    });
+      if (allocateResultEffects.size() < 1)
+        return;
+      // If there is one result only, we will be able to move the allocation and
+      // (possibly existing) deallocation ops.
+      assert(allocateResultEffects.size() == 1 &&
+             "multiple allocations are not allowed");
+      // Get allocation result.
+      auto allocResult = allocateResultEffects[0].getValue().cast<OpResult>();
+      // Find the initial allocation block and register this result.
+      blockMapping.emplace_back(allocResult,
+                                getInitialAllocBlock(allocResult, aliases));
+    });
+  }
 
-public:
-  BufferPlacementAnalysis(Operation *op)
-      : operation(op), liveness(op), dominators(op), postDominators(op),
-        aliases(op) {}
-
-  /// Computes the actual positions to place allocs and deallocs for the given
-  /// value.
-  BufferPlacementPositions
-  computeAllocAndDeallocPositions(OpResult result) const {
-    if (result.use_empty())
-      return BufferPlacementPositions(result.getOwner(), result.getOwner());
-    // Get all possible aliases.
-    auto possibleValues = aliases.resolve(result);
-    return BufferPlacementPositions(getAllocPosition(result, possibleValues),
-                                    getDeallocPosition(result, possibleValues));
-  }
-
-  /// Finds all associated dealloc nodes for the alloc nodes using alias
-  /// information.
-  DeallocSetT findAssociatedDeallocs(OpResult allocResult) const {
-    DeallocSetT result;
-    auto possibleValues = aliases.resolve(allocResult);
-    for (Value alias : possibleValues)
-      for (Operation *op : alias.getUsers()) {
-        // Check for an existing memory effect interface.
-        auto effectInstance = dyn_cast<MemoryEffectOpInterface>(op);
-        if (!effectInstance)
-          continue;
-        // Check whether the associated value will be freed using the current
-        // operation.
-        SmallVector<MemoryEffects::EffectInstance, 2> effects;
-        effectInstance.getEffectsOnValue(alias, effects);
-        if (llvm::any_of(effects, [=](MemoryEffects::EffectInstance &it) {
-              return isa<MemoryEffects::Free>(it.getEffect());
-            }))
-          result.insert(op);
-      }
-    return result;
+  /// Computes a valid allocation position in a dominator (if possible) for the
+  /// given allocation result.
+  Block *getInitialAllocBlock(OpResult result,
+                              BufferPlacementAliasAnalysis aliasAnalysis) {
+    // Check whether this is a known allocation node.
+    // TODO: remove this dependency by introducing a generic interface.
+    if (auto alloc = dyn_cast<AllocOp>(result.getOwner())) {
+      auto dynamicSizes = alloc.getDynamicSizes();
+      if (dynamicSizes.size() < 1)
+        return findPlacementBlock(result, aliasAnalysis.resolve(result),
+                                  dominators);
+
+      // If this node has dependencies, check all dependent nodes with respect
+      // to a common post dominator in which all values are available.
+      BufferPlacementAliasAnalysis::ValueSetT dependencies(
+          ++dynamicSizes.begin(), dynamicSizes.end());
+      return findPlacementBlock(*dynamicSizes.begin(), dependencies,
+                                postDominators);
+    }
+
+    // We do not have access to detailed dependency information. Hence, we place
+    // the alloc operation in the same block to ensure a valid program.
+    return result.getParentBlock();
+  }
+
+  /// Finds correct alloc positions according to the algorithm described at
+  /// the top of the file for all alloc nodes that can be handled by this
+  /// analysis.
+  void placeAllocs() const {
+    for (auto &entry : blockMapping) {
+      Value alloc = entry.first;
+      // Get the actual block to place the alloc and get liveness information
+      // for the placement block.
+      Block *placementBlock = entry.second;
+      // We have to ensure that we place the alloc before its first use in this
+      // block.
+      const LivenessBlockInfo *livenessInfo =
+          liveness.getLiveness(placementBlock);
+      Operation *startOperation = livenessInfo->getStartOperation(alloc);
+      // Check whether the start operation lies in the desired placement block.
+      // If not, we will use the terminator as this is the last operation in
+      // this block.
+      if (startOperation->getBlock() != placementBlock)
+        startOperation = placementBlock->getTerminator();
+
+      // Move the alloc in front of the start operation.
+      Operation *allocOperation = alloc.getDefiningOp();
+      allocOperation->moveBefore(startOperation);
   }
 
-  /// Dumps the buffer placement information to the given stream.
-  void print(raw_ostream &os) const {
-    os << "// ---- Buffer Placement -----\n";
+  /// Introduces required allocs and copy operations to avoid memory leaks.
+  void introduceCopies() {
+    // Initialize the set of block arguments that require a dedicated memory
+    // free operation since their arguments cannot be safely deallocated in a
+    // post dominator.
+    SmallPtrSet<BlockArgument, 8> blockArgsToFree;
+    BufferPlacementAliasAnalysis aliases(operation);
+
+    // Check dominance relation for proper dominance properties. If the given
+    // value node does not dominate an alias, we will have to create a copy in
+    // order to free all buffers that can potentially leak into a post
+    // dominator.
+    auto findUnsafeValues = [&](Value source, Block *definingBlock) {
+      auto it = aliases.find(source);
+      if (it == aliases.end())
+        return;
+      for (Value value : it->second) {
+        auto blockArg = value.cast<BlockArgument>();
+        if (blockArgsToFree.count(blockArg) < 1 &&
+            !dominators.dominates(definingBlock, blockArg.getOwner()))
+          blockArgsToFree.insert(blockArg);
+      }
+    };
 
+    // Detect possibly unsafe aliases starting from all allocations.
+    for (auto &entry : blockMapping)
+      findUnsafeValues(entry.first, entry.second);
+
+    // Detect possibly unsafe aliases starting from all block arguments.
     for (Region &region : operation->getRegions())
       for (Block &block : region)
-        for (Operation &operation : block)
-          for (OpResult result : operation.getResults()) {
-            BufferPlacementPositions positions =
-                computeAllocAndDeallocPositions(result);
-            os << "Positions for ";
-            result.print(os);
-            os << "\n Alloc: ";
-            positions.getAllocPosition()->print(os);
-            os << "\n Dealloc: ";
-            positions.getDeallocPosition()->print(os);
-            os << "\n";
-          }
+        for (BlockArgument arg : block.getArguments())
+          findUnsafeValues(arg, &block);
+
+    // Add new allocs and additional copy operations.
+    for (BlockArgument blockArg : blockArgsToFree) {
+      Block *block = blockArg.getOwner();
+
+      // Allocate a buffer for the current block argument in the block of the
+      // associated value (which will be a predecessor block by definition).
+      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 successorOperand =
+            cast<BranchOpInterface>(terminator)
+                .getMutableSuccessorOperands(it.getSuccessorIndex())
+                .getValue()
+                .slice(blockArg.getArgNumber(), 1);
+        Value sourceValue = ((OperandRange)successorOperand)[0];
+
+        // Create a new alloc at the current location of the terminator.
+        auto memRefType = sourceValue.getType().cast<MemRefType>();
+        OpBuilder builder(terminator);
+
+        // Extract information about dynamically shaped types by extracting
+        // their dynamic dimensions.
+        SmallVector<Value, 4> dynamicOperands;
+        for (auto shapeElement : llvm::enumerate(memRefType.getShape())) {
+          if (!ShapedType::isDynamic(shapeElement.value()))
+            continue;
+          dynamicOperands.push_back(builder.create<DimOp>(
+              terminator->getLoc(), sourceValue, shapeElement.index()));
         }
 
-private:
-  /// Finds a correct placement block to store alloc/dealloc node according to
-  /// the algorithm described at the top of the file. It supports dominator and
-  /// post-dominator analyses via template arguments.
-  template <typename DominatorT>
-  Block *
-  findPlacementBlock(OpResult result,
-                     const BufferPlacementAliasAnalysis::ValueSetT &aliases,
-                     const DominatorT &doms) const {
-    // Start with the current block the value is defined in.
-    Block *dom = result.getOwner()->getBlock();
-    // Iterate over all aliases and their uses to find a safe placement block
-    // according to the given dominator information.
-    for (Value alias : aliases)
-      for (Operation *user : alias.getUsers()) {
-        // Move upwards in the dominator tree to find an appropriate
-        // dominator block that takes the current use into account.
-        dom = doms.findNearestCommonDominator(dom, user->getBlock());
-      }
-    return dom;
+        // TODO: provide a generic interface to create dialect-specific Alloc
+        // and CopyOp nodes.
+        auto alloc = builder.create<AllocOp>(terminator->getLoc(), memRefType,
+                                             dynamicOperands);
+        // Wire new alloc and successor operand.
+        successorOperand.assign(alloc);
+        // Create a new copy operation that copies to contents of the old
+        // allocation to the new one.
+        builder.create<linalg::CopyOp>(terminator->getLoc(), sourceValue,
+                                       alloc);
+      }
+
+      // Register the block argument 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.
+      blockMapping.emplace_back(blockArg, nullptr);
   }
 
-  /// Finds a correct alloc position according to the algorithm described at
-  /// the top of the file.
-  Operation *getAllocPosition(
-      OpResult result,
-      const BufferPlacementAliasAnalysis::ValueSetT &aliases) const {
-    // Determine the actual block to place the alloc and get liveness
-    // information.
-    Block *placementBlock = findPlacementBlock(result, aliases, dominators);
-    const LivenessBlockInfo *livenessInfo =
-        liveness.getLiveness(placementBlock);
+  /// 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 {
+    BufferPlacementAliasAnalysis aliases(operation);
+    // Insert "simple" deallocs first. These allocations will be linked to their
+    // associated allocation nodes since they don't have any aliases that can
+    // (potentially) increase their liveness.
+    for (auto &entry : blockMapping) {
+      Value alloc = entry.first;
+      auto aliasesSet = aliases.resolve(alloc);
+      assert(aliasesSet.size() > 0 && "must contain at least one alias");
 
-    // We have to ensure that the alloc will be before the first use of all
-    // aliases of the given value. We first assume that there are no uses in the
-    // placementBlock and that we can safely place the alloc before the
-    // terminator at the end of the block.
-    Operation *startOperation = placementBlock->getTerminator();
-    // Iterate over all aliases and ensure that the startOperation will point to
-    // the first operation of all potential aliases in the placementBlock.
-    for (Value alias : aliases) {
-      Operation *aliasStartOperation = livenessInfo->getStartOperation(alias);
-      // Check whether the aliasStartOperation lies in the desired block and
-      // whether it is before the current startOperation. If yes, this will be
-      // the new startOperation.
-      if (aliasStartOperation->getBlock() == placementBlock &&
-          aliasStartOperation->isBeforeInBlock(startOperation))
-        startOperation = aliasStartOperation;
-    }
-    // startOperation is the first operation before which we can safely store
-    // the alloc taking all potential aliases into account.
-    return startOperation;
-  }
-
-  /// Finds a correct dealloc position according to the algorithm described at
-  /// the top of the file.
-  Operation *getDeallocPosition(
-      OpResult result,
-      const BufferPlacementAliasAnalysis::ValueSetT &aliases) const {
       // Determine the actual block to place the dealloc and get liveness
       // information.
-    Block *placementBlock = findPlacementBlock(result, aliases, postDominators);
+      Block *placementBlock =
+          findPlacementBlock(alloc, aliasesSet, postDominators);
       const LivenessBlockInfo *livenessInfo =
           liveness.getLiveness(placementBlock);
 
       // We have to ensure that the dealloc will be after the last use of all
-    // aliases of the given value. We first assume that there are no uses in the
-    // placementBlock and that we can safely place the dealloc at the beginning.
+      // aliases of the given value. We first assume that there are no uses in
+      // the placementBlock and that we can safely place the dealloc at the
+      // beginning.
       Operation *endOperation = &placementBlock->front();
       // Iterate over all aliases and ensure that the endOperation will point to
       // the last operation of all potential aliases in the placementBlock.
-    for (Value alias : aliases) {
+      for (Value alias : aliasesSet) {
         Operation *aliasEndOperation =
             livenessInfo->getEndOperation(alias, endOperation);
         // Check whether the aliasEndOperation lies in the desired block and
-      // whether it is behind the current endOperation. If yes, this will be the
-      // new endOperation.
+        // whether it is behind the current endOperation. If yes, this will be
+        // the new endOperation.
         if (aliasEndOperation->getBlock() == placementBlock &&
             endOperation->isBeforeInBlock(aliasEndOperation))
           endOperation = aliasEndOperation;
       }
       // endOperation is the last operation behind which we can safely store the
       // dealloc taking all potential aliases into account.
-    return endOperation;
+      auto nextOp = endOperation->getNextNode();
+      endOperation = nextOp ? nextOp : endOperation;
+      OpBuilder builder(endOperation);
+      builder.create<DeallocOp>(alloc.getLoc(), alloc);
+    }
+  }
+
+  /// Finds a correct placement block to store alloc/dealloc node according to
+  /// the algorithm described at the top of the file. It supports dominator and
+  /// post-dominator analyses via template arguments.
+  template <typename DominatorT>
+  Block *
+  findPlacementBlock(Value value,
+                     const BufferPlacementAliasAnalysis::ValueSetT &aliases,
+                     const DominatorT &doms) const {
+    // Start with the current block the value is defined in.
+    Block *dom = value.getParentBlock();
+    // Iterate over all aliases and their uses to find a safe placement block
+    // according to the given dominator information.
+    for (Value alias : aliases)
+      for (Operation *user : alias.getUsers()) {
+        // Move upwards in the dominator tree to find an appropriate
+        // dominator block that takes the current use into account.
+        dom = doms.findNearestCommonDominator(dom, user->getBlock());
+      }
+    return dom;
   }
 
   /// The operation this transformation was constructed from.
   Operation *operation;
 
+  /// Maps allocation nodes to their associated blocks.
+  AllocEntryList blockMapping;
+
   /// The underlying liveness analysis to compute fine grained information about
   /// alloc and dealloc positions.
   Liveness liveness;
 
-  /// The dominator analysis to place allocs in the appropriate blocks.
+  /// The dominator analysis to place deallocs in the appropriate blocks.
   DominanceInfo dominators;
 
   /// The post dominator analysis to place deallocs in the appropriate blocks.
   PostDominanceInfo postDominators;
-
-  /// The internal alias analysis to ensure that allocs and deallocs take all
-  /// their potential aliases into account.
-  BufferPlacementAliasAnalysis aliases;
 };
 
 //===----------------------------------------------------------------------===//
 // BufferPlacementPass
 //===----------------------------------------------------------------------===//
 
 /// The actual buffer placement pass that moves alloc and dealloc nodes into
-/// the right positions. It uses the algorithm described at the top of the file.
+/// the right positions. It uses the algorithm described at the top of the
+/// file.
 struct BufferPlacementPass
     : mlir::PassWrapper<BufferPlacementPass, FunctionPass> {
-  void runOnFunction() override {
-    // Get required analysis information first.
-    auto &analysis = getAnalysis<BufferPlacementAnalysis>();
 
-    // Compute an initial placement of all nodes.
-    llvm::SmallVector<std::pair<OpResult, BufferPlacementPositions>, 16>
-        placements;
-    getFunction().walk([&](MemoryEffectOpInterface op) {
-      // Try to find a single allocation result.
+  void runOnFunction() override {
+    // Remove all dealloc nodes (if any) since appropriate deallocs will be
+    // placed later on automatically.
+    getFunction().walk([&](MemoryEffectOpInterface opInterface) {
+      // Try to find a free effect that is applied to one of our values or a an
+      // involved block argument that will be automatically freed by our pass.
       SmallVector<MemoryEffects::EffectInstance, 2> effects;
-      op.getEffects(effects);
-
-      SmallVector<MemoryEffects::EffectInstance, 2> allocateResultEffects;
-      llvm::copy_if(effects, std::back_inserter(allocateResultEffects),
-                    [=](MemoryEffects::EffectInstance &it) {
-                      Value value = it.getValue();
-                      return isa<MemoryEffects::Allocate>(it.getEffect()) &&
-                             value && value.isa<OpResult>();
-                    });
-      // If there is one result only, we will be able to move the allocation and
-      // (possibly existing) deallocation ops.
-      if (allocateResultEffects.size() == 1) {
-        // Insert allocation result.
-        auto allocResult = allocateResultEffects[0].getValue().cast<OpResult>();
-        placements.emplace_back(
-            allocResult, analysis.computeAllocAndDeallocPositions(allocResult));
-      }
+      opInterface.getEffects(effects);
+      if (llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) {
+            return isa<MemoryEffects::Free>(it.getEffect());
+          }))
+        opInterface.getOperation()->erase();
     });
 
-    // Move alloc (and dealloc - if any) nodes into the right places and insert
-    // dealloc nodes if necessary.
-    for (auto &entry : placements) {
-      // Find already associated dealloc nodes.
-      OpResult alloc = entry.first;
-      auto deallocs = analysis.findAssociatedDeallocs(alloc);
-      if (deallocs.size() > 1) {
-        emitError(alloc.getLoc(),
-                  "not supported number of associated dealloc operations");
-        return;
-      }
-
-      // Move alloc node to the right place.
-      BufferPlacementPositions &positions = entry.second;
-      Operation *allocOperation = alloc.getOwner();
-      allocOperation->moveBefore(positions.getAllocPosition());
-
-      // If there is an existing dealloc, move it to the right place.
-      Operation *nextOp = positions.getDeallocPosition()->getNextNode();
-      assert(nextOp && "Invalid Dealloc operation position");
-      if (deallocs.size()) {
-        (*deallocs.begin())->moveBefore(nextOp);
-      } else {
-        // If there is no dealloc node, insert one in the right place.
-        OpBuilder builder(nextOp);
-        builder.create<DeallocOp>(allocOperation->getLoc(), alloc);
-      }
-    }
-  };
-};
+    // Place all required alloc, copy and dealloc nodes.
+    BufferPlacement placement(getFunction());
+    placement.place();
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // BufferAssignmentPlacer
 //===----------------------------------------------------------------------===//
 
 /// Creates a new assignment placer.

mlir/test/Transforms/buffer-placement.mlir

 // CHECK-NEXT: %[[ALLOC:.*]] = alloc()
 // CHECK-NEXT: cond_br
 //      CHECK: linalg.copy
 // CHECK-NEXT: dealloc %[[ALLOC]]
 // CHECK-NEXT: return
 
 // -----
 
+// Test Case:
+//    bb0
+//   /   \
+//  bb1  bb2 <- Initial position of AllocOp
+//   \   /
+//    bb3
+// BufferPlacement Expected Behaviour: It should not move the existing AllocOp to any other block
+// since the alloc has a dynamic dependency to block argument %0 in bb2. Since the dynamic type is
+// passed to bb3 via the block argument %2, it is currently required to allocate a temporary buffer
+// for %2 that gets copies of %arg0 and %1 with their appropriate shape dimensions. The copy buffer
+// deallocation will be applied to %2 in block bb3.
+
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: func @condBranchDynamicType
+func @condBranchDynamicType(%arg0: i1, %arg1: memref<?xf32>, %arg2: memref<?xf32>, %arg3: index) {
+  cond_br %arg0, ^bb1, ^bb2(%arg3: index)
+^bb1:
+  br ^bb3(%arg1 : memref<?xf32>)
+^bb2(%0: index):
+  %1 = alloc(%0) : memref<?xf32>
+  linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0], iterator_types = ["parallel"]} %arg1, %1 {
+  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
+    %tmp1 = exp %gen1_arg0 : f32
+    linalg.yield %tmp1 : f32
+  }: memref<?xf32>, memref<?xf32>
+  br ^bb3(%1 : memref<?xf32>)
+^bb3(%2: memref<?xf32>):
+  "linalg.copy"(%2, %arg2) : (memref<?xf32>, memref<?xf32>) -> ()
+  return
+}
+
+// CHECK-NEXT: cond_br
+//      CHECK: %[[DIM0:.*]] = dim
+// CHECK-NEXT: %[[ALLOC0:.*]] = alloc(%[[DIM0]])
+// CHECK-NEXT: linalg.copy(%{{.*}}, %[[ALLOC0]])
+//      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
+// CHECK-NEXT: %[[ALLOC1:.*]] = alloc(%[[IDX]])
+// CHECK-NEXT: linalg.generic
+//      CHECK: %[[DIM1:.*]] = dim %[[ALLOC1]]
+// CHECK-NEXT: %[[ALLOC2:.*]] = alloc(%[[DIM1]])
+// CHECK-NEXT: linalg.copy(%[[ALLOC1]], %[[ALLOC2]])
+// CHECK-NEXT: dealloc %[[ALLOC1]]
+// CHECK-NEXT: br ^bb3
+// CHECK-NEXT: ^bb3(%[[ALLOC02:.*]]:{{.*}})
+//      CHECK: linalg.copy(%[[ALLOC02]],
+// CHECK-NEXT: dealloc %[[ALLOC02]]
+// CHECK-NEXT: return
+
+// -----
+
 // Test Case: Existing AllocOp with no users.
 // BufferPlacement Expected Behaviour: It should insert a DeallocOp right before ReturnOp.
 
 // CHECK-LABEL: func @emptyUsesValue
 func @emptyUsesValue(%arg0: memref<4xf32>) {
   %0 = alloc() : memref<4xf32>
   return
 }

 // CHECK-NEXT: dealloc %[[ALLOC]]
 // CHECK-NEXT: return
 
 // -----
 
 // Test Case:
 //    bb0 <- Initial position of AllocOp
 //   /   \
 //  |    bb1
 //   \   /
 //    bb2
 // BufferPlacement Expected Behaviour: It shouldn't move the alloc position. It only inserts
 // a DeallocOp at the exit block after CopyOp since %1 is an alias for %0 and %arg1.
 
 #map0 = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: func @invCriticalEdge
^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
   ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
     %tmp2 = exp %gen2_arg0 : f32
     linalg.yield %tmp2 : f32
   }: memref<2xf32>, memref<2xf32>
   "linalg.copy"(%7, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
   return
 }
 
 // CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
 // CHECK-NEXT: linalg.generic
+//      CHECK: %[[COPY_ALLOC0:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC0]])
+//      CHECK: %[[COPY_ALLOC1:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC1]])
+//      CHECK: %[[COPY_ALLOC2:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC2]])
+//      CHECK: %[[COPY_ALLOC3:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC3]])
 //      CHECK: %[[SECOND_ALLOC:.*]] = alloc()
-// CHECK-NEXT: linalg.generic
 //      CHECK: dealloc %[[FIRST_ALLOC]]
-// CHECK-NEXT: linalg.copy
+//      CHECK: linalg.generic
+//      CHECK: linalg.copy
 // CHECK-NEXT: dealloc %[[SECOND_ALLOC]]
 // CHECK-NEXT: return
 
 // -----
 
 // Test Case: No users for buffer in if-else CFG
 //    bb0 <- Initial position of AllocOp
 //   /   \
^bb1(%1: memref<2xf32>, %2: memref<2xf32>):
   br ^bb3(%1, %2 : memref<2xf32>, memref<2xf32>)
 ^bb2(%3: memref<2xf32>, %4: memref<2xf32>):
   br ^bb3(%3, %4 : memref<2xf32>, memref<2xf32>)
 ^bb3(%5: memref<2xf32>, %6: memref<2xf32>):
   "linalg.copy"(%arg1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
   return
 }
 
-//      CHECK: dealloc
+// CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
+//      CHECK: dealloc %[[FIRST_ALLOC]]
 // CHECK-NEXT: return
 
 // -----
 
 // Test Case:
 //      bb0 <- Initial position of the first AllocOp
 //     /    \
 //   bb1    bb2
^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
     linalg.yield %tmp2 : f32
   }: memref<2xf32>, memref<2xf32>
   "linalg.copy"(%9, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
   return
 }
 
 // CHECK-NEXT: %[[FIRST_ALLOC:.*]] = alloc()
 // CHECK-NEXT: linalg.generic
+//      CHECK: %[[COPY_ALLOC0:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC0]])
+//      CHECK: %[[COPY_ALLOC1:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC1]])
+//      CHECK: %[[COPY_ALLOC2:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC2]])
+//      CHECK: %[[COPY_ALLOC3:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC3]])
+//      CHECK: %[[COPY_ALLOC4:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC4]])
+//      CHECK: %[[COPY_ALLOC5:.*]] = alloc()
+// CHECK-NEXT: linalg.copy({{.*}}, %[[COPY_ALLOC5]])
 //      CHECK: %[[SECOND_ALLOC:.*]] = alloc()
-// CHECK-NEXT: linalg.generic
 //      CHECK: dealloc %[[FIRST_ALLOC]]
-// CHECK-NEXT: linalg.copy
+//      CHECK: linalg.generic
+//      CHECK: linalg.copy
 // CHECK-NEXT: dealloc %[[SECOND_ALLOC]]
 // CHECK-NEXT: return
 
 // -----
 
 // Test Case: Dead operations in a single block.
 // BufferPlacement Expected Behaviour: It shouldn't move the AllocOps. It only inserts the two missing DeallocOps
 // after the last GenericOp.
func @moving_invalid_dealloc_op(%arg0 : memref<2xf32>, %arg1: memref<2xf32>){
   }: memref<2xf32>, memref<2xf32>
   dealloc %0 : memref<2xf32>
   "linalg.copy"(%0, %arg1) : (memref<2xf32>, memref<2xf32>) -> ()
   return
 }
 
 //      CHECK: linalg.copy
 // CHECK-NEXT: dealloc
 No newline at end of file