[MLIR][Affine] NFC: Move AffineValueMap and MutableAffineMap

[flaub]

Summary

The AffineValueMap is moved into Dialect/AffineOps to prevent a cyclic
dependency between Analysis and Dialect/AffineOps.

[flaub]

• Fix example

[flaub]

arcanist

[mehdi_amini]

There is a bunch of refactoring around AffineValueMap and MutableAffineMap, can you pull this out in a NFC parent revision?

[mehdi_amini]

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

277

represents a hyper-rectangular affine parallel band ?

[flaub]

• Split

[mehdi_amini]

Thanks!

[flaub]

The original diff included the AffineParallelOp, but this has been split out into its own diff based on this one: https://reviews.llvm.org/D74288

[rriddle]

mlir/lib/Dialect/AffineOps/AffineValueMap.cpp

76

We can drop the inline keywords now.

[flaub]

mlir/lib/Dialect/AffineOps/AffineValueMap.cpp

76

Should we drop them from headers too? (IIUC, they're just hints and they are redundant if the definition is already in a header).

[flaub]

Remove inline

mlir/include/mlir/Analysis/AffineStructures.h

 #define MLIR_ANALYSIS_AFFINE_STRUCTURES_H
 
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 
-class AffineApplyOp;
-class AffineBound;
 class AffineCondition;
-class AffineMap;
 class AffineForOp;
+class AffineMap;
+class AffineValueMap;
 class IntegerSet;
 class MLIRContext;
 class Value;
 class HyperRectangularSet;
 class MemRefType;
-
-/// A mutable affine map. Its affine expressions are however unique.
-struct MutableAffineMap {
-public:
-  MutableAffineMap() {}
-  MutableAffineMap(AffineMap map);
-
-  ArrayRef<AffineExpr> getResults() const { return results; }
-  AffineExpr getResult(unsigned idx) const { return results[idx]; }
-  void setResult(unsigned idx, AffineExpr result) { results[idx] = result; }
-  unsigned getNumResults() const { return results.size(); }
-  unsigned getNumDims() const { return numDims; }
-  void setNumDims(unsigned d) { numDims = d; }
-  unsigned getNumSymbols() const { return numSymbols; }
-  void setNumSymbols(unsigned d) { numSymbols = d; }
-  MLIRContext *getContext() const { return context; }
-
-  /// Returns true if the idx'th result expression is a multiple of factor.
-  bool isMultipleOf(unsigned idx, int64_t factor) const;
-
-  /// Resets this MutableAffineMap with 'map'.
-  void reset(AffineMap map);
-
-  /// Simplify the (result) expressions in this map using analysis (used by
-  //-simplify-affine-expr pass).
-  void simplify();
-  /// Get the AffineMap corresponding to this MutableAffineMap. Note that an
-  /// AffineMap will be uniqued and stored in context, while a mutable one
-  /// isn't.
-  AffineMap getAffineMap() const;
-
-private:
-  // Same meaning as AffineMap's fields.
-  SmallVector<AffineExpr, 8> results;
-  unsigned numDims;
-  unsigned numSymbols;
-  /// A pointer to the IR's context to store all newly created
-  /// AffineExprStorage's.
-  MLIRContext *context;
-};
+struct MutableAffineMap;
 
 /// A mutable integer set. Its affine expressions are however unique.
 struct MutableIntegerSet {
 public:
   MutableIntegerSet(IntegerSet set, MLIRContext *context);
 
   /// Create a universal set (no constraints).
   MutableIntegerSet(unsigned numDims, unsigned numSymbols,

 private:
   unsigned numDims;
   unsigned numSymbols;
 
   SmallVector<AffineExpr, 8> constraints;
   SmallVector<bool, 8> eqFlags;
 };
 
-/// An AffineValueMap is an affine map plus its ML value operands and
-/// results for analysis purposes. The structure is still a tree form that is
-/// same as that of an affine map or an AffineApplyOp. However, its operands,
-/// results, and its map can themselves change  as a result of
-/// substitutions, simplifications, and other analysis.
-// An affine value map can readily be constructed from an AffineApplyOp, or an
-// AffineBound of a AffineForOp. It can be further transformed, substituted
-// into, or simplified. Unlike AffineMap's, AffineValueMap's are created and
-// destroyed during analysis. Only the AffineMap expressions that are pointed by
-// them are unique'd. An affine value map, and the operations on it, maintain
-// the invariant that operands are always positionally aligned with the
-// AffineDimExpr and AffineSymbolExpr in the underlying AffineMap.
-// TODO(bondhugula): Some of these classes could go into separate files.
-class AffineValueMap {
-public:
-  // Creates an empty AffineValueMap (users should call 'reset' to reset map
-  // and operands).
-  AffineValueMap() {}
-  AffineValueMap(AffineMap map, ArrayRef<Value> operands,
-                 ArrayRef<Value> results = llvm::None);
-
-  explicit AffineValueMap(AffineApplyOp applyOp);
-  explicit AffineValueMap(AffineBound bound);
-
-  ~AffineValueMap();
-
-  // Resets this AffineValueMap with 'map', 'operands', and 'results'.
-  void reset(AffineMap map, ArrayRef<Value> operands,
-             ArrayRef<Value> results = llvm::None);
-
-  /// Return the value map that is the difference of value maps 'a' and 'b',
-  /// represented as an affine map and its operands. The output map + operands
-  /// are canonicalized and simplified.
-  static void difference(const AffineValueMap &a, const AffineValueMap &b,
-                         AffineValueMap *res);
-
-  /// Return true if the idx^th result can be proved to be a multiple of
-  /// 'factor', false otherwise.
-  inline bool isMultipleOf(unsigned idx, int64_t factor) const;
-
-  /// Return true if the idx^th result depends on 'value', false otherwise.
-  bool isFunctionOf(unsigned idx, Value value) const;
-
-  /// Return true if the result at 'idx' is a constant, false
-  /// otherwise.
-  bool isConstant(unsigned idx) const;
-
-  /// Return true if this is an identity map.
-  bool isIdentity() const;
-
-  void setResult(unsigned i, AffineExpr e) { map.setResult(i, e); }
-  AffineExpr getResult(unsigned i) { return map.getResult(i); }
-  inline unsigned getNumOperands() const { return operands.size(); }
-  inline unsigned getNumDims() const { return map.getNumDims(); }
-  inline unsigned getNumSymbols() const { return map.getNumSymbols(); }
-  inline unsigned getNumResults() const { return map.getNumResults(); }
-
-  Value getOperand(unsigned i) const;
-  ArrayRef<Value> getOperands() const;
-  AffineMap getAffineMap() const;
-
-private:
-  // A mutable affine map.
-  MutableAffineMap map;
-
-  // TODO: make these trailing objects?
-  /// The SSA operands binding to the dim's and symbols of 'map'.
-  SmallVector<Value, 4> operands;
-  /// The SSA results binding to the results of 'map'.
-  SmallVector<Value, 4> results;
-};
-
 /// An IntegerValueSet is an integer set plus its operands.
 // Both, the integer set being pointed to and the operands can change during
 // analysis, simplification, and transformation.
 class IntegerValueSet {
   /// Constructs an integer value set from an affine value map.
   // This will lead to a single equality in 'set'.
   explicit IntegerValueSet(const AffineValueMap &avm);
 

mlir/include/mlir/Dialect/AffineOps/AffineOps.h

public:
   static void build(Builder *builder, OperationState &result, AffineMap map,
                     ValueRange operands);
 
   /// Returns the affine map to be applied by this operation.
   AffineMap getAffineMap() {
     return getAttrOfType<AffineMapAttr>("map").getValue();
   }
 
+  /// Returns the affine value map computed from this operation.
+  AffineValueMap getAffineValueMap();
+
   /// Returns true if the result of this operation can be used as dimension id.
   bool isValidDim();
 
   /// Returns true if the result of this operation is a symbol.
   bool isValidSymbol();
 
   static StringRef getOperationName() { return "affine.apply"; }
 

 /// Modifies both `map` and `operands` in-place so as to:
 /// 1. drop duplicate operands
 /// 2. drop unused dims and symbols from map
 /// 3. promote valid symbols to symbolic operands in case they appeared as
 ///    dimensional operands
 /// 4. propagate constant operands and drop them
 void canonicalizeMapAndOperands(AffineMap *map,
                                 SmallVectorImpl<Value> *operands);
+
 /// Canonicalizes an integer set the same way canonicalizeMapAndOperands does
 /// for affine maps.
 void canonicalizeSetAndOperands(IntegerSet *set,
                                 SmallVectorImpl<Value> *operands);
 
 /// Returns a composed AffineApplyOp by composing `map` and `operands` with
 /// other AffineApplyOps supplying those operands. The operands of the resulting
 /// AffineApplyOp do not change the length of  AffineApplyOp chains.

 /// This class does not own the underlying operands. Instead, it refers
 /// to the operands stored in the AffineForOp. Its life span should not exceed
 /// that of the for operation it refers to.
 class AffineBound {
 public:
   AffineForOp getAffineForOp() { return op; }
   AffineMap getMap() { return map; }
 
-  /// Returns an AffineValueMap representing this bound.
-  AffineValueMap getAsAffineValueMap();
-
   unsigned getNumOperands() { return opEnd - opStart; }
   Value getOperand(unsigned idx) { return op.getOperand(opStart + idx); }
 
   using operand_iterator = AffineForOp::operand_iterator;
   using operand_range = AffineForOp::operand_range;
 
   operand_iterator operand_begin() { return op.operand_begin() + opStart; }
   operand_iterator operand_end() { return op.operand_begin() + opEnd; }

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

let description = [{
     affine map.
 
     Example:
 
        %0 = affine.max (d0) -> (1000, d0 + 512) (%i0) : index
   }];
 }
 
+def AffineParallelOp : Affine_Op<"parallel", [ImplicitAffineTerminator]> {
+  let summary = "multi-index parallel band operation";
+  let description = [{
+    The "affine.parallel" operation represents an 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".  Calling AffineParallelOp::build will create such
+    region and insert the terminator, so will the parsing even in cases if it is
+    absent from the custom format.
+
+    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.
+
+    Example:
+
+      affine.parallel (%i, %j) = [0, 0] to [10, 10] step (1, 1) {
+        ...
+      }
+  }];
+
+  let arguments = (ins
+     AffineMapAttr:$lowerBoundsMap,
+     AffineMapAttr:$upperBoundsMap,
+     I64ArrayAttr:$steps,
+     Variadic<Index>:$mapOperands);
+  let regions = (region SizedRegion<1>:$region);
+
+  let builders = [
+    OpBuilder<"Builder* builder, OperationState& result,"
+              "ArrayRef<int64_t> ranges">,
+    OpBuilder<"Builder* builder, OperationState& result, AffineMap lbMap,"
+              "ValueRange lbArgs, AffineMap ubMap, ValueRange ubArgs">,
+    OpBuilder<"Builder* builder, OperationState& result, 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.
+    llvm::Optional<llvm::SmallVector<int64_t, 8>> getConstantRanges();
+    
+    mlir::Block *getBody();
+    mlir::OpBuilder getBodyBuilder();
+    void setSteps(ArrayRef<int64_t> newSteps);
+
+    static StringRef getLowerBoundsMapAttrName() { return "lowerBoundsMap"; }
+    static StringRef getUpperBoundsMapAttrName() { return "upperBoundsMap"; }
+    static StringRef getStepsAttrName() { return "steps"; }
+  }];
+}
+
 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
     below:
 

mlir/include/mlir/Dialect/AffineOps/AffineValueMap.h

+//===- AffineValueMap.h - MLIR Affine Value Map Class -----------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// An AffineValueMap is an affine map plus its ML value operands and results for
+// analysis purposes.
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_AFFINEOPS_AFFINEVALUEMAP_H
+#define MLIR_DIALECT_AFFINEOPS_AFFINEVALUEMAP_H
+
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/Value.h"
+
+namespace mlir {
+
+/// An AffineValueMap is an affine map plus its ML value operands and
+/// results for analysis purposes. The structure is still a tree form that is
+/// same as that of an affine map or an AffineApplyOp. However, its operands,
+/// results, and its map can themselves change  as a result of
+/// substitutions, simplifications, and other analysis.
+// An affine value map can readily be constructed from an AffineApplyOp, or an
+// AffineBound of a AffineForOp. It can be further transformed, substituted
+// into, or simplified. Unlike AffineMap's, AffineValueMap's are created and
+// destroyed during analysis. Only the AffineMap expressions that are pointed by
+// them are unique'd. An affine value map, and the operations on it, maintain
+// the invariant that operands are always positionally aligned with the
+// AffineDimExpr and AffineSymbolExpr in the underlying AffineMap.
+class AffineValueMap {
+public:
+  // Creates an empty AffineValueMap (users should call 'reset' to reset map
+  // and operands).
+  AffineValueMap() {}
+  AffineValueMap(AffineMap map, ValueRange operands, ValueRange results = {});
+
+  ~AffineValueMap();
+
+  // Resets this AffineValueMap with 'map', 'operands', and 'results'.
+  void reset(AffineMap map, ValueRange operands, ValueRange results = {});
+
+  /// Return the value map that is the difference of value maps 'a' and 'b',
+  /// represented as an affine map and its operands. The output map + operands
+  /// are canonicalized and simplified.
+  static void difference(const AffineValueMap &a, const AffineValueMap &b,
+                         AffineValueMap *res);
+
+  /// Return true if the idx^th result can be proved to be a multiple of
+  /// 'factor', false otherwise.
+  inline bool isMultipleOf(unsigned idx, int64_t factor) const;
+
+  /// Return true if the idx^th result depends on 'value', false otherwise.
+  bool isFunctionOf(unsigned idx, Value value) const;
+
+  /// Return true if the result at 'idx' is a constant, false
+  /// otherwise.
+  bool isConstant(unsigned idx) const;
+
+  /// Return true if this is an identity map.
+  bool isIdentity() const;
+
+  void setResult(unsigned i, AffineExpr e) { map.setResult(i, e); }
+  AffineExpr getResult(unsigned i) { return map.getResult(i); }
+  inline unsigned getNumOperands() const { return operands.size(); }
+  inline unsigned getNumDims() const { return map.getNumDims(); }
+  inline unsigned getNumSymbols() const { return map.getNumSymbols(); }
+  inline unsigned getNumResults() const { return map.getNumResults(); }
+
+  Value getOperand(unsigned i) const;
+  ArrayRef<Value> getOperands() const;
+  AffineMap getAffineMap() const;
+
+private:
+  // A mutable affine map.
+  MutableAffineMap map;
+
+  // TODO: make these trailing objects?
+  /// The SSA operands binding to the dim's and symbols of 'map'.
+  SmallVector<Value, 4> operands;
+  /// The SSA results binding to the results of 'map'.
+  SmallVector<Value, 4> results;
+};
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_AFFINEOPS_AFFINEVALUEMAP_H

mlir/include/mlir/IR/AffineMap.h

 //===- AffineMap.h - MLIR Affine Map Class ----------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Affine maps are mathematical functions which map a list of dimension
 // identifiers and symbols, to multidimensional affine expressions.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_AFFINE_MAP_H
 #define MLIR_IR_AFFINE_MAP_H
 
+#include "mlir/IR/AffineExpr.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMapInfo.h"
 
 namespace mlir {
 
 namespace detail {
 struct AffineMapStorage;
 } // end namespace detail
 
-class AffineExpr;
 class Attribute;
 struct LogicalResult;
 class MLIRContext;
 
 /// A multi-dimensional affine map
 /// Affine map's are immutable like Type's, and they are uniqued.
 /// Eg: (d0, d1) -> (d0/128, d0 mod 128, d1)
 /// The names used (d0, d1) don't matter - it's the mathematical function that
static AffineMap getImpl(unsigned dimCount, unsigned symbolCount,
                            ArrayRef<AffineExpr> results, MLIRContext *context);
 };
 
 // Make AffineExpr hashable.
 inline ::llvm::hash_code hash_value(AffineMap arg) {
   return ::llvm::hash_value(arg.map);
 }
 
+/// A mutable affine map. Its affine expressions are however unique.
+struct MutableAffineMap {
+public:
+  MutableAffineMap() {}
+  MutableAffineMap(AffineMap map);
+
+  ArrayRef<AffineExpr> getResults() const { return results; }
+  AffineExpr getResult(unsigned idx) const { return results[idx]; }
+  void setResult(unsigned idx, AffineExpr result) { results[idx] = result; }
+  unsigned getNumResults() const { return results.size(); }
+  unsigned getNumDims() const { return numDims; }
+  void setNumDims(unsigned d) { numDims = d; }
+  unsigned getNumSymbols() const { return numSymbols; }
+  void setNumSymbols(unsigned d) { numSymbols = d; }
+  MLIRContext *getContext() const { return context; }
+
+  /// Returns true if the idx'th result expression is a multiple of factor.
+  bool isMultipleOf(unsigned idx, int64_t factor) const;
+
+  /// Resets this MutableAffineMap with 'map'.
+  void reset(AffineMap map);
+
+  /// Simplify the (result) expressions in this map using analysis (used by
+  //-simplify-affine-expr pass).
+  void simplify();
+  /// Get the AffineMap corresponding to this MutableAffineMap. Note that an
+  /// AffineMap will be uniqued and stored in context, while a mutable one
+  /// isn't.
+  AffineMap getAffineMap() const;
+
+private:
+  // Same meaning as AffineMap's fields.
+  SmallVector<AffineExpr, 8> results;
+  unsigned numDims;
+  unsigned numSymbols;
+  /// A pointer to the IR's context to store all newly created
+  /// AffineExprStorage's.
+  MLIRContext *context;
+};
+
 /// Simplify an affine map by simplifying its underlying AffineExpr results.
 AffineMap simplifyAffineMap(AffineMap map);
 
 /// Returns a map of codomain to domain dimensions such that the first codomain
 /// dimension for a particular domain dimension is selected.
 /// Returns an empty map if the input map is empty or if `map` is not invertible
 /// (i.e. `map` does not contain a subset that is a permutation of full domain
 /// rank).
inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
   map.print(os);
   return os;
 }
 } // end namespace mlir
 
 namespace llvm {
 
 // AffineExpr hash just like pointers
-template <> struct DenseMapInfo<mlir::AffineMap> {
+template <>
+struct DenseMapInfo<mlir::AffineMap> {
   static mlir::AffineMap getEmptyKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
     return mlir::AffineMap(static_cast<mlir::AffineMap::ImplType *>(pointer));
   }
   static mlir::AffineMap getTombstoneKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return mlir::AffineMap(static_cast<mlir::AffineMap::ImplType *>(pointer));
   }

mlir/lib/Analysis/AffineAnalysis.cpp

 // (expressions, maps, sets), and other utilities relying on such analysis.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/AffineOps/AffineOps.h"
+#include "mlir/Dialect/AffineOps/AffineValueMap.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Support/STLExtras.h"

mlir/lib/Analysis/AffineStructures.cpp

 //===- AffineStructures.cpp - MLIR Affine Structures Class-----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Structures for affine/polyhedral analysis of MLIR functions.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Dialect/AffineOps/AffineOps.h"
+#include "mlir/Dialect/AffineOps/AffineValueMap.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
if (set.getNumConstraints() == 0) {
     localVarCst->reset(set.getNumDims(), set.getNumSymbols());
     return success();
   }
   return ::getFlattenedAffineExprs(set.getConstraints(), set.getNumDims(),
                                    set.getNumSymbols(), flattenedExprs,
                                    localVarCst);
 }
 
-//===----------------------------------------------------------------------===//
-// MutableAffineMap.
-//===----------------------------------------------------------------------===//
-
-MutableAffineMap::MutableAffineMap(AffineMap map)
-    : numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
-      // A map always has at least 1 result by construction
-      context(map.getResult(0).getContext()) {
-  for (auto result : map.getResults())
-    results.push_back(result);
-}
-
-void MutableAffineMap::reset(AffineMap map) {
-  results.clear();
-  numDims = map.getNumDims();
-  numSymbols = map.getNumSymbols();
-  // A map always has at least 1 result by construction
-  context = map.getResult(0).getContext();
-  for (auto result : map.getResults())
-    results.push_back(result);
-}
-
-bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const {
-  if (results[idx].isMultipleOf(factor))
-    return true;
-
-  // TODO(bondhugula): use simplifyAffineExpr and FlatAffineConstraints to
-  // complete this (for a more powerful analysis).
-  return false;
-}
-
-// Simplifies the result affine expressions of this map. The expressions have to
-// be pure for the simplification implemented.
-void MutableAffineMap::simplify() {
-  // Simplify each of the results if possible.
-  // TODO(ntv): functional-style map
-  for (unsigned i = 0, e = getNumResults(); i < e; i++) {
-    results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols);
-  }
-}
-
-AffineMap MutableAffineMap::getAffineMap() const {
-  return AffineMap::get(numDims, numSymbols, results);
-}
-
 MutableIntegerSet::MutableIntegerSet(IntegerSet set, MLIRContext *context)
     : numDims(set.getNumDims()), numSymbols(set.getNumSymbols()) {
   // TODO(bondhugula)
 }
 
 // Universal set.
 MutableIntegerSet::MutableIntegerSet(unsigned numDims, unsigned numSymbols,
                                      MLIRContext *context)
     : numDims(numDims), numSymbols(numSymbols) {}
 
 //===----------------------------------------------------------------------===//
-// AffineValueMap.
-//===----------------------------------------------------------------------===//
-
-AffineValueMap::AffineValueMap(AffineMap map, ArrayRef<Value> operands,
-                               ArrayRef<Value> results)
-    : map(map), operands(operands.begin(), operands.end()),
-      results(results.begin(), results.end()) {}
-
-AffineValueMap::AffineValueMap(AffineApplyOp applyOp)
-    : map(applyOp.getAffineMap()),
-      operands(applyOp.operand_begin(), applyOp.operand_end()) {
-  results.push_back(applyOp.getResult());
-}
-
-AffineValueMap::AffineValueMap(AffineBound bound)
-    : map(bound.getMap()),
-      operands(bound.operand_begin(), bound.operand_end()) {}
-
-void AffineValueMap::reset(AffineMap map, ArrayRef<Value> operands,
-                           ArrayRef<Value> results) {
-  this->map.reset(map);
-  this->operands.assign(operands.begin(), operands.end());
-  this->results.assign(results.begin(), results.end());
-}
-
-void AffineValueMap::difference(const AffineValueMap &a,
-                                const AffineValueMap &b, AffineValueMap *res) {
-  assert(a.getNumResults() == b.getNumResults() && "invalid inputs");
-
-  // Fully compose A's map + operands.
-  auto aMap = a.getAffineMap();
-  SmallVector<Value, 4> aOperands(a.getOperands().begin(),
-                                  a.getOperands().end());
-  fullyComposeAffineMapAndOperands(&aMap, &aOperands);
-
-  // Use the affine apply normalizer to get B's map into A's coordinate space.
-  AffineApplyNormalizer normalizer(aMap, aOperands);
-  SmallVector<Value, 4> bOperands(b.getOperands().begin(),
-                                  b.getOperands().end());
-  auto bMap = b.getAffineMap();
-  normalizer.normalize(&bMap, &bOperands);
-
-  assert(std::equal(bOperands.begin(), bOperands.end(),
-                    normalizer.getOperands().begin()) &&
-         "operands are expected to be the same after normalization");
-
-  // Construct the difference expressions.
-  SmallVector<AffineExpr, 4> diffExprs;
-  diffExprs.reserve(a.getNumResults());
-  for (unsigned i = 0, e = bMap.getNumResults(); i < e; ++i)
-    diffExprs.push_back(normalizer.getAffineMap().getResult(i) -
-                        bMap.getResult(i));
-
-  auto diffMap = AffineMap::get(normalizer.getNumDims(),
-                                normalizer.getNumSymbols(), diffExprs);
-  canonicalizeMapAndOperands(&diffMap, &bOperands);
-  diffMap = simplifyAffineMap(diffMap);
-  res->reset(diffMap, bOperands);
-}
-
-// Returns true and sets 'indexOfMatch' if 'valueToMatch' is found in
-// 'valuesToSearch' beginning at 'indexStart'. Returns false otherwise.
-static bool findIndex(Value valueToMatch, ArrayRef<Value> valuesToSearch,
-                      unsigned indexStart, unsigned *indexOfMatch) {
-  unsigned size = valuesToSearch.size();
-  for (unsigned i = indexStart; i < size; ++i) {
-    if (valueToMatch == valuesToSearch[i]) {
-      *indexOfMatch = i;
-      return true;
-    }
-  }
-  return false;
-}
-
-inline bool AffineValueMap::isMultipleOf(unsigned idx, int64_t factor) const {
-  return map.isMultipleOf(idx, factor);
-}
-
-/// This method uses the invariant that operands are always positionally aligned
-/// with the AffineDimExpr in the underlying AffineMap.
-bool AffineValueMap::isFunctionOf(unsigned idx, Value value) const {
-  unsigned index;
-  if (!findIndex(value, operands, /*indexStart=*/0, &index)) {
-    return false;
-  }
-  auto expr = const_cast<AffineValueMap *>(this)->getAffineMap().getResult(idx);
-  // TODO(ntv): this is better implemented on a flattened representation.
-  // At least for now it is conservative.
-  return expr.isFunctionOfDim(index);
-}
-
-Value AffineValueMap::getOperand(unsigned i) const {
-  return static_cast<Value>(operands[i]);
-}
-
-ArrayRef<Value> AffineValueMap::getOperands() const {
-  return ArrayRef<Value>(operands);
-}
-
-AffineMap AffineValueMap::getAffineMap() const { return map.getAffineMap(); }
-
-AffineValueMap::~AffineValueMap() {}
-
-//===----------------------------------------------------------------------===//
 // FlatAffineConstraints.
 //===----------------------------------------------------------------------===//
 
 // Copy constructor.
 FlatAffineConstraints::FlatAffineConstraints(
     const FlatAffineConstraints &other) {
   numReservedCols = other.numReservedCols;
   numDims = other.getNumDimIds();

mlir/lib/Analysis/LoopAnalysis.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/LoopAnalysis.h"
 
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Analysis/NestedMatcher.h"
 #include "mlir/Dialect/AffineOps/AffineOps.h"
+#include "mlir/Dialect/AffineOps/AffineValueMap.h"
 #include "mlir/Support/MathExtras.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallString.h"
 #include <type_traits>
 
 using namespace mlir;
 
affineApplyOps[0]->emitRemark(
         "CompositionAffineMapsPass must have been run: there should be at most "
         "one AffineApplyOp, returning false conservatively.");
     return false;
   }
 
   auto composeOp = cast<AffineApplyOp>(affineApplyOps[0]);
   // We need yet another level of indirection because the `dim` index of the
   // access may not correspond to the `dim` index of composeOp.
-  return !(AffineValueMap(composeOp).isFunctionOf(0, iv));
+  return !composeOp.getAffineValueMap().isFunctionOf(0, iv);
 }
 
 DenseSet<Value> mlir::getInvariantAccesses(Value iv, ArrayRef<Value> indices) {
   DenseSet<Value> res;
   for (unsigned idx = 0, n = indices.size(); idx < n; ++idx) {
     auto val = indices[idx];
     if (isAccessIndexInvariant(iv, val)) {
       res.insert(val);

mlir/lib/Analysis/Utils.cpp

 // structures.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Utils.h"
 
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Dialect/AffineOps/AffineOps.h"
+#include "mlir/Dialect/AffineOps/AffineValueMap.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "analysis-utils"
 
 using namespace mlir;

mlir/lib/Dialect/AffineOps/AffineOps.cpp

 //===- AffineOps.cpp - MLIR Affine Operations -----------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/AffineOps/AffineOps.h"
+#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Dialect/AffineOps/AffineValueMap.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "mlir/Transforms/SideEffectsInterface.h"
if (auto *op = value.getDefiningOp()) {
     if (auto applyOp = dyn_cast<AffineApplyOp>(op))
       return applyOp.isValidDim();
     // The dim op is okay if its operand memref/tensor is defined at the top
     // level.
     if (auto dimOp = dyn_cast<DimOp>(op))
       return isTopLevelValue(dimOp.getOperand());
     return false;
   }
-  // This value has to be a block argument for a FuncOp or an affine.for.
+  // This value has to be a block argument of a FuncOp, an 'affine.for', or an
+  // 'affine.parallel'.
   auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
-  return isa<FuncOp>(parentOp) || isa<AffineForOp>(parentOp);
+  return isa<FuncOp>(parentOp) || isa<AffineForOp>(parentOp) ||
+         isa<AffineParallelOp>(parentOp);
 }
 
 /// Returns true if the 'index' dimension of the `memref` defined by
 /// `memrefDefOp` is a statically  shaped one or defined using a valid symbol.
 template <typename AnyMemRefDefOp>
 static bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp,
                                     unsigned index) {
   auto memRefType = memrefDefOp.getType();

 
 void AffineApplyOp::build(Builder *builder, OperationState &result,
                           AffineMap map, ValueRange operands) {
   result.addOperands(operands);
   result.types.append(map.getNumResults(), builder->getIndexType());
   result.addAttribute("map", AffineMapAttr::get(map));
 }
 
+AffineValueMap AffineApplyOp::getAffineValueMap() {
+  return AffineValueMap(getAffineMap(), getOperands(), getResult());
+}
+
 ParseResult AffineApplyOp::parse(OpAsmParser &parser, OperationState &result) {
   auto &builder = parser.getBuilder();
   auto indexTy = builder.getIndexType();
 
   AffineMapAttr mapAttr;
   unsigned numDims;
   if (parser.parseAttribute(mapAttr, "map", result.attributes) ||
       parseDimAndSymbolList(parser, result.operands, numDims) ||
static void printBound(AffineMapAttr boundMap,
 
   // Print the map and its operands.
   p << boundMap;
   printDimAndSymbolList(boundOperands.begin(), boundOperands.end(),
                         map.getNumDims(), p);
 }
 
 static void print(OpAsmPrinter &p, AffineForOp op) {
-  p << "affine.for ";
+  p << op.getOperationName() << ' ';
   p.printOperand(op.getBody()->getArgument(0));
   p << " = ";
   printBound(op.getLowerBoundMapAttr(), op.getLowerBoundOperands(), "max", p);
   p << " to ";
   printBound(op.getUpperBoundMapAttr(), op.getUpperBoundOperands(), "min", p);
 
   if (op.getStep() != 1)
     p << " step " << op.getStep();

 
 LogicalResult AffinePrefetchOp::fold(ArrayRef<Attribute> cstOperands,
                                      SmallVectorImpl<OpFoldResult> &results) {
   /// prefetch(memrefcast) -> prefetch
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
+// AffineParallelOp
+//===----------------------------------------------------------------------===//
+
+void AffineParallelOp::build(Builder *builder, OperationState &result,
+                             ArrayRef<int64_t> ranges) {
+  // Default initalize empty maps
+  auto lbMap = AffineMap::get(builder->getContext());
+  auto ubMap = AffineMap::get(builder->getContext());
+  // If ranges, set to [0, N) for each range
+  if (ranges.size()) {
+    SmallVector<AffineExpr, 8> lbExprs;
+    SmallVector<AffineExpr, 8> ubExprs;
+    // Make range expressions for each range
+    for (int64_t range : ranges) {
+      lbExprs.push_back(builder->getAffineConstantExpr(0));
+      ubExprs.push_back(builder->getAffineConstantExpr(range));
+    }
+    lbMap = AffineMap::get(0, 0, lbExprs);
+    ubMap = AffineMap::get(0, 0, ubExprs);
+  }
+  // Fall through
+  build(builder, result, lbMap, {}, ubMap, {});
+}
+
+void AffineParallelOp::build(Builder *builder, OperationState &result,
+                             AffineMap lbMap, ValueRange lbArgs,
+                             AffineMap ubMap, ValueRange ubArgs) {
+  // Verify sizes
+  auto dims = lbMap.getNumResults();
+  assert(dims == ubMap.getNumResults());
+  // Make default step sizes of 1
+  SmallVector<int64_t, 8> steps(dims, 1);
+  // Call through
+  build(builder, result, lbMap, lbArgs, ubMap, ubArgs, steps);
+}
+
+void AffineParallelOp::build(Builder *builder, OperationState &result,
+                             AffineMap lbMap, ValueRange lbArgs,
+                             AffineMap ubMap, ValueRange ubArgs,
+                             ArrayRef<int64_t> steps) {
+  // Verify sizes
+  auto dims = lbMap.getNumResults();
+  assert(dims == ubMap.getNumResults());
+  assert(dims == steps.size());
+  // Set all of the attributes
+  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 args
+  for (unsigned i = 0; i < dims; ++i) {
+    body->addArgument(IndexType::get(builder->getContext()));
+  }
+  bodyRegion->push_back(body);
+  // Terminate
+  ensureTerminator(*bodyRegion, *builder, result.location);
+}
+
+AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() {
+  return {operand_begin(), operand_begin() + lowerBoundsMap().getNumInputs()};
+}
+
+AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() {
+  return {operand_begin() + lowerBoundsMap().getNumInputs(), operand_end()};
+}
+
+unsigned AffineParallelOp::getNumDims() { return steps().size(); }
+
+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() {
+  SmallVector<int64_t, 8> out;
+  // Get the ranges
+  AffineValueMap rangesValueMap = getRangesValueMap();
+  // Try to convert them to constants
+  for (unsigned i = 0, e = rangesValueMap.getNumResults(); i < e; ++i) {
+    auto expr = rangesValueMap.getResult(i);
+    if (auto cst = expr.dyn_cast<mlir::AffineConstantExpr>()) {
+      out.push_back(cst.getValue());
+    } else {
+      return llvm::None;
+    }
+  }
+  return out;
+}
+
+mlir::Block *AffineParallelOp::getBody() { return &region().front(); }
+
+mlir::OpBuilder AffineParallelOp::getBodyBuilder() {
+  return mlir::OpBuilder(getBody(), std::prev(getBody()->end()));
+}
+
+void AffineParallelOp::setSteps(ArrayRef<int64_t> newSteps) {
+  assert(newSteps.size() == getNumDims());
+  setAttr("steps", getBodyBuilder().getI64ArrayAttr(newSteps));
+}
+
+static LogicalResult verify(AffineParallelOp op) {
+  auto dims = op.getNumDims();
+  if (op.lowerBoundsMap().getNumResults() != dims ||
+      op.upperBoundsMap().getNumResults() != dims ||
+      op.steps().size() != dims || op.getBody()->getNumArguments() != dims) {
+    return op.emitOpError("region argument count and num results of upper "
+                          "bounds, lower bounds, and steps must all match");
+  }
+  return success();
+}
+
+static void print(OpAsmPrinter &p, AffineParallelOp op) {
+  p << op.getOperationName() << " [";
+  p.printOperands(op.getBody()->getArguments());
+  p << "] = [";
+  p.printAffineMapOfSSAIds(op.lowerBoundsMapAttr(),
+                           op.getLowerBoundsOperands());
+  p << "] to [";
+  p.printAffineMapOfSSAIds(op.upperBoundsMapAttr(),
+                           op.getUpperBoundsOperands());
+  p << ']';
+  bool elideSteps = true;
+  for (auto attr : op.steps()) {
+    auto step = attr.cast<IntegerAttr>().getValue();
+    if (step != 1) {
+      elideSteps = false;
+      break;
+    }
+  }
+  if (!elideSteps) {
+    p << " step " << op.steps();
+  }
+  p.printRegion(op.region(), /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/false);
+  p.printOptionalAttrDict(
+      op.getAttrs(),
+      /*elidedAttrs=*/{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 `]`
+//
+static ParseResult parseAffineParallelOp(OpAsmParser &parser,
+                                         OperationState &result) {
+  auto &builder = parser.getBuilder();
+  auto indexType = builder.getIndexType();
+  AffineMapAttr lowerBoundsAttr, upperBoundsAttr;
+  SmallVector<OpAsmParser::OperandType, 4> ivs, lowerBoundsMapOperands,
+      upperBoundsMapOperands;
+  if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1,
+                                     OpAsmParser::Delimiter::Square) ||
+      parser.parseEqual() ||
+      parser.parseAffineMapOfSSAIds(
+          lowerBoundsMapOperands, lowerBoundsAttr,
+          AffineParallelOp::getLowerBoundsMapAttrName(), result.attributes) ||
+      parser.resolveOperands(lowerBoundsMapOperands, indexType,
+                             result.operands) ||
+      parser.parseKeyword("to") ||
+      parser.parseAffineMapOfSSAIds(
+          upperBoundsMapOperands, upperBoundsAttr,
+          AffineParallelOp::getUpperBoundsMapAttrName(), result.attributes) ||
+      parser.resolveOperands(upperBoundsMapOperands, indexType,
+                             result.operands))
+    return failure();
+
+  AffineMapAttr stepsMapAttr;
+  SmallVector<NamedAttribute, 1> stepsAttrs;
+  SmallVector<OpAsmParser::OperandType, 4> stepsMapOperands;
+  if (parser.parseOptionalKeyword("step")) {
+    SmallVector<int64_t, 4> steps(ivs.size(), 1);
+    result.addAttribute(AffineParallelOp::getStepsAttrName(),
+                        builder.getI64ArrayAttr(steps));
+  } else {
+    if (parser.parseAffineMapOfSSAIds(stepsMapOperands, stepsMapAttr,
+                                      AffineParallelOp::getStepsAttrName(),
+                                      stepsAttrs))
+      return failure();
+
+    // Convert steps from an AffineMap into an I64ArrayAttr.
+    SmallVector<int64_t, 4> steps;
+    auto stepsMap = stepsMapAttr.getValue();
+    for (const auto &result : stepsMap.getResults()) {
+      auto constExpr = result.dyn_cast<AffineConstantExpr>();
+      if (!constExpr)
+        return parser.emitError(parser.getNameLoc(),
+                                "steps must be constant integers");
+      steps.push_back(constExpr.getValue());
+    }
+    result.addAttribute(AffineParallelOp::getStepsAttrName(),
+                        builder.getI64ArrayAttr(steps));
+  }
+
+  // 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();
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/AffineOps/AffineOps.cpp.inc"

mlir/lib/Dialect/AffineOps/AffineValueMap.cpp

+//===- AffineValueMap.cpp - MLIR Affine Value Map Class -------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/AffineOps/AffineValueMap.h"
+#include "mlir/Dialect/AffineOps/AffineOps.h"
+
+using namespace mlir;
+
+AffineValueMap::AffineValueMap(AffineMap map, ValueRange operands,
+                               ValueRange results)
+    : map(map), operands(operands.begin(), operands.end()),
+      results(results.begin(), results.end()) {}
+
+void AffineValueMap::reset(AffineMap map, ValueRange operands,
+                           ValueRange results) {
+  this->map.reset(map);
+  this->operands.assign(operands.begin(), operands.end());
+  this->results.assign(results.begin(), results.end());
+}
+
+void AffineValueMap::difference(const AffineValueMap &a,
+                                const AffineValueMap &b, AffineValueMap *res) {
+  assert(a.getNumResults() == b.getNumResults() && "invalid inputs");
+
+  // Fully compose A's map + operands.
+  auto aMap = a.getAffineMap();
+  SmallVector<Value, 4> aOperands(a.getOperands().begin(),
+                                  a.getOperands().end());
+  fullyComposeAffineMapAndOperands(&aMap, &aOperands);
+
+  // Use the affine apply normalizer to get B's map into A's coordinate space.
+  AffineApplyNormalizer normalizer(aMap, aOperands);
+  SmallVector<Value, 4> bOperands(b.getOperands().begin(),
+                                  b.getOperands().end());
+  auto bMap = b.getAffineMap();
+  normalizer.normalize(&bMap, &bOperands);
+
+  assert(std::equal(bOperands.begin(), bOperands.end(),
+                    normalizer.getOperands().begin()) &&
+         "operands are expected to be the same after normalization");
+
+  // Construct the difference expressions.
+  SmallVector<AffineExpr, 4> diffExprs;
+  diffExprs.reserve(a.getNumResults());
+  for (unsigned i = 0, e = bMap.getNumResults(); i < e; ++i)
+    diffExprs.push_back(normalizer.getAffineMap().getResult(i) -
+                        bMap.getResult(i));
+
+  auto diffMap = AffineMap::get(normalizer.getNumDims(),
+                                normalizer.getNumSymbols(), diffExprs);
+  canonicalizeMapAndOperands(&diffMap, &bOperands);
+  diffMap = simplifyAffineMap(diffMap);
+  res->reset(diffMap, bOperands);
+}
+
+// Returns true and sets 'indexOfMatch' if 'valueToMatch' is found in
+// 'valuesToSearch' beginning at 'indexStart'. Returns false otherwise.
+static bool findIndex(Value valueToMatch, ArrayRef<Value> valuesToSearch,
+                      unsigned indexStart, unsigned *indexOfMatch) {
+  unsigned size = valuesToSearch.size();
+  for (unsigned i = indexStart; i < size; ++i) {
+    if (valueToMatch == valuesToSearch[i]) {
+      *indexOfMatch = i;
+      return true;
+    }
+  }
+  return false;
+}
+
+inline bool AffineValueMap::isMultipleOf(unsigned idx, int64_t factor) const {
+  return map.isMultipleOf(idx, factor);
+}
+
+/// This method uses the invariant that operands are always positionally aligned
+/// with the AffineDimExpr in the underlying AffineMap.
+bool AffineValueMap::isFunctionOf(unsigned idx, Value value) const {
+  unsigned index;
+  if (!findIndex(value, operands, /*indexStart=*/0, &index)) {
+    return false;
+  }
+  auto expr = const_cast<AffineValueMap *>(this)->getAffineMap().getResult(idx);
+  // TODO(ntv): this is better implemented on a flattened representation.
+  // At least for now it is conservative.
+  return expr.isFunctionOfDim(index);
+}
+
+Value AffineValueMap::getOperand(unsigned i) const {
+  return static_cast<Value>(operands[i]);
+}
+
+ArrayRef<Value> AffineValueMap::getOperands() const {
+  return ArrayRef<Value>(operands);
+}
+
+AffineMap AffineValueMap::getAffineMap() const { return map.getAffineMap(); }
+
+AffineValueMap::~AffineValueMap() {}

mlir/lib/Dialect/AffineOps/CMakeLists.txt

 add_llvm_library(MLIRAffineOps
   AffineOps.cpp
+  AffineValueMap.cpp
   DialectRegistration.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/AffineOps
   )
 add_dependencies(MLIRAffineOps
   MLIRAffineOpsIncGen
   MLIRIR
   MLIRLoopLikeInterfaceIncGen
   MLIRStandardOps)
 target_link_libraries(MLIRAffineOps MLIRIR MLIRStandardOps)
-

mlir/lib/IR/AffineMap.cpp

for (auto m : maps) {
     if (!m)
       continue;
     assert(m.getNumSymbols() == 0 && "expected map without symbols");
     results.append(m.getResults().begin(), m.getResults().end());
     numDims = std::max(m.getNumDims(), numDims);
   }
   return numDims == 0 ? AffineMap() : AffineMap::get(numDims, 0, results);
 }
+
+//===----------------------------------------------------------------------===//
+// MutableAffineMap.
+//===----------------------------------------------------------------------===//
+
+MutableAffineMap::MutableAffineMap(AffineMap map)
+    : numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
+      // A map always has at least 1 result by construction
+      context(map.getResult(0).getContext()) {
+  for (auto result : map.getResults())
+    results.push_back(result);
+}
+
+void MutableAffineMap::reset(AffineMap map) {
+  results.clear();
+  numDims = map.getNumDims();
+  numSymbols = map.getNumSymbols();
+  // A map always has at least 1 result by construction
+  context = map.getResult(0).getContext();
+  for (auto result : map.getResults())
+    results.push_back(result);
+}
+
+bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const {
+  if (results[idx].isMultipleOf(factor))
+    return true;
+
+  // TODO(bondhugula): use simplifyAffineExpr and FlatAffineConstraints to
+  // complete this (for a more powerful analysis).
+  return false;
+}
+
+// Simplifies the result affine expressions of this map. The expressions have to
+// be pure for the simplification implemented.
+void MutableAffineMap::simplify() {
+  // Simplify each of the results if possible.
+  // TODO(ntv): functional-style map
+  for (unsigned i = 0, e = getNumResults(); i < e; i++) {
+    results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols);
+  }
+}
+
+AffineMap MutableAffineMap::getAffineMap() const {
+  return AffineMap::get(numDims, numSymbols, results);
+}

mlir/test/AffineOps/invalid.mlir

 // -----
 
 // CHECK-LABEL: @affine_max
 func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
   // expected-error@+1 {{operand count and affine map dimension and symbol count must match}}
   %0 = affine.max affine_map<(d0) -> (d0)> ()
 
   return
-}
+}
+
+// -----
+
+// CHECK-LABEL: @affine_parallel
+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}}
+  affine.parallel [%i0] = [0, 0] to [100, 100] step [10, 10] {
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @affine_parallel
+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}}
+  affine.parallel [%i0, %j0] = [0] to [100, 100] step [10, 10] {
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @affine_parallel
+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}}
+  affine.parallel [%i0, %j0] = [0, 0] to [100] step [10, 10] {
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @affine_parallel
+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}}
+  affine.parallel [%i0, %j0] = [0, 0] to [100, 100] step [10] {
+  }
+}

mlir/test/AffineOps/ops.mlir

affine.for %arg4 = 0 to %13 step 264 {
       %24 = dim %20, 0 : memref<?x?xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + d1 * s2 + s0)>>
       affine.for %arg5 = 0 to %24 step 768 {
         "foo"() : () -> ()
       }
     }
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: @parallel
+// CHECK-SAME: (%[[N:.*]]: index)
+func @parallel(%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] {
+    }
+  }
+  return
+}