[mlir] Extended BufferPlacement to support nested region control flow.

[dfki-mako]

Summary

The current BufferPlacement implementation does not support nested region control flow. This CL adds support for nested regions via the RegionBranchOpInterface and the detection of branch-like (ReturnLike) terminators inside nested regions.

[herhut]

Can you also add a test with deeper nesting?

mlir/lib/Transforms/BufferPlacement.cpp

81

Another way to describe exitParentRegion would be to only consider immediately nested regions. You can use the approach you have now, or simply use a loop over regions/blocks and get their terminators in the true case. Less traversal work.

87

would terminator->getParentRegion()->getParentOp() == operation do the same?

136–195

detailed information is not very helpful in a comment. What does it query?

150

Maybe llvm::for_each? Drop the {,}.

207

So these are the entry regions for this operation. Maybe write that in the comment.

211

This ties the operands of the op with a region to the block arguments of the target region,

212

Block &successorBlock = regionSuccessor.getSuccessor()->front(); ?

213

for_each or drop the brackets.

217

Now that you have queried the flow-in, you can also query the flow within the op. For this, you can call regionInterface.getSuccessorRegions once for each region of the op. The alias registration is the same as with the initial flow into the op.

Except, when regionSuccessor.getSuccessor() is nulltptr. That signals that the terminator of the region will exit the region. So you have to tie the regionSuccessor.getSuccessorInputs() to parentOp->getResults() in that case.

221

/*exitParentRegion=*/false

221

I don't think this is needed with the above implemented.

302

front()?

446

This copies a region result, so maybe reflect that in the name.

454

This also needs to use the RegionBranchOpInterface to find all the terminators in the operation and their successor inputs (they might be in a different order than the op results). Querying the terminators directly makes assumptions about ordering that do not exist.

mlir/test/Transforms/buffer-placement.mlir

779

There should be a dealloc here, right?

[dfki-mako]

Added extended support for the RegionOpInterface to query successor bindings for successor regions

[rriddle]

mlir/lib/Transforms/BufferPlacement.cpp

81

Please move static functions to the global namespace. anonymous namespace should only really be used by things like classes.

139

nit: Use a lambda instead.

154

nit: I don't think this really saves anything over a normal for loop, if anything it is much less efficient.

163

nit: Use the full name for the type.

188

nit: Use parameter names when passing constants, i.e., /*someName=*/

424

Is this guaranteed to be a RegionBranchOpInterface?

437

You are completely disregarding the fact that getMutableSuccessorOperands can return None, which would cause this to crash. Where are you checking that this is valid?

442

llvm::find?

[herhut]

mlir/lib/Transforms/BufferPlacement.cpp

163

Why is this called operands? These are the results of the overall operation, right?

165

This region does not have a valid successor block if it terminates the parent operation. In that case, we wire its successor operands with the parent operations results.

This could be more obvious in the code if you would pass the parent operation. Please at least rename operands to results.

170

The loop and the zip could be part of the helper function registerAliasFunc.

175

Isn't this the wrong way round? The results alias the successor inputs?

201

This only follows one level of branching. So if the entry region goes to region 2, and that one goes to region 3, the second link will not be seen.

It should be good enough to loop over all regions of an operation and then do this linking to successor regions.

437

The underlying assumption is that if the successor has block arguments, then the branch in the predecessor needs to have operands for those. Can we rely on this?

[rriddle]

mlir/lib/Transforms/BufferPlacement.cpp

437

Branching operations are not required to have a Value already materialized for a block argument. There are certain classes of operations that internally generate the value that is passed to the branch. For example, operations like LLVM Callbr and certain SIL switches.

[dfki-mako]

Refactored implementation and simplified the iteration over all successor regions.

[dfki-mako]

mlir/lib/Transforms/BufferPlacement.cpp

424

The value should either be a BlockArgument or a value resulting from a RegionBranchOpInterface operation.

437

We have changed the code to emit an error message. We prefer the error message over silently ignoring this case as the retrieved alias information can be invalid if an operation passes values implicitly to a block argument.

[herhut]

mlir/lib/Transforms/BufferPlacement.cpp

150

nit: interface

177

entryRegion.getSuccessorInputs returns the inputs of the target region, typically the block arguments of the first block. So this maps them with itself.
Instead, this should map the result from regionInterface().getSuccessorEntryOperands to the entryRegion.getSuccessorInputs.

193

The terminator should also implement the BranchOpInterface, right? So, one should query the inputs to the successor using BranchOpInterface.getSuccessorOperands, correct?

Here, successorRegion.getSuccessorInputs returns the input values to the region, which normally would be the block arguments or, in case this leaves the operation, the results.

[dfki-mako]

Added support for region-region control flow within operations that implement the RegionBranchOpInterface.
Added new test operations to verify more advanced region-region control-flow scenarios.

[herhut]

Cool stuff. Great to see this work with complex region interfaces, as well!

So, this now assumes that there is a 1:1 mapping between return-like op operands and successor inputs and we likely want an interface that makes this configurable. Maybe the returnlike trait could provide a function that enables this. (We discussed this, just writing it down.)

Also, we need a clean-up pass to remove unneeded alloc+copy+dealloc triples. This is underway.

With the nits addressed and some of the comments cleaned up, this is good to go.

mlir/lib/Transforms/BufferPlacement.cpp

178

Mega-nit: entryRegion->getSuccessor reads weird. It read like getting the successor of the entry region but actually it gets the entry region itself :)

Maybe rename this entrySuccessor?

182

Maybe Wire flow between regions and from region exits.?

188

The length of operandAttributes might be wrong here, as it was built for the entry successors.

190

I think this comment is now off. It always wires the terminator operands with the successor inputs. The latter can be block arguments of a region's entry block or the result values, if the terminator exists the op.

301

Should this apply to the case above, as well?

401

Where is this conversion? Or is this comment off?

406

Nit: Doesn't this create a new alloc and a copy? If so, please fix the comment.

435

So is it OK to pass an empty ArrayRef here? If so, why can we not do this in the other cases?

mlir/test/lib/Dialect/Test/TestOps.td

1364

Please use

let description = [{
...
}];

instead of comment.

[herhut]

With my comments addressed, this is good to go.

[dfki-mako]

Addressed reviewer comments.

mlir/lib/Transforms/BufferPlacement.cpp

177

Yeah makes sense to me; we should wire getSuccessorEntryOperands with getSuccessorInputs.

188

According to our interpretation of the comment

operands` is a set of optional attributes that correspond to a constant value for each operand

this should be an array containing values for each operand of the defining operation; although it might make sense to include an attribute for each block argument. However, we should keep this for now, as it is fully compatible with the current implementation of all instantiations of the RegionBrachOpInterface.

193

I guess querying a BranchOpInterface might not make sense in the case of a ReturnLike op, since it does not branch to any block in general (although this can still be expressed via the BranchOpInterface). It feels like that a ReturnLike terminator should provide a method to access all operands similar to the BranchOpInterface.

mlir/lib/Transforms/BufferPlacement.cpp

 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SetOperations.h"
 
 using namespace mlir;
 
-namespace {
+/// Walks over all immediate return-like terminators in the given region.
+template <typename FuncT>
+static void walkReturnOperations(Region *region, const FuncT &func) {
+  for (Block &block : *region)
+    for (Operation &operation : block) {
+      // Skip non-return-like terminators.
+      if (operation.hasTrait<OpTrait::ReturnLike>())
+        func(&operation);
+    }
+}
 
+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()); }
+  BufferPlacementAliasAnalysis(Operation *op) { build(op); }
 
   /// Find 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(); }
 
   /// Find 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;
   }
 
   /// Removes the given values from all alias sets.
-  void remove(const SmallPtrSetImpl<BlockArgument> &aliasValues) {
+  void remove(const SmallPtrSetImpl<Value> &aliasValues) {
     for (auto &entry : aliases)
       llvm::set_subtract(entry.second, aliasValues);
   }
 
 private:
   /// Recursively determines alias information for the given value. It stores
   /// all newly found potential aliases in the given result set.
   void resolveRecursive(Value value, ValueSetT &result) const {
     if (!result.insert(value).second)
       return;
     auto it = aliases.find(value);
     if (it == aliases.end())
       return;
     for (Value alias : it->second)
       resolveRecursive(alias, result);
   }
 
   /// This function constructs a mapping from values to its immediate aliases.
   /// It iterates over all blocks, gets their predecessors, determines the
   /// values that will be passed to the corresponding block arguments and
-  /// inserts them into the underlying map.
-  void build(MutableArrayRef<Region> regions) {
-    for (Region &region : regions) {
-      for (Block &block : region) {
-        // Iterate over all predecessor and get the mapped values to their
-        // corresponding block arguments values.
-        for (auto it = block.pred_begin(), e = block.pred_end(); it != e;
-             ++it) {
-          unsigned successorIndex = it.getSuccessorIndex();
-          // Get the terminator and the values that will be passed to our block.
-          auto branchInterface =
-              dyn_cast<BranchOpInterface>((*it)->getTerminator());
-          if (!branchInterface)
-            continue;
-          // Query the branch op interace to get the successor operands.
-          auto successorOperands =
-              branchInterface.getSuccessorOperands(successorIndex);
-          if (successorOperands.hasValue()) {
-            // Build the actual mapping of values to their immediate aliases.
-            for (auto argPair : llvm::zip(block.getArguments(),
-                                          successorOperands.getValue())) {
-              aliases[std::get<1>(argPair)].insert(std::get<0>(argPair));
-            }
-          }
-        }
+  /// inserts them into the underlying map. Furthermore, it wires successor
+  /// regions and branch-like return operations from nested regions.
+  void build(Operation *op) {
+    // Registers all aliases of the given values.
+    auto registerAliases = [&](auto values, auto aliases) {
+      for (auto entry : llvm::zip(values, aliases))
+        this->aliases[std::get<0>(entry)].insert(std::get<1>(entry));
+    };
+
+    // Query all branch interfaces to link block argument aliases.
+    op->walk([&](BranchOpInterface branchInterface) {
+      Block *parentBlock = branchInterface.getOperation()->getBlock();
+      for (auto it = parentBlock->succ_begin(), e = parentBlock->succ_end();
+           it != e; ++it) {
+        // Query the branch op interface to get the successor operands.
+            branchInterface.getSuccessorOperands(it.getIndex());
+        if (!successorOperands.hasValue())
+          continue;
+        registerAliases(successorOperands.getValue(), (*it)->getArguments());
+    });
+
+    // Query the RegionBranchOpInterface to find potential successor regions.
+    op->walk([&](RegionBranchOpInterface regionInterface) {
+      // Create an empty attribute for each operand to comply with the
+      // `getSuccessorRegions` interface definition that requires a single
+      // attribute per operand.
+      SmallVector<Attribute, 2> operandAttributes(
+          regionInterface.getOperation()->getNumOperands());
+
+      // Extract all entry regions and wire all initial entry successor inputs.
+      SmallVector<RegionSuccessor, 2> entrySuccessors;
+      regionInterface.getSuccessorRegions(/*index=*/llvm::None,
+                                          operandAttributes, entrySuccessors);
+      for (RegionSuccessor &entrySuccessor : entrySuccessors) {
+        // Wire the entry region's successor arguments with the initial
+        // successor inputs.
+        assert(entrySuccessor.getSuccessor() &&
+               "Invalid entry region without an attached successor region");
+        registerAliases(regionInterface.getSuccessorEntryOperands(
+                            entrySuccessor.getSuccessor()->getRegionNumber()),
+                        entrySuccessor.getSuccessorInputs());
+      }
+
+      // Wire flow between regions and from region exits.
+      for (Region &region : regionInterface.getOperation()->getRegions()) {
+        // Iterate over all successor region entries that are reachable from the
+        // current region.
+        SmallVector<RegionSuccessor, 2> successorRegions;
+        regionInterface.getSuccessorRegions(
+            region.getRegionNumber(), operandAttributes, successorRegions);
+        for (RegionSuccessor &successorRegion : successorRegions) {
+          // Iterate over all immediate terminator operations and wire the
+          // successor inputs with the operands of each terminator.
+          walkReturnOperations(&region, [&](Operation *terminator) {
+            registerAliases(terminator->getOperands(),
+                            successorRegion.getSuccessorInputs());
+          });
         }
       }
+    });
   }
 
   /// Maps values to all immediate aliases this value can have.
   ValueMapT aliases;
 };
 
 //===----------------------------------------------------------------------===//
 // BufferPlacement
 //===----------------------------------------------------------------------===//
 
 // The main buffer placement analysis used to place allocs, copies and deallocs.
 class BufferPlacement {
 public:
   using ValueSetT = BufferPlacementAliasAnalysis::ValueSetT;
 
   /// An intermediate representation of a single allocation node.
   struct AllocEntry {
     /// A reference to the associated allocation node.
     Value allocValue;
 
     /// The associated placement block in which the allocation should be
     /// performed.
     Block *placementBlock;
 
     /// The associated dealloc operation (if any).
     Operation *deallocOperation;
   };
 
   using AllocEntryList = SmallVector<AllocEntry, 8>;
 
 public:
void initBlockMapping() {
     });
   }
 
   /// Computes a valid allocation position in a dominator (if possible) for the
   /// given allocation result.
   Block *getInitialAllocBlock(OpResult result) {
     // Get all allocation operands as these operands are important for the
     // allocation operation.
-    auto operands = result.getOwner()->getOperands();
+    Operation *owner = result.getOwner();
+    auto operands = owner->getOperands();
+    Block *dominator;
     if (operands.size() < 1)
-      return findCommonDominator(result, aliases.resolve(result), dominators);
-
+      dominator =
+          findCommonDominator(result, aliases.resolve(result), dominators);
+    else {
       // If this node has dependencies, check all dependent nodes with respect
       // to a common post dominator in which all values are available.
       ValueSetT dependencies(++operands.begin(), operands.end());
-    return findCommonDominator(*operands.begin(), dependencies, postDominators);
+      dominator =
+          findCommonDominator(*operands.begin(), dependencies, postDominators);
+    }
+
+    // Do not move allocs out of their parent regions to keep them local.
+    if (dominator->getParent() != owner->getParentRegion())
+      return &owner->getParentRegion()->front();
+    return dominator;
   }
 
   /// 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 : allocs) {
       Value alloc = entry.allocValue;
       // Get the actual block to place the alloc and get liveness information
for (auto &entry : allocs) {
       // Move the alloc in front of the start operation.
       Operation *allocOperation = alloc.getDefiningOp();
       allocOperation->moveBefore(startOperation);
     }
   }
 
   /// 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;
-    llvm::SmallDenseSet<std::tuple<BlockArgument, Block *>> visitedBlockArgs;
-    SmallVector<std::tuple<BlockArgument, Block *>, 8> toProcess;
+    // 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
     // 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) > 0)
+        if (valuesToFree.count(value) > 0)
           continue;
         // Check whether we have to free this particular block argument.
-        if (!dominators.dominates(definingBlock, blockArg.getOwner())) {
-          toProcess.emplace_back(blockArg, blockArg.getParentBlock());
-          blockArgsToFree.insert(blockArg);
-        } else if (visitedBlockArgs
-                       .insert(std::make_tuple(blockArg, definingBlock))
+        if (!dominators.dominates(definingBlock, value.getParentBlock())) {
+          toProcess.emplace_back(value, value.getParentBlock());
+          valuesToFree.insert(value);
+        } else if (visitedValues.insert(std::make_tuple(value, definingBlock))
                        .second)
-          toProcess.emplace_back(blockArg, definingBlock);
+          toProcess.emplace_back(value, definingBlock);
       }
     };
 
     // Detect possibly unsafe aliases starting from all allocations.
     for (auto &entry : allocs)
       findUnsafeValues(entry.allocValue, entry.placementBlock);
 
     // Try to find block arguments that require an explicit free operation
     // until we reach a fix point.
     while (!toProcess.empty()) {
       auto current = toProcess.pop_back_val();
       findUnsafeValues(std::get<0>(current), std::get<1>(current));
     }
 
     // Update buffer aliases to ensure that we free all buffers and block
     // arguments at the correct locations.
-    aliases.remove(blockArgsToFree);
+    aliases.remove(valuesToFree);
 
     // Add new allocs and additional copy operations.
-    for (BlockArgument blockArg : blockArgsToFree) {
-      Block *block = blockArg.getOwner();
+    for (Value value : valuesToFree) {
+      if (auto blockArg = value.dyn_cast<BlockArgument>())
+        introduceBlockArgCopy(blockArg);
+      else
+        introduceValueCopyForRegionResult(value);
+
+      // 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.push_back({value, nullptr, nullptr});
+    }
+  }
 
-      // 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 branchInterface = cast<BranchOpInterface>(terminator);
-        // Convert the mutable operand range to an immutable range and query the
-        // associated source value.
-        Value sourceValue =
-            branchInterface.getSuccessorOperands(it.getSuccessorIndex())
-                .getValue()[blockArg.getArgNumber()];
-        // Create a new alloc at the current location of the terminator.
-        auto memRefType = sourceValue.getType().cast<MemRefType>();
-        OpBuilder builder(terminator);
+  /// Introduces temporary allocs in all predecessors and copies the source
+  /// values into the newly allocated buffers.
+  void introduceBlockArgCopy(BlockArgument blockArg) {
+    Block *block = blockArg.getOwner();
+    for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
+      // Query the associated source value.
+      // Create a new alloc and copy at the current location of the terminator.
+      Value alloc = introduceBufferCopy(sourceValue, terminator);
+      // Wire new alloc and successor operand.
+      auto mutableOperands =
+          branchInterface.getMutableSuccessorOperands(it.getSuccessorIndex());
+      if (!mutableOperands.hasValue())
+        terminator->emitError() << "terminators with immutable successor "
+                                   "operands are not supported";
+      else
+        mutableOperands.getValue()
+            .slice(blockArg.getArgNumber(), 1)
+            .assign(alloc);
+    }
+
+    // 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();
+    RegionBranchOpInterface regionInterface;
+    if (!argRegion || &argRegion->front() != block ||
+        !(regionInterface =
+              dyn_cast<RegionBranchOpInterface>(argRegion->getParentOp())))
+      return;
+
+    introduceCopiesForRegionSuccessors(
+        regionInterface, argRegion->getParentOp()->getRegions(),
+        [&](RegionSuccessor &successorRegion) {
+          // Find a predecessor of our argRegion.
+          return successorRegion.getSuccessor() == argRegion;
+        },
+        [&](RegionSuccessor &successorRegion) {
+          // The operand index will be the argument number.
+          return blockArg.getArgNumber();
+        });
+  }
+
+  /// Introduces temporary allocs 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);
+    introduceCopiesForRegionSuccessors(
+        regionInterface, operation->getRegions(),
+        [&](RegionSuccessor &successorRegion) {
+          // Determine whether this region has a successor entry that leaves
+          // this region by returning to its parent operation.
+          return !successorRegion.getSuccessor();
+        },
+        [&](RegionSuccessor &successorRegion) {
+          // Find the associated success input index.
+          return llvm::find(successorRegion.getSuccessorInputs(), value)
+              .getIndex();
+        });
+  }
+
+  /// Introduces buffer copies for all terminators in the given regions. The
+  /// regionPredicate is applied to every successor region in order to restrict
+  /// the copies to specific regions. Thereby, the operandProvider is invoked
+  /// for each matching region successor and determines the operand index that
+  /// requires a buffer copy.
+  template <typename TPredicate, typename TOperandProvider>
+  void
+  introduceCopiesForRegionSuccessors(RegionBranchOpInterface regionInterface,
+                                     MutableArrayRef<Region> regions,
+                                     const TPredicate &regionPredicate,
+                                     const TOperandProvider &operandProvider) {
+    // Create an empty attribute for each operand to comply with the
+    // `getSuccessorRegions` interface definition that requires a single
+    // attribute per operand.
+    SmallVector<Attribute, 2> operandAttributes(
+        regionInterface.getOperation()->getNumOperands());
+    for (Region &region : regions) {
+      // Query the regionInterface to get all successor regions of the current
+      // one.
+      SmallVector<RegionSuccessor, 2> successorRegions;
+      regionInterface.getSuccessorRegions(region.getRegionNumber(),
+                                          operandAttributes, 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.
+      auto operandIndex = operandProvider(*regionSuccessor);
+
+      // 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) {
+        // Extract the source value from the current terminator.
+        Value sourceValue = terminator->getOperand(operandIndex);
+        // Create a new alloc at the current location of the terminator.
+        Value alloc = introduceBufferCopy(sourceValue, terminator);
+        // Wire alloc and terminator operand.
+        terminator->setOperand(operandIndex, alloc);
+      });
+    }
+  }
+
+  /// Creates a new memory allocation for the given source value and copies
+  /// its content into the newly allocated buffer. The terminator operation is
+  /// used to insert the alloc and copy operations at the right places.
+  Value introduceBufferCopy(Value sourceValue, Operation *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()));
-        }
 
-        // 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.
-        branchInterface.getMutableSuccessorOperands(it.getSuccessorIndex())
-            .getValue()
-            .slice(blockArg.getArgNumber(), 1)
-            .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);
-      }
+
+    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.
-      allocs.push_back({blockArg, nullptr, nullptr});
-    }
+    return alloc;
   }
 
   /// Finds associated deallocs that can be linked to our allocation nodes (if
   /// any).
   void findDeallocs() {
     for (auto &entry : allocs) {
       auto userIt =
           llvm::find_if(entry.allocValue.getUsers(), [&](Operation *user) {
for (auto &entry : allocs) {
       }
       // endOperation is the last operation behind which we can safely store
       // the dealloc taking all potential aliases into account.
 
       // If there is an existing dealloc, move it to the right place.
       if (entry.deallocOperation) {
         entry.deallocOperation->moveAfter(endOperation);
       } else {
-        // If the Dealloc position is at the terminator operation of the block,
-        // then the value should escape from a deallocation.
+        // 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<DeallocOp>(alloc.getLoc(), alloc);
       }
     }

mlir/test/Transforms/buffer-placement.mlir

 }
 // CHECK: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>, %[[RESULT:.*]]: memref<5xf32>)
 // CHECK: %[[X:.*]] = alloc()
 // CHECK: %[[Y:.*]] = alloc()
 // CHECK: linalg.copy
 // CHECK: dealloc %[[Y]]
 // CHECK: return %[[ARG1]], %[[X]]
 
+// -----
+
+// Test Case: nested region control flow
+// The alloc position of %1 does not need to be changed and flows through
+// both if branches until it is finally returned. Hence, it does not
+// require a specific dealloc operation. However, %3 requires a dealloc.
+
+// CHECK-LABEL: func @nested_region_control_flow
+func @nested_region_control_flow(
+  %arg0 : index,
+  %arg1 : index) -> memref<?x?xf32> {
+  %0 = cmpi "eq", %arg0, %arg1 : index
+  %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
+  %2 = scf.if %0 -> (memref<?x?xf32>) {
+    scf.yield %1 : memref<?x?xf32>
+  } else {
+    %3 = alloc(%arg0, %arg1) : memref<?x?xf32>
+    scf.yield %1 : memref<?x?xf32>
+  }
+  return %2 : memref<?x?xf32>
+}
+
+//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
+// CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
+//      CHECK: scf.yield %[[ALLOC0]]
+//      CHECK: %[[ALLOC2:.*]] = alloc(%arg0, %arg1)
+// CHECK-NEXT: dealloc %[[ALLOC2]]
+// CHECK-NEXT: scf.yield %[[ALLOC0]]
+//      CHECK: return %[[ALLOC1]]
+
+// -----
+
+// Test Case: nested region control flow with a nested buffer allocation in a
+// divergent branch.
+// The alloc positions of %1, %3 does not need to be changed since
+// BufferPlacement does not move allocs out of nested regions at the moment.
+// However, since %3 is allocated and "returned" in a divergent branch, we have
+// to allocate a temporary buffer (like in condBranchDynamicTypeNested).
+
+// CHECK-LABEL: func @nested_region_control_flow_div
+func @nested_region_control_flow_div(
+  %arg0 : index,
+  %arg1 : index) -> memref<?x?xf32> {
+  %0 = cmpi "eq", %arg0, %arg1 : index
+  %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
+  %2 = scf.if %0 -> (memref<?x?xf32>) {
+    scf.yield %1 : memref<?x?xf32>
+  } else {
+    %3 = alloc(%arg0, %arg1) : memref<?x?xf32>
+    scf.yield %3 : memref<?x?xf32>
+  }
+  return %2 : memref<?x?xf32>
+}
+
+//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
+// CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
+//      CHECK: %[[ALLOC2:.*]] = alloc
+// CHECK-NEXT: linalg.copy(%[[ALLOC0]], %[[ALLOC2]])
+//      CHECK: scf.yield %[[ALLOC2]]
+//      CHECK: %[[ALLOC3:.*]] = alloc(%arg0, %arg1)
+//      CHECK: %[[ALLOC4:.*]] = alloc
+// CHECK-NEXT: linalg.copy(%[[ALLOC3]], %[[ALLOC4]])
+//      CHECK: dealloc %[[ALLOC3]]
+//      CHECK: scf.yield %[[ALLOC4]]
+//      CHECK: dealloc %[[ALLOC0]]
+// CHECK-NEXT: return %[[ALLOC1]]
+
+// -----
+
+// Test Case: deeply nested region control flow with a nested buffer allocation
+// in a divergent branch.
+// The alloc positions of %1, %4 and %5 does not need to be changed since
+// BufferPlacement does not move allocs out of nested regions at the moment.
+// However, since %4 is allocated and "returned" in a divergent branch, we have
+// to allocate several temporary buffers (like in condBranchDynamicTypeNested).
+
+// CHECK-LABEL: func @nested_region_control_flow_div_nested
+func @nested_region_control_flow_div_nested(
+  %arg0 : index,
+  %arg1 : index) -> memref<?x?xf32> {
+  %0 = cmpi "eq", %arg0, %arg1 : index
+  %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
+  %2 = scf.if %0 -> (memref<?x?xf32>) {
+    %3 = scf.if %0 -> (memref<?x?xf32>) {
+      scf.yield %1 : memref<?x?xf32>
+    } else {
+      %4 = alloc(%arg0, %arg1) : memref<?x?xf32>
+      scf.yield %4 : memref<?x?xf32>
+    }
+    scf.yield %3 : memref<?x?xf32>
+  } else {
+    %5 = alloc(%arg1, %arg1) : memref<?x?xf32>
+    scf.yield %5 : memref<?x?xf32>
+  }
+  return %2 : memref<?x?xf32>
+}
+//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
+// CHECK-NEXT: %[[ALLOC1:.*]] = scf.if
+// CHECK-NEXT: %[[ALLOC2:.*]] = scf.if
+//      CHECK: %[[ALLOC3:.*]] = alloc
+// CHECK-NEXT: linalg.copy(%[[ALLOC0]], %[[ALLOC3]])
+//      CHECK: scf.yield %[[ALLOC3]]
+//      CHECK: %[[ALLOC4:.*]] = alloc(%arg0, %arg1)
+//      CHECK: %[[ALLOC5:.*]] = alloc
+// CHECK-NEXT: linalg.copy(%[[ALLOC4]], %[[ALLOC5]])
+//      CHECK: dealloc %[[ALLOC4]]
+//      CHECK: scf.yield %[[ALLOC5]]
+//      CHECK: %[[ALLOC6:.*]] = alloc
+// CHECK-NEXT: linalg.copy(%[[ALLOC2]], %[[ALLOC6]])
+//      CHECK: dealloc %[[ALLOC2]]
+//      CHECK: scf.yield %[[ALLOC6]]
+//      CHECK: %[[ALLOC7:.*]] = alloc(%arg1, %arg1)
+//      CHECK: %[[ALLOC8:.*]] = alloc
+// CHECK-NEXT: linalg.copy(%[[ALLOC7]], %[[ALLOC8]])
+//      CHECK: dealloc %[[ALLOC7]]
+//      CHECK: scf.yield %[[ALLOC8]]
+//      CHECK: dealloc %[[ALLOC0]]
+// CHECK-NEXT: return %[[ALLOC1]]
+
+// -----
+
+// Test Case: nested region control flow within a region interface.
+// The alloc positions of %0 does not need to be changed and no copies are
+// required in this case since the allocation finally escapes the method.
+
+// CHECK-LABEL: func @inner_region_control_flow
+func @inner_region_control_flow(%arg0 : index) -> memref<?x?xf32> {
+  %0 = alloc(%arg0, %arg0) : memref<?x?xf32>
+  %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
+    ^bb0(%arg1 : memref<?x?xf32>):
+      test.region_if_yield %arg1 : memref<?x?xf32>
+  } else {
+    ^bb0(%arg1 : memref<?x?xf32>):
+      test.region_if_yield %arg1 : memref<?x?xf32>
+  } join {
+    ^bb0(%arg1 : memref<?x?xf32>):
+      test.region_if_yield %arg1 : memref<?x?xf32>
+  }
+  return %1 : memref<?x?xf32>
+}
+
+//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
+// CHECK-NEXT: %[[ALLOC1:.*]] = test.region_if
+// CHECK-NEXT: ^bb0(%[[ALLOC2:.*]]:{{.*}}):
+// CHECK-NEXT: test.region_if_yield %[[ALLOC2]]
+//      CHECK: ^bb0(%[[ALLOC3:.*]]:{{.*}}):
+// CHECK-NEXT: test.region_if_yield %[[ALLOC3]]
+//      CHECK: ^bb0(%[[ALLOC4:.*]]:{{.*}}):
+// CHECK-NEXT: test.region_if_yield %[[ALLOC4]]
+//      CHECK: return %[[ALLOC1]]
+
+// -----
+
+// Test Case: nested region control flow within a region interface including an
+// allocation in a divergent branch.
+// The alloc positions of %1 and %2 does not need to be changed since
+// BufferPlacement does not move allocs out of nested regions at the moment.
+// However, since %2 is allocated and yielded in a divergent branch, we have
+// to allocate several temporary buffers (like in condBranchDynamicTypeNested).
+
+// CHECK-LABEL: func @inner_region_control_flow_div
+func @inner_region_control_flow_div(
+  %arg0 : index,
+  %arg1 : index) -> memref<?x?xf32> {
+  %0 = alloc(%arg0, %arg0) : memref<?x?xf32>
+  %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
+    ^bb0(%arg2 : memref<?x?xf32>):
+      test.region_if_yield %arg2 : memref<?x?xf32>
+  } else {
+    ^bb0(%arg2 : memref<?x?xf32>):
+      %2 = alloc(%arg0, %arg1) : memref<?x?xf32>
+      test.region_if_yield %2 : memref<?x?xf32>
+  } join {
+    ^bb0(%arg2 : memref<?x?xf32>):
+      test.region_if_yield %arg2 : memref<?x?xf32>
+  }
+  return %1 : memref<?x?xf32>
+}
+
+//      CHECK: %[[ALLOC0:.*]] = alloc(%arg0, %arg0)
+// CHECK-NEXT: %[[ALLOC1:.*]] = test.region_if
+// CHECK-NEXT: ^bb0(%[[ALLOC2:.*]]:{{.*}}):
+//      CHECK: %[[ALLOC3:.*]] = alloc
+// CHECK-NEXT: linalg.copy(%[[ALLOC2]], %[[ALLOC3]])
+// CHECK-NEXT: test.region_if_yield %[[ALLOC3]]
+//      CHECK: ^bb0(%[[ALLOC4:.*]]:{{.*}}):
+//      CHECK: %[[ALLOC5:.*]] = alloc
+//      CHECK: %[[ALLOC6:.*]] = alloc
+// CHECK-NEXT: linalg.copy(%[[ALLOC5]], %[[ALLOC6]])
+// CHECK-NEXT: dealloc %[[ALLOC5]]
+// CHECK-NEXT: test.region_if_yield %[[ALLOC6]]
+//      CHECK: ^bb0(%[[ALLOC7:.*]]:{{.*}}):
+//      CHECK: %[[ALLOC8:.*]] = alloc
+// CHECK-NEXT: linalg.copy(%[[ALLOC7]], %[[ALLOC8]])
+// CHECK-NEXT: dealloc %[[ALLOC7]]
+// CHECK-NEXT: test.region_if_yield %[[ALLOC8]]
+//      CHECK: dealloc %[[ALLOC0]]
+// CHECK-NEXT: return %[[ALLOC1]]

mlir/test/lib/Dialect/Test/TestDialect.cpp

void StringAttrPrettyNameOp::getAsmResultNames(
   auto value = names();
   for (size_t i = 0, e = value.size(); i != e; ++i)
     if (auto str = value[i].dyn_cast<StringAttr>())
       if (!str.getValue().empty())
         setNameFn(getResult(i), str.getValue());
 }
 
 //===----------------------------------------------------------------------===//
+// RegionIfOp
+//===----------------------------------------------------------------------===//
+
+static void print(OpAsmPrinter &p, RegionIfOp op) {
+  p << RegionIfOp::getOperationName() << " ";
+  p.printOperands(op.getOperands());
+  p << ": " << op.getOperandTypes();
+  p.printArrowTypeList(op.getResultTypes());
+  p << " then";
+  p.printRegion(op.thenRegion(),
+                /*printEntryBlockArgs=*/true,
+                /*printBlockTerminators=*/true);
+  p << " else";
+  p.printRegion(op.elseRegion(),
+                /*printEntryBlockArgs=*/true,
+                /*printBlockTerminators=*/true);
+  p << " join";
+  p.printRegion(op.joinRegion(),
+                /*printEntryBlockArgs=*/true,
+                /*printBlockTerminators=*/true);
+}
+
+static ParseResult parseRegionIfOp(OpAsmParser &parser,
+                                   OperationState &result) {
+  SmallVector<OpAsmParser::OperandType, 2> operandInfos;
+  SmallVector<Type, 2> operandTypes;
+
+  result.regions.reserve(3);
+  Region *thenRegion = result.addRegion();
+  Region *elseRegion = result.addRegion();
+  Region *joinRegion = result.addRegion();
+
+  // Parse operand, type and arrow type lists.
+  if (parser.parseOperandList(operandInfos) ||
+      parser.parseColonTypeList(operandTypes) ||
+      parser.parseArrowTypeList(result.types))
+    return failure();
+
+  // Parse all attached regions.
+  if (parser.parseKeyword("then") || parser.parseRegion(*thenRegion, {}, {}) ||
+      parser.parseKeyword("else") || parser.parseRegion(*elseRegion, {}, {}) ||
+      parser.parseKeyword("join") || parser.parseRegion(*joinRegion, {}, {}))
+    return failure();
+
+  return parser.resolveOperands(operandInfos, operandTypes,
+                                parser.getCurrentLocation(), result.operands);
+}
+
+OperandRange RegionIfOp::getSuccessorEntryOperands(unsigned index) {
+  assert(index < 2 && "invalid region index");
+  return getOperands();
+}
+
+void RegionIfOp::getSuccessorRegions(
+    Optional<unsigned> index, ArrayRef<Attribute> operands,
+    SmallVectorImpl<RegionSuccessor> &regions) {
+  // We always branch to the join region.
+  if (index.hasValue()) {
+    if (index.getValue() < 2)
+      regions.push_back(RegionSuccessor(&joinRegion(), getJoinArgs()));
+    else
+      regions.push_back(RegionSuccessor(getResults()));
+    return;
+  }
+
+  // The then and else regions are the entry regions of this op.
+  regions.push_back(RegionSuccessor(&thenRegion(), getThenArgs()));
+  regions.push_back(RegionSuccessor(&elseRegion(), getElseArgs()));
+}
+
+//===----------------------------------------------------------------------===//
 // Dialect Registration
 //===----------------------------------------------------------------------===//
 
 // Static initialization for Test dialect registration.
 static mlir::DialectRegistration<mlir::TestDialect> testDialect;
 
 #include "TestOpEnums.cpp.inc"
 #include "TestOpStructs.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

mlir/test/lib/Dialect/Test/TestOps.td

 // Test SideEffects
 //===----------------------------------------------------------------------===//
 
 def SideEffectOp : TEST_Op<"side_effect_op",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let results = (outs AnyType:$result);
 }
 
+//===----------------------------------------------------------------------===//
+// Test RegionBranchOpInterface
+//===----------------------------------------------------------------------===//
+
+def RegionIfYieldOp : TEST_Op<"region_if_yield",
+      [NoSideEffect, ReturnLike, Terminator]> {
+  let arguments = (ins Variadic<AnyType>:$results);
+  let assemblyFormat = [{
+    $results `:` type($results) attr-dict
+  }];
+}
+
+def RegionIfOp : TEST_Op<"region_if",
+      [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
+       SingleBlockImplicitTerminator<"RegionIfYieldOp">,
+       RecursiveSideEffects]> {
+  let description =[{
+    Represents an abstract if-then-else-join pattern. In this context, the then
+    and else regions jump to the join region, which finally returns to its
+    parent op.
+    }];
+
+  let printer = [{ return ::print(p, *this); }];
+  let parser = [{ return ::parseRegionIfOp(parser, result); }];
+  let arguments = (ins Variadic<AnyType>);
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region SizedRegion<1>:$thenRegion,
+                        AnyRegion:$elseRegion,
+                        AnyRegion:$joinRegion);
+  let extraClassDeclaration = [{
+    Block::BlockArgListType getThenArgs() {
+      return getBody(0)->getArguments();
+    }
+    Block::BlockArgListType getElseArgs() {
+      return getBody(1)->getArguments();
+    }
+    Block::BlockArgListType getJoinArgs() {
+      return getBody(2)->getArguments();
+    }
+    OperandRange getSuccessorEntryOperands(unsigned index);
+  }];
+}
+
 #endif // TEST_OPS