[MLIR] Support for return values in Affine.For yield

[avarmapml]

Summary

• Added support for return values in affine.for yield along the same lines as scf.for and affine.parallel.

Signed-off-by: Abhishek Varma <[e-mail]>

[bondhugula]

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

303

Format fix please.

306

This should return an unsigned.

[bondhugula]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1248–1252

Nit: > 0

1413

operandType

1421

Nit: terminate with full stop.

1483–1484

I think both these methods return unsigned.

mlir/test/Dialect/Affine/invalid.mlir

386

Consider 'defined values' -> 'results'
mismatch in number -> mismatch between the number
?

[dcaballe]

LGTM! Just some nit comments! Thanks for contributing this feature!

This commit does not provide support for reduction var.

What do you mean exactly? The example in the doc is a reduction. Do you mean that this patch doesn't include the analysis to deal with reductions described using iter_args?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

231

some doc here about what this is for would be great

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1179

Same: doc here about bodyBuilder would be useful.

1191

nit: v -> val/value,...

1212

same

mlir/test/Dialect/Affine/invalid.mlir

384

nit: affine_for -> something more descriptive? affine_for_iter_args_mismatch?

mlir/test/Dialect/Affine/ops.mlir

211

nit: affine.if indentation++

221

nit: affine_for_multiple_yield?

[avarmapml]

Tackled initial review comments.

[avarmapml]

Bug fix for build function.

[bondhugula]

I think we definitely need a test for the builder API here to make sure it's working as intended. Please see the other builder API tests for reference in test/EDSC/builder-api-test.cpp.

[avarmapml]

Added test case for build method and bug fix

• getUpperBoundOperands() was not picking up operands correctly. This update fixes that.
• Added a minor test case to check build method for affine.for along with iter_args.
• Updated function affineLoopBuilder definition and updated the same in builder-api-test.cpp.

[avarmapml]

Updated getter/setter functions for lower and upper bound

[bondhugula]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1631–1632

We should update the test cases to make sure these are covered - without this change, print would have been printing the wrong thing. You could add a test case that has upper bound operands as well as iter arguments.

[bondhugula]

Thanks for fixing the issues.

mlir/test/EDSC/builder-api-test.cpp

180 ↗

(On Diff #291719)

builder_for_iter_arg -> builder_affine_for_iter_args?

194 ↗

(On Diff #291719)

Please also match the body with the yield.

725 ↗

(On Diff #291719)

/*iterArgs=*/llvm::None

[bondhugula]

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

179–180

3 loop control values ... <-- this would be inaccurate (copy/paste error).

184

Nit: a -> an

302–303

This would be inaccurate - since the init operands for the iter args shouldn't be counted here. Please check if this accessor is needed and what the intended use is.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1486

You can just do getNumOperands() with getOperation(). Also, compactly:

return getNumOperands() - lbMap.getNumInputs() - ubMap.getNumInputs();

[avarmapml]

Handled review comments and updated test case

[bondhugula]

mlir/test/EDSC/builder-api-test.cpp

194 ↗

(On Diff #291719)

Looks like this hasn't been addressed, but marked done. Did you missing pushing?

[avarmapml]

Updated test case to match loop body/yield

[bondhugula]

mlir/test/EDSC/builder-api-test.cpp

195–196 ↗

(On Diff #292418)

You don't need SmallVector here - instead, ValueRange(.., ..) inline would work.

[bondhugula]

Looks good, thanks!

[avarmapml]

Minor code refactor

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

func @simple_example(%A: memref<?x?xf32>, %B: memref<?x?xf32>) {
           %0 = affine.apply #map57(%j)[%N]
           %tmp = call @F1(%A, %i, %0) : (memref<?x?xf32>, index, index)->(f32)
           call @F2(%tmp, %B, %i, %0) : (f32, memref<?x?xf32>, index, index)->()
         }
       }
       return
     }
     ```
+    `affine.for` can also operate on loop-carried variables and return the final
+    values after loop termination. The initial values of the variables are
+    passed as additional SSA operands to the "affine.for" following the 3 loop
+    control values lower bound, upper bound and step. The operation region has
+    equivalent arguments for each variable representing the value of the variable
+    at the current iteration.
+
+    The region must terminate with a `affine.yield` that passes all the current
+    iteration variables to the next iteration, or to the `affine.for` result, if
+    at the last iteration.
+
+    `affine.for` results hold the final values after the last iteration.
+    For example, to sum-reduce a memref:
+
+     ```mlir
+    func @reduce(%buffer: memref<1024xf32>) -> (f32) {
+      // Initial sum set to 0.
+      %sum_0 = constant 0.0 : f32
+      // iter_args binds initial values to the loop's region arguments.
+      %sum = affine.for %i = 0 to 10 step 2
+          iter_args(%sum_iter = %sum_0) -> (f32) {
+        %t = affine.load %buffer[%i] : memref<1024xf32>
+        %sum_next = addf %sum_iter, %t : f32
+        // Yield current iteration sum to next iteration %sum_iter or to %sum
+        // if final iteration.
+        affine.yield %sum_next : f32
+      }
+      return %sum : f32
+    }
+    ```
+    If the `affine.for` defines any values, a yield terminator must be
+    explicitly present. The number and types of the "affine.for" results must
+    match the initial values in the `iter_args` binding and the yield operands.
   }];
   let arguments = (ins Variadic<AnyType>);
+  let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
               "int64_t lowerBound, int64_t upperBound, int64_t step = 1, "
-              "function_ref<void(OpBuilder &, Location, Value)> bodyBuilder "
-              "    = nullptr">,
+              "ValueRange iterArgs = llvm::None, function_ref<void(OpBuilder "
+              "&, Location, Value, ValueRange)> bodyBuilder = nullptr">,
     OpBuilder<"OpBuilder &builder, OperationState &result, "
               "ValueRange lbOperands, AffineMap lbMap, "
               "ValueRange ubOperands, AffineMap ubMap, "
-              "int64_t step = 1, "
-              "function_ref<void(OpBuilder &, Location, Value)> bodyBuilder "
-              "    = nullptr">
+              "int64_t step = 1, ValueRange iterArgs = llvm::None, "
+              "function_ref<void(OpBuilder &, Location, Value, ValueRange)> "
+              "bodyBuilder = nullptr">
   ];
 
   let extraClassDeclaration = [{
+    /// Defining the function type we use for building the body of affine.for.
+    using BodyBuilderFn =
+        function_ref<void(OpBuilder &, Location, Value, ValueRange)>;
+
     static StringRef getStepAttrName() { return "step"; }
     static StringRef getLowerBoundAttrName() { return "lower_bound"; }
     static StringRef getUpperBoundAttrName() { return "upper_bound"; }
 
     Value getInductionVar() { return getBody()->getArgument(0); }
+    Block::BlockArgListType getRegionIterArgs() {
+      return getBody()->getArguments().drop_front();
+    }
+    Operation::operand_range getIterOperands() {
+      return getOperands().drop_front(getNumControlOperands());
+    }
 
     // TODO: provide iterators for the lower and upper bound operands
     // if the current access via getLowerBound(), getUpperBound() is too slow.
 
     /// Returns operands for the lower bound map.
     operand_range getLowerBoundOperands();
 
     /// Returns operands for the upper bound map.
let extraClassDeclaration = [{
     /// Set loop step.
     void setStep(int64_t step) {
       assert(step > 0 && "step has to be a positive integer constant");
       auto *context = getLowerBoundMap().getContext();
       setAttr(Identifier::get(getStepAttrName(), context),
               IntegerAttr::get(IndexType::get(context), step));
     }
 
+    /// Returns number of region arguments for loop-carried values.
+    unsigned getNumRegionIterArgs() {
+      return getBody()->getNumArguments() - 1;
+    }
+    /// Number of operands controlling the loop: lb and ub.
+    unsigned getNumControlOperands() { return getOperation()->getNumOperands() 
+          - getNumIterOperands(); }
+
+    /// Get the number of loop-carried values.
+    unsigned getNumIterOperands();
+
     /// Returns true if the lower bound is constant.
     bool hasConstantLowerBound();
     /// Returns true if the upper bound is constant.
     bool hasConstantUpperBound();
     /// Returns true if both bounds are constant.
     bool hasConstantBounds() {
       return hasConstantLowerBound() && hasConstantUpperBound();
     }
let description = [{
     Example:
 
     ```mlir
     %0 = affine.max (d0) -> (1000, d0 + 512) (%i0) : index
     ```
   }];
 }
 
-def AffineParallelOp : Affine_Op<"parallel", 
+def AffineParallelOp : Affine_Op<"parallel",
     [ImplicitAffineTerminator, RecursiveSideEffects,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>]> {
   let summary = "multi-index parallel band operation";
   let description = [{
     The "affine.parallel" operation represents a hyper-rectangular affine
     parallel band, defining multiple SSA values for its induction variables. It
     has one region capturing the parallel band body. The induction variables are
     represented as arguments of this region. These SSA values always have type
let description = [{
     the reduction methods defined in the AtomicRMWKind enum.  The order of
     reduction is unspecified, and lowering may produce any valid ordering.
     Loops with a 0 trip count will produce as a result the identity value
     associated with each reduction (i.e. 0.0 for addf, 1.0 for mulf).  Assign
     reductions for loops with a trip count != 1 produces undefined results.
 
     Note: Calling AffineParallelOp::build will create the required region and
     block, and insert the required terminator if it is trivial (i.e. no values
-    are yielded).  Parsing will also create the required region, block, and 
+    are yielded).  Parsing will also create the required region, block, and
     terminator, even when they are missing from the textual representation.
 
     Example (3x3 valid convolution):
 
     ```mlir
     fuction @conv_2d(%D : memref<100x100xf32>, %K : memref<3x3xf32>) -> (memref<98x98xf32>) {
       %O = alloc memref<98x98xf32>
       affine.parallel (%x, %y) = (0, 0) to (98, 98) {

mlir/lib/Dialect/Affine/EDSC/Builders.cpp

void mlir::edsc::affineLoopBuilder(ValueRange lbs, ValueRange ubs, int64_t step,
   assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
   OpBuilder &builder = ScopedContext::getBuilderRef();
   Location loc = ScopedContext::getLocation();
 
   // Create the actual loop and call the body builder, if provided, after
   // updating the scoped context.
   builder.create<AffineForOp>(
       loc, lbs, builder.getMultiDimIdentityMap(lbs.size()), ubs,
-      builder.getMultiDimIdentityMap(ubs.size()), step,
-      [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv) {
+      builder.getMultiDimIdentityMap(ubs.size()), step, /*iterArgs=*/llvm::None,
+      [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,
+          ValueRange iterArgs) {
         if (bodyBuilderFn) {
           ScopedContext nestedContext(nestedBuilder, nestedLoc);
           OpBuilder::InsertionGuard guard(nestedBuilder);
           bodyBuilderFn(iv);
         }
         nestedBuilder.create<AffineYieldOp>(nestedLoc);
       });
 }

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

LogicalResult AffineDmaWaitOp::fold(ArrayRef<Attribute> cstOperands,
   /// dma_wait(memrefcast) -> dma_wait
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // AffineForOp
 //===----------------------------------------------------------------------===//
 
-void AffineForOp::build(
-    OpBuilder &builder, OperationState &result, ValueRange lbOperands,
-    AffineMap lbMap, ValueRange ubOperands, AffineMap ubMap, int64_t step,
-    function_ref<void(OpBuilder &, Location, Value)> bodyBuilder) {
+/// 'bodyBuilder' is used to build the body of affine.for. If iterArgs and
+/// bodyBuilder are empty/null, we include default terminator op.
+void AffineForOp::build(OpBuilder &builder, OperationState &result,
+                        ValueRange lbOperands, AffineMap lbMap,
+                        ValueRange ubOperands, AffineMap ubMap, int64_t step,
+                        ValueRange iterArgs, BodyBuilderFn bodyBuilder) {
   assert(((!lbMap && lbOperands.empty()) ||
           lbOperands.size() == lbMap.getNumInputs()) &&
          "lower bound operand count does not match the affine map");
   assert(((!ubMap && ubOperands.empty()) ||
           ubOperands.size() == ubMap.getNumInputs()) &&
          "upper bound operand count does not match the affine map");
   assert(step > 0 && "step has to be a positive integer constant");
 
+  result.addOperands(iterArgs);
+  for (Value val : iterArgs)
+    result.addTypes(val.getType());
+
   // Add an attribute for the step.
   result.addAttribute(getStepAttrName(),
                       builder.getIntegerAttr(builder.getIndexType(), step));
 
   // Add the lower bound.
   result.addAttribute(getLowerBoundAttrName(), AffineMapAttr::get(lbMap));
   result.addOperands(lbOperands);
 
   // Add the upper bound.
   result.addAttribute(getUpperBoundAttrName(), AffineMapAttr::get(ubMap));
   result.addOperands(ubOperands);
 
   // Create a region and a block for the body.  The argument of the region is
   // the loop induction variable.
   Region *bodyRegion = result.addRegion();
-  Block *body = new Block;
-  Value inductionVar = body->addArgument(IndexType::get(builder.getContext()));
-  bodyRegion->push_back(body);
-  if (bodyBuilder) {
-    OpBuilder::InsertionGuard guard(builder);
-    builder.setInsertionPointToStart(body);
-    bodyBuilder(builder, result.location, inductionVar);
-  } else {
+  bodyRegion->push_back(new Block);
+  Block &bodyBlock = bodyRegion->front();
+  Value inductionVar = bodyBlock.addArgument(builder.getIndexType());
+  for (Value val : iterArgs)
+    bodyBlock.addArgument(val.getType());
+
+  // Create the default terminator if the builder is not provided and if the
+  // iteration arguments are not provided. Otherwise, leave this to the caller
+  // because we don't know which values to return from the loop.
+  if (iterArgs.empty() && !bodyBuilder) {
     ensureTerminator(*bodyRegion, builder, result.location);
+  } else if (bodyBuilder) {
+    OpBuilder::InsertionGuard guard(builder);
+    builder.setInsertionPointToStart(&bodyBlock);
+    bodyBuilder(builder, result.location, inductionVar,
+                bodyBlock.getArguments().drop_front());
   }
 }
 
-void AffineForOp::build(
-    OpBuilder &builder, OperationState &result, int64_t lb, int64_t ub,
-    int64_t step,
-    function_ref<void(OpBuilder &, Location, Value)> bodyBuilder) {
+void AffineForOp::build(OpBuilder &builder, OperationState &result, int64_t lb,
+                        int64_t ub, int64_t step, ValueRange iterArgs,
+                        BodyBuilderFn bodyBuilder) {
   auto lbMap = AffineMap::getConstantMap(lb, builder.getContext());
   auto ubMap = AffineMap::getConstantMap(ub, builder.getContext());
-  return build(builder, result, {}, lbMap, {}, ubMap, step, bodyBuilder);
+  return build(builder, result, {}, lbMap, {}, ubMap, step, iterArgs,
+               bodyBuilder);
 }
 
 static LogicalResult verify(AffineForOp op) {
   // Check that the body defines as single block argument for the induction
   // variable.
   auto *body = op.getBody();
-  if (body->getNumArguments() != 1 || !body->getArgument(0).getType().isIndex())
+  if (body->getNumArguments() == 0 || !body->getArgument(0).getType().isIndex())
     return op.emitOpError(
         "expected body to have a single index argument for the "
         "induction variable");
 
-  // Verify that there are enough operands for the bounds.
-  AffineMap lowerBoundMap = op.getLowerBoundMap(),
-            upperBoundMap = op.getUpperBoundMap();
-  if (op.getNumOperands() !=
-      (lowerBoundMap.getNumInputs() + upperBoundMap.getNumInputs()))
-    return op.emitOpError(
-        "operand count must match with affine map dimension and symbol count");
-
   // Verify that the bound operands are valid dimension/symbols.
   /// Lower bound.
-  if (failed(verifyDimAndSymbolIdentifiers(op, op.getLowerBoundOperands(),
-                                           op.getLowerBoundMap().getNumDims())))
+  if (op.getLowerBoundMap().getNumInputs() > 0)
+    if (failed(
+            verifyDimAndSymbolIdentifiers(op, op.getLowerBoundOperands(),
       return failure();
   /// Upper bound.
-  if (failed(verifyDimAndSymbolIdentifiers(op, op.getUpperBoundOperands(),
-                                           op.getUpperBoundMap().getNumDims())))
+  if (op.getUpperBoundMap().getNumInputs() > 0)
+    if (failed(
+            verifyDimAndSymbolIdentifiers(op, op.getUpperBoundOperands(),
       return failure();
+
+  unsigned opNumResults = op.getNumResults();
+  if (opNumResults == 0)
+    return success();
+
+  // If ForOp defines values, check that the number and types of the defined
+  // values match ForOp initial iter operands and backedge basic block
+  // arguments.
+  if (op.getNumIterOperands() != opNumResults)
+    return op.emitOpError(
+        "mismatch between the number of loop-carried values and results");
+  if (op.getNumRegionIterArgs() != opNumResults)
+    return op.emitOpError(
+        "mismatch between the number of basic block args and results");
+
   return success();
 }
 
 /// Parse a for operation loop bounds.
 static ParseResult parseBound(bool isLower, OperationState &result,
                               OpAsmParser &p) {
   // 'min' / 'max' prefixes are generally syntactic sugar, but are required if
   // the map has multiple results.
if (parser.parseAttribute(stepAttr, builder.getIndexType(),
       return failure();
 
     if (stepAttr.getValue().getSExtValue() < 0)
       return parser.emitError(
           stepLoc,
           "expected step to be representable as a positive signed integer");
   }
 
+  // Parse the optional initial iteration arguments.
+  SmallVector<OpAsmParser::OperandType, 4> regionArgs, operands;
+  SmallVector<Type, 4> argTypes;
+  regionArgs.push_back(inductionVariable);
+
+  if (succeeded(parser.parseOptionalKeyword("iter_args"))) {
+    // Parse assignment list and results type list.
+    if (parser.parseAssignmentList(regionArgs, operands) ||
+        parser.parseArrowTypeList(result.types))
+      return failure();
+    // Resolve input operands.
+    for (auto operandType : llvm::zip(operands, result.types))
+      if (parser.resolveOperand(std::get<0>(operandType),
+                                std::get<1>(operandType), result.operands))
+        return failure();
+  }
+  // Induction variable.
+  Type indexType = builder.getIndexType();
+  argTypes.push_back(indexType);
+  // Loop carried variables.
+  argTypes.append(result.types.begin(), result.types.end());
   // Parse the body region.
   Region *body = result.addRegion();
-  if (parser.parseRegion(*body, inductionVariable, builder.getIndexType()))
+  if (regionArgs.size() != argTypes.size())
+    return parser.emitError(
+        parser.getNameLoc(),
+        "mismatch between the number of loop-carried values and results");
+  if (parser.parseRegion(*body, regionArgs, argTypes))
     return failure();
 
   AffineForOp::ensureTerminator(*body, builder, result.location);
 
   // Parse the optional attribute list.
   return parser.parseOptionalAttrDict(result.attributes);
 }
 
static void printBound(AffineMapAttr boundMap,
   }
 
   // Print the map and its operands.
   p << boundMap;
   printDimAndSymbolList(boundOperands.begin(), boundOperands.end(),
                         map.getNumDims(), p);
 }
 
+unsigned AffineForOp::getNumIterOperands() {
+  AffineMap lbMap = getLowerBoundMapAttr().getValue();
+  AffineMap ubMap = getUpperBoundMapAttr().getValue();
+
+  unsigned lbOperandNum = lbMap.getNumInputs();
+  unsigned ubOperandNum = ubMap.getNumInputs();
+
+  return getOperation()->getNumOperands() - (lbOperandNum + ubOperandNum);
+}
+
 static void print(OpAsmPrinter &p, AffineForOp op) {
   p << op.getOperationName() << ' ';
   p.printOperand(op.getBody()->getArgument(0));
   p << " = ";
   printBound(op.getLowerBoundMapAttr(), op.getLowerBoundOperands(), "max", p);
   p << " to ";
   printBound(op.getUpperBoundMapAttr(), op.getUpperBoundOperands(), "min", p);
 
   if (op.getStep() != 1)
     p << " step " << op.getStep();
+
+  bool printBlockTerminators = false;
+  if (op.getNumIterOperands() > 0) {
+    p << " iter_args(";
+    auto regionArgs = op.getRegionIterArgs();
+    auto operands = op.getIterOperands();
+
+    llvm::interleaveComma(llvm::zip(regionArgs, operands), p, [&](auto it) {
+      p << std::get<0>(it) << " = " << std::get<1>(it);
+    });
+    p << ") -> (" << op.getResultTypes() << ")";
+    printBlockTerminators = true;
+  }
+
   p.printRegion(op.region(),
-                /*printEntryBlockArgs=*/false,
-                /*printBlockTerminators=*/false);
+                /*printEntryBlockArgs=*/false, printBlockTerminators);
   p.printOptionalAttrDict(op.getAttrs(),
                           /*elidedAttrs=*/{op.getLowerBoundAttrName(),
                                            op.getUpperBoundAttrName(),
                                            op.getStepAttrName()});
 }
 
 /// Fold the constant bounds of a loop.
 static LogicalResult foldLoopBounds(AffineForOp forOp) {
AffineBound AffineForOp::getLowerBound() {
   auto lbMap = getLowerBoundMap();
   return AffineBound(AffineForOp(*this), 0, lbMap.getNumInputs(), lbMap);
 }
 
 AffineBound AffineForOp::getUpperBound() {
   auto lbMap = getLowerBoundMap();
   auto ubMap = getUpperBoundMap();
   return AffineBound(AffineForOp(*this), lbMap.getNumInputs(), getNumOperands(),
                      ubMap);
 }
 
 void AffineForOp::setLowerBound(ValueRange lbOperands, AffineMap map) {
   assert(lbOperands.size() == map.getNumInputs());
   assert(map.getNumResults() >= 1 && "bound map has at least one result");
 
   SmallVector<Value, 4> newOperands(lbOperands.begin(), lbOperands.end());
 
   auto ubOperands = getUpperBoundOperands();
if (lbs.empty()) {
     return;
   }
 
   // Create the loops iteratively and store the induction variables.
   SmallVector<Value, 4> ivs;
   ivs.reserve(lbs.size());
   for (unsigned i = 0, e = lbs.size(); i < e; ++i) {
     // Callback for creating the loop body, always creates the terminator.
-    auto loopBody = [&](OpBuilder &nestedBuilder, Location nestedLoc,
-                        Value iv) {
+    auto loopBody = [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,
+                        ValueRange iterArgs) {
       ivs.push_back(iv);
       // In the innermost loop, call the body builder.
       if (i == e - 1 && bodyBuilderFn) {
         OpBuilder::InsertionGuard nestedGuard(nestedBuilder);
         bodyBuilderFn(nestedBuilder, nestedLoc, ivs);
       }
       nestedBuilder.create<AffineYieldOp>(nestedLoc);
     };
 
     // Delegate actual loop creation to the callback in order to dispatch
     // between constant- and variable-bound loops.
     auto loop = loopCreatorFn(builder, loc, lbs[i], ubs[i], steps[i], loopBody);
     builder.setInsertionPointToStart(loop.getBody());
   }
 }
 
 /// Creates an affine loop from the bounds known to be constants.
-static AffineForOp buildAffineLoopFromConstants(
-    OpBuilder &builder, Location loc, int64_t lb, int64_t ub, int64_t step,
-    function_ref<void(OpBuilder &, Location, Value)> bodyBuilderFn) {
-  return builder.create<AffineForOp>(loc, lb, ub, step, bodyBuilderFn);
+static AffineForOp
+buildAffineLoopFromConstants(OpBuilder &builder, Location loc, int64_t lb,
+                             int64_t ub, int64_t step,
+                             AffineForOp::BodyBuilderFn bodyBuilderFn) {
+  return builder.create<AffineForOp>(loc, lb, ub, step, /*iterArgs=*/llvm::None,
+                                     bodyBuilderFn);
 }
 
 /// Creates an affine loop from the bounds that may or may not be constants.
-static AffineForOp buildAffineLoopFromValues(
-    OpBuilder &builder, Location loc, Value lb, Value ub, int64_t step,
-    function_ref<void(OpBuilder &, Location, Value)> bodyBuilderFn) {
+static AffineForOp
+buildAffineLoopFromValues(OpBuilder &builder, Location loc, Value lb, Value ub,
+                          int64_t step,
+                          AffineForOp::BodyBuilderFn bodyBuilderFn) {
   auto lbConst = lb.getDefiningOp<ConstantIndexOp>();
   auto ubConst = ub.getDefiningOp<ConstantIndexOp>();
   if (lbConst && ubConst)
     return buildAffineLoopFromConstants(builder, loc, lbConst.getValue(),
                                         ubConst.getValue(), step,
                                         bodyBuilderFn);
   return builder.create<AffineForOp>(loc, lb, builder.getDimIdentityMap(), ub,
                                      builder.getDimIdentityMap(), step,
-                                     bodyBuilderFn);
+                                     /*iterArgs=*/llvm::None, bodyBuilderFn);
 }
 
 void mlir::buildAffineLoopNest(
     OpBuilder &builder, Location loc, ArrayRef<int64_t> lbs,
     ArrayRef<int64_t> ubs, ArrayRef<int64_t> steps,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
   buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
                           buildAffineLoopFromConstants);

mlir/test/Dialect/Affine/invalid.mlir

affine.if affine_set<(i)[N] : (i - 2 >= 0, 4 - i >= 0)>(%i)[%c]  {
       %w = affine.apply affine_map<(d0,d1)[s0] -> (d0+d1+s0)> (%i, %i) [%N]
   } else {
     ^bb0(%arg:i32):
       %w = affine.apply affine_map<(d0,d1)[s0] -> (d0-d1+s0)> (%i, %i) [%N]
   }
   return
 }
 
+// -----
+
+func @affine_for_iter_args_mismatch(%buffer: memref<1024xf32>) -> f32 {
+  %sum_0 = constant 0.0 : f32
+  // expected-error@+1 {{mismatch between the number of loop-carried values and results}}
+  %res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_0) -> (f32, f32) {
+    %t = affine.load %buffer[%i] : memref<1024xf32>
+    affine.yield %t : f32
+  }
+  return %res : f32
+}

mlir/test/Dialect/Affine/ops.mlir

%0 = affine.if affine_set<() : ()> () -> f32 {
     affine.yield %zero : f32
   } else {
     // CHECK: affine.yield %[[ZERO]] : f32
     affine.yield %zero : f32
   }
   // CHECK: return %[[OUT]] : f32
   return %0 : f32
 }
+
+// -----
+
+//  Test affine.for with yield values.
+
+#set = affine_set<(d0): (d0 - 10 >= 0)>
+
+// CHECK-LABEL: func @yield_loop
+func @yield_loop(%buffer: memref<1024xf32>) -> f32 {
+  %sum_init_0 = constant 0.0 : f32
+  %res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_init_0) -> f32 {
+    %t = affine.load %buffer[%i] : memref<1024xf32>
+    %sum_next = affine.if #set(%i) -> (f32) {
+      %new_sum = addf %sum_iter, %t : f32
+      affine.yield %new_sum : f32
+    } else {
+      affine.yield %sum_iter : f32
+    }
+    affine.yield %sum_next : f32
+  }
+  return %res : f32
+}
+// CHECK:      %[[const_0:.*]] = constant 0.000000e+00 : f32
+// CHECK-NEXT: %[[output:.*]] = affine.for %{{.*}} = 0 to 10 step 2 iter_args(%{{.*}} = %[[const_0]]) -> (f32) {
+// CHECK:        affine.if #set0(%{{.*}}) -> f32 {
+// CHECK:          affine.yield %{{.*}} : f32
+// CHECK-NEXT:   } else {
+// CHECK-NEXT:     affine.yield %{{.*}} : f32
+// CHECK-NEXT:   }
+// CHECK-NEXT:   affine.yield %{{.*}} : f32
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[output]] : f32
+
+// CHECK-LABEL: func @affine_for_multiple_yield
+func @affine_for_multiple_yield(%buffer: memref<1024xf32>) -> (f32, f32) {
+  %init_0 = constant 0.0 : f32
+  %res1, %res2 = affine.for %i = 0 to 10 step 2 iter_args(%iter_arg1 = %init_0, %iter_arg2 = %init_0) -> (f32, f32) {
+    %t = affine.load %buffer[%i] : memref<1024xf32>
+    %ret1 = addf %t, %iter_arg1 : f32
+    %ret2 = addf %t, %iter_arg2 : f32
+    affine.yield %ret1, %ret2 : f32, f32
+  }
+  return %res1, %res2 : f32, f32
+}
+// CHECK:      %[[const_0:.*]] = constant 0.000000e+00 : f32
+// CHECK-NEXT: %[[output:[0-9]+]]:2 = affine.for %{{.*}} = 0 to 10 step 2 iter_args(%[[iter_arg1:.*]] = %[[const_0]], %[[iter_arg2:.*]] = %[[const_0]]) -> (f32, f32) {
+// CHECK:        %[[res1:.*]] = addf %{{.*}}, %[[iter_arg1]] : f32
+// CHECK-NEXT:   %[[res2:.*]] = addf %{{.*}}, %[[iter_arg2]] : f32
+// CHECK-NEXT:   affine.yield %[[res1]], %[[res2]] : f32, f32
+// CHECK-NEXT: }