[MLIR] Move values from FlatAffineValueConstraints to PresburgerSpace

[Groverkss]

Summary

This patch moves values attached to identifiers to PresburgerSpace. This allows
values to be used by any user of PresburgerSpace.

Support for values in IntegerRelation, IntegerPolyhedron, PresburgerRelation,
PresburgerSet will be added in later patches. For now, only
FlatAffineValueConstraints and it's derivatives use these values.

[arjunp]

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

99–100

I would prefer getValue or maybeGetValue

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

1301

Maybe a PresburgerSpace::swapId should be introduced?

1549–1550

This seems more understandable. Also, use getIdKindOffset?

mlir/unittests/Dialect/Affine/Analysis/AffineStructuresParser.cpp

24–25

[bondhugula]

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

9–10

"space" in this doc comment isn't properly qualified. Expand and improve this doc comment?

102–104

Doc comment.

169

This name is too generic.

188

The TODO looks vague here. Did you mean to templatize in the future along
PresburgerSpace<Value>?

Can you state in the commit summary clearly with which users of PresburgerSpace these Value support you'd want? FlatAffineValueConstraints has a Value and you've now moved to a class with a purely mathematical name. It's better to I think think about the future scheme of things and hierarchy instead of multiple changes and moves.

[Groverkss]

Abandoning this patch to explore more design space for this interface. Thank you for your reviews!

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

public:
   /// Removes all equalities and inequalities.
   void clearConstraints();
 
   /// Sets the `values.size()` identifiers starting at `po`s to the specified
   /// values and removes them.
   void setAndEliminate(unsigned pos, ArrayRef<int64_t> values);
 
   /// Replaces the contents of this IntegerRelation with `other`.
-  virtual void clearAndCopyFrom(const IntegerRelation &other);
+  void clearAndCopyFrom(const IntegerRelation &other);
 
   /// Gather positions of all lower and upper bounds of the identifier at `pos`,
   /// and optionally any equalities on it. In addition, the bounds are to be
   /// independent of identifiers in position range [`offset`, `offset` + `num`).
   void
   getLowerAndUpperBoundIndices(unsigned pos,
                                SmallVectorImpl<unsigned> *lbIndices,
                                SmallVectorImpl<unsigned> *ubIndices,
protected:
 
   /// Eliminates the identifier at the specified position using Fourier-Motzkin
   /// variable elimination, but uses Gaussian elimination if there is an
   /// equality involving that identifier. 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.
-  virtual void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false,
-                                       bool *isResultIntegerExact = nullptr);
+  void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false,
 
   /// Tightens inequalities given that we are dealing with integer spaces. This
   /// is similar to the GCD test but applied to inequalities. The constant term
   /// can be reduced to the preceding multiple of the GCD of the coefficients,
   /// i.e.,
   ///  64*i - 100 >= 0  =>  64*i - 128 >= 0 (since 'i' is an integer). This is a
   /// fast method (linear in the number of coefficients).
   void gcdTightenInequalities();
protected:
 
   /// Returns true if the pos^th column is all zero for both inequalities and
   /// equalities.
   bool isColZero(unsigned pos) const;
 
   /// Returns false if the fields corresponding to various identifier counts, or
   /// equality/inequality buffer sizes aren't consistent; true otherwise. This
   /// is meant to be used within an assert internally.
-  virtual bool hasConsistentState() const;
+  bool hasConsistentState() const;
 
   /// Prints the number of constraints, dimensions, symbols and locals in the
   /// IntegerRelation.
   virtual void printSpace(raw_ostream &os) const;
 
   /// Removes identifiers in the column range [idStart, idLimit), and copies any
   /// remaining valid data into place, updates member variables, and resizes
   /// arrays as needed.

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

 //===- PresburgerSpace.h - MLIR PresburgerSpace 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
 //
 //===----------------------------------------------------------------------===//
 //
-// Classes representing space information like number of identifiers and kind of
-// identifiers.
+// Classes representing space information like number of identifiers, kind of
+// identifiers and values attached to identifiers.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
 #define MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
 
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 namespace presburger {
 
 /// Kind of identifier. Implementation wise SetDims are treated as Range
 /// ids, and spaces with no distinction between dimension ids are treated

 ///
 /// Identifiers are stored in the following order:
 ///       [Domain, Range, Symbols, Locals]
 ///
 /// A space with no distinction between types of dimension identifiers can
 /// be implemented as a space with zero domain. IdKind::SetDim should be used
 /// to refer to dimensions in such spaces.
 ///
+/// Each variable in the space can have a mlir::Value attached to it.
+///
 /// Compatibility of two spaces implies that number of identifiers of each kind
 /// other than Locals are equal. Equality of two spaces implies that number of
 /// identifiers of each kind are equal.
 class PresburgerSpace {
 public:
   static PresburgerSpace getRelationSpace(unsigned numDomain = 0,
                                           unsigned numRange = 0,
                                           unsigned numSymbols = 0,
public:
   unsigned getNumDimIds() const { return numDomain + numRange; }
   unsigned getNumDimAndSymbolIds() const {
     return numDomain + numRange + numSymbols;
   }
   unsigned getNumIds() const {
     return numDomain + numRange + numSymbols + numLocals;
   }
 
+  Optional<Value> &atValue(unsigned pos) { return values[pos]; }
+  const Optional<Value> &atValue(unsigned pos) const { return values[pos]; }
+
+  ArrayRef<Optional<Value>> getMaybeValues() const {
+    return {values.data(), values.size()};
+  }
+
   /// Get the number of ids of the specified kind.
   unsigned getNumIdKind(IdKind kind) const;
 
   /// Return the index at which the specified kind of id starts.
   unsigned getIdKindOffset(IdKind kind) const;
 
   /// Return the index at Which the specified kind of id ends.
   unsigned getIdKindEnd(IdKind kind) const;
public:
   /// first added identifier.
   unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1);
 
   /// Removes identifiers of the specified kind in the column range [idStart,
   /// idLimit). The range is relative to the kind of identifier.
   void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
 
   /// Returns true if both the spaces are compatible i.e. if both spaces have
-  /// the same number of identifiers of each kind (excluding locals).
+  /// the same number of identifiers of each kind (excluding locals) and
+  /// have the same values for each identifier (excluding locals).
   bool isCompatible(const PresburgerSpace &other) const;
 
+  /// Returns true if both the spaces are compatible i.e. if both spaces have
+  /// the same number of identifiers of each kind (excluding locals).
+  bool isCompatibleWithoutValues(const PresburgerSpace &other) const;
+
   /// Returns true if both the spaces are equal including local identifiers i.e.
   /// if both spaces have the same number of identifiers of each kind (including
-  /// locals).
+  /// locals) and have the same values for each identifer (excluding locals).
   bool isEqual(const PresburgerSpace &other) const;
 
+  /// Returns true if both the spaces are equal including local identifiers i.e.
+  /// if both spaces have the same number of identifiers of each kind (including
+  /// locals).
+  bool isEqualWithoutValues(const PresburgerSpace &other) const;
+
   /// Changes the partition between dimensions and symbols. Depending on the new
   /// symbol count, either a chunk of dimensional identifiers immediately before
   /// the split become symbols, or some of the symbols immediately after the
   /// split become dimensions.
   void setDimSymbolSeparation(unsigned newSymbolCount);
 
   void print(llvm::raw_ostream &os) const;
   void dump() const;
 
 protected:
   PresburgerSpace(unsigned numDomain = 0, unsigned numRange = 0,
                   unsigned numSymbols = 0, unsigned numLocals = 0)
       : numDomain(numDomain), numRange(numRange), numSymbols(numSymbols),
-        numLocals(numLocals) {}
+        numLocals(numLocals),
+        values(numDomain + numRange + numSymbols + numLocals, None) {}
+
+  /// Checks if the number of values are equal to the number of variables.
+  bool isConsistent() const;
 
 private:
   // Number of identifiers corresponding to domain identifiers.
   unsigned numDomain;
 
   // Number of identifiers corresponding to range identifiers.
   unsigned numRange;
 
   /// Number of identifiers corresponding to symbols (unknown but constant for
   /// analysis).
   unsigned numSymbols;
 
   /// Number of identifers corresponding to locals (identifiers corresponding
   /// to existentially quantified variables).
   unsigned numLocals;
+
+  /// Values corresponding to the identifiers of this space. Temporary ones or
+  /// those that aren't associated with any Value are set to llvm::None.
+  /// TODO: Remove dependency of Presburger on MLIRIR by making this general.
+  SmallVector<Optional<Value>, 8> values;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H

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

FlatAffineValueConstraints(unsigned numReservedInequalities,
                              unsigned numSymbols, unsigned numLocals,
                              ArrayRef<Optional<Value>> valArgs = {})
       : IntegerPolyhedron(numReservedInequalities, numReservedEqualities,
                           numReservedCols,
                           presburger::PresburgerSpace::getSetSpace(
                               numDims, numSymbols, numLocals)) {
     assert(numReservedCols >= getNumIds() + 1);
     assert(valArgs.empty() || valArgs.size() == getNumIds());
-    values.reserve(numReservedCols);
-    if (valArgs.empty())
-      values.resize(getNumIds(), None);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+
+    // Set values.
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      atValue(i) = valArgs[i];
   }
 
   /// Constructs a constraint system with the specified number of
   /// dimensions and symbols.
   FlatAffineValueConstraints(unsigned numDims = 0, unsigned numSymbols = 0,
                              unsigned numLocals = 0,
                              ArrayRef<Optional<Value>> valArgs = {})
       : FlatAffineValueConstraints(/*numReservedInequalities=*/0,
                                    /*numReservedEqualities=*/0,
                                    /*numReservedCols=*/numDims + numSymbols +
                                        numLocals + 1,
                                    numDims, numSymbols, numLocals, valArgs) {}
 
   FlatAffineValueConstraints(const IntegerPolyhedron &fac,
                              ArrayRef<Optional<Value>> valArgs = {})
       : IntegerPolyhedron(fac) {
     assert(valArgs.empty() || valArgs.size() == getNumIds());
-    if (valArgs.empty())
-      values.resize(getNumIds(), None);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+
+    // Set values.
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      atValue(i) = valArgs[i];
   }
 
   /// 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);
   explicit FlatAffineValueConstraints(ArrayRef<const AffineValueMap *> avmRef);
 
   /// Creates an affine constraint system from an IntegerSet.
public:
   unsigned insertSymbolId(unsigned pos, unsigned num = 1) {
     return insertId(IdKind::Symbol, pos, num);
   }
   unsigned insertLocalId(unsigned pos, unsigned num = 1) {
     return insertId(IdKind::Local, pos, num);
   }
   unsigned insertDimId(unsigned pos, ValueRange vals);
   unsigned insertSymbolId(unsigned pos, ValueRange vals);
-  unsigned insertId(presburger::IdKind kind, unsigned pos,
-                    unsigned num = 1) override;
   unsigned insertId(presburger::IdKind kind, unsigned pos, ValueRange vals);
+  using IntegerPolyhedron::insertId;
 
   /// Append identifiers of the specified kind after the last identifier of that
   /// kind. The coefficient columns corresponding to the added identifiers are
   /// initialized to zero. `vals` are the Values corresponding to the
   /// identifiers. Return the position of the first added column.
   ///
   /// Note: Empty Values are allowed in `vals`.
   unsigned appendDimId(ValueRange vals);
   unsigned appendSymbolId(ValueRange vals);
   unsigned appendDimId(unsigned num = 1) {
     return appendId(IdKind::SetDim, num);
   }
   unsigned appendSymbolId(unsigned num = 1) {
     return appendId(IdKind::Symbol, num);
   }
   unsigned appendLocalId(unsigned num = 1) {
     return appendId(IdKind::Local, num);
   }
 
-  /// 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(presburger::IdKind kind, unsigned idStart,
-                     unsigned idLimit) override;
-  using IntegerPolyhedron::removeIdRange;
-
   /// Add the specified values as a dim or symbol id depending on its nature, if
   /// it already doesn't exist in the system. `val` has to be either a terminal
   /// symbol or a loop IV, i.e., it cannot be the result affine.apply of any
   /// symbols or loop IVs. The identifier is added to the end of the existing
   /// dims or symbols. Additional information on the identifier is extracted
   /// from the IR and added to the constraint system.
   void addInductionVarOrTerminalSymbol(Value val);
 
public:
   void mergeAndAlignIdsWithOther(unsigned offset,
                                  FlatAffineValueConstraints *other);
 
   /// Returns true if this constraint system and `other` are in the same
   /// space, i.e., if they are associated with the same set of identifiers,
   /// appearing in the same order. Returns false otherwise.
   bool areIdsAlignedWithOther(const FlatAffineValueConstraints &other);
 
-  /// Replaces the contents of this FlatAffineValueConstraints with `other`.
-  void clearAndCopyFrom(const IntegerRelation &other) override;
+  Optional<Value> &atValue(unsigned pos) { return space.atValue(pos); }
+  const Optional<Value> &atValue(unsigned pos) const {
+    return space.atValue(pos);
+  }
 
   /// Returns the Value associated with the pos^th identifier. Asserts if
   /// no Value identifier was associated.
   inline Value getValue(unsigned pos) const {
     assert(hasValue(pos) && "identifier's Value not set");
-    return values[pos].getValue();
+    return atValue(pos).getValue();
   }
 
   /// Returns true if the pos^th identifier has an associated Value.
-  inline bool hasValue(unsigned pos) const { return values[pos].hasValue(); }
+  inline bool hasValue(unsigned pos) const { return atValue(pos).hasValue(); }
 
   /// Returns true if at least one identifier has an associated Value.
   bool hasValues() const;
 
   /// Returns the Values associated with identifiers in range [start, end).
   /// Asserts if no Value was associated with one of these identifiers.
   inline void getValues(unsigned start, unsigned end,
                         SmallVectorImpl<Value> *values) const {
     assert((start < getNumIds() || start == end) && "invalid start position");
     assert(end <= getNumIds() && "invalid end position");
     values->clear();
     values->reserve(end - start);
     for (unsigned i = start; i < end; i++)
       values->push_back(getValue(i));
   }
   inline void getAllValues(SmallVectorImpl<Value> *values) const {
     getValues(0, getNumIds(), values);
   }
 
   inline ArrayRef<Optional<Value>> getMaybeValues() const {
-    return {values.data(), values.size()};
+    return space.getMaybeValues();
   }
 
   inline ArrayRef<Optional<Value>> getMaybeDimValues() const {
-    return {values.data(), getNumDimIds()};
+    return getMaybeValues().slice(space.getIdKindOffset(IdKind::SetDim),
+                                  space.getNumIdKind(IdKind::SetDim));
   }
 
   inline ArrayRef<Optional<Value>> getMaybeSymbolValues() const {
-    return {values.data() + getNumDimIds(), getNumSymbolIds()};
+    return getMaybeValues().slice(space.getIdKindOffset(IdKind::Symbol),
+                                  space.getNumIdKind(IdKind::Symbol));
   }
 
   inline ArrayRef<Optional<Value>> getMaybeDimAndSymbolValues() const {
-    return {values.data(), getNumDimIds() + getNumSymbolIds()};
+    return getMaybeValues().slice(space.getIdKindOffset(IdKind::SetDim),
+                                  space.getNumDimAndSymbolIds());
   }
 
   /// Sets the Value associated with the pos^th identifier.
   inline void setValue(unsigned pos, Value val) {
     assert(pos < getNumIds() && "invalid id position");
-    values[pos] = val;
+    atValue(pos) = val;
   }
 
   /// Sets the Values associated with the identifiers in the range [start, end).
   void setValues(unsigned start, unsigned end, ArrayRef<Value> values) {
     assert((start < getNumIds() || end == start) && "invalid start position");
     assert(end <= getNumIds() && "invalid end position");
     assert(values.size() == end - start);
     for (unsigned i = start; i < end; ++i)
       setValue(i, values[i - start]);
   }
 
   /// Merge and align symbols of `this` and `other` such that both get union of
   /// of symbols that are unique. Symbols in `this` and `other` should be
   /// unique. Symbols with Value as `None` are considered to be inequal to all
   /// other symbols.
   void mergeSymbolIds(FlatAffineValueConstraints &other);
 
 protected:
   using IdKind = presburger::IdKind;
 
-  /// Returns false if the fields corresponding to various identifier 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;
-
   /// Given an affine map that is aligned with this constraint system:
   /// * Flatten the map.
   /// * Add newly introduced local columns at the beginning of this constraint
   ///   system (local column pos 0).
   /// * Add equalities that define the new local columns to this constraint
   ///   system.
   /// * Return the flattened expressions via `flattenedExprs`.
   ///
   /// Note: This is a shared helper function of `addLowerOrUpperBound` and
   ///       `composeMatchingMap`.
   LogicalResult flattenAlignedMapAndMergeLocals(
       AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs);
 
-  /// Eliminates the identifier at the specified position using Fourier-Motzkin
-  /// variable elimination, but uses Gaussian elimination if there is an
-  /// equality involving that identifier. 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;
-
   /// Prints the number of constraints, dimensions, symbols and locals in the
   /// FlatAffineConstraints. Also, prints for each identifier whether there is
   /// an SSA Value attached to it.
   void printSpace(raw_ostream &os) const override;
-
-  /// Values corresponding to the (column) identifiers of this constraint
-  /// system appearing in the order the identifiers correspond to columns.
-  /// Temporary ones or those that aren't associated with any Value are set to
-  /// None.
-  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 {

mlir/lib/Analysis/Presburger/CMakeLists.txt

 add_mlir_library(MLIRPresburger
   IntegerRelation.cpp
   LinearTransform.cpp
   Matrix.cpp
   PresburgerRelation.cpp
   PresburgerSpace.cpp
   PWMAFunction.cpp
   Simplex.cpp
   Utils.cpp
 
   LINK_LIBS PUBLIC
   MLIRSupport
+  MLIRIR
   )

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

std::unique_ptr<IntegerRelation> IntegerRelation::clone() const {
   return std::make_unique<IntegerRelation>(*this);
 }
 
 std::unique_ptr<IntegerPolyhedron> IntegerPolyhedron::clone() const {
   return std::make_unique<IntegerPolyhedron>(*this);
 }
 
 void IntegerRelation::append(const IntegerRelation &other) {
-  assert(space.isEqual(other.getSpace()) && "Spaces must be equal.");
+  assert(space.isEqualWithoutValues(other.getSpace()) &&
+         "Spaces must be equal.");
 
   inequalities.reserveRows(inequalities.getNumRows() +
                            other.getNumInequalities());
   equalities.reserveRows(equalities.getNumRows() + other.getNumEqualities());
 
   for (unsigned r = 0, e = other.getNumInequalities(); r < e; r++) {
     addInequality(other.getInequality(r));
   }
   for (unsigned r = 0, e = other.getNumEqualities(); r < e; r++) {
     addEquality(other.getEquality(r));
   }
 }
 
 IntegerRelation IntegerRelation::intersect(IntegerRelation other) const {
   IntegerRelation result = *this;
   result.mergeLocalIds(other);
   result.append(other);
   return result;
 }
 
 bool IntegerRelation::isEqual(const IntegerRelation &other) const {
-  assert(space.isEqual(other.getSpace()) && "Spaces must be equal.");
+  assert(space.isEqualWithoutValues(other.getSpace()) &&
+         "Spaces must be equal.");
   return PresburgerRelation(*this).isEqual(PresburgerRelation(other));
 }
 
 bool IntegerRelation::isSubsetOf(const IntegerRelation &other) const {
-  assert(space.isEqual(other.getSpace()) && "Spaces must be equal.");
+  assert(space.isEqualWithoutValues(other.getSpace()) &&
+         "Spaces must be equal.");
   return PresburgerRelation(*this).isSubsetOf(PresburgerRelation(other));
 }
 
 MaybeOptimum<SmallVector<Fraction, 8>>
 IntegerRelation::findRationalLexMin() const {
   assert(getNumSymbolIds() == 0 && "Symbols are not supported!");
   MaybeOptimum<SmallVector<Fraction, 8>> maybeLexMin =
       LexSimplex(*this).findRationalLexMin();

 ///
 /// To detect local ids that always take the same in both sets, each local id is
 /// represented as a floordiv with constant denominator in terms of other ids.
 /// After extracting these divisions, local ids with the same division
 /// representation are considered duplicate and are merged. It is possible that
 /// division representation for some local id cannot be obtained, and thus these
 /// local ids are not considered for detecting duplicates.
 void IntegerRelation::mergeLocalIds(IntegerRelation &other) {
-  assert(space.isCompatible(other.getSpace()) &&
+  assert(space.isCompatibleWithoutValues(other.getSpace()) &&
          "Spaces should be compatible.");
 
   IntegerRelation &relA = *this;
   IntegerRelation &relB = other;
 
   // Merge local ids of relA and relB without using division information,
   // i.e. append local ids of `relB` to `relA` and insert local ids of `relA`
   // to `relB` at start of its local ids.
for (unsigned r = 0, e = a.getNumEqualities(); r < e; ++r) {
     }
   }
 }
 
 // Computes the bounding box with respect to 'other' by finding the min of the
 // lower bounds and the max of the upper bounds along each of the dimensions.
 LogicalResult
 IntegerRelation::unionBoundingBox(const IntegerRelation &otherCst) {
-  assert(space.isEqual(otherCst.getSpace()) && "Spaces should match.");
+  assert(space.isEqualWithoutValues(otherCst.getSpace()) &&
+         "Spaces should match.");
   assert(getNumLocalIds() == 0 && "local ids not supported yet here");
 
   // Get the constraints common to both systems; these will be added as is to
   // the union.
   IntegerRelation commonCst(PresburgerSpace::getRelationSpace());
   getCommonConstraints(*this, otherCst, commonCst);
 
   std::vector<SmallVector<int64_t, 8>> boundingLbs;

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

void MultiAffineFunction::print(raw_ostream &os) const {
   os << "Output:\n";
   output.print(os);
   os << "\n";
 }
 
 void MultiAffineFunction::dump() const { print(llvm::errs()); }
 
 bool MultiAffineFunction::isEqual(const MultiAffineFunction &other) const {
-  return space.isCompatible(other.getSpace()) &&
+  return space.isCompatibleWithoutValues(other.getSpace()) &&
          getDomain().isEqual(other.getDomain()) &&
          isEqualWhereDomainsOverlap(other);
 }
 
 unsigned MultiAffineFunction::insertId(IdKind kind, unsigned pos,
                                        unsigned num) {
   assert((kind != IdKind::Domain || num == 0) &&
          "Domain has to be zero in a set");
void PWMAFunction::truncateOutput(unsigned count) {
   assert(count <= numOutputs);
   for (MultiAffineFunction &piece : pieces)
     piece.truncateOutput(count);
   numOutputs = count;
 }
 
 bool MultiAffineFunction::isEqualWhereDomainsOverlap(
     MultiAffineFunction other) const {
-  if (!space.isCompatible(other.getSpace()))
+  if (!space.isCompatibleWithoutValues(other.getSpace()))
     return false;
 
   // `commonFunc` has the same output as `this`.
   MultiAffineFunction commonFunc = *this;
   // After this merge, `commonFunc` and `other` have the same local ids; they
   // are merged.
   commonFunc.mergeLocalIds(other);
   // After this, the domain of `commonFunc` will be the intersection of the
bool MultiAffineFunction::isEqualWhereDomainsOverlap(
   // If the whole common domain is a subset of commonDomainMatching, then they
   // are equal and the two functions match on the whole common domain.
   return commonFunc.getDomain().isSubsetOf(commonDomainMatching);
 }
 
 /// Two PWMAFunctions are equal if they have the same dimensionalities,
 /// the same domain, and take the same value at every point in the domain.
 bool PWMAFunction::isEqual(const PWMAFunction &other) const {
-  if (!space.isCompatible(other.space))
+  if (!space.isCompatibleWithoutValues(other.space))
     return false;
 
   if (!this->getDomain().isEqual(other.getDomain()))
     return false;
 
   // Check if, whenever the domains of a piece of `this` and a piece of `other`
   // overlap, they take the same output value. If `this` and `other` have the
   // same domain (checked above), then this check passes iff the two functions
   // have the same output at every point in the domain.
   for (const MultiAffineFunction &aPiece : this->pieces)
     for (const MultiAffineFunction &bPiece : other.pieces)
       if (!aPiece.isEqualWhereDomainsOverlap(bPiece))
         return false;
   return true;
 }
 
 void PWMAFunction::addPiece(const MultiAffineFunction &piece) {
-  assert(space.isCompatible(piece.getSpace()) &&
+  assert(space.isCompatibleWithoutValues(piece.getSpace()) &&
          "Piece to be added is not compatible with this PWMAFunction!");
   assert(piece.isConsistent() && "Piece is internally inconsistent!");
   assert(this->getDomain()
              .intersect(PresburgerSet(piece.getDomain()))
              .isIntegerEmpty() &&
          "New piece's domain overlaps with that of existing pieces!");
   pieces.push_back(piece);
 }

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

 const IntegerRelation &PresburgerRelation::getDisjunct(unsigned index) const {
   assert(index < disjuncts.size() && "index out of bounds!");
   return disjuncts[index];
 }
 
 /// Mutate this set, turning it into the union of this set and the given
 /// IntegerRelation.
 void PresburgerRelation::unionInPlace(const IntegerRelation &disjunct) {
-  assert(space.isCompatible(disjunct.getSpace()) && "Spaces should match");
+  assert(space.isCompatibleWithoutValues(disjunct.getSpace()) &&
+         "Spaces should match");
   disjuncts.push_back(disjunct);
 }
 
 /// Mutate this set, turning it into the union of this set and the given set.
 ///
 /// This is accomplished by simply adding all the disjuncts of the given set
 /// to this set.
 void PresburgerRelation::unionInPlace(const PresburgerRelation &set) {
-  assert(space.isCompatible(set.getSpace()) && "Spaces should match");
+  assert(space.isCompatibleWithoutValues(set.getSpace()) &&
+         "Spaces should match");
   for (const IntegerRelation &disjunct : set.disjuncts)
     unionInPlace(disjunct);
 }
 
 /// Return the union of this set and the given set.
 PresburgerRelation
 PresburgerRelation::unionSet(const PresburgerRelation &set) const {
-  assert(space.isCompatible(set.getSpace()) && "Spaces should match");
+  assert(space.isCompatibleWithoutValues(set.getSpace()) &&
+         "Spaces should match");
   PresburgerRelation result = *this;
   result.unionInPlace(set);
   return result;
 }
 
 /// A point is contained in the union iff any of the parts contain the point.
 bool PresburgerRelation::containsPoint(ArrayRef<int64_t> point) const {
   return llvm::any_of(disjuncts, [&](const IntegerRelation &disjunct) {

 //
 // We directly compute (S_1 or S_2 ...) and (T_1 or T_2 ...)
 // as (S_1 and T_1) or (S_1 and T_2) or ...
 //
 // If S_i or T_j have local variables, then S_i and T_j contains the local
 // variables of both.
 PresburgerRelation
 PresburgerRelation::intersect(const PresburgerRelation &set) const {
-  assert(space.isCompatible(set.getSpace()) && "Spaces should match");
+  assert(space.isCompatibleWithoutValues(set.getSpace()) &&
+         "Spaces should match");
 
   PresburgerRelation result(getSpace());
   for (const IntegerRelation &csA : disjuncts) {
     for (const IntegerRelation &csB : set.disjuncts) {
       IntegerRelation intersection = csA.intersect(csB);
       if (!intersection.isEmpty())
         result.unionInPlace(intersection);
     }

 /// As a heuristic, we try adding all the constraints and check if simplex
 /// says that the intersection is empty. If it is, then subtracting this
 /// disjuncts is a no-op and we just skip it. Also, in the process we find out
 /// that some constraints are redundant. These redundant constraints are
 /// ignored.
 ///
 static PresburgerRelation getSetDifference(IntegerRelation b,
                                            const PresburgerRelation &s) {
-  assert(b.getSpace().isCompatible(s.getSpace()) && "Spaces should match");
+  assert(b.getSpace().isCompatibleWithoutValues(s.getSpace()) &&
+         "Spaces should match");
   if (b.isEmptyByGCDTest())
     return PresburgerRelation::getEmpty(b.getSpaceWithoutLocals());
 
   // Remove duplicate divs up front here to avoid existing
   // divs disappearing in the call to mergeLocalIds below.
   b.removeDuplicateDivs();
 
   PresburgerRelation result =

 PresburgerRelation PresburgerRelation::complement() const {
   return getSetDifference(IntegerRelation::getUniverse(getSpace()), *this);
 }
 
 /// Return the result of subtract the given set from this set, i.e.,
 /// return `this \ set`.
 PresburgerRelation
 PresburgerRelation::subtract(const PresburgerRelation &set) const {
-  assert(space.isCompatible(set.getSpace()) && "Spaces should match");
+  assert(space.isCompatibleWithoutValues(set.getSpace()) &&
+         "Spaces should match");
   PresburgerRelation result(getSpace());
   // We compute (U_i t_i) \ (U_i set_i) as U_i (t_i \ V_i set_i).
   for (const IntegerRelation &disjunct : disjuncts)
     result.unionInPlace(getSetDifference(disjunct, set));
   return result;
 }
 
 /// T is a subset of S iff T \ S is empty, since if T \ S contains a
 /// point then this is a point that is contained in T but not S, and
 /// if T contains a point that is not in S, this also lies in T \ S.
 bool PresburgerRelation::isSubsetOf(const PresburgerRelation &set) const {
   return this->subtract(set).isIntegerEmpty();
 }
 
 /// Two sets are equal iff they are subsets of each other.
 bool PresburgerRelation::isEqual(const PresburgerRelation &set) const {
-  assert(space.isCompatible(set.getSpace()) && "Spaces should match");
+  assert(space.isCompatibleWithoutValues(set.getSpace()) &&
+         "Spaces should match");
   return this->isSubsetOf(set) && set.isSubsetOf(*this);
 }
 
 /// Return true if all the sets in the union are known to be integer empty,
 /// false otherwise.
 bool PresburgerRelation::isIntegerEmpty() const {
   // The set is empty iff all of the disjuncts are empty.
   return llvm::all_of(disjuncts, std::mem_fn(&IntegerRelation::isIntegerEmpty));

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

if (kind == IdKind::Domain)
     numDomain += num;
   else if (kind == IdKind::Range)
     numRange += num;
   else if (kind == IdKind::Symbol)
     numSymbols += num;
   else
     numLocals += num;
 
+  // Insert values for newly added variables.
+  values.insert(values.begin() + absolutePos, num, None);
+
+  assert(isConsistent() && "Space must be consistent.");
+
   return absolutePos;
 }
 
 void PresburgerSpace::removeIdRange(IdKind kind, unsigned idStart,
                                     unsigned idLimit) {
   assert(idLimit <= getNumIdKind(kind) && "invalid id limit");
 
   if (idStart >= idLimit)
     return;
 
   unsigned numIdsEliminated = idLimit - idStart;
   if (kind == IdKind::Domain)
     numDomain -= numIdsEliminated;
   else if (kind == IdKind::Range)
     numRange -= numIdsEliminated;
   else if (kind == IdKind::Symbol)
     numSymbols -= numIdsEliminated;
   else
     numLocals -= numIdsEliminated;
+
+  // Remove values for removed variables.
+  unsigned offset = getIdKindOffset(kind);
+  values.erase(values.begin() + offset + idStart,
+               values.begin() + offset + idLimit);
+
+  assert(isConsistent() && "Space must be consistent.");
+}
+
+bool PresburgerSpace::isConsistent() const {
+  return values.size() == getNumIds();
 }
 
 bool PresburgerSpace::isCompatible(const PresburgerSpace &other) const {
+  assert(isConsistent() && "Space must be consistent.");
+  return isCompatibleWithoutValues(other) &&
+         getMaybeValues().slice(0, getNumDimAndSymbolIds()) ==
+             other.getMaybeValues().slice(0, other.getNumDimAndSymbolIds());
+}
+
+bool PresburgerSpace::isCompatibleWithoutValues(
+    const PresburgerSpace &other) const {
+  assert(isConsistent() && "Space must be consistent.");
   return getNumDomainIds() == other.getNumDomainIds() &&
          getNumRangeIds() == other.getNumRangeIds() &&
          getNumSymbolIds() == other.getNumSymbolIds();
 }
 
 bool PresburgerSpace::isEqual(const PresburgerSpace &other) const {
-  return isCompatible(other) && getNumLocalIds() == other.getNumLocalIds();
+  assert(isConsistent() && "Space must be consistent.");
+  return isEqualWithoutValues(other) &&
+         getMaybeValues() == other.getMaybeValues();
+}
+
+bool PresburgerSpace::isEqualWithoutValues(const PresburgerSpace &other) const {
+  assert(isConsistent() && "Space must be consistent.");
+  return isCompatibleWithoutValues(other) &&
+         getNumLocalIds() == other.getNumLocalIds();
 }
 
 void PresburgerSpace::setDimSymbolSeparation(unsigned newSymbolCount) {
+  assert(isConsistent() && "Space must be consistent.");
   assert(newSymbolCount <= getNumDimAndSymbolIds() &&
          "invalid separation position");
+  // No need to modify `values`, since only the seperation is changed.
   numRange = numRange + numSymbols - newSymbolCount;
   numSymbols = newSymbolCount;
 }
 
 void PresburgerSpace::print(llvm::raw_ostream &os) const {
   os << "Domain: " << getNumDomainIds() << ", "
      << "Range: " << getNumRangeIds() << ", "
      << "Symbols: " << getNumSymbolIds() << ", "
      << "Locals: " << getNumLocalIds() << "\n";
 }
 
 void PresburgerSpace::dump() const { print(llvm::errs()); }

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

 // Construct from an IntegerSet.
 FlatAffineValueConstraints::FlatAffineValueConstraints(IntegerSet set)
     : IntegerPolyhedron(set.getNumInequalities(), set.getNumEqualities(),
                         set.getNumDims() + set.getNumSymbols() + 1,
                         PresburgerSpace::getSetSpace(set.getNumDims(),
                                                      set.getNumSymbols(),
                                                      /*numLocals=*/0)) {
 
-  // Resize values.
-  values.resize(getNumIds(), None);
-
   // Flatten expressions and add them to the constraint system.
   std::vector<SmallVector<int64_t, 8>> flatExprs;
   FlatAffineValueConstraints localVarCst;
   if (failed(getFlattenedAffineExprs(set, &flatExprs, &localVarCst))) {
     assert(false && "flattening unimplemented for semi-affine integer sets");
     return;
   }
   assert(flatExprs.size() == set.getNumConstraints());

 }
 
 unsigned FlatAffineValueConstraints::insertSymbolId(unsigned pos,
                                                     ValueRange vals) {
   return insertId(IdKind::Symbol, pos, vals);
 }
 
 unsigned FlatAffineValueConstraints::insertId(IdKind kind, unsigned pos,
-                                              unsigned num) {
-  unsigned absolutePos = IntegerPolyhedron::insertId(kind, pos, num);
-  values.insert(values.begin() + absolutePos, num, None);
-  assert(values.size() == getNumIds());
-  return absolutePos;
-}
-
-unsigned FlatAffineValueConstraints::insertId(IdKind kind, unsigned pos,
                                               ValueRange vals) {
   assert(!vals.empty() && "expected ValueRange with Values");
   unsigned num = vals.size();
   unsigned absolutePos = IntegerPolyhedron::insertId(kind, pos, num);
 
   // If a Value is provided, insert it; otherwise use None.
   for (unsigned i = 0; i < num; ++i)
-    values.insert(values.begin() + absolutePos + i,
-                  vals[i] ? Optional<Value>(vals[i]) : None);
+    atValue(absolutePos + i) = vals[i] ? Optional<Value>(vals[i]) : None;
 
-  assert(values.size() == getNumIds());
   return absolutePos;
 }
 
 bool FlatAffineValueConstraints::hasValues() const {
-  return llvm::find_if(values, [](Optional<Value> id) {
-           return id.hasValue();
-         }) != values.end();
+  for (unsigned i = 0, e = getNumIds(); i < e; ++i)
+    if (atValue(i).hasValue())
+      return true;
+  return false;
 }
 
 /// Checks if two constraint systems are in the same space, i.e., if they are
 /// associated with the same set of identifiers, appearing in the same order.
 static bool areIdsAligned(const FlatAffineValueConstraints &a,
                           const FlatAffineValueConstraints &b) {
   return a.getNumDimIds() == b.getNumDimIds() &&
          a.getNumSymbolIds() == b.getNumSymbolIds() &&
void FlatAffineValueConstraints::addAffineIfOpDomain(AffineIfOp ifOp) {
 
   // Merge the constraints from ifOp to the current domain. We need first merge
   // and align the IDs from both constraints, and then append the constraints
   // from the ifOp into the current one.
   mergeAndAlignIdsWithOther(0, &cst);
   append(cst);
 }
 
-bool FlatAffineValueConstraints::hasConsistentState() const {
-  return IntegerPolyhedron::hasConsistentState() &&
-         values.size() == getNumIds();
-}
-
-void FlatAffineValueConstraints::removeIdRange(IdKind kind, unsigned idStart,
-                                               unsigned idLimit) {
-  IntegerPolyhedron::removeIdRange(kind, idStart, idLimit);
-  unsigned offset = getIdKindOffset(kind);
-  values.erase(values.begin() + idStart + offset,
-               values.begin() + idLimit + offset);
-}
-
 // Determine whether the identifier at 'pos' (say id_r) can be expressed as
 // modulo of another known identifier (say id_n) w.r.t a constant. For example,
 // if the following constraints hold true:
 // ```
 // 0 <= id_r <= divisor - 1
 // id_n - (divisor * q_expr) = id_r
 // ```
 // where `id_n` is a known identifier (called dividend), and `q_expr` is an
#else
   SmallVector<Value> *newSymsPtr = nullptr;
 #endif // NDEBUG
 
   dims.reserve(getNumDimIds());
   syms.reserve(getNumSymbolIds());
   for (unsigned i = getIdKindOffset(IdKind::SetDim),
                 e = getIdKindEnd(IdKind::SetDim);
        i < e; ++i)
-    dims.push_back(values[i] ? *values[i] : Value());
+    dims.push_back(atValue(i) ? *atValue(i) : Value());
   for (unsigned i = getIdKindOffset(IdKind::Symbol),
                 e = getIdKindEnd(IdKind::Symbol);
        i < e; ++i)
-    syms.push_back(values[i] ? *values[i] : Value());
+    syms.push_back(atValue(i) ? *atValue(i) : Value());
 
   AffineMap alignedMap =
       alignAffineMapWithValues(map, operands, dims, syms, newSymsPtr);
   // All symbols are already part of this FlatAffineConstraints.
   assert(syms.size() == newSymsPtr->size() && "unexpected new/missing symbols");
   assert(std::equal(syms.begin(), syms.end(), newSymsPtr->begin()) &&
          "unexpected new/missing symbols");
   return alignedMap;
if (lbMap && lbMap.getNumResults() != 0 && ubMap &&
       if (failed(this->addAffineForOpDomain(loop)))
         return failure();
     }
   }
   return success();
 }
 
 bool FlatAffineValueConstraints::findId(Value val, unsigned *pos) const {
-  unsigned i = 0;
-  for (const auto &mayBeId : values) {
-    if (mayBeId.hasValue() && mayBeId.getValue() == val) {
+  for (unsigned i = 0, e = getNumIds(); i < e; ++i) {
+    const Optional<Value> &spaceVal = atValue(i);
+    if (spaceVal.hasValue() && spaceVal.getValue() == val) {
       *pos = i;
       return true;
     }
-    i++;
   }
+
   return false;
 }
 
 bool FlatAffineValueConstraints::containsId(Value val) const {
-  return llvm::any_of(values, [&](const Optional<Value> &mayBeId) {
-    return mayBeId.hasValue() && mayBeId.getValue() == val;
-  });
+  unsigned pos;
+  return findId(val, &pos);
 }
 
 void FlatAffineValueConstraints::swapId(unsigned posA, unsigned posB) {
   IntegerPolyhedron::swapId(posA, posB);
-  std::swap(values[posA], values[posB]);
+  std::swap(atValue(posA), atValue(posB));
 }
 
 void FlatAffineValueConstraints::addBound(BoundType type, Value val,
                                           int64_t value) {
   unsigned pos;
   if (!findId(val, &pos))
     // This is a pre-condition for this method.
     assert(0 && "id not found");
   addBound(type, pos, value);
 }
 
 void FlatAffineValueConstraints::printSpace(raw_ostream &os) const {
   IntegerPolyhedron::printSpace(os);
   os << "(";
   for (unsigned i = 0, e = getNumIds(); i < e; i++) {
     if (hasValue(i))
       os << "Value ";
     else
       os << "None ";
   }
   os << " const)\n";
 }
 
-void FlatAffineValueConstraints::clearAndCopyFrom(
-    const IntegerRelation &other) {
-
-  if (auto *otherValueSet =
-          dyn_cast<const FlatAffineValueConstraints>(&other)) {
-    *this = *otherValueSet;
-  } else {
-    *static_cast<IntegerRelation *>(this) = other;
-    values.clear();
-    values.resize(getNumIds(), None);
-  }
-}
-
-void FlatAffineValueConstraints::fourierMotzkinEliminate(
-    unsigned pos, bool darkShadow, bool *isResultIntegerExact) {
-  SmallVector<Optional<Value>, 8> newVals;
-  newVals.reserve(getNumIds() - 1);
-  newVals.append(values.begin(), values.begin() + pos);
-  newVals.append(values.begin() + pos + 1, values.end());
-  // Note: Base implementation discards all associated Values.
-  IntegerPolyhedron::fourierMotzkinEliminate(pos, darkShadow,
-                                             isResultIntegerExact);
-  values = newVals;
-  assert(values.size() == getNumIds());
-}
-
 void FlatAffineValueConstraints::projectOut(Value val) {
   unsigned pos;
   bool ret = findId(val, &pos);
   assert(ret);
   (void)ret;
   fourierMotzkinEliminate(pos);
 }
 
FlatAffineValueConstraints FlatAffineRelation::getRangeSet() const {
   // Convert all domain variables to local variables.
   range.convertToLocal(IdKind::SetDim, 0, getNumDomainDims());
   return range;
 }
 
 void FlatAffineRelation::compose(const FlatAffineRelation &other) {
   assert(getNumDomainDims() == other.getNumRangeDims() &&
          "Domain of this and range of other do not match");
-  assert(std::equal(values.begin(), values.begin() + getNumDomainDims(),
-                    other.values.begin() + other.getNumDomainDims()) &&
+  assert(getMaybeValues().take_front(getNumDomainDims()) ==
+             other.getMaybeValues().slice(other.getNumDomainDims(),
+                                          getNumDomainDims()) &&
          "Domain of this and range of other do not match");
 
   FlatAffineRelation rel = other;
 
   // Convert `rel` from
   //    [otherDomain] -> [otherRange]
   // to
   //    [otherDomain] -> [otherRange thisRange]

mlir/unittests/Dialect/Affine/Analysis/AffineStructuresParser.cpp

 FailureOr<IntegerPolyhedron>
 mlir::parseIntegerSetToFAC(llvm::StringRef str, MLIRContext *context,
                            bool printDiagnosticInfo) {
   IntegerSet set = parseIntegerSet(str, context, printDiagnosticInfo);
 
   if (!set)
     return failure();
 
-  return FlatAffineValueConstraints(set);
+  // Set values in FAC to llvm::None.
+  // TODO: Check if we can also introduce use Value in tests.
+  FlatAffineValueConstraints fac(set);
+  for (unsigned i = 0, e = fac.getNumIds(); i < e; ++i)
+    fac.atValue(i) = llvm::None;
+
+  return fac;
 }