[MLIR][Presburger] Use PresburgerSpace in constructors

[Groverkss]

Summary

This patch modifies IntegerPolyhedron, IntegerRelation, PresburgerRelation,
PresburgerSet, PWMAFunction, constructors to take PresburgerSpace instead of
dimensions. This allows information present in PresburgerSpace to be carried
better and allows for a general interface.

[arjunp]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

1881–1883

Why this change? Maybe they wanted to avoid mallocs on copying, though I don't know why RVO wouldn't take place...

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

458–460

Can just store a space here and then pass it below instead. (Not necessarily for this patch.)

[Groverkss]

• Revert change to getCommonConstraints

[Groverkss]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

1881–1883

I will send this change as a patch after this patch and we can discuss this. Reverting this back for now.

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

458–460

I'll send it as patch after this patch.

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

enum class Kind {
     IntegerRelation,
     IntegerPolyhedron,
   };
 
   /// Constructs a relation reserving memory for the specified number
   /// of constraints and identifiers.
   IntegerRelation(unsigned numReservedInequalities,
                   unsigned numReservedEqualities, unsigned numReservedCols,
-                  unsigned numDomain, unsigned numRange, unsigned numSymbols,
-                  unsigned numLocals)
-      : PresburgerSpace(numDomain, numRange, numSymbols, numLocals),
+                  const PresburgerSpace &space)
+      : PresburgerSpace(space),
         equalities(0, getNumIds() + 1, numReservedEqualities, numReservedCols),
         inequalities(0, getNumIds() + 1, numReservedInequalities,
                      numReservedCols) {
     assert(numReservedCols >= getNumIds() + 1);
   }
 
   /// Constructs a relation with the specified number of dimensions and symbols.
-  IntegerRelation(unsigned numDomain = 0, unsigned numRange = 0,
-                  unsigned numSymbols = 0, unsigned numLocals = 0)
+  IntegerRelation(const PresburgerSpace &space)
       : IntegerRelation(/*numReservedInequalities=*/0,
                         /*numReservedEqualities=*/0,
-                        /*numReservedCols=*/numDomain + numRange + numSymbols +
-                            numLocals + 1,
-                        numDomain, numRange, numSymbols, numLocals) {}
+                        /*numReservedCols=*/space.getNumIds() + 1, space) {}
 
   /// Return a system with no constraints, i.e., one which is satisfied by all
   /// points.
-  static IntegerRelation getUniverse(unsigned numDomain = 0,
-                                     unsigned numRange = 0,
-                                     unsigned numSymbols = 0) {
-    return IntegerRelation(numDomain, numRange, numSymbols);
+  static IntegerRelation getUniverse(const PresburgerSpace &space) {
+    return IntegerRelation(space);
   }
 
   /// Return the kind of this IntegerRelation.
   virtual Kind getKind() const { return Kind::IntegerRelation; }
 
   static bool classof(const IntegerRelation *cst) { return true; }
 
   // Clones this object.

 /// Since IntegerPolyhedron does not make a distinction between kinds of
 /// dimensions, IdKind::SetDim should be used to refer to dimension identifiers.
 class IntegerPolyhedron : public IntegerRelation {
 public:
   /// Constructs a set reserving memory for the specified number
   /// of constraints and identifiers.
   IntegerPolyhedron(unsigned numReservedInequalities,
                     unsigned numReservedEqualities, unsigned numReservedCols,
-                    unsigned numDims, unsigned numSymbols, unsigned numLocals)
+                    const PresburgerSpace &space)
       : IntegerRelation(numReservedInequalities, numReservedEqualities,
-                        numReservedCols, /*numDomain=*/0, /*numRange=*/numDims,
-                        numSymbols, numLocals) {}
+                        numReservedCols, space) {
+    assert(space.getNumDomainIds() == 0 &&
+           "Number of domain id's should be zero in Set kind space.");
+  }
 
-  /// Constructs a relation with the specified number of dimensions and symbols.
-  IntegerPolyhedron(unsigned numDims = 0, unsigned numSymbols = 0,
-                    unsigned numLocals = 0)
+  /// Constructs a relation with the specified number of dimensions and
+  /// symbols.
+  IntegerPolyhedron(const PresburgerSpace &space)
       : IntegerPolyhedron(/*numReservedInequalities=*/0,
                           /*numReservedEqualities=*/0,
-                          /*numReservedCols=*/numDims + numSymbols + numLocals +
-                              1,
-                          numDims, numSymbols, numLocals) {}
+                          /*numReservedCols=*/space.getNumIds() + 1, space) {}
 
   /// Return a system with no constraints, i.e., one which is satisfied by all
   /// points.
-  static IntegerPolyhedron getUniverse(unsigned numDims = 0,
-                                       unsigned numSymbols = 0) {
-    return IntegerPolyhedron(numDims, numSymbols);
+  static IntegerPolyhedron getUniverse(const PresburgerSpace &space) {
+    return IntegerPolyhedron(space);
   }
 
   /// Return the kind of this IntegerRelation.
   Kind getKind() const override { return Kind::IntegerPolyhedron; }
 
   static bool classof(const IntegerRelation *cst) {
     return cst->getKind() == Kind::IntegerPolyhedron;
   }

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

public:
   using IntegerPolyhedron::getNumIds;
   using IntegerPolyhedron::getNumLocalIds;
   using IntegerPolyhedron::getNumSymbolIds;
   using PresburgerSpace::isSpaceCompatible;
   using PresburgerSpace::isSpaceEqual;
 
   MultiAffineFunction(const IntegerPolyhedron &domain, const Matrix &output)
       : IntegerPolyhedron(domain), output(output) {}
-  MultiAffineFunction(const Matrix &output, unsigned numDims,
-                      unsigned numSymbols = 0, unsigned numLocals = 0)
-      : IntegerPolyhedron(numDims, numSymbols, numLocals), output(output) {}
+  MultiAffineFunction(const Matrix &output, const PresburgerSpace &space)
+      : IntegerPolyhedron(space), output(output) {}
 
   ~MultiAffineFunction() override = default;
   Kind getKind() const override { return Kind::MultiAffineFunction; }
   bool classof(const IntegerRelation *rel) const {
     return rel->getKind() == Kind::MultiAffineFunction;
   }
 
   unsigned getNumInputs() const { return getNumDimAndSymbolIds(); }

 ///
 /// As in PresburgerSets, the input ids are partitioned into dimension ids and
 /// symbolic ids.
 ///
 /// Support is provided to compare equality of two such functions as well as
 /// finding the value of the function at a point.
 class PWMAFunction : public PresburgerSpace {
 public:
-  PWMAFunction(unsigned numDims, unsigned numSymbols, unsigned numOutputs)
-      : PresburgerSpace(/*numDomain=*/0, /*numRange=*/numDims, numSymbols,
-                        /*numLocals=*/0),
-        numOutputs(numOutputs) {
+  PWMAFunction(const PresburgerSpace &space, unsigned numOutputs)
+      : PresburgerSpace(space), numOutputs(numOutputs) {
+    assert(getNumDomainIds() == 0 && "Set type space should zero domain ids.");
+    assert(getNumLocalIds() == 0 && "PWMAFunction cannot have local ids.");
     assert(numOutputs >= 1 && "The function must output something!");
   }
 
   void addPiece(const MultiAffineFunction &piece);
   void addPiece(const IntegerPolyhedron &domain, const Matrix &output);
 
   const MultiAffineFunction &getPiece(unsigned i) const { return pieces[i]; }
   unsigned getNumPieces() const { return pieces.size(); }

mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h

 /// the empty set.
 ///
 /// Note that there are no invariants guaranteed on the list of relations
 /// other than that they are all in the same PresburgerSpace. For example, the
 /// relations may overlap with each other.
 class PresburgerRelation : public PresburgerSpace {
 public:
   /// Return a universe set of the specified type that contains all points.
-  static PresburgerRelation getUniverse(unsigned numDomain, unsigned numRange,
-                                        unsigned numSymbols);
+  static PresburgerRelation getUniverse(const PresburgerSpace &space);
 
   /// Return an empty set of the specified type that contains no points.
-  static PresburgerRelation getEmpty(unsigned numDomain = 0,
-                                     unsigned numRange = 0,
-                                     unsigned numSymbols = 0);
+  static PresburgerRelation getEmpty(const PresburgerSpace &space);
 
   explicit PresburgerRelation(const IntegerRelation &disjunct);
 
   /// Return the number of Disjuncts in the union.
   unsigned getNumDisjuncts() const;
 
   /// Return a reference to the list of IntegerRelations.
   ArrayRef<IntegerRelation> getAllDisjuncts() const;
public:
 
   /// Print the set's internal state.
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
   /// Construct an empty PresburgerRelation with the specified number of
   /// dimension and symbols.
-  PresburgerRelation(unsigned numDomain = 0, unsigned numRange = 0,
-                     unsigned numSymbols = 0)
-      : PresburgerSpace(numDomain, numRange, numSymbols, /*numLocals=*/0) {}
+  PresburgerRelation(const PresburgerSpace &space) : PresburgerSpace(space) {
+    assert(space.getNumLocalIds() == 0 &&
+           "PresburgerRelation cannot have local ids.");
+  }
 
   /// The list of disjuncts that this set is the union of.
   SmallVector<IntegerRelation, 2> integerRelations;
 
   friend class SetCoalescer;
 };
 
 class PresburgerSet : public PresburgerRelation {
 public:
   /// Return a universe set of the specified type that contains all points.
-  static PresburgerSet getUniverse(unsigned numDims = 0,
-                                   unsigned numSymbols = 0);
+  static PresburgerSet getUniverse(const PresburgerSpace &space);
 
   /// Return an empty set of the specified type that contains no points.
-  static PresburgerSet getEmpty(unsigned numDims = 0, unsigned numSymbols = 0);
+  static PresburgerSet getEmpty(const PresburgerSpace &space);
 
   /// Create a set from a relation.
   explicit PresburgerSet(const IntegerPolyhedron &disjunct);
   explicit PresburgerSet(const PresburgerRelation &set);
 
   /// These operations are the same as the ones in PresburgeRelation, they just
   /// forward the arguement and return the result as a set instead of a
   /// relation.
   PresburgerSet unionSet(const PresburgerRelation &set) const;
   PresburgerSet intersect(const PresburgerRelation &set) const;
   PresburgerSet complement() const;
   PresburgerSet subtract(const PresburgerRelation &set) const;
   PresburgerSet coalesce() const;
 
 protected:
   /// Construct an empty PresburgerRelation with the specified number of
   /// dimension and symbols.
-  PresburgerSet(unsigned numDims = 0, unsigned numSymbols = 0)
-      : PresburgerRelation(/*numDomain=*/0, numDims, numSymbols) {}
+  PresburgerSet(const PresburgerSpace &space) : PresburgerRelation(space) {
+    assert(space.getNumDomainIds() == 0 && "Set type cannot have domain ids.");
+    assert(space.getNumLocalIds() == 0 &&
+           "PresburgerRelation cannot have local ids.");
+  }
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERRELATION_H

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

 /// be implemented as a space with zero domain. IdKind::SetDim should be used
 /// to refer to dimensions in such spaces.
 ///
 /// 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:
-  PresburgerSpace(unsigned numDomain = 0, unsigned numRange = 0,
-                  unsigned numSymbols = 0, unsigned numLocals = 0)
-      : numDomain(numDomain), numRange(numRange), numSymbols(numSymbols),
-        numLocals(numLocals) {}
+  static PresburgerSpace getRelationSpace(unsigned numDomain = 0,
+                                          unsigned numRange = 0,
+                                          unsigned numSymbols = 0,
+                                          unsigned numLocals = 0) {
+    return PresburgerSpace(numDomain, numRange, numSymbols, numLocals);
+  }
+
+  static PresburgerSpace getSetSpace(unsigned numDims = 0,
+                                     unsigned numSymbols = 0,
+                                     unsigned numLocals = 0) {
+    return PresburgerSpace(/*numDomain=*/0, /*numRange=*/numDims, numSymbols,
+                           numLocals);
+  }
+
+  PresburgerSpace getSpace() const { return *this; }
+  PresburgerSpace getCompatibleSpace() const {
+    return PresburgerSpace(numDomain, numRange, numSymbols);
+  }
 
   virtual ~PresburgerSpace() = default;
 
   unsigned getNumDomainIds() const { return numDomain; }
   unsigned getNumRangeIds() const { return numRange; }
   unsigned getNumSetDimIds() const { return numRange; }
   unsigned getNumSymbolIds() const { return numSymbols; }
   unsigned getNumLocalIds() const { return numLocals; }
public:
   /// Return the index at Which the specified kind of id ends.
   unsigned getIdKindEnd(IdKind kind) const;
 
   /// Get the number of elements of the specified kind in the range
   /// [idStart, idLimit).
   unsigned getIdKindOverlap(IdKind kind, unsigned idStart,
                             unsigned idLimit) const;
 
+  /// Return the IdKind of the id at the specified position.
+  IdKind getIdKindAt(unsigned pos) const;
+
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. Return the absolute
   /// column position (i.e., not relative to the kind of identifier) of the
   /// first added identifier.
   virtual 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.
public:
   /// 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) {}
+
 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

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

 class FlatAffineConstraints : public presburger::IntegerPolyhedron {
 public:
   /// Constructs a constraint system reserving memory for the specified number
   /// of constraints and identifiers.
   FlatAffineConstraints(unsigned numReservedInequalities,
                         unsigned numReservedEqualities,
                         unsigned numReservedCols, unsigned numDims,
                         unsigned numSymbols, unsigned numLocals)
-      : IntegerPolyhedron(numReservedInequalities, numReservedEqualities,
-                          numReservedCols, numDims, numSymbols, numLocals) {}
+      : IntegerPolyhedron(
+            numReservedInequalities, numReservedEqualities, numReservedCols,
+            PresburgerSpace::getSetSpace(numDims, numSymbols, numLocals)) {}
 
   /// Constructs a constraint system with the specified number of
   /// dimensions and symbols.
   FlatAffineConstraints(unsigned numDims = 0, unsigned numSymbols = 0,
                         unsigned numLocals = 0)
-      : IntegerPolyhedron(/*numReservedInequalities=*/0,
+      : FlatAffineConstraints(/*numReservedInequalities=*/0,
                               /*numReservedEqualities=*/0,
-                          /*numReservedCols=*/numDims + numSymbols + numLocals +
-                              1,
+                              /*numReservedCols=*/numDims + numSymbols +
+                                  numLocals + 1,
                               numDims, numSymbols, numLocals) {}
 
   explicit FlatAffineConstraints(const IntegerPolyhedron &poly)
       : IntegerPolyhedron(poly) {}
 
   /// Return a system with no constraints, i.e., one which is satisfied by all
   /// points.
   static FlatAffineConstraints getUniverse(unsigned numDims = 0,
                                            unsigned numSymbols = 0) {

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

if (atIneq(r, pos) == 0) {
       // Lower bound.
       lbIndices.push_back(r);
     } else {
       // Upper bound.
       ubIndices.push_back(r);
     }
   }
 
-  // Set the number of dimensions, symbols, locals in the resulting system.
-  unsigned newNumDomain =
-      getNumDomainIds() - getIdKindOverlap(IdKind::Domain, pos, pos + 1);
-  unsigned newNumRange =
-      getNumRangeIds() - getIdKindOverlap(IdKind::Range, pos, pos + 1);
-  unsigned newNumSymbols =
-      getNumSymbolIds() - getIdKindOverlap(IdKind::Symbol, pos, pos + 1);
-  unsigned newNumLocals =
-      getNumLocalIds() - getIdKindOverlap(IdKind::Local, pos, pos + 1);
+  PresburgerSpace newSpace = getSpace();
+  IdKind idKindRemove = newSpace.getIdKindAt(pos);
+  unsigned relativePos = pos - newSpace.getIdKindOffset(idKindRemove);
+  newSpace.removeIdRange(idKindRemove, relativePos, relativePos + 1);
 
   /// Create the new system which has one identifier less.
   IntegerRelation newRel(lbIndices.size() * ubIndices.size() + nbIndices.size(),
-                         getNumEqualities(), getNumCols() - 1, newNumDomain,
-                         newNumRange, newNumSymbols, newNumLocals);
+                         getNumEqualities(), getNumCols() - 1, newSpace);
 
   // This will be used to check if the elimination was integer exact.
   unsigned lcmProducts = 1;
 
   // Let x be the variable we are eliminating.
   // For each lower bound, lb <= c_l*x, and each upper bound c_u*x <= ub, (note
   // that c_l, c_u >= 1) we have:
   // lb*lcm(c_l, c_u)/c_l <= lcm(c_l, c_u)*x <= ub*lcm(c_l, c_u)/c_u
static BoundCmpResult compareBounds(ArrayRef<int64_t> a, ArrayRef<int64_t> b) {
 
   return a.back() < b.back() ? Less : Greater;
 }
 } // namespace
 
 // Returns constraints that are common to both A & B.
 static void getCommonConstraints(const IntegerRelation &a,
                                  const IntegerRelation &b, IntegerRelation &c) {
-  c = IntegerRelation(a.getNumDomainIds(), a.getNumRangeIds(),
-                      a.getNumSymbolIds(), a.getNumLocalIds());
+  c = IntegerRelation(a.getSpace());
   // a naive O(n^2) check should be enough here given the input sizes.
   for (unsigned r = 0, e = a.getNumInequalities(); r < e; ++r) {
     for (unsigned s = 0, f = b.getNumInequalities(); s < f; ++s) {
       if (a.getInequality(r) == b.getInequality(s)) {
         c.addInequality(a.getInequality(r));
         break;
       }
     }
   }
   for (unsigned r = 0, e = a.getNumEqualities(); r < e; ++r) {
     for (unsigned s = 0, f = b.getNumEqualities(); s < f; ++s) {
       if (a.getEquality(r) == b.getEquality(s)) {
         c.addEquality(a.getEquality(r));
         break;
       }
     }
   }
 }
 
 // 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(isSpaceEqual(otherCst) && "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;
+  IntegerRelation commonCst(PresburgerSpace::getRelationSpace());
   getCommonConstraints(*this, otherCst, commonCst);
 
   std::vector<SmallVector<int64_t, 8>> boundingLbs;
   std::vector<SmallVector<int64_t, 8>> boundingUbs;
   boundingLbs.reserve(2 * getNumDimIds());
   boundingUbs.reserve(2 * getNumDimIds());
 
   // To hold lower and upper bounds for each dimension.

mlir/lib/Analysis/Presburger/LinearTransform.cpp

for (unsigned row = 0; row < m.getNumRows(); ++row) {
 
     ++echelonCol;
   }
 
   return {echelonCol, LinearTransform(std::move(resultMatrix))};
 }
 
 IntegerRelation LinearTransform::applyTo(const IntegerRelation &rel) const {
-  IntegerRelation result(rel.getNumDomainIds(), rel.getNumRangeIds(),
-                         rel.getNumSymbolIds(), rel.getNumLocalIds());
+  IntegerRelation result(rel.getSpace());
 
   for (unsigned i = 0, e = rel.getNumEqualities(); i < e; ++i) {
     ArrayRef<int64_t> eq = rel.getEquality(i);
 
     int64_t c = eq.back();
 
     SmallVector<int64_t, 8> newEq = preMultiplyWithRow(eq.drop_back());
     newEq.push_back(c);

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

static SmallVector<int64_t, 8> subtract(ArrayRef<int64_t> vecA,
   SmallVector<int64_t, 8> result;
   result.reserve(vecA.size());
   for (unsigned i = 0, e = vecA.size(); i < e; ++i)
     result.push_back(vecA[i] - vecB[i]);
   return result;
 }
 
 PresburgerSet PWMAFunction::getDomain() const {
-  PresburgerSet domain =
-      PresburgerSet::getEmpty(getNumDimIds(), getNumSymbolIds());
+  PresburgerSet domain = PresburgerSet::getEmpty(getSpace());
   for (const MultiAffineFunction &piece : pieces)
     domain.unionInPlace(piece.getDomain());
   return domain;
 }
 
 Optional<SmallVector<int64_t, 8>>
 MultiAffineFunction::valueAt(ArrayRef<int64_t> point) const {
   assert(point.size() == getNumDimAndSymbolIds() &&

mlir/lib/Analysis/Presburger/PresburgerRelation.cpp

 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 using namespace presburger;
 
 PresburgerRelation::PresburgerRelation(const IntegerRelation &disjunct)
-    : PresburgerSpace(disjunct) {
+    : PresburgerSpace(disjunct.getCompatibleSpace()) {
   unionInPlace(disjunct);
 }
 
 unsigned PresburgerRelation::getNumDisjuncts() const {
   return integerRelations.size();
 }
 
 ArrayRef<IntegerRelation> PresburgerRelation::getAllDisjuncts() const {

 
 /// 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(integerRelations, [&](const IntegerRelation &disjunct) {
     return (disjunct.containsPoint(point));
   });
 }
 
-PresburgerRelation PresburgerRelation::getUniverse(unsigned numDomain,
-                                                   unsigned numRange,
-                                                   unsigned numSymbols) {
-  PresburgerRelation result(numDomain, numRange, numSymbols);
-  result.unionInPlace(
-      IntegerRelation::getUniverse(numDomain, numRange, numSymbols));
+PresburgerRelation
+PresburgerRelation::getUniverse(const PresburgerSpace &space) {
+  PresburgerRelation result(space);
+  result.unionInPlace(IntegerRelation::getUniverse(space));
   return result;
 }
 
-PresburgerRelation PresburgerRelation::getEmpty(unsigned numDomain,
-                                                unsigned numRange,
-                                                unsigned numSymbols) {
-  return PresburgerRelation(numDomain, numRange, numSymbols);
+PresburgerRelation PresburgerRelation::getEmpty(const PresburgerSpace &space) {
+  return PresburgerRelation(space);
 }
 
 // Return the intersection of this set with the given set.
 //
 // 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(isSpaceCompatible(set) && "Spaces should match");
 
-  PresburgerRelation result(getNumDomainIds(), getNumRangeIds(),
-                            getNumSymbolIds());
+  PresburgerRelation result(getSpace());
   for (const IntegerRelation &csA : integerRelations) {
     for (const IntegerRelation &csB : set.integerRelations) {
       IntegerRelation intersection = csA.intersect(csB);
       if (!intersection.isEmpty())
         result.unionInPlace(intersection);
     }
   }
   return result;

 ///
 /// The disjunct here is modified in subtractRecursively, so it cannot be a
 /// const reference even though it is restored to its original state before
 /// returning from that function.
 static PresburgerRelation getSetDifference(IntegerRelation disjunct,
                                            const PresburgerRelation &set) {
   assert(disjunct.isSpaceCompatible(set) && "Spaces should match");
   if (disjunct.isEmptyByGCDTest())
-    return PresburgerRelation::getEmpty(disjunct.getNumDomainIds(),
-                                        disjunct.getNumRangeIds(),
-                                        disjunct.getNumSymbolIds());
-
-  PresburgerRelation result = PresburgerRelation::getEmpty(
-      disjunct.getNumDomainIds(), disjunct.getNumRangeIds(),
-      disjunct.getNumSymbolIds());
+    return PresburgerRelation::getEmpty(disjunct.getCompatibleSpace());
+
+  PresburgerRelation result =
+      PresburgerRelation::getEmpty(disjunct.getCompatibleSpace());
   Simplex simplex(disjunct);
   subtractRecursively(disjunct, simplex, set, 0, result);
   return result;
 }
 
 /// Return the complement of this set.
 PresburgerRelation PresburgerRelation::complement() const {
-  return getSetDifference(IntegerRelation::getUniverse(getNumDomainIds(),
-                                                       getNumRangeIds(),
-                                                       getNumSymbolIds()),
-                          *this);
+  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(isSpaceCompatible(set) && "Spaces should match");
-  PresburgerRelation result(getNumDomainIds(), getNumRangeIds(),
-                            getNumSymbolIds());
+  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 : integerRelations)
     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
for (unsigned i = 0; i < disjuncts.size();) {
     if (simp.isEmpty()) {
       disjuncts[i] = disjuncts[disjuncts.size() - 1];
       disjuncts.pop_back();
       continue;
     }
     ++i;
     simplices.push_back(simp);
   }
   numDomainIds = s.getNumDomainIds();
   numRangeIds = s.getNumRangeIds();
   numSymbolIds = s.getNumSymbolIds();
 }
 
 /// Simplifies the representation of a PresburgerSet.
 PresburgerRelation SetCoalescer::coalesce() {
   // For all tuples of IntegerRelations, check whether they can be
   // coalesced. When coalescing is successful, the contained IntegerRelation
   // is swapped with the last element of `disjuncts` and subsequently erased
   // and similarly for simplices.
for (unsigned i = 0; i < disjuncts.size();) {
     // Only if the inner loop was not broken, i is incremented. This is
     // required as otherwise, if a coalescing occurs, the IntegerRelation
     // now at position i is not compared.
     if (!broken)
       ++i;
   }
 
   PresburgerRelation newSet =
-      PresburgerRelation::getEmpty(numDomainIds, numRangeIds, numSymbolIds);
+      PresburgerRelation::getEmpty(PresburgerSpace::getRelationSpace(
+          numDomainIds, numRangeIds, numSymbolIds));
   for (unsigned i = 0, e = disjuncts.size(); i < e; ++i)
     newSet.unionInPlace(disjuncts[i]);
 
   return newSet;
 }
 
 /// Given a Simplex `simp` and one of its inequalities `ineq`, check
 /// that all inequalities of `cuttingIneqsB` are redundant for the facet of
LogicalResult SetCoalescer::coalescePairCutCase(unsigned i, unsigned j) {
   /// All inequalities of `b` need to be redundant. We already know that the
   /// redundant ones are, so only the cutting ones remain to be checked.
   Simplex &simp = simplices[i];
   IntegerRelation &disjunct = disjuncts[i];
   if (llvm::any_of(cuttingIneqsA, [this, &simp](ArrayRef<int64_t> curr) {
         return !isFacetContained(curr, simp);
       }))
     return failure();
-  IntegerRelation newSet(disjunct.getNumDomainIds(), disjunct.getNumRangeIds(),
-                         disjunct.getNumSymbolIds(), disjunct.getNumLocalIds());
+  IntegerRelation newSet(disjunct.getSpace());
 
   for (ArrayRef<int64_t> curr : redundantIneqsA)
     newSet.addInequality(curr);
 
   for (ArrayRef<int64_t> curr : redundantIneqsB)
     newSet.addInequality(curr);
 
   addCoalescedDisjunct(i, j, newSet);
void PresburgerRelation::print(raw_ostream &os) const {
   for (const IntegerRelation &disjunct : integerRelations) {
     disjunct.print(os);
     os << '\n';
   }
 }
 
 void PresburgerRelation::dump() const { print(llvm::errs()); }
 
-PresburgerSet PresburgerSet::getUniverse(unsigned numDims,
-                                         unsigned numSymbols) {
-  PresburgerSet result(numDims, numSymbols);
-  result.unionInPlace(IntegerPolyhedron::getUniverse(numDims, numSymbols));
+PresburgerSet PresburgerSet::getUniverse(const PresburgerSpace &space) {
+  PresburgerSet result(space);
+  result.unionInPlace(IntegerPolyhedron::getUniverse(space));
   return result;
 }
 
-PresburgerSet PresburgerSet::getEmpty(unsigned numDims, unsigned numSymbols) {
-  return PresburgerSet(numDims, numSymbols);
+PresburgerSet PresburgerSet::getEmpty(const PresburgerSpace &space) {
+  return PresburgerSet(space);
 }
 
 PresburgerSet::PresburgerSet(const IntegerPolyhedron &disjunct)
     : PresburgerRelation(disjunct) {}
 
 PresburgerSet::PresburgerSet(const PresburgerRelation &set)
     : PresburgerRelation(set) {}
 

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

unsigned PresburgerSpace::getIdKindOverlap(IdKind kind, unsigned idStart,
   unsigned overlapStart = std::max(idStart, idRangeStart);
   unsigned overlapEnd = std::min(idLimit, idRangeEnd);
 
   if (overlapStart > overlapEnd)
     return 0;
   return overlapEnd - overlapStart;
 }
 
+IdKind PresburgerSpace::getIdKindAt(unsigned pos) const {
+  assert(pos < getNumIds() && "`pos` should represent a valid id position");
+  if (pos < getIdKindEnd(IdKind::Domain))
+    return IdKind::Domain;
+  if (pos < getIdKindEnd(IdKind::Range))
+    return IdKind::Range;
+  if (pos < getIdKindEnd(IdKind::Symbol))
+    return IdKind::Symbol;
+  if (pos < getIdKindEnd(IdKind::Local))
+    return IdKind::Local;
+  llvm_unreachable("`pos` should represent a valid id position");
+}
+
 unsigned PresburgerSpace::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
   unsigned absolutePos = getIdKindOffset(kind) + pos;
 
   if (kind == IdKind::Domain)
     numDomain += num;
   else if (kind == IdKind::Range)

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

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

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

 enum class TestFunction { Sample, Empty };
 
 /// Construct a IntegerPolyhedron from a set of inequality and
 /// equality constraints.
 static IntegerPolyhedron
 makeSetFromConstraints(unsigned ids, ArrayRef<SmallVector<int64_t, 4>> ineqs,
                        ArrayRef<SmallVector<int64_t, 4>> eqs,
                        unsigned syms = 0) {
-  IntegerPolyhedron set(ineqs.size(), eqs.size(), ids + 1, ids - syms, syms,
-                        /*numLocals=*/0);
+  IntegerPolyhedron set(
+      ineqs.size(), eqs.size(), ids + 1,
+      PresburgerSpace::getSetSpace(ids - syms, syms, /*numLocals=*/0));
   for (const auto &eq : eqs)
     set.addEquality(eq);
   for (const auto &ineq : ineqs)
     set.addInequality(ineq);
   return set;
 }
 
 static void dump(ArrayRef<int64_t> vec) {
TEST(IntegerPolyhedronTest, clearConstraints) {
   set.clearConstraints();
 
   set.addInequality({1, 0});
   EXPECT_EQ(set.atIneq(0, 0), 1);
   EXPECT_EQ(set.atIneq(0, 1), 0);
 }
 
 TEST(IntegerPolyhedronTest, removeIdRange) {
-  IntegerPolyhedron set(3, 2, 1);
+  IntegerPolyhedron set(PresburgerSpace::getSetSpace(3, 2, 1));
 
   set.addInequality({10, 11, 12, 20, 21, 30, 40});
   set.removeId(IdKind::Symbol, 1);
   EXPECT_THAT(set.getInequality(0),
               testing::ElementsAre(10, 11, 12, 20, 30, 40));
 
   set.removeIdRange(IdKind::SetDim, 0, 2);
   EXPECT_THAT(set.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
TEST(IntegerPolyhedronTest, removeRedundantConstraintsTest) {
   SmallVector<int64_t, 8> redundantConstraint = {0, 1, 0};
   for (unsigned i = 0; i < 3; ++i) {
     // Ensure that the removed constraint was the redundant constraint [3].
     EXPECT_NE(poly5.getInequality(i), ArrayRef<int64_t>(redundantConstraint));
   }
 }
 
 TEST(IntegerPolyhedronTest, addConstantUpperBound) {
-  IntegerPolyhedron poly(2);
+  IntegerPolyhedron poly(PresburgerSpace::getSetSpace(2));
   poly.addBound(IntegerPolyhedron::UB, 0, 1);
   EXPECT_EQ(poly.atIneq(0, 0), -1);
   EXPECT_EQ(poly.atIneq(0, 1), 0);
   EXPECT_EQ(poly.atIneq(0, 2), 1);
 
   poly.addBound(IntegerPolyhedron::UB, {1, 2, 3}, 1);
   EXPECT_EQ(poly.atIneq(1, 0), -1);
   EXPECT_EQ(poly.atIneq(1, 1), -2);
   EXPECT_EQ(poly.atIneq(1, 2), -2);
 }
 
 TEST(IntegerPolyhedronTest, addConstantLowerBound) {
-  IntegerPolyhedron poly(2);
+  IntegerPolyhedron poly(PresburgerSpace::getSetSpace(2));
   poly.addBound(IntegerPolyhedron::LB, 0, 1);
   EXPECT_EQ(poly.atIneq(0, 0), 1);
   EXPECT_EQ(poly.atIneq(0, 1), 0);
   EXPECT_EQ(poly.atIneq(0, 2), -1);
 
   poly.addBound(IntegerPolyhedron::LB, {1, 2, 3}, 1);
   EXPECT_EQ(poly.atIneq(1, 0), 1);
   EXPECT_EQ(poly.atIneq(1, 1), 2);
static void checkDivisionRepresentation(
   for (unsigned i = 0, e = dividends.size(); i < e; ++i) {
     if (denominators[i] != 0) {
       EXPECT_TRUE(expectedDividends[i] == dividends[i]);
     }
   }
 }
 
 TEST(IntegerPolyhedronTest, computeLocalReprSimple) {
-  IntegerPolyhedron poly(1);
+  IntegerPolyhedron poly(PresburgerSpace::getSetSpace(1));
 
   poly.addLocalFloorDiv({1, 4}, 10);
   poly.addLocalFloorDiv({1, 0, 100}, 10);
 
   std::vector<SmallVector<int64_t, 8>> divisions = {{1, 0, 0, 4},
                                                     {1, 0, 0, 100}};
   SmallVector<unsigned, 8> denoms = {10, 10};
 
   // Check if floordivs can be computed when no other inequalities exist
   // and floor divs do not depend on each other.
   checkDivisionRepresentation(poly, divisions, denoms);
 }
 
 TEST(IntegerPolyhedronTest, computeLocalReprConstantFloorDiv) {
-  IntegerPolyhedron poly(4);
+  IntegerPolyhedron poly(PresburgerSpace::getSetSpace(4));
 
   poly.addInequality({1, 0, 3, 1, 2});
   poly.addInequality({1, 2, -8, 1, 10});
   poly.addEquality({1, 2, -4, 1, 10});
 
   poly.addLocalFloorDiv({0, 0, 0, 0, 100}, 30);
   poly.addLocalFloorDiv({0, 0, 0, 0, 0, 206}, 101);
 
   std::vector<SmallVector<int64_t, 8>> divisions = {{0, 0, 0, 0, 0, 0, 3},
                                                     {0, 0, 0, 0, 0, 0, 2}};
   SmallVector<unsigned, 8> denoms = {1, 1};
 
   // Check if floordivs with constant numerator can be computed.
   checkDivisionRepresentation(poly, divisions, denoms);
 }
 
 TEST(IntegerPolyhedronTest, computeLocalReprRecursive) {
-  IntegerPolyhedron poly(4);
+  IntegerPolyhedron poly(PresburgerSpace::getSetSpace(4));
   poly.addInequality({1, 0, 3, 1, 2});
   poly.addInequality({1, 2, -8, 1, 10});
   poly.addEquality({1, 2, -4, 1, 10});
 
   poly.addLocalFloorDiv({0, -2, 7, 2, 10}, 3);
   poly.addLocalFloorDiv({3, 0, 9, 2, 2, 10}, 5);
   poly.addLocalFloorDiv({0, 1, -123, 2, 0, -4, 10}, 3);
 
TEST(IntegerPolyhedronTest, computeLocalReprNegConstNormalize) {
   //   = floor((-1 - x)/2).
   std::vector<SmallVector<int64_t, 8>> divisions = {{-1, 0, -1}};
   SmallVector<unsigned, 8> denoms = {2};
   checkDivisionRepresentation(poly, divisions, denoms);
 }
 
 TEST(IntegerPolyhedronTest, simplifyLocalsTest) {
   // (x) : (exists y: 2x + y = 1 and y = 2).
-  IntegerPolyhedron poly(1, 0, 1);
+  IntegerPolyhedron poly(PresburgerSpace::getSetSpace(1, 0, 1));
   poly.addEquality({2, 1, -1});
   poly.addEquality({0, 1, -2});
 
   EXPECT_TRUE(poly.isEmpty());
 
   // (x) : (exists y, z, w: 3x + y = 1 and 2y = z and 3y = w and z = w).
-  IntegerPolyhedron poly2(1, 0, 3);
+  IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1, 0, 3));
   poly2.addEquality({3, 1, 0, 0, -1});
   poly2.addEquality({0, 2, -1, 0, 0});
   poly2.addEquality({0, 3, 0, -1, 0});
   poly2.addEquality({0, 0, 1, -1, 0});
 
   EXPECT_TRUE(poly2.isEmpty());
 
   // (x) : (exists y: x >= y + 1 and 2x + y = 0 and y >= -1).
-  IntegerPolyhedron poly3(1, 0, 1);
+  IntegerPolyhedron poly3(PresburgerSpace::getSetSpace(1, 0, 1));
   poly3.addInequality({1, -1, -1});
   poly3.addInequality({0, 1, 1});
   poly3.addEquality({2, 1, 0});
 
   EXPECT_TRUE(poly3.isEmpty());
 }
 
 TEST(IntegerPolyhedronTest, mergeDivisionsSimple) {
   {
     // (x) : (exists z, y  = [x / 2] : x = 3y and x + z + 1 >= 0).
-    IntegerPolyhedron poly1(1, 0, 1);
+    IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1, 0, 1));
     poly1.addLocalFloorDiv({1, 0, 0}, 2); // y = [x / 2].
     poly1.addEquality({1, 0, -3, 0});     // x = 3y.
     poly1.addInequality({1, 1, 0, 1});    // x + z + 1 >= 0.
 
     // (x) : (exists y = [x / 2], z : x = 5y).
-    IntegerPolyhedron poly2(1);
+    IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
     poly2.addLocalFloorDiv({1, 0}, 2); // y = [x / 2].
     poly2.addEquality({1, -5, 0});     // x = 5y.
     poly2.appendId(IdKind::Local);     // Add local id z.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 1 division should be matched + 2 unmatched local ids.
     EXPECT_EQ(poly1.getNumLocalIds(), 3u);
     EXPECT_EQ(poly2.getNumLocalIds(), 3u);
   }
 
   {
     // (x) : (exists z = [x / 5], y = [x / 2] : x = 3y).
-    IntegerPolyhedron poly1(1);
+    IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1));
     poly1.addLocalFloorDiv({1, 0}, 5);    // z = [x / 5].
     poly1.addLocalFloorDiv({1, 0, 0}, 2); // y = [x / 2].
     poly1.addEquality({1, 0, -3, 0});     // x = 3y.
 
     // (x) : (exists y = [x / 2], z = [x / 5]: x = 5z).
-    IntegerPolyhedron poly2(1);
+    IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
     poly2.addLocalFloorDiv({1, 0}, 2);    // y = [x / 2].
     poly2.addLocalFloorDiv({1, 0, 0}, 5); // z = [x / 5].
     poly2.addEquality({1, 0, -5, 0});     // x = 5z.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(poly1.getNumLocalIds(), 2u);
     EXPECT_EQ(poly2.getNumLocalIds(), 2u);
   }
 
   {
     // Division Normalization test.
     // (x) : (exists z, y  = [x / 2] : x = 3y and x + z + 1 >= 0).
-    IntegerPolyhedron poly1(1, 0, 1);
+    IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1, 0, 1));
     // This division would be normalized.
     poly1.addLocalFloorDiv({3, 0, 0}, 6); // y = [3x / 6] -> [x/2].
     poly1.addEquality({1, 0, -3, 0});     // x = 3z.
     poly1.addInequality({1, 1, 0, 1});    // x + y + 1 >= 0.
 
     // (x) : (exists y = [x / 2], z : x = 5y).
-    IntegerPolyhedron poly2(1);
+    IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
     poly2.addLocalFloorDiv({1, 0}, 2); // y = [x / 2].
     poly2.addEquality({1, -5, 0});     // x = 5y.
     poly2.appendId(IdKind::Local);     // Add local id z.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // One division should be matched + 2 unmatched local ids.
     EXPECT_EQ(poly1.getNumLocalIds(), 3u);
     EXPECT_EQ(poly2.getNumLocalIds(), 3u);
   }
 }
 
 TEST(IntegerPolyhedronTest, mergeDivisionsNestedDivsions) {
   {
     // (x) : (exists y = [x / 2], z = [x + y / 3]: y + z >= x).
-    IntegerPolyhedron poly1(1);
+    IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1));
     poly1.addLocalFloorDiv({1, 0}, 2);    // y = [x / 2].
     poly1.addLocalFloorDiv({1, 1, 0}, 3); // z = [x + y / 3].
     poly1.addInequality({-1, 1, 1, 0});   // y + z >= x.
 
     // (x) : (exists y = [x / 2], z = [x + y / 3]: y + z <= x).
-    IntegerPolyhedron poly2(1);
+    IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
     poly2.addLocalFloorDiv({1, 0}, 2);    // y = [x / 2].
     poly2.addLocalFloorDiv({1, 1, 0}, 3); // z = [x + y / 3].
     poly2.addInequality({1, -1, -1, 0});  // y + z <= x.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(poly1.getNumLocalIds(), 2u);
     EXPECT_EQ(poly2.getNumLocalIds(), 2u);
   }
 
   {
     // (x) : (exists y = [x / 2], z = [x + y / 3], w = [z + 1 / 5]: y + z >= x).
-    IntegerPolyhedron poly1(1);
+    IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1));
     poly1.addLocalFloorDiv({1, 0}, 2);       // y = [x / 2].
     poly1.addLocalFloorDiv({1, 1, 0}, 3);    // z = [x + y / 3].
     poly1.addLocalFloorDiv({0, 0, 1, 1}, 5); // w = [z + 1 / 5].
     poly1.addInequality({-1, 1, 1, 0, 0});   // y + z >= x.
 
     // (x) : (exists y = [x / 2], z = [x + y / 3], w = [z + 1 / 5]: y + z <= x).
-    IntegerPolyhedron poly2(1);
+    IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
     poly2.addLocalFloorDiv({1, 0}, 2);       // y = [x / 2].
     poly2.addLocalFloorDiv({1, 1, 0}, 3);    // z = [x + y / 3].
     poly2.addLocalFloorDiv({0, 0, 1, 1}, 5); // w = [z + 1 / 5].
     poly2.addInequality({1, -1, -1, 0, 0});  // y + z <= x.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 3 divisions should be matched.
     EXPECT_EQ(poly1.getNumLocalIds(), 3u);
     EXPECT_EQ(poly2.getNumLocalIds(), 3u);
   }
   {
     // (x) : (exists y = [x / 2], z = [x + y / 3]: y + z >= x).
-    IntegerPolyhedron poly1(1);
+    IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1));
     poly1.addLocalFloorDiv({2, 0}, 4);    // y = [2x / 4] -> [x / 2].
     poly1.addLocalFloorDiv({1, 1, 0}, 3); // z = [x + y / 3].
     poly1.addInequality({-1, 1, 1, 0});   // y + z >= x.
 
     // (x) : (exists y = [x / 2], z = [x + y / 3]: y + z <= x).
-    IntegerPolyhedron poly2(1);
+    IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
     poly2.addLocalFloorDiv({1, 0}, 2); // y = [x / 2].
     // This division would be normalized.
     poly2.addLocalFloorDiv({3, 3, 0}, 9); // z = [3x + 3y / 9] -> [x + y / 3].
     poly2.addInequality({1, -1, -1, 0});  // y + z <= x.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(poly1.getNumLocalIds(), 2u);
     EXPECT_EQ(poly2.getNumLocalIds(), 2u);
   }
 }
 
 TEST(IntegerPolyhedronTest, mergeDivisionsConstants) {
   {
     // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z >= x).
-    IntegerPolyhedron poly1(1);
+    IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1));
     poly1.addLocalFloorDiv({1, 1}, 2);    // y = [x + 1 / 2].
     poly1.addLocalFloorDiv({1, 0, 2}, 3); // z = [x + 2 / 3].
     poly1.addInequality({-1, 1, 1, 0});   // y + z >= x.
 
     // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z <= x).
-    IntegerPolyhedron poly2(1);
+    IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
     poly2.addLocalFloorDiv({1, 1}, 2);    // y = [x + 1 / 2].
     poly2.addLocalFloorDiv({1, 0, 2}, 3); // z = [x + 2 / 3].
     poly2.addInequality({1, -1, -1, 0});  // y + z <= x.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(poly1.getNumLocalIds(), 2u);
     EXPECT_EQ(poly2.getNumLocalIds(), 2u);
   }
   {
     // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z >= x).
-    IntegerPolyhedron poly1(1);
+    IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1));
     poly1.addLocalFloorDiv({1, 1}, 2); // y = [x + 1 / 2].
     // Normalization test.
     poly1.addLocalFloorDiv({3, 0, 6}, 9); // z = [3x + 6 / 9] -> [x + 2 / 3].
     poly1.addInequality({-1, 1, 1, 0});   // y + z >= x.
 
     // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z <= x).
-    IntegerPolyhedron poly2(1);
+    IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
     // Normalization test.
     poly2.addLocalFloorDiv({2, 2}, 4);    // y = [2x + 2 / 4] -> [x + 1 / 2].
     poly2.addLocalFloorDiv({1, 0, 2}, 3); // z = [x + 2 / 3].
     poly2.addInequality({1, -1, -1, 0});  // y + z <= x.
 
     poly1.mergeLocalIds(poly2);
 
     // Local space should be same.
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(poly1.getNumLocalIds(), 2u);
     EXPECT_EQ(poly2.getNumLocalIds(), 2u);
   }
 }
 
 TEST(IntegerPolyhedronTest, negativeDividends) {
   // (x) : (exists y = [-x + 1 / 2], z = [-x - 2 / 3]: y + z >= x).
-  IntegerPolyhedron poly1(1);
+  IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1));
   poly1.addLocalFloorDiv({-1, 1}, 2); // y = [x + 1 / 2].
   // Normalization test with negative dividends
   poly1.addLocalFloorDiv({-3, 0, -6}, 9); // z = [3x + 6 / 9] -> [x + 2 / 3].
   poly1.addInequality({-1, 1, 1, 0});     // y + z >= x.
 
   // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z <= x).
-  IntegerPolyhedron poly2(1);
+  IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
   // Normalization test.
   poly2.addLocalFloorDiv({-2, 2}, 4);     // y = [-2x + 2 / 4] -> [-x + 1 / 2].
   poly2.addLocalFloorDiv({-1, 0, -2}, 3); // z = [-x - 2 / 3].
   poly2.addInequality({1, -1, -1, 0});    // y + z <= x.
 
   poly1.mergeLocalIds(poly2);
 
   // Merging triggers normalization.
TEST(IntegerPolyhedronTest, containsPointNoLocal) {
   IntegerPolyhedron poly3 = parsePoly("(x, y) : (2*x - y >= 0, y - 3*x >= 0)");
   EXPECT_TRUE(poly3.containsPointNoLocal({0, 0}));
   EXPECT_FALSE(poly3.containsPointNoLocal({1, 0}));
 }
 
 TEST(IntegerPolyhedronTest, truncateEqualityRegressionTest) {
   // IntegerRelation::truncate was truncating inequalities to the number of
   // equalities.
-  IntegerRelation set(1);
+  IntegerRelation set(PresburgerSpace::getSetSpace(1));
   IntegerRelation::CountsSnapshot snapshot = set.getCounts();
   set.addEquality({1, 0});
   set.truncate(snapshot);
   EXPECT_EQ(set.getNumEqualities(), 0u);
 }

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

 }
 
 /// Construct a PresburgerSet having `numDims` dimensions and no symbols from
 /// the given list of IntegerPolyhedron. Each Poly in `polys` should also have
 /// `numDims` dimensions and no symbols, although it can have any number of
 /// local ids.
 static PresburgerSet makeSetFromPoly(unsigned numDims,
                                      ArrayRef<IntegerPolyhedron> polys) {
-  PresburgerSet set = PresburgerSet::getEmpty(numDims);
+  PresburgerSet set =
+      PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(numDims));
   for (const IntegerPolyhedron &poly : polys)
     set.unionInPlace(poly);
   return set;
 }
 
 TEST(SetTest, containsPoint) {
   PresburgerSet setA = parsePresburgerSetFromPolyStrings(
       1,

 }
 
 TEST(SetTest, Union) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       1,
       {"(x) : (x - 2 >= 0, -x + 8 >= 0)", "(x) : (x - 10 >= 0, -x + 20 >= 0)"});
 
   // Universe union set.
-  testUnionAtPoints(PresburgerSet::getUniverse(1), set,
-                    {{1}, {2}, {8}, {9}, {10}, {20}, {21}});
+  testUnionAtPoints(PresburgerSet::getUniverse(PresburgerSpace::getSetSpace(1)),
+                    set, {{1}, {2}, {8}, {9}, {10}, {20}, {21}});
 
   // empty set union set.
-  testUnionAtPoints(PresburgerSet::getEmpty(1), set,
-                    {{1}, {2}, {8}, {9}, {10}, {20}, {21}});
+  testUnionAtPoints(PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(1)),
+                    set, {{1}, {2}, {8}, {9}, {10}, {20}, {21}});
 
   // empty set union Universe.
-  testUnionAtPoints(PresburgerSet::getEmpty(1), PresburgerSet::getUniverse(1),
+  testUnionAtPoints(PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(1)),
+                    PresburgerSet::getUniverse(PresburgerSpace::getSetSpace(1)),
                     {{1}, {2}, {0}, {-1}});
 
   // Universe union empty set.
-  testUnionAtPoints(PresburgerSet::getUniverse(1), PresburgerSet::getEmpty(1),
+  testUnionAtPoints(PresburgerSet::getUniverse(PresburgerSpace::getSetSpace(1)),
+                    PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(1)),
                     {{1}, {2}, {0}, {-1}});
 
   // empty set union empty set.
-  testUnionAtPoints(PresburgerSet::getEmpty(1), PresburgerSet::getEmpty(1),
+  testUnionAtPoints(PresburgerSet::getEmpty(PresburgerSpace::getSetSpace((1))),
+                    PresburgerSet::getEmpty(PresburgerSpace::getSetSpace((1))),
                     {{1}, {2}, {0}, {-1}});
 }
 
 TEST(SetTest, Intersect) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       1,
       {"(x) : (x - 2 >= 0, -x + 8 >= 0)", "(x) : (x - 10 >= 0, -x + 20 >= 0)"});
 
   // Universe intersection set.
-  testIntersectAtPoints(PresburgerSet::getUniverse(1), set,
-                        {{1}, {2}, {8}, {9}, {10}, {20}, {21}});
+  testIntersectAtPoints(
+      PresburgerSet::getUniverse(PresburgerSpace::getSetSpace((1))), set,
 
   // empty set intersection set.
-  testIntersectAtPoints(PresburgerSet::getEmpty(1), set,
-                        {{1}, {2}, {8}, {9}, {10}, {20}, {21}});
+  testIntersectAtPoints(
+      PresburgerSet::getEmpty(PresburgerSpace::getSetSpace((1))), set,
 
   // empty set intersection Universe.
-  testIntersectAtPoints(PresburgerSet::getEmpty(1),
-                        PresburgerSet::getUniverse(1), {{1}, {2}, {0}, {-1}});
+  testIntersectAtPoints(
+      PresburgerSet::getEmpty(PresburgerSpace::getSetSpace((1))),
+      PresburgerSet::getUniverse(PresburgerSpace::getSetSpace((1))),
+      {{1}, {2}, {0}, {-1}});
 
   // Universe intersection empty set.
-  testIntersectAtPoints(PresburgerSet::getUniverse(1),
-                        PresburgerSet::getEmpty(1), {{1}, {2}, {0}, {-1}});
+  testIntersectAtPoints(
+      PresburgerSet::getUniverse(PresburgerSpace::getSetSpace((1))),
+      PresburgerSet::getEmpty(PresburgerSpace::getSetSpace((1))),
+      {{1}, {2}, {0}, {-1}});
 
   // Universe intersection Universe.
-  testIntersectAtPoints(PresburgerSet::getUniverse(1),
-                        PresburgerSet::getUniverse(1), {{1}, {2}, {0}, {-1}});
+  testIntersectAtPoints(
+      PresburgerSet::getUniverse(PresburgerSpace::getSetSpace((1))),
+      PresburgerSet::getUniverse(PresburgerSpace::getSetSpace((1))),
+      {{1}, {2}, {0}, {-1}});
 }
 
 TEST(SetTest, Subtract) {
   // The interval [2, 8] minus the interval [10, 20].
   testSubtractAtPoints(
       parsePresburgerSetFromPolyStrings(1, {"(x) : (x - 2 >= 0, -x + 8 >= 0)"}),
       parsePresburgerSetFromPolyStrings(1,
                                         {"(x) : (x - 10 >= 0, -x + 20 >= 0)"}),
testSubtractAtPoints(
        {6},
        {7},
        {8}});
 }
 
 TEST(SetTest, Complement) {
   // Complement of universe.
   testComplementAtPoints(
-      PresburgerSet::getUniverse(1),
+      PresburgerSet::getUniverse(PresburgerSpace::getSetSpace((1))),
       {{-1}, {-2}, {-8}, {1}, {2}, {8}, {9}, {10}, {20}, {21}});
 
   // Complement of empty set.
   testComplementAtPoints(
-      PresburgerSet::getEmpty(1),
+      PresburgerSet::getEmpty(PresburgerSpace::getSetSpace((1))),
       {{-1}, {-2}, {-8}, {1}, {2}, {8}, {9}, {10}, {20}, {21}});
 
   testComplementAtPoints(
       parsePresburgerSetFromPolyStrings(2, {"(x,y) : (x - 2 >= 0, y - 2 >= 0, "
                                             "-x + 10 >= 0, -y + 10 >= 0)"}),
       {{1, 1},
        {2, 1},
        {1, 2},
        {2, 2},
        {2, 3},
        {3, 2},
        {10, 10},
        {10, 11},
        {11, 10},
        {2, 10},
        {2, 11},
        {1, 10}});
 }
 
 TEST(SetTest, isEqual) {
   // set = [2, 8] U [10, 20].
-  PresburgerSet universe = PresburgerSet::getUniverse(1);
-  PresburgerSet emptySet = PresburgerSet::getEmpty(1);
+  PresburgerSet universe =
+      PresburgerSet::getUniverse(PresburgerSpace::getSetSpace((1)));
+  PresburgerSet emptySet =
+      PresburgerSet::getEmpty(PresburgerSpace::getSetSpace((1)));
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       1,
       {"(x) : (x - 2 >= 0, -x + 8 >= 0)", "(x) : (x - 10 >= 0, -x + 20 >= 0)"});
 
   // universe != emptySet.
   EXPECT_FALSE(universe.isEqual(emptySet));
   // emptySet != universe.
   EXPECT_FALSE(emptySet.isEqual(universe));
TEST(SetTest, divisions) {
   PresburgerSet multiples3{parsePoly("(x) : (x - 3 * (x floordiv 3) == 0)")};
 
   // multiples6 = {x : exists q, x = 6q}.
   PresburgerSet multiples6{parsePoly("(x) : (x - 6 * (x floordiv 6) == 0)")};
 
   // evens /\ odds = empty.
   expectEmpty(PresburgerSet(evens).intersect(PresburgerSet(odds)));
   // evens U odds = universe.
-  expectEqual(evens.unionSet(odds), PresburgerSet::getUniverse(1));
+  expectEqual(evens.unionSet(odds),
+              PresburgerSet::getUniverse(PresburgerSpace::getSetSpace((1))));
   expectEqual(evens.complement(), odds);
   expectEqual(odds.complement(), evens);
   // even multiples of 3 = multiples of 6.
   expectEqual(multiples3.intersect(evens), multiples6);
 
   PresburgerSet setA{parsePoly("(x) : (-x >= 0)")};
   PresburgerSet setB{parsePoly("(x) : (x floordiv 2 - 4 >= 0)")};
   EXPECT_TRUE(setA.subtract(setB).isEqual(setA));

mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp

 //===- PresburgerSpaceTest.cpp - Tests for PresburgerSpace ----------------===//
 //
 // 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/Analysis/Presburger/PresburgerSpace.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 using namespace mlir;
 using namespace presburger;
 
 TEST(PresburgerSpaceTest, insertId) {
-  PresburgerSpace space(2, 2, 1);
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 2, 1);
 
   // Try inserting 2 domain ids.
   space.insertId(IdKind::Domain, 0, 2);
   EXPECT_EQ(space.getNumDomainIds(), 4u);
 
   // Try inserting 1 range ids.
   space.insertId(IdKind::Range, 0, 1);
   EXPECT_EQ(space.getNumRangeIds(), 3u);
 }
 
 TEST(PresburgerSpaceTest, insertIdSet) {
-  PresburgerSpace space(0, 2, 1);
+  PresburgerSpace space = PresburgerSpace::getSetSpace(2, 1);
 
   // Try inserting 2 dimension ids. The space should have 4 range ids since
   // spaces which do not distinguish between domain, range are implemented like
   // this.
   space.insertId(IdKind::SetDim, 0, 2);
   EXPECT_EQ(space.getNumRangeIds(), 4u);
 }
 
 TEST(PresburgerSpaceTest, removeIdRange) {
-  PresburgerSpace space(2, 1, 3);
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3);
 
   // Remove 1 domain identifier.
   space.removeIdRange(IdKind::Domain, 0, 1);
   EXPECT_EQ(space.getNumDomainIds(), 1u);
 
   // Remove 1 symbol and 1 range identifier.
   space.removeIdRange(IdKind::Symbol, 0, 1);
   space.removeIdRange(IdKind::Range, 0, 1);
   EXPECT_EQ(space.getNumDomainIds(), 1u);
   EXPECT_EQ(space.getNumRangeIds(), 0u);
   EXPECT_EQ(space.getNumSymbolIds(), 2u);
 }

mlir/unittests/Analysis/Presburger/Utils.h

 }
 
 /// Parse a list of StringRefs to IntegerRelation and combine them into a
 /// PresburgerSet be using the union operation. It is expected that the strings
 /// are all valid IntegerSet representation and that all of them have the same
 /// number of dimensions as is specified by the numDims argument.
 inline PresburgerSet
 parsePresburgerSetFromPolyStrings(unsigned numDims, ArrayRef<StringRef> strs) {
-  PresburgerSet set = PresburgerSet::getEmpty(numDims);
+  PresburgerSet set =
+      PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(numDims));
   for (StringRef str : strs)
     set.unionInPlace(parsePoly(str));
   return set;
 }
 
 inline Matrix makeMatrix(unsigned numRow, unsigned numColumns,
                          ArrayRef<SmallVector<int64_t, 8>> matrix) {
   Matrix results(numRow, numColumns);

 /// and a 2D array that represents the output.
 inline PWMAFunction parsePWMAF(
     unsigned numInputs, unsigned numOutputs,
     ArrayRef<std::pair<StringRef, SmallVector<SmallVector<int64_t, 8>, 8>>>
         data,
     unsigned numSymbols = 0) {
   static MLIRContext context;
 
-  PWMAFunction result(numInputs - numSymbols, numSymbols, numOutputs);
+  PWMAFunction result(
+      PresburgerSpace::getSetSpace(numInputs - numSymbols, numSymbols),
+      numOutputs);
   for (const auto &pair : data) {
     IntegerPolyhedron domain = parsePoly(pair.first);
 
     result.addPiece(
         domain, makeMatrix(numOutputs, domain.getNumIds() + 1, pair.second));
   }
   return result;
 }