[mlir][bufferization] Support unstructured control flow

[springerm]

Summary

This revision adds support for unstructured control flow to the bufferization infrastructure. In particular: regions with multiple blocks, cf.br, cf.cond_br.

Two helper templates are added to BufferizableOpInterface.h, which can be implemented by ops that supported unstructured control flow in their regions (e.g., func.func) and ops that branch to another block (e.g., cf.br).

A block signature is always bufferized together with the op that owns the block.

Depends On: D158154

[springerm]

rebase

[springerm]

fix bug

[maerhart]

Can we have an invalid-test checking for multiple return operations inside a function?
There is no looping test where the buffer type has to resort to the fully dynamic type.

Thanks for adding a lot of comments, it makes understanding the code a lot easier!

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

417

Should there be some max number of revisits such that the compiler does not run into an infinite loop?

735

Those helper templates are a lot of code in this public API. Maybe we should move them to a separate header file UnstructuredControlFlowBufferizableOpImpl.h or so?

756

Shouldn't this be an emitError or assertion? Why is it safe to just skip it?

758

insert "not"

763–766

I'm curious, why does result.push_back(operands.getForwardedOperands()[bbArg.getArgNumber()]) not work?

776

Nit: Should the body of this function be outlined into a protected helper function that is called here with an assertion in front that checks for tensor results to make sure the user of this helper template uses it correctly?

794

But it calls defaultGetBufferType for OpResults below, so is it also the default implementation for those and it's not for block arguments only?

806

Where do we get the guarantee from that this will never be an unranked tensor?

813

Are we still guaranteed that it will end up in the same memory space?

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

416–417

It's not clear to me what you mean by just reading this. Should the Op containing the region with unstructured control-flow bufferize the terminators as well or just the block arguments and the terminators bufferize themselves? It becomes very clear when reading the rest of the PR, but I think the phrasing here could be a bit improved.

mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h

82

Nit: maybe add this parenthesized part right after the 'callers' in the first line of the comment?

mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp

346

Should this check if the function is external/just a declaration? Should we then just use isExternal or isDeclaration instead?

mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp

19

Do we really need this include? I don't immediately see what is used from it.

48

"without"?

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

103

"no"?

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

665

If there is nothing to test here, why does the test then exist? Just to check that the canonicalizer is run beforehand?

mlir/test/Dialect/ControlFlow/one-shot-bufferize.mlir

15

Wouldn't it be better to have the scf.execute_region tests in the SCF directory? And focus on tests where the CF ops are directly inside a func.func here instead?

mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir

20–21

Does scf.execute_region verify that there exists at least one yield or should we add a test-case without any yield?

[springerm]

In D158094#4606232, @maerhart wrote:

There is no looping test where the buffer type has to resort to the fully dynamic type.

This is tested by the @cond_br test case. (It is not looping, but there are two SSA values with different buffer types coming into ^bb1.)

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

417

There is currently no way to keep to track of the number of visits. This would require a change to the signature of findValueInReverseUseDefChain. At the moment, this flag is always false and it is expected that (in the future) a user of findValueInReverseUseDefChain performs extra checks in the lambda function to avoid such looping.

756

Good point. We should assert. (Maybe an error in the future if there are such ops. At the moment it is not possible to emit an error here.)

763–766

There could be more block arguments than forwarded operands. E.g., the branch op could create its own values and ("produced operands") and make them available as a block argument. Such cases are not supported at the moment, but I already wanted to account for this where possible.

813

We could have different memory spaces on incoming op operands. In that case an error is produced below: incoming operands of block argument have inconsistent memory spaces

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

416–417

This is actually up to the users. Either way works but bufferizing the terminator together with the block signatures ensures that the IR will stay valid at any point of time. I extended the documentation here.

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

665

I meant "nothing to check". The test is here so that unstructured control flow in a FuncOp is tested.

mlir/test/Dialect/ControlFlow/one-shot-bufferize.mlir

15

It doesn't really matter which one we use. In either case, we have to use an op from another dialect.

[springerm]

address comments

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 //===- BufferizableOpInterface.h - Bufferizable Ops -------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_
 #define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_
 
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMapInfoVariant.h"
 #include "llvm/ADT/SetVector.h"
 #include <optional>
 
 #include "mlir/Dialect/Bufferization/IR/BufferizationEnums.h.inc"
 
 namespace mlir {
struct TraversalConfig {
 
   /// Specifies whether unknown/non-bufferizable/ops not included in the
   /// OpFilter of BufferizationOptions should be followed.
   bool followUnknownOps = false;
 
   /// Specifies whether OpOperands with a different type that are not the result
   /// of a CastOpInterface op should be followed.
   bool followSameTypeOrCastsOnly = false;
+
+  /// Specifies whether already visited values should be visited again.
+  /// (Note: This can result in infinite looping.)
+  bool revisitAlreadyVisitedValues = false;
 };
 
 /// AnalysisState provides a variety of helper functions for dealing with
 /// tensor values.
 class AnalysisState {
 public:
   /// Determine which OpOperand* will alias with `value` if the op is
   /// bufferized in place. Return all tensor OpOperand* if the op is not

 MLIR_DECLARE_EXPLICIT_TYPE_ID(mlir::bufferization::AnalysisState)
 
 //===----------------------------------------------------------------------===//
 // Bufferization Interfaces
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
+//===----------------------------------------------------------------------===//
+// Helpers for Unstructured Control Flow
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+namespace bufferization {
+
+/// A template that provides a default implementation of `getAliasingOpOperands`
+/// for ops that support unstructured control flow within their regions.
+template <typename ConcreteModel, typename ConcreteOp>
+struct OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel
+    : public BufferizableOpInterface::ExternalModel<ConcreteModel, ConcreteOp> {
+
+  /// Return a list of operands that are forwarded to the given block argument.
+  /// I.e., find all predecessors of the block argument's owner and gather the
+  /// operands that are equivalent to the block argument.
+  static SmallVector<OpOperand *> getCallerOpOperands(BlockArgument bbArg) {
+    SmallVector<OpOperand *> result;
+    Block *block = bbArg.getOwner();
+    for (Operation *caller : block->getUsers()) {
+      auto branchOp = dyn_cast<BranchOpInterface>(caller);
+      if (!branchOp)
+        continue;
+      auto it = llvm::find(caller->getSuccessors(), block);
+      assert(it != caller->getSuccessors().end() && "could find successor");
+      int64_t successorIdx = std::distance(caller->getSuccessors().begin(), it);
+      SuccessorOperands operands = branchOp.getSuccessorOperands(successorIdx);
+      assert(operands.getProducedOperandCount() == 0 &&
+             "produced operands not supported");
+      int64_t operandIndex =
+          operands.getForwardedOperands().getBeginOperandIndex() +
+          bbArg.getArgNumber();
+      result.push_back(&caller->getOpOperand(operandIndex));
+    }
+    return result;
+  }
+
+  AliasingOpOperandList
+  getAliasingOpOperands(Operation *op, Value value,
+                        const AnalysisState &state) const {
+    // This is a default implementation for block arguments only. If the op also
+    // has tensor results, the external model should define its own
+    // implementation and call this base function in case of a block argument.
+    AliasingOpOperandList result;
+    auto bbArg = dyn_cast<BlockArgument>(value);
+    if (!bbArg)
+      return result;
+
+    // Gather aliasing OpOperands of all operations (callers) that link to
+    // this block.
+    for (OpOperand *opOperand : getCallerOpOperands(bbArg))
+      result.addAlias(
+          {opOperand, BufferRelation::Equivalent, /*isDefinite=*/false});
+
+    return result;
+  }
+
+  FailureOr<BaseMemRefType>
+  getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+                SmallVector<Value> &invocationStack) const {
+    // This is a default implementation for block arguments only. If the op also
+    // has tensor results, the external model should define its own
+    // implementation and call this base function in case of a block argument.
+    if (isa<OpResult>(value))
+      return bufferization::detail::defaultGetBufferType(value, options,
+                                                         invocationStack);
+
+    // Compute the buffer type of the block argument by computing the bufferized
+    // operand types of all forwarded values. If these are all the same type,
+    // take that type. Otherwise, take only the memory space and fall back to a
+    // buffer type with a fully dynamic layout map.
+    BaseMemRefType bufferType;
+    auto tensorType = cast<RankedTensorType>(value.getType());
+    for (OpOperand *opOperand :
+         getCallerOpOperands(cast<BlockArgument>(value))) {
+
+      // If the forwarded operand is already on the invocation stack, we ran
+      // into a loop and this operand cannot be used to compute the bufferized
+      // type.
+      if (llvm::find(invocationStack, opOperand->get()) !=
+          invocationStack.end())
+        continue;
+
+      // Compute the bufferized type of the forwarded operand.
+      BaseMemRefType callerType;
+      if (auto memrefType =
+              dyn_cast<BaseMemRefType>(opOperand->get().getType())) {
+        // The operand was already bufferized. Take its type directly.
+        callerType = memrefType;
+      } else {
+        FailureOr<BaseMemRefType> maybeCallerType =
+            bufferization::getBufferType(opOperand->get(), options,
+                                         invocationStack);
+        if (failed(maybeCallerType))
+          return failure();
+        callerType = *maybeCallerType;
+      }
+
+      if (!bufferType) {
+        // This is the first buffer type that we computed.
+        bufferType = callerType;
+        continue;
+      }
+
+      if (bufferType == callerType)
+        continue;
+
+      // If the computed buffer type does not match the computed buffer type of
+      // the earlier forwarded operands, fall back to a buffer type with a fully
+      // dynamic layout map.
+#ifndef NDEBUG
+      assert(llvm::all_equal({bufferType.getShape(), callerType.getShape(),
+                              tensorType.getShape()}) &&
+             "expected same shape");
+#endif // NDEBUG
+
+      if (bufferType.getMemorySpace() != callerType.getMemorySpace())
+        return op->emitOpError("incoming operands of block argument have "
+                               "inconsistent memory spaces");
+
+      bufferType = getMemRefTypeWithFullyDynamicLayout(
+          tensorType, bufferType.getMemorySpace());
+    }
+
+    if (!bufferType)
+      return op->emitOpError("could not infer buffer type of block argument");
+
+    return bufferType;
+  }
+};
+
+/// A template that provides a default implementation of `getAliasingValues`
+/// for ops that implement the `BranchOpInterface`.
+template <typename ConcreteModel, typename ConcreteOp>
+struct BranchOpBufferizableOpInterfaceExternalModel
+    : public BufferizableOpInterface::ExternalModel<ConcreteModel, ConcreteOp> {
+  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
+                                      const AnalysisState &state) const {
+    AliasingValueList result;
+    auto branchOp = cast<BranchOpInterface>(op);
+    auto operandNumber = opOperand.getOperandNumber();
+
+    // Gather aliasing block arguments of blocks to which this op may branch to.
+    for (const auto &it : llvm::enumerate(op->getSuccessors())) {
+      Block *block = it.value();
+      SuccessorOperands operands = branchOp.getSuccessorOperands(it.index());
+      assert(operands.getProducedOperandCount() == 0 &&
+             "produced operands not supported");
+      // The first and last operands that are forwarded to this successor.
+      int64_t firstOperandIndex =
+          operands.getForwardedOperands().getBeginOperandIndex();
+      int64_t lastOperandIndex =
+          firstOperandIndex + operands.getForwardedOperands().size();
+      bool matchingDestination = operandNumber >= firstOperandIndex &&
+                                 operandNumber < lastOperandIndex;
+      // A branch op may have multiple successors. Find the ones that correspond
+      // to this OpOperand. (There is usually only one.)
+      if (!matchingDestination)
+        continue;
+      // Compute the matching block argument of the destination block.
+      BlockArgument bbArg =
+          block->getArgument(operandNumber - firstOperandIndex);
+      result.addAlias(
+          {bbArg, BufferRelation::Equivalent, /*isDefinite=*/false});
+    }
+
+    return result;
+  }
+};
+
+} // namespace bufferization
+} // namespace mlir
+
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

let methods = [
           tensor operand/result.
 
           The return value of this method indicates whether there was an error
           while bufferizing this op (such as failing to create a new buffer
           allocation op). The bufferization driver immediately stops bufferizing
           the input IR and returns `failure` in that case. If this op is
           expected to survive bufferization, `success` should be returned
           (together with `allow-unknown-ops` enabled).
+
+          Note: If this op supports unstructured control flow in its regions,
+          then this function should also bufferize all block signatures that
+          belong to this op. Branch ops (that branch to a block) are typically
+          bufferized together with the block signature.
         }],
         /*retType=*/"::mlir::LogicalResult",
         /*methodName=*/"bufferize",
         /*args=*/(ins "::mlir::RewriterBase &":$rewriter,
                       "const ::mlir::bufferization::BufferizationOptions &":$options),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           llvm_unreachable("bufferize not implemented");

mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h

 ///
 /// Note: This function bufferizes ops without utilizing analysis results. It
 /// can be used to implement partial bufferization passes.
 LogicalResult bufferizeOp(Operation *op, const BufferizationOptions &options,
                           bool copyBeforeWrite = true,
                           const OpFilter *opFilter = nullptr,
                           BufferizationStatistics *statistics = nullptr);
 
+/// Bufferize the signature of `block` and its callers. All block argument types
+/// are changed to memref types. All corresponding operands of all callers
+/// (i.e., ops that have the given block as a successor) are wrapped in
+/// bufferization.to_memref ops.
+///
+/// It is expected that all callers implement the `BranchOpInterface`.
+/// Otherwise, this function will fail. The `BranchOpInterface` is used to query
+/// the range of operands that are forwarded to this block.
+///
+/// It is expected that the parent op of this block implements the
+/// `BufferizableOpInterface`. The buffer types of tensor block arugments are
+/// computed with `BufferizableOpIntercace::getBufferType`.
+LogicalResult bufferizeBlockSignature(Block *block, RewriterBase &rewriter,
+                                      const BufferizationOptions &options);
+
 BufferizationOptions getPartialBufferizationOptions();
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

mlir/include/mlir/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.h

+//===- BufferizableOpInterfaceImpl.h - Impl. of BufferizableOpInterface ---===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_CONTROLFLOW_BUFFERIZABLEOPINTERFACEIMPL_H
+#define MLIR_DIALECT_CONTROLFLOW_BUFFERIZABLEOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace cf {
+void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace cf
+} // namespace mlir
+
+#endif // MLIR_DIALECT_CONTROLFLOW_BUFFERIZABLEOPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/ArmSME/IR/ArmSME.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
+#include "mlir/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/TransformOps/FuncTransformOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
 #include "mlir/Dialect/IRDL/IR/IRDL.h"
 #include "mlir/Dialect/Index/IR/IndexDialect.h"
inline void registerAllDialects(DialectRegistry &registry) {
 
   // Register all external models.
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   arith::registerValueBoundsOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   builtin::registerCastOpInterfaceExternalModels(registry);
+  cf::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
   linalg::registerValueBoundsOpInterfaceExternalModels(registry);
   memref::registerBufferizableOpInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   memref::registerValueBoundsOpInterfaceExternalModels(registry);
   memref::registerMemorySlotExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

for (OpOperand &opOperand : op->getOpOperands()) {
     // entirely? In either case, mark the allocation as "escaping", so that it
     // will not be deallocated.
     bool escape = !state.getOptions().createDeallocs ||
                   llvm::any_of(aliasingValues, [&](AliasingValue a) {
                     return state.isTensorYielded(a.value);
                   });
 
     if (aliasingValues.getNumAliases() == 1 &&
+        isa<OpResult>(aliasingValues.getAliases()[0].value) &&
         !state.bufferizesToMemoryWrite(opOperand) &&
         state.getAliasingOpOperands(aliasingValues.getAliases()[0].value)
                 .getNumAliases() == 1 &&
         !isa<UnrankedTensorType>(
             aliasingValues.getAliases()[0].value.getType())) {
       // The op itself does not write but may create exactly one alias. Instead
       // of copying the OpOperand, copy the OpResult. The OpResult can sometimes
       // be smaller than the OpOperand (e.g., in the case of an extract_slice,

 }
 
 /// Return true if the given value is read by an op that bufferizes to a memory
 /// read. Also takes into account ops that create an alias but do not read by
 /// themselves (e.g., ExtractSliceOp).
 bool AnalysisState::isValueRead(Value value) const {
   assert(llvm::isa<TensorType>(value.getType()) && "expected TensorType");
   SmallVector<OpOperand *> workingSet;
+  DenseSet<OpOperand *> visited;
   for (OpOperand &use : value.getUses())
     workingSet.push_back(&use);
 
   while (!workingSet.empty()) {
     OpOperand *uMaybeReading = workingSet.pop_back_val();
+    if (visited.contains(uMaybeReading))
+      continue;
+    visited.insert(uMaybeReading);
+
     // Skip over all ops that neither read nor write (but create an alias).
     if (bufferizesToAliasOnly(*uMaybeReading))
       for (AliasingValue alias : getAliasingValues(*uMaybeReading))
         for (OpOperand &use : alias.value.getUses())
           workingSet.push_back(&use);
     if (bufferizesToMemoryRead(*uMaybeReading))
       return true;
   }
 
   return false;
 }
 
 // Starting from `value`, follow the use-def chain in reverse, always selecting
 // the aliasing OpOperands. Find and return Values for which `condition`
 // evaluates to true. OpOperands of such matching Values are not traversed any
 // further.
 llvm::SetVector<Value> AnalysisState::findValueInReverseUseDefChain(
     Value value, llvm::function_ref<bool(Value)> condition,
     TraversalConfig config) const {
+  llvm::DenseSet<Value> visited;
   llvm::SetVector<Value> result, workingSet;
   workingSet.insert(value);
 
   while (!workingSet.empty()) {
     Value value = workingSet.pop_back_val();
+
+    if (!config.revisitAlreadyVisitedValues && visited.contains(value)) {
+      // Stop traversal if value was already visited.
+      if (config.alwaysIncludeLeaves)
+        result.insert(value);
+      continue;
+    }
+    visited.insert(value);
+
     if (condition(value)) {
       result.insert(value);
       continue;
     }
 
     if (!config.followUnknownOps && !options.dynCastBufferizableOp(value)) {
       // Stop iterating if `followUnknownOps` is unset and the op is either
       // not bufferizable or excluded in the OpFilter.
bool AnalysisState::isTensorYielded(Value tensor) const {
   // In the absence of analysis information, the conservative answer is "true".
   if (!tensor.getDefiningOp<AllocTensorOp>())
     return true;
 
   // For AllocTensorOp results, we can do better: They do not alias with any
   // preceding value, so we can follow SSA use-def chains and do a simple
   // analysis.
   SmallVector<OpOperand *> worklist;
+  DenseSet<OpOperand *> visited;
   for (OpOperand &use : tensor.getUses())
     worklist.push_back(&use);
 
   while (!worklist.empty()) {
     OpOperand *operand = worklist.pop_back_val();
+    if (visited.contains(operand))
+      continue;
+    visited.insert(operand);
     Operation *op = operand->getOwner();
 
     // If the op is not bufferizable, we can safely assume that the value is not
     // yielded. (When bufferizing that op, it must handle such cases.)
     if (!options.dynCastBufferizableOp(op))
       continue;
 
     // We cannot analyze through ToMemrefOps, so we have to conservatively

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 #include <optional>
 
 namespace mlir {
 namespace bufferization {

 //===----------------------------------------------------------------------===//
 // BufferizableOpInterface-based Bufferization
 //===----------------------------------------------------------------------===//
 
 static bool isaTensor(Type t) { return isa<TensorType>(t); }
 
 /// Return true if the given op has a tensor result or a tensor operand.
 static bool hasTensorSemantics(Operation *op) {
+  bool hasTensorBlockArgument = any_of(op->getRegions(), [](Region &r) {
+    return any_of(r.getBlocks(), [](Block &b) {
+      return any_of(b.getArguments(), [](BlockArgument bbArg) {
+        return isaTensor(bbArg.getType());
+      });
+    });
+  });
+  if (hasTensorBlockArgument)
+    return true;
+
   if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
     bool hasTensorArg = any_of(funcOp.getArgumentTypes(), isaTensor);
     bool hasTensorResult = any_of(funcOp.getResultTypes(), isaTensor);
     return hasTensorArg || hasTensorResult;
   }
 
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
for (Operation *op : worklist) {
     if (isa<ToTensorOp, ToMemrefOp>(op))
       continue;
     return op->emitError("op was not bufferized");
   }
 
   return success();
 }
 
+LogicalResult
+bufferization::bufferizeBlockSignature(Block *block, RewriterBase &rewriter,
+                                       const BufferizationOptions &options) {
+  OpBuilder::InsertionGuard g(rewriter);
+  auto bufferizableOp = options.dynCastBufferizableOp(block->getParentOp());
+  if (!bufferizableOp)
+    return failure();
+
+  // Compute the new signature.
+  SmallVector<Type> newTypes;
+  for (BlockArgument &bbArg : block->getArguments()) {
+    auto tensorType = dyn_cast<TensorType>(bbArg.getType());
+    if (!tensorType) {
+      newTypes.push_back(bbArg.getType());
+      continue;
+    }
+
+    FailureOr<BaseMemRefType> memrefType =
+        bufferization::getBufferType(bbArg, options);
+    if (failed(memrefType))
+      return failure();
+    newTypes.push_back(*memrefType);
+  }
+
+  // Change the type of all block arguments.
+  for (auto [bbArg, type] : llvm::zip(block->getArguments(), newTypes)) {
+    if (bbArg.getType() == type)
+      continue;
+
+    // Collect all uses of the bbArg.
+    SmallVector<OpOperand *> bbArgUses;
+    for (OpOperand &use : bbArg.getUses())
+      bbArgUses.push_back(&use);
+
+    // Change the bbArg type to memref.
+    bbArg.setType(type);
+
+    // Replace all uses of the original tensor bbArg.
+    rewriter.setInsertionPointToStart(block);
+    if (!bbArgUses.empty()) {
+      Value toTensorOp =
+          rewriter.create<bufferization::ToTensorOp>(bbArg.getLoc(), bbArg);
+      for (OpOperand *use : bbArgUses)
+        use->set(toTensorOp);
+    }
+  }
+
+  // Bufferize callers of the block.
+  for (Operation *op : block->getUsers()) {
+    auto branchOp = dyn_cast<BranchOpInterface>(op);
+    if (!branchOp)
+      return op->emitOpError("cannot bufferize ops with block references that "
+                             "do not implement BranchOpInterface");
+
+    auto it = llvm::find(op->getSuccessors(), block);
+    assert(it != op->getSuccessors().end() && "could find successor");
+    int64_t successorIdx = std::distance(op->getSuccessors().begin(), it);
+
+    SuccessorOperands operands = branchOp.getSuccessorOperands(successorIdx);
+    SmallVector<Value> newOperands;
+    for (auto [operand, type] :
+         llvm::zip(operands.getForwardedOperands(), newTypes)) {
+      if (operand.getType() == type) {
+        // Not a tensor type. Nothing to do for this operand.
+        newOperands.push_back(operand);
+        continue;
+      }
+      FailureOr<BaseMemRefType> operandBufferType =
+          bufferization::getBufferType(operand, options);
+      if (failed(operandBufferType))
+        return failure();
+      rewriter.setInsertionPointAfterValue(operand);
+      Value bufferizedOperand = rewriter.create<bufferization::ToMemrefOp>(
+          operand.getLoc(), *operandBufferType, operand);
+      // A cast is needed if the operand and the block argument have different
+      // bufferized types.
+      if (type != *operandBufferType)
+        bufferizedOperand = rewriter.create<memref::CastOp>(
+            operand.getLoc(), type, bufferizedOperand);
+      newOperands.push_back(bufferizedOperand);
+    }
+    operands.getMutableForwardedOperands().assign(newOperands);
+  }
+
+  return success();
+}
+
 BufferizationOptions bufferization::getPartialBufferizationOptions() {
   BufferizationOptions options;
   options.allowUnknownOps = true;
   options.createDeallocs = false;
   options.enforceAliasingInvariants = false;
   options.unknownTypeConverterFn = [](Value value, Attribute memorySpace,
                                       const BufferizationOptions &options) {
     return getMemRefTypeWithStaticIdentityLayout(
         cast<TensorType>(value.getType()), memorySpace);
   };
   options.opFilter.allowDialect<BufferizationDialect>();
   return options;
 }

mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp

 //===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 #include <optional>
 
 namespace mlir {

 #endif // NDEBUG
 
     // ReturnOps are bufferized as part of FuncOps.
     return success();
   }
 };
 
 struct FuncOpInterface
-    : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
+    : public OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel<
+          FuncOpInterface, FuncOp> {
+
+  static bool supportsUnstructuredControlFlow() { return true; }
+
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
                 SmallVector<Value> &invocationStack) const {
     auto funcOp = cast<FuncOp>(op);
     auto bbArg = cast<BlockArgument>(value);
-    // Unstructured control flow is not supported.
-    assert(bbArg.getOwner() == &funcOp.getBody().front() &&
-           "expected that block argument belongs to first block");
-    return getBufferizedFunctionArgType(funcOp, bbArg.getArgNumber(), options);
+
+    // Function arguments are special.
+    if (bbArg.getOwner() == &funcOp.getBody().front())
+      return getBufferizedFunctionArgType(funcOp, bbArg.getArgNumber(),
+                                          options);
+
+    return OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel::
+        getBufferType(op, value, options, invocationStack);
+  }
+
+  LogicalResult verifyAnalysis(Operation *op,
+                               const AnalysisState &state) const {
+    auto funcOp = cast<func::FuncOp>(op);
+    // TODO: func.func with multiple returns are not supported.
+    if (!getAssumedUniqueReturnOp(funcOp) && !funcOp.getBody().empty())
+      return op->emitOpError("op without unique func.return is not supported");
+    const auto &options =
+        static_cast<const OneShotBufferizationOptions &>(state.getOptions());
+    // allow-return-allocs is required for ops with multiple blocks.
+    if (options.allowReturnAllocs || funcOp.getRegion().getBlocks().size() <= 1)
+      return success();
+    return op->emitOpError(
+        "op cannot be bufferized without allow-return-allocs");
   }
 
   /// Rewrite function bbArgs and return values into buffer form. This function
   /// bufferizes the function signature and the ReturnOp. When the entire
   /// function body has been bufferized, function return types can be switched
   /// to more concise memref types as part of `foldMemRefCasts`.
   ///
   /// All function bbArgs are writable unless they are explicitly marked as
if (funcOp.getBody().empty()) {
       return success();
     }
 
     // TODO: Support functions with multiple returns.
     func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
     assert(returnOp && "expected func with single return op");
     Location loc = returnOp.getLoc();
 
-    // 1. Rewrite the bbArgs. Turn every tensor bbArg into a memref bbArg.
-    Block &frontBlock = funcOp.getBody().front();
-    for (BlockArgument &bbArg : frontBlock.getArguments()) {
-      auto tensorType = dyn_cast<TensorType>(bbArg.getType());
-      // Non-tensor types stay the same.
-      if (!tensorType)
-        continue;
-
-      // Collect all uses of the bbArg.
-      SmallVector<OpOperand *> bbArgUses;
-      for (OpOperand &use : bbArg.getUses())
-        bbArgUses.push_back(&use);
-
-      // Change the bbArg type to memref.
-      FailureOr<BaseMemRefType> memrefType =
-          bufferization::getBufferType(bbArg, options);
-      if (failed(memrefType))
+    // 1. Bufferize every block.
+    for (Block &block : funcOp.getBody())
+      if (failed(bufferization::bufferizeBlockSignature(&block, rewriter,
+                                                        options)))
         return failure();
-      bbArg.setType(*memrefType);
-
-      // Replace all uses of the original tensor bbArg.
-      rewriter.setInsertionPointToStart(&frontBlock);
-      if (!bbArgUses.empty()) {
-        // Insert to_tensor because the remaining function body has not been
-        // bufferized yet.
-        Value toTensorOp =
-            rewriter.create<bufferization::ToTensorOp>(funcOp.getLoc(), bbArg);
-        for (OpOperand *use : bbArgUses)
-          use->set(toTensorOp);
-      }
-    }
 
     // 2. For each result, keep track of which inplace argument it reuses.
     SmallVector<Value> returnValues;
     for (OpOperand &returnOperand : returnOp->getOpOperands()) {
       Value returnVal = returnOperand.get();
       auto tensorType = dyn_cast<TensorType>(returnVal.getType());
       rewriter.setInsertionPoint(returnOp);
 
struct FuncOpInterface
 
   /// Return `true` if the given function argument is writable.
   bool isWritable(Operation *op, Value value,
                   const AnalysisState &state) const {
     auto funcOp = cast<FuncOp>(op);
     BlockArgument bbArg = dyn_cast<BlockArgument>(value);
     assert(bbArg && "expected BlockArgument");
 
+    // Non-entry block arguments are always writable. (They may alias with
+    // values that are not writable, which will turn them into read-only.)
+    if (bbArg.getOwner() != &funcOp.getBody().front())
+      return true;
+
     // "bufferization.writable" overrides other writability decisions. This is
     // currently used for testing only.
     if (BoolAttr writable = funcOp.getArgAttrOfType<BoolAttr>(
             bbArg.getArgNumber(), BufferizationDialect::kWritableAttrName))
       return writable.getValue();
 
     // All function arguments are writable by default.
     return true;

mlir/lib/Dialect/ControlFlow/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(Transforms)

mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp

+//===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+using namespace mlir::bufferization;
+
+namespace mlir {
+namespace cf {
+namespace {
+
+template <typename ConcreteModel, typename ConcreteOp>
+struct BranchLikeOpInterface
+    : public BranchOpBufferizableOpInterfaceExternalModel<ConcreteModel,
+                                                          ConcreteOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  LogicalResult verifyAnalysis(Operation *op,
+                               const AnalysisState &state) const {
+    const auto &options =
+        static_cast<const OneShotBufferizationOptions &>(state.getOptions());
+    if (options.allowReturnAllocs)
+      return success();
+    return op->emitOpError("op cannot be bufferized with allow-return-allocs");
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options) const {
+    // The operands of this op are bufferized together with the block signature.
+    return success();
+  }
+};
+
+/// Bufferization of cf.br.
+struct BranchOpInterface
+    : public BranchLikeOpInterface<BranchOpInterface, cf::BranchOp> {};
+
+/// Bufferization of cf.cond_br.
+struct CondBranchOpInterface
+    : public BranchLikeOpInterface<CondBranchOpInterface, cf::CondBranchOp> {};
+
+} // namespace
+} // namespace cf
+} // namespace mlir
+
+void mlir::cf::registerBufferizableOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, cf::ControlFlowDialect *dialect) {
+    cf::BranchOp::attachInterface<BranchOpInterface>(*ctx);
+    cf::CondBranchOp::attachInterface<CondBranchOpInterface>(*ctx);
+  });
+}

mlir/lib/Dialect/ControlFlow/Transforms/CMakeLists.txt

+add_mlir_dialect_library(MLIRControlFlowTransforms
+  BufferizableOpInterfaceImpl.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ControlFlow/Transforms
+
+  LINK_LIBS PUBLIC
+  MLIRBufferizationDialect
+  MLIRBufferizationTransforms
+  MLIRControlFlowDialect
+  MLIRMemRefDialect
+  MLIRIR
+  )

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

 //===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h"
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     }
 
     replaceOpWithNewBufferizedOp<scf::ConditionOp>(
         rewriter, op, conditionOp.getCondition(), newArgs);
     return success();
   }
 };
 
+/// Return the unique scf.yield op. If there are multiple or none scf.yield ops,
+/// return an empty op.
+static scf::YieldOp getUniqueYieldOp(scf::ExecuteRegionOp executeRegionOp) {
+  scf::YieldOp result;
+  for (Block &block : executeRegionOp.getRegion()) {
+    if (auto yieldOp = dyn_cast<scf::YieldOp>(block.getTerminator())) {
+      if (result)
+        return {};
+      result = yieldOp;
+    }
+  }
+  return result;
+}
+
 /// Bufferization of scf.execute_region. Can be analyzed, but bufferization not
 /// fully implemented at the moment.
 struct ExecuteRegionOpInterface
-    : public BufferizableOpInterface::ExternalModel<ExecuteRegionOpInterface,
-                                                    scf::ExecuteRegionOp> {
+    : public OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel<
+          ExecuteRegionOpInterface, scf::ExecuteRegionOp> {
+
+  static bool supportsUnstructuredControlFlow() { return true; }
+
+  bool isWritable(Operation *op, Value value,
+                  const AnalysisState &state) const {
+    return true;
+  }
+
+  LogicalResult verifyAnalysis(Operation *op,
+                               const AnalysisState &state) const {
+    auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
+    // TODO: scf.execute_region with multiple yields are not supported.
+    if (!getUniqueYieldOp(executeRegionOp))
+      return op->emitOpError("op without unique scf.yield is not supported");
+    const auto &options =
+        static_cast<const OneShotBufferizationOptions &>(state.getOptions());
+    // allow-return-allocs is required for ops with multiple blocks.
+    if (options.allowReturnAllocs ||
+        executeRegionOp.getRegion().getBlocks().size() == 1)
+      return success();
+    return op->emitOpError(
+        "op cannot be bufferized without allow-return-allocs");
+  }
+
   AliasingOpOperandList
   getAliasingOpOperands(Operation *op, Value value,
                         const AnalysisState &state) const {
+    if (isa<BlockArgument>(value))
+      return OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel::
+          getAliasingOpOperands(op, value, state);
+
     // ExecuteRegionOps do not have tensor OpOperands. The yielded value can be
     // any SSA value that is in scope. To allow for use-def chain traversal
     // through ExecuteRegionOps in the analysis, the corresponding yield value
     // is considered to be aliasing with the result.
     auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
-    size_t resultNum = std::distance(op->getOpResults().begin(),
-                                     llvm::find(op->getOpResults(), value));
-    // TODO: Support multiple blocks.
-    assert(executeRegionOp.getRegion().getBlocks().size() == 1 &&
-           "expected exactly 1 block");
-    auto yieldOp = dyn_cast<scf::YieldOp>(
-        executeRegionOp.getRegion().front().getTerminator());
-    assert(yieldOp && "expected scf.yield terminator in scf.execute_region");
+    auto it = llvm::find(op->getOpResults(), value);
+    assert(it != op->getOpResults().end() && "invalid value");
+    size_t resultNum = std::distance(op->getOpResults().begin(), it);
+    auto yieldOp = getUniqueYieldOp(executeRegionOp);
+    // Note: If there is no unique scf.yield op, `verifyAnalysis` will fail.
+    if (!yieldOp)
+      return {};
     return {{&yieldOp->getOpOperand(resultNum), BufferRelation::Equivalent}};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
-    assert(executeRegionOp.getRegion().getBlocks().size() == 1 &&
-           "only 1 block supported");
-    auto yieldOp =
-        cast<scf::YieldOp>(executeRegionOp.getRegion().front().getTerminator());
+    auto yieldOp = getUniqueYieldOp(executeRegionOp);
     TypeRange newResultTypes(yieldOp.getResults());
 
     // Create new op and move over region.
     auto newOp =
         rewriter.create<scf::ExecuteRegionOp>(op->getLoc(), newResultTypes);
     newOp.getRegion().takeBody(executeRegionOp.getRegion());
 
+    // Bufferize every block.
+    for (Block &block : newOp.getRegion())
+      if (failed(bufferization::bufferizeBlockSignature(&block, rewriter,
+                                                        options)))
+        return failure();
+
     // Update all uses of the old op.
     rewriter.setInsertionPointAfter(newOp);
     SmallVector<Value> newResults;
     for (const auto &it : llvm::enumerate(executeRegionOp->getResultTypes())) {
       if (isa<TensorType>(it.value())) {
         newResults.push_back(rewriter.create<bufferization::ToTensorOp>(
             executeRegionOp.getLoc(), newOp->getResult(it.index())));
       } else {

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

func.func @to_memref_op_unsupported(
   "test.foo"(%0) : (memref<?xf32>) -> ()
 
   // CHECK: vector.transfer_read %[[arg0]]
   %cst = arith.constant 0.0 : f32
   %r1 = vector.transfer_read %t1[%idx3], %cst : tensor<?xf32>, vector<5xf32>
 
   return %r1 : vector<5xf32>
 }
+
+// -----
+
+// Note: The cf.br canonicalizes away, so there's nothing to test here. There is
+// a detailed test in ControlFlow/bufferize.mlir.
+
+// CHECK-LABEL: func @br_in_func(
+func.func @br_in_func(%t: tensor<5xf32>) -> tensor<5xf32> {
+  cf.br ^bb1(%t : tensor<5xf32>)
+^bb1(%arg1 : tensor<5xf32>):
+  func.return %arg1 : tensor<5xf32>
+}

mlir/test/Dialect/ControlFlow/one-shot-bufferize-analysis.mlir

+// RUN: mlir-opt -one-shot-bufferize="allow-return-allocs test-analysis-only dump-alias-sets bufferize-function-boundaries" -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: func @single_branch(
+//  CHECK-SAME:     {__bbarg_alias_set_attr__ = [{{\[}}[{{\[}}"%[[arg1:.*]]", "%[[t:.*]]"]], [{{\[}}"%[[arg1]]", "%[[t]]"]]]]}
+func.func @single_branch(%t: tensor<5xf32>) -> tensor<5xf32> {
+// CHECK: cf.br
+// CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
+  cf.br ^bb1(%t : tensor<5xf32>)
+// CHECK: ^{{.*}}(%[[arg1]]: tensor<5xf32>)
+^bb1(%arg1 : tensor<5xf32>):
+  func.return %arg1 : tensor<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @diamond_branch(
+//  CHECK-SAME:     %{{.*}}: i1, %[[t0:.*]]: tensor<5xf32> {{.*}}, %[[t1:.*]]: tensor<5xf32> {{.*}}) -> tensor<5xf32>
+//  CHECK-SAME:     {__bbarg_alias_set_attr__ = [{{\[}}[{{\[}}"%[[arg1:.*]]", "%[[arg3:.*]]", "%[[arg2:.*]]", "%[[t0]]", "%[[t1]]"], [
+func.func @diamond_branch(%c: i1, %t0: tensor<5xf32>, %t1: tensor<5xf32>) -> tensor<5xf32> {
+// CHECK: cf.cond_br
+// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "true"]}
+  cf.cond_br %c, ^bb1(%t0 : tensor<5xf32>), ^bb2(%t1 : tensor<5xf32>)
+// CHECK: ^{{.*}}(%[[arg1]]: tensor<5xf32>):
+^bb3(%arg1 : tensor<5xf32>):
+  func.return %arg1 : tensor<5xf32>
+// CHECK: ^{{.*}}(%[[arg2]]: tensor<5xf32>):
+^bb1(%arg2 : tensor<5xf32>):
+// CHECK: cf.br
+// CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
+  cf.br ^bb3(%arg2 : tensor<5xf32>)
+// CHECK: ^{{.*}}(%[[arg3]]: tensor<5xf32>):
+^bb2(%arg3 : tensor<5xf32>):
+// CHECK: cf.br
+// CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
+  cf.br ^bb3(%arg3 : tensor<5xf32>)
+}
+
+// -----
+
+// CHECK-LABEL: func @looping_branches(
+//  CHECK-SAME:     {__bbarg_alias_set_attr__ = [{{\[}}[], [{{\[}}"%[[arg2:.*]]", "%[[arg1:.*]]", "%[[inserted:.*]]", "%[[empty:.*]]"]], [
+func.func @looping_branches() -> tensor<5xf32> {
+// CHECK: %[[empty]] = tensor.empty()
+  %0 = tensor.empty() : tensor<5xf32>
+// CHECK: cf.br
+// CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
+  cf.br ^bb1(%0: tensor<5xf32>)
+// CHECK: ^{{.*}}(%[[arg1]]: tensor<5xf32>):
+^bb1(%arg1: tensor<5xf32>):
+  %pos = "test.foo"() : () -> (index)
+  %val = "test.bar"() : () -> (f32)
+// CHECK: %[[inserted]] = tensor.insert
+// CHECK-SAME: __inplace_operands_attr__ = ["none", "true", "none"]
+  %inserted = tensor.insert %val into %arg1[%pos] : tensor<5xf32>
+  %cond = "test.qux"() : () -> (i1)
+// CHECK: cf.cond_br
+// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "true"]}
+  cf.cond_br %cond, ^bb1(%inserted: tensor<5xf32>), ^bb2(%inserted: tensor<5xf32>)
+^bb2(%arg2: tensor<5xf32>):
+  func.return %arg2 : tensor<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @looping_branches_with_conflict(
+func.func @looping_branches_with_conflict(%f: f32) -> tensor<5xf32> {
+  %0 = tensor.empty() : tensor<5xf32>
+  %filled = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
+// CHECK: cf.br
+// CHECK-SAME: {__inplace_operands_attr__ = ["false"]}
+  cf.br ^bb1(%filled: tensor<5xf32>)
+^bb2(%arg2: tensor<5xf32>):
+  %pos2 = "test.foo"() : () -> (index)
+  // One OpOperand cannot bufferize in-place because an "old" value is read.
+  %element = tensor.extract %filled[%pos2] : tensor<5xf32>
+  func.return %arg2 : tensor<5xf32>
+^bb1(%arg1: tensor<5xf32>):
+  %pos = "test.foo"() : () -> (index)
+  %val = "test.bar"() : () -> (f32)
+// CHECK: tensor.insert
+// CHECK-SAME: __inplace_operands_attr__ = ["none", "true", "none"]
+  %inserted = tensor.insert %val into %arg1[%pos] : tensor<5xf32>
+  %cond = "test.qux"() : () -> (i1)
+// CHECK: cf.cond_br
+// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "true"]}
+  cf.cond_br %cond, ^bb1(%inserted: tensor<5xf32>), ^bb2(%inserted: tensor<5xf32>)
+}

mlir/test/Dialect/ControlFlow/one-shot-bufferize-invalid.mlir

+// RUN: mlir-opt -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -split-input-file %s -verify-diagnostics
+
+// expected-error @below{{failed to bufferize op}}
+// expected-error @below{{incoming operands of block argument have inconsistent memory spaces}}
+func.func @inconsistent_memory_space() -> tensor<5xf32> {
+  %0 = bufferization.alloc_tensor() {memory_space = 0 : ui64} : tensor<5xf32>
+  cf.br ^bb1(%0: tensor<5xf32>)
+^bb1(%arg1: tensor<5xf32>):
+  func.return %arg1 : tensor<5xf32>
+^bb2():
+  %1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<5xf32>
+  cf.br ^bb1(%1: tensor<5xf32>)
+}
+
+// -----
+
+// expected-error @below{{failed to bufferize op}}
+// expected-error @below{{could not infer buffer type of block argument}}
+func.func @cannot_infer_type() {
+  return
+  // The type of the block argument cannot be inferred.
+^bb1(%t: tensor<5xf32>):
+  cf.br ^bb1(%t: tensor<5xf32>)
+}

mlir/test/Dialect/ControlFlow/one-shot-bufferize.mlir

+// RUN: mlir-opt -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -one-shot-bufferize="allow-return-allocs" -split-input-file %s | FileCheck %s --check-prefix=CHECK-NO-FUNC
+
+// CHECK-NO-FUNC-LABEL: func @br(
+//  CHECK-NO-FUNC-SAME:     %[[t:.*]]: tensor<5xf32>)
+//       CHECK-NO-FUNC:   %[[m:.*]] = bufferization.to_memref %[[t]] : memref<5xf32, strided<[?], offset: ?>>
+//       CHECK-NO-FUNC:   %[[r:.*]] = scf.execute_region -> memref<5xf32, strided<[?], offset: ?>> {
+//       CHECK-NO-FUNC:     cf.br ^[[block:.*]](%[[m]]
+//       CHECK-NO-FUNC:   ^[[block]](%[[arg1:.*]]: memref<5xf32, strided<[?], offset: ?>>):
+//       CHECK-NO-FUNC:     scf.yield %[[arg1]]
+//       CHECK-NO-FUNC:   }
+//       CHECK-NO-FUNC:   return
+func.func @br(%t: tensor<5xf32>) {
+  %0 = scf.execute_region -> tensor<5xf32> {
+    cf.br ^bb1(%t : tensor<5xf32>)
+  ^bb1(%arg1 : tensor<5xf32>):
+    scf.yield %arg1 : tensor<5xf32>
+  }
+  return
+}
+
+// -----
+
+// CHECK-NO-FUNC-LABEL: func @cond_br(
+//  CHECK-NO-FUNC-SAME:     %[[t1:.*]]: tensor<5xf32>,
+//       CHECK-NO-FUNC:   %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<5xf32, strided<[?], offset: ?>>
+//       CHECK-NO-FUNC:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
+//       CHECK-NO-FUNC:   %[[r:.*]] = scf.execute_region -> memref<5xf32, strided<[?], offset: ?>> {
+//       CHECK-NO-FUNC:     cf.cond_br %{{.*}}, ^[[block1:.*]](%[[m1]] : {{.*}}), ^[[block2:.*]](%[[alloc]] : {{.*}})
+//       CHECK-NO-FUNC:   ^[[block1]](%[[arg1:.*]]: memref<5xf32, strided<[?], offset: ?>>):
+//       CHECK-NO-FUNC:     scf.yield %[[arg1]]
+//       CHECK-NO-FUNC:   ^[[block2]](%[[arg2:.*]]: memref<5xf32>):
+//       CHECK-NO-FUNC:     %[[cast:.*]] = memref.cast %[[arg2]] : memref<5xf32> to memref<5xf32, strided<[?], offset: ?>
+//       CHECK-NO-FUNC:     cf.br ^[[block1]](%[[cast]] : {{.*}})
+//       CHECK-NO-FUNC:   }
+//       CHECK-NO-FUNC:   return
+func.func @cond_br(%t1: tensor<5xf32>, %c: i1) {
+  // Use an alloc for the second block instead of a function block argument.
+  // A cast must be inserted because the two will have different layout maps.
+  %t0 = bufferization.alloc_tensor() : tensor<5xf32>
+  %0 = scf.execute_region -> tensor<5xf32> {
+    cf.cond_br %c, ^bb1(%t1 : tensor<5xf32>), ^bb2(%t0 : tensor<5xf32>)
+  ^bb1(%arg1 : tensor<5xf32>):
+    scf.yield %arg1 : tensor<5xf32>
+  ^bb2(%arg2 : tensor<5xf32>):
+    cf.br ^bb1(%arg2 : tensor<5xf32>)
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @looping_branches(
+func.func @looping_branches() -> tensor<5xf32> {
+// CHECK: %[[alloc:.*]] = memref.alloc
+  %0 = bufferization.alloc_tensor() : tensor<5xf32>
+// CHECK: cf.br {{.*}}(%[[alloc]]
+  cf.br ^bb1(%0: tensor<5xf32>)
+// CHECK: ^{{.*}}(%[[arg1:.*]]: memref<5xf32>):
+^bb1(%arg1: tensor<5xf32>):
+  %pos = "test.foo"() : () -> (index)
+  %val = "test.bar"() : () -> (f32)
+// CHECK: memref.store %{{.*}}, %[[arg1]]
+  %inserted = tensor.insert %val into %arg1[%pos] : tensor<5xf32>
+  %cond = "test.qux"() : () -> (i1)
+// CHECK: cf.cond_br {{.*}}(%[[arg1]] {{.*}}(%[[arg1]]
+  cf.cond_br %cond, ^bb1(%inserted: tensor<5xf32>), ^bb2(%inserted: tensor<5xf32>)
+// CHECK: ^{{.*}}(%[[arg2:.*]]: memref<5xf32>):
+^bb2(%arg2: tensor<5xf32>):
+// CHECK: return %[[arg2]]
+  func.return %arg2 : tensor<5xf32>
+}

mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir

func.func @inconsistent_memory_space_scf_if(%c: i1) -> tensor<10xf32> {
   } else {
     scf.yield %1 : tensor<10xf32>
   }
   func.return %r : tensor<10xf32>
 }
 
 // -----
 
-func.func @execute_region_multiple_blocks(%t: tensor<5xf32>) -> tensor<5xf32> {
-  // expected-error @below{{op or BufferizableOpInterface implementation does not support unstructured control flow, but at least one region has multiple blocks}}
+func.func @execute_region_multiple_yields(%t: tensor<5xf32>) -> tensor<5xf32> {
+  // expected-error @below{{op op without unique scf.yield is not supported}}
   %0 = scf.execute_region -> tensor<5xf32> {
-    cf.br ^bb1(%t : tensor<5xf32>)
+    scf.yield %t : tensor<5xf32>
   ^bb1(%arg1 : tensor<5xf32>):
     scf.yield %arg1 : tensor<5xf32>
   }
   func.return %0 : tensor<5xf32>
 }
 
 // -----
 

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":IR",
         ":SideEffectInterfaces",
         ":Support",
         "//llvm:Support",
     ],
 )
 
 cc_library(
+    name = "ControlFlowTransforms",
+    srcs = glob([
+        "lib/Dialect/ControlFlow/Transforms/*.cpp",
+    ]),
+    hdrs = glob([
+        "include/mlir/Dialect/ControlFlow/Transforms/*.h",
+    ]),
+    includes = ["include"],
+    deps = [
+        ":BufferizationDialect",
+        ":BufferizationTransforms",
+        ":ControlFlowDialect",
+        ":IR",
+        ":MemRefDialect",
+    ],
+)
+
+cc_library(
     name = "FuncDialect",
     srcs = glob(
         [
             "lib/Dialect/Func/IR/*.cpp",
             "lib/Dialect/Func/IR/*.h",
             "lib/Dialect/Func/Utils/*.cpp",
         ],
     ),
deps = [
         ":BufferizationTransformOps",
         ":BufferizationTransforms",
         ":CastInterfaces",
         ":ComplexDialect",
         ":ComplexToLLVM",
         ":ComplexToLibm",
         ":ComplexToSPIRV",
         ":ControlFlowDialect",
+        ":ControlFlowTransforms",
         ":ConversionPasses",
         ":ConvertToLLVM",
         ":DLTIDialect",
         ":EmitCDialect",
         ":FuncDialect",
         ":FuncToLLVM",
         ":FuncToSPIRV",
         ":FuncTransformOps",