[mlir] Add traits for verifying the number of successors and providing relevant accessors.

[rriddle]

Summary

This allows for simplifying OpDefGen, as well providing specializing accessors for the different successor counts. This mirrors the existing traits for operands and results.

[antiagainst]

Nice!

[rriddle]

Fix formatting.

[rriddle]

Rebase

mlir/include/mlir/IR/OpDefinition.h

 LogicalResult verifySameOperandsAndResultShape(Operation *op);
 LogicalResult verifySameOperandsElementType(Operation *op);
 LogicalResult verifySameOperandsAndResultElementType(Operation *op);
 LogicalResult verifySameOperandsAndResultType(Operation *op);
 LogicalResult verifyResultsAreBoolLike(Operation *op);
 LogicalResult verifyResultsAreFloatLike(Operation *op);
 LogicalResult verifyResultsAreSignlessIntegerLike(Operation *op);
 LogicalResult verifyIsTerminator(Operation *op);
+LogicalResult verifyZeroSuccessor(Operation *op);
+LogicalResult verifyOneSuccessor(Operation *op);
+LogicalResult verifyNSuccessors(Operation *op, unsigned numSuccessors);
+LogicalResult verifyAtLeastNSuccessors(Operation *op, unsigned numSuccessors);
 LogicalResult verifyOperandSizeAttr(Operation *op, StringRef sizeAttrName);
 LogicalResult verifyResultSizeAttr(Operation *op, StringRef sizeAttrName);
 } // namespace impl
 
 /// Helper class for implementing traits.  Clients are not expected to interact
 /// with this directly, so its members are all protected.
 template <typename ConcreteType, template <typename> class TraitType>
 class TraitBase {
protected:
   /// Provide default implementations of trait hooks.  This allows traits to
   /// provide exactly the overrides they care about.
   static LogicalResult verifyTrait(Operation *op) { return success(); }
   static AbstractOperation::OperationProperties getTraitProperties() {
     return 0;
   }
 };
 
+//===----------------------------------------------------------------------===//
+// Operand Traits
+
 namespace detail {
 /// Utility trait base that provides accessors for derived traits that have
 /// multiple operands.
 template <typename ConcreteType, template <typename> class TraitType>
 struct MultiOperandTraitBase : public TraitBase<ConcreteType, TraitType> {
   using operand_iterator = Operation::operand_iterator;
   using operand_range = Operation::operand_range;
   using operand_type_iterator = Operation::operand_type_iterator;

 };
 
 /// This class provides the API for ops which have an unknown number of
 /// SSA operands.
 template <typename ConcreteType>
 class VariadicOperands
     : public detail::MultiOperandTraitBase<ConcreteType, VariadicOperands> {};
 
+//===----------------------------------------------------------------------===//
+// Result Traits
+
 /// This class provides return value APIs for ops that are known to have
 /// zero results.
 template <typename ConcreteType>
 class ZeroResult : public TraitBase<ConcreteType, ZeroResult> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifyZeroResult(op);
   }

 };
 
 /// This class provides the API for ops which have an unknown number of
 /// results.
 template <typename ConcreteType>
 class VariadicResults
     : public detail::MultiResultTraitBase<ConcreteType, VariadicResults> {};
 
+//===----------------------------------------------------------------------===//
+// Terminator Traits
+
+/// This class provides the API for ops that are known to be terminators.
+template <typename ConcreteType>
+class IsTerminator : public TraitBase<ConcreteType, IsTerminator> {
+public:
+  static AbstractOperation::OperationProperties getTraitProperties() {
+    return static_cast<AbstractOperation::OperationProperties>(
+        OperationProperty::Terminator);
+  }
+  static LogicalResult verifyTrait(Operation *op) {
+    return impl::verifyIsTerminator(op);
+  }
+
+  unsigned getNumSuccessorOperands(unsigned index) {
+    return this->getOperation()->getNumSuccessorOperands(index);
+  }
+};
+
+/// This class provides verification for ops that are known to have zero
+/// successors.
+template <typename ConcreteType>
+class ZeroSuccessor : public TraitBase<ConcreteType, ZeroSuccessor> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    return impl::verifyZeroSuccessor(op);
+  }
+};
+
+namespace detail {
+/// Utility trait base that provides accessors for derived traits that have
+/// multiple successors.
+template <typename ConcreteType, template <typename> class TraitType>
+struct MultiSuccessorTraitBase : public TraitBase<ConcreteType, TraitType> {
+  using succ_iterator = Operation::succ_iterator;
+  using succ_range = SuccessorRange;
+
+  /// Return the number of successors.
+  unsigned getNumSuccessors() {
+    return this->getOperation()->getNumSuccessors();
+  }
+
+  /// Return the successor at `index`.
+  Block *getSuccessor(unsigned i) {
+    return this->getOperation()->getSuccessor(i);
+  }
+
+  /// Set the successor at `index`.
+  void setSuccessor(Block *block, unsigned i) {
+    return this->getOperation()->setSuccessor(block, i);
+  }
+
+  /// Successor iterator access.
+  succ_iterator succ_begin() { return this->getOperation()->succ_begin(); }
+  succ_iterator succ_end() { return this->getOperation()->succ_end(); }
+  succ_range getSuccessors() { return this->getOperation()->getSuccessors(); }
+};
+} // end namespace detail
+
+/// This class provides APIs for ops that are known to have a single successor.
+template <typename ConcreteType>
+class OneSuccessor : public TraitBase<ConcreteType, OneSuccessor> {
+public:
+  Block *getSuccessor() { return this->getOperation()->getSuccessor(0); }
+  void setSuccessor(Block *succ) {
+    this->getOperation()->setSuccessor(succ, 0);
+  }
+
+  static LogicalResult verifyTrait(Operation *op) {
+    return impl::verifyOneSuccessor(op);
+  }
+};
+
+/// This class provides the API for ops that are known to have a specified
+/// number of successors.
+template <unsigned N>
+class NSuccessors {
+public:
+  static_assert(N > 1, "use ZeroSuccessor/OneSuccessor for N < 2");
+
+  template <typename ConcreteType>
+  class Impl : public detail::MultiSuccessorTraitBase<ConcreteType,
+                                                      NSuccessors<N>::Impl> {
+  public:
+    static LogicalResult verifyTrait(Operation *op) {
+      return impl::verifyNSuccessors(op, N);
+    }
+  };
+};
+
+/// This class provides APIs for ops that are known to have at least a specified
+/// number of successors.
+template <unsigned N>
+class AtLeastNSuccessors {
+public:
+  template <typename ConcreteType>
+  class Impl
+      : public detail::MultiSuccessorTraitBase<ConcreteType,
+                                               AtLeastNSuccessors<N>::Impl> {
+  public:
+    static LogicalResult verifyTrait(Operation *op) {
+      return impl::verifyAtLeastNSuccessors(op, N);
+    }
+  };
+};
+
+/// This class provides the API for ops which have an unknown number of
+/// successors.
+template <typename ConcreteType>
+class VariadicSuccessors
+    : public detail::MultiSuccessorTraitBase<ConcreteType, VariadicSuccessors> {
+};
+
+//===----------------------------------------------------------------------===//
+// Misc Traits
+
 /// This class provides verification for ops that are known to have the same
 /// operand shape: all operands are scalars, vectors/tensors of the same
 /// shape.
 template <typename ConcreteType>
 class SameOperandsShape : public TraitBase<ConcreteType, SameOperandsShape> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifySameOperandsShape(op);

 template <typename ConcreteType>
 class SameTypeOperands : public TraitBase<ConcreteType, SameTypeOperands> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifySameTypeOperands(op);
   }
 };
 
-/// This class provides the API for ops that are known to be terminators.
-template <typename ConcreteType>
-class IsTerminator : public TraitBase<ConcreteType, IsTerminator> {
-public:
-  static AbstractOperation::OperationProperties getTraitProperties() {
-    return static_cast<AbstractOperation::OperationProperties>(
-        OperationProperty::Terminator);
-  }
-  static LogicalResult verifyTrait(Operation *op) {
-    return impl::verifyIsTerminator(op);
-  }
-
-  unsigned getNumSuccessors() {
-    return this->getOperation()->getNumSuccessors();
-  }
-  unsigned getNumSuccessorOperands(unsigned index) {
-    return this->getOperation()->getNumSuccessorOperands(index);
-  }
-
-  Block *getSuccessor(unsigned index) {
-    return this->getOperation()->getSuccessor(index);
-  }
-
-  void setSuccessor(Block *block, unsigned index) {
-    return this->getOperation()->setSuccessor(block, index);
-  }
-
-  void addSuccessorOperand(unsigned index, Value value) {
-    return this->getOperation()->addSuccessorOperand(index, value);
-  }
-  void addSuccessorOperands(unsigned index, ArrayRef<Value> values) {
-    return this->getOperation()->addSuccessorOperand(index, values);
-  }
-};
-
 /// This class provides the API for ops that are known to be isolated from
 /// above.
 template <typename ConcreteType>
 class IsIsolatedFromAbove
     : public TraitBase<ConcreteType, IsIsolatedFromAbove> {
 public:
   static AbstractOperation::OperationProperties getTraitProperties() {
     return static_cast<AbstractOperation::OperationProperties>(

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

 /// Returns true if the given `srcBlock` contains only one `spv.Branch` to the
 /// given `dstBlock`.
 static inline bool hasOneBranchOpTo(Block &srcBlock, Block &dstBlock) {
   // Check that there is only one op in the `srcBlock`.
   if (!has_single_element(srcBlock))
     return false;
 
   auto branchOp = dyn_cast<spirv::BranchOp>(srcBlock.back());
-  return branchOp && branchOp.getSuccessor(0) == &dstBlock;
+  return branchOp && branchOp.getSuccessor() == &dstBlock;
 }
 
 static LogicalResult verify(spirv::LoopOp loopOp) {
   auto *op = loopOp.getOperation();
 
   // We need to verify that the blocks follow the following layout:
   //
   //                     +-------------+

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

PatternMatchResult matchAndRewrite(BranchOp op,
     // Merge the successor into the current block and erase the branch.
     rewriter.mergeBlocks(succ, opParent, op.getOperands());
     rewriter.eraseOp(op);
     return matchSuccess();
   }
 };
 } // end anonymous namespace.
 
-Block *BranchOp::getDest() { return getSuccessor(0); }
+Block *BranchOp::getDest() { return getSuccessor(); }
 
-void BranchOp::setDest(Block *block) { return setSuccessor(block, 0); }
+void BranchOp::setDest(Block *block) { return setSuccessor(block); }
 
 void BranchOp::eraseOperand(unsigned index) {
   getOperation()->eraseSuccessorOperand(0, index);
 }
 
 void BranchOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                            MLIRContext *context) {
   results.insert<SimplifyBrToBlockWithSinglePred>(context);

mlir/lib/IR/Operation.cpp

for (auto opType : op->getOperandTypes()) {
     if (getElementTypeOrSelf(opType) != elementType ||
         failed(verifyCompatibleShape(opType, type)))
       return op->emitOpError()
              << "requires the same type for all operands and results";
   }
   return success();
 }
 
+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 "
for (unsigned i = 0, e = op->getNumSuccessors(); i != e; ++i) {
     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::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");
-
-  // Verify the state of the successor blocks.
-  if (op->getNumSuccessors() != 0 && failed(verifyTerminatorSuccessors(op)))
-    return failure();
+LogicalResult OpTrait::impl::verifyZeroSuccessor(Operation *op) {
+  if (op->getNumSuccessors() != 0) {
+    return op->emitOpError("requires 0 successors but found ")
+           << op->getNumSuccessors();
+  }
   return success();
 }
 
+LogicalResult OpTrait::impl::verifyOneSuccessor(Operation *op) {
+  if (op->getNumSuccessors() != 1) {
+    return op->emitOpError("requires 1 successor but found ")
+           << op->getNumSuccessors();
+  }
+  return verifyTerminatorSuccessors(op);
+}
+LogicalResult OpTrait::impl::verifyNSuccessors(Operation *op,
+                                               unsigned numSuccessors) {
+  if (op->getNumSuccessors() != numSuccessors) {
+    return op->emitOpError("requires ")
+           << numSuccessors << " successors but found "
+           << op->getNumSuccessors();
+  }
+  return verifyTerminatorSuccessors(op);
+}
+LogicalResult OpTrait::impl::verifyAtLeastNSuccessors(Operation *op,
+                                                      unsigned numSuccessors) {
+  if (op->getNumSuccessors() < numSuccessors) {
+    return op->emitOpError("requires at least ")
+           << numSuccessors << " successors but found "
+           << op->getNumSuccessors();
+  }
+  return verifyTerminatorSuccessors(op);
+}
+
 LogicalResult OpTrait::impl::verifyResultsAreBoolLike(Operation *op) {
   for (auto resultType : op->getResultTypes()) {
     auto elementType = getTensorOrVectorElementType(resultType);
     bool isBoolType = elementType.isInteger(1);
     if (!isBoolType)
       return op->emitOpError() << "requires a bool result type";
   }
 

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

for (auto operand : lookupValues(lpOp.getOperands())) {
         lpi->addClause(constOperand);
     }
     return success();
   }
 
   // Emit branches.  We need to look up the remapped blocks and ignore the block
   // arguments that were transformed into PHI nodes.
   if (auto brOp = dyn_cast<LLVM::BrOp>(opInst)) {
-    builder.CreateBr(blockMapping[brOp.getSuccessor(0)]);
+    builder.CreateBr(blockMapping[brOp.getSuccessor()]);
     return success();
   }
   if (auto condbrOp = dyn_cast<LLVM::CondBrOp>(opInst)) {
     builder.CreateCondBr(valueMapping.lookup(condbrOp.getOperand(0)),
                          blockMapping[condbrOp.getSuccessor(0)],
                          blockMapping[condbrOp.getSuccessor(1)]);
     return success();
   }

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

func @branch_argument() -> () {
   spv.Branch ^next(%zero, %zero: i32, i32)
 ^next(%arg0: i32, %arg1: i32):
   spv.Return
 }
 
 // -----
 
 func @missing_accessor() -> () {
-  // expected-error @+1 {{has incorrect number of successors: expected 1 but found 0}}
+  // expected-error @+1 {{requires 1 successor but found 0}}
   spv.Branch
 }
 
 // -----
 
 func @wrong_accessor_count() -> () {
   %true = spv.constant true
-  // expected-error @+1 {{incorrect number of successors: expected 1 but found 2}}
+  // expected-error @+1 {{requires 1 successor 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 {{incorrect number of successors: expected 2 but found 1}}
+  // expected-error @+1 {{requires 2 successors but found 1}}
   "spv.BranchConditional"(%true)[^one] : (i1) -> ()
 ^one:
   spv.Return
 ^two:
   spv.Return
 }
 
 // -----

mlir/test/mlir-tblgen/op-decl.td

 // CHECK:   AOpOperandAdaptor(ArrayRef<Value> values);
 // CHECK:   ArrayRef<Value> getODSOperands(unsigned index);
 // CHECK:   Value a();
 // CHECK:   ArrayRef<Value> b();
 // CHECK: private:
 // CHECK:   ArrayRef<Value> tblgen_operands;
 // CHECK: };
 
-// CHECK: class AOp : public Op<AOp, OpTrait::AtLeastNResults<1>::Impl, OpTrait::HasNoSideEffect, OpTrait::AtLeastNOperands<1>::Impl
+// CHECK: class AOp : public Op<AOp, OpTrait::AtLeastNResults<1>::Impl, OpTrait::ZeroSuccessor, OpTrait::HasNoSideEffect, OpTrait::AtLeastNOperands<1>::Impl
 // CHECK: public:
 // CHECK:   using Op::Op;
 // CHECK:   using OperandAdaptor = AOpOperandAdaptor;
 // CHECK:   static StringRef getOperationName();
 // CHECK:   Operation::operand_range getODSOperands(unsigned index);
 // CHECK:   Value a();
 // CHECK:   Operation::operand_range b();
 // CHECK:   Operation::result_range getODSResults(unsigned index);

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

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.
-  unsigned numVariadicSuccessors = op.getNumVariadicSuccessors();
-  if (numVariadicSuccessors == 0) {
-    body << formatv(checkSuccessorSizeCode, "!=", numSuccessors, "");
-  } else if (numVariadicSuccessors != numSuccessors) {
-    body << formatv(checkSuccessorSizeCode, "<",
-                    numSuccessors - numVariadicSuccessors, " at least");
-  }
-
   // If we have no successors, there is nothing more to do.
+  unsigned numSuccessors = op.getNumSuccessors();
   if (numSuccessors == 0)
     return;
 
   body << "{\n";
   body << "    unsigned index = 0; (void)index;\n";
 
   for (unsigned i = 0; i < numSuccessors; ++i) {
     const auto &successor = op.getSuccessor(i);
body << formatv(
         "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");
+/// Add a size count trait to the given operation class.
+static void addSizeCountTrait(OpClass &opClass, StringRef traitKind,
+                              int numNonVariadic, int numVariadic) {
+  if (numVariadic != 0) {
+    if (numNonVariadic == numVariadic)
+      opClass.addTrait("OpTrait::Variadic" + traitKind + "s");
     else
-      opClass.addTrait("OpTrait::AtLeastNResults<" +
-                       Twine(numResults - numVariadicResults) + ">::Impl");
-  } else {
-    switch (numResults) {
-    case 0:
-      opClass.addTrait("OpTrait::ZeroResult");
-      break;
-    case 1:
-      opClass.addTrait("OpTrait::OneResult");
-      break;
-    default:
-      opClass.addTrait("OpTrait::NResults<" + Twine(numResults) + ">::Impl");
-      break;
-    }
-  }
+      opClass.addTrait("OpTrait::AtLeastN" + traitKind + "s<" +
+                       Twine(numNonVariadic - numVariadic) + ">::Impl");
+    return;
+  }
+  switch (numNonVariadic) {
+    opClass.addTrait("OpTrait::Zero" + traitKind);
+    opClass.addTrait("OpTrait::One" + traitKind);
+    opClass.addTrait("OpTrait::N" + traitKind + "s<" + Twine(numNonVariadic) +
+                     ">::Impl");
 
+void OpEmitter::genTraits() {
+  int numResults = op.getNumResults();
+  int numVariadicResults = op.getNumVariadicResults();
+
+  // Add return size trait.
+  addSizeCountTrait(opClass, "Result", numResults, numVariadicResults);
+
+  // Add successor size trait.
+  unsigned numSuccessors = op.getNumSuccessors();
+  unsigned numVariadicSuccessors = op.getNumVariadicSuccessors();
+  addSizeCountTrait(opClass, "Successor", numSuccessors, numVariadicSuccessors);
+
+  // Add the native and interface traits.
   for (const auto &trait : op.getTraits()) {
     if (auto opTrait = dyn_cast<tblgen::NativeOpTrait>(&trait))
       opClass.addTrait(opTrait->getTrait());
     else if (auto opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait))
       opClass.addTrait(opTrait->getTrait());
   }
 
   // Add variadic size trait and normal op traits.