[MLIR] Generalize Affine dependence analysis using Affine Relations

This patch removes code very specific to affine dependence analysis and
refactors it as a FlatAfffineRelation.

A FlatAffineRelation represents a set of ordered pairs (domain -> range) where
"domain" and "range" are tuples of identifiers. These relations are used to
represent an "access relation" for memory access on a memref.  An access
relation maps elements of an iteration domain to the element(s) of an array
domain accessed by that iteration of the associated statement through some
array reference.  The dependence relation representing the dependence
constraints between two memory accesses can be built by composing the access
relation of the destination access by the inverse of the access relation of
source access.

This patch does not change the functionality of the existing dependence
analysis in checkMemrefAccessDependence, but refactors it to use
FlatAffineRelations to deduplicate code and enable code reuse for future
development of features like scheduling, value-based dependence analysis, etc.

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D110563

Author: Groverkss

mlir/include/mlir/Analysis/AffineAnalysis.h

@@ -25,6 +25,7 @@ namespace mlir {
 class AffineApplyOp;
 class AffineForOp;
 class AffineValueMap;
+class FlatAffineRelation;
 class FlatAffineValueConstraints;
 class Operation;
 
@@ -89,6 +90,30 @@ struct MemRefAccess {
   // Returns true if this access is of a store op.
   bool isStore() const;
 
+  /// Creates an access relation for the access. An access relation maps
+  /// elements of an iteration domain to the element(s) of an array domain
+  /// accessed by that iteration of the associated statement through some array
+  /// reference. For example, given the MLIR code:
+  ///
+  /// affine.for %i0 = 0 to 10 {
+  ///   affine.for %i1 = 0 to 10 {
+  ///     %a = affine.load %arr[%i0 + %i1, %i0 + 2 * %i1] : memref<100x100xf32>
+  ///   }
+  /// }
+  ///
+  /// The access relation, assuming that the memory locations for %arr are
+  /// represented as %m0, %m1 would be:
+  ///
+  ///   (%i0, %i1) -> (%m0, %m1)
+  ///   %m0 = %i0 + %i1
+  ///   %m1 = %i0 + 2 * %i1
+  ///   0  <= %i0 < 10
+  ///   0  <= %i1 < 10
+  ///
+  /// Returns failure for yet unimplemented/unsupported cases (see docs of
+  /// mlir::getIndexSet and mlir::getRelationFromMap for these cases).
+  LogicalResult getAccessRelation(FlatAffineRelation &accessRel) const;
+
   /// Populates 'accessMap' with composition of AffineApplyOps reachable from
   /// 'indices'.
   void getAccessMap(AffineValueMap *accessMap) const;

mlir/include/mlir/Analysis/AffineStructures.h

@@ -426,6 +426,10 @@ class FlatAffineConstraints {
   /// O(VC) time.
   void removeRedundantConstraints();
 
+  /// Converts identifiers in the column range [idStart, idLimit) to local
+  /// variables.
+  void convertDimToLocal(unsigned dimStart, unsigned dimLimit);
+
   /// Merge local ids of `this` and `other`. This is done by appending local ids
   /// of `other` to `this` and inserting local ids of `this` to `other` at start
   /// of its local ids.
@@ -581,6 +585,16 @@ class FlatAffineValueConstraints : public FlatAffineConstraints {
                                        numLocals + 1,
                                    numDims, numSymbols, numLocals, valArgs) {}
 
+  FlatAffineValueConstraints(const FlatAffineConstraints &fac,
+                             ArrayRef<Optional<Value>> valArgs = {})
+      : FlatAffineConstraints(fac) {
+    assert(valArgs.empty() || valArgs.size() == numIds);
+    if (valArgs.empty())
+      values.resize(numIds, None);
+    else
+      values.append(valArgs.begin(), valArgs.end());
+  }
+
   /// Create a flat affine constraint system from an AffineValueMap or a list of
   /// these. The constructed system will only include equalities.
   explicit FlatAffineValueConstraints(const AffineValueMap &avm);
@@ -721,7 +735,8 @@ class FlatAffineValueConstraints : public FlatAffineConstraints {
   using FlatAffineConstraints::insertDimId;
   unsigned insertSymbolId(unsigned pos, ValueRange vals);
   using FlatAffineConstraints::insertSymbolId;
-  unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
+  virtual unsigned insertId(IdKind kind, unsigned pos,
+                            unsigned num = 1) override;
   unsigned insertId(IdKind kind, unsigned pos, ValueRange vals);
 
   /// Append identifiers of the specified kind after the last identifier of that
@@ -882,7 +897,7 @@ class FlatAffineValueConstraints : public FlatAffineConstraints {
   /// Removes identifiers in the column range [idStart, idLimit), and copies any
   /// remaining valid data into place, updates member variables, and resizes
   /// arrays as needed.
-  void removeIdRange(unsigned idStart, unsigned idLimit) override;
+  virtual void removeIdRange(unsigned idStart, unsigned idLimit) override;
 
   /// Eliminates the identifier at the specified position using Fourier-Motzkin
   /// variable elimination, but uses Gaussian elimination if there is an
@@ -901,6 +916,83 @@ class FlatAffineValueConstraints : public FlatAffineConstraints {
   SmallVector<Optional<Value>, 8> values;
 };
 
+/// A FlatAffineRelation represents a set of ordered pairs (domain -> range)
+/// where "domain" and "range" are tuples of identifiers. The relation is
+/// represented as a FlatAffineValueConstraints with separation of dimension
+/// identifiers into domain and  range. The identifiers are stored as:
+/// [domainIds, rangeIds, symbolIds, localIds, constant].
+class FlatAffineRelation : public FlatAffineValueConstraints {
+public:
+  FlatAffineRelation(unsigned numReservedInequalities,
+                     unsigned numReservedEqualities, unsigned numReservedCols,
+                     unsigned numDomainDims, unsigned numRangeDims,
+                     unsigned numSymbols, unsigned numLocals,
+                     ArrayRef<Optional<Value>> valArgs = {})
+      : FlatAffineValueConstraints(
+            numReservedInequalities, numReservedEqualities, numReservedCols,
+            numDomainDims + numRangeDims, numSymbols, numLocals, valArgs),
+        numDomainDims(numDomainDims), numRangeDims(numRangeDims) {}
+
+  FlatAffineRelation(unsigned numDomainDims = 0, unsigned numRangeDims = 0,
+                     unsigned numSymbols = 0, unsigned numLocals = 0)
+      : FlatAffineValueConstraints(numDomainDims + numRangeDims, numSymbols,
+                                   numLocals),
+        numDomainDims(numDomainDims), numRangeDims(numRangeDims) {}
+
+  FlatAffineRelation(unsigned numDomainDims, unsigned numRangeDims,
+                     FlatAffineValueConstraints &fac)
+      : FlatAffineValueConstraints(fac), numDomainDims(numDomainDims),
+        numRangeDims(numRangeDims) {}
+
+  FlatAffineRelation(unsigned numDomainDims, unsigned numRangeDims,
+                     FlatAffineConstraints &fac)
+      : FlatAffineValueConstraints(fac), numDomainDims(numDomainDims),
+        numRangeDims(numRangeDims) {}
+
+  /// Returns a set corresponding to the domain/range of the affine relation.
+  FlatAffineValueConstraints getDomainSet() const;
+  FlatAffineValueConstraints getRangeSet() const;
+
+  /// Returns the number of identifiers corresponding to domain/range of
+  /// relation.
+  inline unsigned getNumDomainDims() const { return numDomainDims; }
+  inline unsigned getNumRangeDims() const { return numRangeDims; }
+
+  /// Given affine relation `other: (domainOther -> rangeOther)`, this operation
+  /// takes the composition of `other` on `this: (domainThis -> rangeThis)`.
+  /// The resulting relation represents tuples of the form: `domainOther ->
+  /// rangeThis`.
+  void compose(const FlatAffineRelation &other);
+
+  /// Swap domain and range of the relation.
+  /// `(domain -> range)` is converted to `(range -> domain)`.
+  void inverse();
+
+  /// Insert `num` identifiers of the specified kind after the `pos` identifier
+  /// of that kind. The coefficient columns corresponding to the added
+  /// identifiers are initialized to zero.
+  void insertDomainId(unsigned pos, unsigned num = 1);
+  void insertRangeId(unsigned pos, unsigned num = 1);
+
+  /// Append `num` identifiers of the specified kind after the last identifier
+  /// of that kind. The coefficient columns corresponding to the added
+  /// identifiers are initialized to zero.
+  void appendDomainId(unsigned num = 1);
+  void appendRangeId(unsigned num = 1);
+
+protected:
+  // Number of dimension identifers corresponding to domain identifers.
+  unsigned numDomainDims;
+
+  // Number of dimension identifers corresponding to range identifers.
+  unsigned numRangeDims;
+
+  /// Removes identifiers in the column range [idStart, idLimit), and copies any
+  /// remaining valid data into place, updates member variables, and resizes
+  /// arrays as needed.
+  void removeIdRange(unsigned idStart, unsigned idLimit) override;
+};
+
 /// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the
 /// dimensions, symbols, and additional variables that represent floor divisions
 /// of dimensions, symbols, and in turn other floor divisions.  Returns failure
@@ -957,6 +1049,26 @@ AffineMap alignAffineMapWithValues(AffineMap map, ValueRange operands,
                                    ValueRange dims, ValueRange syms,
                                    SmallVector<Value> *newSyms = nullptr);
 
+/// Builds a relation from the given AffineMap/AffineValueMap `map`, containing
+/// all pairs of the form `operands -> result` that satisfy `map`. `rel` is set
+/// to the relation built. For example, give the AffineMap:
+///
+///   (d0, d1)[s0] -> (d0 + s0, d0 - s0)
+///
+/// the resulting relation formed is:
+///
+///   (d0, d1) -> (r1, r2)
+///   [d0  d1  r1  r2  s0  const]
+///    1   0   -1   0  1     0     = 0
+///    0   1    0  -1  -1    0     = 0
+///
+/// For AffineValueMap, the domain and symbols have Value set corresponding to
+/// the Value in `map`. Returns failure if the AffineMap could not be flattened
+/// (i.e., semi-affine is not yet handled).
+LogicalResult getRelationFromMap(AffineMap &map, FlatAffineRelation &rel);
+LogicalResult getRelationFromMap(const AffineValueMap &map,
+                                 FlatAffineRelation &rel);
+
 } // end namespace mlir.
 
 #endif // MLIR_ANALYSIS_AFFINESTRUCTURES_H