[mlir] support max/min lower/upper bounds in affine.parallel

[ftynse]

Summary

This enables to express more complex parallel loops in the affine framework,
for example, in cases of tiling by sizes not dividing loop trip counts perfectly
or inner wavefront parallelism, among others. One can't use affine.max/min
and supply values to the nested loop bounds since the results of such
affine.max/min operations aren't valid symbols. Making them valid symbols
isn't an option since they would introduce selection trees into memref
subscript arithmetic as an unintended and undesired consequence. Also
add support for converting such loops to SCF. Drop some API that isn't used in
the core repo from AffineParallelOp since its semantics becomes ambiguous in
presence of max/min bounds. Loop normalization is currently unavailable for
such loops.

Depends On D101171

[bondhugula]

This is fantastic! Thanks @ftynse The fact that affine.parallel didn't allow min/max for lb/ub meant that dim/symbol requirements would be broken whenever an affine.max/min couldn't be at the top-level (since the result of an affine.max and affine.min isn't a valid symbol (unless at top level) for good reasons). Having this support also means the max/min is unified/composed in the op itself now consistent with other designs. In retrospect, I think this should have been in the design on day 0 itself because otherwise the design just goes in an undesired direction and is invasive to fix later. I'll be happy to review this revision.

[bondhugula]

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

669

Missing updates to the documentation.

733

Some of these should probably have doc comments - likewise above.

mlir/include/mlir/IR/AffineMap.h

265–266

Nit: length -> num ?

mlir/include/mlir/IR/OpImplementation.h

690

Missing doc comment - you can just refer to the one above here / make it relative.

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

434

I wonder when this fails. Is it the floodiv's/ceildiv/mod w.r.t negative values and symbols (semi-affine). This can be addressed separately but floordiv/ceildiv/mod RHS are always expected to be positive - it's UB otherwise. And so you can freely use the same operation to divide as in the case of positive constants.

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

2622–2644

Assert message please although trivial.

2655

same repeated.

2930

Nit: '('.

mlir/lib/IR/AsmPrinter.cpp

2607

Likewise.

[bondhugula]

In D101172#2714274, @bondhugula wrote:

This is fantastic! Thanks @ftynse The fact that affine.parallel didn't allow min/max for lb/ub meant that dim/symbol requirements would be broken whenever an affine.max/min couldn't be at the top-level (since the result of an affine.max and affine.min isn't a valid symbol (unless at top level) for good reasons). Having this support also means the max/min is unified/composed in the op itself now consistent with other designs. In retrospect, I think this should have been in the design on day 0 itself because otherwise the design just goes in an undesired direction and is invasive to fix later. I'll be happy to review this revision.

@ftynse You could actually augment your commit summary. It's not just imperfect tiling or wavefront parallelism, but also plain loop tiling when the trip counts aren't multiple of tile sizes. One can't use affine.max/min and supply values to the intra-tile bounds since the results of such affine.max/min operations aren't valid symbols. Making them valid symbols isn't an option since they'd introduce selection trees into memref subscript arithmetic as an unintended and undesired consequence.

[bondhugula]

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

2992–3000

I think you can do:

pos = llvm::find(...) - uniqueOperands.begin();
if (pos == uniqueOperands.size())
   uniqueOperands.push_back(operand);

2993

Nit: Consider switching to * it - although we use auto it in the codebase, the clang-tidy warnings mean those using in-editor syntax errors (via clangd and clang-tidy checks) would see a highlighted issue that can't be distinguished immediately from other "must fix" warnings.

[bondhugula]

mlir/include/mlir/IR/AffineMap.h

265–266

I think length is fine (as in slice length).

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

2655

same -> space.

2659

assert(!maps.empty() && ..)

2677–2680

lbGroups and ubGroups can't be empty here. An assertion missing somewhere.

3024

Always good to also have an example op here.

3054

Nit: Consider hoisting this out and having an mapOperands.clear() here to avoid repeated allocation.

3079

Good to leave a blank line right below here.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

160

{lowerBoundMap} won't work?

[ftynse]

Address review.

[ftynse]

Thanks for the detailed review!

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

434

Yes, it fails in the case of negative or symbolic RHS for div/mod. This has been around for a while, I remember having written prototype code with selects that supports Euclidean division by a value that can be either positive or negative, with @albertcohen. Not sure if there was a use case for that.

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

2677–2680

They won't be empty if lbMaps, ubMaps are not given how they are constructed in concatMapsSameInput. And we have an assertion for lbMaps, ubMaps.

2992–3000

Nice, thanks!

2993

Yeah, I find this specific clang-tidy suggestion borderline incorrect. auto it intends to say that we don't care about the specific iterator type, auto *it effective forces the iterator to be implemented as typedef elementTy *iterator, which is almost the exact opposite of the original intent.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

160

Nope, template type deduction doesn't work with initializer lists forwarded to implicit constructors.

Candidate template ignored: substitution failure [with OpTy = mlir::AffineParallelOp]: deduced incomplete pack <mlir::ValueTypeRange<mlir::ValueRange>, llvm::SmallVector<mlir::AtomicRMWKind, 6> &, (no value), mlir::ValueRange &, llvm::ArrayRef<mlir::AffineMap>, mlir::ValueRange &, llvm::ArrayRef<long>> for template parameter 'Args'

this is one of the reasons why I kept proposing arguably better interfaces for IR construction, but abandoned given the unreasonable amount of pushback. Building ops feels like one of the worst developer experience solutions I have ever seen.

[ftynse]

In D101172#2714294, @bondhugula wrote:

In D101172#2714274, @bondhugula wrote:

This is fantastic! Thanks @ftynse The fact that affine.parallel didn't allow min/max for lb/ub meant that dim/symbol requirements would be broken whenever an affine.max/min couldn't be at the top-level (since the result of an affine.max and affine.min isn't a valid symbol (unless at top level) for good reasons). Having this support also means the max/min is unified/composed in the op itself now consistent with other designs. In retrospect, I think this should have been in the design on day 0 itself because otherwise the design just goes in an undesired direction and is invasive to fix later. I'll be happy to review this revision.

@ftynse You could actually augment your commit summary. It's not just imperfect tiling or wavefront parallelism, but also plain loop tiling when the trip counts aren't multiple of tile sizes. One can't use affine.max/min and supply values to the intra-tile bounds since the results of such affine.max/min operations aren't valid symbols. Making them valid symbols isn't an option since they'd introduce selection trees into memref subscript arithmetic as an unintended and undesired consequence.

I borrowed parts of your text if you don't mind.

[bondhugula]

Looks great.

[bondhugula]

@ftynse This can be landed.

[ftynse]

No it cannot because it depends on the other commit blocked in review.

[bondhugula]

In D101172#2724879, @ftynse wrote:

No it cannot because it depends on the other commit blocked in review.

It wasn't clear how this revision which adds additional IR support logically depends on the one that was detecting and parallelizing reductions. I realize now this was done after and so affine.parallel builder changes. But if it's not a major change, this can go in first. Anyway, the other one is getting close to landing as well.

[ftynse]

This is a non-trivial rebase which would require splitting out the part of this commit that actually depends on the changes of the previous commit. I have a very strong preference of not wasting my time on such things. If you want, feel free to do the rebase yourself and commit it on my behalf.

[ftynse]

Rebase.

[bondhugula]

In D101172#2724903, @ftynse wrote:

This is a non-trivial rebase which would require splitting out the part of this commit that actually depends on the changes of the previous commit. I have a very strong preference of not wasting my time on such things. If you want, feel free to do the rebase yourself and commit it on my behalf.

There is really no need to - nor was I recommending it. I was just checking whether the rebase was a trivial one.

[ftynse]

More rebase.

[flaub]

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

2699–2705

Perhaps we can keep this helper, but fail if hasMinMaxBounds is true? This helper is used in some downstream code:
https://github.com/plaidml/plaidml/blob/plaidml-v1/pmlc/dialect/pxa/analysis/strides.cc#L80

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

let description = [{
     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.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.
+    values to dimensions and symbols. The list of expressions in each map is
+    interpreted according to the respective bounds group attribute. If a single
+    expression belongs to the group, then the result of this expression is taken
+    as a lower(upper) bound of the corresponding loop induction variable. If
+    multiple expressions belong to the group, then the lower(upper) bound is the
+    max(min) of these values obtained from these expressions. The loop band has
+    as many loops as elements in the group bounds attributes.
 
     Each value yielded by affine.yield will be accumulated/reduced via one of
     the reduction methods defined in the AtomicRMWKind enum.  The order of
     reduction is unspecified, and lowering may produce any valid ordering.
     Loops with a 0 trip count will produce as a result the identity value
     associated with each reduction (i.e. 0.0 for addf, 1.0 for mulf).  Assign
     reductions for loops with a trip count != 1 produces undefined results.
 
func @conv_2d(%D : memref<100x100xf32>, %K : memref<3x3xf32>) -> (memref<98x98xf32>) {
           %3 = mulf %1, %2 : f32
           affine.yield %3 : f32
         }
         affine.store %0, O[%x, %y] : memref<98x98xf32>
       }
       return %O
     }
     ```
+
+    Example (tiling by potentially imperfectly dividing sizes):
+
+    ```mlir
+    affine.parallel (%ii, %jj) = (0, 0) to (%N, %M) step (32, 32) {
+      affine.parallel (%i, %j) = (%ii, %jj)
+                              to (min(%ii + 32, %N), min(%jj + 32, %M)) {
+        call @f(%i, %j) : (index, index) -> ()
+      }
+    }
+    ```
   }];
 
   let arguments = (ins
      TypedArrayAttrBase<AtomicRMWKindAttr, "Reduction ops">:$reductions,
      AffineMapAttr:$lowerBoundsMap,
+     I32ElementsAttr:$lowerBoundsGroups,
      AffineMapAttr:$upperBoundsMap,
+     I32ElementsAttr:$upperBoundsGroups,
      I64ArrayAttr:$steps,
      Variadic<Index>:$mapOperands);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
   let builders = [
     OpBuilder<(ins "TypeRange":$resultTypes,
       "ArrayRef<AtomicRMWKind>":$reductions, "ArrayRef<int64_t>":$ranges)>,
     OpBuilder<(ins "TypeRange":$resultTypes,
-      "ArrayRef<AtomicRMWKind>":$reductions, "AffineMap":$lbMap,
-      "ValueRange":$lbArgs, "AffineMap":$ubMap, "ValueRange":$ubArgs)>,
-    OpBuilder<(ins "TypeRange":$resultTypes,
-      "ArrayRef<AtomicRMWKind>":$reductions, "AffineMap":$lbMap,
-      "ValueRange":$lbArgs, "AffineMap":$ubMap, "ValueRange":$ubArgs,
+      "ArrayRef<AtomicRMWKind>":$reductions, "ArrayRef<AffineMap>":$lbMaps,
+      "ValueRange":$lbArgs, "ArrayRef<AffineMap>":$ubMaps, "ValueRange":$ubArgs,
       "ArrayRef<int64_t>":$steps)>
   ];
 
   let extraClassDeclaration = [{
     /// Get the number of dimensions.
     unsigned getNumDims();
 
-    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();
     }
 
+    /// Returns elements of the loop lower bound.
+    AffineMap getLowerBoundMap(unsigned pos);
     operand_range getLowerBoundsOperands();
     AffineValueMap getLowerBoundsValueMap();
+
+    /// Sets elements of the loop lower bound.
     void setLowerBounds(ValueRange operands, AffineMap map);
     void setLowerBoundsMap(AffineMap map);
 
+    /// Returns elements of the loop upper bound.
+    AffineMap getUpperBoundMap(unsigned pos);
     operand_range getUpperBoundsOperands();
     AffineValueMap getUpperBoundsValueMap();
+
+    /// Sets elements fo the loop upper bound.
     void setUpperBounds(ValueRange operands, AffineMap map);
     void setUpperBoundsMap(AffineMap map);
 
     SmallVector<int64_t, 8> getSteps();
     void setSteps(ArrayRef<int64_t> newSteps);
 
+    /// Returns attribute names to use in op construction. Not expected to be
+    /// used directly.
     static StringRef getReductionsAttrName() { return "reductions"; }
     static StringRef getLowerBoundsMapAttrName() { return "lowerBoundsMap"; }
+    static StringRef getLowerBoundsGroupsAttrName() {
+      return "lowerBoundsGroups";
+    }
     static StringRef getUpperBoundsMapAttrName() { return "upperBoundsMap"; }
+    static StringRef getUpperBoundsGroupsAttrName() {
+      return "upperBoundsGroups";
+    }
     static StringRef getStepsAttrName() { return "steps"; }
+
+    /// Returns `true` if the loop bounds have min/max expressions.
+    bool hasMinMaxBounds() {
+      return lowerBoundsMap().getNumResults() != getNumDims() ||
+             upperBoundsMap().getNumResults() != getNumDims();
+    }
   }];
 
   let hasFolder = 1;
 }
 
 def AffinePrefetchOp : Affine_Op<"prefetch",
   [DeclareOpInterfaceMethods<AffineMapAccessInterface>]> {
   let summary = "affine prefetch operation";

mlir/include/mlir/IR/AffineMap.h

public:
   bool isProjectedPermutation() const;
 
   /// Returns true if the AffineMap represents a symbol-less permutation map.
   bool isPermutation() const;
 
   /// Returns the map consisting of the `resultPos` subset.
   AffineMap getSubMap(ArrayRef<unsigned> resultPos) const;
 
+  /// Returns the map consisting of `length` expressions starting from `start`.
+  AffineMap getSliceMap(unsigned start, unsigned length) const;
+
   /// Returns the map consisting of the most major `numResults` results.
   /// Returns the null AffineMap if `numResults` == 0.
   /// Returns `*this` if `numResults` >= `this->getNumResults()`.
   AffineMap getMajorSubMap(unsigned numResults) const;
 
   /// Returns the map consisting of the most minor `numResults` results.
   /// Returns the null AffineMap if `numResults` == 0.
   /// Returns `*this` if `numResults` >= `this->getNumResults()`.

mlir/include/mlir/IR/OpImplementation.h

public:
 
   /// Prints an affine map of SSA ids, where SSA id names are used in place
   /// of dims/symbols.
   /// Operand values must come from single-result sources, and be valid
   /// dimensions/symbol identifiers according to mlir::isValidDim/Symbol.
   virtual void printAffineMapOfSSAIds(AffineMapAttr mapAttr,
                                       ValueRange operands) = 0;
 
+  /// Prints an affine expression of SSA ids with SSA id names used instead of
+  /// dims and symbols.
+  /// Operand values must come from single-result sources, and be valid
+  /// dimensions/symbol identifiers according to mlir::isValidDim/Symbol.
+  virtual void printAffineExprOfSSAIds(AffineExpr expr, ValueRange dimOperands,
+                                       ValueRange symOperands) = 0;
+
   /// Print an optional arrow followed by a type list.
   template <typename TypeRange>
   void printOptionalArrowTypeList(TypeRange &&types) {
     if (types.begin() != types.end())
       printArrowTypeList(types);
   }
   template <typename TypeRange>
   void printArrowTypeList(TypeRange &&types) {
public:
   /// Parses an affine map attribute where dims and symbols are SSA operands.
   /// Operand values must come from single-result sources, and be valid
   /// dimensions/symbol identifiers according to mlir::isValidDim/Symbol.
   virtual ParseResult
   parseAffineMapOfSSAIds(SmallVectorImpl<OperandType> &operands, Attribute &map,
                          StringRef attrName, NamedAttrList &attrs,
                          Delimiter delimiter = Delimiter::Square) = 0;
 
+  /// Parses an affine expression where dims and symbols are SSA operands.
+  /// Operand values must come from single-result sources, and be valid
+  /// dimensions/symbol identifiers according to mlir::isValidDim/Symbol.
+  virtual ParseResult
+  parseAffineExprOfSSAIds(SmallVectorImpl<OperandType> &dimOperands,
+                          SmallVectorImpl<OperandType> &symbOperands,
+                          AffineExpr &expr) = 0;
+
   //===--------------------------------------------------------------------===//
   // Region Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parses a region. Any parsed blocks are appended to 'region' and must be
   /// moved to the op regions after the op is created. The first block of the
   /// region takes 'arguments' of types 'argTypes'. If 'enableNameShadowing' is
   /// set to true, the argument names are allowed to shadow the names of other

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

public:
 
   LogicalResult matchAndRewrite(AffineParallelOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     SmallVector<Value, 8> steps;
     SmallVector<Value, 8> upperBoundTuple;
     SmallVector<Value, 8> lowerBoundTuple;
     SmallVector<Value, 8> identityVals;
-    // Finding lower and upper bound by expanding the map expression.
-    // Checking if expandAffineMap is not giving NULL.
-    Optional<SmallVector<Value, 8>> lowerBound = expandAffineMap(
-        rewriter, loc, op.lowerBoundsMap(), op.getLowerBoundsOperands());
-    Optional<SmallVector<Value, 8>> upperBound = expandAffineMap(
-        rewriter, loc, op.upperBoundsMap(), op.getUpperBoundsOperands());
-    if (!lowerBound || !upperBound)
-      return failure();
-    upperBoundTuple = *upperBound;
-    lowerBoundTuple = *lowerBound;
+    // Emit IR computing the lower and upper bound by expanding the map
+    // expression.
+    lowerBoundTuple.reserve(op.getNumDims());
+    upperBoundTuple.reserve(op.getNumDims());
+    for (unsigned i = 0, e = op.getNumDims(); i < e; ++i) {
+      Value lower = lowerAffineMapMax(rewriter, loc, op.getLowerBoundMap(i),
+                                      op.getLowerBoundsOperands());
+      if (!lower)
+        return rewriter.notifyMatchFailure(op, "couldn't convert lower bounds");
+      lowerBoundTuple.push_back(lower);
+
+      Value upper = lowerAffineMapMin(rewriter, loc, op.getUpperBoundMap(i),
+                                      op.getUpperBoundsOperands());
+      if (!upper)
+        return rewriter.notifyMatchFailure(op, "couldn't convert upper bounds");
+      upperBoundTuple.push_back(upper);
+    }
     steps.reserve(op.steps().size());
     for (Attribute step : op.steps())
       steps.push_back(rewriter.create<ConstantIndexOp>(
           loc, step.cast<IntegerAttr>().getInt()));
+
     // Get the terminator op.
     Operation *affineParOpTerminator = op.getBody()->getTerminator();
     scf::ParallelOp parOp;
     if (op.results().empty()) {
       // Case with no reduction operations/return values.
       parOp = rewriter.create<scf::ParallelOp>(loc, lowerBoundTuple,
                                                upperBoundTuple, steps,
                                                /*bodyBuilderFn=*/nullptr);

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

 //===----------------------------------------------------------------------===//
 // AffineParallelOp
 //===----------------------------------------------------------------------===//
 
 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
                              TypeRange resultTypes,
                              ArrayRef<AtomicRMWKind> reductions,
                              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, resultTypes, reductions, lbMap, /*lbArgs=*/{}, ubMap,
-        /*ubArgs=*/{});
+  SmallVector<AffineMap> lbs(ranges.size(), builder.getConstantAffineMap(0));
+  auto ubs = llvm::to_vector<4>(llvm::map_range(ranges, [&](int64_t value) {
+    return builder.getConstantAffineMap(value);
+  }));
+  SmallVector<int64_t> steps(ranges.size(), 1);
+  build(builder, result, resultTypes, reductions, lbs, /*lbArgs=*/{}, ubs,
+        /*ubArgs=*/{}, steps);
 }
 
 void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
                              TypeRange resultTypes,
                              ArrayRef<AtomicRMWKind> reductions,
-                             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, resultTypes, reductions, lbMap, lbArgs, ubMap, ubArgs,
-        steps);
-}
-
-void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
-                             TypeRange resultTypes,
-                             ArrayRef<AtomicRMWKind> reductions,
-                             AffineMap lbMap, ValueRange lbArgs,
-                             AffineMap ubMap, ValueRange ubArgs,
+                             ArrayRef<AffineMap> lbMaps, ValueRange lbArgs,
+                             ArrayRef<AffineMap> ubMaps, 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");
+  assert(!lbMaps.empty() && "expected the lower bound map to be non-empty");
+  assert(!ubMaps.empty() && "expected the upper bound map to be non-empty");
+  assert(llvm::all_of(lbMaps,
+                      [lbMaps](AffineMap m) {
+                        return m.getNumDims() == lbMaps[0].getNumDims() &&
+                               m.getNumSymbols() == lbMaps[0].getNumSymbols();
+                      }) &&
+         "expected all lower bounds maps to have the same number of dimensions "
+         "and symbols");
+  assert(llvm::all_of(ubMaps,
+                      [ubMaps](AffineMap m) {
+                        return m.getNumDims() == ubMaps[0].getNumDims() &&
+                               m.getNumSymbols() == ubMaps[0].getNumSymbols();
+                      }) &&
+         "expected all upper bounds maps to have the same number of dimensions "
+         "and symbols");
+  assert(lbMaps[0].getNumInputs() == lbArgs.size() &&
+         "expected lower bound maps to have as many inputs as lower bound "
+         "operands");
+  assert(ubMaps[0].getNumInputs() == ubArgs.size() &&
+         "expected upper bound maps to have as many inputs as upper bound "
+         "operands");
 
   result.addTypes(resultTypes);
+
   // Convert the reductions to integer attributes.
   SmallVector<Attribute, 4> reductionAttrs;
   for (AtomicRMWKind reduction : reductions)
     reductionAttrs.push_back(
         builder.getI64IntegerAttr(static_cast<int64_t>(reduction)));
   result.addAttribute(getReductionsAttrName(),
                       builder.getArrayAttr(reductionAttrs));
+
+  // Concatenates maps defined in the same input space (same dimensions and
+  // symbols), assumes there is at least one map.
+  auto concatMapsSameInput = [](ArrayRef<AffineMap> maps,
+                                SmallVectorImpl<int32_t> &groups) {
+    SmallVector<AffineExpr> exprs;
+    groups.reserve(groups.size() + maps.size());
+    exprs.reserve(maps.size());
+    for (AffineMap m : maps) {
+      llvm::append_range(exprs, m.getResults());
+      groups.push_back(m.getNumResults());
+    }
+    assert(!maps.empty() && "expected a non-empty list of maps");
+    return AffineMap::get(maps[0].getNumDims(), maps[0].getNumSymbols(), exprs,
+                          maps[0].getContext());
+  };
+
+  // Set up the bounds.
+  SmallVector<int32_t> lbGroups, ubGroups;
+  AffineMap lbMap = concatMapsSameInput(lbMaps, lbGroups);
+  AffineMap ubMap = concatMapsSameInput(ubMaps, ubGroups);
   result.addAttribute(getLowerBoundsMapAttrName(), AffineMapAttr::get(lbMap));
+  result.addAttribute(getLowerBoundsGroupsAttrName(),
+                      builder.getI32VectorAttr(lbGroups));
   result.addAttribute(getUpperBoundsMapAttrName(), AffineMapAttr::get(ubMap));
+  result.addAttribute(getUpperBoundsGroupsAttrName(),
+                      builder.getI32VectorAttr(ubGroups));
   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)
+  for (unsigned i = 0, e = steps.size(); i < e; ++i)
     body->addArgument(IndexType::get(builder.getContext()));
   bodyRegion->push_back(body);
   if (resultTypes.empty())
     ensureTerminator(*bodyRegion, builder, result.location);
 }
 
 Region &AffineParallelOp::getLoopBody() { return region(); }
 

 AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() {
   return getOperands().take_front(lowerBoundsMap().getNumInputs());
 }
 
 AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() {
   return getOperands().drop_front(lowerBoundsMap().getNumInputs());
 }
 
+AffineMap AffineParallelOp::getLowerBoundMap(unsigned pos) {
+  unsigned start = 0;
+  for (unsigned i = 0; i < pos; ++i)
+    start += lowerBoundsGroups().getValue<int32_t>(i);
+  return lowerBoundsMap().getSliceMap(
+      start, lowerBoundsGroups().getValue<int32_t>(pos));
+}
+
+AffineMap AffineParallelOp::getUpperBoundMap(unsigned pos) {
+  unsigned start = 0;
+  for (unsigned i = 0; i < pos; ++i)
+    start += upperBoundsGroups().getValue<int32_t>(i);
+  return upperBoundsMap().getSliceMap(
+      start, upperBoundsGroups().getValue<int32_t>(pos));
+}
+
 AffineValueMap AffineParallelOp::getLowerBoundsValueMap() {
   return AffineValueMap(lowerBoundsMap(), getLowerBoundsOperands());
 }
 
 AffineValueMap AffineParallelOp::getUpperBoundsValueMap() {
   return AffineValueMap(upperBoundsMap(), getUpperBoundsOperands());
 }
 
-AffineValueMap AffineParallelOp::getRangesValueMap() {
-  AffineValueMap out;
-  AffineValueMap::difference(getUpperBoundsValueMap(), getLowerBoundsValueMap(),
-                             &out);
-  return out;
-}
-
 Optional<SmallVector<int64_t, 8>> AffineParallelOp::getConstantRanges() {
+  if (hasMinMaxBounds())
+    return llvm::None;
+
   // Try to convert all the ranges to constant expressions.
   SmallVector<int64_t, 8> out;
-  AffineValueMap rangesValueMap = getRangesValueMap();
+  AffineValueMap rangesValueMap;
+  AffineValueMap::difference(getUpperBoundsValueMap(), getLowerBoundsValueMap(),
+                             &rangesValueMap);
   out.reserve(rangesValueMap.getNumResults());
   for (unsigned i = 0, e = rangesValueMap.getNumResults(); i < e; ++i) {
     auto expr = rangesValueMap.getResult(i);
     auto cst = expr.dyn_cast<AffineConstantExpr>();
     if (!cst)
       return llvm::None;
     out.push_back(cst.getValue());
   }

 }
 
 void AffineParallelOp::setSteps(ArrayRef<int64_t> newSteps) {
   stepsAttr(getBodyBuilder().getI64ArrayAttr(newSteps));
 }
 
 static LogicalResult verify(AffineParallelOp op) {
   auto numDims = op.getNumDims();
-  if (op.lowerBoundsMap().getNumResults() != numDims ||
-      op.upperBoundsMap().getNumResults() != numDims ||
+  if (op.lowerBoundsGroups().getNumElements() != numDims ||
+      op.upperBoundsGroups().getNumElements() != numDims ||
       op.steps().size() != numDims ||
-      op.getBody()->getNumArguments() != numDims)
-    return op.emitOpError("region argument count and num results of upper "
-                          "bounds, lower bounds, and steps must all match");
+      op.getBody()->getNumArguments() != numDims) {
+    return op.emitOpError()
+           << "the number of region arguments ("
+           << op.getBody()->getNumArguments()
+           << ") and the number of map groups for lower ("
+           << op.lowerBoundsGroups().getNumElements() << ") and upper bound ("
+           << op.upperBoundsGroups().getNumElements()
+           << "), and the number of steps (" << op.steps().size()
+           << ") must all match";
+  }
+
+  unsigned expectedNumLBResults = 0;
+  for (APInt v : op.lowerBoundsGroups())
+    expectedNumLBResults += v.getZExtValue();
+  if (expectedNumLBResults != op.lowerBoundsMap().getNumResults())
+    return op.emitOpError() << "expected lower bounds map to have "
+                            << expectedNumLBResults << " results";
+  unsigned expectedNumUBResults = 0;
+  for (APInt v : op.upperBoundsGroups())
+    expectedNumUBResults += v.getZExtValue();
+  if (expectedNumUBResults != op.upperBoundsMap().getNumResults())
+    return op.emitOpError() << "expected upper bounds map to have "
+                            << expectedNumUBResults << " results";
 
   if (op.reductions().size() != op.getNumResults())
     return op.emitOpError("a reduction must be specified for each output");
 
   // Verify reduction  ops are all valid
   for (Attribute attr : op.reductions()) {
     auto intAttr = attr.dyn_cast<IntegerAttr>();
     if (!intAttr || !symbolizeAtomicRMWKind(intAttr.getInt()))
static LogicalResult canonicalizeLoopBounds(AffineParallelOp op) {
   return success();
 }
 
 LogicalResult AffineParallelOp::fold(ArrayRef<Attribute> operands,
                                      SmallVectorImpl<OpFoldResult> &results) {
   return canonicalizeLoopBounds(*this);
 }
 
+/// Prints a lower(upper) bound of an affine parallel loop with max(min)
+/// conditions in it. `mapAttr` is a flat list of affine expressions and `group`
+/// identifies which of the those expressions form max/min groups. `operands`
+/// are the SSA values of dimensions and symbols and `keyword` is either "min"
+/// or "max".
+static void printMinMaxBound(OpAsmPrinter &p, AffineMapAttr mapAttr,
+                             DenseIntElementsAttr group, ValueRange operands,
+                             StringRef keyword) {
+  AffineMap map = mapAttr.getValue();
+  unsigned numDims = map.getNumDims();
+  ValueRange dimOperands = operands.take_front(numDims);
+  ValueRange symOperands = operands.drop_front(numDims);
+  unsigned start = 0;
+  for (llvm::APInt groupSize : group) {
+    if (start != 0)
+      p << ", ";
+
+    unsigned size = groupSize.getZExtValue();
+    if (size == 1) {
+      p.printAffineExprOfSSAIds(map.getResult(start), dimOperands, symOperands);
+      ++start;
+    } else {
+      p << keyword << '(';
+      AffineMap submap = map.getSliceMap(start, size);
+      p.printAffineMapOfSSAIds(AffineMapAttr::get(submap), operands);
+      p << ')';
+      start += size;
+    }
+  }
+}
+
 static void print(OpAsmPrinter &p, AffineParallelOp op) {
   p << op.getOperationName() << " (" << op.getBody()->getArguments() << ") = (";
-  p.printAffineMapOfSSAIds(op.lowerBoundsMapAttr(),
-                           op.getLowerBoundsOperands());
+  printMinMaxBound(p, op.lowerBoundsMapAttr(), op.lowerBoundsGroupsAttr(),
+                   op.getLowerBoundsOperands(), "max");
   p << ") to (";
-  p.printAffineMapOfSSAIds(op.upperBoundsMapAttr(),
-                           op.getUpperBoundsOperands());
+  printMinMaxBound(p, op.upperBoundsMapAttr(), op.upperBoundsGroupsAttr(),
+                   op.getUpperBoundsOperands(), "min");
   p << ')';
   SmallVector<int64_t, 8> steps = op.getSteps();
   bool elideSteps = llvm::all_of(steps, [](int64_t step) { return step == 1; });
   if (!elideSteps) {
     p << " step (";
     llvm::interleaveComma(steps, p);
     p << ')';
   }
   if (op.getNumResults()) {
     p << " reduce (";
     llvm::interleaveComma(op.reductions(), p, [&](auto &attr) {
       AtomicRMWKind sym =
           *symbolizeAtomicRMWKind(attr.template cast<IntegerAttr>().getInt());
       p << "\"" << stringifyAtomicRMWKind(sym) << "\"";
     });
     p << ") -> (" << op.getResultTypes() << ")";
   }
 
   p.printRegion(op.region(), /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/op.getNumResults());
   p.printOptionalAttrDict(
       op->getAttrs(),
       /*elidedAttrs=*/{AffineParallelOp::getReductionsAttrName(),
                        AffineParallelOp::getLowerBoundsMapAttrName(),
+                       AffineParallelOp::getLowerBoundsGroupsAttrName(),
                        AffineParallelOp::getUpperBoundsMapAttrName(),
+                       AffineParallelOp::getUpperBoundsGroupsAttrName(),
                        AffineParallelOp::getStepsAttrName()});
 }
 
+/// Given a list of lists of parsed operands, populates `uniqueOperands` with
+/// unique operands. Also populates `replacements with affine expressions of
+/// `kind` that can be used to update affine maps previously accepting a
+/// `operands` to accept `uniqueOperands` instead.
+static void deduplicateAndResolveOperands(
+    OpAsmParser &parser,
+    ArrayRef<SmallVector<OpAsmParser::OperandType>> operands,
+    SmallVectorImpl<Value> &uniqueOperands,
+    SmallVectorImpl<AffineExpr> &replacements, AffineExprKind kind) {
+  assert((kind == AffineExprKind::DimId || kind == AffineExprKind::SymbolId) &&
+         "expected operands to be dim or symbol expression");
+
+  Type indexType = parser.getBuilder().getIndexType();
+  for (const auto &list : operands) {
+    SmallVector<Value> valueOperands;
+    parser.resolveOperands(list, indexType, valueOperands);
+    for (Value operand : valueOperands) {
+      unsigned pos = std::distance(uniqueOperands.begin(),
+                                   llvm::find(uniqueOperands, operand));
+      if (pos == uniqueOperands.size())
+        uniqueOperands.push_back(operand);
+      replacements.push_back(
+          kind == AffineExprKind::DimId
+              ? getAffineDimExpr(pos, parser.getBuilder().getContext())
+              : getAffineSymbolExpr(pos, parser.getBuilder().getContext()));
+    }
+  }
+}
+
+namespace {
+enum class MinMaxKind { Min, Max };
+} // namespace
+
+/// Parses an affine map that can contain a min/max for groups of its results,
+/// e.g., max(expr-1, expr-2), expr-3, max(expr-4, expr-5, expr-6). Populates
+/// `result` attributes with the map (flat list of expressions) and the grouping
+/// (list of integers that specify how many expressions to put into each
+/// min/max) attributes. Deduplicates repeated operands.
+///
+/// parallel-bound       ::= `(` parallel-group-list `)`
+/// parallel-group-list  ::= parallel-group (`,` parallel-group-list)?
+/// parallel-group       ::= simple-group | min-max-group
+/// simple-group         ::= expr-of-ssa-ids
+/// min-max-group        ::= ( `min` | `max` ) `(` expr-of-ssa-ids-list `)`
+/// expr-of-ssa-ids-list ::= expr-of-ssa-ids (`,` expr-of-ssa-id-list)?
+///
+/// Examples:
+///   (%0, min(%1 + %2, %3), %4, min(%5 floordiv 32, %6))
+///   (%0, max(%1 - 2 * %2))
+static ParseResult parseAffineMapWithMinMax(OpAsmParser &parser,
+                                            OperationState &result,
+                                            MinMaxKind kind) {
+  constexpr llvm::StringLiteral tmpAttrName = "__pseudo_bound_map";
+
+  StringRef mapName = kind == MinMaxKind::Min
+                          ? AffineParallelOp::getUpperBoundsMapAttrName()
+                          : AffineParallelOp::getLowerBoundsMapAttrName();
+  StringRef groupsName = kind == MinMaxKind::Min
+                             ? AffineParallelOp::getUpperBoundsGroupsAttrName()
+                             : AffineParallelOp::getLowerBoundsGroupsAttrName();
+
+  if (failed(parser.parseLParen()))
+    return failure();
+
+  if (succeeded(parser.parseOptionalRParen())) {
+    result.addAttribute(
+        mapName, AffineMapAttr::get(parser.getBuilder().getEmptyAffineMap()));
+    result.addAttribute(groupsName, parser.getBuilder().getI32VectorAttr({}));
+    return success();
+  }
+
+  SmallVector<AffineExpr> flatExprs;
+  SmallVector<SmallVector<OpAsmParser::OperandType>> flatDimOperands;
+  SmallVector<SmallVector<OpAsmParser::OperandType>> flatSymOperands;
+  SmallVector<int32_t> numMapsPerGroup;
+  SmallVector<OpAsmParser::OperandType> mapOperands;
+  do {
+    if (succeeded(parser.parseOptionalKeyword(
+            kind == MinMaxKind::Min ? "min" : "max"))) {
+      mapOperands.clear();
+      AffineMapAttr map;
+      if (failed(parser.parseAffineMapOfSSAIds(mapOperands, map, tmpAttrName,
+                                               result.attributes,
+                                               OpAsmParser::Delimiter::Paren)))
+        return failure();
+      result.attributes.erase(tmpAttrName);
+      llvm::append_range(flatExprs, map.getValue().getResults());
+      auto operandsRef = llvm::makeArrayRef(mapOperands);
+      auto dimsRef = operandsRef.take_front(map.getValue().getNumDims());
+      SmallVector<OpAsmParser::OperandType> dims(dimsRef.begin(),
+                                                 dimsRef.end());
+      auto symsRef = operandsRef.drop_front(map.getValue().getNumDims());
+      SmallVector<OpAsmParser::OperandType> syms(symsRef.begin(),
+                                                 symsRef.end());
+      flatDimOperands.append(map.getValue().getNumResults(), dims);
+      flatSymOperands.append(map.getValue().getNumResults(), syms);
+      numMapsPerGroup.push_back(map.getValue().getNumResults());
+    } else {
+      if (failed(parser.parseAffineExprOfSSAIds(flatDimOperands.emplace_back(),
+                                                flatSymOperands.emplace_back(),
+                                                flatExprs.emplace_back())))
+        return failure();
+      numMapsPerGroup.push_back(1);
+    }
+  } while (succeeded(parser.parseOptionalComma()));
+
+  if (failed(parser.parseRParen()))
+    return failure();
+
+  unsigned totalNumDims = 0;
+  unsigned totalNumSyms = 0;
+  for (unsigned i = 0, e = flatExprs.size(); i < e; ++i) {
+    unsigned numDims = flatDimOperands[i].size();
+    unsigned numSyms = flatSymOperands[i].size();
+    flatExprs[i] = flatExprs[i]
+                       .shiftDims(numDims, totalNumDims)
+                       .shiftSymbols(numSyms, totalNumSyms);
+    totalNumDims += numDims;
+    totalNumSyms += numSyms;
+  }
+
+  // Deduplicate map operands.
+  SmallVector<Value> dimOperands, symOperands;
+  SmallVector<AffineExpr> dimRplacements, symRepacements;
+  deduplicateAndResolveOperands(parser, flatDimOperands, dimOperands,
+                                dimRplacements, AffineExprKind::DimId);
+  deduplicateAndResolveOperands(parser, flatSymOperands, symOperands,
+                                symRepacements, AffineExprKind::SymbolId);
+
+  result.operands.append(dimOperands.begin(), dimOperands.end());
+  result.operands.append(symOperands.begin(), symOperands.end());
+
+  Builder &builder = parser.getBuilder();
+  auto flatMap = AffineMap::get(totalNumDims, totalNumSyms, flatExprs,
+                                parser.getBuilder().getContext());
+  flatMap = flatMap.replaceDimsAndSymbols(
+      dimRplacements, symRepacements, dimOperands.size(), symOperands.size());
+
+  result.addAttribute(mapName, AffineMapAttr::get(flatMap));
+  result.addAttribute(groupsName, builder.getI32VectorAttr(numMapsPerGroup));
+  return success();
+}
+
 //
-// operation ::= `affine.parallel` `(` ssa-ids `)` `=` `(` map-of-ssa-ids `)`
-//               `to` `(` map-of-ssa-ids `)` steps? region attr-dict?
+// operation ::= `affine.parallel` `(` ssa-ids `)` `=` parallel-bound
+//               `to` parallel-bound steps? region attr-dict?
 // steps     ::= `steps` `(` integer-literals `)`
 //
 static ParseResult parseAffineParallelOp(OpAsmParser &parser,
                                          OperationState &result) {
   auto &builder = parser.getBuilder();
   auto indexType = builder.getIndexType();
-  AffineMapAttr lowerBoundsAttr, upperBoundsAttr;
   SmallVector<OpAsmParser::OperandType, 4> ivs;
-  SmallVector<OpAsmParser::OperandType, 4> lowerBoundsMapOperands;
-  SmallVector<OpAsmParser::OperandType, 4> upperBoundsMapOperands;
   if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1,
                                      OpAsmParser::Delimiter::Paren) ||
       parser.parseEqual() ||
-      parser.parseAffineMapOfSSAIds(
-          lowerBoundsMapOperands, lowerBoundsAttr,
-          AffineParallelOp::getLowerBoundsMapAttrName(), result.attributes,
-          OpAsmParser::Delimiter::Paren) ||
-      parser.resolveOperands(lowerBoundsMapOperands, indexType,
-                             result.operands) ||
+      parseAffineMapWithMinMax(parser, result, MinMaxKind::Max) ||
       parser.parseKeyword("to") ||
-      parser.parseAffineMapOfSSAIds(
-          upperBoundsMapOperands, upperBoundsAttr,
-          AffineParallelOp::getUpperBoundsMapAttrName(), result.attributes,
-          OpAsmParser::Delimiter::Paren) ||
-      parser.resolveOperands(upperBoundsMapOperands, indexType,
-                             result.operands))
+      parseAffineMapWithMinMax(parser, result, MinMaxKind::Min))
     return failure();
 
   AffineMapAttr stepsMapAttr;
   NamedAttrList stepsAttrs;
   SmallVector<OpAsmParser::OperandType, 4> stepsMapOperands;
   if (failed(parser.parseOptionalKeyword("step"))) {
     SmallVector<int64_t, 4> steps(ivs.size(), 1);
     result.addAttribute(AffineParallelOp::getStepsAttrName(),

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

 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/LoopUtils.h"
 
 using namespace mlir;
 
 void mlir::normalizeAffineParallel(AffineParallelOp op) {
+  // Loops with min/max in bounds are not normalized at the moment.
+  if (op.hasMinMaxBounds())
+    return;
+
   AffineMap lbMap = op.lowerBoundsMap();
   SmallVector<int64_t, 8> steps = op.getSteps();
   // No need to do any work if the parallel op is already normalized.
   bool isAlreadyNormalized =
       llvm::all_of(llvm::zip(steps, lbMap.getResults()), [](auto tuple) {
         int64_t step = std::get<0>(tuple);
         auto lbExpr =
             std::get<1>(tuple).template dyn_cast<AffineConstantExpr>();
         return lbExpr && lbExpr.getValue() == 0 && step == 1;
       });
   if (isAlreadyNormalized)
     return;
 
-  AffineValueMap ranges = op.getRangesValueMap();
+  AffineValueMap ranges;
+  AffineValueMap::difference(op.getUpperBoundsValueMap(),
+                             op.getLowerBoundsValueMap(), &ranges);
   auto builder = OpBuilder::atBlockBegin(op.getBody());
   auto zeroExpr = builder.getAffineConstantExpr(0);
   SmallVector<AffineExpr, 8> lbExprs;
   SmallVector<AffineExpr, 8> ubExprs;
   for (unsigned i = 0, e = steps.size(); i < e; ++i) {
     int64_t step = steps[i];
 
     // Adjust the lower bound to be 0.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

mlir::affineParallelize(AffineForOp forOp,
                         ArrayRef<LoopReduction> parallelReductions) {
   // Fail early if there are iter arguments that are not reductions.
   unsigned numReductions = parallelReductions.size();
   if (numReductions != forOp.getNumIterOperands())
     return failure();
 
   Location loc = forOp.getLoc();
   OpBuilder outsideBuilder(forOp);
-
-  // If a loop has a 'max' in the lower bound, emit it outside the parallel loop
-  // as it does not have implicit 'max' behavior.
   AffineMap lowerBoundMap = forOp.getLowerBoundMap();
   ValueRange lowerBoundOperands = forOp.getLowerBoundOperands();
   AffineMap upperBoundMap = forOp.getUpperBoundMap();
   ValueRange upperBoundOperands = forOp.getUpperBoundOperands();
 
-  bool needsMax = lowerBoundMap.getNumResults() > 1;
-  bool needsMin = upperBoundMap.getNumResults() > 1;
-  AffineMap identityMap;
-  if (needsMax || needsMin) {
-    if (forOp->getParentOp() &&
-        !forOp->getParentOp()->hasTrait<OpTrait::AffineScope>())
-      return failure();
-
-    identityMap = AffineMap::getMultiDimIdentityMap(1, loc->getContext());
-  }
-  if (needsMax) {
-    auto maxOp = outsideBuilder.create<AffineMaxOp>(loc, lowerBoundMap,
-                                                    lowerBoundOperands);
-    lowerBoundMap = identityMap;
-    lowerBoundOperands = maxOp->getResults();
-  }
-
-  // Same for the upper bound.
-  if (needsMin) {
-    auto minOp = outsideBuilder.create<AffineMinOp>(loc, upperBoundMap,
-                                                    upperBoundOperands);
-    upperBoundMap = identityMap;
-    upperBoundOperands = minOp->getResults();
-  }
-
   // Creating empty 1-D affine.parallel op.
   auto reducedValues = llvm::to_vector<4>(llvm::map_range(
       parallelReductions, [](const LoopReduction &red) { return red.value; }));
   auto reductionKinds = llvm::to_vector<4>(llvm::map_range(
       parallelReductions, [](const LoopReduction &red) { return red.kind; }));
   AffineParallelOp newPloop = outsideBuilder.create<AffineParallelOp>(
-      loc, ValueRange(reducedValues).getTypes(), reductionKinds, lowerBoundMap,
-      lowerBoundOperands, upperBoundMap, upperBoundOperands);
+      loc, ValueRange(reducedValues).getTypes(), reductionKinds,
+      llvm::makeArrayRef(lowerBoundMap), lowerBoundOperands,
+      llvm::makeArrayRef(upperBoundMap), upperBoundOperands,
+      llvm::makeArrayRef(forOp.getStep()));
   // Steal the body of the old affine for op.
   newPloop.region().takeBody(forOp.region());
   Operation *yieldOp = &newPloop.getBody()->back();
 
   // Handle the initial values of reductions because the parallel loop always
   // starts from the neutral value.
   SmallVector<Value> newResults;
   newResults.reserve(numReductions);

mlir/lib/IR/AffineMap.cpp

 AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) const {
   SmallVector<AffineExpr, 4> exprs;
   exprs.reserve(resultPos.size());
   for (auto idx : resultPos)
     exprs.push_back(getResult(idx));
   return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
 }
 
+AffineMap AffineMap::getSliceMap(unsigned start, unsigned length) const {
+  return AffineMap::get(getNumDims(), getNumSymbols(),
+                        getResults().slice(start, length), getContext());
+}
+
 AffineMap AffineMap::getMajorSubMap(unsigned numResults) const {
   if (numResults == 0)
     return AffineMap();
   if (numResults > getNumResults())
     return *this;
   return getSubMap(llvm::to_vector<4>(llvm::seq<unsigned>(0, numResults)));
 }
 

mlir/lib/IR/AsmPrinter.cpp

private:
 
   /// Return a null stream as the output stream, this will ignore any data fed
   /// to it.
   raw_ostream &getStream() const override { return os; }
 
   /// The following are hooks of `OpAsmPrinter` that are not necessary for
   /// determining potential aliases.
   void printAffineMapOfSSAIds(AffineMapAttr, ValueRange) override {}
+  void printAffineExprOfSSAIds(AffineExpr, ValueRange, ValueRange) override {}
   void printNewline() override {}
   void printOperand(Value) override {}
   void printOperand(Value, raw_ostream &os) override {
     // Users expect the output string to have at least the prefixed % to signal
     // a value name. To maintain this invariant, emit a name even if it is
     // guaranteed to go unused.
     os << "%";
   }
void shadowRegionArgs(Region &region, ValueRange namesToUse) override {
     state->getSSANameState().shadowRegionArgs(region, namesToUse);
   }
 
   /// Print the given affine map with the symbol and dimension operands printed
   /// inline with the map.
   void printAffineMapOfSSAIds(AffineMapAttr mapAttr,
                               ValueRange operands) override;
 
+  /// Print the given affine expression with the symbol and dimension operands
+  /// printed inline with the expression.
+  void printAffineExprOfSSAIds(AffineExpr expr, ValueRange dimOperands,
+                               ValueRange symOperands) override;
+
   /// Print the given string as a symbol reference.
   void printSymbolName(StringRef symbolRef) override {
     ::printSymbolReference(symbolRef, os);
   }
 
 private:
   /// The number of spaces used for indenting nested operations.
   const static unsigned indentWidth = 2;
if (isSymbol)
       os << ')';
   };
 
   interleaveComma(map.getResults(), [&](AffineExpr expr) {
     printAffineExpr(expr, printValueName);
   });
 }
 
+void OperationPrinter::printAffineExprOfSSAIds(AffineExpr expr,
+                                               ValueRange dimOperands,
+                                               ValueRange symOperands) {
+  auto printValueName = [&](unsigned pos, bool isSymbol) {
+    if (!isSymbol)
+      return printValueID(dimOperands[pos]);
+    os << "symbol(";
+    printValueID(symOperands[pos]);
+    os << ')';
+  };
+  printAffineExpr(expr, printValueName);
+}
+
 //===----------------------------------------------------------------------===//
 // print and dump methods
 //===----------------------------------------------------------------------===//
 
 void Attribute::print(raw_ostream &os) const {
   ModulePrinter(os).printAttribute(*this);
 }
 

mlir/lib/Parser/AffineParser.cpp

AffineParser(ParserState &state, bool allowParsingSSAIds = false,
       : Parser(state), allowParsingSSAIds(allowParsingSSAIds),
         parseElement(parseElement), numDimOperands(0), numSymbolOperands(0) {}
 
   AffineMap parseAffineMapRange(unsigned numDims, unsigned numSymbols);
   ParseResult parseAffineMapOrIntegerSetInline(AffineMap &map, IntegerSet &set);
   IntegerSet parseIntegerSetConstraints(unsigned numDims, unsigned numSymbols);
   ParseResult parseAffineMapOfSSAIds(AffineMap &map,
                                      OpAsmParser::Delimiter delimiter);
+  ParseResult parseAffineExprOfSSAIds(AffineExpr &expr);
   void getDimsAndSymbolSSAIds(SmallVectorImpl<StringRef> &dimAndSymbolSSAIds,
                               unsigned &numDims);
 
 private:
   // Binary affine op parsing.
   AffineLowPrecOp consumeIfLowPrecOp();
   AffineHighPrecOp consumeIfHighPrecOp();
 
if (parseCommaSeparatedListUntil(rightToken, parseElt,
                                    /*allowEmptyList=*/true))
     return failure();
   // Parsed a valid affine map.
   map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
                        exprs, getContext());
   return success();
 }
 
+/// Parse an AffineExpr where the dim and symbol identifiers are SSA ids.
+ParseResult AffineParser::parseAffineExprOfSSAIds(AffineExpr &expr) {
+  expr = parseAffineExpr();
+  return success(expr != nullptr);
+}
+
 /// Parse the range and sizes affine map definition inline.
 ///
 ///  affine-map ::= dim-and-symbol-id-lists `->` multi-dim-affine-expr
 ///
 ///  multi-dim-affine-expr ::= `(` `)`
 ///  multi-dim-affine-expr ::= `(` affine-expr (`,` affine-expr)* `)`
 AffineMap AffineParser::parseAffineMapRange(unsigned numDims,
                                             unsigned numSymbols) {

 /// parse SSA value uses encountered while parsing affine expressions.
 ParseResult
 Parser::parseAffineMapOfSSAIds(AffineMap &map,
                                function_ref<ParseResult(bool)> parseElement,
                                OpAsmParser::Delimiter delimiter) {
   return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
       .parseAffineMapOfSSAIds(map, delimiter);
 }
+
+/// Parse an AffineExpr of SSA ids. The callback `parseElement` is used to parse
+/// SSA value uses encountered while parsing.
+ParseResult
+Parser::parseAffineExprOfSSAIds(AffineExpr &expr,
+                                function_ref<ParseResult(bool)> parseElement) {
+  return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
+      .parseAffineExprOfSSAIds(expr);
+}

mlir/lib/Parser/Parser.h

public:
   ParseResult parseIntegerSetReference(IntegerSet &set);
 
   /// Parse an AffineMap where the dim and symbol identifiers are SSA ids.
   ParseResult
   parseAffineMapOfSSAIds(AffineMap &map,
                          function_ref<ParseResult(bool)> parseElement,
                          OpAsmParser::Delimiter delimiter);
 
+  /// Parse an AffineExpr where dim and symbol identifiers are SSA ids.
+  ParseResult
+  parseAffineExprOfSSAIds(AffineExpr &expr,
+                          function_ref<ParseResult(bool)> parseElement);
+
 protected:
   /// The Parser is subclassed and reinstantiated.  Do not add additional
   /// non-trivial state here, add it to the ParserState class.
   ParserState &state;
 };
 } // end namespace detail
 } // end namespace mlir
 
 #endif // MLIR_LIB_PARSER_PARSER_H

mlir/lib/Parser/Parser.cpp

ParseResult parseAffineMapOfSSAIds(SmallVectorImpl<OperandType> &operands,
     }
 
     // Add dim operands before symbol operands in 'operands'.
     operands.assign(dimOperands.begin(), dimOperands.end());
     operands.append(symOperands.begin(), symOperands.end());
     return success();
   }
 
+  /// Parse an AffineExpr of SSA ids.
+  ParseResult
+  parseAffineExprOfSSAIds(SmallVectorImpl<OperandType> &dimOperands,
+                          SmallVectorImpl<OperandType> &symbOperands,
+                          AffineExpr &expr) override {
+    auto parseElement = [&](bool isSymbol) -> ParseResult {
+      OperandType operand;
+      if (parseOperand(operand))
+        return failure();
+      if (isSymbol)
+        symbOperands.push_back(operand);
+      else
+        dimOperands.push_back(operand);
+      return success();
+    };
+
+    return parser.parseAffineExprOfSSAIds(expr, parseElement);
+  }
+
   //===--------------------------------------------------------------------===//
   // Region Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a region that takes `arguments` of `argTypes` types.  This
   /// effectively defines the SSA values of `arguments` and assigns their type.
   ParseResult parseRegion(Region &region, ArrayRef<OperandType> arguments,
                           ArrayRef<Type> argTypes,

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

affine.parallel (%kx, %ky) = (0, 0) to (2, 2) {
       %2 = affine.load %arg1[%kx, %ky] : memref<3x3xf32>
       %3 = mulf %1, %2 : f32
       affine.store %3, %O[%kx, %ky] : memref<3x3xf32>
   }
   return %O : memref<3x3xf32>
 }
 // CHECK-LABEL: func @affine_parallel_simple
 // CHECK:         %[[LOWER_1:.*]] = constant 0 : index
-// CHECK-NEXT:    %[[LOWER_2:.*]] = constant 0 : index
 // CHECK-NEXT:    %[[UPPER_1:.*]] = constant 2 : index
+// CHECK-NEXT:    %[[LOWER_2:.*]] = constant 0 : index
 // CHECK-NEXT:    %[[UPPER_2:.*]] = constant 2 : index
 // CHECK-NEXT:    %[[STEP_1:.*]] = constant 1 : index
 // CHECK-NEXT:    %[[STEP_2:.*]] = constant 1 : index
 // CHECK-NEXT:    scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[LOWER_1]], %[[LOWER_2]]) to (%[[UPPER_1]], %[[UPPER_2]]) step (%[[STEP_1]], %[[STEP_2]]) {
 // CHECK-NEXT:      %[[VAL_1:.*]] = memref.load
 // CHECK-NEXT:      %[[VAL_2:.*]] = memref.load
 // CHECK-NEXT:      %[[PRODUCT:.*]] = mulf
 // CHECK-NEXT:      store
%0:2 = affine.parallel (%kx, %ky) = (0, 0) to (2, 2) reduce ("addf", "mulf") -> (f32, f32) {
             %3 = mulf %1, %2 : f32
             %4 = addf %1, %2 : f32
             affine.yield %3, %4 : f32, f32
           }
   return %0#0, %0#1 : f32, f32
 }
 // CHECK-LABEL: func @affine_parallel_with_reductions
 // CHECK:         %[[LOWER_1:.*]] = constant 0 : index
-// CHECK-NEXT:    %[[LOWER_2:.*]] = constant 0 : index
 // CHECK-NEXT:    %[[UPPER_1:.*]] = constant 2 : index
+// CHECK-NEXT:    %[[LOWER_2:.*]] = constant 0 : index
 // CHECK-NEXT:    %[[UPPER_2:.*]] = constant 2 : index
 // CHECK-NEXT:    %[[STEP_1:.*]] = constant 1 : index
 // CHECK-NEXT:    %[[STEP_2:.*]] = constant 1 : index
 // CHECK-NEXT:    %[[INIT_1:.*]] = constant 0.000000e+00 : f32
 // CHECK-NEXT:    %[[INIT_2:.*]] = constant 1.000000e+00 : f32
 // CHECK-NEXT:    %[[RES:.*]] = scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[LOWER_1]], %[[LOWER_2]]) to (%[[UPPER_1]], %[[UPPER_2]]) step (%[[STEP_1]], %[[STEP_2]]) init (%[[INIT_1]], %[[INIT_2]]) -> (f32, f32) {
 // CHECK-NEXT:      %[[VAL_1:.*]] = memref.load
 // CHECK-NEXT:      %[[VAL_2:.*]] = memref.load
%0:2 = affine.parallel (%kx, %ky) = (0, 0) to (2, 2) reduce ("addf", "mulf") -> (f64, f64) {
             %3 = mulf %1, %2 : f64
             %4 = addf %1, %2 : f64
             affine.yield %3, %4 : f64, f64
           }
   return %0#0, %0#1 : f64, f64
 }
 // CHECK-LABEL: @affine_parallel_with_reductions_f64
 // CHECK:  %[[LOWER_1:.*]] = constant 0 : index
-// CHECK:  %[[LOWER_2:.*]] = constant 0 : index
 // CHECK:  %[[UPPER_1:.*]] = constant 2 : index
+// CHECK:  %[[LOWER_2:.*]] = constant 0 : index
 // CHECK:  %[[UPPER_2:.*]] = constant 2 : index
 // CHECK:  %[[STEP_1:.*]] = constant 1 : index
 // CHECK:  %[[STEP_2:.*]] = constant 1 : index
 // CHECK:  %[[INIT_1:.*]] = constant 0.000000e+00 : f64
 // CHECK:  %[[INIT_2:.*]] = constant 1.000000e+00 : f64
 // CHECK:  %[[RES:.*]] = scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[LOWER_1]], %[[LOWER_2]]) to (%[[UPPER_1]], %[[UPPER_2]]) step (%[[STEP_1]], %[[STEP_2]]) init (%[[INIT_1]], %[[INIT_2]]) -> (f64, f64) {
 // CHECK:    %[[VAL_1:.*]] = memref.load
 // CHECK:    %[[VAL_2:.*]] = memref.load
%0:2 = affine.parallel (%kx, %ky) = (0, 0) to (2, 2) reduce ("addi", "muli") -> (i64, i64) {
             %3 = muli %1, %2 : i64
             %4 = addi %1, %2 : i64
             affine.yield %3, %4 : i64, i64
           }
   return %0#0, %0#1 : i64, i64
 }
 // CHECK-LABEL: @affine_parallel_with_reductions_i64
 // CHECK:  %[[LOWER_1:.*]] = constant 0 : index
-// CHECK:  %[[LOWER_2:.*]] = constant 0 : index
 // CHECK:  %[[UPPER_1:.*]] = constant 2 : index
+// CHECK:  %[[LOWER_2:.*]] = constant 0 : index
 // CHECK:  %[[UPPER_2:.*]] = constant 2 : index
 // CHECK:  %[[STEP_1:.*]] = constant 1 : index
 // CHECK:  %[[STEP_2:.*]] = constant 1 : index
 // CHECK:  %[[INIT_1:.*]] = constant 0 : i64
 // CHECK:  %[[INIT_2:.*]] = constant 1 : i64
 // CHECK:  %[[RES:.*]] = scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[LOWER_1]], %[[LOWER_2]]) to (%[[UPPER_1]], %[[UPPER_2]]) step (%[[STEP_1]], %[[STEP_2]]) init (%[[INIT_1]], %[[INIT_2]]) -> (i64, i64) {
 // CHECK:    %[[VAL_1:.*]] = memref.load
 // CHECK:    %[[VAL_2:.*]] = memref.load

mlir/test/Dialect/Affine/invalid.mlir

func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
   %0 = affine.max affine_map<(d0) -> (d0)> ()
 
   return
 }
 
 // -----
 
 func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) {
-  // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}}
+  // expected-error@+1 {{the number of region arguments (1) and the number of map groups for lower (2) and upper bound (2), and the number of steps (2) must all match}}
   affine.parallel (%i) = (0, 0) to (100, 100) step (10, 10) {
   }
 }
 
 // -----
 
 func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) {
-  // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}}
+  // expected-error@+1 {{the number of region arguments (2) and the number of map groups for lower (1) and upper bound (2), and the number of steps (2) must all match}}
   affine.parallel (%i, %j) = (0) to (100, 100) step (10, 10) {
   }
 }
 
 // -----
 
 func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) {
-  // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}}
+  // expected-error@+1 {{the number of region arguments (2) and the number of map groups for lower (2) and upper bound (1), and the number of steps (2) must all match}}
   affine.parallel (%i, %j) = (0, 0) to (100) step (10, 10) {
   }
 }
 
 // -----
 
 func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) {
-  // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}}
+  // expected-error@+1 {{the number of region arguments (2) and the number of map groups for lower (2) and upper bound (2), and the number of steps (1) must all match}}
   affine.parallel (%i, %j) = (0, 0) to (100, 100) step (10) {
   }
 }
 
 // -----
 
 func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) {
   affine.for %x = 0 to 7 {

mlir/test/Dialect/Affine/ops.mlir

%0:2 = affine.parallel (%i1, %j1) = (%i0, %j0) to (%i0 + 10, %j0 + 10) reduce ("minf", "maxf") -> (f32, f32) {
       affine.yield %2, %2 : f32, f32
     }
   }
   return
 }
 
 // -----
 
+// CHECK-LABEL: @parallel_min_max
+// CHECK: %[[A:.*]]: index, %[[B:.*]]: index, %[[C:.*]]: index, %[[D:.*]]: index
+func @parallel_min_max(%a: index, %b: index, %c: index, %d: index) {
+  // CHECK: affine.parallel (%{{.*}}, %{{.*}}, %{{.*}}) =
+  // CHECK:                 (max(%[[A]], %[[B]])
+  // CHECK:              to (%[[C]], min(%[[C]], %[[D]]), %[[B]])
+  affine.parallel (%i, %j, %k) = (max(%a, %b), %b, max(%a, %c))
+                              to (%c, min(%c, %d), %b) {
+    affine.yield
+  }
+  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

mlir/test/Dialect/Affine/parallelize.mlir

// CHECK:  affine.for %{{.*}} = 0 to 100 {
     memref.load %0[%i] : memref<100 x f32>
   }
   return
 }
 
 // CHECK-LABEL: for_with_minmax
 func @for_with_minmax(%m: memref<?xf32>, %lb0: index, %lb1: index,
                       %ub0: index, %ub1: index) {
-  // CHECK: %[[lb:.*]] = affine.max
-  // CHECK: %[[ub:.*]] = affine.min
-  // CHECK: affine.parallel (%{{.*}}) = (%[[lb]]) to (%[[ub]])
+  // CHECK: affine.parallel (%{{.*}}) = (max(%{{.*}}, %{{.*}})) to (min(%{{.*}}, %{{.*}}))
   affine.for %i = max affine_map<(d0, d1) -> (d0, d1)>(%lb0, %lb1)
           to min affine_map<(d0, d1) -> (d0, d1)>(%ub0, %ub1) {
     affine.load %m[%i] : memref<?xf32>
   }
   return
 }
 
 // CHECK-LABEL: nested_for_with_minmax
 func @nested_for_with_minmax(%m: memref<?xf32>, %lb0: index,
                              %ub0: index, %ub1: index) {
-  // CHECK: affine.parallel
+  // CHECK: affine.parallel (%[[I:.*]]) =
   affine.for %j = 0 to 10 {
-    // Cannot parallelize the inner loop because we would need to compute
-    // affine.max for its lower bound inside the loop, and that is not (yet)
-    // considered as a valid affine dimension.
-    // CHECK: affine.for
+    // CHECK: affine.parallel (%{{.*}}) = (max(%{{.*}}, %[[I]])) to (min(%{{.*}}, %{{.*}}))
     affine.for %i = max affine_map<(d0, d1) -> (d0, d1)>(%lb0, %j)
             to min affine_map<(d0, d1) -> (d0, d1)>(%ub0, %ub1) {
       affine.load %m[%i] : memref<?xf32>
     }
   }
   return
 }
 
func @use_in_backward_slice() {
   // REDUCE-NOT: affine.parallel
   affine.for %i = 0 to 10 iter_args(%it1 = %cst1, %it2 = %cst2) -> (f32, f32) {
     %0 = "test.some_modification"(%it2) : (f32) -> f32
     %1 = addf %it1, %0 : f32
     affine.yield %1, %1 : f32, f32
   }
   return
 }
+
+// REDUCE-LABEL: @nested_min_max
+// CHECK-LABEL: @nested_min_max
+// CHECK: (%{{.*}}, %[[LB0:.*]]: index, %[[UB0:.*]]: index, %[[UB1:.*]]: index)
+func @nested_min_max(%m: memref<?xf32>, %lb0: index,
+                     %ub0: index, %ub1: index) {
+  // CHECK: affine.parallel (%[[J:.*]]) =
+  affine.for %j = 0 to 10 {
+    // CHECK: affine.parallel (%{{.*}}) = (max(%[[LB0]], %[[J]]))
+    // CHECK:                          to (min(%[[UB0]], %[[UB1]]))
+    affine.for %i = max affine_map<(d0, d1) -> (d0, d1)>(%lb0, %j)
+            to min affine_map<(d0, d1) -> (d0, d1)>(%ub0, %ub1) {
+      affine.load %m[%i] : memref<?xf32>
+    }
+  }
+  return
+}