[MLIR] Move Presburger Math from FlatAffineConstraints to Presburger/IntegerPolyhedron

[Groverkss]

Summary

This patch factors out math functionality that is a subset of Presburger arithmetic and moves it from FlatAffineConstraints to Presburger/IntegerPolyhedron. This patch only moves some parts of the functionality planned to be moved, with subsequent patches moving more functionality. There are three main reasons for this:

1. This split makes the Presburger Library easier and more flexible to use across MLIR, by not depending on IR.
2. This split allows the Presburger library to be developed independently from Affine Analysis, with Affine Analysis using this library.
3. With more functionality being upstreamed to the Presburger Library, the mlir/Analysis directory will be cluttered with Presburger library components since they depend on math functionality from FlatAffineConstraints. Moving this functionality to the Presburger directory allows keeping the new functionality in the Presburger directory.

This patch is part of an ongoing effort to make the Presburger Library easier to use. The motivation for this effort is the feedback received at the LLVM conference from Mehdi and others.

[bondhugula]

mlir/include/mlir/Analysis/Presburger/PresburgerBasicSet.h

1 ↗

(On Diff #390196)

While refactoring this out looks fine, PresburgerBasicSet is too heavy a name for what this comprises. I can't immediately think of a better one though.

mlir/lib/Analysis/Presburger/PresburgerBasicSet.cpp

1 ↗

(On Diff #390196)

Wrong file extension here. Also, no need of the C++.

[Groverkss]

• Fix Presburger.cpp file header

mlir/include/mlir/Analysis/Presburger/PresburgerBasicSet.h

1 ↗

(On Diff #390196)

I think there is some misunderstanding. PresburgerBasicSet is meant to represent all functionality of FlatAffineConstraints that does not interact with the IR. This patch only moves some functionality so as to not make it harder to review. Subsequent patches will move more functionality to PresburgerBasicSet from FlatAffineConstraints.

[grosser]

mlir/include/mlir/Analysis/Presburger/PresburgerBasicSet.h

1 ↗

(On Diff #390196)

What about just calling it "BasicSet". This could help to clarify that this is not just for Presburger arithmetic?

Best,
Tobias

[Groverkss]

mlir/include/mlir/Analysis/Presburger/PresburgerBasicSet.h

1 ↗

(On Diff #390196)

I do not have any preference for what we call it :)

We can go with BasicSet, IntegerPolyhedra, Polyhedra, PresburgerBasicSet, or any other name you may have in mind.

@bondhugula Which one do you think will be good?

[ftynse]

mlir/include/mlir/Analysis/Presburger/PresburgerBasicSet.h

1 ↗

(On Diff #390196)

Fly-by comment: polyhedra is plural of polyhedron. So if you chose that name, please use the singular form unless you explicitly refer to plurality.

[bondhugula]

mlir/include/mlir/Analysis/Presburger/PresburgerBasicSet.h

1 ↗

(On Diff #390196)

BasicSet, IntegerPolyhedron, and Polyhedron are all candidates. BasicSet may only make sense when there is another one named Set - otherwise, Basic is only adding ambiguity or noise. Polyhedron or IntegerPolyhedron?

[Groverkss]

mlir/include/mlir/Analysis/Presburger/PresburgerBasicSet.h

1 ↗

(On Diff #390196)

IntegerPolyhedron sounds good to me.

[Groverkss]

Change name from PresburgerBasicSet to IntegerPolyhedron

[bondhugula]

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

14

...INTEGERPOLYHEDRON_H

mlir/lib/Analysis/Presburger/CMakeLists.txt

4

Please maintain these in sorted order - order above should be fixed as well.

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

16

These shouldn't be enclosed in a namespace this way. Instead, just using namespace mlir;.

[bondhugula]

Address clang-format issues.

[Groverkss]

• Address bondhugula's comments

[arjunp]

SGTM. Some nits and minor comments below.

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

31

40

42–43

It seems strange to say that local variables represent existential quantification. I would rather say that local variables are existentially quantified.

44

May as well use similar constraints 1 <= x <= 7, x = 2q to the set above. Also, it would be good to not break the set across two lines. Also, I think you can move the example to the bottom.

For example, consider the set: (x) : (exists q : 1 <= x <= 7, x = 2q). This set is is equivalent to {2, 4, 6}.`

Also, I think you can mention that mathematically, existential quantification can be thought of as the result of projection. In the above 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, 3)}.

If you use this then remember to change y to q in the first example

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

164–170

I think these have some tests. You can probably move them to an IntegerPolyhedronTest.

[Groverkss]

• Update docs
• Move tests for IntegerPolyhedron to Presburger/IntegerPolyhedronTest

[bondhugula]

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

10

polyhedra is plural. Drop "an".

[Groverkss]

• Update docs

mlir/include/mlir/Analysis/AffineStructures.h

 //===- AffineStructures.h - MLIR Affine Structures 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Structures for affine/polyhedral analysis of ML functions.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_AFFINESTRUCTURES_H
 #define MLIR_ANALYSIS_AFFINESTRUCTURES_H
 
+#include "mlir/Analysis/Presburger/IntegerPolyhedron.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 
 class AffineCondition;

 /// representational power. Each local identifier is always (by construction) a
 /// floordiv of a pure add/mul affine function of dimensional, symbolic, and
 /// other local identifiers, in a non-mutually recursive way. Hence, every local
 /// identifier can ultimately always be recovered as an affine function of
 /// dimensional and symbolic identifiers (involving floordiv's); note however
 /// that some floordiv combinations are converted to mod's by AffineExpr
 /// construction.
 ///
-class FlatAffineConstraints {
+class FlatAffineConstraints : public IntegerPolyhedron {
 public:
   /// All derived classes of FlatAffineConstraints.
   enum class Kind { FlatAffineConstraints, FlatAffineValueConstraints };
 
-  /// Kind of identifier (column).
-  enum IdKind { Dimension, Symbol, Local };
-
   /// 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)
-      : numIds(numDims + numSymbols + numLocals), numDims(numDims),
-        numSymbols(numSymbols),
-        equalities(0, numIds + 1, numReservedEqualities, numReservedCols),
-        inequalities(0, numIds + 1, numReservedInequalities, numReservedCols) {
-    assert(numReservedCols >= numIds + 1);
-  }
+      : IntegerPolyhedron(numReservedInequalities, numReservedEqualities,
+                          numReservedCols, 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)
-      : FlatAffineConstraints(/*numReservedInequalities=*/0,
-                              /*numReservedEqualities=*/0,
-                              /*numReservedCols=*/numDims + numSymbols +
-                                  numLocals + 1,
-                              numDims, numSymbols, numLocals) {}
+      : IntegerPolyhedron(/*numReservedInequalities=*/0,
+                          /*numReservedCols=*/numDims + numSymbols + numLocals +
+                              1,
 
   /// Return a system with no constraints, i.e., one which is satisfied by all
   /// points.
   static FlatAffineConstraints getUniverse(unsigned numDims = 0,
                                            unsigned numSymbols = 0) {
     return FlatAffineConstraints(numDims, numSymbols);
   }
 
   /// Creates an affine constraint system from an IntegerSet.
   explicit FlatAffineConstraints(IntegerSet set);
 
   FlatAffineConstraints(const MutableAffineMap &map);
 
   virtual ~FlatAffineConstraints() = default;
 
   /// Return the kind of this FlatAffineConstraints.
   virtual Kind getKind() const { return Kind::FlatAffineConstraints; }
 
   static bool classof(const FlatAffineConstraints *cst) { return true; }
 
-  /// Clears any existing data and reserves memory for the specified
-  /// constraints.
-  virtual void reset(unsigned numReservedInequalities,
-                     unsigned numReservedEqualities, unsigned numReservedCols,
-                     unsigned numDims, unsigned numSymbols,
-                     unsigned numLocals = 0);
-
-  void reset(unsigned numDims = 0, unsigned numSymbols = 0,
-             unsigned numLocals = 0);
-
-  /// Appends constraints from `other` into `this`. This is equivalent to an
-  /// intersection with no simplification of any sort attempted.
-  void append(const FlatAffineConstraints &other);
-
   /// Checks for emptiness by performing variable elimination on all
   /// identifiers, running the GCD test on each equality constraint, and
   /// checking for invalid constraints. Returns true if the GCD test fails for
   /// any equality, or if any invalid constraints are discovered on any row.
   /// Returns false otherwise.
   bool isEmpty() const;
 
   /// Runs the GCD test on all equality constraints. Returns true if this test
public:
   void getLocalReprs(
       std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const;
   void getLocalReprs(std::vector<SmallVector<int64_t, 8>> &dividends,
                      SmallVector<unsigned, 4> &denominators) const;
 
   // Clones this object.
   std::unique_ptr<FlatAffineConstraints> clone() const;
 
-  /// Returns the value at the specified equality row and column.
-  inline int64_t atEq(unsigned i, unsigned j) const { return equalities(i, j); }
-  inline int64_t &atEq(unsigned i, unsigned j) { return equalities(i, j); }
-
-  /// Returns the value at the specified inequality row and column.
-  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); }
-
-  /// Returns the number of columns in the constraint system.
-  inline unsigned getNumCols() const { return numIds + 1; }
-
-  inline unsigned getNumEqualities() const { return equalities.getNumRows(); }
-
-  inline unsigned getNumInequalities() const {
-    return inequalities.getNumRows();
-  }
-
-  inline unsigned getNumReservedEqualities() const {
-    return equalities.getNumReservedRows();
-  }
-
-  inline unsigned getNumReservedInequalities() const {
-    return inequalities.getNumReservedRows();
-  }
-
-  inline ArrayRef<int64_t> getEquality(unsigned idx) const {
-    return equalities.getRow(idx);
-  }
-
-  inline ArrayRef<int64_t> getInequality(unsigned idx) const {
-    return inequalities.getRow(idx);
-  }
-
   /// The type of bound: equal, lower bound or upper bound.
   enum BoundType { EQ, LB, UB };
 
   /// Adds a bound for the identifier at the specified position with constraints
   /// being drawn from the specified bound map. In case of an EQ bound, the
   /// bound map is expected to have exactly one result. In case of a LB/UB, the
   /// bound map may have more than one result, for each of which an inequality
   /// is added.
public:
   /// identifiers (dimensional and symbolic). This method is able to detect
   /// identifiers as floordiv's and mod's of affine expressions of other
   /// identifiers with respect to (positive) constants. Sets bound map to a
   /// null AffineMap if such a bound can't be found (or yet unimplemented).
   void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context,
                       SmallVectorImpl<AffineMap> *lbMaps,
                       SmallVectorImpl<AffineMap> *ubMaps);
 
-  /// Adds an inequality (>= 0) from the coefficients specified in `inEq`.
-  void addInequality(ArrayRef<int64_t> inEq);
-  /// Adds an equality from the coefficients specified in `eq`.
-  void addEquality(ArrayRef<int64_t> eq);
-
   /// Adds a new local identifier as the floordiv of an affine function of other
   /// identifiers, the coefficients of which are provided in `dividend` and with
   /// respect to a positive constant `divisor`. Two constraints are added to the
   /// system to capture equivalence with the floordiv:
   /// q = dividend floordiv c    <=>   c*q <= dividend <= c*q + c - 1.
   void addLocalFloorDiv(ArrayRef<int64_t> dividend, int64_t divisor);
 
-  /// Swap the posA^th identifier with the posB^th identifier.
-  virtual void swapId(unsigned posA, unsigned posB);
-
-  /// 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);
-
-  /// 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);
-
   /// Composes an affine map whose dimensions and symbols match one to one with
   /// the dimensions and symbols of this FlatAffineConstraints. The results of
   /// the map `other` are added as the leading dimensions of this constraint
   /// system. Returns failure if `other` is a semi-affine map.
   LogicalResult composeMatchingMap(AffineMap other);
 
   /// 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); }
 
-  /// Removes identifiers of the specified kind with the specified pos (or
-  /// within the specified range) from the system. The specified location is
-  /// relative to the first identifier of the specified kind.
-  void removeId(IdKind kind, unsigned pos);
-  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
-
-  /// Removes the specified identifier from the system.
-  void removeId(unsigned pos);
-
-  void removeEquality(unsigned pos);
-  void removeInequality(unsigned pos);
-
-  /// Remove the (in)equalities at positions [start, end).
-  void removeEqualityRange(unsigned start, unsigned end);
-  void removeInequalityRange(unsigned start, unsigned end);
-
   /// Sets the `values.size()` identifiers starting at `po`s to the specified
   /// values and removes them.
   void setAndEliminate(unsigned pos, ArrayRef<int64_t> values);
 
   /// 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);
public:
   /// 2) this   = {s0 + 5 <= d0 <= s0 + 20},
   ///    other  = {s0 + 1 <= d0 <= s0 + 9},
   ///    output = {s0 + 1 <= d0 <= s0 + 20}
   /// 3) this   = {0 <= d0 <= 5, 1 <= d1 <= 9}
   ///    other  = {2 <= d0 <= 6, 5 <= d1 <= 15},
   ///    output = {0 <= d0 <= 6, 1 <= d1 <= 15}
   LogicalResult unionBoundingBox(const FlatAffineConstraints &other);
 
-  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;
-  }
-
   /// Replaces the contents of this FlatAffineConstraints with `other`.
   virtual void clearAndCopyFrom(const FlatAffineConstraints &other);
 
   /// Returns the smallest known constant bound for the extent of the specified
   /// identifier (pos^th), i.e., the smallest known constant that is greater
   /// than or equal to 'exclusive upper bound' - 'lower bound' of the
   /// identifier. This constant bound is guaranteed to be non-negative. Returns
   /// None if it's not a constant. This method employs trivial (low complexity /
public:
   /// one set has a known division representation, then the `i^th` local
   /// variable in the other set either has the same division representation or
   /// no known division representation.
   ///
   /// The number of dimensions and symbol ids in `this` and `other` should
   /// match.
   void mergeLocalIds(FlatAffineConstraints &other);
 
-  /// Removes all equalities and inequalities.
-  void clearConstraints();
-
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
-  /// Return the index at which the specified kind of id starts.
-  unsigned getIdKindOffset(IdKind kind) const;
-
-  /// Assert that `value` is at most the number of ids of the specified kind.
-  void assertAtMostNumIdKind(unsigned value, IdKind kind) const;
-
   /// Returns false if the fields corresponding to various identifier counts, or
   /// equality/inequality buffer sizes aren't consistent; true otherwise. This
   /// is meant to be used within an assert internally.
   virtual bool hasConsistentState() const;
 
   /// Checks all rows of equality/inequality constraints for trivial
   /// contradictions (for example: 1 == 0, 0 >= 1), which may have surfaced
   /// after elimination. Returns true if an invalid constraint is found;
protected:
   /// i.e.,
   ///  64*i - 100 >= 0  =>  64*i - 128 >= 0 (since 'i' is an integer). This is a
   /// fast method (linear in the number of coefficients).
   void gcdTightenInequalities();
 
   /// Normalized each constraints by the GCD of its coefficients.
   void normalizeConstraintsByGCD();
 
-  /// 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);
-
-  /// 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;
-
   /// 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 FlatAffineConstraints.
   // 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

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

+//===- IntegerPolyhedron.h - MLIR IntegerPolyhedron 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A class to represent an integer polyhedron.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_INTEGERPOLYHEDRON_H
+#define MLIR_ANALYSIS_PRESBURGER_INTEGERPOLYHEDRON_H
+
+#include "mlir/Analysis/Presburger/Matrix.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:
+///
+/// Inequality: c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n >= 0
+/// Equality  : c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n == 0
+///
+/// where c_0, c_1, ..., c_n are integers.
+///
+/// Such a set corresponds to the set of integer points lying in a convex
+/// polyhedron. For example, consider the set: (x, y) : (1 <= x <= 7, x = 2y).
+/// This set contains the points (2, 1), (4, 2), and (6, 3).
+///
+/// The integer-valued variables are distinguished into 3 types of:
+///
+/// Dimension: Ordinary variables over which the set is represented.
+///
+/// Symbol: Symbol variables correspond to fixed but unknown values.
+/// Mathematically, an integer polyhedron with symbolic variables is like a
+/// family of integer polyhedra indexed by the symbolic variables.
+///
+/// Local: Local variables correspond to existentially quantified variables. For
+/// 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 {
+public:
+  /// 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);
+  }
+
+  /// 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) {}
+
+  virtual ~IntegerPolyhedron() = default;
+
+  // Clones this object.
+  std::unique_ptr<IntegerPolyhedron> clone() const;
+
+  /// Clears any existing data and reserves memory for the specified
+  /// constraints.
+  virtual void reset(unsigned numReservedInequalities,
+                     unsigned numReservedEqualities, unsigned numReservedCols,
+                     unsigned numDims, unsigned numSymbols,
+                     unsigned numLocals = 0);
+
+  void reset(unsigned numDims = 0, unsigned numSymbols = 0,
+             unsigned numLocals = 0);
+
+  /// Appends constraints from `other` into `this`. This is equivalent to an
+  /// intersection with no simplification of any sort attempted.
+  void append(const IntegerPolyhedron &other);
+
+  /// Returns the value at the specified equality row and column.
+  inline int64_t atEq(unsigned i, unsigned j) const { return equalities(i, j); }
+  inline int64_t &atEq(unsigned i, unsigned j) { return equalities(i, j); }
+
+  /// Returns the value at the specified inequality row and column.
+  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 getNumEqualities() const { return equalities.getNumRows(); }
+
+  inline unsigned getNumInequalities() const {
+    return inequalities.getNumRows();
+  }
+
+  inline unsigned getNumReservedEqualities() const {
+    return equalities.getNumReservedRows();
+  }
+
+  inline unsigned getNumReservedInequalities() const {
+    return inequalities.getNumReservedRows();
+  }
+
+  inline ArrayRef<int64_t> getEquality(unsigned idx) const {
+    return equalities.getRow(idx);
+  }
+
+  inline ArrayRef<int64_t> getInequality(unsigned idx) const {
+    return inequalities.getRow(idx);
+  }
+
+  /// 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);
+
+  /// 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);
+
+  /// Adds an inequality (>= 0) from the coefficients specified in `inEq`.
+  void addInequality(ArrayRef<int64_t> inEq);
+  /// Adds an equality from the coefficients specified in `eq`.
+  void addEquality(ArrayRef<int64_t> eq);
+
+  /// Removes identifiers of the specified kind with the specified pos (or
+  /// within the specified range) from the system. The specified location is
+  /// relative to the first identifier of the specified kind.
+  void removeId(IdKind kind, unsigned pos);
+  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
+
+  /// Removes the specified identifier from the system.
+  void removeId(unsigned pos);
+
+  void removeEquality(unsigned pos);
+  void removeInequality(unsigned pos);
+
+  /// Remove the (in)equalities at positions [start, end).
+  void removeEqualityRange(unsigned start, unsigned end);
+  void removeInequalityRange(unsigned start, unsigned end);
+
+  /// Swap the posA^th identifier with the posB^th identifier.
+  virtual void swapId(unsigned posA, unsigned posB);
+
+  /// Removes all equalities and inequalities.
+  void clearConstraints();
+
+protected:
+  /// Return the index at which the specified kind of id starts.
+  unsigned getIdKindOffset(IdKind kind) const;
+
+  /// Assert that `value` is at most the number of ids of the specified kind.
+  void assertAtMostNumIdKind(unsigned value, IdKind kind) 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);
+
+  /// 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/lib/Analysis/AffineStructures.cpp

 
 std::unique_ptr<FlatAffineValueConstraints>
 FlatAffineValueConstraints::clone() const {
   return std::make_unique<FlatAffineValueConstraints>(*this);
 }
 
 // Construct from an IntegerSet.
 FlatAffineConstraints::FlatAffineConstraints(IntegerSet set)
-    : numIds(set.getNumDims() + set.getNumSymbols()), numDims(set.getNumDims()),
-      numSymbols(set.getNumSymbols()),
-      equalities(0, numIds + 1, set.getNumEqualities(), numIds + 1),
-      inequalities(0, numIds + 1, set.getNumInequalities(), numIds + 1) {
+    : IntegerPolyhedron(set.getNumInequalities(), set.getNumEqualities(),
+                        set.getNumDims() + set.getNumSymbols() + 1,
+                        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;
   }
   assert(flatExprs.size() == set.getNumConstraints());
for (int ivIdx = 0, e = nIvs; ivIdx < e; ++ivIdx) {
     AffineMap ub = AffineMap::get(/*dimCount=*/3 * nIvs, /*symbolCount=*/0,
                                   getAffineDimExpr(ubsStart + ivIdx, ctx));
     if (failed(res.addBound(BoundType::UB, ivIdx, ub)))
       llvm_unreachable("Unexpected FlatAffineValueConstraints creation error");
   }
   return res;
 }
 
-void FlatAffineConstraints::reset(unsigned numReservedInequalities,
-                                  unsigned numReservedEqualities,
-                                  unsigned newNumReservedCols,
-                                  unsigned newNumDims, unsigned newNumSymbols,
-                                  unsigned newNumLocals) {
-  assert(newNumReservedCols >= newNumDims + newNumSymbols + newNumLocals + 1 &&
-         "minimum 1 column");
-  *this = FlatAffineConstraints(numReservedInequalities, numReservedEqualities,
-                                newNumReservedCols, newNumDims, newNumSymbols,
-                                newNumLocals);
-}
-
 void FlatAffineValueConstraints::reset(unsigned numReservedInequalities,
                                        unsigned numReservedEqualities,
                                        unsigned newNumReservedCols,
                                        unsigned newNumDims,
                                        unsigned newNumSymbols,
                                        unsigned newNumLocals) {
   reset(numReservedInequalities, numReservedEqualities, newNumReservedCols,
         newNumDims, newNumSymbols, newNumLocals, /*valArgs=*/{});
void FlatAffineValueConstraints::reset(
   if (!valArgs.empty())
     newVals.assign(valArgs.begin(), valArgs.end());
 
   *this = FlatAffineValueConstraints(
       numReservedInequalities, numReservedEqualities, newNumReservedCols,
       newNumDims, newNumSymbols, newNumLocals, newVals);
 }
 
-void FlatAffineConstraints::reset(unsigned newNumDims, unsigned newNumSymbols,
-                                  unsigned newNumLocals) {
-  reset(0, 0, newNumDims + newNumSymbols + newNumLocals + 1, newNumDims,
-        newNumSymbols, newNumLocals);
-}
-
 void FlatAffineValueConstraints::reset(unsigned newNumDims,
                                        unsigned newNumSymbols,
                                        unsigned newNumLocals,
                                        ArrayRef<Value> valArgs) {
   reset(0, 0, newNumDims + newNumSymbols + newNumLocals + 1, newNumDims,
         newNumSymbols, newNumLocals, valArgs);
 }
 
-void FlatAffineConstraints::append(const FlatAffineConstraints &other) {
-  assert(other.getNumCols() == getNumCols());
-  assert(other.getNumDimIds() == getNumDimIds());
-  assert(other.getNumSymbolIds() == getNumSymbolIds());
-
-  inequalities.reserveRows(inequalities.getNumRows() +
-                           other.getNumInequalities());
-  equalities.reserveRows(equalities.getNumRows() + other.getNumEqualities());
-
-  for (unsigned r = 0, e = other.getNumInequalities(); r < e; r++) {
-    addInequality(other.getInequality(r));
-  }
-  for (unsigned r = 0, e = other.getNumEqualities(); r < e; r++) {
-    addEquality(other.getEquality(r));
-  }
-}
-
-unsigned FlatAffineConstraints::appendDimId(unsigned num) {
-  unsigned pos = getNumDimIds();
-  insertId(IdKind::Dimension, pos, num);
-  return pos;
-}
-
 unsigned FlatAffineValueConstraints::appendDimId(ValueRange vals) {
   unsigned pos = getNumDimIds();
   insertId(IdKind::Dimension, pos, vals);
   return pos;
 }
 
-unsigned FlatAffineConstraints::appendSymbolId(unsigned num) {
-  unsigned pos = getNumSymbolIds();
-  insertId(IdKind::Symbol, pos, num);
-  return pos;
-}
-
 unsigned FlatAffineValueConstraints::appendSymbolId(ValueRange vals) {
   unsigned pos = getNumSymbolIds();
   insertId(IdKind::Symbol, pos, vals);
   return pos;
 }
 
-unsigned FlatAffineConstraints::appendLocalId(unsigned num) {
-  unsigned pos = getNumLocalIds();
-  insertId(IdKind::Local, pos, num);
-  return pos;
-}
-
-unsigned FlatAffineConstraints::insertDimId(unsigned pos, unsigned num) {
-  return insertId(IdKind::Dimension, pos, num);
-}
-
 unsigned FlatAffineValueConstraints::insertDimId(unsigned pos,
                                                  ValueRange vals) {
   return insertId(IdKind::Dimension, pos, vals);
 }
 
-unsigned FlatAffineConstraints::insertSymbolId(unsigned pos, unsigned num) {
-  return insertId(IdKind::Symbol, pos, num);
-}
-
 unsigned FlatAffineValueConstraints::insertSymbolId(unsigned pos,
                                                     ValueRange vals) {
   return insertId(IdKind::Symbol, pos, vals);
 }
 
-unsigned FlatAffineConstraints::insertLocalId(unsigned pos, unsigned num) {
-  return insertId(IdKind::Local, pos, num);
-}
-
-unsigned FlatAffineConstraints::insertId(IdKind kind, unsigned pos,
-                                         unsigned num) {
-  assertAtMostNumIdKind(pos, 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;
-}
-
-void FlatAffineConstraints::assertAtMostNumIdKind(unsigned val,
-                                                  IdKind kind) const {
-  if (kind == IdKind::Dimension)
-    assert(val <= getNumDimIds());
-  else if (kind == IdKind::Symbol)
-    assert(val <= getNumSymbolIds());
-  else if (kind == IdKind::Local)
-    assert(val <= getNumLocalIds());
-  else
-    llvm_unreachable("IdKind expected to be Dimension, Symbol or Local!");
-}
-
-unsigned FlatAffineConstraints::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 FlatAffineValueConstraints::insertId(IdKind kind, unsigned pos,
                                               unsigned num) {
   unsigned absolutePos = FlatAffineConstraints::insertId(kind, pos, num);
   values.insert(values.begin() + absolutePos, num, None);
   assert(values.size() == getNumIds());
   return absolutePos;
 }
 

 }
 
 bool FlatAffineValueConstraints::hasValues() const {
   return llvm::find_if(values, [](Optional<Value> id) {
            return id.hasValue();
          }) != values.end();
 }
 
-void FlatAffineConstraints::removeId(IdKind kind, unsigned pos) {
-  removeIdRange(kind, pos, pos + 1);
-}
-
-void FlatAffineConstraints::removeIdRange(IdKind kind, unsigned idStart,
-                                          unsigned idLimit) {
-  assertAtMostNumIdKind(idLimit, kind);
-  removeIdRange(getIdKindOffset(kind) + idStart,
-                getIdKindOffset(kind) + idLimit);
-}
-
 /// Checks if two constraint systems are in the same space, i.e., if they are
 /// associated with the same set of identifiers, appearing in the same order.
 static bool areIdsAligned(const FlatAffineValueConstraints &a,
                           const FlatAffineValueConstraints &b) {
   return a.getNumDimIds() == b.getNumDimIds() &&
          a.getNumSymbolIds() == b.getNumSymbolIds() &&
          a.getNumIds() == b.getNumIds() &&
          a.getMaybeValues().equals(b.getMaybeValues());
if (j >= elimColStart && j < pivotCol)
       continue;
     int64_t v = pivotMultiplier * constraints->atEq(pivotRow, j) +
                 rowMultiplier * at(rowIdx, j);
     isEq ? constraints->atEq(rowIdx, j) = v
          : constraints->atIneq(rowIdx, j) = v;
   }
 }
 
-void FlatAffineConstraints::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 < numIds && "invalid idStart position");
-
-  // TODO: Make 'removeIdRange' a lambda called from here.
-  // Remove eliminated identifiers from the constraints..
-  equalities.removeColumns(idStart, idLimit - idStart);
-  inequalities.removeColumns(idStart, idLimit - idStart);
-
-  // Update members numDims, numSymbols and numIds.
-  unsigned numDimsEliminated = 0;
-  unsigned numLocalsEliminated = 0;
-  unsigned numColsEliminated = idLimit - idStart;
-  if (idStart < numDims) {
-    numDimsEliminated = std::min(numDims, idLimit) - idStart;
-  }
-  // Check how many local id's were removed. Note that our identifier order is
-  // [dims, symbols, locals]. Local id start at position numDims + numSymbols.
-  if (idLimit > numDims + numSymbols) {
-    numLocalsEliminated = std::min(
-        idLimit - std::max(idStart, numDims + numSymbols), getNumLocalIds());
-  }
-  unsigned numSymbolsEliminated =
-      numColsEliminated - numDimsEliminated - numLocalsEliminated;
-
-  numDims -= numDimsEliminated;
-  numSymbols -= numSymbolsEliminated;
-  numIds = numIds - numColsEliminated;
-}
-
 void FlatAffineValueConstraints::removeIdRange(unsigned idStart,
                                                unsigned idLimit) {
   FlatAffineConstraints::removeIdRange(idStart, idLimit);
   values.erase(values.begin() + idStart, values.begin() + idLimit);
 }
 
 /// Returns the position of the identifier that has the minimum <number of lower
 /// bounds> times <number of upper bounds> from the specified range of
if (lbMap && lbMap.getNumResults() != 0 && ubMap &&
       auto loop = getForInductionVarOwner(values[i]);
       if (failed(this->addAffineForOpDomain(loop)))
         return failure();
     }
   }
   return success();
 }
 
-void FlatAffineConstraints::addEquality(ArrayRef<int64_t> eq) {
-  assert(eq.size() == getNumCols());
-  unsigned row = equalities.appendExtraRow();
-  for (unsigned i = 0, e = eq.size(); i < e; ++i)
-    equalities(row, i) = eq[i];
-}
-
-void FlatAffineConstraints::addInequality(ArrayRef<int64_t> inEq) {
-  assert(inEq.size() == getNumCols());
-  unsigned row = inequalities.appendExtraRow();
-  for (unsigned i = 0, e = inEq.size(); i < e; ++i)
-    inequalities(row, i) = inEq[i];
-}
-
 void FlatAffineConstraints::addBound(BoundType type, unsigned pos,
                                      int64_t value) {
   assert(pos < getNumCols());
   if (type == BoundType::EQ) {
     unsigned row = equalities.appendExtraRow();
     equalities(row, pos) = 1;
     equalities(row, getNumCols() - 1) = -value;
   } else {

 }
 
 bool FlatAffineValueConstraints::containsId(Value val) const {
   return llvm::any_of(values, [&](const Optional<Value> &mayBeId) {
     return mayBeId.hasValue() && mayBeId.getValue() == val;
   });
 }
 
-void FlatAffineConstraints::swapId(unsigned posA, unsigned posB) {
-  assert(posA < getNumIds() && "invalid position A");
-  assert(posB < getNumIds() && "invalid position B");
-
-  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));
-}
-
 void FlatAffineValueConstraints::swapId(unsigned posA, unsigned posB) {
   FlatAffineConstraints::swapId(posA, posB);
   std::swap(values[posA], values[posB]);
 }
 
 void FlatAffineConstraints::setDimSymbolSeparation(unsigned newSymbolCount) {
   assert(newSymbolCount <= numDims + numSymbols &&
          "invalid separation position");
   numDims = numDims + numSymbols - newSymbolCount;
   numSymbols = newSymbolCount;
 }
 
 void FlatAffineValueConstraints::addBound(BoundType type, Value val,
                                           int64_t value) {
   unsigned pos;
   if (!findId(val, &pos))
     // This is a pre-condition for this method.
     assert(0 && "id not found");
   addBound(type, pos, value);
 }
 
-void FlatAffineConstraints::removeEquality(unsigned pos) {
-  equalities.removeRow(pos);
-}
-
-void FlatAffineConstraints::removeInequality(unsigned pos) {
-  inequalities.removeRow(pos);
-}
-
-void FlatAffineConstraints::removeEqualityRange(unsigned begin, unsigned end) {
-  if (begin >= end)
-    return;
-  equalities.removeRows(begin, end - begin);
-}
-
-void FlatAffineConstraints::removeInequalityRange(unsigned begin,
-                                                  unsigned end) {
-  if (begin >= end)
-    return;
-  inequalities.removeRows(begin, end - begin);
-}
-
 /// 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 FlatAffineConstraints &cst,
                                   unsigned pos, bool symbolic = false) {
   assert(pos < cst.getNumIds() && "invalid position");
if (auto *otherValueSet =
     *this = *otherValueSet;
   } else {
     *static_cast<FlatAffineConstraints *>(this) = other;
     values.clear();
     values.resize(numIds, None);
   }
 }
 
-void FlatAffineConstraints::removeId(unsigned pos) {
-  removeIdRange(pos, pos + 1);
-}
-
 static std::pair<unsigned, unsigned>
 getNewNumDimsSymbols(unsigned pos, const FlatAffineConstraints &cst) {
   unsigned numDims = cst.getNumDimIds();
   unsigned numSymbols = cst.getNumSymbolIds();
   unsigned newNumDims, newNumSymbols;
   if (pos < numDims) {
     newNumDims = numDims - 1;
     newNumSymbols = numSymbols;

 void FlatAffineValueConstraints::projectOut(Value val) {
   unsigned pos;
   bool ret = findId(val, &pos);
   assert(ret);
   (void)ret;
   fourierMotzkinEliminate(pos);
 }
 
-void FlatAffineConstraints::clearConstraints() {
-  equalities.resizeVertically(0);
-  inequalities.resizeVertically(0);
-}
-
 namespace {
 
 enum BoundCmpResult { Greater, Less, Equal, Unknown };
 
 /// Compares two affine bounds whose coefficients are provided in 'first' and
 /// 'second'. The last coefficient is the constant term.
 static BoundCmpResult compareBounds(ArrayRef<int64_t> a, ArrayRef<int64_t> b) {
   assert(a.size() == b.size());

mlir/lib/Analysis/Presburger/CMakeLists.txt

 add_mlir_library(MLIRPresburger
-  Simplex.cpp
+  IntegerPolyhedron.cpp
   Matrix.cpp
+  Simplex.cpp
 
   DEPENDS
   MLIRBuiltinLocationAttributesIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   )

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

+//===- IntegerPolyhedron.cpp - MLIR IntegerPolyhedron 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A class to represent an integer polyhedron.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/Presburger/IntegerPolyhedron.h"
+
+using namespace mlir;
+
+std::unique_ptr<IntegerPolyhedron> IntegerPolyhedron::clone() const {
+  return std::make_unique<IntegerPolyhedron>(*this);
+}
+
+void IntegerPolyhedron::reset(unsigned numReservedInequalities,
+                              unsigned numReservedEqualities,
+                              unsigned newNumReservedCols, unsigned newNumDims,
+                              unsigned newNumSymbols, unsigned newNumLocals) {
+  assert(newNumReservedCols >= newNumDims + newNumSymbols + newNumLocals + 1 &&
+         "minimum 1 column");
+  *this = IntegerPolyhedron(numReservedInequalities, numReservedEqualities,
+                            newNumReservedCols, newNumDims, newNumSymbols,
+                            newNumLocals);
+}
+
+void IntegerPolyhedron::reset(unsigned newNumDims, unsigned newNumSymbols,
+                              unsigned newNumLocals) {
+  reset(0, 0, newNumDims + newNumSymbols + newNumLocals + 1, newNumDims,
+        newNumSymbols, newNumLocals);
+}
+
+void IntegerPolyhedron::append(const IntegerPolyhedron &other) {
+  assert(other.getNumCols() == getNumCols());
+  assert(other.getNumDimIds() == getNumDimIds());
+  assert(other.getNumSymbolIds() == getNumSymbolIds());
+
+  inequalities.reserveRows(inequalities.getNumRows() +
+                           other.getNumInequalities());
+  equalities.reserveRows(equalities.getNumRows() + other.getNumEqualities());
+
+  for (unsigned r = 0, e = other.getNumInequalities(); r < e; r++) {
+    addInequality(other.getInequality(r));
+  }
+  for (unsigned r = 0, e = other.getNumEqualities(); r < e; r++) {
+    addEquality(other.getEquality(r));
+  }
+}
+
+unsigned IntegerPolyhedron::insertDimId(unsigned pos, unsigned num) {
+  return insertId(IdKind::Dimension, pos, num);
+}
+
+unsigned IntegerPolyhedron::insertSymbolId(unsigned pos, unsigned num) {
+  return insertId(IdKind::Symbol, pos, num);
+}
+
+unsigned IntegerPolyhedron::insertLocalId(unsigned pos, unsigned num) {
+  return insertId(IdKind::Local, pos, num);
+}
+
+unsigned IntegerPolyhedron::insertId(IdKind kind, unsigned pos, unsigned num) {
+  assertAtMostNumIdKind(pos, 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;
+}
+
+unsigned IntegerPolyhedron::appendDimId(unsigned num) {
+  unsigned pos = getNumDimIds();
+  insertId(IdKind::Dimension, pos, num);
+  return pos;
+}
+
+unsigned IntegerPolyhedron::appendSymbolId(unsigned num) {
+  unsigned pos = getNumSymbolIds();
+  insertId(IdKind::Symbol, pos, num);
+  return pos;
+}
+
+unsigned IntegerPolyhedron::appendLocalId(unsigned num) {
+  unsigned pos = getNumLocalIds();
+  insertId(IdKind::Local, pos, num);
+  return pos;
+}
+
+void IntegerPolyhedron::addEquality(ArrayRef<int64_t> eq) {
+  assert(eq.size() == getNumCols());
+  unsigned row = equalities.appendExtraRow();
+  for (unsigned i = 0, e = eq.size(); i < e; ++i)
+    equalities(row, i) = eq[i];
+}
+
+void IntegerPolyhedron::addInequality(ArrayRef<int64_t> inEq) {
+  assert(inEq.size() == getNumCols());
+  unsigned row = inequalities.appendExtraRow();
+  for (unsigned i = 0, e = inEq.size(); i < e; ++i)
+    inequalities(row, i) = inEq[i];
+}
+
+void IntegerPolyhedron::removeId(IdKind kind, unsigned pos) {
+  removeIdRange(kind, pos, pos + 1);
+}
+
+void IntegerPolyhedron::removeId(unsigned pos) { removeIdRange(pos, pos + 1); }
+
+void IntegerPolyhedron::removeIdRange(IdKind kind, unsigned idStart,
+                                      unsigned idLimit) {
+  assertAtMostNumIdKind(idLimit, 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 < numIds && "invalid idStart position");
+
+  // TODO: Make 'removeIdRange' a lambda called from here.
+  // Remove eliminated identifiers from the constraints..
+  equalities.removeColumns(idStart, idLimit - idStart);
+  inequalities.removeColumns(idStart, idLimit - idStart);
+
+  // Update members numDims, numSymbols and numIds.
+  unsigned numDimsEliminated = 0;
+  unsigned numLocalsEliminated = 0;
+  unsigned numColsEliminated = idLimit - idStart;
+  if (idStart < numDims) {
+    numDimsEliminated = std::min(numDims, idLimit) - idStart;
+  }
+  // Check how many local id's were removed. Note that our identifier order is
+  // [dims, symbols, locals]. Local id start at position numDims + numSymbols.
+  if (idLimit > numDims + numSymbols) {
+    numLocalsEliminated = std::min(
+        idLimit - std::max(idStart, numDims + numSymbols), getNumLocalIds());
+  }
+  unsigned numSymbolsEliminated =
+      numColsEliminated - numDimsEliminated - numLocalsEliminated;
+
+  numDims -= numDimsEliminated;
+  numSymbols -= numSymbolsEliminated;
+  numIds = numIds - numColsEliminated;
+}
+
+void IntegerPolyhedron::removeEquality(unsigned pos) {
+  equalities.removeRow(pos);
+}
+
+void IntegerPolyhedron::removeInequality(unsigned pos) {
+  inequalities.removeRow(pos);
+}
+
+void IntegerPolyhedron::removeEqualityRange(unsigned begin, unsigned end) {
+  if (begin >= end)
+    return;
+  equalities.removeRows(begin, end - begin);
+}
+
+void IntegerPolyhedron::removeInequalityRange(unsigned begin, unsigned end) {
+  if (begin >= end)
+    return;
+  inequalities.removeRows(begin, end - begin);
+}
+
+void IntegerPolyhedron::swapId(unsigned posA, unsigned posB) {
+  assert(posA < getNumIds() && "invalid position A");
+  assert(posB < getNumIds() && "invalid position B");
+
+  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!");
+}
+
+void IntegerPolyhedron::assertAtMostNumIdKind(unsigned val, IdKind kind) const {
+  if (kind == IdKind::Dimension)
+    assert(val <= getNumDimIds());
+  else if (kind == IdKind::Symbol)
+    assert(val <= getNumSymbolIds());
+  else if (kind == IdKind::Local)
+    assert(val <= getNumLocalIds());
+  else
+    llvm_unreachable("IdKind expected to be Dimension, Symbol or Local!");
+}
+
+void IntegerPolyhedron::clearConstraints() {
+  equalities.resizeVertically(0);
+  inequalities.resizeVertically(0);
+}

mlir/unittests/Analysis/AffineStructuresTest.cpp

TEST(FlatAffineConstraintsTest, addConstantLowerBound) {
   EXPECT_EQ(fac.atIneq(0, 2), -1);
 
   fac.addBound(FlatAffineConstraints::LB, {1, 2, 3}, 1);
   EXPECT_EQ(fac.atIneq(1, 0), 1);
   EXPECT_EQ(fac.atIneq(1, 1), 2);
   EXPECT_EQ(fac.atIneq(1, 2), 2);
 }
 
-TEST(FlatAffineConstraintsTest, removeInequality) {
-  FlatAffineConstraints fac =
-      makeFACFromConstraints(1, {{0, 0}, {1, 1}, {2, 2}, {3, 3}, {4, 4}}, {});
-
-  fac.removeInequalityRange(0, 0);
-  EXPECT_EQ(fac.getNumInequalities(), 5u);
-
-  fac.removeInequalityRange(1, 3);
-  EXPECT_EQ(fac.getNumInequalities(), 3u);
-  EXPECT_THAT(fac.getInequality(0), ElementsAre(0, 0));
-  EXPECT_THAT(fac.getInequality(1), ElementsAre(3, 3));
-  EXPECT_THAT(fac.getInequality(2), ElementsAre(4, 4));
-
-  fac.removeInequality(1);
-  EXPECT_EQ(fac.getNumInequalities(), 2u);
-  EXPECT_THAT(fac.getInequality(0), ElementsAre(0, 0));
-  EXPECT_THAT(fac.getInequality(1), ElementsAre(4, 4));
-}
-
-TEST(FlatAffineConstraintsTest, removeEquality) {
-  FlatAffineConstraints fac =
-      makeFACFromConstraints(1, {}, {{0, 0}, {1, 1}, {2, 2}, {3, 3}, {4, 4}});
-
-  fac.removeEqualityRange(0, 0);
-  EXPECT_EQ(fac.getNumEqualities(), 5u);
-
-  fac.removeEqualityRange(1, 3);
-  EXPECT_EQ(fac.getNumEqualities(), 3u);
-  EXPECT_THAT(fac.getEquality(0), ElementsAre(0, 0));
-  EXPECT_THAT(fac.getEquality(1), ElementsAre(3, 3));
-  EXPECT_THAT(fac.getEquality(2), ElementsAre(4, 4));
-
-  fac.removeEquality(1);
-  EXPECT_EQ(fac.getNumEqualities(), 2u);
-  EXPECT_THAT(fac.getEquality(0), ElementsAre(0, 0));
-  EXPECT_THAT(fac.getEquality(1), ElementsAre(4, 4));
-}
-
-TEST(FlatAffineConstraintsTest, clearConstraints) {
-  FlatAffineConstraints fac = makeFACFromConstraints(1, {}, {});
-
-  fac.addInequality({1, 0});
-  EXPECT_EQ(fac.atIneq(0, 0), 1);
-  EXPECT_EQ(fac.atIneq(0, 1), 0);
-
-  fac.clearConstraints();
-
-  fac.addInequality({1, 0});
-  EXPECT_EQ(fac.atIneq(0, 0), 1);
-  EXPECT_EQ(fac.atIneq(0, 1), 0);
-}
-
 /// Check if the expected division representation of local variables matches the
 /// computed representation. The expected division representation is given as
 /// a vector of expressions set in `expectedDividends` and the corressponding
 /// denominator in `expectedDenominators`. The `denominators` and `dividends`
 /// obtained through `getLocalRepr` function is verified against the
 /// `expectedDenominators` and `expectedDividends` respectively.
 static void checkDivisionRepresentation(
     FlatAffineConstraints &fac,
std::vector<SmallVector<int64_t, 8>> divisions = {
       {0, 1, -123, 2, 0, -4, 0, 10}};
 
   SmallVector<unsigned, 8> denoms = {3, 5, 3};
 
   // Check if floordivs which may depend on other floordivs can be computed.
   checkDivisionRepresentation(fac, divisions, denoms);
 }
 
-TEST(FlatAffineConstraintsTest, removeIdRange) {
-  FlatAffineConstraints fac(3, 2, 1);
-
-  fac.addInequality({10, 11, 12, 20, 21, 30, 40});
-  fac.removeId(FlatAffineConstraints::IdKind::Symbol, 1);
-  EXPECT_THAT(fac.getInequality(0),
-              testing::ElementsAre(10, 11, 12, 20, 30, 40));
-
-  fac.removeIdRange(FlatAffineConstraints::IdKind::Dimension, 0, 2);
-  EXPECT_THAT(fac.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
-
-  fac.removeIdRange(FlatAffineConstraints::IdKind::Local, 1, 1);
-  EXPECT_THAT(fac.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
-
-  fac.removeIdRange(FlatAffineConstraints::IdKind::Local, 0, 1);
-  EXPECT_THAT(fac.getInequality(0), testing::ElementsAre(12, 20, 40));
-}
-
 TEST(FlatAffineConstraintsTest, simplifyLocalsTest) {
   // (x) : (exists y: 2x + y = 1 and y = 2).
   FlatAffineConstraints fac(1, 0, 1);
   fac.addEquality({2, 1, -1});
   fac.addEquality({0, 1, -2});
 
   EXPECT_TRUE(fac.isEmpty());
 

mlir/unittests/Analysis/Presburger/CMakeLists.txt

 add_mlir_unittest(MLIRPresburgerTests
+  IntegerPolyhedronTest.cpp
   MatrixTest.cpp
   SimplexTest.cpp
   ../AffineStructuresParser.cpp
 )
 
 target_link_libraries(MLIRPresburgerTests
   PRIVATE MLIRPresburger
   MLIRLoopAnalysis
   MLIRParser
   )

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

+//===- IntegerPolyhedron.cpp - Tests for IntegerPolyhedron 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/IntegerPolyhedron.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace mlir {
+
+using testing::ElementsAre;
+
+/// 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);
+  for (const auto &eq : eqs)
+    set.addEquality(eq);
+  for (const auto &ineq : ineqs)
+    set.addInequality(ineq);
+  return set;
+}
+
+TEST(IntegerPolyhedronTest, removeInequality) {
+  IntegerPolyhedron set =
+      makeSetFromConstraints(1, {{0, 0}, {1, 1}, {2, 2}, {3, 3}, {4, 4}}, {});
+
+  set.removeInequalityRange(0, 0);
+  EXPECT_EQ(set.getNumInequalities(), 5u);
+
+  set.removeInequalityRange(1, 3);
+  EXPECT_EQ(set.getNumInequalities(), 3u);
+  EXPECT_THAT(set.getInequality(0), ElementsAre(0, 0));
+  EXPECT_THAT(set.getInequality(1), ElementsAre(3, 3));
+  EXPECT_THAT(set.getInequality(2), ElementsAre(4, 4));
+
+  set.removeInequality(1);
+  EXPECT_EQ(set.getNumInequalities(), 2u);
+  EXPECT_THAT(set.getInequality(0), ElementsAre(0, 0));
+  EXPECT_THAT(set.getInequality(1), ElementsAre(4, 4));
+}
+
+TEST(IntegerPolyhedronTest, removeEquality) {
+  IntegerPolyhedron set =
+      makeSetFromConstraints(1, {}, {{0, 0}, {1, 1}, {2, 2}, {3, 3}, {4, 4}});
+
+  set.removeEqualityRange(0, 0);
+  EXPECT_EQ(set.getNumEqualities(), 5u);
+
+  set.removeEqualityRange(1, 3);
+  EXPECT_EQ(set.getNumEqualities(), 3u);
+  EXPECT_THAT(set.getEquality(0), ElementsAre(0, 0));
+  EXPECT_THAT(set.getEquality(1), ElementsAre(3, 3));
+  EXPECT_THAT(set.getEquality(2), ElementsAre(4, 4));
+
+  set.removeEquality(1);
+  EXPECT_EQ(set.getNumEqualities(), 2u);
+  EXPECT_THAT(set.getEquality(0), ElementsAre(0, 0));
+  EXPECT_THAT(set.getEquality(1), ElementsAre(4, 4));
+}
+
+TEST(IntegerPolyhedronTest, clearConstraints) {
+  IntegerPolyhedron set = makeSetFromConstraints(1, {}, {});
+
+  set.addInequality({1, 0});
+  EXPECT_EQ(set.atIneq(0, 0), 1);
+  EXPECT_EQ(set.atIneq(0, 1), 0);
+
+  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);
+
+  set.addInequality({10, 11, 12, 20, 21, 30, 40});
+  set.removeId(IntegerPolyhedron::IdKind::Symbol, 1);
+  EXPECT_THAT(set.getInequality(0),
+              testing::ElementsAre(10, 11, 12, 20, 30, 40));
+
+  set.removeIdRange(IntegerPolyhedron::IdKind::Dimension, 0, 2);
+  EXPECT_THAT(set.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
+
+  set.removeIdRange(IntegerPolyhedron::IdKind::Local, 1, 1);
+  EXPECT_THAT(set.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
+
+  set.removeIdRange(IntegerPolyhedron::IdKind::Local, 0, 1);
+  EXPECT_THAT(set.getInequality(0), testing::ElementsAre(12, 20, 40));
+}
+
+} // namespace mlir