[MLIR] Add type checking capability to RegionBranchOpInterface

[jurahul]

Summary

• Add function verifyTypes that Op's can call to do type checking verification along the control flow edges described the Op's RegionBranchOpInterface.
• We cannot rely on the verify methods on the OpInterface because this interface might in turn invoke methods in other interfaces when the Op is still under verification causing code crash (example seen was an invalid scf.for can invoke SingleBlockImplicitTerminator::getBody() as a part of implementing the RegionBranchOpInterface and crash).

[jurahul]

Hi all,

This is outcome of my attempt to implement the OpResultsMatchRegionResults trait that was discussed on MLIR discourse (https://llvm.discourse.group/t/rfc-extending-singleblockimplicitterminator-to-support-additional-type-checking/1278). Since the trait was implemented entirely on top of the RegionBranchOpInterface, I felt appropriate to just fold it into the interface and then add type checking for not just region results -> parent op results flow but for all of the control flow described by that interface. As mentioned in the description, doing this verification during the RegionBranchOpInterface verification proved impossible as the may not be fully verified, hence I just provided a verify method that Op's that want this type checking are supposed to call towards the end their verification when all the basic structural properties have been verified.

Thanks.

[jurahul]

Fix clang-tidy warnings

[jurahul]

Full sentence comments

[jurahul]

Actually fail verifyTypes() if verifyTypesAlongAllSuccessors() fails

[jurahul]

mlir/lib/Dialect/SCF/SCF.cpp

501

This represents a slight interface change, where empty constant operands is considered to be unknown.

[ftynse]

Deferring to @rriddle for approval.

mlir/include/mlir/Dialect/SCF/SCFOps.td

422

Nit: add one more space before ParentOneOf

441

You can just omit this variable then.

mlir/lib/Dialect/SCF/SCF.cpp

600

No need to duplicate the operation name (yield), it will be printed by emitOpError(). Consider something like "not allowed to have operands inside '" << getOperationName() << "'", which will also be more future proof to renamings

869

Don't drop this comment, there is still code related to YieldOp below.

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

93–102

It feels like this would have been more readable using llvm::raw_string_ostream and operator <<. Also, llvm::SmallString<256> is probably more efficient here.

108

Why the leading colon?

133

Nit: outline the lambda into a variable and save horizontal space

149

Should this rather be an assertion? E.g., do we expect this to ever trigger for valid IR?

[jurahul]

Address review comments.

mlir/include/mlir/Dialect/SCF/SCFOps.td

441

SCF_Op defines let verifier = [{ return ::verify(*this); }]; and YieldOp derived from that, so I have to explicitly null it out.

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

93–102

I changed this to not require building the string but to print it directly to the InFlightDiagnostic object.

108

Removed.

118

Removed leading colon here as well.

149

Can regions attached to an Op be optional? For example, the else region for scf.if? Looks like the region is always present, but its empty, as opposed to the region itself being not present. In which case, I agree that this can be an assert.

[ftynse]

mlir/include/mlir/Dialect/SCF/SCFOps.td

441

Let's have a comment that explains exactly that: "Override default verifier logic, no custom verification necessary.".

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

149

I think it's technically possible to construct such operation, but we don't have any upstream. Anyway, "verifyTypesAlongControlFlowEdges" sounds like a wrong place to put this check if we need it.

[jurahul]

Address review comments

[jurahul]

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

149

Agreed. I have made that an assert.

[rriddle]

Thanks, took a pass over.

mlir/include/mlir/Dialect/SCF/SCFOps.td

443

nit: The common way to do this is let verifier = ?;

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

45

nit: Use /// for top level comments.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

137

Can you put this in the trait verifier(let verify = [{...}];) instead? That would remove the need to call it from each of the operation verifiers.

138

nit: Drop the mlir::

mlir/lib/Dialect/SCF/SCF.cpp

141

nit: return RegionBranchOpInterface::verifyTypes(op);

502

When is this empty?

599

nit: Add a space after inside

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

76

nit: Please use /// for top-level comments.

86

nit: Please add a parameter comment(/*...=*/) for the {}

88

nit: Spell out auto here.

89

nit: I think it would be cleaner to just break this:

Optional<unsigned> succRegionNo;
if (Region *succRegion = succ. getSuccessor())
  succRegionNo = ...;

104

nit: The .getStringRef shouldn't be necessary here and above.

139

nit: Please add braces here, the body is no longer trivial.

168

nit: Spell out auto here.

187

nit: All

189

Why not use a default constructed TypeRange here?

[jurahul]

Address review comments

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

137

I tried that first, but it did not work.

"We cannot rely on the verify methods on the OpInterface because this interface might in turn invoke methods in other interfaces when the Op is still under verification causing code crash (example seen was an invalid scf.for can invoke SingleBlockImplicitTerminator::getBody() as a part of implementing the RegionBranchOpInterface and crash)."

mlir/lib/Dialect/SCF/SCF.cpp

141

I thought I saw some similar code not doing this (parseIfOp). But I see that parseOptionalAttrDict is not returning a LogicalResult directly. Changed.

417

Did the same thing here as well.

502

When exercised from the verifier (specifically in verifyTypesAlongAllEdges). No constant values are know, so with this I avoid creating an array of null Attribute values of the size of the number of inputs etc. As I said below (line numbers got moved) this is a slight interface change to interpret an empty operands value as all operands unknown.

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

502

IMO we should have the same constraint here as the fold methods, i.e. operands always has the same number of elements as the operands to the operation. The verify method could just initialize them all to Attribute().

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

157

You can use a static_assert in the verify method of the interface for this(something like !OpType::hasTrait<OpTrait::ZeroRegion>()). This is fairly common practice for traits/interfaces.

[jurahul]

mlir/lib/Dialect/SCF/SCF.cpp

502

That sounds fine as well. I'll update the change accordingly.

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

157

Agreed. I'll change that.

[jurahul]

Address River's comments.

[jurahul]

Adopt Region::getNumArguments()

[jurahul]

Also undo changes to IfOp::getSuccessorRegions() since we are now passing an array of null attributes during verification.

mlir/include/mlir/Dialect/SCF/SCFOps.td

let description = [{
     scf.reduce.return %res : f32
     ```
   }];
 
   let arguments = (ins AnyType:$result);
   let assemblyFormat = "$result attr-dict `:` type($result)";
 }
 
-def YieldOp : SCF_Op<"yield", [NoSideEffect, ReturnLike, Terminator]> {
+def YieldOp : SCF_Op<"yield", [NoSideEffect, ReturnLike, Terminator,
+                               ParentOneOf<["IfOp, ForOp", "ParallelOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "scf.yield" yields an SSA value from the SCF dialect op region and
     terminates the regions. The semantics of how the values are yielded is
     defined by the parent operation.
     If "scf.yield" has any operands, the operands must match the parent
     operation's results.
     If the parent operation defines no values, then the "scf.yield" may be
     left out in the custom syntax and the builders will insert one implicitly.
     Otherwise, it has to be present in the syntax to indicate which values are
     yielded.
   }];
 
   let arguments = (ins Variadic<AnyType>:$results);
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result",
               [{ /* nothing to do */ }]>
   ];
+  // Override default verifier (defined in SCF_Op), no custom verification
+  // needed.
+  let verifier = ?;
 }
 #endif // MLIR_DIALECT_SCF_SCFOPS

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

 
 #ifndef MLIR_INTERFACES_CONTROLFLOWINTERFACES_H
 #define MLIR_INTERFACES_CONTROLFLOWINTERFACES_H
 
 #include "mlir/IR/OpDefinition.h"
 
 namespace mlir {
 class BranchOpInterface;
+class RegionBranchOpInterface;
 
 //===----------------------------------------------------------------------===//
 // BranchOpInterface
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// Return the `BlockArgument` corresponding to operand `operandIndex` in some
 /// successor if `operandIndex` is within the range of `operands`, or None if
 /// `operandIndex` isn't a successor operand index.
 Optional<BlockArgument>
 getBranchSuccessorArgument(Optional<OperandRange> operands,
                            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
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
 //===----------------------------------------------------------------------===//
 
+namespace detail {
+/// Verify that types match along control flow edges described the given op.
+LogicalResult verifyTypesAlongControlFlowEdges(Operation *op);
+} //  namespace detail
+
 /// This class represents a successor of a region. A region successor can either
 /// be another region, or the parent operation. If the successor is a region,
-/// this class accepts the destination region, as well as a set of arguments
+/// this class represents the destination region, as well as a set of arguments
 /// from that region that will be populated by values from the current region.
-/// If the successor is the parent operation, this class accepts an optional set
-/// of results that will be populated by values from the current region.
+/// If the successor is the parent operation, this class represents an optional
+/// set of results that will be populated by values from the current region.
+///
+/// This interface assumes that the values from the current region that are used
+/// to populate the successor inputs are the operands of the return-like
+/// terminator operations in the blocks within this region.
 class RegionSuccessor {
 public:
   /// Initialize a successor that branches to another region of the parent
   /// operation.
   RegionSuccessor(Region *region, Block::BlockArgListType regionInputs = {})
       : region(region), inputs(regionInputs) {}
   /// Initialize a successor that branches back to/out of the parent operation.
   RegionSuccessor(Optional<Operation::result_range> results = {})
       : region(nullptr), inputs(results ? ValueRange(*results) : ValueRange()) {
   }
 
   /// Return the given region successor. Returns nullptr if the successor is the
   /// parent operation.
   Region *getSuccessor() const { return region; }
 
+  /// Return true if the successor is the parent operation.
+  bool isParent() const { return region == nullptr; }
+
   /// Return the inputs to the successor that are remapped by the exit values of
   /// the current region.
   ValueRange getSuccessorInputs() const { return inputs; }
 
 private:
   Region *region;
   ValueRange inputs;
 };

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

let description = [{
     branching behavior between held regions, i.e. this interface allows for
     expressing control flow information for region holding operations.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
         Returns the operands of this operation used as the entry arguments when
-        entering the region at `index`, which was specified as a successor by
-        `getSuccessorRegions`. These operands should correspond 1-1 with the
-        successor inputs specified in `getSuccessorRegions`, and may corre
+        entering the region at `index`, which was specified as a successor of this
+        operation by `getSuccessorRegions`. These operands should correspond 1-1
+        with the successor inputs specified in `getSuccessorRegions`.
       }],
       "OperandRange", "getSuccessorEntryOperands",
       (ins "unsigned":$index), [{}], /*defaultImplementation=*/[{
         auto operandEnd = this->getOperation()->operand_end();
         return OperandRange(operandEnd, operandEnd);
       }]
     >,
     InterfaceMethod<[{
InterfaceMethod<[{
         are executed after having executed another region of the parent op. The
         successor region must be non-empty.
       }],
       "void", "getSuccessorRegions",
       (ins "Optional<unsigned>":$index, "ArrayRef<Attribute>":$operands,
            "SmallVectorImpl<RegionSuccessor> &":$regions)
     >
   ];
+
+  let verify = [{
+    static_assert(!ConcreteOpType::template hasTrait<OpTrait::ZeroRegion>(),
+                  "expected operation to have non-zero regions");
+    return success();
+  }];
+
+  let extraClassDeclaration = [{
+    /// Verify types along control flow edges described by this interface.
+    static LogicalResult verifyTypes(Operation *op) {
+      return detail::verifyTypesAlongControlFlowEdges(op);
+    }
+  }];
 }
 
 //===----------------------------------------------------------------------===//
 // ControlFlow Traits
 //===----------------------------------------------------------------------===//
 
 // Op is "return-like".
 def ReturnLike : NativeOpTrait<"ReturnLike">;
 
 #endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES

mlir/lib/Dialect/SCF/SCF.cpp

if (std::get<0>(e).getType() != std::get<2>(e).getType())
       return op.emitOpError() << "types mismatch between " << i
                               << "th iter operand and defined value";
     if (std::get<1>(e).getType() != std::get<2>(e).getType())
       return op.emitOpError() << "types mismatch between " << i
                               << "th iter region arg and defined value";
 
     i++;
   }
-  return success();
+
+  return RegionBranchOpInterface::verifyTypes(op);
 }
 
 static void print(OpAsmPrinter &p, ForOp op) {
   bool printBlockTerminators = false;
   p << op.getOperationName() << " " << op.getInductionVar() << " = "
     << op.lowerBound() << " to " << op.upperBound() << " step " << op.step();
 
   if (op.hasIterOperands()) {
void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  function_ref<void(OpBuilder &, Location)> elseBuilder) {
   build(builder, result, TypeRange(), cond, thenBuilder, elseBuilder);
 }
 
 static LogicalResult verify(IfOp op) {
   if (op.getNumResults() != 0 && op.elseRegion().empty())
     return op.emitOpError("must have an else block if defining values");
 
-  return success();
+  return RegionBranchOpInterface::verifyTypes(op);
 }
 
 static ParseResult parseIfOp(OpAsmParser &parser, OperationState &result) {
   // Create the regions for 'then'.
   result.regions.reserve(2);
   Region *thenRegion = result.addRegion();
   Region *elseRegion = result.addRegion();
 
void IfOp::getSuccessorRegions(Optional<unsigned> index,
   Region *elseRegion = &this->elseRegion();
   if (elseRegion->empty())
     elseRegion = nullptr;
 
   // Otherwise, the successor is dependent on the condition.
   bool condition;
   if (auto condAttr = operands.front().dyn_cast_or_null<IntegerAttr>()) {
     condition = condAttr.getValue().isOneValue();
   } else {
     // If the condition isn't constant, both regions may be executed.
     regions.push_back(RegionSuccessor(&thenRegion()));
     regions.push_back(RegionSuccessor(elseRegion));
     return;
   }
 
   // Add the successor regions using the condition.
   regions.push_back(RegionSuccessor(condition ? &thenRegion() : elseRegion));
 }
return op.emitOpError()
            << "expects the same number of induction variables: "
            << body->getNumArguments()
            << " as bound and step values: " << stepValues.size();
   for (auto arg : body->getArguments())
     if (!arg.getType().isIndex())
       return op.emitOpError(
           "expects arguments for the induction variable to be of index type");
 
+  // Check that the yield has no results
+  Operation *yield = body->getTerminator();
+  if (yield->getNumOperands() != 0)
+    return yield->emitOpError() << "not allowed to have operands inside '"
+                                << ParallelOp::getOperationName() << "'";
+
   // Check that the number of results is the same as the number of ReduceOps.
   SmallVector<ReduceOp, 4> reductions(body->getOps<ReduceOp>());
   auto resultsSize = op.results().size();
   auto reductionsSize = reductions.size();
   auto initValsSize = op.initVals().size();
   if (resultsSize != reductionsSize)
     return op.emitOpError()
            << "expects number of results: " << resultsSize
static LogicalResult verify(ReduceReturnOp op) {
   auto reduceOp = cast<ReduceOp>(op.getParentOp());
   Type reduceType = reduceOp.operand().getType();
   if (reduceType != op.result().getType())
     return op.emitOpError() << "needs to have type " << reduceType
                             << " (the type of the enclosing ReduceOp)";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
-static LogicalResult verify(YieldOp op) {
-  auto parentOp = op.getParentOp();
-  auto results = parentOp->getResults();
-  auto operands = op.getOperands();
-
-  if (isa<IfOp, ForOp>(parentOp)) {
-    if (parentOp->getNumResults() != op.getNumOperands())
-      return op.emitOpError() << "parent of yield must have same number of "
-                                 "results as the yield operands";
-    for (auto e : llvm::zip(results, operands)) {
-      if (std::get<0>(e).getType() != std::get<1>(e).getType())
-        return op.emitOpError()
-               << "types mismatch between yield op and its parent";
-    }
-  } else if (isa<ParallelOp>(parentOp)) {
-    if (op.getNumOperands() != 0)
-      return op.emitOpError()
-             << "yield inside scf.parallel is not allowed to have operands";
-  } else {
-    return op.emitOpError()
-           << "yield only terminates If, For or Parallel regions";
-  }
-
-  return success();
-}
 
 static ParseResult parseYieldOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 4> operands;
   SmallVector<Type, 4> types;
   llvm::SMLoc loc = parser.getCurrentLocation();
   // Parse variadic operands list, their types, and resolve operands to SSA
   // values.
   if (parser.parseOperandList(operands) ||

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

 //===- ControlFlowInterfaces.cpp - 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/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/IR/StandardTypes.h"
+#include "llvm/ADT/SmallPtrSet.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ControlFlowInterfaces
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/ControlFlowInterfaces.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // BranchOpInterface
 //===----------------------------------------------------------------------===//
 
 /// Returns the `BlockArgument` corresponding to operand `operandIndex` in some
 /// successor if 'operandIndex' is within the range of 'operands', or None if
 /// `operandIndex` isn't a successor operand index.
-Optional<BlockArgument> mlir::detail::getBranchSuccessorArgument(
-    Optional<OperandRange> operands, unsigned operandIndex, Block *successor) {
+Optional<BlockArgument>
+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) {
+detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
   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();
 }
+
+//===----------------------------------------------------------------------===//
+// RegionBranchOpInterface
+//===----------------------------------------------------------------------===//
+
+/// Verify that types match along all region control flow edges originating from
+/// `sourceNo` (region # if source is a region, llvm::None if source is parent
+/// op). `getInputsTypesForRegion` is a function that returns the types of the
+/// inputs that flow from `sourceIndex' to the given region.
+static LogicalResult verifyTypesAlongAllEdges(
+    Operation *op, Optional<unsigned> sourceNo,
+    function_ref<TypeRange(Optional<unsigned>)> getInputsTypesForRegion) {
+  auto regionInterface = cast<RegionBranchOpInterface>(op);
+
+  SmallVector<RegionSuccessor, 2> successors;
+  unsigned numInputs;
+  if (sourceNo) {
+    Region &srcRegion = op->getRegion(sourceNo.getValue());
+    numInputs = srcRegion.getNumArguments();
+  } else {
+    numInputs = op->getNumOperands();
+  }
+  SmallVector<Attribute, 2> operands(numInputs, nullptr);
+  regionInterface.getSuccessorRegions(sourceNo, operands, successors);
+
+  for (RegionSuccessor &succ : successors) {
+    Optional<unsigned> succRegionNo;
+    if (!succ.isParent())
+      succRegionNo = succ.getSuccessor()->getRegionNumber();
+
+    auto printEdgeName = [&](InFlightDiagnostic &diag) -> InFlightDiagnostic & {
+      diag << "from ";
+      if (sourceNo)
+        diag << "Region #" << sourceNo.getValue();
+      else
+        diag << op->getName();
+
+      diag << " to ";
+      if (succRegionNo)
+        diag << "Region #" << succRegionNo.getValue();
+      else
+        diag << op->getName();
+      return diag;
+    };
+
+    TypeRange sourceTypes = getInputsTypesForRegion(succRegionNo);
+    TypeRange succInputsTypes = succ.getSuccessorInputs().getTypes();
+    if (sourceTypes.size() != succInputsTypes.size()) {
+      InFlightDiagnostic diag = op->emitOpError(" region control flow edge ");
+      return printEdgeName(diag)
+             << " has " << sourceTypes.size()
+             << " source operands, but target successor needs "
+             << succInputsTypes.size();
+    }
+
+    for (auto typesIdx :
+         llvm::enumerate(llvm::zip(sourceTypes, succInputsTypes))) {
+      Type sourceType = std::get<0>(typesIdx.value());
+      Type inputType = std::get<1>(typesIdx.value());
+      if (sourceType != inputType) {
+        InFlightDiagnostic diag = op->emitOpError(" along control flow edge ");
+        return printEdgeName(diag)
+               << " source #" << typesIdx.index() << " type " << sourceType
+               << " should match input #" << typesIdx.index() << " type "
+               << inputType;
+      }
+    }
+  }
+  return success();
+}
+
+/// Verify that types match along control flow edges described the given op.
+LogicalResult detail::verifyTypesAlongControlFlowEdges(Operation *op) {
+  auto regionInterface = cast<RegionBranchOpInterface>(op);
+
+  auto inputTypesFromParent = [&](Optional<unsigned> regionNo) -> TypeRange {
+    if (regionNo.hasValue()) {
+      return regionInterface.getSuccessorEntryOperands(regionNo.getValue())
+          .getTypes();
+    }
+
+    // If the successor of a parent op is the parent itself
+    // RegionBranchOpInterface does not have an API to query what the entry
+    // operands will be in that case. Vend out the result types of the op in
+    // that case so that type checking succeeds for this case.
+    return op->getResultTypes();
+  };
+
+  // Verify types along control flow edges originating from the parent.
+  if (failed(verifyTypesAlongAllEdges(op, llvm::None, inputTypesFromParent)))
+    return failure();
+
+  // RegionBranchOpInterface should not be implemented by Ops that do not have
+  // attached regions.
+  assert(op->getNumRegions() != 0);
+
+  // Verify types along control flow edges originating from each region.
+  for (unsigned regionNo : llvm::seq(0U, op->getNumRegions())) {
+    Region &region = op->getRegion(regionNo);
+
+    // Since the interface cannnot distinguish between different ReturnLike
+    // ops within the region branching to different successors, all ReturnLike
+    // ops in this region should have the same operand types. We will then use
+    // one of them as the representative for type matching.
+
+    Operation *regionReturn = nullptr;
+    for (Block &block : region) {
+      Operation *terminator = block.getTerminator();
+      if (!terminator->hasTrait<OpTrait::ReturnLike>())
+        continue;
+
+      if (!regionReturn) {
+        regionReturn = terminator;
+        continue;
+      }
+
+      // Found more than one ReturnLike terminator. Make sure the operand types
+      // match with the first one.
+      if (regionReturn->getOperandTypes() != terminator->getOperandTypes())
+        return op->emitOpError("Region #")
+               << regionNo
+               << " operands mismatch between return-like terminators";
+    }
+
+    auto inputTypesFromRegion = [&](Optional<unsigned> regionNo) -> TypeRange {
+      // All successors get the same set of operands.
+      return regionReturn ? TypeRange(regionReturn->getOperands().getTypes())
+                          : TypeRange();
+    };
+
+    if (failed(verifyTypesAlongAllEdges(op, regionNo, inputTypesFromRegion)))
+      return failure();
+  }
+
+  return success();
+}

mlir/test/Dialect/SCF/invalid.mlir

func @reduceReturn_not_inside_reduce(%arg0 : f32) {
   }): () -> ()
   return
 }
 
 // -----
 
 func @std_if_incorrect_yield(%arg0: i1, %arg1: f32)
 {
+  // expected-error@+1 {{region control flow edge from Region #0 to scf.if has 1 source operands, but target successor needs 2}}
   %x, %y = scf.if %arg0 -> (f32, f32) {
     %0 = addf %arg1, %arg1 : f32
-    // expected-error@+1 {{parent of yield must have same number of results as the yield operands}}
     scf.yield %0 : f32
   } else {
     %0 = subf %arg1, %arg1 : f32
-    scf.yield %0 : f32
+    scf.yield %0, %0 : f32, f32
   }
   return
 }
 
 // -----
 
 func @std_if_missing_else(%arg0: i1, %arg1: f32)
 {
%result1:2 = scf.for %i0 = %arg0 to %arg1 step %arg2
     %tn = addf %ti, %ti : i32
     scf.yield %sn, %tn : i32, i32
   }
   return
 }
 
 // -----
 
+func @std_for_operands_mismatch_4(%arg0 : index, %arg1 : index, %arg2 : index) {
+  %s0 = constant 0.0 : f32
+  %t0 = constant 1.0 : f32
+  // expected-error @+1 {{along control flow edge from Region #0 to Region #0 source #1 type 'i32' should match input #1 type 'f32'}}
+  %result1:2 = scf.for %i0 = %arg0 to %arg1 step %arg2
+                    iter_args(%si = %s0, %ti = %t0) -> (f32, f32) {
+    %sn = addf %si, %si : f32
+    %ic = constant 1 : i32
+    scf.yield %sn, %ic : f32, i32
+  }
+  return
+}
+
+
+// -----
+
 func @parallel_invalid_yield(
     %arg0: index, %arg1: index, %arg2: index) {
   scf.parallel (%i0) = (%arg0) to (%arg1) step (%arg2) {
     %c0 = constant 1.0 : f32
-    // expected-error@+1 {{yield inside scf.parallel is not allowed to have operands}}
+    // expected-error@+1 {{'scf.yield' op not allowed to have operands inside 'scf.parallel'}}
     scf.yield %c0 : f32
   }
   return
 }
+
+// -----
+func @yield_invalid_parent_op() {
+  "my.op"() ({
+   // expected-error@+1 {{'scf.yield' op expects parent op to be one of 'scf.if, scf.for, scf.parallel'}}
+   scf.yield
+  }) : () -> ()
+  return
+}