[MLIR][Presburger] Use Identifiers outside Presburger library (#77316)

The pull request includes the following changes.
1. Refactors the interface to `PresburgerSpace::identifiers` to `setId` and a
const `getId`, instead of previous `getId` which returned a mutable
reference. `resetIds` does not need to be called to use identifiers, `setId`
calls `resetIds` if identifiers are not enabled.
2. Deprecates `FlatAffineRelation` by refactoring all usages of
`FlatAffineRelation` to `IntegerRelation`. To achieve this,
`FlatAffineRelation::compose` is refactored into
`IntegerRelation::mergeAndCompose`.
3. Deletes unneeded overrides of virtual functions `hasConsistentState`,
`clearAndCopyFrom` and `fourierMotzkinEliminate` from
`FlatLinearValueConstraints` as these were only used through
`FlatAffineRelation` and we now use `IntegerRelation`'s member functions
instead.
4. Fixes an existing bug in FlatLinearValueConstraints' constructor
which caused
identifiers set by superclass FlatLinearConstraints' constructor to be
erased.
5. Fixes `IntegerRelation::convertVarKind` not preserving identifiers.

Author: iambrj

mlir/include/mlir/Analysis/FlatLinearValueConstraints.h

@@ -205,6 +205,10 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
 /// where each non-local variable can have an SSA Value attached to it.
 class FlatLinearValueConstraints : public FlatLinearConstraints {
 public:
+  /// The SSA Values attached to each non-local variable are stored as
+  /// identifiers in the constraint system's space.
+  using Identifier = presburger::Identifier;
+
   /// Constructs a constraint system reserving memory for the specified number
   /// of constraints and variables. `valArgs` are the optional SSA values
   /// associated with each dimension/symbol. These must either be empty or match
@@ -217,11 +221,9 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
       : FlatLinearConstraints(numReservedInequalities, numReservedEqualities,
                               numReservedCols, numDims, numSymbols, numLocals) {
     assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolVars());
-    values.reserve(numReservedCols);
-    if (valArgs.empty())
-      values.resize(getNumDimAndSymbolVars(), std::nullopt);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      if (valArgs[i])
+        setValue(i, *valArgs[i]);
   }
 
   /// Constructs a constraint system reserving memory for the specified number
@@ -236,11 +238,9 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
       : FlatLinearConstraints(numReservedInequalities, numReservedEqualities,
                               numReservedCols, numDims, numSymbols, numLocals) {
     assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolVars());
-    values.reserve(numReservedCols);
-    if (valArgs.empty())
-      values.resize(getNumDimAndSymbolVars(), std::nullopt);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      if (valArgs[i])
+        setValue(i, valArgs[i]);
   }
 
   /// Constructs a constraint system with the specified number of dimensions
@@ -272,11 +272,12 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   FlatLinearValueConstraints(const IntegerPolyhedron &fac,
                              ArrayRef<std::optional<Value>> valArgs = {})
       : FlatLinearConstraints(fac) {
-    assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolVars());
     if (valArgs.empty())
-      values.resize(getNumDimAndSymbolVars(), std::nullopt);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+      return;
+    assert(valArgs.size() == getNumDimAndSymbolVars());
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      if (valArgs[i])
+        setValue(i, *valArgs[i]);
   }
 
   /// Creates an affine constraint system from an IntegerSet.
@@ -290,9 +291,6 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
            cst->getKind() <= Kind::FlatAffineRelation;
   }
 
-  /// Replaces the contents of this FlatLinearValueConstraints with `other`.
-  void clearAndCopyFrom(const IntegerRelation &other) override;
-
   /// Adds a constant bound for the variable associated with the given Value.
   void addBound(presburger::BoundType type, Value val, int64_t value);
   using FlatLinearConstraints::addBound;
@@ -302,7 +300,9 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   inline Value getValue(unsigned pos) const {
     assert(pos < getNumDimAndSymbolVars() && "Invalid position");
     assert(hasValue(pos) && "variable's Value not set");
-    return *values[pos];
+    VarKind kind = getVarKindAt(pos);
+    unsigned relativePos = pos - getVarKindOffset(kind);
+    return space.getId(kind, relativePos).getValue<Value>();
   }
 
   /// Returns the Values associated with variables in range [start, end).
@@ -313,25 +313,44 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
     assert(start <= end && "invalid start position");
     values->clear();
     values->reserve(end - start);
-    for (unsigned i = start; i < end; i++)
+    for (unsigned i = start; i < end; ++i)
       values->push_back(getValue(i));
   }
 
-  inline ArrayRef<std::optional<Value>> getMaybeValues() const {
-    return {values.data(), values.size()};
+  inline SmallVector<std::optional<Value>> getMaybeValues() const {
+    SmallVector<std::optional<Value>> maybeValues;
+    maybeValues.reserve(getNumDimAndSymbolVars());
+    for (unsigned i = 0, e = getNumDimAndSymbolVars(); i < e; ++i)
+      if (hasValue(i)) {
+        maybeValues.push_back(getValue(i));
+      } else {
+        maybeValues.push_back(std::nullopt);
+      }
+    return maybeValues;
   }
 
-  inline ArrayRef<std::optional<Value>>
+  inline SmallVector<std::optional<Value>>
   getMaybeValues(presburger::VarKind kind) const {
     assert(kind != VarKind::Local &&
            "Local variables do not have any value attached to them.");
-    return {values.data() + getVarKindOffset(kind), getNumVarKind(kind)};
+    SmallVector<std::optional<Value>> maybeValues;
+    maybeValues.reserve(getNumVarKind(kind));
+    const unsigned offset = space.getVarKindOffset(kind);
+    for (unsigned i = 0, e = getNumVarKind(kind); i < e; ++i) {
+      if (hasValue(offset + i))
+        maybeValues.push_back(getValue(offset + i));
+      else
+        maybeValues.push_back(std::nullopt);
+    }
+    return maybeValues;
   }
 
   /// Returns true if the pos^th variable has an associated Value.
   inline bool hasValue(unsigned pos) const {
     assert(pos < getNumDimAndSymbolVars() && "Invalid position");
-    return values[pos].has_value();
+    VarKind kind = getVarKindAt(pos);
+    unsigned relativePos = pos - getVarKindOffset(kind);
+    return space.getId(kind, relativePos).hasValue();
   }
 
   unsigned appendDimVar(ValueRange vals);
@@ -358,9 +377,12 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   using IntegerPolyhedron::removeVarRange;
 
   /// Sets the Value associated with the pos^th variable.
+  /// Stores the Value in the space's identifiers.
   inline void setValue(unsigned pos, Value val) {
     assert(pos < getNumDimAndSymbolVars() && "invalid var position");
-    values[pos] = val;
+    VarKind kind = getVarKindAt(pos);
+    unsigned relativePos = pos - getVarKindOffset(kind);
+    space.setId(kind, relativePos, presburger::Identifier(val));
   }
 
   /// Sets the Values associated with the variables in the range [start, end).
@@ -387,9 +409,6 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   void projectOut(Value val);
   using IntegerPolyhedron::projectOut;
 
-  /// Swap the posA^th variable with the posB^th variable.
-  void swapVar(unsigned posA, unsigned posB) override;
-
   /// Prints the number of constraints, dimensions, symbols and locals in the
   /// FlatAffineValueConstraints. Also, prints for each variable whether there
   /// is an SSA Value attached to it.
@@ -444,28 +463,6 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   ///    output = {0 <= d0 <= 6, 1 <= d1 <= 15}
   LogicalResult unionBoundingBox(const FlatLinearValueConstraints &other);
   using IntegerPolyhedron::unionBoundingBox;
-
-protected:
-  /// Eliminates the variable at the specified position using Fourier-Motzkin
-  /// variable elimination, but uses Gaussian elimination if there is an
-  /// equality involving that variable. If the result of the elimination is
-  /// integer exact, `*isResultIntegerExact` is set to true. If `darkShadow` is
-  /// set to true, a potential under approximation (subset) of the rational
-  /// shadow / exact integer shadow is computed.
-  // See implementation comments for more details.
-  void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false,
-                               bool *isResultIntegerExact = nullptr) override;
-
-  /// Returns false if the fields corresponding to various variable counts, or
-  /// equality/inequality buffer sizes aren't consistent; true otherwise. This
-  /// is meant to be used within an assert internally.
-  bool hasConsistentState() const override;
-
-  /// Values corresponding to the (column) non-local variables of this
-  /// constraint system appearing in the order the variables correspond to
-  /// columns. Variables that aren't associated with any Value are set to
-  /// std::nullopt.
-  SmallVector<std::optional<Value>, 8> values;
 };
 
 /// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the

mlir/include/mlir/Analysis/Presburger/IntegerRelation.h

@@ -127,6 +127,12 @@ class IntegerRelation {
   /// the variable.
   void setId(VarKind kind, unsigned i, Identifier id);
 
+  void resetIds() { space.resetIds(); }
+
+  /// Get the identifiers for the variables of specified varKind. Calls resetIds
+  /// on the relations space if identifiers are not enabled.
+  ArrayRef<Identifier> getIds(VarKind kind);
+
   /// Returns a copy of the space without locals.
   PresburgerSpace getSpaceWithoutLocals() const {
     return PresburgerSpace::getRelationSpace(space.getNumDomainVars(),
@@ -674,6 +680,11 @@ class IntegerRelation {
   /// this for uniformity with `applyDomain`.
   void applyRange(const IntegerRelation &rel);
 
+  /// Given a relation `other: (A -> B)`, this operation merges the symbol and
+  /// local variables and then takes the composition of `other` on `this: (B ->
+  /// C)`. The resulting relation represents tuples of the form: `A -> C`.
+  void mergeAndCompose(const IntegerRelation &other);
+
   /// Compute an equivalent representation of the same set, such that all local
   /// vars in all disjuncts have division representations. This representation
   /// may involve local vars that correspond to divisions, and may also be a

mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h

@@ -250,21 +250,34 @@ class PresburgerSpace {
   /// locals).
   bool isEqual(const PresburgerSpace &other) const;
 
-  /// Get the identifier of the specified variable.
-  Identifier &getId(VarKind kind, unsigned pos) {
-    assert(kind != VarKind::Local && "Local variables have no identifiers");
-    return identifiers[getVarKindOffset(kind) + pos];
-  }
+  /// Get the identifier of pos^th variable of the specified kind.
   Identifier getId(VarKind kind, unsigned pos) const {
     assert(kind != VarKind::Local && "Local variables have no identifiers");
+    if (!usingIds)
+      return Identifier();
     return identifiers[getVarKindOffset(kind) + pos];
   }
 
   ArrayRef<Identifier> getIds(VarKind kind) const {
     assert(kind != VarKind::Local && "Local variables have no identifiers");
+    assert(usingIds && "Identifiers not enabled for space");
     return {identifiers.data() + getVarKindOffset(kind), getNumVarKind(kind)};
   }
 
+  ArrayRef<Identifier> getIds() const {
+    assert(usingIds && "Identifiers not enabled for space");
+    return identifiers;
+  }
+
+  /// Set the identifier of pos^th variable of the specified kind. Calls
+  /// resetIds if identifiers are not enabled.
+  void setId(VarKind kind, unsigned pos, Identifier id) {
+    assert(kind != VarKind::Local && "Local variables have no identifiers");
+    if (!usingIds)
+      resetIds();
+    identifiers[getVarKindOffset(kind) + pos] = id;
+  }
+
   /// Returns if identifiers are being used.
   bool isUsingIds() const { return usingIds; }
 

mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h

@@ -15,6 +15,7 @@
 #ifndef MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H
 #define MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H
 
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/SmallVector.h"
@@ -115,7 +116,7 @@ struct MemRefAccess {
   ///
   /// Returns failure for yet unimplemented/unsupported cases (see docs of
   /// mlir::getIndexSet and mlir::getRelationFromMap for these cases).
-  LogicalResult getAccessRelation(FlatAffineRelation &accessRel) const;
+  LogicalResult getAccessRelation(presburger::IntegerRelation &accessRel) const;
 
   /// Populates 'accessMap' with composition of AffineApplyOps reachable from
   /// 'indices'.

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

@@ -154,6 +154,14 @@ class FlatAffineValueConstraints : public FlatLinearValueConstraints {
 /// represented as a FlatAffineValueConstraints with separation of dimension
 /// variables into domain and  range. The variables are stored as:
 /// [domainVars, rangeVars, symbolVars, localVars, constant].
+///
+/// Deprecated: use IntegerRelation and store SSA Values in the PresburgerSpace
+/// of the relation using PresburgerSpace::identifiers. Note that
+/// FlatAffineRelation::numDomainDims and FlatAffineRelation::numRangeDims are
+/// independent of numDomain and numRange of the relation's space. In
+/// particular, operations such as FlatAffineRelation::compose do not ensure
+/// consistency between numDomainDims/numRangeDims and numDomain/numRange which
+/// may lead to unexpected behaviour.
 class FlatAffineRelation : public FlatAffineValueConstraints {
 public:
   FlatAffineRelation(unsigned numReservedInequalities,
@@ -251,9 +259,10 @@ class FlatAffineRelation : public FlatAffineValueConstraints {
 /// 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(AffineMap &map,
+                                 presburger::IntegerRelation &rel);
 LogicalResult getRelationFromMap(const AffineValueMap &map,
-                                 FlatAffineRelation &rel);
+                                 presburger::IntegerRelation &rel);
 
 } // namespace affine
 } // namespace mlir