[mlir] Add a generic while/do-while loop to the SCF dialect

[ftynse]

Summary

The new construct represents a generic loop with two regions: one executed
before the loop condition is verifier and another after that. This construct
can be used to express both a "while" loop and a "do-while" loop, depending on
where the main payload is located. It is intended as an intermediate
abstraction for lowering, which will be added later. This form is relatively
easy to target from higher-level abstractions and supports transformations such
as loop rotation and LICM.

[silvas]

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

498

Unclear what "match" means in this first sentence (match with each other? match with the outer op?). Maybe make it explicit "The two regions need not have the same argument list."

mlir/lib/Dialect/SCF/SCF.cpp

1029

A lot of this verification is subsumed by implementing RegionBranchOpInterface and calling RegionBranchOpInterface::verifyTypes (this would also reduce the testing burden, as we generally don't require testing properties verified by traits/interfaces)

[silvas]

Thanks for pushing on this! This is looking really good :)

[aartbik]

Thanks for adding the while!

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

438

.... represents a generic ...
reads better

469

Label this as "After" region

mlir/lib/Dialect/SCF/SCF.cpp

1032

You switch a bit between first/second and before/after region in your terminology. Perhaps use one name consistently throughout doc and messages?

[mehdi_amini]

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

440–441

Random question: we tend to easily add this single-block restriction, but I'm not sure why is it the right thing to do?

[rriddle]

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

435

Can you add the RegionControlFlowInterface to this operation?

448

nit: The two regions ?

mlir/lib/Dialect/SCF/SCF.cpp

976

parser.resolveOperands?

1005

nit: Add braces to the outer loop if the inner statement also has braces.

1043

nit: Add braces around this, it isn't really trivial.

[nicolasvasilache]

Nice, thanks for implementing this ! :)

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

466

You probably want to be very loud in this part of the doc to make the textual part more intuitive.

/* Compute values for "after" block*/
%foo#2 = call @compute(%arg1) : (f32) -> (i42, vector<3xf32>)

/* "Before" region. */
%condition = call @evaluate_condition(%arg1, %foo#1) : (f32, vector<3xf32>) -> i1

/* Forward the argument (as result or "after" region argument). */
scf.yield %condition, %foo#0, %foo#1, %arg1 : i1, i42, vector<3xf32>, f32

473

I'd add a comment: /* Must match enclosing scf.while operand types */ for in the same spirit of making the textual part more intuitive.

mlir/lib/Dialect/SCF/SCF.cpp

1029

+1

[ftynse]

Partially address review

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

440–441

I am considering to relax this, but this should be done with a proper RFC and apply to all SCF loops IMO. Allowing scf.yield in multiple blocks is equivalent to implementing a continue, and we may also want a break. Loop transformations currently rely on the single-block property.

466

Added another example. I want to keep the simple examples simple.

mlir/lib/Dialect/SCF/SCF.cpp

1029

Not as much actually as I would have liked. We cannot use it to verify the yield of the 'before' block because it takes the condition operand that it does not forward.

Looking at this closer, I am considering to introduce a different terminator for this region, something like scf.while_condition(%i1) %arg1, %arg2 : types to avoid breaking the "return-like" semantics of scf.yield. WDYT?

[silvas]

mlir/lib/Dialect/SCF/SCF.cpp

1029

Interesting. I thought this was supported, but perhaps not. @rriddle FYI

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1029

About the yield "return semantics": I see the terminator as always terminating the region and transferring the control to the parent op, yield fits here to me just like other loops.
The parent op then internally look at one of the result, and depending on the value forward the other results to the other region or to its control successor.

I was also expecting the region branch op interface to support this, but it may haven't been implemented yet. @jurahul worked on this at some point?

[aartbik]

lowering next? ;-)

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

438

The a is now at the wrong place :-)

This operation represents a generic ....

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

1029

If I read the code correctly, there is an assumption that all operands of a ReturnLike terminator are forwarded somewhere.

I see the terminator as always terminating the region and transferring the control to the parent op,

This is what we wrote in the spec. But I interpret the RegionBranchOpInterface API as "after receiving the flow of control from region X, the op can forwarded it to region Y or to its natural control-flow successor (next op in the block)". In the former case, there is an expectation that types match exactly AFAICS.

[ftynse]

Added scf.condition and made WhileOp use RegionBranchOpInterface

[ftynse]

PTAL

[jurahul]

mlir/lib/Dialect/SCF/SCF.cpp

1029

Yes, I had spec'd an extension to RegionBranchOpInterface to support more extensive type checking, specifically around "internal" values that op's may produce and pass into attached regions. The main feature there was the interface declares, across all control flow edges, the types of values that flow across them, separately from the values. So its then becomes possible to typecheck say a region inputs when they feed not from the operation inputs but some internally generated values (which have no direct IR representation). And similarly typecheck region outputs that are consumed internally by the operation (and not directly feeding into the output).

However, I have not yet had time to actually follow that up with a discourse RFC and actual implementation.

[ftynse]

Rebase

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1029

I don't understand how you quote from the RegionBranchOpInterface contradicts the quote you're answering from me? On the contrary it seems to say the same thing here?
In particular my comment was about the introduction of another terminator "to avoid breaking the "return-like" semantics of scf.yield", and I don't quite get the relationship with the types right now?

Can we have a TODO somewhere for the missing RegionBranchOpInterface feature?

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

1029

I don't think there is a contradiction on RegionBranchOpInterface.

However, there is a "ReturnLike" trait that RegionBranchOpInterface uses to understand that all operands are sent back to the parent op for dispatch. A potentially useful addition could be an interface method that specifies which operands instead of just taking all of them. That being said, "ReturnLike" seems to have some other semantics attached to it, and I am not willing to revise those semantics here or remove them from YieldOp for the purpose of this patch. Hence the new op.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1029

A potentially useful addition could be an interface method that specifies which operands instead of just taking all of them. That being said, "ReturnLike" seems to have some other semantics attached to it, and I am not willing to revise those semantics here or remove them from YieldOp for the purpose of this patch. Hence the new op.

I still don't understand why yield can't be directly used here?

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

1029

Because it has the ReturnLike trait, which disallows to take the leading i1 operand without the next region also having the leading i1 argument. ReturnLike is also connected to other concepts, so I am reluctant to change it without a proper discussion.

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

class SCF_Op<string mnemonic, list<OpTrait> traits = []> :
   //   * ParseResult parse${C++ class of Op}(OpAsmParser &parser,
   //                                         OperationState &result)
   // functions.
   let printer = [{ return ::print(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
+def ConditionOp : SCF_Op<"condition",
+                         [HasParent<"WhileOp">, NoSideEffect, Terminator]> {
+  let summary = "loop continuation condition";
+  let description = [{
+    This operation accepts the continuation (i.e., inverse of exit) condition
+    of the `scf.while` construct. If its first argument is true, the "after"
+    region of `scf.while` is executed, with the remaining arguments forwarded
+    to the entry block of the region. Otherwise, the loop terminates.
+  }];
+
+  let arguments = (ins I1:$condition, Variadic<AnyType>:$args);
+
+  let assemblyFormat =
+      [{ `(` $condition `)` attr-dict ($args^ `:` type($args))? }];
+
+  // Override the default verifier, everything is checked by traits.
+  let verifier = ?;
+}
+
 def ForOp : SCF_Op<"for",
       [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveSideEffects]> {
   let summary = "for operation";
   let description = [{
     The "scf.for" operation represents a loop taking 3 SSA value as operands
let description = [{
     scf.reduce.return %res : f32
     ```
   }];
 
   let arguments = (ins AnyType:$result);
   let assemblyFormat = "$result attr-dict `:` type($result)";
 }
 
+def WhileOp : SCF_Op<"while",
+    [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
+     RecursiveSideEffects]> {
+  let summary = "a generic 'while' loop";
+  let description = [{
+    This operation represents a generic "while"/"do-while" loop that keeps
+    iterating as long as a condition is satisfied. There is no restriction on
+    the complexity of the condition. It consists of two regions (with single
+    block each): "before" region and "after" region. The names of regions
+    indicates whether they execute before or after the condition check.
+    Therefore, if the main loop payload is located in the "before" region, the
+    operation is a "do-while" loop. Otherwise, it is a "while" loop.
+
+    The "before" region terminates with a special operation, `scf.condition`,
+    that accepts as its first operand an `i1` value indicating whether to
+    proceed to the "after" region (value is `true`) or not. The two regions
+    communicate by means of region arguments. Initially, the "before" region
+    accepts as arguments the operands of the `scf.while` operation and uses them
+    to evaluate the condition. It forwards the trailing, non-condition operands
+    of the `scf.condition` terminator either to the "after" region if the
+    control flow is transferred there or to results of the `scf.while` operation
+    otherwise. The "after" region takes as arguments the values produced by the
+    "before" region and uses `scf.yield` to supply new arguments for the "after"
+    region, into which it transfers the control flow unconditionally.
+
+    A simple "while" loop can be represented as follows.
+
+    ```mlir
+    %res = scf.while (%arg1 = %init1) : (f32) -> f32 {
+      /* "Before" region.
+       * In a "while" loop, this region computes the condition. */
+      %condition = call @evaluate_condition(%arg1) : (f32) -> i1
+
+      /* Forward the argument (as result or "after" region argument). */
+      scf.condition(%condition) %arg1 : f32
+
+    } do {
+    ^bb0(%arg2: f32):
+      /* "After region.
+       * In a "while" loop, this region is the loop body. */
+      %next = call @payload(%arg2) : (f32) -> f32
+
+      /* Forward the new value to the "before" region.
+       * The operand types must match the types of the `scf.while` operands. */
+      scf.yield %next : f32
+    }
+    ```
+
+    A simple "do-while" loop can be represented by reducing the "after" block
+    to a simple forwarder.
+
+    ```mlir
+    %res = scf.while (%arg1 = %init1) : (f32) -> f32 {
+      /* "Before" region.
+       * In a "do-while" loop, this region contains the loop body. */
+      %next = call @payload(%arg1) : (f32) -> f32
+
+      /* And also evalutes the condition. */
+      %condition = call @evaluate_condition(%arg1) : (f32) -> i1
+
+      /* Loop through the "after" region. */
+      scf.condition(%condition) %next : f32
+
+    } do {
+    ^bb0(%arg2: f32):
+      /* "After" region.
+       * Forwards the values back to "before" region unmodified. */
+      scf.yield %arg2 : f32
+    }
+    ```
+
+    Note that the types of region arguments need not to match with each other.
+    The op expects the operand types to match with argument types of the
+    "before" region"; the result types to match with the trailing operand types
+    of the terminator of the "before" region, and with the argument types of the
+    "after" region. The following scheme can be used to share the results of
+    some operations executed in the "before" region with the "after" region,
+    avoiding the need to recompute them.
+
+    ```mlir
+    %res = scf.while (%arg1 = %init1) : (f32) -> i64 {
+      /* One can perform some computations, e.g., necessary to evaluate the
+       * condition, in the "before" region and forward their results to the
+       * "after" region. */
+      %shared = call @shared_compute(%arg1) : (f32) -> i64
+
+      /* Evalute the condition. */
+      %condition = call @evaluate_condition(%arg1, %shared) : (f32, i64) -> i1
+
+      /* Forward the result of the shared computation to the "after" region.
+       * The types must match the arguments of the "after" region as well as
+       * those of the `scf.while` results. */
+      scf.condition(%condition) %shared : i64
+
+    } do {
+    ^bb0(%arg2: i64) {
+      /* Use the partial result to compute the rest of the payload in the
+       * "after" region. */
+      %res = call @payload(%arg2) : (i64) -> f32
+
+      /* Forward the new value to the "before" region.
+       * The operand types must match the types of the `scf.while` operands. */
+      scf.yield %res : f32
+    }
+    ```
+
+    The custom syntax for this operation is as follows.
+
+    ```
+    op ::= `scf.while` assignments `:` function-type region `do` region
+           `attributes` attribute-dict
+    initializer ::= /* empty */ | `(` assignment-list `)`
+    assignment-list ::= assignment | assignment `,` assignment-list
+    assignment ::= ssa-value `=` ssa-value
+    ```
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$inits);
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region SizedRegion<1>:$before, SizedRegion<1>:$after);
+
+  let extraClassDeclaration = [{
+    OperandRange getSuccessorEntryOperands(unsigned index);
+  }];
+}
+
 def YieldOp : SCF_Op<"yield", [NoSideEffect, ReturnLike, Terminator,
-                               ParentOneOf<["IfOp, ForOp", "ParallelOp"]>]> {
+                               ParentOneOf<["IfOp, ForOp", "ParallelOp",
+                                            "WhileOp"]>]> {
   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 = [OpBuilderDAG<(ins), [{ /* nothing to do */ }]>];
   // Override default verifier (defined in SCF_Op), no custom verification
   // needed.
   let verifier = ?;
 }
+
 #endif // MLIR_DIALECT_SCF_SCFOPS

mlir/include/mlir/IR/OpImplementation.h

public:
   virtual ParseResult parseColonTypeList(SmallVectorImpl<Type> &result) = 0;
 
   /// Parse an optional colon followed by a type list, which if present must
   /// have at least one type.
   virtual ParseResult
   parseOptionalColonTypeList(SmallVectorImpl<Type> &result) = 0;
 
   /// Parse a list of assignments of the form
-  /// (%x1 = %y1 : type1, %x2 = %y2 : type2, ...).
-  /// The list must contain at least one entry
-  virtual ParseResult
-  parseAssignmentList(SmallVectorImpl<OperandType> &lhs,
+  ///   (%x1 = %y1, %x2 = %y2, ...)
+  ParseResult parseAssignmentList(SmallVectorImpl<OperandType> &lhs,
+                                  SmallVectorImpl<OperandType> &rhs) {
+    OptionalParseResult result = parseOptionalAssignmentList(lhs, rhs);
+    if (!result.hasValue())
+      return emitError(getCurrentLocation(), "expected '('");
+    return result.getValue();
+  }
+
+  virtual OptionalParseResult
+  parseOptionalAssignmentList(SmallVectorImpl<OperandType> &lhs,
                               SmallVectorImpl<OperandType> &rhs) = 0;
 
   /// Parse a keyword followed by a type.
   ParseResult parseKeywordType(const char *keyword, Type &result) {
     return failure(parseKeyword(keyword) || parseType(result));
   }
 
   /// Add the specified type to the end of the specified type list and return
   /// success.  This is a helper designed to allow parse methods to be simple

mlir/lib/Dialect/SCF/SCF.cpp

if (std::get<1>(e).getType() != std::get<2>(e).getType())
                               << "th iter region arg and defined value";
 
     i++;
   }
 
   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()) {
-    p << " iter_args(";
-    auto regionArgs = op.getRegionIterArgs();
-    auto operands = op.getIterOperands();
+/// Prints the initialization list in the form of
+///   <prefix>(%inner = %outer, %inner2 = %outer2, <...>)
+/// where 'inner' values are assumed to be region arguments and 'outer' values
+/// are regular SSA values.
+static void printInitializationList(OpAsmPrinter &p,
+                                    Block::BlockArgListType blocksArgs,
+                                    ValueRange initializers,
+                                    StringRef prefix = "") {
+  assert(blocksArgs.size() == initializers.size() &&
+         "expected same length of arguments and initializers");
+  if (initializers.empty())
+    return;
 
-    llvm::interleaveComma(llvm::zip(regionArgs, operands), p, [&](auto it) {
-      p << std::get<0>(it) << " = " << std::get<1>(it);
-    });
-    p << ")";
-    p << " -> (" << op.getResultTypes() << ")";
-    printBlockTerminators = true;
-  }
+  p << prefix << '(';
+  llvm::interleaveComma(llvm::zip(blocksArgs, initializers), p, [&](auto it) {
+
+static void print(OpAsmPrinter &p, ForOp op) {
+  p << op.getOperationName() << " " << op.getInductionVar() << " = "
+    << op.lowerBound() << " to " << op.upperBound() << " step " << op.step();
+
+  printInitializationList(p, op.getRegionIterArgs(), op.getIterOperands(),
+                          " iter_args");
+  if (!op.getIterOperands().empty())
+    p << " -> (" << op.getIterOperands().getTypes() << ')';
   p.printRegion(op.region(),
                 /*printEntryBlockArgs=*/false,
-                /*printBlockTerminators=*/printBlockTerminators);
+                /*printBlockTerminators=*/op.hasIterOperands());
   p.printOptionalAttrDict(op.getAttrs());
 }
 
 static ParseResult parseForOp(OpAsmParser &parser, OperationState &result) {
   auto &builder = parser.getBuilder();
   OpAsmParser::OperandType inductionVariable, lb, ub, step;
   // Parse the induction variable followed by '='.
   if (parser.parseRegionArgument(inductionVariable) || parser.parseEqual())
static LogicalResult verify(ReduceReturnOp op) {
   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();
 }
 
 //===----------------------------------------------------------------------===//
+// WhileOp
+//===----------------------------------------------------------------------===//
+
+OperandRange WhileOp::getSuccessorEntryOperands(unsigned index) {
+  assert(index == 0 &&
+         "WhileOp is expected to branch only to the first region");
+
+  return inits();
+}
+
+void WhileOp::getSuccessorRegions(Optional<unsigned> index,
+                                  ArrayRef<Attribute> operands,
+                                  SmallVectorImpl<RegionSuccessor> &regions) {
+  (void)operands;
+
+  if (!index.hasValue()) {
+    regions.emplace_back(&before(), before().getArguments());
+    return;
+  }
+
+  assert(*index < 2 && "there are only two regions in a WhileOp");
+  if (*index == 0) {
+    regions.emplace_back(&after(), after().getArguments());
+    regions.emplace_back(getResults());
+    return;
+  }
+
+  regions.emplace_back(&before(), before().getArguments());
+}
+
+/// Parses a `while` op.
+///
+/// op ::= `scf.while` assignments `:` function-type region `do` region
+///         `attributes` attribute-dict
+/// initializer ::= /* empty */ | `(` assignment-list `)`
+/// assignment-list ::= assignment | assignment `,` assignment-list
+/// assignment ::= ssa-value `=` ssa-value
+static ParseResult parseWhileOp(OpAsmParser &parser, OperationState &result) {
+  SmallVector<OpAsmParser::OperandType, 4> regionArgs, operands;
+  Region *before = result.addRegion();
+  Region *after = result.addRegion();
+
+  OptionalParseResult listResult =
+      parser.parseOptionalAssignmentList(regionArgs, operands);
+  if (listResult.hasValue() && failed(listResult.getValue()))
+    return failure();
+
+  FunctionType functionType;
+  llvm::SMLoc typeLoc = parser.getCurrentLocation();
+  if (failed(parser.parseColonType(functionType)))
+    return failure();
+
+  result.addTypes(functionType.getResults());
+
+  if (functionType.getNumInputs() != operands.size()) {
+    return parser.emitError(typeLoc)
+           << "expected as many input types as operands "
+           << "(expected " << operands.size() << " got "
+           << functionType.getNumInputs() << ")";
+  }
+
+  // Resolve input operands.
+  if (failed(parser.resolveOperands(operands, functionType.getInputs(),
+                                    parser.getCurrentLocation(),
+                                    result.operands)))
+    return failure();
+
+  return failure(
+      parser.parseRegion(*before, regionArgs, functionType.getInputs()) ||
+      parser.parseKeyword("do") || parser.parseRegion(*after) ||
+      parser.parseOptionalAttrDictWithKeyword(result.attributes));
+}
+
+/// Prints a `while` op.
+static void print(OpAsmPrinter &p, scf::WhileOp op) {
+  p << op.getOperationName();
+  printInitializationList(p, op.before().front().getArguments(), op.inits(),
+                          " ");
+  p << " : ";
+  p.printFunctionalType(op.inits().getTypes(), op.results().getTypes());
+  p.printRegion(op.before(), /*printEntryBlockArgs=*/false);
+  p << " do";
+  p.printRegion(op.after());
+  p.printOptionalAttrDictWithKeyword(op.getAttrs());
+}
+
+/// Verifies that two ranges of types match, i.e. have the same number of
+/// entries and that types are pairwise equals. Reports errors on the given
+/// operation in case of mismatch.
+template <typename OpTy>
+static LogicalResult verifyTypeRangesMatch(OpTy op, TypeRange left,
+                                           TypeRange right, StringRef message) {
+  if (left.size() != right.size())
+    return op.emitOpError("expects the same number of ") << message;
+
+  for (unsigned i = 0, e = left.size(); i < e; ++i) {
+    if (left[i] != right[i]) {
+      InFlightDiagnostic diag = op.emitOpError("expects the same types for ")
+                                << message;
+      diag.attachNote() << "for argument " << i << ", found " << left[i]
+                        << " and " << right[i];
+      return diag;
+    }
+  }
+
+  return success();
+}
+
+/// Verifies that the first block of the given `region` is terminated by a
+/// YieldOp. Reports errors on the given operation if it is not the case.
+template <typename TerminatorTy>
+static TerminatorTy verifyAndGetTerminator(scf::WhileOp op, Region &region,
+                                           StringRef errorMessage) {
+  Operation *terminatorOperation = region.front().getTerminator();
+  if (auto yield = dyn_cast_or_null<TerminatorTy>(terminatorOperation))
+    return yield;
+
+  auto diag = op.emitOpError(errorMessage);
+  if (terminatorOperation)
+    diag.attachNote(terminatorOperation->getLoc()) << "terminator here";
+  return nullptr;
+}
+
+static LogicalResult verify(scf::WhileOp op) {
+  if (failed(RegionBranchOpInterface::verifyTypes(op)))
+    return failure();
+
+  auto beforeTerminator = verifyAndGetTerminator<scf::ConditionOp>(
+      op, op.before(),
+      "expects the 'before' region to terminate with 'scf.condition'");
+  if (!beforeTerminator)
+    return failure();
+
+  TypeRange trailingTerminatorOperands = beforeTerminator.args().getTypes();
+  if (failed(verifyTypeRangesMatch(op, trailingTerminatorOperands,
+                                   op.after().getArgumentTypes(),
+                                   "trailing operands of the 'before' block "
+                                   "terminator and 'after' region arguments")))
+    return failure();
+
+  if (failed(verifyTypeRangesMatch(
+          op, trailingTerminatorOperands, op.getResultTypes(),
+          "trailing operands of the 'before' block terminator and op results")))
+    return failure();
+
+  auto afterTerminator = verifyAndGetTerminator<scf::YieldOp>(
+      op, op.after(),
+      "expects the 'after' region to terminate with 'scf.yield'");
+  return success(afterTerminator != nullptr);
+}
+
+//===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 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

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

 
 //===----------------------------------------------------------------------===//
 // 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) {
+/// inputs that flow from `sourceIndex' to the given region, or llvm::None if
+/// the exact type match verification is not necessary (e.g., if the Op verifies
+/// the match itself).
+static LogicalResult
+verifyTypesAlongAllEdges(Operation *op, Optional<unsigned> sourceNo,
+                         function_ref<Optional<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 {
auto printEdgeName = [&](InFlightDiagnostic &diag) -> InFlightDiagnostic & {
       diag << " to ";
       if (succRegionNo)
         diag << "Region #" << succRegionNo.getValue();
       else
         diag << "parent results";
       return diag;
     };
 
-    TypeRange sourceTypes = getInputsTypesForRegion(succRegionNo);
+    Optional<TypeRange> sourceTypes = getInputsTypesForRegion(succRegionNo);
+    if (!sourceTypes.hasValue())
+      continue;
+
     TypeRange succInputsTypes = succ.getSuccessorInputs().getTypes();
-    if (sourceTypes.size() != succInputsTypes.size()) {
+    if (sourceTypes->size() != succInputsTypes.size()) {
       InFlightDiagnostic diag = op->emitOpError(" region control flow edge ");
-      return printEdgeName(diag) << ": source has " << sourceTypes.size()
+      return printEdgeName(diag) << ": source has " << sourceTypes->size()
                                  << " operands, but target successor needs "
                                  << succInputsTypes.size();
     }
 
     for (auto typesIdx :
-         llvm::enumerate(llvm::zip(sourceTypes, succInputsTypes))) {
+         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 type #" << typesIdx.index() << " " << sourceType
                << " should match input type #" << typesIdx.index() << " "
                << inputType;
for (Block &block : region) {
       // 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 {
+    auto inputTypesFromRegion =
+        [&](Optional<unsigned> regionNo) -> Optional<TypeRange> {
+      // If there is no return-like terminator, the op itself should verify
+      // type consistency.
+      if (!regionReturn)
+        return llvm::None;
+
       // All successors get the same set of operands.
-      return regionReturn ? TypeRange(regionReturn->getOperands().getTypes())
-                          : TypeRange();
+      return TypeRange(regionReturn->getOperands().getTypes());
     };
 
     if (failed(verifyTypesAlongAllEdges(op, regionNo, inputTypesFromRegion)))
       return failure();
   }
 
   return success();
 }

mlir/lib/Parser/Parser.cpp

public:
   ParseResult
   parseOptionalColonTypeList(SmallVectorImpl<Type> &result) override {
     if (!parser.consumeIf(Token::colon))
       return success();
     return parser.parseTypeListNoParens(result);
   }
 
   /// Parse a list of assignments of the form
-  /// (%x1 = %y1 : type1, %x2 = %y2 : type2, ...).
-  /// The list must contain at least one entry
-  ParseResult parseAssignmentList(SmallVectorImpl<OperandType> &lhs,
-                                  SmallVectorImpl<OperandType> &rhs) override {
+  ///   (%x1 = %y1, %x2 = %y2, ...).
+  OptionalParseResult
+  parseOptionalAssignmentList(SmallVectorImpl<OperandType> &lhs,
+    if (failed(parseOptionalLParen()))
+      return llvm::None;
+
     auto parseElt = [&]() -> ParseResult {
       OperandType regionArg, operand;
       if (parseRegionArgument(regionArg) || parseEqual() ||
           parseOperand(operand))
         return failure();
       lhs.push_back(regionArg);
       rhs.push_back(operand);
       return success();
     };
-    if (parseLParen())
-      return failure();
     return parser.parseCommaSeparatedListUntil(Token::r_paren, parseElt);
   }
 
 private:
   /// The source location of the operation name.
   SMLoc nameLoc;
 
   /// Information about the result name specifiers.

mlir/test/Dialect/SCF/invalid.mlir

scf.parallel (%i0) = (%arg0) to (%arg1) step (%arg2) {
     %c0 = constant 1.0 : f32
     // 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'}}
+   // expected-error@+1 {{'scf.yield' op expects parent op to be one of 'scf.if, scf.for, scf.parallel, scf.while'}}
    scf.yield
   }) : () -> ()
   return
 }
+
+// -----
+
+func @while_parser_type_mismatch() {
+  %true = constant true
+  // expected-error@+1 {{expected as many input types as operands (expected 0 got 1)}}
+  scf.while : (i32) -> () {
+    scf.condition(%true)
+  } do {
+    scf.yield
+  }
+}
+
+// -----
+
+func @while_bad_terminator() {
+  // expected-error@+1 {{expects the 'before' region to terminate with 'scf.condition'}}
+  scf.while : () -> () {
+    // expected-note@+1 {{terminator here}}
+    "some.other_terminator"() : () -> ()
+  } do {
+    scf.yield
+  }
+}
+
+// -----
+
+func @while_cross_region_type_mismatch() {
+  %true = constant true
+  // expected-error@+1 {{expects the same number of trailing operands of the 'before' block terminator and 'after' region arguments}}
+  scf.while : () -> () {
+    scf.condition(%true)
+  } do {
+  ^bb0(%arg0: i32):
+    scf.yield
+  }
+}
+
+// -----
+
+func @while_cross_region_type_mismatch() {
+  %true = constant true
+  // expected-error@+2 {{expects the same types for trailing operands of the 'before' block terminator and 'after' region arguments}}
+  // expected-note@+1 {{for argument 0, found 'i1' and 'i32}}
+  scf.while : () -> () {
+    scf.condition(%true) %true : i1
+  } do {
+  ^bb0(%arg0: i32):
+    scf.yield
+  }
+}
+
+// -----
+
+func @while_result_type_mismatch() {
+  %true = constant true
+  // expected-error@+1 {{expects the same number of trailing operands of the 'before' block terminator and op results}}
+  scf.while : () -> () {
+    scf.condition(%true) %true : i1
+  } do {
+  ^bb0(%arg0: i1):
+    scf.yield
+  }
+}
+
+// -----
+
+func @while_bad_terminator() {
+  %true = constant true
+  // expected-error@+1 {{expects the 'after' region to terminate with 'scf.yield'}}
+  scf.while : () -> () {
+    scf.condition(%true)
+  } do {
+    // expected-note@+1 {{terminator here}}
+    "some.other_terminator"() : () -> ()
+  }
+}

mlir/test/Dialect/SCF/ops.mlir

 //  CHECK-NEXT: %[[THENRES:.*]] = addf %[[ITER]], %[[T]]
 //  CHECK-NEXT: scf.yield %[[THENRES]] : f32
 //  CHECK-NEXT: } else {
 //  CHECK-NEXT: scf.yield %[[ITER]] : f32
 //  CHECK-NEXT: }
 //  CHECK-NEXT: scf.yield %[[IFRES]] : f32
 //  CHECK-NEXT: }
 //  CHECK-NEXT: return %[[RESULT]]
+
+// CHECK-LABEL: @while
+func @while() {
+  %0 = "test.get_some_value"() : () -> i32
+  %1 = "test.get_some_value"() : () -> f32
+
+  // CHECK: = scf.while (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) : (i32, f32) -> (i64, f64) {
+  %2:2 = scf.while (%arg0 = %0, %arg1 = %1) : (i32, f32) -> (i64, f64) {
+    %3:2 = "test.some_operation"(%arg0, %arg1) : (i32, f32) -> (i64, f64)
+    %4 = "test.some_condition"(%arg0, %arg1) : (i32, f32) -> i1
+    // CHECK: scf.condition(%{{.*}}) %{{.*}}, %{{.*}} : i64, f64
+    scf.condition(%4) %3#0, %3#1 : i64, f64
+  // CHECK: } do {
+  } do {
+  // CHECK: ^{{.*}}(%{{.*}}: i64, %{{.*}}: f64):
+  ^bb0(%arg2: i64, %arg3: f64):
+    %5:2 = "test.some_operation"(%arg2, %arg3): (i64, f64) -> (i32, f32)
+    // CHECK: scf.yield %{{.*}}, %{{.*}} : i32, f32
+    scf.yield %5#0, %5#1 : i32, f32
+  // CHECK: attributes {foo = "bar"}
+  } attributes {foo="bar"}
+  return
+}
+
+// CHECK-LABEL: @infinite_while
+func @infinite_while() {
+  %true = constant true
+
+  // CHECK: scf.while  : () -> () {
+  scf.while : () -> () {
+    // CHECK: scf.condition(%{{.*}})
+    scf.condition(%true)
+  // CHECK: } do {
+  } do {
+    // CHECK: scf.yield
+    scf.yield
+  }
+  return
+}