IR changes to add yield semantics for affine.if and affine.parallel, additional affine.parallel canonicalization

[jbruestle]

Summary

IR and parsing changes to allow yield semantics similar to SCF for affine.if, and to provide a method to perform reductions via affine.parallel. Currently this commit does not modify lowering for affine.if (although lowerings for affine.if that do not yield values are still correct). Additionally, affine.parallel does not currently have a lowering.

[rriddle]

Please re-upload with full context.
https://llvm.org/docs/Phabricator.html#requesting-a-review-via-the-web-interface

[jbruestle]

Attempting to fix patch to add context.

[rriddle]

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

1

Please make sure that all lines are wrapped at 80 characters. For this line specifically, you need to shorten it and not wrap.

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

1916

Please drop trivial braces, also why can't you just do result.addTypes(resultTypes)?

1923

nit: Please use builder.createBlock instead.

1929

Same here.

1937

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

2433

nit: Drop this comment, the code is self-explanatory.

2434

Same comment as above.

2439

nit: Please spell out the type instead of auto, and drop the trivial braces.

2442

nit: The code is self-explanatory, so this comment can be dropped. Also, please use complete sentences in comments with correct grammar/punctuation.

2467

nit: Spell out auto here.

2531

Please drop these trivial braces, here and throughout the commit.

2535

nit: Spell out auto here.

2568

Drop the mlir:: here and many other places.

2570

Please do not use replaceAllUsesWith directly, use the pattern rewriter instead.

2792

Please do not use e as an iteration variable, it is very confusing.

[rriddle]

Thanks for reuploading!

[jbruestle]

Updates from river's review.

[flaub]

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

572

To be consistent, use OperationState &result (or whatever is the standard in this file).

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

2404

Parameter comments would be nice here too.

2584

Update comment to match code. Or drop if it's trivial.

[nicolasvasilache]

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

2548

The canonicalization really feels like it should be in a separate revision?

[jbruestle]

Code review changes from flaub, and removed additional canonicalizations which are unrelated to the IR change for a separate diff later.

[dcaballe]

Thanks Jeremy! Some minor comments.

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

540

Since this is about SSA value reductions, should we rename "AtomicRMWKind" to something that doesn't imply memory or an atomic implementation? These reductions could be implemented following different approaches, not only using atomic ops, right? Maybe ReductionKind would work?

545

I'm trying to understand if this is intended to achieve the same as scf.reduce.return. Couldn't we do something similar and keep both SCF and Affine aligned? Or is this intentionally different for some reason?

607

I would suggest using 'reduction' terminology instead of aggregates, to align with scf, OpenMp, and others.

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

2648

Same. I would suggest reduce/reduction to align terminology with scf, OpenMp and others.

[jbruestle]

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

540

I do think ReductionKind would be a better name, but I wanted to minimize unrelated changes such as in atomicrmw, so I was considering do a NFC commit afterwards for both atomicrmw and affine.parallel. Actually, I was considering if it might be good to rename atomicrmw to atomic_reduce or something similar since there is already another more general atomic rmw that has an arbitrary interior and is typically lowered though atomic compare and exchange, and the enum based one is all reductions.

545

SCF uses a somewhat more general mechanism for reductions, but that adds a fair amount of complexity, in particular regarding initialization. I'd rather keep reductions for affine.parallel to be a closed set with known identity elements until we have a reason not to, since it's easier to pattern match to specialized hardware, and makes a number of the transforms easier as well. In practice this set of reduction operations seems to cover all of the ML use cases we've run into.

607

Totally agree, I will rename.

[jbruestle]

Naming change for better consistency as suggested by @dcaballe.

[flaub]

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

539

Typo: acccumulated

561

Multiple plurals, seems a bit confusing.

[jbruestle]

Typo fixes

[bondhugula]

It's nice to see this. Sorry about the delay - I missed this patch. I think this is an important design update to be mentioned on discourse. Could you please post a note there - just the op doc descriptions + the example snippets you have here should be sufficient to start with.

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

307–312

Missing doc description here on the meaning on the yield operands. Note that the semantics of the terminator depend on the parent op holding it, and so this has to be described here. Please also update the example to include affine.yield.

[jbruestle]

Improve docs, rebase off of most recent commit.

[bondhugula]

Could you please add a single line guard in hoistAffineIfOp with a TODO to bail out on any affine.if ops that have results? The loop unswitching will no longer be correct with such affine.if ops.

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

793

This has to be marked ReturnLike.

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

330

Nit: Yield ops (terminators for affine regions) are removed.

[bondhugula]

Could you please update the commit summary to reflect this revision? "reduce" isn't mentioned there. You can drop the "... I do not intend .." part from the summary. I think the last line is also outdated.

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

2568

Please avoid auto here.

[bondhugula]

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

2643–2666

Some comments here for this whole block code.

2647

Comment here.

2654

Avoid auto here please.

2689–2690

You can use variadic isa here.

[jbruestle]

Changes requested by @bondhugula.

[bondhugula]

Update commit summary to say you are adding reduction support for affine.parallel?

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

342

Nit: load -> affine.load (won't parse otherwise).

[bondhugula]

Nit: I think the commit title could be shortened to start with "Add yield semantics ..." ('IR changes to' is obvious).

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

 #define MLIR_DIALECT_AFFINE_IR_AFFINEOPS_H
 
 #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"

 class AffineBound;
 class AffineDimExpr;
 class AffineValueMap;
-class AffineTerminatorOp;
+class AffineYieldOp;
 class FlatAffineConstraints;
 class OpBuilder;
 

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

 #define AFFINE_OPS
 
 include "mlir/Dialect/Affine/IR/AffineOpsBase.td"
+include "mlir/Dialect/StandardOps/IR/StandardOpsBase.td"
 include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"

   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
-// Require regions to have affine terminator.
+// Require regions to have affine.yield.
 def ImplicitAffineTerminator
-    : SingleBlockImplicitTerminator<"AffineTerminatorOp">;
+    : SingleBlockImplicitTerminator<"AffineYieldOp">;
 
 def AffineApplyOp : Affine_Op<"apply", [NoSideEffect]> {
   let summary = "affine apply operation";

 
     The `affine.for` operation represents an affine loop nest. It has one region
     containing its body. This region must contain one block that terminates with
-    [`affine.terminator`](#affineterminator-affineterminatorop). *Note:* when
+    [`affine.yield`](#affineyield-affineyieldop). *Note:* when
     `affine.for` is printed in custom format, the terminator is omitted. The
     block has one argument of [`index`](../LangRef.md#index-type) type that
     represents the induction variable of the loop.

     clauses. The latter may be empty (i.e. contain no blocks), meaning the
     absence of the else clause. When non-empty, both regions must contain
     exactly one block terminating with
-    [`affine.terminator`](#affineterminator-affineterminatorop). *Note:* when `affine.if`
+    [`affine.yield`](#affineyield-operation). *Note:* when `affine.if`
     is printed in custom format, the terminator is omitted. These blocks must
     not have any arguments.
 

     ```
   }];
   let arguments = (ins Variadic<AnyType>);
+  let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion);
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
-              "IntegerSet set, ValueRange args, bool withElseRegion">
+              "IntegerSet set, ValueRange args, bool withElseRegion">,
+    OpBuilder<"OpBuilder &builder, OperationState &result, "
+              "TypeRange resultTypes, IntegerSet set, ValueRange args,"
+              "bool withElseRegion">,
   ];
 
   let extraClassDeclaration = [{

   }];
 }
 
-def AffineParallelOp : Affine_Op<"parallel", [ImplicitAffineTerminator]> {
+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

     steps, are positive constant integers which defaults to "1" if not present.
     The lower and upper bounds specify a half-open range: the range includes the
     lower bound but does not include the upper bound. The body region must
-    contain exactly one block that terminates with "affine.terminator".
+    contain exactly one block that terminates with "affine.yield".
 
     The lower and upper bounds of a parallel operation are represented as an
     application of an affine mapping to a list of SSA values passed to the map.
     The same restrictions hold for these SSA values as for all bindings of SSA
     values to dimensions and symbols.
 
+    Each value yielded by affine.yield will be accumulated by one of the
+    aggregation ops defined in the AtomicRMWKind enum.  The order of
+    accumulation 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 accumulation (i.e. 0.0 for addf, 1.0 for mulf).
+    'assign' accumulation 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. Parsing will also create the
-    required region, block, and terminator, even when they are missing from the
-    textual representation.
+    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 terminator, even when they are missing from the textual
+    representation.
 
     Example:
 

   }];
 
   let arguments = (ins
+     TypedArrayAttrBase<AtomicRMWKindAttr, "list of aggregation ops">:$aggOps,
      AffineMapAttr:$lowerBoundsMap,
      AffineMapAttr:$upperBoundsMap,
      I64ArrayAttr:$steps,
      Variadic<Index>:$mapOperands);
+  let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
   let builders = [
-    OpBuilder<"OpBuilder &builder, OperationState &result,"
+    OpBuilder<"OpBuilder &builder, OperationState& result, "
+              "ArrayRef<Type> resultTypes, ArrayRef<AtomicRMWKind> aggOps, "
               "ArrayRef<int64_t> ranges">,
-    OpBuilder<"OpBuilder &builder, OperationState &result, AffineMap lbMap,"
-              "ValueRange lbArgs, AffineMap ubMap, ValueRange ubArgs">,
-    OpBuilder<"OpBuilder &builder, OperationState &result, AffineMap lbMap,"
-              "ValueRange lbArgs, AffineMap ubMap, ValueRange ubArgs,"
+    OpBuilder<"OpBuilder &builder, OperationState& result, "
+              "ArrayRef<Type> resultTypes, ArrayRef<AtomicRMWKind> aggOps, "
+              "AffineMap lbMap, ValueRange lbArgs, "
+              "AffineMap ubMap, ValueRange ubArgs">,
+    OpBuilder<"OpBuilder &builder, OperationState& result, "
+              "ArrayRef<Type> resultTypes, ArrayRef<AtomicRMWKind> aggOps, "
+              "AffineMap lbMap, ValueRange lbArgs, "
+              "AffineMap ubMap, ValueRange ubArgs, "
               "ArrayRef<int64_t> steps">
   ];
 

     }
     void setSteps(ArrayRef<int64_t> newSteps);
 
+    static StringRef getAggOpsAttrName() { return "aggOps"; }
     static StringRef getLowerBoundsMapAttrName() { return "lowerBoundsMap"; }
     static StringRef getUpperBoundsMapAttrName() { return "upperBoundsMap"; }
     static StringRef getStepsAttrName() { return "steps"; }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def AffinePrefetchOp : Affine_Op<"prefetch"> {

   let hasFolder = 1;
 }
 
-def AffineTerminatorOp :
-    Affine_Op<"terminator", [NoSideEffect, Terminator]> {
-  let summary = "affine terminator operation";
+def AffineYieldOp : Affine_Op<"yield", [NoSideEffect, Terminator]> {
+  let summary = "Yield values to parent operation";
   let description = [{
-    Syntax:
-
-    ```
-    operation ::= `"affine.terminator"() : () -> ()`
+    "affine.yield" yields zero or more SSA values from an affine op region and
+    terminates the region. The semantics of how the values yielded are used
+    is defined by the parent operation.
+    If "affine.yield" has any operands, the operands must match the parent
+    operation's results.
+    If the parent operation defines no values, then the "affine.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.
     ```
-
-    Affine terminator is a special terminator operation for blocks inside affine
-    loops ([`affine.for`](#affinefor-affineforop)) and branches
-    ([`affine.if`](#affineif-affineifop)). It unconditionally transmits the
-    control flow to the successor of the operation enclosing the region.
-
-    *Rationale*: bodies of affine operations are [blocks](../LangRef.md#blocks)
-    that must have terminators. Loops and branches represent structured control
-    flow and should not accept arbitrary branches as terminators.
-
-    This operation does _not_ have a custom syntax. However, affine control
-    operations omit the terminator in their custom syntax for brevity.
   }];
 
-  // No custom parsing/printing form.
-  let parser = ?;
-  let printer = ?;
+  let arguments = (ins Variadic<AnyType>:$operands);
+
+  let builders = [OpBuilder<
+    "OpBuilder &b, OperationState &result", [{ build(b, result, llvm::None); }]
+  >];
 
-  // Fully specified by traits.
-  let verifier = ?;
+  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
 def AffineVectorLoadOp : AffineLoadOpBase<"vector_load"> {

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

 #ifndef STANDARD_OPS
 #define STANDARD_OPS
 
+include "mlir/Dialect/StandardOps/IR/StandardOpsBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"

   let assemblyFormat = "$memref `,` $alignment attr-dict `:` type($memref)";
 }
 
-//===----------------------------------------------------------------------===//
-// AtomicRMWOp
-//===----------------------------------------------------------------------===//
-
-def ATOMIC_RMW_KIND_ADDF    : I64EnumAttrCase<"addf", 0>;
-def ATOMIC_RMW_KIND_ADDI    : I64EnumAttrCase<"addi", 1>;
-def ATOMIC_RMW_KIND_ASSIGN  : I64EnumAttrCase<"assign", 2>;
-def ATOMIC_RMW_KIND_MAXF    : I64EnumAttrCase<"maxf", 3>;
-def ATOMIC_RMW_KIND_MAXS    : I64EnumAttrCase<"maxs", 4>;
-def ATOMIC_RMW_KIND_MAXU    : I64EnumAttrCase<"maxu", 5>;
-def ATOMIC_RMW_KIND_MINF    : I64EnumAttrCase<"minf", 6>;
-def ATOMIC_RMW_KIND_MINS    : I64EnumAttrCase<"mins", 7>;
-def ATOMIC_RMW_KIND_MINU    : I64EnumAttrCase<"minu", 8>;
-def ATOMIC_RMW_KIND_MULF    : I64EnumAttrCase<"mulf", 9>;
-def ATOMIC_RMW_KIND_MULI    : I64EnumAttrCase<"muli", 10>;
-
-def AtomicRMWKindAttr : I64EnumAttr<
-    "AtomicRMWKind", "",
-    [ATOMIC_RMW_KIND_ADDF, ATOMIC_RMW_KIND_ADDI, ATOMIC_RMW_KIND_ASSIGN,
-     ATOMIC_RMW_KIND_MAXF, ATOMIC_RMW_KIND_MAXS, ATOMIC_RMW_KIND_MAXU,
-     ATOMIC_RMW_KIND_MINF, ATOMIC_RMW_KIND_MINS, ATOMIC_RMW_KIND_MINU,
-     ATOMIC_RMW_KIND_MULF, ATOMIC_RMW_KIND_MULI]> {
-  let cppNamespace = "::mlir";
-}
-
 def AtomicRMWOp : Std_Op<"atomic_rmw", [
       AllTypesMatch<["value", "result"]>,
       TypesMatchWith<"value type matches element type of memref",

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

+//===- StandardOpsBase.td - Standard operation definitions -----*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines base support for standard operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef STANDARD_OPS_BASE
+#define STANDARD_OPS_BASE
+
+include "mlir/IR/OpBase.td"
+
+def ATOMIC_RMW_KIND_ADDF    : I64EnumAttrCase<"addf", 0>;
+def ATOMIC_RMW_KIND_ADDI    : I64EnumAttrCase<"addi", 1>;
+def ATOMIC_RMW_KIND_ASSIGN  : I64EnumAttrCase<"assign", 2>;
+def ATOMIC_RMW_KIND_MAXF    : I64EnumAttrCase<"maxf", 3>;
+def ATOMIC_RMW_KIND_MAXS    : I64EnumAttrCase<"maxs", 4>;
+def ATOMIC_RMW_KIND_MAXU    : I64EnumAttrCase<"maxu", 5>;
+def ATOMIC_RMW_KIND_MINF    : I64EnumAttrCase<"minf", 6>;
+def ATOMIC_RMW_KIND_MINS    : I64EnumAttrCase<"mins", 7>;
+def ATOMIC_RMW_KIND_MINU    : I64EnumAttrCase<"minu", 8>;
+def ATOMIC_RMW_KIND_MULF    : I64EnumAttrCase<"mulf", 9>;
+def ATOMIC_RMW_KIND_MULI    : I64EnumAttrCase<"muli", 10>;
+
+def AtomicRMWKindAttr : I64EnumAttr<
+    "AtomicRMWKind", "",
+    [ATOMIC_RMW_KIND_ADDF, ATOMIC_RMW_KIND_ADDI, ATOMIC_RMW_KIND_ASSIGN,
+     ATOMIC_RMW_KIND_MAXF, ATOMIC_RMW_KIND_MAXS, ATOMIC_RMW_KIND_MAXU,
+     ATOMIC_RMW_KIND_MINF, ATOMIC_RMW_KIND_MINS, ATOMIC_RMW_KIND_MINU,
+     ATOMIC_RMW_KIND_MULF, ATOMIC_RMW_KIND_MULI]> {
+  let cppNamespace = "::mlir";
+}
+
+#endif // STANDARD_OPS_BASE

mlir/lib/Analysis/Utils.cpp

     if (isa<AffineReadOpInterface>(opInst) ||
         isa<AffineWriteOpInterface>(opInst))
       loadAndStoreOpInsts.push_back(opInst);
-    else if (!isa<AffineForOp>(opInst) && !isa<AffineTerminatorOp>(opInst) &&
+    else if (!isa<AffineForOp>(opInst) && !isa<AffineYieldOp>(opInst) &&
              !isa<AffineIfOp>(opInst) &&
              !MemoryEffectOpInterface::hasNoEffect(opInst))
       return WalkResult::interrupt();

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

 };
 
 /// Affine terminators are removed.
-class AffineTerminatorLowering : public OpRewritePattern<AffineTerminatorOp> {
+class AffineTerminatorLowering : public OpRewritePattern<AffineYieldOp> {
 public:
-  using OpRewritePattern<AffineTerminatorOp>::OpRewritePattern;
+  using OpRewritePattern<AffineYieldOp>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(AffineTerminatorOp op,
+  LogicalResult matchAndRewrite(AffineYieldOp op,
                                 PatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<scf::YieldOp>(op);
     return success();

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

           OpBuilder::InsertionGuard guard(nestedBuilder);
           bodyBuilderFn(iv);
         }
-        nestedBuilder.create<AffineTerminatorOp>(nestedLoc);
+        nestedBuilder.create<AffineYieldOp>(nestedLoc);
       });
 }
 

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

 
   LogicalResult matchAndRewrite(AffineForOp forOp,
                                 PatternRewriter &rewriter) const override {
-    // Check that the body only contains a terminator.
+    // Check that the body only contains a yield.
     if (!llvm::hasSingleElement(*forOp.getBody()))
       return failure();
     rewriter.eraseOp(forOp);

         OpBuilder::InsertionGuard nestedGuard(nestedBuilder);
         bodyBuilderFn(nestedBuilder, nestedLoc, ivs);
       }
-      nestedBuilder.create<AffineTerminatorOp>(nestedLoc);
+      nestedBuilder.create<AffineYieldOp>(nestedLoc);
     };
 
     // Delegate actual loop creation to the callback in order to dispatch

 //===----------------------------------------------------------------------===//
 
 namespace {
-/// Remove else blocks that have nothing other than the terminator.
+/// Remove else blocks that have nothing other than a zero value yield.
 struct SimplifyDeadElse : public OpRewritePattern<AffineIfOp> {
   using OpRewritePattern<AffineIfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AffineIfOp ifOp,
                                 PatternRewriter &rewriter) const override {
     if (ifOp.elseRegion().empty() ||
-        !llvm::hasSingleElement(*ifOp.getElseBlock()))
+        !llvm::hasSingleElement(*ifOp.getElseBlock()) || ifOp.getNumResults())
       return failure();
 
     rewriter.startRootUpdate(ifOp);

         parser.getNameLoc(),
         "symbol operand count and integer set symbol count must match");
 
+  if (parser.parseOptionalArrowTypeList(result.types))
+    return failure();
+
   // Create the regions for 'then' and 'else'.  The latter must be created even
   // if it remains empty for the validity of the operation.
   result.regions.reserve(2);

   p << "affine.if " << conditionAttr;
   printDimAndSymbolList(op.operand_begin(), op.operand_end(),
                         conditionAttr.getValue().getNumDims(), p);
+  p.printOptionalArrowTypeList(op.getResultTypes());
   p.printRegion(op.thenRegion(),
                 /*printEntryBlockArgs=*/false,
-                /*printBlockTerminators=*/false);
+                /*printBlockTerminators=*/op.getNumResults());
 
   // Print the 'else' regions if it has any blocks.
   auto &elseRegion = op.elseRegion();

     p << " else";
     p.printRegion(elseRegion,
                   /*printEntryBlockArgs=*/false,
-                  /*printBlockTerminators=*/false);
+                  /*printBlockTerminators=*/op.getNumResults());
   }
 
   // Print the attribute list.

 }
 
 void AffineIfOp::build(OpBuilder &builder, OperationState &result,
-                       IntegerSet set, ValueRange args, bool withElseRegion) {
+                       TypeRange resultTypes, IntegerSet set, ValueRange args,
+                       bool withElseRegion) {
+  assert(resultTypes.empty() || withElseRegion);
+  for (Type type : resultTypes) {
+    result.addTypes(type);
+  }
   result.addOperands(args);
   result.addAttribute(getConditionAttrName(), IntegerSetAttr::get(set));
+
   Region *thenRegion = result.addRegion();
+  thenRegion->push_back(new Block());
+  if (resultTypes.empty())
+    AffineIfOp::ensureTerminator(*thenRegion, builder, result.location);
+
   Region *elseRegion = result.addRegion();
-  AffineIfOp::ensureTerminator(*thenRegion, builder, result.location);
-  if (withElseRegion)
-    AffineIfOp::ensureTerminator(*elseRegion, builder, result.location);
+  if (withElseRegion) {
+    elseRegion->push_back(new Block());
+    if (resultTypes.empty())
+      AffineIfOp::ensureTerminator(*elseRegion, builder, result.location);
+  }
+}
+
+void AffineIfOp::build(OpBuilder &builder, OperationState &result,
+                       IntegerSet set, ValueRange args, bool withElseRegion) {
+  AffineIfOp::build(builder, result, {}, set, args, withElseRegion);
 }
 
 /// Canonicalize an affine if op's conditional (integer set + operands).

 //===----------------------------------------------------------------------===//
 
 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
+                             ArrayRef<Type> resultTypes,
+                             ArrayRef<AtomicRMWKind> aggOps,
                              ArrayRef<int64_t> ranges) {
   SmallVector<AffineExpr, 8> lbExprs(ranges.size(),
                                      builder.getAffineConstantExpr(0));

   for (int64_t range : ranges)
     ubExprs.push_back(builder.getAffineConstantExpr(range));
   auto ubMap = AffineMap::get(0, 0, ubExprs, builder.getContext());
-  build(builder, result, lbMap, {}, ubMap, {});
+  build(builder, result, resultTypes, aggOps, lbMap, {}, ubMap, {});
 }
 
 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
-                             AffineMap lbMap, ValueRange lbArgs,
-                             AffineMap ubMap, ValueRange ubArgs) {
+                             ArrayRef<Type> resultTypes,
+                             ArrayRef<AtomicRMWKind> aggOps, AffineMap lbMap,
+                             ValueRange lbArgs, AffineMap ubMap,
+                             ValueRange ubArgs) {
   auto numDims = lbMap.getNumResults();
   // Verify that the dimensionality of both maps are the same.
   assert(numDims == ubMap.getNumResults() &&
          "num dims and num results mismatch");
   // Make default step sizes of 1.
   SmallVector<int64_t, 8> steps(numDims, 1);
-  build(builder, result, lbMap, lbArgs, ubMap, ubArgs, steps);
+  build(builder, result, resultTypes, aggOps, lbMap, lbArgs, ubMap, ubArgs,
+        steps);
 }
 
 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
-                             AffineMap lbMap, ValueRange lbArgs,
-                             AffineMap ubMap, ValueRange ubArgs,
-                             ArrayRef<int64_t> steps) {
+                             ArrayRef<Type> resultTypes,
+                             ArrayRef<AtomicRMWKind> aggOps, AffineMap lbMap,
+                             ValueRange lbArgs, AffineMap ubMap,
+                             ValueRange ubArgs, ArrayRef<int64_t> steps) {
   auto numDims = lbMap.getNumResults();
   // Verify that the dimensionality of the maps matches the number of steps.
   assert(numDims == ubMap.getNumResults() &&
          "num dims and num results mismatch");
   assert(numDims == steps.size() && "num dims and num steps mismatch");
+  // Add the types
+  for (Type type : resultTypes) {
+    result.addTypes(type);
+  }
+  // Convert the aggOps to integer attributes
+  SmallVector<Attribute, 4> aggOpAttrs;
+  for (auto agg : aggOps) {
+    aggOpAttrs.push_back(builder.getI64IntegerAttr(static_cast<int64_t>(agg)));
+  }
+  // Add the attributes
+  result.addAttribute(getAggOpsAttrName(), builder.getArrayAttr(aggOpAttrs));
   result.addAttribute(getLowerBoundsMapAttrName(), AffineMapAttr::get(lbMap));
   result.addAttribute(getUpperBoundsMapAttrName(), AffineMapAttr::get(ubMap));
   result.addAttribute(getStepsAttrName(), builder.getI64ArrayAttr(steps));

   for (unsigned i = 0; i < numDims; ++i)
     body->addArgument(IndexType::get(builder.getContext()));
   bodyRegion->push_back(body);
-  ensureTerminator(*bodyRegion, builder, result.location);
+  if (resultTypes.empty())
+    ensureTerminator(*bodyRegion, builder, result.location);
+}
+
+Region &AffineParallelOp::getLoopBody() { return region(); }
+
+bool AffineParallelOp::isDefinedOutsideOfLoop(Value value) {
+  return !region().isAncestor(value.getParentRegion());
+}
+
+LogicalResult AffineParallelOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
+  for (auto *op : ops)
+    op->moveBefore(*this);
+  return success();
 }
 
 unsigned AffineParallelOp::getNumDims() { return steps().size(); }

     return op.emitOpError("region argument count and num results of upper "
                           "bounds, lower bounds, and steps must all match");
   }
+  if (op.aggOps().size() != op.getNumResults()) {
+    return op.emitOpError("aggregation must be specified for each output");
+  }
+  // Verify agg ops are all valid
+  for (auto attr : op.aggOps()) {
+    auto intAttr = attr.dyn_cast<IntegerAttr>();
+    if (!intAttr || !symbolizeAtomicRMWKind(intAttr.getInt())) {
+      return op.emitOpError("invalid aggregation attribute");
+    }
+  }
+
   // Verify that the bound operands are valid dimension/symbols.
   /// Lower bounds.
   if (failed(verifyDimAndSymbolIdentifiers(op, op.getLowerBoundsOperands(),

   return success();
 }
 
+namespace {
+/// This pattern removes affine.parallel ops with no induction variables.
+struct AffineParallelRank0LoopRemover
+    : public OpRewritePattern<AffineParallelOp> {
+  using OpRewritePattern<AffineParallelOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AffineParallelOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check that there are no induction variables
+    if (op.lowerBoundsMap().getNumResults())
+      return failure();
+    // Remove the affine.parallel wrapper, retain the body in the same location
+    auto &parentOps = rewriter.getInsertionBlock()->getOperations();
+    auto &parallelBodyOps = op.region().front().getOperations();
+    auto yield = mlir::cast<AffineYieldOp>(std::prev(parallelBodyOps.end()));
+    for (auto it : zip(op.getResults(), yield.operands())) {
+      std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
+    }
+    parentOps.splice(mlir::Block::iterator(op), parallelBodyOps,
+                     parallelBodyOps.begin(), std::prev(parallelBodyOps.end()));
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+/// This pattern removes indexes that go over an empty range.
+struct AffineParallelTripCount1IndexRemover
+    : public OpRewritePattern<AffineParallelOp> {
+  using OpRewritePattern<AffineParallelOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AffineParallelOp op,
+                                PatternRewriter &rewriter) const override {
+    auto ranges = op.getRangesValueMap();
+    auto origNumArgs = op.getBody()->getArguments().size();
+    size_t curArgNum = 0;
+    SmallVector<AffineExpr, 6> newLowerBounds;
+    SmallVector<AffineExpr, 6> newUpperBounds;
+    SmallVector<int64_t, 6> newSteps;
+    for (unsigned i = 0; i < origNumArgs; i++) {
+      // Is the range a constant value of 1?
+      auto constExpr = ranges.getResult(i).dyn_cast<AffineConstantExpr>();
+      int64_t step = op.steps()[i].template cast<IntegerAttr>().getInt();
+      if (constExpr && constExpr.getValue() == step) {
+        // Remove argument and replace with lower bound
+        auto curArg = op.getBody()->getArgument(curArgNum);
+        auto lowerBoundValue = rewriter.create<AffineApplyOp>(
+            op.getLoc(), op.lowerBoundsMap().getSubMap({i}),
+            op.getLowerBoundsOperands());
+        curArg.replaceAllUsesWith(lowerBoundValue);
+        op.getBody()->eraseArgument(curArgNum);
+      } else {
+        // Keep argument
+        newLowerBounds.push_back(op.lowerBoundsMap().getResult(i));
+        newUpperBounds.push_back(op.upperBoundsMap().getResult(i));
+        newSteps.push_back(step);
+        curArgNum++;
+      }
+    }
+    // If no arguments were removed, return failure to match
+    if (newLowerBounds.size() == op.lowerBoundsMap().getNumResults())
+      return failure();
+    // Update attributes and return success
+    auto newLower = AffineMap::get(op.lowerBoundsMap().getNumDims(),
+                                   op.lowerBoundsMap().getNumSymbols(),
+                                   newLowerBounds, op.getContext());
+    auto newUpper = AffineMap::get(op.upperBoundsMap().getNumDims(),
+                                   op.upperBoundsMap().getNumSymbols(),
+                                   newUpperBounds, op.getContext());
+    op.setAttr(AffineParallelOp::getLowerBoundsMapAttrName(),
+               AffineMapAttr::get(newLower));
+    op.setAttr(AffineParallelOp::getUpperBoundsMapAttrName(),
+               AffineMapAttr::get(newUpper));
+    op.setAttr(AffineParallelOp::getStepsAttrName(),
+               rewriter.getI64ArrayAttr(newSteps));
+    return success();
+  }
+};
+
+} // end anonymous namespace
+
+void AffineParallelOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<AffineParallelRank0LoopRemover,
+                 AffineParallelTripCount1IndexRemover>(context);
+}
+
 static void print(OpAsmPrinter &p, AffineParallelOp op) {
   p << op.getOperationName() << " (" << op.getBody()->getArguments() << ") = (";
   p.printAffineMapOfSSAIds(op.lowerBoundsMapAttr(),

     llvm::interleaveComma(steps, p);
     p << ')';
   }
+  if (op.getNumResults()) {
+    p << " agg (";
+    llvm::interleaveComma(op.aggOps(), p, [&](auto &attr) {
+      auto sym =
+          *symbolizeAtomicRMWKind(attr.template cast<IntegerAttr>().getInt());
+      p << "\"" << stringifyAtomicRMWKind(sym) << "\"";
+    });
+    p << ") -> (" << op.getResultTypes() << ")";
+  }
+
   p.printRegion(op.region(), /*printEntryBlockArgs=*/false,
-                /*printBlockTerminators=*/false);
+                /*printBlockTerminators=*/op.getNumResults());
   p.printOptionalAttrDict(
       op.getAttrs(),
-      /*elidedAttrs=*/{AffineParallelOp::getLowerBoundsMapAttrName(),
+      /*elidedAttrs=*/{AffineParallelOp::getAggOpsAttrName(),
+                       AffineParallelOp::getLowerBoundsMapAttrName(),
                        AffineParallelOp::getUpperBoundsMapAttrName(),
                        AffineParallelOp::getStepsAttrName()});
 }

     result.addAttribute(AffineParallelOp::getStepsAttrName(),
                         builder.getI64ArrayAttr(steps));
   }
+  SmallVector<Attribute, 4> aggOps;
+  if (succeeded(parser.parseOptionalKeyword("agg"))) {
+    if (parser.parseLParen())
+      return failure();
+    do {
+      StringAttr attrVal;
+      NamedAttrList attrStorage;
+      auto loc = parser.getCurrentLocation();
+      if (parser.parseAttribute(attrVal, builder.getNoneType(), "agg",
+                                attrStorage))
+        return failure();
+      auto attrOptional = symbolizeAtomicRMWKind(attrVal.getValue());
+      if (!attrOptional)
+        return parser.emitError(loc, "invalid aggOp value: ") << attrVal;
+      aggOps.push_back(builder.getI64IntegerAttr(
+          static_cast<int64_t>(attrOptional.getValue())));
+    } while (succeeded(parser.parseOptionalComma()));
+    if (parser.parseRParen())
+      return failure();
+  }
+  if (parser.parseOptionalArrowTypeList(result.types))
+    return failure();
+  result.addAttribute(AffineParallelOp::getAggOpsAttrName(),
+                      builder.getArrayAttr(aggOps));
 
   // Now parse the body.
   Region *body = result.addRegion();

   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// AffineYieldOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(AffineYieldOp op) {
+  auto parentOp = op.getParentOp();
+  auto results = parentOp->getResults();
+  auto operands = op.getOperands();
+
+  if (!isa<AffineParallelOp>(parentOp) && !isa<AffineIfOp>(parentOp) &&
+      !isa<AffineForOp>(parentOp)) {
+    return op.emitOpError()
+           << "affine.terminate only terminates If, For or Parallel regions";
+  }
+  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";
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // AffineVectorLoadOp
 //===----------------------------------------------------------------------===//

   if (memrefType.getElementType() != vectorType.getElementType())
     return op->emitOpError(
         "requires memref and vector types of the same elemental type");
-
   return success();
 }
 

mlir/lib/Dialect/Affine/Transforms/AffineDataCopyGeneration.cpp

   if (curBegin != block->end()) {
     // Can't be a terminator because it would have been skipped above.
     assert(!curBegin->isKnownTerminator() && "can't be a terminator");
-    // Exclude the affine terminator - hence, the std::prev.
+    // Exclude the affine.yield - hence, the std::prev.
     affineDataCopyGenerate(/*begin=*/curBegin, /*end=*/std::prev(block->end()),
                            copyOptions, /*filterMemRef=*/llvm::None, copyNests);
   }

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

     // Insert this op in the defined ops list.
     definedOps.insert(&op);
 
-    if (op.getNumOperands() == 0 && !isa<AffineTerminatorOp>(op)) {
+    if (op.getNumOperands() == 0 && !isa<AffineYieldOp>(op)) {
       LLVM_DEBUG(llvm::dbgs() << "\nNon-constant op with 0 operands\n");
       return false;
     }

   for (auto &op : *loopBody) {
     // We don't hoist for loops.
     if (!isa<AffineForOp>(op)) {
-      if (!isa<AffineTerminatorOp>(op)) {
+      if (!isa<AffineYieldOp>(op)) {
         if (isOpLoopInvariant(op, indVar, definedOps, opsToHoist)) {
           opsToMove.push_back(&op);
         }

mlir/test/Dialect/Affine/canonicalize.mlir

   }
   return
 }
+
+// -----
+// CHECK-LABEL: func @affine_parallel_rank0
+func @affine_parallel_rank0() {
+  // CHECK-NEXT: constant
+  %cst = constant 1.0 : f32
+  // CHECK-NEXT: alloc
+  %0 = alloc() : memref<f32>
+  // CHECK-NEXT: affine.store
+  affine.parallel () = () to () {
+    affine.store %cst, %0[] : memref<f32>
+  }
+  // CHECK-NEXT: return
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @affine_parallel_range1
+func @affine_parallel_range1() {
+  // CHECK-NEXT: constant
+  %cst = constant 1.0 : f32
+  // CHECK-NEXT: alloc
+  %0 = alloc() : memref<2x4xf32>
+  // CHECK-NEXT: affine.store
+  affine.parallel (%i, %j) = (0, 1) to (1, 2) {
+    affine.store %cst, %0[%i, %j] : memref<2x4xf32>
+  }
+  // CHECK-NEXT: return
+  return
+}
+
+// -----
+// CHECK-LABEL: func @affine_parallel_range_step
+func @affine_parallel_range_step() {
+  // CHECK-NEXT: constant
+  %cst = constant 1.0 : f32
+  // CHECK-NEXT: alloc
+  %0 = alloc() : memref<100xf32>
+  // CHECK-NEXT: affine.store
+  affine.parallel (%i) = (0) to (100) step (100) {
+    affine.store %cst, %0[%i] : memref<100xf32>
+  }
+  // CHECK-NEXT: return
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @affine_parallel_partial_range1
+func @affine_parallel_partial_range1() {
+  // CHECK-NEXT: constant
+  %cst = constant 1.0 : f32
+  // CHECK-NEXT: alloc
+  %0 = alloc() : memref<2x4xf32>
+  // CHECK-NEXT: affine.parallel (%{{.*}}) = (0) to (10)
+  affine.parallel (%i, %j) = (0, 1) to (10, 2) {
+    // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}, 1]
+    affine.store %cst, %0[%i, %j] : memref<2x4xf32>
+  }
+  // CHECK: return
+  return
+}

mlir/test/Dialect/Affine/invalid.mlir

 
 // -----
 
+// CHECK-LABEL: @affine_parallel
+
+func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) {
+  %0 = alloc() : memref<100x100xf32>
+  //  expected-error@+1 {{op aggregation must be specified for each output}}
+  %1 = affine.parallel (%i, %j) = (0, 0) to (100, 100) step (10, 10) -> (f32) {
+    %2 = affine.load %0[%i, %j] : memref<100x100xf32>
+    affine.yield %2 : f32
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @affine_parallel
+
+func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) {
+  %0 = alloc() : memref<100x100xf32>
+  //  expected-error@+1 {{invalid aggOp value: "bad"}}
+  %1 = affine.parallel (%i, %j) = (0, 0) to (100, 100) step (10, 10) agg ("bad") -> (f32) {
+    %2 = affine.load %0[%i, %j] : memref<100x100xf32>
+    affine.yield %2 : f32
+  }
+  return
+}
+
+// -----
+// CHECK-LABEL: @affine_parallel
+
+func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) {
+  %0 = alloc() : memref<100x100xi32>
+  %1 = affine.parallel (%i, %j) = (0, 0) to (100, 100) step (10, 10) agg ("minf") -> (f32) {
+    %2 = affine.load %0[%i, %j] : memref<100x100xi32>
+    //  expected-error@+1 {{types mismatch between yield op and its parent}}
+    affine.yield %2 : i32
+  }
+  return
+}
+
+// -----
+
 func @vector_load_invalid_vector_type() {
   %0 = alloc() : memref<100xf32>
   affine.for %i0 = 0 to 16 step 8 {

mlir/test/Dialect/Affine/ops.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -mlir-print-op-generic | FileCheck -check-prefix=GENERIC %s
 
 // Check that the attributes for the affine operations are round-tripped.
-// Check that `affine.terminator` is visible in the generic form.
+// Check that `affine.yield` is visible in the generic form.
 // CHECK-LABEL: @empty
 func @empty() {
   // CHECK: affine.for

   //
   // GENERIC:      "affine.for"()
   // GENERIC-NEXT: ^bb0(%{{.*}}: index):
-  // GENERIC-NEXT:   "affine.terminator"() : () -> ()
+  // GENERIC-NEXT:   "affine.yield"() : () -> ()
   // GENERIC-NEXT: })
   affine.for %i = 0 to 10 {
   } {some_attr = true}

   // CHECK-NEXT: } {some_attr = true}
   //
   // GENERIC:      "affine.if"()
-  // GENERIC-NEXT:   "affine.terminator"() : () -> ()
+  // GENERIC-NEXT:   "affine.yield"() : () -> ()
   // GENERIC-NEXT: },  {
   // GENERIC-NEXT: })
   affine.if affine_set<() : ()> () {

   // CHECK: } {some_attr = true}
   //
   // GENERIC:      "affine.if"()
-  // GENERIC-NEXT:   "affine.terminator"() : () -> ()
+  // GENERIC-NEXT:   "affine.yield"() : () -> ()
   // GENERIC-NEXT: },  {
   // GENERIC-NEXT:   "foo"() : () -> ()
-  // GENERIC-NEXT:   "affine.terminator"() : () -> ()
+  // GENERIC-NEXT:   "affine.yield"() : () -> ()
   // GENERIC-NEXT: })
   affine.if affine_set<() : ()> () {
   } else {

   return
 }
 
-// Check that an explicit affine terminator is not printed in custom format.
+// Check that an explicit affine.yield is not printed in custom format.
 // Check that no extra terminator is introduced.
-// CHECK-LABEL: @affine_terminator
-func @affine_terminator() {
+// CHECK-LABEL: @affine.yield
+func @affine.yield() {
   // CHECK: affine.for
   // CHECK-NEXT: }
   //
   // GENERIC:      "affine.for"() ( {
   // GENERIC-NEXT: ^bb0(%{{.*}}: index):  // no predecessors
-  // GENERIC-NEXT:   "affine.terminator"() : () -> ()
+  // GENERIC-NEXT:   "affine.yield"() : () -> ()
   // GENERIC-NEXT: }) {lower_bound = #map0, step = 1 : index, upper_bound = #map1} : () -> ()
   affine.for %i = 0 to 10 {
-    "affine.terminator"() : () -> ()
+    "affine.yield"() : () -> ()
   }
   return
 }

 
 // -----
 
-// CHECK-LABEL: @parallel
-// CHECK-SAME: (%[[N:.*]]: index)
-func @parallel(%N : index) {
+// CHECK-LABEL: func @parallel
+// CHECK-SAME: (%[[A:.*]]: memref<100x100xf32>, %[[N:.*]]: index)
+func @parallel(%A : memref<100x100xf32>, %N : index) {
   // CHECK: affine.parallel (%[[I0:.*]], %[[J0:.*]]) = (0, 0) to (symbol(%[[N]]), 100) step (10, 10)
   affine.parallel (%i0, %j0) = (0, 0) to (symbol(%N), 100) step (10, 10) {
-    // CHECK-NEXT: affine.parallel (%{{.*}}, %{{.*}}) = (%[[I0]], %[[J0]]) to (%[[I0]] + 10, %[[J0]] + 10)
-    affine.parallel (%i1, %j1) = (%i0, %j0) to (%i0 + 10, %j0 + 10) {
+    // CHECK: affine.parallel (%{{.*}}, %{{.*}}) = (%[[I0]], %[[J0]]) to (%[[I0]] + 10, %[[J0]] + 10) agg ("minf", "maxf") -> (f32, f32)
+    %0:2 = affine.parallel (%i1, %j1) = (%i0, %j0) to (%i0 + 10, %j0 + 10) agg ("minf", "maxf") -> (f32, f32) {
+      %2 = affine.load %A[%i0 + %i0, %j0 + %j1] : memref<100x100xf32>
+      affine.yield %2, %2 : f32, f32
     }
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @affine_if
+func @affine_if() -> f32 {
+  // CHECK: %[[ZERO:.*]] = constant {{.*}} : f32
+  %zero = constant 0.0 : f32
+  // CHECK: %[[OUT:.*]] = affine.if {{.*}}() -> f32 {
+  %0 = affine.if affine_set<() : ()> () -> f32 {
+    // CHECK: affine.yield %[[ZERO]] : f32
+    affine.yield %zero : f32
+  } else {
+    // CHECK: affine.yield %[[ZERO]] : f32
+    affine.yield %zero : f32
+  }
+  // CHECK: return %[[OUT]] : f32
+  return %0 : f32
+}

mlir/test/IR/invalid.mlir

 
 func @return_inside_loop() {
   affine.for %i = 1 to 100 {
-    // expected-error@-1 {{op expects regions to end with 'affine.terminator', found 'std.return'}}
+    // expected-error@-1 {{op expects regions to end with 'affine.yield', found 'std.return'}}
     // expected-note@-2 {{in custom textual format, the absence of terminator implies}}
     return
   }

mlir/test/IR/parser.mlir

   "affine.if"(%c, %N, %c) ({
     // CHECK-NEXT: "add"
     %y = "add"(%c, %N) : (index, index) -> index
-    "affine.terminator"() : () -> ()
+    "affine.yield"() : () -> ()
     // CHECK-NEXT: } else {
   }, { // The else region.
     // CHECK-NEXT: "add"
     %z = "add"(%c, %c) : (index, index) -> index
-    "affine.terminator"() : () -> ()
+    "affine.yield"() : () -> ()
   })
   { condition = #set0 } : (index, index, index) -> ()
   return