[MLIR][Presburger] Factor out space information to PresburgerSpace

[Groverkss]

Summary

This patch factors out space information from IntegerPolyhedron, PresburgerSet
and PWMAFunction to PresburgerSpace and its extension with local variables,
PresburgerLocalSpace.

Generally any new data structure additions in Presburger library will require
space information. This patch removes the need to duplicate the space
information.

[arjunp]

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

96

Why is this removed?

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

37

Can we avoid the duplication between the two classes by just putting numLocals in PresburgerSpace and setting it to zero in the constructor?

64

This can just be called insertId right?

67

This can just be called removeIdRange right?

71

I think it should be IDs immediately before the split become symbols or IDs immediately after become dims. Talk about leading/trailing sounds like the first or last IDs in the space, which is not what happens I think

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

105

You can factor out this repeated part into a free function

[Groverkss]

• Address Arjun's comments

[Groverkss]

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

96

Only the base class destructor needs to be virtual right?

[arjunp]

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

96

AFAIU, this will cause undefined behaviour if we ever have e.g. a unique_ptr<IntegerPolyhedron>

https://stackoverflow.com/questions/461203/when-to-use-virtual-destructors

[Groverkss]

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

96

"Inheriting from a base class using virtual destructor makes the destructor of the inheriting class automatically virtual, too (so you do not have to retype 'virtual' in the inheriting class destructor declaration)."

PresburgerSpace already has a virtual destructor. I can't find anything there that says derived classes also need to specify destructors as virtual.

[arjunp]

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

54

You can remove this line

96

Yeah, true. I think PresburgerSpace technically doesn't have to be virtual in this patch, but we need it to be later anyway so it's fine.

[Groverkss]

• Fix doc comment for IntegerPolyhedron

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_INTEGERPOLYHEDRON_H
 #define MLIR_ANALYSIS_PRESBURGER_INTEGERPOLYHEDRON_H
 
 #include "mlir/Analysis/Presburger/Fraction.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Analysis/Presburger/PresburgerSpace.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 
 /// An integer polyhedron is the set of solutions to a list of affine
 /// constraints over n integer-valued variables/identifiers. Affine constraints
 /// can be inequalities or equalities in the form:

 /// example, consider the set: (x) : (exists q : 1 <= x <= 7, x = 2q). An
 /// assignment to symbolic and dimension variables is valid if there exists some
 /// assignment to the local variable `q` satisfying these constraints. For this
 /// example, the set is equivalent to {2, 4, 6}. Mathematically, existential
 /// quantification can be thought of as the result of projection. In this
 /// example, `q` is existentially quantified. This can be thought of as the
 /// result of projecting out `q` from the previous example, i.e. we obtained {2,
 /// 4, 6} by projecting out the second dimension from {(2, 1), (4, 2), (6, 2)}.
-class IntegerPolyhedron {
+class IntegerPolyhedron : public PresburgerLocalSpace {
 public:
   /// All derived classes of IntegerPolyhedron.
   enum class Kind {
     FlatAffineConstraints,
     FlatAffineValueConstraints,
     MultiAffineFunction,
     IntegerPolyhedron
   };
 
-  /// Kind of identifier (column).
-  enum IdKind { Dimension, Symbol, Local };
-
   /// Constructs a constraint system reserving memory for the specified number
   /// of constraints and identifiers.
   IntegerPolyhedron(unsigned numReservedInequalities,
                     unsigned numReservedEqualities, unsigned numReservedCols,
                     unsigned numDims, unsigned numSymbols, unsigned numLocals)
-      : numIds(numDims + numSymbols + numLocals), numDims(numDims),
-        numSymbols(numSymbols),
-        equalities(0, numIds + 1, numReservedEqualities, numReservedCols),
-        inequalities(0, numIds + 1, numReservedInequalities, numReservedCols) {
-    assert(numReservedCols >= numIds + 1);
+      : PresburgerLocalSpace(numDims, numSymbols, numLocals),
+        equalities(0, getNumIds() + 1, numReservedEqualities, numReservedCols),
+        inequalities(0, getNumIds() + 1, numReservedInequalities,
+                     numReservedCols) {
+    assert(numReservedCols >= getNumIds() + 1);
   }
 
   /// Constructs a constraint system with the specified number of
   /// dimensions and symbols.
   IntegerPolyhedron(unsigned numDims = 0, unsigned numSymbols = 0,
                     unsigned numLocals = 0)
       : IntegerPolyhedron(/*numReservedInequalities=*/0,
                           /*numReservedEqualities=*/0,
                           /*numReservedCols=*/numDims + numSymbols + numLocals +
                               1,
                           numDims, numSymbols, numLocals) {}
 
   /// 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);
   }
 
-  virtual ~IntegerPolyhedron() = default;
-
   /// Return the kind of this IntegerPolyhedron.
   virtual Kind getKind() const { return Kind::IntegerPolyhedron; }
 
   static bool classof(const IntegerPolyhedron *cst) { return true; }
 
   // Clones this object.
   std::unique_ptr<IntegerPolyhedron> clone() const;
 
public:
   inline int64_t atIneq(unsigned i, unsigned j) const {
     return inequalities(i, j);
   }
   inline int64_t &atIneq(unsigned i, unsigned j) { return inequalities(i, j); }
 
   unsigned getNumConstraints() const {
     return getNumInequalities() + getNumEqualities();
   }
-  inline unsigned getNumIds() const { return numIds; }
-  inline unsigned getNumDimIds() const { return numDims; }
-  inline unsigned getNumSymbolIds() const { return numSymbols; }
-  inline unsigned getNumDimAndSymbolIds() const { return numDims + numSymbols; }
-  inline unsigned getNumLocalIds() const {
-    return numIds - numDims - numSymbols;
-  }
 
   /// Returns the number of columns in the constraint system.
-  inline unsigned getNumCols() const { return numIds + 1; }
+  inline unsigned getNumCols() const { return getNumIds() + 1; }
 
   inline unsigned getNumEqualities() const { return equalities.getNumRows(); }
 
   inline unsigned getNumInequalities() const {
     return inequalities.getNumRows();
   }
 
   inline unsigned getNumReservedEqualities() const {
public:
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. The coefficient columns
   /// corresponding to the added identifiers are initialized to zero. Return the
   /// absolute column position (i.e., not relative to the kind of identifier)
   /// of the first added identifier.
   unsigned insertDimId(unsigned pos, unsigned num = 1);
   unsigned insertSymbolId(unsigned pos, unsigned num = 1);
   unsigned insertLocalId(unsigned pos, unsigned num = 1);
-  virtual unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1);
+  unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
 
   /// Append `num` identifiers of the specified kind after the last identifier.
   /// of that kind. Return the position of the first appended column. The
   /// coefficient columns corresponding to the added identifiers are initialized
   /// to zero.
   unsigned appendDimId(unsigned num = 1);
   unsigned appendSymbolId(unsigned num = 1);
   unsigned appendLocalId(unsigned num = 1);
public:
   /// Projects out (aka eliminates) `num` identifiers starting at position
   /// `pos`. The resulting constraint system is the shadow along the dimensions
   /// that still exist. This method may not always be integer exact.
   // TODO: deal with integer exactness when necessary - can return a value to
   // mark exactness for example.
   void projectOut(unsigned pos, unsigned num);
   inline void projectOut(unsigned pos) { return projectOut(pos, 1); }
 
-  /// Changes the partition between dimensions and symbols. Depending on the new
-  /// symbol count, either a chunk of trailing dimensional identifiers becomes
-  /// symbols, or some of the leading symbols become dimensions.
-  void setDimSymbolSeparation(unsigned newSymbolCount);
-
   /// Tries to fold the specified identifier to a constant using a trivial
   /// equality detection; if successful, the constant is substituted for the
   /// identifier everywhere in the constraint system and then removed from the
   /// system.
   LogicalResult constantFoldId(unsigned pos);
 
   /// This method calls `constantFoldId` for the specified range of identifiers,
   /// `num` identifiers starting at position `pos`.
protected:
   /// equality/inequality buffer sizes aren't consistent; true otherwise. This
   /// is meant to be used within an assert internally.
   virtual bool hasConsistentState() const;
 
   /// Prints the number of constraints, dimensions, symbols and locals in the
   /// IntegerPolyhedron.
   virtual void printSpace(raw_ostream &os) 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;
-
-  /// Get the number of ids of the specified kind.
-  unsigned getNumIdKind(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;
-
   /// Removes identifiers in the column range [idStart, idLimit), and copies any
   /// remaining valid data into place, updates member variables, and resizes
   /// arrays as needed.
-  virtual void removeIdRange(unsigned idStart, unsigned idLimit);
+  void removeIdRange(unsigned idStart, unsigned idLimit) override;
 
   /// A parameter that controls detection of an unrealistic number of
   /// constraints. If the number of constraints is this many times the number of
   /// variables, we consider such a system out of line with the intended use
   /// case of IntegerPolyhedron.
   // The rationale for 32 is that in the typical simplest of cases, an
   // identifier is expected to have one lower bound and one upper bound
   // constraint. With a level of tiling or a connection to another identifier
   // through a div or mod, an extra pair of bounds gets added. As a limit, we
   // don't expect an identifier to have more than 32 lower/upper/equality
   // constraints. This is conservatively set low and can be raised if needed.
   constexpr static unsigned kExplosionFactor = 32;
 
-  /// Total number of identifiers.
-  unsigned numIds;
-
-  /// Number of identifiers corresponding to real dimensions.
-  unsigned numDims;
-
-  /// Number of identifiers corresponding to symbols (unknown but constant for
-  /// analysis).
-  unsigned numSymbols;
-
   /// Coefficients of affine equalities (in == 0 form).
   Matrix equalities;
 
   /// Coefficients of affine inequalities (in >= 0 form).
   Matrix inequalities;
 };
 
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_INTEGERPOLYHEDRON_H

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

public:
   ~MultiAffineFunction() override = default;
   Kind getKind() const override { return Kind::MultiAffineFunction; }
   bool classof(const IntegerPolyhedron *poly) const {
     return poly->getKind() == Kind::MultiAffineFunction;
   }
 
   unsigned getNumInputs() const { return getNumDimAndSymbolIds(); }
   unsigned getNumOutputs() const { return output.getNumRows(); }
-  bool isConsistent() const { return output.getNumColumns() == numIds + 1; }
+  bool isConsistent() const {
+    return output.getNumColumns() == getNumIds() + 1;
+  }
   const IntegerPolyhedron &getDomain() const { return *this; }
 
   bool hasCompatibleDimensions(const MultiAffineFunction &f) const;
 
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. The coefficient columns
   /// corresponding to the added identifiers are initialized to zero. Return the
   /// absolute column position (i.e., not relative to the kind of identifier)

 /// this represents a partial function and so could be undefined at some points.
 ///
 /// 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. Note that local ids in the
 /// piece are not supported for the latter.
-class PWMAFunction {
+class PWMAFunction : PresburgerSpace {
 public:
   PWMAFunction(unsigned numDims, unsigned numSymbols, unsigned numOutputs)
-      : numDims(numDims), numSymbols(numSymbols), numOutputs(numOutputs) {
+      : PresburgerSpace(numDims, numSymbols), numOutputs(numOutputs) {
     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(); }
   unsigned getNumOutputs() const { return numOutputs; }
-  unsigned getNumInputs() const { return numDims + numSymbols; }
-  unsigned getNumDimIds() const { return numDims; }
-  unsigned getNumSymbolIds() const { return numSymbols; }
+  unsigned getNumInputs() const { return getNumIds(); }
   MultiAffineFunction &getPiece(unsigned i) { return pieces[i]; }
 
   /// Return the domain of this piece-wise MultiAffineFunction. This is the
   /// union of the domains of all the pieces.
   PresburgerSet getDomain() const;
 
   /// Check whether the `this` and the given function have compatible
   /// dimensions, i.e., the same number of dimension inputs, symbol inputs, and
public:
 
   void print(raw_ostream &os) const;
   void dump() const;
 
 private:
   /// The list of pieces in this piece-wise MultiAffineFunction.
   SmallVector<MultiAffineFunction, 4> pieces;
 
-  /// The number of dimensions ids in the domains.
-  unsigned numDims;
-  /// The number of symbol ids in the domains.
-  unsigned numSymbols;
   /// The number of output ids.
   unsigned numOutputs;
 };
 
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PWMAFUNCTION_H

mlir/include/mlir/Analysis/Presburger/PresburgerSet.h

 ///
 /// The IntegerPolyhedrons (Polys) are stored in a vector, and the set
 /// represents the union of these Polys. An empty list corresponds to the empty
 /// set.
 ///
 /// Note that there are no invariants guaranteed on the list of Poly other than
 /// that they are all in the same space, i.e., they all have the same number of
 /// dimensions and symbols. For example, the Polys may overlap each other.
-class PresburgerSet {
+class PresburgerSet : public PresburgerSpace {
 public:
   explicit PresburgerSet(const IntegerPolyhedron &poly);
 
   /// Return the number of Polys in the union.
   unsigned getNumPolys() const;
 
-  /// Return the number of real dimensions.
-  unsigned getNumDimIds() const;
-
-  /// Return the number of symbolic dimensions.
-  unsigned getNumSymbolIds() const;
-
   /// Return a reference to the list of IntegerPolyhedrons.
   ArrayRef<IntegerPolyhedron> getAllIntegerPolyhedron() const;
 
   /// Return the IntegerPolyhedron at the specified index.
   const IntegerPolyhedron &getIntegerPolyhedron(unsigned index) const;
 
   /// Mutate this set, turning it into the union of this set and the given
   /// IntegerPolyhedron.
public:
   ///
   /// In particular, removes all Polys which are subsets of other Polys in the
   /// union.
   PresburgerSet coalesce() const;
 
 private:
   /// Construct an empty PresburgerSet.
   PresburgerSet(unsigned numDims = 0, unsigned numSymbols = 0)
-      : numDims(numDims), numSymbols(numSymbols) {}
+      : PresburgerSpace(numDims, numSymbols) {}
 
   /// Return the set difference poly \ set.
   static PresburgerSet getSetDifference(IntegerPolyhedron poly,
                                         const PresburgerSet &set);
 
-  /// Number of identifiers corresponding to real dimensions.
-  unsigned numDims;
-
-  /// Number of symbolic dimensions, unknown but constant for analysis, as in
-  /// IntegerPolyhedron.
-  unsigned numSymbols;
-
   /// The list of integerPolyhedrons that this set is the union of.
   SmallVector<IntegerPolyhedron, 2> integerPolyhedrons;
 };
 
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSET_H

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.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
+#define MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
+
+#include "llvm/Support/ErrorHandling.h"
+
+namespace mlir {
+
+class PresburgerLocalSpace;
+
+/// PresburgerSpace is a tuple of identifiers with information about what kind
+/// they correspond to. The identifiers can be split into three types:
+///
+/// Dimension: Ordinary variables over which the space is represented.
+///
+/// Symbol: Symbol identifiers correspond to fixed but unknown values.
+/// Mathematically, a space with symbolic identifiers is like a
+/// family of spaces indexed by the symbolic identifiers.
+///
+/// Local: Local identifiers correspond to existentially quantified variables.
+///
+/// PresburgerSpace only supports identifiers of kind Dimension and Symbol.
+class PresburgerSpace {
+  friend PresburgerLocalSpace;
+
+public:
+  /// Kind of identifier (column).
+  enum IdKind { Dimension, Symbol, Local };
+
+  PresburgerSpace(unsigned numDims, unsigned numSymbols)
+      : numDims(numDims), numSymbols(numSymbols), numLocals(0) {}
+
+  virtual ~PresburgerSpace() = default;
+
+  unsigned getNumIds() const { return numDims + numSymbols + numLocals; }
+  unsigned getNumDimIds() const { return numDims; }
+  unsigned getNumSymbolIds() const { return numSymbols; }
+  unsigned getNumDimAndSymbolIds() const { return numDims + numSymbols; }
+
+  /// 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;
+
+  /// Get the number of elements of the specified kind in the range
+  /// [idStart, idLimit).
+  unsigned getIdKindOverlap(IdKind kind, unsigned idStart,
+                            unsigned idLimit) 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 in the column range [idStart, idLimit).
+  virtual void removeIdRange(unsigned idStart, unsigned idLimit);
+
+  /// 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);
+
+private:
+  PresburgerSpace(unsigned numDims, unsigned numSymbols, unsigned numLocals)
+      : numDims(numDims), numSymbols(numSymbols), numLocals(numLocals) {}
+
+  /// Number of identifiers corresponding to real dimensions.
+  unsigned numDims;
+
+  /// Number of identifiers corresponding to symbols (unknown but constant for
+  /// analysis).
+  unsigned numSymbols;
+
+  /// Total number of identifiers.
+  unsigned numLocals;
+};
+
+/// Extension of PresburgerSpace supporting Local identifiers.
+class PresburgerLocalSpace : public PresburgerSpace {
+public:
+  PresburgerLocalSpace(unsigned numDims, unsigned numSymbols,
+                       unsigned numLocals)
+      : PresburgerSpace(numDims, numSymbols, numLocals) {}
+
+  unsigned getNumLocalIds() const { return numLocals; }
+
+  /// 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.
+  unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
+
+  /// Removes identifiers in the column range [idStart, idLimit).
+  void removeIdRange(unsigned idStart, unsigned idLimit) override;
+};
+
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H

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

public:
   /// of constraints and identifiers.
   FlatAffineValueConstraints(unsigned numReservedInequalities,
                              unsigned numReservedEqualities,
                              unsigned numReservedCols, unsigned numDims,
                              unsigned numSymbols, unsigned numLocals,
                              ArrayRef<Optional<Value>> valArgs = {})
       : FlatAffineConstraints(numReservedInequalities, numReservedEqualities,
                               numReservedCols, numDims, numSymbols, numLocals) {
-    assert(numReservedCols >= numIds + 1);
-    assert(valArgs.empty() || valArgs.size() == numIds);
+    assert(numReservedCols >= getNumIds() + 1);
+    assert(valArgs.empty() || valArgs.size() == getNumIds());
     values.reserve(numReservedCols);
     if (valArgs.empty())
-      values.resize(numIds, None);
+      values.resize(getNumIds(), None);
     else
       values.append(valArgs.begin(), valArgs.end());
   }
 
   /// 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 FlatAffineConstraints &fac,
                              ArrayRef<Optional<Value>> valArgs = {})
       : FlatAffineConstraints(fac) {
-    assert(valArgs.empty() || valArgs.size() == numIds);
+    assert(valArgs.empty() || valArgs.size() == getNumIds());
     if (valArgs.empty())
-      values.resize(numIds, None);
+      values.resize(getNumIds(), None);
     else
       values.append(valArgs.begin(), valArgs.end());
   }
 
   /// Create a flat affine constraint system from an AffineValueMap or a list of
   /// these. The constructed system will only include equalities.
   explicit FlatAffineValueConstraints(const AffineValueMap &avm);
   explicit FlatAffineValueConstraints(ArrayRef<const AffineValueMap *> avmRef);
public:
 
   /// 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 < numIds || start == end) && "invalid start position");
-    assert(end <= numIds && "invalid end position");
+    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, numIds, values);
+    getValues(0, getNumIds(), values);
   }
 
   inline ArrayRef<Optional<Value>> getMaybeValues() const {
     return {values.data(), values.size()};
   }
 
   inline ArrayRef<Optional<Value>> getMaybeDimValues() const {
     return {values.data(), getNumDimIds()};
   }
 
   inline ArrayRef<Optional<Value>> getMaybeSymbolValues() const {
     return {values.data() + getNumDimIds(), getNumSymbolIds()};
   }
 
   inline ArrayRef<Optional<Value>> getMaybeDimAndSymbolValues() const {
     return {values.data(), getNumDimIds() + getNumSymbolIds()};
   }
 
   /// Sets the Value associated with the pos^th identifier.
   inline void setValue(unsigned pos, Value val) {
-    assert(pos < numIds && "invalid id position");
+    assert(pos < getNumIds() && "invalid id position");
     values[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 < numIds || end == start) && "invalid start position");
-    assert(end <= numIds && "invalid end position");
+    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

mlir/lib/Analysis/Presburger/CMakeLists.txt

 add_mlir_library(MLIRPresburger
   IntegerPolyhedron.cpp
   LinearTransform.cpp
   Matrix.cpp
   PresburgerSet.cpp
+  PresburgerSpace.cpp
   PWMAFunction.cpp
   Simplex.cpp
   Utils.cpp
 
   LINK_LIBS PUBLIC
   MLIRSupport
   )

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

 unsigned IntegerPolyhedron::insertLocalId(unsigned pos, unsigned num) {
   return insertId(IdKind::Local, pos, num);
 }
 
 unsigned IntegerPolyhedron::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
   unsigned absolutePos = getIdKindOffset(kind) + pos;
-  if (kind == IdKind::Dimension)
-    numDims += num;
-  else if (kind == IdKind::Symbol)
-    numSymbols += num;
-  numIds += num;
-
   inequalities.insertColumns(absolutePos, num);
   equalities.insertColumns(absolutePos, num);
 
-  return absolutePos;
+  return PresburgerLocalSpace::insertId(kind, pos, num);
 }
 
 unsigned IntegerPolyhedron::appendDimId(unsigned num) {
   unsigned pos = getNumDimIds();
   insertId(IdKind::Dimension, pos, num);
   return pos;
 }
 

 void IntegerPolyhedron::removeIdRange(IdKind kind, unsigned idStart,
                                       unsigned idLimit) {
   assert(idLimit <= getNumIdKind(kind));
   removeIdRange(getIdKindOffset(kind) + idStart,
                 getIdKindOffset(kind) + idLimit);
 }
 
 void IntegerPolyhedron::removeIdRange(unsigned idStart, unsigned idLimit) {
-  assert(idLimit < getNumCols() && "invalid id limit");
-
-  if (idStart >= idLimit)
-    return;
-
-  // We are going to be removing one or more identifiers from the range.
-  assert(idStart < getNumIds() && "invalid idStart position");
+  // Update space paramaters.
+  PresburgerLocalSpace::removeIdRange(idStart, idLimit);
 
   // Remove eliminated identifiers from the constraints..
   equalities.removeColumns(idStart, idLimit - idStart);
   inequalities.removeColumns(idStart, idLimit - idStart);
-
-  // Update members numDims, numSymbols and numIds.
-  unsigned numDimsEliminated =
-      getIdKindOverlap(IdKind::Dimension, idStart, idLimit);
-  unsigned numSymbolsEliminated =
-      getIdKindOverlap(IdKind::Symbol, idStart, idLimit);
-
-  numDims -= numDimsEliminated;
-  numSymbols -= numSymbolsEliminated;
-  numIds -= (idLimit - idStart);
 }
 
 void IntegerPolyhedron::removeEquality(unsigned pos) {
   equalities.removeRow(pos);
 }
 
 void IntegerPolyhedron::removeInequality(unsigned pos) {
   inequalities.removeRow(pos);
if (posA == posB)
     return;
 
   for (unsigned r = 0, e = getNumInequalities(); r < e; r++)
     std::swap(atIneq(r, posA), atIneq(r, posB));
   for (unsigned r = 0, e = getNumEqualities(); r < e; r++)
     std::swap(atEq(r, posA), atEq(r, posB));
 }
 
-unsigned IntegerPolyhedron::getIdKindOffset(IdKind kind) const {
-  if (kind == IdKind::Dimension)
-    return 0;
-  if (kind == IdKind::Symbol)
-    return getNumDimIds();
-  if (kind == IdKind::Local)
-    return getNumDimAndSymbolIds();
-  llvm_unreachable("IdKind expected to be Dimension, Symbol or Local!");
-}
-
-unsigned IntegerPolyhedron::getIdKindEnd(IdKind kind) const {
-  return getIdKindOffset(kind) + getNumIdKind(kind);
-}
-
-unsigned IntegerPolyhedron::getNumIdKind(IdKind kind) const {
-  if (kind == IdKind::Dimension)
-    return getNumDimIds();
-  if (kind == IdKind::Symbol)
-    return getNumSymbolIds();
-  if (kind == IdKind::Local)
-    return getNumLocalIds();
-  llvm_unreachable("IdKind expected to be Dimension, Symbol or Local!");
-}
-
-unsigned IntegerPolyhedron::getIdKindOverlap(IdKind kind, unsigned idStart,
-                                             unsigned idLimit) const {
-  unsigned idRangeStart = getIdKindOffset(kind);
-  unsigned idRangeEnd = getIdKindEnd(kind);
-
-  // Compute number of elements in intersection of the ranges [idStart, idLimit)
-  // and [idRangeStart, idRangeEnd).
-  unsigned overlapStart = std::max(idStart, idRangeStart);
-  unsigned overlapEnd = std::min(idLimit, idRangeEnd);
-
-  if (overlapStart > overlapEnd)
-    return 0;
-  return overlapEnd - overlapStart;
-}
-
 void IntegerPolyhedron::clearConstraints() {
   equalities.resizeVertically(0);
   inequalities.resizeVertically(0);
 }
 
 /// Gather 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 IntegerPolyhedron::addLocalFloorDiv(ArrayRef<int64_t> dividend,
 
   // -dividend +qdivisor * q + divisor - 1 >= 0
   std::transform(bound.begin(), bound.end(), bound.begin(),
                  std::negate<int64_t>());
   bound[bound.size() - 1] += divisor - 1;
   addInequality(bound);
 }
 
-void IntegerPolyhedron::setDimSymbolSeparation(unsigned newSymbolCount) {
-  assert(newSymbolCount <= getNumDimAndSymbolIds() &&
-         "invalid separation position");
-  numDims = numDims + numSymbols - newSymbolCount;
-  numSymbols = newSymbolCount;
-}
-
 /// Finds an equality that equates the specified identifier to a constant.
 /// Returns the position of the equality row. If 'symbolic' is set to true,
 /// symbols are also treated like a constant, i.e., an affine function of the
 /// symbols is also treated like a constant. Returns -1 if such an equality
 /// could not be found.
 static int findEqualityToConstant(const IntegerPolyhedron &cst, unsigned pos,
                                   bool symbolic = false) {
   assert(pos < cst.getNumIds() && "invalid position");

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 using namespace presburger_utils;
 
 PresburgerSet::PresburgerSet(const IntegerPolyhedron &poly)
-    : numDims(poly.getNumDimIds()), numSymbols(poly.getNumSymbolIds()) {
+    : PresburgerSpace(poly) {
   unionPolyInPlace(poly);
 }
 
 unsigned PresburgerSet::getNumPolys() const {
   return integerPolyhedrons.size();
 }
 
-unsigned PresburgerSet::getNumDimIds() const { return numDims; }
-
-unsigned PresburgerSet::getNumSymbolIds() const { return numSymbols; }
-
 ArrayRef<IntegerPolyhedron> PresburgerSet::getAllIntegerPolyhedron() const {
   return integerPolyhedrons;
 }
 
 const IntegerPolyhedron &
 PresburgerSet::getIntegerPolyhedron(unsigned index) const {
   assert(index < integerPolyhedrons.size() && "index out of bounds!");
   return integerPolyhedrons[index];

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

+//===- PresburgerSpace.cpp - MLIR PresburgerSpace Class -------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/Presburger/PresburgerSpace.h"
+#include <algorithm>
+#include <cassert>
+
+using namespace mlir;
+
+unsigned PresburgerSpace::getNumIdKind(IdKind kind) const {
+  if (kind == IdKind::Dimension)
+    return getNumDimIds();
+  if (kind == IdKind::Symbol)
+    return getNumSymbolIds();
+  if (kind == IdKind::Local)
+    return numLocals;
+  llvm_unreachable("IdKind does not exit!");
+}
+
+unsigned PresburgerSpace::getIdKindOffset(IdKind kind) const {
+  if (kind == IdKind::Dimension)
+    return 0;
+  if (kind == IdKind::Symbol)
+    return getNumDimIds();
+  if (kind == IdKind::Local)
+    return getNumDimAndSymbolIds();
+  llvm_unreachable("IdKind does not exit!");
+}
+
+unsigned PresburgerSpace::getIdKindEnd(IdKind kind) const {
+  return getIdKindOffset(kind) + getNumIdKind(kind);
+}
+
+unsigned PresburgerSpace::getIdKindOverlap(IdKind kind, unsigned idStart,
+                                           unsigned idLimit) const {
+  unsigned idRangeStart = getIdKindOffset(kind);
+  unsigned idRangeEnd = getIdKindEnd(kind);
+
+  // Compute number of elements in intersection of the ranges [idStart, idLimit)
+  // and [idRangeStart, idRangeEnd).
+  unsigned overlapStart = std::max(idStart, idRangeStart);
+  unsigned overlapEnd = std::min(idLimit, idRangeEnd);
+
+  if (overlapStart > overlapEnd)
+    return 0;
+  return overlapEnd - overlapStart;
+}
+
+unsigned PresburgerSpace::insertId(IdKind kind, unsigned pos, unsigned num) {
+  assert(pos <= getNumIdKind(kind));
+
+  unsigned absolutePos = getIdKindOffset(kind) + pos;
+
+  if (kind == IdKind::Dimension)
+    numDims += num;
+  else if (kind == IdKind::Symbol)
+    numSymbols += num;
+  else
+    llvm_unreachable(
+        "PresburgerSpace only supports Dimensions and Symbol identifiers!");
+
+  return absolutePos;
+}
+
+void PresburgerSpace::removeIdRange(unsigned idStart, unsigned idLimit) {
+  assert(idLimit <= getNumIds() && "invalid id limit");
+
+  if (idStart >= idLimit)
+    return;
+
+  // We are going to be removing one or more identifiers from the range.
+  assert(idStart < getNumIds() && "invalid idStart position");
+
+  // Update members numDims, numSymbols and numIds.
+  unsigned numDimsEliminated =
+      getIdKindOverlap(IdKind::Dimension, idStart, idLimit);
+  unsigned numSymbolsEliminated =
+      getIdKindOverlap(IdKind::Symbol, idStart, idLimit);
+
+  numDims -= numDimsEliminated;
+  numSymbols -= numSymbolsEliminated;
+}
+
+unsigned PresburgerLocalSpace::insertId(IdKind kind, unsigned pos,
+                                        unsigned num) {
+  if (kind == IdKind::Local) {
+    numLocals += num;
+    return getIdKindOffset(IdKind::Local) + pos;
+  }
+  return PresburgerSpace::insertId(kind, pos, num);
+}
+
+void PresburgerLocalSpace::removeIdRange(unsigned idStart, unsigned idLimit) {
+  assert(idLimit <= getNumIds() && "invalid id limit");
+
+  if (idStart >= idLimit)
+    return;
+
+  // We are going to be removing one or more identifiers from the range.
+  assert(idStart < getNumIds() && "invalid idStart position");
+
+  unsigned numLocalsEliminated =
+      getIdKindOverlap(IdKind::Local, idStart, idLimit);
+
+  // Update space parameters.
+  PresburgerSpace::removeIdRange(
+      idStart, std::min(idLimit, PresburgerSpace::getNumIds()));
+
+  // Update local ids.
+  numLocals -= numLocalsEliminated;
+}
+
+void PresburgerSpace::setDimSymbolSeparation(unsigned newSymbolCount) {
+  assert(newSymbolCount <= getNumDimAndSymbolIds() &&
+         "invalid separation position");
+  numDims = numDims + numSymbols - newSymbolCount;
+  numSymbols = newSymbolCount;
+}