[mlir][ODS] Add support for specifying the successors of an operation.

[rriddle]

Summary

This revision add support in ODS for specifying the successors of an operation. Successors are specified via the successors list:

let successors = (successor AnySuccessor:$target, AnySuccessor:$otherTarget);

Depends On D74683

[jpienaar]

NIce

mlir/lib/TableGen/Operator.cpp

316

Where does/could 'successorsOp->getDef()->getName() != "successor"' happen?

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1370

Wouldn't formatv(...).str() also give this?

Feels a little weird to have a variable called name with spaces ...

1378

Why not move this up and then you can emit the name directly in the body rather than concatting string above?

1442

It is a little bit difficult to correlate variadic operands with successors, could you explain that here?

[rriddle]

Resolve comments

[rriddle]

mlir/lib/TableGen/Operator.cpp

316

I'm assuming this is intended to capture cases like : let successors = (outs ...); or let successors = (regions ...);. I followed suit with what the regions/results/etc. lists were verifying.

[rriddle]

Rebase

[antiagainst]

Cool! Thanks River for adding this! Can we also update the doc ?

mlir/include/mlir/TableGen/Operator.h

212

I've default to 0. I think typically ops don't have successors anyway.

[rriddle]

Resolve comments.

[bondhugula]

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

235 ↗

(On Diff #246016)

Where did 'arguments' disappear here? (i.e., where are the operands associated with successors?) I don't see the documentation being updated on this.

mlir/docs/OpDefinitions.md

 ```
 
 #### Variadic results
 
 Similar to variadic operands, `Variadic<...>` can also be used for results.
 And similarly, `SameVariadicResultSize` for multiple variadic results in the
 same operation.
 
+### Operation successors
+
+For terminator operations, the successors are specified inside of the
+`dag`-typed `successors`, led by `successor`:
+
+```tablegen
+let successors = (successor
+  <successor-constraint>:$<successor-name>,
+  ...
+);
+```
+
+#### Variadic successors
+
+Similar to the `Variadic` class used for variadic operands and results,
+`VariadicSuccessor<...>` can be used for successors. Variadic successors can
+currently only be specified as the last successor in the successor list.
+
 ### Operation traits and constraints
 
 Traits are operation properties that affect syntax or semantics. MLIR C++
 models various traits in the `mlir::OpTrait` namespace.
 
 Both operation traits, [interfaces](#operation-interfaces), and constraints
 involving multiple operands/attributes/results are provided as the second
 template parameter to the `Op` class. They should be deriving from the `OpTrait`

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

def LLVM_Type : Type<CPred<"$_self.isa<::mlir::LLVM::LLVMType>()">,
                      "LLVM dialect type">;
 
 // Type constraint accepting only wrapped LLVM integer types.
 def LLVMInt : TypeConstraint<
   And<[LLVM_Type.predicate,
        CPred<"$_self.cast<::mlir::LLVM::LLVMType>().isIntegerTy()">]>,
   "LLVM dialect integer">;
 
+def LLVMIntBase : TypeConstraint<
+  And<[LLVM_Type.predicate,
+       CPred<"$_self.cast<::mlir::LLVM::LLVMType>().isIntegerTy()">]>,
+  "LLVM dialect integer">;
+
+// Integer type of a specific width.
+class LLVMI<int width>
+    : Type<And<[
+        LLVM_Type.predicate,
+        CPred<
+         "$_self.cast<::mlir::LLVM::LLVMType>().isIntegerTy(" # width # ")">]>,
+         "LLVM dialect " # width # "-bit integer">,
+      BuildableType<
+        "::mlir::LLVM::LLVMType::getIntNTy("
+        "$_builder.getContext()->getRegisteredDialect<LLVM::LLVMDialect>(),"
+         # width # ")">;
+
+def LLVMI1 : LLVMI<1>;
+
 // Base class for LLVM operations. Defines the interface to the llvm::IRBuilder
 // used to translate to LLVM IR proper.
 class LLVM_OpBase<Dialect dialect, string mnemonic, list<OpTrait> traits = []> :
     Op<dialect, mnemonic, traits> {
   // A pattern for constructing the LLVM IR Instruction (or other Value) that
   // corresponds to this op.  This pattern can use `builder` to refer to an
   // `llvm::IRBuilder<>` instance, $-names of arguments and results and the
   // following special variable names:

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

 // Base class for LLVM operations with zero results.
 class LLVM_ZeroResultOp<string mnemonic, list<OpTrait> traits = []> :
     LLVM_Op<mnemonic, traits>, Results<(outs)>,
     LLVM_TwoBuilders<LLVM_VoidResultTypeOpBuilder, LLVM_ZeroResultOpBuilder>;
 
 // Base class for LLVM terminator operations.  All terminator operations have
 // zero results and an optional list of successors.
 class LLVM_TerminatorOp<string mnemonic, list<OpTrait> traits = []> :
-    LLVM_Op<mnemonic, !listconcat(traits, [Terminator])>,
-    Arguments<(ins Variadic<LLVM_Type>:$args)>, Results<(outs)> {
+    LLVM_Op<mnemonic, !listconcat(traits, [Terminator])> {
   let builders = [
     OpBuilder<
       "Builder *, OperationState &result, "
       "ValueRange properOperands, "
       "ArrayRef<Block *> destinations, "
       "ArrayRef<ValueRange> operands, "
       "ArrayRef<NamedAttribute> attributes = {}",
       [{

 def LLVM_FPExtOp : LLVM_CastOp<"fpext", "CreateFPExt">;
 def LLVM_FPTruncOp : LLVM_CastOp<"fptrunc", "CreateFPTrunc">;
 
 // Call-related operations.
 def LLVM_InvokeOp : LLVM_Op<"invoke", [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, "
-    "FlatSymbolRefAttr callee, ValueRange ops, Block* normal, "
-    "ValueRange normalOps, Block* unwind, ValueRange unwindOps",
-    [{
-      result.addAttribute("callee", callee);
-      build(b, result, tys, ops, normal, normalOps, unwind, unwindOps);
-    }]>,
-    OpBuilder<
-    "Builder *b, OperationState &result, ArrayRef<Type> tys, "
     "ValueRange ops, Block* normal, "
     "ValueRange normalOps, Block* unwind, ValueRange unwindOps",
     [{
       result.addTypes(tys);
       result.addOperands(ops);
       result.addSuccessor(normal, normalOps);
       result.addSuccessor(unwind, unwindOps);
     }]>];
let builders = [OpBuilder<
     "Value rhs", [{
       build(b, result, lhs.getType(), condition, lhs, rhs);
     }]>];
   let assemblyFormat = "operands attr-dict `:` type($condition) `,` type($res)";
 }
 
 // Terminators.
 def LLVM_BrOp : LLVM_TerminatorOp<"br", []> {
+  let successors = (successor AnySuccessor:$dest);
   let parser = [{ return parseBrOp(parser, result); }];
   let printer = [{ printBrOp(p, *this); }];
 }
 def LLVM_CondBrOp : LLVM_TerminatorOp<"cond_br", []> {
-  let verifier = [{
-    if (getNumSuccessors() != 2)
-      return emitOpError("expected exactly two successors");
-    return success();
-  }];
+  let arguments = (ins LLVMI1:$condition);
+  let successors = (successor AnySuccessor:$trueDest, AnySuccessor:$falseDest);
+
   let parser = [{ return parseCondBrOp(parser, result); }];
   let printer = [{ printCondBrOp(p, *this); }];
 }
-def LLVM_ReturnOp : LLVM_TerminatorOp<"return", []> {
+def LLVM_ReturnOp : LLVM_TerminatorOp<"return", []>,
+                    Arguments<(ins Variadic<LLVM_Type>:$args)> {
   string llvmBuilder = [{
     if ($_numOperands != 0)
       builder.CreateRet($args[0]);
     else
       builder.CreateRetVoid();
   }];
 
   let verifier = [{

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

let description = [{
     For example:
 
     ```
     spv.Branch ^target
     spv.Branch ^target(%0, %1: i32, f32)
     ```
   }];
 
-  let arguments = (ins
-    Variadic<AnyType>:$block_arguments
-  );
+  let arguments = (ins);
 
   let results = (outs);
 
+  let successors = (successor AnySuccessor:$target);
+
+  let verifier = [{ return success(); }];
+
   let builders = [
     OpBuilder<
       "Builder *, OperationState &state, "
       "Block *successor, ValueRange arguments = {}", [{
         state.addSuccessor(successor, arguments);
       }]
     >
   ];
 
   let skipDefaultBuilders = 1;
 
   let extraClassDeclaration = [{
     /// Returns the branch target block.
-    Block *getTarget() { return getOperation()->getSuccessor(0); }
+    Block *getTarget() { return target(); }
 
     /// Returns the block arguments.
-    operand_range getBlockArguments() {
-      return getOperation()->getSuccessorOperands(0);
-    }
+    operand_range getBlockArguments() { return targetOperands(); }
   }];
 
   let autogenSerialization = 0;
 }
 
 // -----
 
 def SPV_BranchConditionalOp : SPV_Op<"BranchConditional",
let description = [{
     ```
     spv.BranchConditional %condition, ^true_branch, ^false_branch
     spv.BranchConditional %condition, ^true_branch(%0: i32), ^false_branch(%1: i32)
     ```
   }];
 
   let arguments = (ins
     SPV_Bool:$condition,
-    Variadic<AnyType>:$branch_arguments,
     OptionalAttr<I32ArrayAttr>:$branch_weights
   );
 
   let results = (outs);
 
+  let successors = (successor AnySuccessor:$trueTarget,
+                               AnySuccessor:$falseTarget);
+
   let builders = [
     OpBuilder<
       "Builder *builder, OperationState &state, Value condition, "
       "Block *trueBlock, ValueRange trueArguments, "
       "Block *falseBlock, ValueRange falseArguments, "
       "Optional<std::pair<uint32_t, uint32_t>> weights = {}",
       [{
         state.addOperands(condition);

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

let description = [{
     the block successor. For example:
 
       ^bb2:
         %2 = call @someFn()
         br ^bb3(%2 : tensor<*xf32>)
       ^bb3(%3: tensor<*xf32>):
   }];
 
-  let arguments = (ins Variadic<AnyType>:$operands);
+  let successors = (successor AnySuccessor:$dest);
 
-  let builders = [OpBuilder<
-    "Builder *, OperationState &result, Block *dest,"
-    "ValueRange operands = {}", [{
-      result.addSuccessor(dest, operands);
+  let builders = [OpBuilder<"Builder *, OperationState &result, Block *dest", [{
+    result.addSuccessor(dest, llvm::None);
   }]>];
 
   // BranchOp is fully verified by traits.
   let verifier = ?;
 
   let extraClassDeclaration = [{
     Block *getDest();
     void setDest(Block *block);
arguments of the block successor. For example:
          %0 = extract_element %arg0[] : tensor<i1>
          cond_br %0, ^bb1, ^bb2
       ^bb1:
          ...
       ^bb2:
          ...
   }];
 
-  let arguments = (ins I1:$condition, Variadic<AnyType>:$branchOperands);
-
-  let builders = [OpBuilder<
-    "Builder *, OperationState &result, Value condition,"
-    "Block *trueDest, ValueRange trueOperands,"
-    "Block *falseDest, ValueRange falseOperands", [{
-      result.addOperands(condition);
-      result.addSuccessor(trueDest, trueOperands);
-      result.addSuccessor(falseDest, falseOperands);
-  }]>];
+  let arguments = (ins I1:$condition);
+  let successors = (successor AnySuccessor:$trueDest, AnySuccessor:$falseDest);
 
   // CondBranchOp is fully verified by traits.
   let verifier = ?;
 
   let extraClassDeclaration = [{
     // These are the indices into the dests list.
     enum { trueIndex = 0, falseIndex = 1 };
 

mlir/include/mlir/IR/OpBase.td

 // Subclass for constraints on an attribute.
 class AttrConstraint<Pred predicate, string description = ""> :
     Constraint<predicate, description>;
 
 // Subclass for constraints on a region.
 class RegionConstraint<Pred predicate, string description = ""> :
     Constraint<predicate, description>;
 
+// Subclass for constraints on a successor.
+class SuccessorConstraint<Pred predicate, string description = ""> :
+    Constraint<predicate, description>;
+
 // How to use these constraint categories:
 //
 // * Use TypeConstraint to specify
 //   * Constraints on an op's operand/result definition
 //   * Further constraints to match an op's operand/result in source pattern
 //
 // * Use Attr (a subclass for AttrConstraint) for
 //   * Constraints on an op's attribute definition

 def AnyRegion : Region<CPred<"true">, "any region">;
 
 // A region with the given number of blocks.
 class SizedRegion<int numBlocks> : Region<
   CPred<"$_self.getBlocks().size() == " # numBlocks>,
   "region with " # numBlocks # " blocks">;
 
 //===----------------------------------------------------------------------===//
+// Successor definitions
+//===----------------------------------------------------------------------===//
+
+class Successor<Pred condition, string descr = ""> :
+    SuccessorConstraint<condition, descr>;
+
+// Any successor.
+def AnySuccessor : Successor<?, "any successor">;
+
+// A variadic successor constraint. It expands to zero or more of the base
+// successor.
+class VariadicSuccessor<Successor successor>
+  : Successor<successor.predicate, successor.description>;
+
+//===----------------------------------------------------------------------===//
 // OpTrait definitions
 //===----------------------------------------------------------------------===//
 
 // OpTrait represents a trait regarding an op.
 class OpTrait;
 
 // NativeOpTrait corresponds to the MLIR C++ OpTrait mechanism. The
 // purpose to wrap around C++ symbol string with this class is to make

 //===----------------------------------------------------------------------===//
 
 // Marker used to identify the result list for an op.
 def outs;
 
 // Marker used to identify the region list for an op.
 def region;
 
+// Marker used to identify the successor list for an op.
+def successor;
+
 // Class for defining a custom builder.
 //
 // TableGen generates several generic builders for each op by default (see
 // comment in the `Op` class). If the default generated ones cannot cover
 // some use case, custom builders can be defined using instances of this class.
 //
 // The signature of the builder is always
 //
class Op<Dialect dialect, string mnemonic, list<OpTrait> props = []> {
   dag arguments = (ins);
 
   // The list of results of the op. Default to 0 results.
   dag results = (outs);
 
   // The list of regions of the op. Default to 0 regions.
   dag regions = (region);
 
+  // The list of successors of the op. Default to 0 successors.
+  dag successors = (successor);
+
   // Attribute getters can be added to the op by adding an Attr member
   // with the name and type of the attribute. E.g., adding int attribute
   // with name "value" and type "i32":
   //   I32Attr value;
 
   // Define the hooks used for building, parsing, printing, verification.
 
   // Custom builder.

mlir/include/mlir/TableGen/Constraint.h

public:
   // mlir::Attribute.
   std::string getConditionTemplate() const;
 
   // Returns the user-readable description of this constraint. If the
   // description is not provided, returns the TableGen def name.
   StringRef getDescription() const;
 
   // Constraint kind
-  enum Kind { CK_Attr, CK_Region, CK_Type, CK_Uncategorized };
+  enum Kind { CK_Attr, CK_Region, CK_Successor, CK_Type, CK_Uncategorized };
 
   Kind getKind() const { return kind; }
 
 protected:
   Constraint(Kind kind, const llvm::Record *record);
 
   // The TableGen definition of this constraint.
   const llvm::Record *def;

mlir/include/mlir/TableGen/Operator.h

 #define MLIR_TABLEGEN_OPERATOR_H_
 
 #include "mlir/Support/LLVM.h"
 #include "mlir/TableGen/Argument.h"
 #include "mlir/TableGen/Attribute.h"
 #include "mlir/TableGen/Dialect.h"
 #include "mlir/TableGen/OpTrait.h"
 #include "mlir/TableGen/Region.h"
+#include "mlir/TableGen/Successor.h"
 #include "mlir/TableGen/Type.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/SMLoc.h"
 
 namespace llvm {
 class CodeInit;
public:
   // requiring the raw MLIR trait here.
   const OpTrait *getTrait(llvm::StringRef trait) const;
 
   // Returns the number of regions.
   unsigned getNumRegions() const;
   // Returns the `index`-th region.
   const NamedRegion &getRegion(unsigned index) const;
 
+  // Successors.
+  using const_successor_iterator = const NamedSuccessor *;
+  const_successor_iterator successor_begin() const;
+  const_successor_iterator successor_end() const;
+  llvm::iterator_range<const_successor_iterator> getSuccessors() const;
+
+  // Returns the number of successors.
+  unsigned getNumSuccessors() const;
+  // Returns the `index`-th successor.
+  const NamedSuccessor &getSuccessor(unsigned index) const;
+
+  // Returns the number of variadic successors in this operation.
+  unsigned getNumVariadicSuccessors() const;
+
   // Trait.
   using const_trait_iterator = const OpTrait *;
   const_trait_iterator trait_begin() const;
   const_trait_iterator trait_end() const;
   llvm::iterator_range<const_trait_iterator> getTraits() const;
 
   ArrayRef<llvm::SMLoc> getLoc() const;
 
private:
   SmallVector<NamedAttribute, 4> attributes;
 
   // The arguments of the op (operands and native attributes).
   SmallVector<Argument, 4> arguments;
 
   // The results of the op.
   SmallVector<NamedTypeConstraint, 4> results;
 
+  // The successors of this op.
+  SmallVector<NamedSuccessor, 0> successors;
+
   // The traits of the op.
   SmallVector<OpTrait, 4> traits;
 
   // The regions of this op.
   SmallVector<NamedRegion, 1> regions;
 
   // The number of native attributes stored in the leading positions of
   // `attributes`.

mlir/include/mlir/TableGen/Successor.h

+//===- Successor.h - TableGen successor definitions -------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TABLEGEN_SUCCESSOR_H_
+#define MLIR_TABLEGEN_SUCCESSOR_H_
+
+#include "mlir/Support/LLVM.h"
+#include "mlir/TableGen/Constraint.h"
+
+namespace mlir {
+namespace tblgen {
+
+// Wrapper class providing helper methods for accessing Successor defined in
+// TableGen.
+class Successor : public Constraint {
+public:
+  using Constraint::Constraint;
+
+  static bool classof(const Constraint *c) {
+    return c->getKind() == CK_Successor;
+  }
+
+  // Returns true if this successor is variadic.
+  bool isVariadic() const;
+};
+
+// A struct bundling a successor's constraint and its name.
+struct NamedSuccessor {
+  // Returns true if this successor is variadic.
+  bool isVariadic() const { return constraint.isVariadic(); }
+
+  StringRef name;
+  Successor constraint;
+};
+
+} // end namespace tblgen
+} // end namespace mlir
+
+#endif // MLIR_TABLEGEN_SUCCESSOR_H_

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

 
 ///===----------------------------------------------------------------------===//
 /// Verifying/Printing/Parsing for LLVM::InvokeOp.
 ///===----------------------------------------------------------------------===//
 
 static LogicalResult verify(InvokeOp op) {
   if (op.getNumResults() > 1)
     return op.emitOpError("must have 0 or 1 result");
-  if (op.getNumSuccessors() != 2)
-    return op.emitOpError("must have normal and unwind destinations");
 
-  if (op.getSuccessor(1)->empty())
+  Block *unwindDest = op.unwindDest();
+  if (unwindDest->empty())
     return op.emitError(
         "must have at least one operation in unwind destination");
 
   // In unwind destination, first operation must be LandingpadOp
-  if (!isa<LandingpadOp>(op.getSuccessor(1)->front()))
+  if (!isa<LandingpadOp>(unwindDest->front()))
     return op.emitError("first operation in unwind destination should be a "
                         "llvm.landingpad operation");
 
   return success();
 }
 
 static void printInvokeOp(OpAsmPrinter &p, InvokeOp &op) {
   auto callee = op.callee();

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

static ParseResult parseBranchOp(OpAsmParser &parser, OperationState &state) {
   return success();
 }
 
 static void print(spirv::BranchOp branchOp, OpAsmPrinter &printer) {
   printer << spirv::BranchOp::getOperationName() << ' ';
   printer.printSuccessorAndUseList(branchOp.getOperation(), /*index=*/0);
 }
 
-static LogicalResult verify(spirv::BranchOp branchOp) {
-  auto *op = branchOp.getOperation();
-  if (op->getNumSuccessors() != 1)
-    branchOp.emitOpError("must have exactly one successor");
-
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // spv.BranchConditionalOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseBranchConditionalOp(OpAsmParser &parser,
                                             OperationState &state) {
   auto &builder = parser.getBuilder();
   OpAsmParser::OperandType condInfo;
static void print(spirv::BranchConditionalOp branchOp, OpAsmPrinter &printer) {
   printer.printSuccessorAndUseList(branchOp.getOperation(),
                                    spirv::BranchConditionalOp::kTrueIndex);
   printer << ", ";
   printer.printSuccessorAndUseList(branchOp.getOperation(),
                                    spirv::BranchConditionalOp::kFalseIndex);
 }
 
 static LogicalResult verify(spirv::BranchConditionalOp branchOp) {
-  auto *op = branchOp.getOperation();
-  if (op->getNumSuccessors() != 2)
-    return branchOp.emitOpError("must have exactly two successors");
-
   if (auto weights = branchOp.branch_weights()) {
     if (weights->getValue().size() != 2) {
       return branchOp.emitOpError("must have exactly two branch weights");
     }
     if (llvm::all_of(*weights, [](Attribute attr) {
           return attr.cast<IntegerAttr>().getValue().isNullValue();
         }))
       return branchOp.emitOpError("branch weights cannot both be zero");

mlir/lib/TableGen/CMakeLists.txt

 add_llvm_library(LLVMMLIRTableGen
   Argument.cpp
   Attribute.cpp
   Constraint.cpp
   Dialect.cpp
   Format.cpp
   Operator.cpp
   OpClass.cpp
   OpInterfaces.cpp
   OpTrait.cpp
   Pattern.cpp
   Predicate.cpp
+  Successor.cpp
   Type.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/TableGen
   )
 target_link_libraries(LLVMMLIRTableGen LLVMSupport LLVMTableGen)

mlir/lib/TableGen/Constraint.cpp

 Constraint::Constraint(const llvm::Record *record)
     : def(record), kind(CK_Uncategorized) {
   if (record->isSubClassOf("TypeConstraint")) {
     kind = CK_Type;
   } else if (record->isSubClassOf("AttrConstraint")) {
     kind = CK_Attr;
   } else if (record->isSubClassOf("RegionConstraint")) {
     kind = CK_Region;
+  } else if (record->isSubClassOf("SuccessorConstraint")) {
+    kind = CK_Successor;
   } else {
     assert(record->isSubClassOf("Constraint"));
   }
 }
 
 Constraint::Constraint(Kind kind, const llvm::Record *record)
     : def(record), kind(kind) {}
 

mlir/lib/TableGen/Operator.cpp

 }
 
 unsigned tblgen::Operator::getNumRegions() const { return regions.size(); }
 
 const tblgen::NamedRegion &tblgen::Operator::getRegion(unsigned index) const {
   return regions[index];
 }
 
+auto tblgen::Operator::successor_begin() const -> const_successor_iterator {
+  return successors.begin();
+}
+auto tblgen::Operator::successor_end() const -> const_successor_iterator {
+  return successors.end();
+}
+auto tblgen::Operator::getSuccessors() const
+    -> llvm::iterator_range<const_successor_iterator> {
+  return {successor_begin(), successor_end()};
+}
+
+unsigned tblgen::Operator::getNumSuccessors() const {
+  return successors.size();
+}
+
+const tblgen::NamedSuccessor &
+tblgen::Operator::getSuccessor(unsigned index) const {
+  return successors[index];
+}
+
+unsigned tblgen::Operator::getNumVariadicSuccessors() const {
+  return llvm::count_if(successors,
+                        [](const NamedSuccessor &c) { return c.isVariadic(); });
+}
+
 auto tblgen::Operator::trait_begin() const -> const_trait_iterator {
   return traits.begin();
 }
 auto tblgen::Operator::trait_end() const -> const_trait_iterator {
   return traits.end();
 }
 auto tblgen::Operator::getTraits() const
     -> llvm::iterator_range<const_trait_iterator> {
for (unsigned i = 0, e = resultsDag->getNumArgs(); i < e; ++i) {
     auto *resultDef = dyn_cast<DefInit>(resultsDag->getArg(i));
     if (!resultDef) {
       PrintFatalError(def.getLoc(),
                       Twine("undefined type for result #") + Twine(i));
     }
     results.push_back({name, TypeConstraint(resultDef)});
   }
 
+  // Handle successors
+  auto *successorsDag = def.getValueAsDag("successors");
+  auto *successorsOp = dyn_cast<DefInit>(successorsDag->getOperator());
+  if (!successorsOp || successorsOp->getDef()->getName() != "successor") {
+    PrintFatalError(def.getLoc(),
+                    "'successors' must have 'successor' directive");
+  }
+
+  for (unsigned i = 0, e = successorsDag->getNumArgs(); i < e; ++i) {
+    auto name = successorsDag->getArgNameStr(i);
+    auto *successorInit = dyn_cast<DefInit>(successorsDag->getArg(i));
+    if (!successorInit) {
+      PrintFatalError(def.getLoc(),
+                      Twine("undefined kind for successor #") + Twine(i));
+    }
+    Successor successor(successorInit->getDef());
+
+    // Only support variadic successors if it is the last one for now.
+    if (i != e - 1 && successor.isVariadic())
+      PrintFatalError(def.getLoc(), "only the last successor can be variadic");
+    successors.push_back({name, successor});
+  }
+
   // Create list of traits, skipping over duplicates: appending to lists in
   // tablegen is easy, making them unique less so, so dedupe here.
   if (auto traitList = def.getValueAsListInit("traits")) {
     // This is uniquing based on pointers of the trait.
     SmallPtrSet<const llvm::Init *, 32> traitSet;
     traits.reserve(traitSet.size());
     for (auto traitInit : *traitList) {
       // Keep traits in the same order while skipping over duplicates.

mlir/lib/TableGen/Successor.cpp

+//===- Successor.cpp - Successor class ------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Successor wrapper to simplify using TableGen Record defining a MLIR
+// Successor.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/TableGen/Successor.h"
+#include "mlir/ADT/TypeSwitch.h"
+#include "llvm/TableGen/Record.h"
+
+using namespace mlir;
+using namespace mlir::tblgen;
+
+// Returns true if this successor is variadic.
+bool Successor::isVariadic() const {
+  return def->isSubClassOf("VariadicSuccessor");
+}

mlir/test/Dialect/SPIRV/control-flow-ops.mlir

func @missing_accessor() -> () {
   spv.Branch
   // expected-error @+1 {{expected block name}}
 }
 
 // -----
 
 func @wrong_accessor_count() -> () {
   %true = spv.constant true
-  // expected-error @+1 {{must have exactly one successor}}
+  // expected-error @+1 {{incorrect number of successors: expected 1 but found 2}}
   "spv.Branch"()[^one, ^two] : () -> ()
 ^one:
   spv.Return
 ^two:
   spv.Return
 }
 
 // -----

 ^two:
   spv.Return
 }
 
 // -----
 
 func @wrong_accessor_count() -> () {
   %true = spv.constant true
-  // expected-error @+1 {{must have exactly two successors}}
+  // expected-error @+1 {{incorrect number of successors: expected 2 but found 1}}
   "spv.BranchConditional"(%true)[^one] : (i1) -> ()
 ^one:
   spv.Return
 ^two:
   spv.Return
 }
 
 // -----

mlir/test/lib/TestDialect/TestOps.td

 
 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]> {
-  let arguments = (ins Variadic<AnyType>:$operands);
+  let successors = (successor AnySuccessor:$target);
 }
 
 def AttrSizedOperandOp : TEST_Op<"attr_sized_operands",
                                  [AttrSizedOperandSegments]> {
   let arguments = (ins
     Variadic<I32>:$a,
     Variadic<I32>:$b,
     I32:$c,

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

private:
   void genNamedOperandGetters();
 
   // Generates getters for named results.
   void genNamedResultGetters();
 
   // Generates getters for named regions.
   void genNamedRegionGetters();
 
+  // Generates getters for named successors.
+  void genNamedSuccessorGetters();
+
   // Generates builder methods for the operation.
   void genBuilder();
 
   // Generates the build() method that takes each operand/attribute
   // as a stand-alone parameter.
   void genSeparateArgParamBuilder();
 
   // Generates the build() method that takes each operand/attribute as a
private:
   void genOperandResultVerifier(OpMethodBody &body,
                                 Operator::value_range values,
                                 StringRef valueKind);
 
   // Generates verify statements for regions in the operation.
   // The generated code will be attached to `body`.
   void genRegionVerifier(OpMethodBody &body);
 
+  // Generates verify statements for successors in the operation.
+  // The generated code will be attached to `body`.
+  void genSuccessorVerifier(OpMethodBody &body);
+
   // Generates the traits used by the object.
   void genTraits();
 
   // Generate the OpInterface methods.
   void genOpInterfaceMethods();
 
 private:
   // The TableGen record for this op.
OpEmitter::OpEmitter(const Operator &op)
   genTraits();
   // Generate C++ code for various op methods. The order here determines the
   // methods in the generated file.
   genOpAsmInterface();
   genOpNameGetter();
   genNamedOperandGetters();
   genNamedResultGetters();
   genNamedRegionGetters();
+  genNamedSuccessorGetters();
   genAttrGetters();
   genBuilder();
   genParser();
   genPrinter();
   genVerifier();
   genCanonicalizerDecls();
   genFolderDecls();
   genOpInterfaceMethods();
for (unsigned i = 0; i < numRegions; ++i) {
     const auto &region = op.getRegion(i);
     if (!region.name.empty()) {
       auto &m = opClass.newMethod("Region &", region.name);
       m.body() << formatv("  return this->getOperation()->getRegion({0});", i);
     }
   }
 }
 
+void OpEmitter::genNamedSuccessorGetters() {
+  unsigned numSuccessors = op.getNumSuccessors();
+  for (unsigned i = 0; i < numSuccessors; ++i) {
+    const NamedSuccessor &successor = op.getSuccessor(i);
+    if (successor.name.empty())
+      continue;
+
+    // Generate the accessors for a variadic successor.
+    if (successor.isVariadic()) {
+      // Generate the getter.
+      auto &m = opClass.newMethod("SuccessorRange", successor.name);
+      m.body() << formatv(
+          "  return {std::next(this->getOperation()->successor_begin(), {0}), "
+          "this->getOperation()->successor_end()};",
+          i);
+      continue;
+    }
+
+    // Generate the block getter.
+    auto &m = opClass.newMethod("Block *", successor.name);
+    m.body() << formatv("  return this->getOperation()->getSuccessor({0});", i);
+
+    // Generate the all-operands getter.
+    auto &operandsMethod = opClass.newMethod(
+        "Operation::operand_range", (successor.name + "Operands").str());
+    operandsMethod.body() << formatv(
+        " return this->getOperation()->getSuccessorOperands({0});", i);
+
+    // Generate the individual-operand getter.
+    auto &operandMethod = opClass.newMethod(
+        "Value", (successor.name + "Operand").str(), "unsigned index");
+    operandMethod.body() << formatv(
+        " return this->getOperation()->getSuccessorOperand({0}, index);", i);
+  }
+}
+
 static bool canGenerateUnwrappedBuilder(Operator &op) {
   // If this op does not have native attributes at all, return directly to avoid
   // redefining builders.
   if (op.getNumNativeAttributes() == 0)
     return false;
 
   bool canGenerate = false;
   // We are generating builders that take raw values for attributes. We need to
if (numVariadicResults == 0 || numNonVariadicResults != 0)
     body << "  assert(resultTypes.size()"
          << (numVariadicResults != 0 ? " >= " : " == ") << numNonVariadicResults
          << "u && \"mismatched number of return types\");\n";
   body << "  " << builderOpState << ".addTypes(resultTypes);\n";
 
   // Generate builder that infers type too.
   // TODO(jpienaar): Subsume this with general checking if type can be infered
   // automatically.
-  // TODO(jpienaar): Expand to handle regions.
-  if (op.getTrait("InferTypeOpInterface::Trait") && op.getNumRegions() == 0)
+  // TODO(jpienaar): Expand to handle regions and successors.
+  if (op.getTrait("InferTypeOpInterface::Trait") && op.getNumRegions() == 0 &&
+      op.getNumSuccessors() == 0)
     genInferedTypeCollectiveParamBuilder();
 }
 
 void OpEmitter::buildParamList(std::string &paramList,
                                SmallVectorImpl<std::string> &resultTypeNames,
                                TypeParamKind typeParamKind,
                                AttrParamKind attrParamKind) {
   resultTypeNames.clear();
if (argument.is<tblgen::NamedTypeConstraint *>()) {
           paramList.append("\"");
         paramList.append(std::string(attr.getDefaultValue()));
         if (isString)
           paramList.append("\"");
       }
       ++numAttrs;
     }
   }
+
+  /// Insert parameters for the block and operands for each successor.
+  const char *variadicSuccCode =
+      ", ArrayRef<Block *> {0}, ArrayRef<ValueRange> {0}Operands";
+  const char *succCode = ", Block *{0}, ValueRange {0}Operands";
+  for (const NamedSuccessor &namedSuccessor : op.getSuccessors()) {
+    if (namedSuccessor.isVariadic())
+      paramList += llvm::formatv(variadicSuccCode, namedSuccessor.name).str();
+    else
+      paramList += llvm::formatv(succCode, namedSuccessor.name).str();
+  }
 }
 
 void OpEmitter::genCodeForAddingArgAndRegionForBuilder(OpMethodBody &body,
                                                        bool isRawValueAttr) {
-  // Push all operands to the result
+  // Push all operands to the result.
   for (int i = 0, e = op.getNumOperands(); i < e; ++i) {
     body << "  " << builderOpState << ".addOperands(" << getArgumentName(op, i)
          << ");\n";
   }
 
-  // Push all attributes to the result
+  // Push all attributes to the result.
   for (const auto &namedAttr : op.getAttributes()) {
     auto &attr = namedAttr.attr;
     if (!attr.isDerivedAttr()) {
       bool emitNotNullCheck = attr.isOptional();
       if (emitNotNullCheck) {
         body << formatv("  if ({0}) ", namedAttr.name) << "{\n";
       }
       if (isRawValueAttr && canUseUnwrappedRawValue(attr)) {
if (!attr.isDerivedAttr()) {
                         namedAttr.name);
       }
       if (emitNotNullCheck) {
         body << "  }\n";
       }
     }
   }
 
-  // Create the correct number of regions
+  // Create the correct number of regions.
   if (int numRegions = op.getNumRegions()) {
     for (int i = 0; i < numRegions; ++i)
       body << "  (void)" << builderOpState << ".addRegion();\n";
   }
+
+  // Push all successors to the result.
+  for (const NamedSuccessor &namedSuccessor : op.getSuccessors()) {
+    if (namedSuccessor.isVariadic()) {
+      body << formatv("  for (int i = 0, e = {1}.size(); i != e; ++i)\n"
+                      "    {0}.addSuccessor({1}[i], {1}Operands[i]);\n",
+                      builderOpState, namedSuccessor.name);
+      continue;
+    }
+
+    body << formatv("  {0}.addSuccessor({1}, {1}Operands);\n", builderOpState,
+                    namedSuccessor.name);
+  }
 }
 
 void OpEmitter::genCanonicalizerDecls() {
   if (!def.getValueAsBit("hasCanonicalizer"))
     return;
 
   const char *const params =
       "OwningRewritePatternList &results, MLIRContext *context";
if (auto *t = dyn_cast<tblgen::PredOpTrait>(&trait)) {
       body << tgfmt("  if (!($0)) {\n    "
                     "return emitOpError(\"failed to verify that $1\");\n  }\n",
                     &verifyCtx, tgfmt(t->getPredTemplate(), &verifyCtx),
                     t->getDescription());
     }
   }
 
   genRegionVerifier(body);
+  genSuccessorVerifier(body);
 
   if (hasCustomVerify) {
     FmtContext fctx;
     fctx.addSubst("cppClass", opClass.getClassName());
     auto printer = codeInit->getValue().ltrim().rtrim(" \t\v\f\r");
     body << "  " << tgfmt(printer, &fctx);
   } else {
     body << "  return mlir::success();\n";
for (unsigned i = 0; i < numRegions; ++i) {
 
     body << formatv("  if (!({0})) {\n    "
                     "return emitOpError(\"region {1} failed to verify "
                     "constraint: {2}\");\n  }\n",
                     constraint, name, region.constraint.getDescription());
   }
 }
 
+void OpEmitter::genSuccessorVerifier(OpMethodBody &body) {
+  unsigned numSuccessors = op.getNumSuccessors();
+
+  const char *checkSuccessorSizeCode = R"(
+  if (this->getOperation()->getNumSuccessors() {0} {1}) {
+    return emitOpError("has incorrect number of successors: expected{2} {1}"
+                       " but found ")
+             << this->getOperation()->getNumSuccessors();
+  }
+  )";
+
+  // Verify this op has the correct number of successors.
+  if (unsigned numVariadicSuccessors = op.getNumVariadicSuccessors()) {
+    body << formatv(checkSuccessorSizeCode, "<",
+                    numSuccessors - numVariadicSuccessors, " at least");
+  } else {
+    body << formatv(checkSuccessorSizeCode, "!=", numSuccessors, "");
+  }
+
+  body << "{\n";
+  body << "    unsigned index = 0; (void)index;\n";
+
+  for (unsigned i = 0; i < numSuccessors; ++i) {
+    const auto &successor = op.getSuccessor(i);
+    if (successor.constraint.getPredicate().isNull())
+      continue;
+
+    body << "    for (Block *successor : ";
+    body << formatv(successor.isVariadic() ? "{0}()"
+                                           : "ArrayRef<Block *>({0}())",
+                    successor.name);
+    body << ") {\n";
+    auto constraint = tgfmt(successor.constraint.getConditionTemplate(),
+                            &verifyCtx.withSelf("successor"))
+                          .str();
+
+    body << formatv(
+        "      (void)successor;\n"
+        "      if (!({0})) {\n        "
+        "return emitOpError(\"successor #\") << index << \"('{2}') failed to "
+        "verify constraint: {3}\";\n      }\n",
+        constraint, i, successor.name, successor.constraint.getDescription());
+    body << "    }\n";
+  }
+  body << "  }\n";
+}
+
 void OpEmitter::genTraits() {
   int numResults = op.getNumResults();
   int numVariadicResults = op.getNumVariadicResults();
 
   // Add return size trait.
   if (numVariadicResults != 0) {
     if (numResults == numVariadicResults)
       opClass.addTrait("OpTrait::VariadicResults");
else if (auto opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait))
       opClass.addTrait(opTrait->getTrait());
   }
 
   // Add variadic size trait and normal op traits.
   int numOperands = op.getNumOperands();
   int numVariadicOperands = op.getNumVariadicOperands();
 
   // Add operand size trait.
-  if (numVariadicOperands != 0) {
+  // Note: Successor operands are also included in the operation's operand list,
+  // so we always need to use VariadicOperands in the presence of successors.
+  if (numVariadicOperands != 0 || op.getNumSuccessors()) {
     if (numOperands == numVariadicOperands)
       opClass.addTrait("OpTrait::VariadicOperands");
     else
       opClass.addTrait("OpTrait::AtLeastNOperands<" +
                        Twine(numOperands - numVariadicOperands) + ">::Impl");
   } else {
     switch (numOperands) {
     case 0:
       opClass.addTrait("OpTrait::ZeroOperands");