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

2

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.

2535

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

2539

nit: Spell out auto here.

2562

Drop the mlir:: here and many other places.

2564

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

2783

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

2730

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

563

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

762–763

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

2654

Please avoid auto here.

[bondhugula]

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

2729–2752

Some comments here for this whole block code.

2733

Comment here.

2740

Avoid auto here please.

2775–2776

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

339

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

 // in the MLIR operation set.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef 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"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 class AffineApplyOp;
 class AffineBound;
 class AffineDimExpr;
 class AffineValueMap;
-class AffineTerminatorOp;
+class AffineYieldOp;
 class FlatAffineConstraints;
 class OpBuilder;
 
 /// A utility function to check if a value is defined at the top level of an
 /// op with trait `AffineScope` or is a region argument for such an op. A value
 /// of index type defined at the top level is always a valid symbol for all its
 /// uses.
 bool isTopLevelValue(Value value);

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

 //===- AffineOps.td - Affine 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 MLIR affine operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef AFFINE_OPS
 #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"
 
 def Affine_Dialect : Dialect {
   let name = "affine";
   let cppNamespace = "";
   let hasConstantMaterializer = 1;
 }
 
 // Base class for Affine dialect ops.
 class Affine_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<Affine_Dialect, mnemonic, traits> {
   // For every affine op, there needs to be a:
   //   * void print(OpAsmPrinter &p, ${C++ class of Op} op)
   //   * LogicalResult verify(${C++ class of Op} op)
   //   * 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); }];
 }
 
-// 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";
   let description = [{
     The affine.apply operation applies an [affine mapping](#affine-expressions)
     to a list of SSA values, yielding a single SSA value. The number of
     dimension and symbol arguments to `affine.apply` must be equal to the
     respective number of dimensional and symbolic inputs to the affine mapping;
let description = [{
 
     lower-bound ::= `max`? affine-map-attribute dim-and-symbol-use-list | shorthand-bound
     upper-bound ::= `min`? affine-map-attribute dim-and-symbol-use-list | shorthand-bound
     shorthand-bound ::= ssa-id | `-`? integer-literal
     ```
 
     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.
 
     The `affine.for` operation executes its body a number of times iterating
     from a lower bound to an upper bound by a stride. The stride, represented by
     `step`, is a positive constant integer which defaults to "1" if not present.
     The lower and upper bounds specify a half-open range: the range includes the
let description = [{
     The [same restrictions](#restrictions-on-dimensions-and-symbols) hold for
     these SSA values as for all bindings of SSA values to dimensions and
     symbols.
 
     The `affine.if` operation contains two regions for the "then" and "else"
     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.
 
     Example:
 
     ```mlir
     #set = affine_set<(d0, d1)[s0]: (d0 - 10 >= 0, s0 - d0 - 9 >= 0,
                                      d1 - 10 >= 0, s0 - d1 - 9 >= 0)>
     func @reduced_domain_example(%A, %X, %N) : (memref<10xi32>, i32, i32) {
       affine.for %i = 0 to %N {
          affine.for %j = 0 to %N {
            %0 = affine.apply #map42(%j)
            %tmp = call @S1(%X, %i, %0)
            affine.if #set(%i, %j)[%N] {
               %1 = affine.apply #map43(%i, %j)
               call @S2(%tmp, %A, %i, %1)
            }
         }
       }
       return
     }
     ```
   }];
   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 = [{
     static StringRef getConditionAttrName() { return "condition"; }
 
     IntegerSet getIntegerSet();
     void setIntegerSet(IntegerSet newSet);
 
let description = [{
     Example:
 
     ```mlir
     %0 = affine.max (d0) -> (1000, d0 + 512) (%i0) : index
     ```
   }];
 }
 
-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
     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
     index, which is the size of the machine word. The strides, represented by
     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:
 
     ```mlir
     affine.parallel (%i, %j) = (0, 0) to (10, 10) step (1, 1) {
       ...
     }
     ```
   }];
 
   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">
   ];
 
   let extraClassDeclaration = [{
     /// Get the number of dimensions.
     unsigned getNumDims();
 
     operand_range getLowerBoundsOperands();
     operand_range getUpperBoundsOperands();
 
     AffineValueMap getLowerBoundsValueMap();
     AffineValueMap getUpperBoundsValueMap();
     AffineValueMap getRangesValueMap();
 
     /// Get ranges as constants, may fail in dynamic case.
     Optional<SmallVector<int64_t, 8>> getConstantRanges();
 
     Block *getBody();
     OpBuilder getBodyBuilder();
     MutableArrayRef<BlockArgument> getIVs() {
       return getBody()->getArguments();
     }
     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 summary = "affine prefetch operation";
   let description = [{
     The "affine.prefetch" op prefetches data from a memref location described
     with an affine subscript similar to affine.load, and has three attributes:
     a read/write specifier, a locality hint, and a cache type specifier as shown
def AffineStoreOp : AffineStoreOpBase<"store"> {
   ];
 
   let extraClassDeclaration = extraClassDeclarationBase;
 
   let hasCanonicalizer = 1;
   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:
-
+    "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.
     ```
-    operation ::= `"affine.terminator"() : () -> ()`
-    ```
-
-    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"> {
   let summary = "affine vector load operation";
   let description = [{
     The "affine.vector_load" is the vector counterpart of
     [affine.load](#affineload-operation). It reads a slice from a
     [MemRef](../LangRef.md#memref-type), supplied as its first operand,

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

 //===- Ops.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 some MLIR standard operations.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 def StandardOps_Dialect : Dialect {
   let name = "std";
def AssumeAlignmentOp : Std_Op<"assume_alignment"> {
   }];
   let arguments = (ins AnyMemRef:$memref,
                        Confined<I32Attr, [IntPositive]>:$alignment);
   let results = (outs);
 
   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",
                      "memref", "value",
                      "$_self.cast<MemRefType>().getElementType()">
     ]> {
   let summary = "atomic read-modify-write operation";
   let description = [{

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

+//===- StandardOpsBase.td - Standard ops 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

 /// Returns true if 'forOp' is parallel.
 bool mlir::isLoopParallel(AffineForOp forOp) {
   // Collect all load and store ops in loop nest rooted at 'forOp'.
   SmallVector<Operation *, 8> loadAndStoreOpInsts;
   auto walkResult = forOp.walk([&](Operation *opInst) -> WalkResult {
     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();
 
     return WalkResult::advance();
   });
 
   // Stop early if the loop has unknown ops with side effects.

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

LogicalResult matchAndRewrite(AffineMaxOp op,
     if (!reduced)
       return failure();
 
     rewriter.replaceOp(op, reduced);
     return success();
   }
 };
 
 /// 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();
   }
 };
 
 class AffineForLowering : public OpRewritePattern<AffineForOp> {
 public:

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

builder.create<AffineForOp>(
       loc, lbs, builder.getMultiDimIdentityMap(lbs.size()), ubs,
       builder.getMultiDimIdentityMap(ubs.size()), step,
       [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv) {
         if (bodyBuilderFn) {
           ScopedContext nestedContext(nestedBuilder, nestedLoc);
           OpBuilder::InsertionGuard guard(nestedBuilder);
           bodyBuilderFn(iv);
         }
-        nestedBuilder.create<AffineTerminatorOp>(nestedLoc);
+        nestedBuilder.create<AffineYieldOp>(nestedLoc);
       });
 }
 
 static std::pair<AffineExpr, Value>
 categorizeValueByAffineType(MLIRContext *context, Value val, unsigned &numDims,
                             unsigned &numSymbols) {
   AffineExpr d;
   Value resultVal = nullptr;

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

 
 namespace {
 /// This is a pattern to fold trivially empty loops.
 struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
   using OpRewritePattern<AffineForOp>::OpRewritePattern;
 
   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);
     return success();
   }
 };
 } // end anonymous namespace
 
for (unsigned i = 0, e = lbs.size(); i < e; ++i) {
     auto loopBody = [&](OpBuilder &nestedBuilder, Location nestedLoc,
                         Value iv) {
       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<AffineTerminatorOp>(nestedLoc);
+      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());
   }
 }
buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
                           buildAffineLoopFromValues);
 }
 
 //===----------------------------------------------------------------------===//
 // AffineIfOp
 //===----------------------------------------------------------------------===//
 
 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);
     rewriter.eraseBlock(ifOp.getElseBlock());
     rewriter.finalizeRootUpdate(ifOp);
     return success();
   }
 };
if (set.getNumDims() != numDims)
     return parser.emitError(
         parser.getNameLoc(),
         "dim operand count and integer set dim count must match");
   if (numDims + set.getNumSymbols() != result.operands.size())
     return parser.emitError(
         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);
   Region *thenRegion = result.addRegion();
   Region *elseRegion = result.addRegion();
 
   // Parse the 'then' region.
   if (parser.parseRegion(*thenRegion, {}, {}))

 }
 
 static void print(OpAsmPrinter &p, AffineIfOp op) {
   auto conditionAttr =
       op.getAttrOfType<IntegerSetAttr>(op.getConditionAttrName());
   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();
   if (!elseRegion.empty()) {
     p << " else";
     p.printRegion(elseRegion,
                   /*printEntryBlockArgs=*/false,
-                  /*printBlockTerminators=*/false);
+                  /*printBlockTerminators=*/op.getNumResults());
   }
 
   // Print the attribute list.
   p.printOptionalAttrDict(op.getAttrs(),
                           /*elidedAttrs=*/op.getConditionAttrName());
 }
 
 IntegerSet AffineIfOp::getIntegerSet() {
   return getAttrOfType<IntegerSetAttr>(getConditionAttrName()).getValue();
 }
 void AffineIfOp::setIntegerSet(IntegerSet newSet) {
   setAttr(getConditionAttrName(), IntegerSetAttr::get(newSet));
 }
 
 void AffineIfOp::setConditional(IntegerSet set, ValueRange operands) {
   setIntegerSet(set);
   getOperation()->setOperands(operands);
 }
 
 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);
+  result.addTypes(resultTypes);
   result.addOperands(args);
   result.addAttribute(getConditionAttrName(), IntegerSetAttr::get(set));
+
   Region *thenRegion = result.addRegion();
-  Region *elseRegion = result.addRegion();
+  builder.createBlock(thenRegion);
+  if (resultTypes.empty())
     AffineIfOp::ensureTerminator(*thenRegion, builder, result.location);
-  if (withElseRegion)
+
+  Region *elseRegion = result.addRegion();
+  if (withElseRegion) {
+    builder.createBlock(elseRegion);
+    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, /*resultTypes=*/{}, set, args,
+                    withElseRegion);
+}
 
 /// Canonicalize an affine if op's conditional (integer set + operands).
 LogicalResult AffineIfOp::fold(ArrayRef<Attribute>,
                                SmallVectorImpl<OpFoldResult> &) {
   auto set = getIntegerSet();
   SmallVector<Value, 4> operands(getOperands());
   canonicalizeSetAndOperands(&set, &operands);
 
LogicalResult AffinePrefetchOp::fold(ArrayRef<Attribute> cstOperands,
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // AffineParallelOp
 //===----------------------------------------------------------------------===//
 
 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));
   auto lbMap = AffineMap::get(0, 0, lbExprs, builder.getContext());
   SmallVector<AffineExpr, 8> ubExprs;
   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");
+  result.addTypes(resultTypes);
+  // Convert the aggOps to integer attributes.
+  SmallVector<Attribute, 4> aggOpAttrs;
+  for (AtomicRMWKind agg : aggOps)
+    aggOpAttrs.push_back(builder.getI64IntegerAttr(static_cast<int64_t>(agg)));
+  result.addAttribute(getAggOpsAttrName(), builder.getArrayAttr(aggOpAttrs));
   result.addAttribute(getLowerBoundsMapAttrName(), AffineMapAttr::get(lbMap));
   result.addAttribute(getUpperBoundsMapAttrName(), AffineMapAttr::get(ubMap));
   result.addAttribute(getStepsAttrName(), builder.getI64ArrayAttr(steps));
   result.addOperands(lbArgs);
   result.addOperands(ubArgs);
   // Create a region and a block for the body.
   auto bodyRegion = result.addRegion();
   auto body = new Block();
   // Add all the block arguments.
   for (unsigned i = 0; i < numDims; ++i)
     body->addArgument(IndexType::get(builder.getContext()));
   bodyRegion->push_back(body);
+  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 (Operation *op : ops)
+    op->moveBefore(*this);
+  return success();
+}
+
 unsigned AffineParallelOp::getNumDims() { return steps().size(); }
 
 AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() {
   return getOperands().take_front(lowerBoundsMap().getNumInputs());
 }
 
 AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() {
   return getOperands().drop_front(lowerBoundsMap().getNumInputs());
 }
 
void AffineParallelOp::setSteps(ArrayRef<int64_t> newSteps) {
   setAttr(getStepsAttrName(), getBodyBuilder().getI64ArrayAttr(newSteps));
 }
 
 static LogicalResult verify(AffineParallelOp op) {
   auto numDims = op.getNumDims();
   if (op.lowerBoundsMap().getNumResults() != numDims ||
       op.upperBoundsMap().getNumResults() != numDims ||
       op.steps().size() != numDims ||
-      op.getBody()->getNumArguments() != numDims) {
+      op.getBody()->getNumArguments() != numDims)
     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 (Attribute 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(),
                                            op.lowerBoundsMap().getNumDims())))
     return failure();
   /// Upper bounds.
   if (failed(verifyDimAndSymbolIdentifiers(op, op.getUpperBoundsOperands(),
                                            op.upperBoundsMap().getNumDims())))
     return failure();
   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 = cast<AffineYieldOp>(std::prev(parallelBodyOps.end()));
+    parentOps.splice(Block::iterator(op), parallelBodyOps,
+                     parallelBodyOps.begin(), std::prev(parallelBodyOps.end()));
+    rewriter.replaceOp(op, yield.operands());
+    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(),
                            op.getLowerBoundsOperands());
   p << ") to (";
   p.printAffineMapOfSSAIds(op.upperBoundsMapAttr(),
                            op.getUpperBoundsOperands());
   p << ')';
   SmallVector<int64_t, 4> steps;
   bool elideSteps = true;
   for (auto attr : op.steps()) {
     auto step = attr.cast<IntegerAttr>().getInt();
     elideSteps &= (step == 1);
     steps.push_back(step);
   }
   if (!elideSteps) {
     p << " step (";
     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()});
 }
 
 //
 // operation ::= `affine.parallel` `(` ssa-ids `)` `=` `(` map-of-ssa-ids `)`
 //               `to` `(` map-of-ssa-ids `)` steps? region attr-dict?
 // steps     ::= `steps` `(` integer-literals `)`
for (const auto &result : stepsMap.getResults()) {
       if (!constExpr)
         return parser.emitError(parser.getNameLoc(),
                                 "steps must be constant integers");
       steps.push_back(constExpr.getValue());
     }
     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();
   SmallVector<Type, 4> types(ivs.size(), indexType);
   if (parser.parseRegion(*body, ivs, types) ||
       parser.parseOptionalAttrDict(result.attributes))
     return failure();
 
   // Add a terminator if none was parsed.
   AffineParallelOp::ensureTerminator(*body, builder, result.location);
   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 it : llvm::zip(results, operands)) {
+    if (std::get<0>(it).getType() != std::get<1>(it).getType())
+      return op.emitOpError()
+             << "types mismatch between yield op and its parent";
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // AffineVectorLoadOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseAffineVectorLoadOp(OpAsmParser &parser,
                                            OperationState &result) {
   auto &builder = parser.getBuilder();
   auto indexTy = builder.getIndexType();
 

 
 /// Verify common invariants of affine.vector_load and affine.vector_store.
 static LogicalResult verifyVectorMemoryOp(Operation *op, MemRefType memrefType,
                                           VectorType vectorType) {
   // Check that memref and vector element types match.
   if (memrefType.getElementType() != vectorType.getElementType())
     return op->emitOpError(
         "requires memref and vector types of the same elemental type");
-
   return success();
 }
 
 static LogicalResult verify(AffineVectorLoadOp op) {
   MemRefType memrefType = op.getMemRefType();
   if (failed(verifyMemoryOpIndexing(
           op.getOperation(),
           op.getAttrOfType<AffineMapAttr>(op.getMapAttrName()),

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

if ((forOp = dyn_cast<AffineForOp>(&*it)) && copyNests.count(forOp) == 0) {
       ++it;
     }
   }
 
   // Generate the copy for the final block range.
   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);
   }
 
   return success();
 }
 
 void AffineDataCopyGeneration::runOnFunction() {

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

if (isMemRefDereferencingOp(op)) {
           }
         }
       }
     }
 
     // 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 (unsigned int i = 0; i < op.getNumOperands(); ++i) {
       auto *operandSrc = op.getOperand(i).getDefiningOp();
 
       LLVM_DEBUG(
           op.getOperand(i).print(llvm::dbgs() << "\nIterating on operand\n"));
void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
   OpBuilder b(forOp.getOperation());
 
   SmallPtrSet<Operation *, 8> opsToHoist;
   SmallVector<Operation *, 8> opsToMove;
 
   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);
         }
       }
     }
   }
 
   // For all instructions that we found to be invariant, place sequentially

mlir/test/Dialect/Affine/canonicalize.mlir

 
 func @drop_duplicate_bounds(%N : index) {
   // affine.for %i = max #lb(%arg0) to min #ub(%arg0)
   affine.for %i = max affine_map<(d0) -> (d0, d0)>(%N) to min affine_map<(d0) -> (d0 + 2, d0 + 2)>(%N) {
     "foo"() : () -> ()
   }
   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

affine.for %x = 0 to 7 {
     affine.parallel (%i, %j) = (0, 0) to (symbol(%y), 100) step (10, 10) {
     }
   }
   return
 }
 
 // -----
 
+// 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 {
     // expected-error@+1 {{requires memref and vector types of the same elemental type}}
     %1 = affine.vector_load %0[%i0] : memref<100xf32>, vector<8xf64>
   }
   return
 }

mlir/test/Dialect/Affine/ops.mlir

 // RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s | FileCheck %s
 // 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
   // CHECK-NEXT: } {some_attr = true}
   //
   // 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: affine.if
   // CHECK-NEXT: } {some_attr = true}
   //
   // GENERIC:      "affine.if"()
-  // GENERIC-NEXT:   "affine.terminator"() : () -> ()
+  // GENERIC-NEXT:   "affine.yield"() : () -> ()
   // GENERIC-NEXT: },  {
   // GENERIC-NEXT: })
   affine.if affine_set<() : ()> () {
   } {some_attr = true}
 
   // CHECK: } else {
   // 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 {
     "foo"() : () -> ()
   } {some_attr = true}
 
   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-DAG: #[[$MAP0:map[0-9]+]] = affine_map<(d0)[s0] -> (1000, d0 + 512, s0)>
 // CHECK-DAG: #[[$MAP1:map[0-9]+]] = affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)>

 "test.func"() ({
 ^bb0(%A : memref<?xf32>):
   affine.load %A[%idx] : memref<?xf32>
   "terminate"() : () -> ()
 }) : () -> ()
 
 // -----
 
-// 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_type_mismatch() -> i32 {
   %0 = "foo"() : ()->f32
   return %0 : f32  // expected-error {{type of return operand 0 ('f32') doesn't match function result type ('i32') in function @return_type_mismatch}}
 }
 
 // -----
 
 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
   }
   return
 }
 
 // -----
 

mlir/test/IR/parser.mlir

 // CHECK-LABEL: func @verbose_if(
 func @verbose_if(%N: index) {
   %c = constant 200 : index
 
   // CHECK: affine.if #set{{.*}}(%{{.*}})[%{{.*}}, %{{.*}}] {
   "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
 }
 
 // CHECK-LABEL: func @terminator_with_regions
 func @terminator_with_regions() {
   // Combine successors and regions in the same operation.