[mlir] Add support for operation-produced successor arguments in BranchOpInterface

[zero9178]

Summary

This patch revamps the BranchOpInterface a bit and allows a proper implementation of what was previously getMutableSuccessorOperands for operations, which internally produce arguments to some of the block arguments. A motivating example for this would be an invoke op with a error handling path:

invoke %function(%0)
  label ^success ^error(%1 : i32)

^error(%e: !error, %arg0 : i32):
  ...

The advantages of this are that any users of BranchOpInterface can still argue over remaining block argument operands (such as %1 in the example above), as well as make use of the modifying capabilities to add more operands, erase an operand etc.

The way this patch implements that functionality is via a new class called SuccessorOperands, which is now returned by getSuccessorOperands. It basically contains an unsigned denoting how many operator produced operands exist, as well as a MutableOperandRange, which are the usual forwarded operands we are used to. The produced operands are assumed to the first few block arguments, followed by the forwarded operands afterwards. The role of SuccessorOperands is to provide various utility functions to modify and query the successor arguments from a BranchOpInterface.

---

Some sort of breaking changes that are worth discussing:

• According to the documentation the old return type of Optional<MutableOperandRange> was simply made to temporarily accommodate above cases. Since these would now be properly handled I did not use Optional<SuccessorOperands>. That is technically a loss of functionality as an Op could previously implement BranchOpInterface without having any successor operands at all. I am not sure whether that is an important use case.
• For now I removed getMutableSuccessorOperands in favour of using getSuccessorOperands. The history as to why both of these exist is very unclear to me and I would be open to separating them again and then probably make a templated version of SuccessorOperands that operates of OperandRange instead, if desired.

[Mogball]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

25

From the description, I wasn't able to easily tell the difference between "produced" arguments and "forwarded" block arguments.

29

Can you include the example you have in the description?

34

Add a constructor SuccessorOperands(MutableOperandRange)?

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

67–68

The example and description you gave in the patch description should be included here.

[rriddle]

Seems like a good approach in general and a nice cleanup. It does, however, place restrictions on how operations implementing BranchOpInterface must handle inner-produced vs. forward operands, and that will need to be explicitly conveyed in the interface documentation.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

25

I would make it explicit the order of produced and fowarded operands, and add IR examples here.

34

+1, not having produced operands is overwhelmingly the common case.

43

Add a doc comment here.

45

Prefer active voice for functions.

61

Why the static cast here?

65

Missing punctuation at the end.

84

Please add punctutation.

mlir/test/Transforms/sccp.mlir

204

Please indent this.

[zero9178]

Address reviewer comments:

• Add constructor taking just a forwarded operands range
• Add more documentation to both the SuccessorOperands class and the BranchOpInterface
• Address comment style issues.

[zero9178]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

61

MutableOperandRange does not have a operator[]. Doing the static_cast calls the conversion operator to OperandRange however, which does.

[rriddle]

Nice work, thanks!

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

61

Can you just add an operator[] to MutableOperandRange?

[zero9178]

Define and make use of operator[] in MutableOperandRange

flang/include/flang/Optimizer/Dialect/FIROps.td

llvm::Optional<llvm::ArrayRef<mlir::Value>> getCompareOperands(
         llvm::ArrayRef<mlir::Value> operands, unsigned cond);
     llvm::Optional<mlir::ValueRange> getCompareOperands(
         mlir::ValueRange operands, unsigned cond);
 
     llvm::Optional<llvm::ArrayRef<mlir::Value>> getSuccessorOperands(
         llvm::ArrayRef<mlir::Value> operands, unsigned cond);
     llvm::Optional<mlir::ValueRange> getSuccessorOperands(
         mlir::ValueRange operands, unsigned cond);
-    using BranchOpInterfaceTrait::getSuccessorOperands;
 
     // Helper function to deal with Optional operand forms
     void printSuccessorAtIndex(mlir::OpAsmPrinter &p, unsigned i) {
       auto *succ = getSuccessor(i);
       auto ops = getSuccessorOperands(i);
-      if (ops.hasValue())
-        p.printSuccessorAndUseList(succ, ops.getValue());
-      else
-        p.printSuccessor(succ);
+      p.printSuccessorAndUseList(succ, ops.getForwardedOperands());
     }
 
     mlir::ArrayAttr getCases() {
       return (*this)->getAttrOfType<mlir::ArrayAttr>(getCasesAttr());
     }
 
     unsigned targetOffsetSize();
   }];

flang/lib/Optimizer/Dialect/FIROps.cpp

llvm::Optional<mlir::OperandRange> fir::SelectOp::getCompareOperands(unsigned) {
   return {};
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectOp::getCompareOperands(llvm::ArrayRef<mlir::Value>, unsigned) {
   return {};
 }
 
-llvm::Optional<mlir::MutableOperandRange>
-fir::SelectOp::getMutableSuccessorOperands(unsigned oper) {
-  return ::getMutableSuccessorOperands(oper, getTargetArgsMutable(),
-                                       getTargetOffsetAttr());
+mlir::SuccessorOperands fir::SelectOp::getSuccessorOperands(unsigned oper) {
+  return mlir::SuccessorOperands(::getMutableSuccessorOperands(
+      oper, getTargetArgsMutable(), getTargetOffsetAttr()));
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                     unsigned oper) {
   auto a =
       (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
fir::SelectCaseOp::getCompareOperands(mlir::ValueRange operands,
                                       unsigned cond) {
   auto a = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
       getCompareOffsetAttr());
   auto segments = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
       getOperandSegmentSizeAttr());
   return {getSubOperands(cond, getSubOperands(1, operands, segments), a)};
 }
 
-llvm::Optional<mlir::MutableOperandRange>
-fir::SelectCaseOp::getMutableSuccessorOperands(unsigned oper) {
-  return ::getMutableSuccessorOperands(oper, getTargetArgsMutable(),
-                                       getTargetOffsetAttr());
+mlir::SuccessorOperands fir::SelectCaseOp::getSuccessorOperands(unsigned oper) {
+  return mlir::SuccessorOperands(::getMutableSuccessorOperands(
+      oper, getTargetArgsMutable(), getTargetOffsetAttr()));
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectCaseOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                         unsigned oper) {
   auto a =
       (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
fir::SelectRankOp::getCompareOperands(unsigned) {
   return {};
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectRankOp::getCompareOperands(llvm::ArrayRef<mlir::Value>, unsigned) {
   return {};
 }
 
-llvm::Optional<mlir::MutableOperandRange>
-fir::SelectRankOp::getMutableSuccessorOperands(unsigned oper) {
-  return ::getMutableSuccessorOperands(oper, getTargetArgsMutable(),
-                                       getTargetOffsetAttr());
+mlir::SuccessorOperands fir::SelectRankOp::getSuccessorOperands(unsigned oper) {
+  return mlir::SuccessorOperands(::getMutableSuccessorOperands(
+      oper, getTargetArgsMutable(), getTargetOffsetAttr()));
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectRankOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                         unsigned oper) {
   auto a =
       (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
fir::SelectTypeOp::getCompareOperands(unsigned) {
   return {};
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectTypeOp::getCompareOperands(llvm::ArrayRef<mlir::Value>, unsigned) {
   return {};
 }
 
-llvm::Optional<mlir::MutableOperandRange>
-fir::SelectTypeOp::getMutableSuccessorOperands(unsigned oper) {
-  return ::getMutableSuccessorOperands(oper, getTargetArgsMutable(),
-                                       getTargetOffsetAttr());
+mlir::SuccessorOperands fir::SelectTypeOp::getSuccessorOperands(unsigned oper) {
+  return mlir::SuccessorOperands(::getMutableSuccessorOperands(
+      oper, getTargetArgsMutable(), getTargetOffsetAttr()));
 }
 
 llvm::Optional<llvm::ArrayRef<mlir::Value>>
 fir::SelectTypeOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                         unsigned oper) {
   auto a =
       (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

 #define MLIR_INTERFACES_CONTROLFLOWINTERFACES_H
 
 #include "mlir/IR/OpDefinition.h"
 
 namespace mlir {
 class BranchOpInterface;
 class RegionBranchOpInterface;
 
+/// This class models how operands are forwarded to block arguments in control
+/// flow. It consists of a number, denoting how many of the successors block
+/// arguments are produced by the operation, followed by a range of operands
+/// that are forwarded. The produced operands are passed to the first few
+/// block arguments of the successor, followed by the forwarded operands.
+/// It is unsupported to pass them in a different order.
+///
+/// An example operation with both of these concepts would be a branch-on-error
+/// operation, that internally produces an error object on the error path:
+///
+///   invoke %function(%0)
+///     label ^success ^error(%1 : i32)
+///
+///     ^error(%e: !error, %arg0 : i32):
+///       ...
+///
+/// This operation would return an instance of SuccessorOperands with a produced
+/// operand count of 1 (mapped to %e in the successor) and a forwarded
+/// operands range consisting of %1 in the example above (mapped to %arg0 in the
+/// successor).
+class SuccessorOperands {
+public:
+  /// Constructs a SuccessorOperands with no produced operands that simply
+  /// forwards operands to the successor.
+  explicit SuccessorOperands(MutableOperandRange forwardedOperands);
+
+  /// Constructs a SuccessorOperands with the given amount of produced operands
+  /// and forwarded operands.
+  SuccessorOperands(unsigned producedOperandCount,
+                    MutableOperandRange forwardedOperands);
+
+  /// Returns the amount of operands passed to the successor. This consists both
+  /// of produced operands by the operation as well as forwarded ones.
+  unsigned size() const {
+    return producedOperandCount + forwardedOperands.size();
+  }
+
+  /// Returns true if there are no successor operands.
+  bool empty() const { return size() == 0; }
+
+  /// Returns the amount of operands that are produced internally by the
+  /// operation. These are passed to the first few block arguments.
+  unsigned getProducedOperandCount() const { return producedOperandCount; }
+
+  /// Returns true if the successor operand denoted by `index` is produced by
+  /// the operation.
+  bool isOperandProduced(unsigned index) const {
+    return index < producedOperandCount;
+  }
+
+  /// Returns the Value that is passed to the successors block argument denoted
+  /// by `index`. If it is produced by the operation, no such value exists and
+  /// a null Value is returned.
+  Value operator[](unsigned index) const {
+    if (isOperandProduced(index))
+      return Value();
+    return static_cast<OperandRange>(
+        forwardedOperands)[index - producedOperandCount];
+  }
+
+  /// Get the range of operands that are simply forwarded to the successor.
+  OperandRange getForwardedOperands() const { return forwardedOperands; }
+
+  /// Get a slice of the operands forwarded to the successor. The given range
+  /// must not contain any operands produced by the operation.
+  MutableOperandRange slice(unsigned subStart, unsigned subLen) const {
+    assert(!isOperandProduced(subStart) &&
+           "can't slice operands produced by the operation");
+    return forwardedOperands.slice(subStart - producedOperandCount, subLen);
+  }
+
+  /// Erase operands forwarded to the successor. The given range must
+  /// not contain any operands produced by the operation.
+  void erase(unsigned subStart, unsigned subLen = 1) {
+    assert(!isOperandProduced(subStart) &&
+           "can't erase operands produced by the operation");
+    forwardedOperands.erase(subStart - producedOperandCount, subLen);
+  }
+
+  /// Add new operands that are forwarded to the successor.
+  void append(ValueRange valueRange) { forwardedOperands.append(valueRange); }
+
+  /// Gets the index of the forwarded operand within the operation which maps
+  /// to the block argument denoted by `blockArgumentIndex`. The block argument
+  /// must be mapped to a forwarded operand.
+  unsigned getOperandIndex(unsigned blockArgumentIndex) const {
+    assert(!isOperandProduced(blockArgumentIndex) &&
+           "can't map operand produced by the operation");
+    return static_cast<mlir::OperandRange>(forwardedOperands)
+               .getBeginOperandIndex() +
+           (blockArgumentIndex - producedOperandCount);
+  }
+
+private:
+  /// Amount of operands that are produced internally within the operation and
+  /// passed to the first few block arguments.
+  unsigned producedOperandCount;
+  /// Range of operands that are forwarded to the remaining block arguments.
+  MutableOperandRange forwardedOperands;
+};
+
 //===----------------------------------------------------------------------===//
 // BranchOpInterface
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// Return the `BlockArgument` corresponding to operand `operandIndex` in some
 /// successor if `operandIndex` is within the range of `operands`, or None if
 /// `operandIndex` isn't a successor operand index.
 Optional<BlockArgument>
-getBranchSuccessorArgument(Optional<OperandRange> operands,
+getBranchSuccessorArgument(const SuccessorOperands &operands,
                            unsigned operandIndex, Block *successor);
 
 /// Verify that the given operands match those of the given successor block.
 LogicalResult verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
-                                            Optional<OperandRange> operands);
+                                            const SuccessorOperands &operands);
 } // namespace detail
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// Verify that types match along control flow edges described the given op.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

let description = [{
     successor block arguments may have different types. In such cases,
     `areTypesCompatible` can be implemented to compare types along control-flow
     edges. By default, type equality is used.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
-        Returns a mutable range of operands that correspond to the arguments of
-        successor at the given index. Returns None if the operands to the
-        successor are non-materialized values, i.e. they are internal to the
-        operation.
+        Returns the operands that correspond to the arguments of the successor
+        at the given index. It consists of a number of operands that are
+        internally produced by the operation, followed by a range of operands
+        that are forwarded. An example operation making use of produced
+        operands would be:
+
+        ```mlir
+        invoke %function(%0)
+            label ^success ^error(%1 : i32)
+
+        ^error(%e: !error, %arg0: i32):
+            ...
+        ```
+
+        The operand that would map to the `^error`s `%e` operand is produced
+        by the `invoke` operation, while `%1` is a forwarded operand that maps
+        to `%arg0` in the successor.
+
+        Produced operands always map to the first few block arguments of the
+        successor, followed by the forwarded operands. Mapping them in any
+        other order is not supported by the interface.
+
+        By having the forwarded operands last allows users of the interface
+        to append more forwarded operands to the branch operation without
+        interfering with other successor operands.
       }],
-      "::mlir::Optional<::mlir::MutableOperandRange>", "getMutableSuccessorOperands",
+      "::mlir::SuccessorOperands", "getSuccessorOperands",
       (ins "unsigned":$index)
     >,
     InterfaceMethod<[{
-        Returns a range of operands that correspond to the arguments of
-        successor at the given index. Returns None if the operands to the
-        successor are non-materialized values, i.e. they are internal to the
-        operation.
-      }],
-      "::mlir::Optional<::mlir::OperandRange>", "getSuccessorOperands",
-      (ins "unsigned":$index), [{}], [{
-        auto operands = $_op.getMutableSuccessorOperands(index);
-        return operands ? ::mlir::Optional<::mlir::OperandRange>(*operands) : ::llvm::None;
-      }]
-    >,
-    InterfaceMethod<[{
         Returns the `BlockArgument` corresponding to operand `operandIndex` in
         some successor, or None if `operandIndex` isn't a successor operand
         index.
       }],
       "::mlir::Optional<::mlir::BlockArgument>", "getSuccessorBlockArgument",
       (ins "unsigned":$operandIndex), [{
         ::mlir::Operation *opaqueOp = $_op;
         for (unsigned i = 0, e = opaqueOp->getNumSuccessors(); i != e; ++i) {
InterfaceMethod<[{
       (ins "::mlir::Type":$lhs, "::mlir::Type":$rhs), [{}],
        [{ return lhs == rhs; }]
     >,
   ];
 
   let verify = [{
     auto concreteOp = ::mlir::cast<ConcreteOp>($_op);
     for (unsigned i = 0, e = $_op->getNumSuccessors(); i != e; ++i) {
-      ::mlir::Optional<OperandRange> operands = concreteOp.getSuccessorOperands(i);
+      ::mlir::SuccessorOperands operands = concreteOp.getSuccessorOperands(i);
       if (::mlir::failed(::mlir::detail::verifyBranchSuccessorOperands($_op, i, operands)))
         return ::mlir::failure();
     }
     return ::mlir::success();
   }];
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp

for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
       if (!branch) {
         // We can't analyze the control flow, so bail out early.
         output.push_back(arg);
         return;
       }
 
       // Try to get the operand passed for this argument.
       unsigned index = it.getSuccessorIndex();
-      Optional<OperandRange> operands = branch.getSuccessorOperands(index);
-      if (!operands) {
+      Value operand = branch.getSuccessorOperands(index)[argNumber];
+      if (!operand) {
         // We can't analyze the control flow, so bail out early.
         output.push_back(arg);
         return;
       }
-      collectUnderlyingAddressValues((*operands)[argNumber], maxDepth, visited,
-                                     output);
+      collectUnderlyingAddressValues(operand, maxDepth, visited, output);
     }
     return;
   }
 
   // Otherwise, check to see if we can reason about the control flow of this op.
   Region *region = block->getParent();
   Operation *op = region->getParentOp();
   if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) {

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

void BufferViewFlowAnalysis::build(Operation *op) {
   // Query all branch interfaces to link block argument dependencies.
   op->walk([&](BranchOpInterface branchInterface) {
     Block *parentBlock = branchInterface->getBlock();
     for (auto it = parentBlock->succ_begin(), e = parentBlock->succ_end();
          it != e; ++it) {
       // Query the branch op interface to get the successor operands.
       auto successorOperands =
           branchInterface.getSuccessorOperands(it.getIndex());
-      if (!successorOperands.hasValue())
-        continue;
       // Build the actual mapping of values to their immediate dependencies.
-      registerDependencies(successorOperands.getValue(), (*it)->getArguments());
+      registerDependencies(successorOperands.getForwardedOperands(),
+                           (*it)->getArguments().drop_front(
+                               successorOperands.getProducedOperandCount()));
     }
   });
 
   // Query the RegionBranchOpInterface to find potential successor regions.
   op->walk([&](RegionBranchOpInterface regionInterface) {
     // Extract all entry regions and wire all initial entry successor inputs.
     SmallVector<RegionSuccessor, 2> entrySuccessors;
     regionInterface.getSuccessorRegions(/*index=*/llvm::None, entrySuccessors);

mlir/lib/Analysis/DataFlowAnalysis.cpp

for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
 
     // We only care about this predecessor if it is going to execute.
     if (!isEdgeExecutable(pred, block))
       continue;
 
     // Try to get the operand forwarded by the predecessor. If we can't reason
     // about the terminator of the predecessor, mark as having reached a
     // fixpoint.
-    Optional<OperandRange> branchOperands;
-    if (auto branch = dyn_cast<BranchOpInterface>(pred->getTerminator()))
-      branchOperands = branch.getSuccessorOperands(it.getSuccessorIndex());
-    if (!branchOperands) {
+    auto branch = dyn_cast<BranchOpInterface>(pred->getTerminator());
+    if (!branch) {
+      updatedLattice |= argLattice.markPessimisticFixpoint();
+      break;
+    }
+    Value operand = branch.getSuccessorOperands(it.getSuccessorIndex())[i];
+    if (!operand) {
       updatedLattice |= argLattice.markPessimisticFixpoint();
       break;
     }
 
     // If the operand hasn't been resolved, it is uninitialized and can merge
     // with anything.
     AbstractLatticeElement *operandLattice =
-        analysis.lookupLatticeElement((*branchOperands)[i]);
+        analysis.lookupLatticeElement(operand);
     if (!operandLattice)
       continue;
 
     // Otherwise, join the operand lattice into the argument lattice.
     updatedLattice |= argLattice.join(*operandLattice);
     if (argLattice.isAtFixpoint())
       break;
   }

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

LogicalResult introduceBlockArgCopy(BlockArgument blockArg) {
     // the associated value (which will be a predecessor block by
     // definition).
     Block *block = blockArg.getOwner();
     for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
       // Get the terminator and the value that will be passed to our
       // argument.
       Operation *terminator = (*it)->getTerminator();
       auto branchInterface = cast<BranchOpInterface>(terminator);
+      SuccessorOperands operands =
+          branchInterface.getSuccessorOperands(it.getSuccessorIndex());
+
       // Query the associated source value.
-      Value sourceValue =
-          branchInterface.getSuccessorOperands(it.getSuccessorIndex())
-              .getValue()[blockArg.getArgNumber()];
-      // Wire new clone and successor operand.
-      auto mutableOperands =
-          branchInterface.getMutableSuccessorOperands(it.getSuccessorIndex());
-      if (!mutableOperands) {
-        terminator->emitError() << "terminators with immutable successor "
-                                   "operands are not supported";
-        continue;
+      Value sourceValue = operands[blockArg.getArgNumber()];
+      if (!sourceValue) {
+        return failure();
       }
+      // Wire new clone and successor operand.
       // Create a new clone at the current location of the terminator.
       auto clone = introduceCloneBuffers(sourceValue, terminator);
       if (failed(clone))
         return failure();
-      mutableOperands.getValue()
-          .slice(blockArg.getArgNumber(), 1)
-          .assign(*clone);
+      operands.slice(blockArg.getArgNumber(), 1).assign(*clone);
     }
 
     // Check whether the block argument has implicitly defined predecessors via
     // the RegionBranchOpInterface. This can be the case if the current block
     // argument belongs to the first block in a region and the parent operation
     // implements the RegionBranchOpInterface.
     Region *argRegion = block->getParent();
     Operation *parentOp = argRegion->getParentOp();

mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp

LogicalResult BranchOp::canonicalize(BranchOp op, PatternRewriter &rewriter) {
   return success(succeeded(simplifyBrToBlockWithSinglePred(op, rewriter)) ||
                  succeeded(simplifyPassThroughBr(op, rewriter)));
 }
 
 void BranchOp::setDest(Block *block) { return setSuccessor(block); }
 
 void BranchOp::eraseOperand(unsigned index) { (*this)->eraseOperand(index); }
 
-Optional<MutableOperandRange>
-BranchOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands BranchOp::getSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
-  return getDestOperandsMutable();
+  return SuccessorOperands(getDestOperandsMutable());
 }
 
 Block *BranchOp::getSuccessorForOperands(ArrayRef<Attribute>) {
   return getDest();
 }
 
 //===----------------------------------------------------------------------===//
 // CondBranchOp

 void CondBranchOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
   results.add<SimplifyConstCondBranchPred, SimplifyPassThroughCondBranch,
               SimplifyCondBranchIdenticalSuccessors,
               SimplifyCondBranchFromCondBranchOnSameCondition,
               CondBranchTruthPropagation>(context);
 }
 
-Optional<MutableOperandRange>
-CondBranchOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands CondBranchOp::getSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
-  return index == trueIndex ? getTrueDestOperandsMutable()
-                            : getFalseDestOperandsMutable();
+  return SuccessorOperands(index == trueIndex ? getTrueDestOperandsMutable()
+                                              : getFalseDestOperandsMutable());
 }
 
 Block *CondBranchOp::getSuccessorForOperands(ArrayRef<Attribute> operands) {
   if (IntegerAttr condAttr = operands.front().dyn_cast_or_null<IntegerAttr>())
     return condAttr.getValue().isOneValue() ? getTrueDest() : getFalseDest();
   return nullptr;
 }
 
if (caseValues &&
       caseValues->size() != static_cast<int64_t>(caseDestinations.size()))
     return emitOpError() << "number of case values (" << caseValues->size()
                          << ") should match number of "
                             "case destinations ("
                          << caseDestinations.size() << ")";
   return success();
 }
 
-Optional<MutableOperandRange>
-SwitchOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands SwitchOp::getSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
-  return index == 0 ? getDefaultOperandsMutable()
-                    : getCaseOperandsMutable(index - 1);
+  return SuccessorOperands(index == 0 ? getDefaultOperandsMutable()
+                                      : getCaseOperandsMutable(index - 1));
 }
 
 Block *SwitchOp::getSuccessorForOperands(ArrayRef<Attribute> operands) {
   Optional<DenseIntElementsAttr> caseValues = getCaseValues();
 
   if (!caseValues)
     return getDefaultDestination();
 

mlir/lib/Dialect/Func/Transforms/FuncConversions.cpp

public:
   LogicalResult
   matchAndRewrite(BranchOpInterface op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     // For a branch operation, only some operands go to the target blocks, so
     // only rewrite those.
     SmallVector<Value, 4> newOperands(op->operand_begin(), op->operand_end());
     for (int succIdx = 0, succEnd = op->getBlock()->getNumSuccessors();
          succIdx < succEnd; ++succIdx) {
-      auto successorOperands = op.getSuccessorOperands(succIdx);
-      if (!successorOperands || successorOperands->empty())
+      OperandRange forwardedOperands =
+          op.getSuccessorOperands(succIdx).getForwardedOperands();
+      if (forwardedOperands.empty())
         continue;
 
-      for (int idx = successorOperands->getBeginOperandIndex(),
-               eidx = idx + successorOperands->size();
+      for (int idx = forwardedOperands.getBeginOperandIndex(),
+               eidx = idx + forwardedOperands.size();
            idx < eidx; ++idx) {
         if (!shouldConvertBranchOperand || shouldConvertBranchOperand(op, idx))
           newOperands[idx] = operands[idx];
       }
     }
     rewriter.updateRootInPlace(
         op, [newOperands, op]() { op->setOperands(newOperands); });
     return success();

 }
 
 bool mlir::isLegalForBranchOpInterfaceTypeConversionPattern(
     Operation *op, TypeConverter &converter) {
   // All successor operands of branch like operations must be rewritten.
   if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
     for (int p = 0, e = op->getBlock()->getNumSuccessors(); p < e; ++p) {
       auto successorOperands = branchOp.getSuccessorOperands(p);
-      if (successorOperands.hasValue() &&
-          !converter.isLegal(successorOperands.getValue().getTypes()))
+      if (!converter.isLegal(
+              successorOperands.getForwardedOperands().getTypes()))
         return false;
     }
     return true;
   }
 
   return false;
 }
 

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

ParseResult AllocaOp::parse(OpAsmParser &parser, OperationState &result) {
   result.addTypes({funcType.getResult(0)});
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // LLVM::BrOp
 //===----------------------------------------------------------------------===//
 
-Optional<MutableOperandRange>
-BrOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands BrOp::getSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
-  return getDestOperandsMutable();
+  return SuccessorOperands(getDestOperandsMutable());
 }
 
 //===----------------------------------------------------------------------===//
 // LLVM::CondBrOp
 //===----------------------------------------------------------------------===//
 
-Optional<MutableOperandRange>
-CondBrOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands CondBrOp::getSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
-  return index == 0 ? getTrueDestOperandsMutable()
-                    : getFalseDestOperandsMutable();
+  return SuccessorOperands(index == 0 ? getTrueDestOperandsMutable()
+                                      : getFalseDestOperandsMutable());
 }
 
 //===----------------------------------------------------------------------===//
 // LLVM::SwitchOp
 //===----------------------------------------------------------------------===//
 
 void SwitchOp::build(OpBuilder &builder, OperationState &result, Value value,
                      Block *defaultDestination, ValueRange defaultOperands,
return emitOpError("expects number of case values to match number of "
                        "case destinations");
   if (getBranchWeights() && getBranchWeights()->size() != getNumSuccessors())
     return emitError("expects number of branch weights to match number of "
                      "successors: ")
            << getBranchWeights()->size() << " vs " << getNumSuccessors();
   return success();
 }
 
-Optional<MutableOperandRange>
-SwitchOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands SwitchOp::getSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
-  return index == 0 ? getDefaultOperandsMutable()
-                    : getCaseOperandsMutable(index - 1);
+  return SuccessorOperands(index == 0 ? getDefaultOperandsMutable()
+                                      : getCaseOperandsMutable(index - 1));
 }
 
 //===----------------------------------------------------------------------===//
 // Code for LLVM::GEPOp.
 //===----------------------------------------------------------------------===//
 
 constexpr int GEPOp::kDynamicIndex;
 
ParseResult StoreOp::parse(OpAsmParser &parser, OperationState &result) {
 
   return success();
 }
 
 ///===---------------------------------------------------------------------===//
 /// LLVM::InvokeOp
 ///===---------------------------------------------------------------------===//
 
-Optional<MutableOperandRange>
-InvokeOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands InvokeOp::getSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
-  return index == 0 ? getNormalDestOperandsMutable()
-                    : getUnwindDestOperandsMutable();
+  return SuccessorOperands(index == 0 ? getNormalDestOperandsMutable()
+                                      : getUnwindDestOperandsMutable());
 }
 
 LogicalResult InvokeOp::verify() {
   if (getNumResults() > 1)
     return emitOpError("must have 0 or 1 result");
 
   Block *unwindDest = getUnwindDest();
   if (unwindDest->empty())

mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp

static DenseMap<Operation *, DenseSet<int>> computeBranchOpDetensoring(
 
         for (PredecessorIterator pred = block->pred_begin();
              pred != block->pred_end(); ++pred) {
           BranchOpInterface terminator =
               dyn_cast<BranchOpInterface>((*pred)->getTerminator());
           auto blockOperands =
               terminator.getSuccessorOperands(pred.getSuccessorIndex());
 
-          if (!blockOperands || blockOperands->empty())
+          if (blockOperands.empty() ||
+              blockOperands.isOperandProduced(blockArgumentElem.getArgNumber()))
             continue;
 
           detensorableBranchOps[terminator].insert(
-              blockOperands->getBeginOperandIndex() +
-              blockArgumentElem.getArgNumber());
+              blockOperands.getOperandIndex(blockArgumentElem.getArgNumber()));
         }
       }
 
       return detensorableBranchOps;
     }
   };
 
   /// Detensorize linalg ops involved in control-flow within a function.
void compute(FunctionOpInterface func,
               opsToDetensor.clear();
               blockArgsToDetensor.clear();
               return;
             }
 
             auto ownerBlockOperands =
                 predTerminator.getSuccessorOperands(pred.getSuccessorIndex());
 
-            if (!ownerBlockOperands || ownerBlockOperands->empty())
+            if (ownerBlockOperands.empty() ||
+                ownerBlockOperands.isOperandProduced(
+                    currentItemBlockArgument.getArgNumber()))
               continue;
 
             // For each predecessor, add the value it passes to that argument to
             // workList to find out how it's computed.
             workList.push_back(
-                ownerBlockOperands
-                    .getValue()[currentItemBlockArgument.getArgNumber()]);
+                ownerBlockOperands[currentItemBlockArgument.getArgNumber()]);
           }
 
           continue;
         }
 
         Operation *currentItemDefiningOp = currentItem.getDefiningOp();
 
         if (!visitedOps.insert(currentItemDefiningOp).second)
void compute(FunctionOpInterface func,
         // correpsonding operands in predecessor blocks.
         for (PredecessorIterator pred = block->pred_begin();
              pred != block->pred_end(); ++pred) {
           BranchOpInterface terminator =
               dyn_cast<BranchOpInterface>((*pred)->getTerminator());
           auto blockOperands =
               terminator.getSuccessorOperands(pred.getSuccessorIndex());
 
-          if (!blockOperands || blockOperands->empty())
+          if (blockOperands.empty() ||
+              blockOperands.isOperandProduced(blockArg.getArgNumber()))
             continue;
 
           Operation *definingOp =
-              terminator
-                  ->getOperand(blockOperands->getBeginOperandIndex() +
-                               blockArg.getArgNumber())
-                  .getDefiningOp();
+              blockOperands[blockArg.getArgNumber()].getDefiningOp();
 
           // If the operand is defined by a GenericOp that will not be
           // detensored, then do not detensor the corresponding block argument.
-          if (dyn_cast_or_null<GenericOp>(definingOp) &&
+          if (isa_and_nonnull<GenericOp>(definingOp) &&
               opsToDetensor.count(definingOp) == 0) {
             blockArgsToRemove.insert(blockArg);
             break;
           }
         }
       }
 
       for (auto &blockArg : blockArgsToRemove) {

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

LogicalResult spirv::BitcastOp::verify() {
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // spv.BranchOp
 //===----------------------------------------------------------------------===//
 
-Optional<MutableOperandRange>
-spirv::BranchOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands spirv::BranchOp::getSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
-  return targetOperandsMutable();
+  return SuccessorOperands(0, targetOperandsMutable());
 }
 
 //===----------------------------------------------------------------------===//
 // spv.BranchConditionalOp
 //===----------------------------------------------------------------------===//
 
-Optional<MutableOperandRange>
-spirv::BranchConditionalOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands
+spirv::BranchConditionalOp::getSuccessorOperands(unsigned index) {
   assert(index < 2 && "invalid successor index");
-  return index == kTrueIndex ? trueTargetOperandsMutable()
-                             : falseTargetOperandsMutable();
+  return SuccessorOperands(index == kTrueIndex ? trueTargetOperandsMutable()
+                                               : falseTargetOperandsMutable());
 }
 
 ParseResult spirv::BranchConditionalOp::parse(OpAsmParser &parser,
                                               OperationState &state) {
   auto &builder = parser.getBuilder();
   OpAsmParser::UnresolvedOperand condInfo;
   Block *dest;
 

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ControlFlowInterfaces
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/ControlFlowInterfaces.cpp.inc"
 
+SuccessorOperands::SuccessorOperands(MutableOperandRange forwardedOperands)
+    : producedOperandCount(0), forwardedOperands(forwardedOperands) {}
+
+SuccessorOperands::SuccessorOperands(unsigned int producedOperandCount,
+                                     MutableOperandRange forwardedOperands)
+    : producedOperandCount(producedOperandCount),
+      forwardedOperands(std::move(forwardedOperands)) {}
+
 //===----------------------------------------------------------------------===//
 // BranchOpInterface
 //===----------------------------------------------------------------------===//
 
 /// Returns the `BlockArgument` corresponding to operand `operandIndex` in some
 /// successor if 'operandIndex' is within the range of 'operands', or None if
 /// `operandIndex` isn't a successor operand index.
 Optional<BlockArgument>
-detail::getBranchSuccessorArgument(Optional<OperandRange> operands,
+detail::getBranchSuccessorArgument(const SuccessorOperands &operands,
                                    unsigned operandIndex, Block *successor) {
+  OperandRange forwardedOperands = operands.getForwardedOperands();
   // Check that the operands are valid.
-  if (!operands || operands->empty())
+  if (forwardedOperands.empty())
     return llvm::None;
 
   // Check to ensure that this operand is within the range.
-  unsigned operandsStart = operands->getBeginOperandIndex();
+  unsigned operandsStart = forwardedOperands.getBeginOperandIndex();
   if (operandIndex < operandsStart ||
-      operandIndex >= (operandsStart + operands->size()))
+      operandIndex >= (operandsStart + forwardedOperands.size()))
     return llvm::None;
 
   // Index the successor.
-  unsigned argIndex = operandIndex - operandsStart;
+  unsigned argIndex =
+      operands.getProducedOperandCount() + operandIndex - operandsStart;
   return successor->getArgument(argIndex);
 }
 
 /// Verify that the given operands match those of the given successor block.
 LogicalResult
 detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
-                                      Optional<OperandRange> operands) {
-  if (!operands)
-    return success();
-
+                                      const SuccessorOperands &operands) {
   // Check the count.
-  unsigned operandCount = operands->size();
+  unsigned operandCount = operands.size();
   Block *destBB = op->getSuccessor(succNo);
   if (operandCount != destBB->getNumArguments())
     return op->emitError() << "branch has " << operandCount
                            << " operands for successor #" << succNo
                            << ", but target block has "
                            << destBB->getNumArguments();
 
   // Check the types.
-  auto operandIt = operands->begin();
-  for (unsigned i = 0; i != operandCount; ++i, ++operandIt) {
+  for (unsigned i = operands.getProducedOperandCount(); i != operandCount;
+       ++i) {
     if (!cast<BranchOpInterface>(op).areTypesCompatible(
-            (*operandIt).getType(), destBB->getArgument(i).getType()))
+            operands[i].getType(), destBB->getArgument(i).getType()))
       return op->emitError() << "type mismatch for bb argument #" << i
                              << " of successor #" << succNo;
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

static Value getPHISourceValue(Block *current, Block *pred,
   if (isa<LLVM::BrOp>(terminator))
     return terminator.getOperand(index);
 
 #ifndef NDEBUG
   llvm::SmallPtrSet<Block *, 4> seenSuccessors;
   for (unsigned i = 0, e = terminator.getNumSuccessors(); i < e; ++i) {
     Block *successor = terminator.getSuccessor(i);
     auto branch = cast<BranchOpInterface>(terminator);
-    Optional<OperandRange> successorOperands = branch.getSuccessorOperands(i);
+    SuccessorOperands successorOperands = branch.getSuccessorOperands(i);
     assert(
-        (!seenSuccessors.contains(successor) ||
-         (successorOperands && successorOperands->empty())) &&
+        (!seenSuccessors.contains(successor) || successorOperands.empty()) &&
         "successors with arguments in LLVM branches must be different blocks");
     seenSuccessors.insert(successor);
   }
 #endif
 
   // For instructions that branch based on a condition value, we need to take
   // the operands for the branch that was taken.
   if (auto condBranchOp = dyn_cast<LLVM::CondBrOp>(terminator)) {

mlir/lib/Transforms/Utils/RegionUtils.cpp

static void propagateTerminatorLiveness(Operation *op, LiveMap &liveMap) {
   BranchOpInterface branchInterface = dyn_cast<BranchOpInterface>(op);
   if (!branchInterface) {
     for (Block *successor : op->getSuccessors())
       for (BlockArgument arg : successor->getArguments())
         liveMap.setProvedLive(arg);
     return;
   }
 
-  // If we can't reason about the operands to a successor, conservatively mark
-  // all arguments as live.
+  // If we can't reason about the operand to a successor, conservatively mark
+  // it as live.
   for (unsigned i = 0, e = op->getNumSuccessors(); i != e; ++i) {
-    if (!branchInterface.getMutableSuccessorOperands(i))
-      for (BlockArgument arg : op->getSuccessor(i)->getArguments())
-        liveMap.setProvedLive(arg);
+    SuccessorOperands successorOperands =
+        branchInterface.getSuccessorOperands(i);
+    for (unsigned opI = 0, opE = successorOperands.getProducedOperandCount();
+         opI != opE; ++opI)
+      liveMap.setProvedLive(op->getSuccessor(i)->getArgument(opI));
   }
 }
 
 static void propagateLiveness(Operation *op, LiveMap &liveMap) {
   // Recurse on any regions the op has.
   for (Region &region : op->getRegions())
     propagateLiveness(region, liveMap);
 
static void eraseTerminatorSuccessorOperands(Operation *terminator,
 
   for (unsigned succI = 0, succE = terminator->getNumSuccessors();
        succI < succE; succI++) {
     // Iterating successors in reverse is not strictly needed, since we
     // aren't erasing any successors. But it is slightly more efficient
     // since it will promote later operands of the terminator being erased
     // first, reducing the quadratic-ness.
     unsigned succ = succE - succI - 1;
-    Optional<MutableOperandRange> succOperands =
-        branchOp.getMutableSuccessorOperands(succ);
-    if (!succOperands)
-      continue;
+    SuccessorOperands succOperands = branchOp.getSuccessorOperands(succ);
     Block *successor = terminator->getSuccessor(succ);
 
-    for (unsigned argI = 0, argE = succOperands->size(); argI < argE; ++argI) {
+    for (unsigned argI = 0, argE = succOperands.size(); argI < argE; ++argI) {
       // Iterating args in reverse is needed for correctness, to avoid
       // shifting later args when earlier args are erased.
       unsigned arg = argE - argI - 1;
       if (!liveMap.wasProvenLive(successor->getArgument(arg)))
-        succOperands->erase(arg);
+        succOperands.erase(arg);
     }
   }
 }
 
 static LogicalResult deleteDeadness(RewriterBase &rewriter,
                                     MutableArrayRef<Region> regions,
                                     LiveMap &liveMap) {
   bool erasedAnything = false;
LogicalResult BlockMergeCluster::addToCluster(BlockEquivalenceData &blockData) {
   blocksToMerge.insert(blockData.block);
   return success();
 }
 
 /// Returns true if the predecessor terminators of the given block can not have
 /// their operands updated.
 static bool ableToUpdatePredOperands(Block *block) {
   for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
-    auto branch = dyn_cast<BranchOpInterface>((*it)->getTerminator());
-    if (!branch || !branch.getMutableSuccessorOperands(it.getSuccessorIndex()))
+    if (!isa<BranchOpInterface>((*it)->getTerminator()))
       return false;
   }
   return true;
 }
 
 LogicalResult BlockMergeCluster::merge(RewriterBase &rewriter) {
   // Don't consider clusters that don't have blocks to merge.
   if (blocksToMerge.empty())
for (const auto &it : llvm::enumerate(operandsToMerge)) {
       }
     }
     // Update the predecessors for each of the blocks.
     auto updatePredecessors = [&](Block *block, unsigned clusterIndex) {
       for (auto predIt = block->pred_begin(), predE = block->pred_end();
            predIt != predE; ++predIt) {
         auto branch = cast<BranchOpInterface>((*predIt)->getTerminator());
         unsigned succIndex = predIt.getSuccessorIndex();
-        branch.getMutableSuccessorOperands(succIndex)->append(
+        branch.getSuccessorOperands(succIndex).append(
             newArguments[clusterIndex]);
       }
     };
     updatePredecessors(leaderBlock, /*clusterIndex=*/0);
     for (unsigned i = 0, e = blocksToMerge.size(); i != e; ++i)
       updatePredecessors(blocksToMerge[i], /*clusterIndex=*/i + 1);
   }
 

mlir/test/Transforms/sccp.mlir

^bb2(%x: i1):
   // CHECK: ^bb2(%[[X:.*]]: i1):
   cf.br ^bb3(%x : i1)
 
 ^bb3(%y: i1):
   // CHECK: ^bb3(%[[Y:.*]]: i1):
   // CHECK: return %[[X]], %[[Y]]
   return %x, %y : i1, i1
 }
+
+// CHECK-LABEL: func @simple_produced_operand
+func @simple_produced_operand() -> (i32, i32) {
+  // CHECK: %[[ONE:.*]] = arith.constant 1
+  %1 = arith.constant 1 : i32
+  "test.internal_br"(%1) [^bb1, ^bb2] {
+    operand_segment_sizes = dense<[0, 1]> : vector<2 x i32>
+  } : (i32) -> ()
+
+^bb1:
+  cf.br ^bb2(%1, %1 : i32, i32)
+
+^bb2(%arg1 : i32, %arg2 : i32):
+  // CHECK: ^bb2(%[[ARG:.*]]: i32, %{{.*}}: i32):
+  // CHECK: return %[[ARG]], %[[ONE]] : i32, i32
+
+  return %arg1, %arg2 : i32, i32
+}

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

TestDialect::getOperationPrinter(Operation *op) const {
   }
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // TestBranchOp
 //===----------------------------------------------------------------------===//
 
-Optional<MutableOperandRange>
-TestBranchOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands TestBranchOp::getSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
-  return getTargetOperandsMutable();
+  return SuccessorOperands(getTargetOperandsMutable());
 }
 
 //===----------------------------------------------------------------------===//
 // TestProducingBranchOp
 //===----------------------------------------------------------------------===//
 
-Optional<MutableOperandRange>
-TestProducingBranchOp::getMutableSuccessorOperands(unsigned index) {
+SuccessorOperands TestProducingBranchOp::getSuccessorOperands(unsigned index) {
   assert(index <= 1 && "invalid successor index");
   if (index == 1)
-    return getFirstOperandsMutable();
-  return getSecondOperandsMutable();
+    return SuccessorOperands(getFirstOperandsMutable());
+  return SuccessorOperands(getSecondOperandsMutable());
+}
+
+//===----------------------------------------------------------------------===//
+// TestProducingBranchOp
+//===----------------------------------------------------------------------===//
+
+SuccessorOperands TestInternalBranchOp::getSuccessorOperands(unsigned index) {
+  assert(index <= 1 && "invalid successor index");
+  if (index == 0)
+    return SuccessorOperands(0, getSuccessOperandsMutable());
+  return SuccessorOperands(1, getErrorOperandsMutable());
 }
 
 //===----------------------------------------------------------------------===//
 // TestDialectCanonicalizerOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult
 dialectCanonicalizationPattern(TestDialectCanonicalizerOp op,

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

def TestProducingBranchOp : TEST_Op<"producing_br",
     [DeclareOpInterfaceMethods<BranchOpInterface>, Terminator,
      AttrSizedOperandSegments]> {
   let arguments = (ins Variadic<AnyType>:$firstOperands,
                        Variadic<AnyType>:$secondOperands);
   let results = (outs I32:$dummy);
   let successors = (successor AnySuccessor:$first,AnySuccessor:$second);
 }
 
+// Produces an error value on the error path
+def TestInternalBranchOp : TEST_Op<"internal_br",
+  [DeclareOpInterfaceMethods<BranchOpInterface>, Terminator,
+   AttrSizedOperandSegments]> {
+
+  let arguments = (ins Variadic<AnyType>:$successOperands,
+                       Variadic<AnyType>:$errorOperands);
+
+  let successors = (successor AnySuccessor:$successPath, AnySuccessor:$errorPath);
+}
+
 def AttrSizedOperandOp : TEST_Op<"attr_sized_operands",
                                  [AttrSizedOperandSegments]> {
   let arguments = (ins
     Variadic<I32>:$a,
     Variadic<I32>:$b,
     I32:$c,
     Variadic<I32>:$d,
     I32ElementsAttr:$operand_segment_sizes