[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 ↗

(On Diff #243355)

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

75

We can drop the inline keywords now.

[flaub]

mlir/lib/Dialect/AffineOps/AffineValueMap.cpp

75

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/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"

 
 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();

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);
+}