[mlir] Add a new BranchOpInterface to allow for opaquely interfacing with branching terminator operations.

[rriddle]

Summary

This interface contains the necessary components to provide the same builtin behavior that terminators have. This will be used in future revisions to remove many of the hardcoded constraints placed on successors and successor operands. The interface initially contains three methods:

c++
// Return a set of values corresponding to the operands for successor 'index', or None if the operands do not correspond to materialized values.
Optional<OperandRange> getSuccessorOperands(unsigned index);

// Return true if this terminator can have it's successor operands erased.
bool canEraseSuccessorOperand();

// Erase the operand of a successor. This is only valid to call if 'canEraseSuccessorOperand' returns true.
void eraseSuccessorOperand(unsigned succIdx, unsigned opIdx);

Depends On D75313

[mravishankar]

Nice. This cleans up successor access. Thanks River!

[antiagainst]

Cool! I just have a few nits.

mlir/include/mlir/Analysis/ControlFlowInterfaces.h

25

Use back tick for quotes to be consistent with others?

26

Minor issue: I see we have a few range concepts: the whole range for all the operands and the partial range for a specific operand. We use index to mean indexing into both. What about using "offset" for the latter to differentiate?

mlir/include/mlir/Analysis/ControlFlowInterfaces.td

11

remove "what".

mlir/include/mlir/Dialect/SPIRV/SPIRVControlFlowOps.td

24

Nit: can you sort traits alphabetically?

mlir/lib/IR/OperationSupport.cpp

188

Nit: Number typically makes me think about count or something. What about Index? But I see we also have getOperandNumber(). So meh. Up to you. :)

[rriddle]

Address comments

[rriddle]

Thanks for the review Lei!

[rriddle]

Rebase

[stephenneuendorffer]

This commit introduces a circular library dependence with the dialects.
Can this be part of the Standard dialect instead, or maybe lib/IR if it is fundamental?

[rriddle]

In D75314#1908993, @stephenneuendorffer wrote:

This commit introduces a circular library dependence with the dialects.

Yes thanks, all of the interfaces in analysis are going to be moved out of Analysis in a followup.

Can this be part of the Standard dialect instead,

It doesn't really make sense to put an interface definition inside of a specific dialect.

or maybe lib/IR if it is fundamental?

IR/ should be kept clean, otherwise it will become a dumping ground for interfaces(like it is for traits ATM unfortunately).

mlir/include/mlir/Analysis/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS CallInterfaces.td)
 mlir_tablegen(CallInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(CallInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRCallOpInterfacesIncGen)
 
+set(LLVM_TARGET_DEFINITIONS ControlFlowInterfaces.td)
+mlir_tablegen(ControlFlowInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(ControlFlowInterfaces.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRControlFlowInterfacesIncGen)
+
 set(LLVM_TARGET_DEFINITIONS InferTypeOpInterface.td)
 mlir_tablegen(InferTypeOpInterface.h.inc -gen-op-interface-decls)
 mlir_tablegen(InferTypeOpInterface.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRTypeInferOpInterfaceIncGen)

mlir/include/mlir/Analysis/ControlFlowInterfaces.h

+//===- ControlFlowInterfaces.h - ControlFlow Interfaces ---------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the definitions of the branch interfaces defined in
+// `ControlFlowInterfaces.td`.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_CONTROLFLOWINTERFACES_H
+#define MLIR_ANALYSIS_CONTROLFLOWINTERFACES_H
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+class BranchOpInterface;
+
+namespace detail {
+/// Erase an operand from a branch operation that is used as a successor
+/// operand. `operandIndex` is the operand within `operands` to be erased.
+void eraseBranchSuccessorOperand(OperandRange operands, unsigned operandIndex,
+                                 Operation *op);
+
+/// 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,
+                           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);
+} // namespace detail
+
+#include "mlir/Analysis/ControlFlowInterfaces.h.inc"
+} // end namespace mlir
+
+#endif // MLIR_ANALYSIS_CONTROLFLOWINTERFACES_H

mlir/include/mlir/Analysis/ControlFlowInterfaces.td

+//===-- ControlFlowInterfaces.td - ControlFlow Interfaces --*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a set of interfaces that can be used to define information
+// about control flow operations, e.g. branches.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_CONTROLFLOWINTERFACES
+#define MLIR_ANALYSIS_CONTROLFLOWINTERFACES
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// BranchOpInterface
+//===----------------------------------------------------------------------===//
+
+def BranchOpInterface : OpInterface<"BranchOpInterface"> {
+  let description = [{
+    This interface provides information for branching terminator operations,
+    i.e. terminator operations with successors.
+  }];
+  let methods = [
+    InterfaceMethod<[{
+        Returns a set of values that correspond to the arguments to the
+        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.
+      }],
+      "Optional<OperandRange>", "getSuccessorOperands", (ins "unsigned":$index)
+    >,
+    InterfaceMethod<[{
+        Return true if this operation can erase an operand to a successor block.
+      }],
+      "bool", "canEraseSuccessorOperand"
+    >,
+    InterfaceMethod<[{
+        Erase the operand at `operandIndex` from the `index`-th successor. This
+        should only be called if `canEraseSuccessorOperand` returns true.
+      }],
+      "void", "eraseSuccessorOperand",
+      (ins "unsigned":$index, "unsigned":$operandIndex), [{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp *op = static_cast<ConcreteOp *>(this);
+        Optional<OperandRange> operands = op->getSuccessorOperands(index);
+        assert(operands && "unable to query operands for successor");
+        detail::eraseBranchSuccessorOperand(*operands, operandIndex, *op);
+      }]
+    >,
+    InterfaceMethod<[{
+        Returns the `BlockArgument` corresponding to operand `operandIndex` in
+        some successor, or None if `operandIndex` isn't a successor operand
+        index.
+      }],
+      "Optional<BlockArgument>", "getSuccessorBlockArgument",
+      (ins "unsigned":$operandIndex), [{
+        Operation *opaqueOp = op;
+        for (unsigned i = 0, e = opaqueOp->getNumSuccessors(); i != e; ++i) {
+          if (Optional<BlockArgument> arg = detail::getBranchSuccessorArgument(
+                op.getSuccessorOperands(i), operandIndex,
+                opaqueOp->getSuccessor(i)))
+            return arg;
+        }
+        return llvm::None;
+      }]
+    >
+  ];
+
+  let verify = [{
+    auto concreteOp = cast<ConcreteOpType>($_op);
+    for (unsigned i = 0, e = $_op->getNumSuccessors(); i != e; ++i) {
+      Optional<OperandRange> operands = concreteOp.getSuccessorOperands(i);
+      if (failed(detail::verifyBranchSuccessorOperands($_op, i, operands)))
+        return failure();
+    }
+    return success();
+  }];
+}
+
+#endif // MLIR_ANALYSIS_CONTROLFLOWINTERFACES

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

 //===- LLVMDialect.h - MLIR LLVM IR dialect ---------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the LLVM IR dialect in MLIR, containing LLVM operations and
 // LLVM type system.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_LLVMDIALECT_H_
 #define MLIR_DIALECT_LLVMIR_LLVMDIALECT_H_
 
+#include "mlir/Analysis/ControlFlowInterfaces.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeSupport.h"
 #include "mlir/IR/Types.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/LLVMContext.h"

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 //===-- LLVMOps.td - LLVM IR dialect op definition file ----*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This is the LLVM IR operation definition file.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVMIR_OPS
 #define LLVMIR_OPS
 
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+include "mlir/Analysis/ControlFlowInterfaces.td"
 
 class LLVM_Builder<string builder> {
   string llvmBuilder = builder;
 }
 
 def LLVM_OneResultOpBuilder : OpBuilder<
   "Builder *, OperationState &result, Type resultType, "
   "ValueRange operands, ArrayRef<NamedAttribute> attributes = {}",

 def LLVM_SIToFPOp : LLVM_CastOp<"sitofp", "CreateSIToFP">;
 def LLVM_UIToFPOp : LLVM_CastOp<"uitofp", "CreateUIToFP">;
 def LLVM_FPToSIOp : LLVM_CastOp<"fptosi", "CreateFPToSI">;
 def LLVM_FPToUIOp : LLVM_CastOp<"fptoui", "CreateFPToUI">;
 def LLVM_FPExtOp : LLVM_CastOp<"fpext", "CreateFPExt">;
 def LLVM_FPTruncOp : LLVM_CastOp<"fptrunc", "CreateFPTrunc">;
 
 // Call-related operations.
-def LLVM_InvokeOp : LLVM_Op<"invoke", [Terminator]>,
+def LLVM_InvokeOp : LLVM_Op<"invoke", [
+                      DeclareOpInterfaceMethods<BranchOpInterface>, Terminator
+                    ]>,
                     Arguments<(ins OptionalAttr<FlatSymbolRefAttr>:$callee,
                                Variadic<LLVM_Type>)>,
                     Results<(outs Variadic<LLVM_Type>)> {
   let successors = (successor AnySuccessor:$normalDest,
                               AnySuccessor:$unwindDest);
 
   let builders = [OpBuilder<
     "Builder *b, OperationState &result, ArrayRef<Type> tys, "

 }
 def LLVM_FreezeOp : LLVM_OneResultOp<"freeze", [SameOperandsAndResultType]>,
                     Arguments<(ins LLVM_Type:$val)> {
   let assemblyFormat = "$val attr-dict `:` type($val)";
   string llvmBuilder = "builder.CreateFreeze($val);";
 }
 
 // Terminators.
-def LLVM_BrOp : LLVM_TerminatorOp<"br", []> {
+def LLVM_BrOp : LLVM_TerminatorOp<"br",
+    [DeclareOpInterfaceMethods<BranchOpInterface>]> {
   let successors = (successor AnySuccessor:$dest);
   let assemblyFormat = "$dest attr-dict";
 }
-def LLVM_CondBrOp : LLVM_TerminatorOp<"cond_br", []> {
+def LLVM_CondBrOp : LLVM_TerminatorOp<"cond_br",
+    [DeclareOpInterfaceMethods<BranchOpInterface>]> {
   let arguments = (ins LLVMI1:$condition);
   let successors = (successor AnySuccessor:$trueDest, AnySuccessor:$falseDest);
   let assemblyFormat = "$condition `,` successors attr-dict";
 }
 def LLVM_ReturnOp : LLVM_TerminatorOp<"return", []>,
                     Arguments<(ins Variadic<LLVM_Type>:$args)> {
   string llvmBuilder = [{
     if ($_numOperands != 0)

mlir/include/mlir/Dialect/SPIRV/SPIRVControlFlowOps.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPIRV_CONTROLFLOW_OPS
 #define SPIRV_CONTROLFLOW_OPS
 
 include "mlir/Dialect/SPIRV/SPIRVBase.td"
 include "mlir/Analysis/CallInterfaces.td"
+include "mlir/Analysis/ControlFlowInterfaces.td"
 
 // -----
 
-def SPV_BranchOp : SPV_Op<"Branch", [InFunctionScope, Terminator]> {
+def SPV_BranchOp : SPV_Op<"Branch",[
+    DeclareOpInterfaceMethods<BranchOpInterface>, InFunctionScope,
+    Terminator]> {
   let summary = "Unconditional branch to target block.";
 
   let description = [{
     This instruction must be the last instruction in a block.
 
     ### Custom assembly form
 
     ```
def SPV_BranchOp : SPV_Op<"Branch",[
 
   let autogenSerialization = 0;
 
   let assemblyFormat = "successors attr-dict";
 }
 
 // -----
 
-def SPV_BranchConditionalOp : SPV_Op<"BranchConditional",
-                                     [InFunctionScope, Terminator]> {
+def SPV_BranchConditionalOp : SPV_Op<"BranchConditional", [
+    DeclareOpInterfaceMethods<BranchOpInterface>, InFunctionScope,
+    Terminator]> {
   let summary = [{
     If Condition is true, branch to true block, otherwise branch to false
     block.
   }];
 
   let description = [{
     Condition must be a Boolean type scalar.
 

mlir/include/mlir/Dialect/SPIRV/SPIRVOps.h

 //===- SPIRVOps.h - MLIR SPIR-V operations ----------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the operations in the SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPIRV_SPIRVOPS_H_
 #define MLIR_DIALECT_SPIRV_SPIRVOPS_H_
 
+#include "mlir/Analysis/ControlFlowInterfaces.h"
 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/IR/Function.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 class OpBuilder;
 
 namespace spirv {

mlir/include/mlir/Dialect/StandardOps/IR/Ops.h

 // in the MLIR operation set.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_STANDARDOPS_IR_OPS_H
 #define MLIR_DIALECT_STANDARDOPS_IR_OPS_H
 
 #include "mlir/Analysis/CallInterfaces.h"
+#include "mlir/Analysis/ControlFlowInterfaces.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/StandardTypes.h"
 
 // Pull in all enum type definitions and utility function declarations.
 #include "mlir/Dialect/StandardOps/IR/OpsEnums.h.inc"
 

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 //===- Ops.td - Standard operation definitions -------------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Defines some MLIR standard operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef STANDARD_OPS
 #define STANDARD_OPS
 
 include "mlir/Analysis/CallInterfaces.td"
+include "mlir/Analysis/ControlFlowInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 
 def Std_Dialect : Dialect {
   let name = "std";
   let cppNamespace = "";
 }
 
 // Base class for Standard dialect ops.
let extraClassDeclaration = [{
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // BranchOp
 //===----------------------------------------------------------------------===//
 
-def BranchOp : Std_Op<"br", [Terminator]> {
+def BranchOp : Std_Op<"br",
+    [DeclareOpInterfaceMethods<BranchOpInterface>, Terminator]> {
   let summary = "branch operation";
   let description = [{
     The "br" operation represents a branch operation in a function.
     The operation takes variable number of operands and produces no results.
     The operand number and types for each successor must match the arguments of
     the block successor. For example:
 
       ^bb2:
def CmpIOp : Std_Op<"cmpi",
 
   let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)";
 }
 
 //===----------------------------------------------------------------------===//
 // CondBranchOp
 //===----------------------------------------------------------------------===//
 
-def CondBranchOp : Std_Op<"cond_br", [Terminator]> {
+def CondBranchOp : Std_Op<"cond_br",
+    [DeclareOpInterfaceMethods<BranchOpInterface>, Terminator]> {
   let summary = "conditional branch operation";
   let description = [{
     The "cond_br" operation represents a conditional branch operation in a
     function. The operation takes variable number of operands and produces
     no results. The operand number and types for each successor must match the
     arguments of the block successor. For example:
 
       ^bb0:

mlir/include/mlir/IR/OperationSupport.h

public:
   OperandRange(Operation *op);
 
   /// Returns the types of the values within this range.
   using type_iterator = ValueTypeIterator<iterator>;
   using type_range = ValueTypeRange<OperandRange>;
   type_range getTypes() const { return {begin(), end()}; }
   auto getType() const { return getTypes(); }
 
+  /// Return the operand index of the first element of this range. The range
+  /// must not be empty.
+  unsigned getBeginOperandIndex() const;
+
 private:
   /// See `detail::indexed_accessor_range_base` for details.
   static OpOperand *offset_base(OpOperand *object, ptrdiff_t index) {
     return object + index;
   }
   /// See `detail::indexed_accessor_range_base` for details.
   static Value dereference_iterator(OpOperand *object, ptrdiff_t index) {
     return object[index].get();

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AffineAnalysis.cpp
   AffineStructures.cpp
   CallGraph.cpp
+  ControlFlowInterfaces.cpp
   Dominance.cpp
   InferTypeOpInterface.cpp
   Liveness.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   SliceAnalysis.cpp
   Utils.cpp
   Verifier.cpp
   )
 
 add_llvm_library(MLIRAnalysis
   CallGraph.cpp
+  ControlFlowInterfaces.cpp
   InferTypeOpInterface.cpp
   Liveness.cpp
   SliceAnalysis.cpp
   Dominance.cpp
   Verifier.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
   )
 add_dependencies(MLIRAnalysis
   MLIRAffineOps
   MLIRCallOpInterfacesIncGen
+  MLIRControlFlowInterfacesIncGen
   MLIRTypeInferOpInterfaceIncGen
   MLIRLoopOps
   )
 
 target_link_libraries(MLIRAnalysis MLIRAffineOps MLIRLoopOps)
 
 add_llvm_library(MLIRLoopAnalysis
   AffineAnalysis.cpp
   AffineStructures.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
   )
 add_dependencies(MLIRLoopAnalysis
   MLIRAffineOps
   MLIRCallOpInterfacesIncGen
+  MLIRControlFlowInterfacesIncGen
   MLIRTypeInferOpInterfaceIncGen
   MLIRLoopOps
   )
 
 target_link_libraries(MLIRLoopAnalysis MLIRAffineOps MLIRLoopOps)

mlir/lib/Analysis/ControlFlowInterfaces.cpp

+//===- ControlFlowInterfaces.h - ControlFlow Interfaces -------------------===//
+//
+// 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/Analysis/ControlFlowInterfaces.h"
+#include "mlir/IR/StandardTypes.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// ControlFlowInterfaces
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/ControlFlowInterfaces.cpp.inc"
+
+//===----------------------------------------------------------------------===//
+// BranchOpInterface
+//===----------------------------------------------------------------------===//
+
+/// Erase an operand from a branch operation that is used as a successor
+/// operand. 'operandIndex' is the operand within 'operands' to be erased.
+void mlir::detail::eraseBranchSuccessorOperand(OperandRange operands,
+                                               unsigned operandIndex,
+                                               Operation *op) {
+  assert(operandIndex < operands.size() &&
+         "invalid index for successor operands");
+
+  // Erase the operand from the operation.
+  size_t fullOperandIndex = operands.getBeginOperandIndex() + operandIndex;
+  op->eraseOperand(fullOperandIndex);
+
+  // If this operation has an OperandSegmentSizeAttr, keep it up to date.
+  auto operandSegmentAttr =
+      op->getAttrOfType<DenseElementsAttr>("operand_segment_sizes");
+  if (!operandSegmentAttr)
+    return;
+
+  // Find the segment containing the full operand index and decrement it.
+  // TODO: This seems like a general utility that could be added somewhere.
+  SmallVector<int32_t, 4> values(operandSegmentAttr.getValues<int32_t>());
+  unsigned currentSize = 0;
+  for (unsigned i = 0, e = values.size(); i != e; ++i) {
+    currentSize += values[i];
+    if (fullOperandIndex < currentSize) {
+      --values[i];
+      break;
+    }
+  }
+  op->setAttr("operand_segment_sizes",
+              DenseIntElementsAttr::get(operandSegmentAttr.getType(), values));
+}
+
+/// 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> mlir::detail::getBranchSuccessorArgument(
+    Optional<OperandRange> operands, unsigned operandIndex, Block *successor) {
+  // Check that the operands are valid.
+  if (!operands || operands->empty())
+    return llvm::None;
+
+  // Check to ensure that this operand is within the range.
+  unsigned operandsStart = operands->getBeginOperandIndex();
+  if (operandIndex < operandsStart ||
+      operandIndex >= (operandsStart + operands->size()))
+    return llvm::None;
+
+  // Index the successor.
+  unsigned argIndex = operandIndex - operandsStart;
+  return successor->getArgument(argIndex);
+}
+
+/// Verify that the given operands match those of the given successor block.
+LogicalResult
+mlir::detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
+                                            Optional<OperandRange> operands) {
+  if (!operands)
+    return success();
+
+  // Check the count.
+  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) {
+    if ((*operandIt).getType() != destBB->getArgument(i).getType())
+      return op->emitError() << "type mismatch for bb argument #" << i
+                             << " of successor #" << succNo;
+  }
+  return success();
+}

mlir/lib/Dialect/LLVMIR/CMakeLists.txt

 add_mlir_dialect_library(MLIRLLVMIR
   IR/LLVMDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/LLVMIR
   )
-add_dependencies(MLIRLLVMIR MLIRLLVMOpsIncGen MLIRLLVMConversionsIncGen MLIROpenMP LLVMFrontendOpenMP LLVMAsmParser LLVMCore LLVMSupport)
+add_dependencies(MLIRLLVMIR MLIRControlFlowInterfacesIncGen MLIRLLVMOpsIncGen MLIRLLVMConversionsIncGen MLIROpenMP LLVMFrontendOpenMP LLVMAsmParser LLVMCore LLVMSupport)
 target_link_libraries(MLIRLLVMIR LLVMAsmParser LLVMCore LLVMSupport LLVMFrontendOpenMP MLIROpenMP MLIRIR)
 
 add_mlir_dialect_library(MLIRNVVMIR
   IR/NVVMDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/LLVMIR
   )

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

static ParseResult parseAllocaOp(OpAsmParser &parser, OperationState &result) {
   if (parser.resolveOperand(arraySize, funcType.getInput(0), result.operands))
     return failure();
 
   result.addTypes({funcType.getResult(0)});
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// LLVM::BrOp
+//===----------------------------------------------------------------------===//
+
+Optional<OperandRange> BrOp::getSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return getOperands();
+}
+
+bool BrOp::canEraseSuccessorOperand() { return true; }
+
+//===----------------------------------------------------------------------===//
+// LLVM::CondBrOp
+//===----------------------------------------------------------------------===//
+
+Optional<OperandRange> CondBrOp::getSuccessorOperands(unsigned index) {
+  assert(index < getNumSuccessors() && "invalid successor index");
+  return index == 0 ? trueDestOperands() : falseDestOperands();
+}
+
+bool CondBrOp::canEraseSuccessorOperand() { return true; }
+
+//===----------------------------------------------------------------------===//
 // Printing/parsing for LLVM::LoadOp.
 //===----------------------------------------------------------------------===//
 
 static void printLoadOp(OpAsmPrinter &p, LoadOp &op) {
   p << op.getOperationName() << ' ' << op.addr();
   p.printOptionalAttrDict(op.getAttrs());
   p << " : " << op.addr().getType();
 }
static ParseResult parseStoreOp(OpAsmParser &parser, OperationState &result) {
   if (parser.resolveOperand(value, elemTy, result.operands) ||
       parser.resolveOperand(addr, type, result.operands))
     return failure();
 
   return success();
 }
 
 ///===----------------------------------------------------------------------===//
-/// Verifying/Printing/Parsing for LLVM::InvokeOp.
+/// LLVM::InvokeOp
 ///===----------------------------------------------------------------------===//
 
+Optional<OperandRange> InvokeOp::getSuccessorOperands(unsigned index) {
+  assert(index < getNumSuccessors() && "invalid successor index");
+  return index == 0 ? normalDestOperands() : unwindDestOperands();
+}
+
+bool InvokeOp::canEraseSuccessorOperand() { return true; }
+
 static LogicalResult verify(InvokeOp op) {
   if (op.getNumResults() > 1)
     return op.emitOpError("must have 0 or 1 result");
 
   Block *unwindDest = op.unwindDest();
   if (unwindDest->empty())
     return op.emitError(
         "must have at least one operation in unwind destination");

mlir/lib/Dialect/SPIRV/CMakeLists.txt

add_mlir_dialect_library(MLIRSPIRV
   SPIRVTypes.cpp
   TargetAndABI.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SPIRV
   )
 
 add_dependencies(MLIRSPIRV
+  MLIRControlFlowInterfacesIncGen
   MLIRSPIRVAvailabilityIncGen
   MLIRSPIRVCanonicalizationIncGen
   MLIRSPIRVEnumAvailabilityIncGen
   MLIRSPIRVEnumsIncGen
   MLIRSPIRVOpsIncGen
   MLIRSPIRVOpUtilsGen
   MLIRSPIRVTargetAndABIIncGen)
 
 target_link_libraries(MLIRSPIRV
   MLIRIR
   MLIRParser
   MLIRSupport
   MLIRTransforms)
 
 add_subdirectory(Serialization)
 add_subdirectory(Transforms)

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

if (operandBitWidth != resultBitWidth) {
     return bitcastOp.emitOpError("mismatch in result type bitwidth ")
            << resultBitWidth << " and operand type bitwidth "
            << operandBitWidth;
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// spv.BranchOp
+//===----------------------------------------------------------------------===//
+
+Optional<OperandRange> spirv::BranchOp::getSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return getOperands();
+}
+
+bool spirv::BranchOp::canEraseSuccessorOperand() { return true; }
+
+//===----------------------------------------------------------------------===//
 // spv.BranchConditionalOp
 //===----------------------------------------------------------------------===//
 
+Optional<OperandRange>
+spirv::BranchConditionalOp::getSuccessorOperands(unsigned index) {
+  assert(index < 2 && "invalid successor index");
+  return index == kTrueIndex ? getTrueBlockArguments()
+                             : getFalseBlockArguments();
+}
+
+bool spirv::BranchConditionalOp::canEraseSuccessorOperand() { return true; }
+
 static ParseResult parseBranchConditionalOp(OpAsmParser &parser,
                                             OperationState &state) {
   auto &builder = parser.getBuilder();
   OpAsmParser::OperandType condInfo;
   Block *dest;
   SmallVector<Value, 4> destOperands;
 
   // Parse the condition.

mlir/lib/Dialect/StandardOps/CMakeLists.txt

 add_mlir_dialect_library(MLIRStandardOps
   IR/Ops.cpp
   EDSC/Builders.cpp
   EDSC/Intrinsics.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps
   )
 add_dependencies(MLIRStandardOps
 
   MLIRCallOpInterfacesIncGen
+  MLIRControlFlowInterfacesIncGen
   MLIREDSC
   MLIRIR
   MLIRStandardOpsIncGen
   LLVMSupport
   )
 target_link_libraries(MLIRStandardOps
 
   MLIREDSC
   MLIRIR
   LLVMSupport
   )

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 };
 } // end anonymous namespace.
 
 Block *BranchOp::getDest() { return getSuccessor(); }
 
 void BranchOp::setDest(Block *block) { return setSuccessor(block); }
 
 void BranchOp::eraseOperand(unsigned index) {
-  getOperation()->eraseSuccessorOperand(0, index);
+  getOperation()->eraseOperand(index);
 }
 
 void BranchOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                            MLIRContext *context) {
   results.insert<SimplifyBrToBlockWithSinglePred>(context);
 }
 
+Optional<OperandRange> BranchOp::getSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return getOperands();
+}
+
+bool BranchOp::canEraseSuccessorOperand() { return true; }
+
 //===----------------------------------------------------------------------===//
 // CallOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(CallOp op) {
   // Check that the callee attribute was specified.
   auto fnAttr = op.getAttrOfType<FlatSymbolRefAttr>("callee");
   if (!fnAttr)

 };
 } // end anonymous namespace.
 
 void CondBranchOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<SimplifyConstCondBranchPred>(context);
 }
 
+Optional<OperandRange> CondBranchOp::getSuccessorOperands(unsigned index) {
+  assert(index < getNumSuccessors() && "invalid successor index");
+  return index == trueIndex ? getTrueOperands() : getFalseOperands();
+}
+
+bool CondBranchOp::canEraseSuccessorOperand() { return true; }
+
 //===----------------------------------------------------------------------===//
 // Constant*Op
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, ConstantOp &op) {
   p << "constant ";
   p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"value"});
 

mlir/lib/IR/Operation.cpp

 LogicalResult OpTrait::impl::verifyIsTerminator(Operation *op) {
   Block *block = op->getBlock();
   // Verify that the operation is at the end of the respective parent block.
   if (!block || &block->back() != op)
     return op->emitOpError("must be the last operation in the parent block");
   return success();
 }
 
-static LogicalResult verifySuccessor(Operation *op, unsigned succNo) {
-  Operation::operand_range operands = op->getSuccessorOperands(succNo);
-  unsigned operandCount = op->getNumSuccessorOperands(succNo);
-  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();
-
-  auto operandIt = operands.begin();
-  for (unsigned i = 0, e = operandCount; i != e; ++i, ++operandIt) {
-    if ((*operandIt).getType() != destBB->getArgument(i).getType())
-      return op->emitError() << "type mismatch for bb argument #" << i
-                             << " of successor #" << succNo;
-  }
-
-  return success();
-}
-
 static LogicalResult verifyTerminatorSuccessors(Operation *op) {
   auto *parent = op->getParentRegion();
 
   // Verify that the operands lines up with the BB arguments in the successor.
-  for (unsigned i = 0, e = op->getNumSuccessors(); i != e; ++i) {
-    auto *succ = op->getSuccessor(i);
+  for (Block *succ : op->getSuccessors())
     if (succ->getParent() != parent)
       return op->emitError("reference to block defined in another region");
-    if (failed(verifySuccessor(op, i)))
-      return failure();
-  }
   return success();
 }
 
 LogicalResult OpTrait::impl::verifyZeroSuccessor(Operation *op) {
   if (op->getNumSuccessors() != 0) {
     return op->emitOpError("requires 0 successors but found ")
            << op->getNumSuccessors();
   }

mlir/lib/IR/OperationSupport.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // OperandRange
 
 OperandRange::OperandRange(Operation *op)
     : OperandRange(op->getOpOperands().data(), op->getNumOperands()) {}
 
+/// Return the operand index of the first element of this range. The range
+/// must not be empty.
+unsigned OperandRange::getBeginOperandIndex() const {
+  assert(!empty() && "range must not be empty");
+  return base->getOperandNumber();
+}
+
 //===----------------------------------------------------------------------===//
 // ResultRange
 
 ResultRange::ResultRange(Operation *op)
     : ResultRange(op, /*startIndex=*/0, op->getNumResults()) {}
 
 ArrayRef<Type> ResultRange::getTypes() const {
   return getBase()->getResultTypes();

mlir/test/lib/TestDialect/CMakeLists.txt

 mlir_tablegen(TestPatterns.inc -gen-rewriters)
 add_public_tablegen_target(MLIRTestOpsIncGen)
 
 add_llvm_library(MLIRTestDialect
   TestDialect.cpp
   TestPatterns.cpp
 )
 add_dependencies(MLIRTestDialect
+  MLIRControlFlowInterfacesIncGen
   MLIRTestOpsIncGen
   MLIRTypeInferOpInterfaceIncGen
 )
 target_link_libraries(MLIRTestDialect
   MLIRDialect
   MLIRIR
   MLIRLinalgTransforms
   MLIRPass
   MLIRTransforms
   MLIRTransformUtils
   LLVMSupport
 )

mlir/test/lib/TestDialect/TestDialect.h

 // that do not have a respective counterpart in the main source directories.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTDIALECT_H
 #define MLIR_TESTDIALECT_H
 
 #include "mlir/Analysis/CallInterfaces.h"
+#include "mlir/Analysis/ControlFlowInterfaces.h"
 #include "mlir/Analysis/InferTypeOpInterface.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/SymbolTable.h"
 

mlir/test/lib/TestDialect/TestDialect.cpp

TestDialect::verifyRegionResultAttribute(Operation *op, unsigned regionIndex,
                                          unsigned resultIndex,
                                          NamedAttribute namedAttr) {
   if (namedAttr.first == "test.invalid_attr")
     return op->emitError() << "invalid to use 'test.invalid_attr'";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// TestBranchOp
+//===----------------------------------------------------------------------===//
+
+Optional<OperandRange> TestBranchOp::getSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return getOperands();
+}
+
+bool TestBranchOp::canEraseSuccessorOperand() { return true; }
+
+//===----------------------------------------------------------------------===//
 // Test IsolatedRegionOp - parse passthrough region arguments.
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseIsolatedRegionOp(OpAsmParser &parser,
                                          OperationState &result) {
   OpAsmParser::OperandType argInfo;
   Type argType = parser.getBuilder().getIndexType();
 

mlir/test/lib/TestDialect/TestOps.td

 //===-- TestOps.td - Test dialect operation definitions ----*- tablegen -*-===//
 //
 // 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 TEST_OPS
 #define TEST_OPS
 
 include "mlir/IR/OpBase.td"
 include "mlir/IR/OpAsmInterface.td"
+include "mlir/Analysis/ControlFlowInterfaces.td"
 include "mlir/Analysis/CallInterfaces.td"
 include "mlir/Analysis/InferTypeOpInterface.td"
 
 def TEST_Dialect : Dialect {
   let name = "test";
   let cppNamespace = "";
 }
 
[
      // The op will have methods implementing the ShapedType type infer interface.
      DeclareOpInterfaceMethods<InferShapedTypeOpInterface>,
      // The op produces tensors and will use the ShapedType type infer interface
      // along with knowledge that it is producing Tensors to infer shape.
      InferTensorType
    ]> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
+
+  let extraClassDeclaration = [{
+    LogicalResult reifyReturnTypeShapes(OpBuilder &builder,
+                                        SmallVectorImpl<Value> &shapes);
+  }];
 }
 
 def IsNotScalar : Constraint<CPred<"$0.getType().getRank() != 0">>;
 
 def UpdateAttr : Pat<(I32ElementsAttrOp $attr),
                      (I32ElementsAttrOp ConstantAttr<I32ElementsAttr, "0">),
                      [(IsNotScalar $attr)]>;
 
-def TestBranchOp : TEST_Op<"br", [Terminator]> {
+def TestBranchOp : TEST_Op<"br",
+    [DeclareOpInterfaceMethods<BranchOpInterface>, Terminator]> {
   let successors = (successor AnySuccessor:$target);
 }
 
 def AttrSizedOperandOp : TEST_Op<"attr_sized_operands",
                                  [AttrSizedOperandSegments]> {
   let arguments = (ins
     Variadic<I32>:$a,
     Variadic<I32>:$b,

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

 
 void OpEmitter::genOpInterfaceMethods() {
   for (const auto &trait : op.getTraits()) {
     auto opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait);
     if (!opTrait || !opTrait->shouldDeclareMethods())
       continue;
     auto interface = opTrait->getOpInterface();
     for (auto method : interface.getMethods()) {
-      // Don't declare if the method has a body.
-      if (method.getBody())
+      // Don't declare if the method has a body or a default implementation.
+      if (method.getBody() || method.getDefaultImplementation())
         continue;
       std::string args;
       llvm::raw_string_ostream os(args);
       mlir::interleaveComma(method.getArguments(), os,
                             [&](const OpInterfaceMethod::Argument &arg) {
                               os << arg.type << " " << arg.name;
                             });
       opClass.newMethod(method.getReturnType(), method.getName(), os.str(),