[MLIR][Presburger] IntegerPolyhedron: add support for symbolic integer lexmin

[arjunp]

Summary

Add support for computing the symbolic integer lexmin of a polyhedron.
This finds, for every assignment to the symbols, the lexicographically
minimum value attained by the dimensions. For example, the symbolic lexmin
of the set

(x, y)[a, b, c] : (a <= x, b <= x, x <= c)

can be written as

x = a if b <= a, a <= c
x = b if a <  b, b <= c

This also finds the set of assignments to the symbols that make the lexmin unbounded.

[arjunp]

Fix some issues I just noticed

[Groverkss]

Some initial comments which may change the implementation a lot (in case we agree).

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

618

I would prefer not treating the relation as a set while producing this lexmin. Can we instead also treat the domain variables as parameters and have something like:

(d0, d1) -> (r1, r2)[a, b, c] : ...

gives

Domain: (d0, d1)[a, b, c] : ...
Values of range: r1 = .., r2 = ..

...

I think this is a more general case of what the current implementation produces.

620–623

I find the documentation incomplete. What is symbolDomain? Does it only contain symbols? Does it have to belong to the same space as the relation?

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

111

175

[Groverkss]

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

618

Just to clarify, I wanted this to be moved to IntegerRelation as a general case. I'm not sure if these changes can be easily accommodated later.

[arjunp]

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

618

Yeah, that can be done "easily", i.e., without changes to the Simplex. I mostly put it in the set and not relation for now because I didn't want to think of exactly what interface would be exposed.

The only thing is, I think we might have to change MultiAffineFunction and PWMAFunction to use IntegerRelation and that's a pretty big find-replace change. These are interface changes that do not affect the underlying algorithm. You can tell SymbolicLexSimplex which variables are to be treated as "symbols" and which as "non-symbols". Currently, only the symbols in the space are treated as symbols and the rest as non-symbols, but this can easily be changed later.

[arjunp]

Remove unneeded overload

[arjunp]

Add some documentation and asserts to the SymbolicLexSimplex constructor

[arjunp]

Fix some documentation

[Groverkss]

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

618

Why would you need to change MAF and PWMA to IntegerRelation? Right now your domain of function will look something like:

()[a, b, c] : ...

After this change, it will be:

(d0, d1)[a, b, c] : ...

618

Can you also tell which IdKind the symbol was after finding it's lexmin? I think that would be needed to distinguish domain variables vs symbol variables.

[arjunp]

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

618

Right, if we go with defining the domain like that it'll be fine I guess. At the moment the symbols of the set are the dimensions of the function's domain.

618

The order of symbol IDs will be the same on output, so first the domain IDs, then the symbol, then the locals created by the symbolic lexmin over these. So we can always find out. Also, we could just put symbols as symbols and domain as dims in the space of the symbolDomain provided to SymbolicLexMin, so the resulting function's domain will have the same space

[Groverkss]

Another pass of comments.

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

618

Can we for this patch have the domain of PWMA treat symbols as symbols only instead of converting them to dimensions?

mlir/include/mlir/Analysis/Presburger/Simplex.h

219–221

Instead of passing these parameters, you could also pass a PresburgerSpace to this and extract this information from it.

581–582

I don't like this conversion to dim id. Why can't they just be symbols?

[arjunp]

mlir/include/mlir/Analysis/Presburger/Simplex.h

219–221

The "symbols" here are not necessarily the symbols of the input set. For example, if we want the output to be a function of the domains and symbols of the input, we would pass both domain and symbols as symbols here. In general, these may not even form a a continuous range, in which case we can consider either passing a BitVector denoting which vars are symbols, or indeed passing a PresburgerSpace and a list of IdKinds to be considered symbols. These things are easily changed later.

581–582

It seems more natural to me to think of the result as being a function with its domain corresponding to the symbols of the input set and its range corresponding to the non-symbols of the input set, as opposed to the result being a family of functions parametrized by the symbols of the input sets and with no domain.

[arjunp]

TODO -> TODO:

[Groverkss]

Another pass of comments.

mlir/include/mlir/Analysis/Presburger/Matrix.h

154–155

This is a bit confusing. Can you rephrase it to say that the length of elems should be equal to nColumns?

mlir/include/mlir/Analysis/Presburger/Simplex.h

394

Are you missing a full stop (.) here?

mlir/lib/Analysis/Presburger/Simplex.cpp

346–348

You can maybe make the divisor unsigned and add an assert to check if it positive.

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

1144

Remove this debug statement please.

[Groverkss]

mlir/include/mlir/Analysis/Presburger/Simplex.h

219–221

Ok it seems like we can do these things after this patch.

581–582

I guess it's fine for now and we can discuss this again later as it seems to be easy to change.

[arjunp]

Address comments

mlir/lib/Analysis/Presburger/Simplex.cpp

346–348

I added the assert in isRangeDivisbleBy since I think it's better for it to assert itself than make its users assert (users are typically going to have an int64_t)

[arjunp]

Remove use of auto

[arjunp]

fix expression in symbolic cut documentation

[Groverkss]

Another pass of comments.

mlir/include/mlir/Analysis/Presburger/Simplex.h

339

Should this be isSymbol?

397

mlir/lib/Analysis/Presburger/Simplex.cpp

34

The assert should be at top of function body I think.

346–348

Sorry if I was ambiguous in what I mean. I meant the same thing that you did here :P

389–412

Could you add some inline explanation (above each logical block of code maybe) of what is happening here?

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

1159–1162

Why std::cerr instead of llvm::errs?

[arjunp]

Address comments and add a bit more doc

[Groverkss]

LGTM.

[arjunp]

lint

[mehdi_amini]

Please include revert reason in the commit message when you revert in the future.

[arjunp]

In D122985#3428131, @mehdi_amini wrote:

Please include revert reason in the commit message when you revert in the future.

Okay, will do. Sorry about that!

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

protected:
 
   /// Coefficients of affine equalities (in == 0 form).
   Matrix equalities;
 
   /// Coefficients of affine inequalities (in >= 0 form).
   Matrix inequalities;
 };
 
+struct SymbolicLexMin;
 /// An IntegerPolyhedron is a PresburgerSpace subject to affine
 /// constraints. 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 and n is the total number of
public:
   // Clones this object.
   std::unique_ptr<IntegerPolyhedron> clone() 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.
   unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
+
+  /// Compute the symbolic integer lexmin of the polyhedron.
+  /// This finds, for every assignment to the symbols, the lexicographically
+  /// minimum value attained by the dimensions. For example, the symbolic lexmin
+  /// of the set
+  ///
+  /// (x, y)[a, b, c] : (a <= x, b <= x, x <= c)
+  ///
+  /// can be written as
+  ///
+  /// x = a if b <= a, a <= c
+  /// x = b if a <  b, b <= c
+  ///
+  /// This function is stored in the `lexmin` function in the result.
+  /// Some assignments to the symbols might make the set empty.
+  /// Such points are not part of the function's domain.
+  /// In the above example, this happens when max(a, b) > c.
+  ///
+  /// For some values of the symbols, the lexmin may be unbounded.
+  /// `SymbolicLexMin` stores these parts of the symbolic domain in a separate
+  /// `PresburgerSet`, `unboundedDomain`.
+  SymbolicLexMin findSymbolicIntegerLexMin() const;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_INTEGERRELATION_H

mlir/include/mlir/Analysis/Presburger/Matrix.h

public:
   /// is less expensive than increasing the number of columns beyond
   /// nReservedColumns.
   void resize(unsigned newNRows, unsigned newNColumns);
   void resizeHorizontally(unsigned newNColumns);
   void resizeVertically(unsigned newNRows);
 
   /// Add an extra row at the bottom of the matrix and return its position.
   unsigned appendExtraRow();
+  /// Same as above, but copy the given elements into the row. The length of
+  /// `elems` must be equal to the number of columns.
+  unsigned appendExtraRow(ArrayRef<int64_t> elems);
 
   /// Print the matrix.
   void print(raw_ostream &os) const;
   void dump() const;
 
   /// Return whether the Matrix is in a consistent state with all its
   /// invariants satisfied.
   bool hasConsistentState() const;

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

public:
   /// they lie in the same space, i.e. have the same dimensions, and their
   /// domains are identical and their outputs are equal on their domain.
   bool isEqual(const MultiAffineFunction &other) const;
 
   /// Get the value of the function at the specified point. If the point lies
   /// outside the domain, an empty optional is returned.
   Optional<SmallVector<int64_t, 8>> valueAt(ArrayRef<int64_t> point) const;
 
+  /// Truncate the output dimensions to the first `count` dimensions.
+  ///
+  /// TODO: refactor so that this can be accomplished through removeIdRange.
+  void truncateOutput(unsigned count);
+
   void print(raw_ostream &os) const;
   void dump() const;
 
 private:
   /// The function's output is a tuple of integers, with the ith element of the
   /// tuple defined by the affine expression given by the ith row of this output
   /// matrix.
   Matrix output;
public:
   /// point does not lie in the domain.
   Optional<SmallVector<int64_t, 8>> valueAt(ArrayRef<int64_t> point) const;
 
   /// Return whether `this` and `other` are equal as PWMAFunctions, i.e. whether
   /// they have the same dimensions, the same domain and they take the same
   /// value at every point in the domain.
   bool isEqual(const PWMAFunction &other) const;
 
+  /// Truncate the output dimensions to the first `count` dimensions.
+  ///
+  /// TODO: refactor so that this can be accomplished through removeIdRange.
+  void truncateOutput(unsigned count);
+
   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 output ids.
   unsigned numOutputs;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PWMAFUNCTION_H

mlir/include/mlir/Analysis/Presburger/Simplex.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H
 #define MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H
 
 #include "mlir/Analysis/Presburger/Fraction.h"
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Analysis/Presburger/PWMAFunction.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 namespace presburger {
 
 class GBRSimplex;
 
 /// The Simplex class implements a version of the Simplex and Generalized Basis
 /// Reduction algorithms, which can perform analysis of integer sets with affine
 /// inequalities and equalities. A Simplex can be constructed
 /// by specifying the dimensionality of the set. It supports adding affine
 /// inequalities and equalities, and can perform emptiness checks, i.e., it can
 /// find a solution to the set of constraints if one exists, or say that the
 /// set is empty if no solution exists. Furthermore, it can find a subset of
 /// these constraints that are redundant, i.e. a subset of constraints that
 /// doesn't constrain the affine set further after adding the non-redundant
 /// constraints. The LexSimplex class provides support for computing the
-/// lexicographical minimum of an IntegerRelation. Both these classes can be
-/// constructed from an IntegerRelation, and both inherit common
+/// lexicographic minimum of an IntegerRelation. The SymbolicLexMin class
+/// provides support for computing symbolic lexicographic minimums. All of these
+/// classes can be constructed from an IntegerRelation, and all inherit common
 /// functionality from SimplexBase.
 ///
 /// The implementations of the Simplex and SimplexBase classes, other than the
 /// functionality for obtaining an integer sample, are based on the paper
 /// "Simplify: A Theorem Prover for Program Checking"
 /// by D. Detlefs, G. Nelson, J. B. Saxe.
 ///
 /// We define variables, constraints, and unknowns. Consider the example of a

 /// since 3/6 = 1/2, 4/6 = 2/3, and 18/6 = 3.
 ///
 /// Every row and column except the first and second columns is associated with
 /// an unknown and every unknown is associated with a row or column. An unknown
 /// associated with a row or column is said to be in row or column orientation
 /// respectively. As described above, the first column is the common
 /// denominator. The second column represents the constant term, explained in
 /// more detail below. These two are _fixed columns_; they always retain their
-/// position as the first and second columns. Additionally, LexSimplex stores
-/// a so-call big M parameter (explained below) in the third column, so
-/// LexSimplex has three fixed columns.
+/// position as the first and second columns. Additionally, LexSimplexBase
+/// stores a so-call big M parameter (explained below) in the third column, so
+/// LexSimplexBase has three fixed columns. Finally, SymbolicLexSimplex has
+/// `nSymbol` variables designated as symbols. These occupy the next `nSymbol`
+/// columns, viz. the columns [3, 3 + nSymbol). For more information on symbols,
+/// see LexSimplexBase and SymbolicLexSimplex.
 ///
-/// LexSimplex does not directly support variables which can be negative, so we
-/// introduce the so-called big M parameter, an artificial variable that is
+/// LexSimplexBase does not directly support variables which can be negative, so
+/// we introduce the so-called big M parameter, an artificial variable that is
 /// considered to have an arbitrarily large value. We then transform the
 /// variables, say x, y, z, ... to M, M + x, M + y, M + z. Since M has been
 /// added to these variables, they are now known to have non-negative values.
-/// For more details, see the documentation for LexSimplex. The big M parameter
-/// is not considered a real unknown and is not stored in the `var` data
-/// structure; rather the tableau just has an extra fixed column for it just
-/// like the constant term.
+/// For more details, see the documentation for LexSimplexBase. The big M
+/// parameter is not considered a real unknown and is not stored in the `var`
+/// data structure; rather the tableau just has an extra fixed column for it
+/// just like the constant term.
 ///
 /// The vectors var and con store information about the variables and
 /// constraints respectively, namely, whether they are in row or column
 /// position, which row or column they are associated with, and whether they
 /// correspond to a variable or a constraint.
 ///
 /// An unknown is addressed by its index. If the index i is non-negative, then
 /// the variable var[i] is being addressed. If the index i is negative, then

 ///
 /// This SimplexBase class also supports taking snapshots of the current state
 /// and rolling back to prior snapshots. This works by maintaining an undo log
 /// of operations. Snapshots are just pointers to a particular location in the
 /// log, and rolling back to a snapshot is done by reverting each log entry's
 /// operation from the end until we reach the snapshot's location. SimplexBase
 /// also supports taking a snapshot including the exact set of basis unknowns;
 /// if this functionality is used, then on rolling back the exact basis will
-/// also be restored. This is used by LexSimplex because its algorithm, unlike
-/// Simplex, is sensitive to the exact basis used at a point.
+/// also be restored. This is used by LexSimplexBase because the lex algorithm,
+/// unlike `Simplex`, is sensitive to the exact basis used at a point.
 class SimplexBase {
 public:
   SimplexBase() = delete;
   virtual ~SimplexBase() = default;
 
   /// Returns true if the tableau is empty (has conflicting constraints),
   /// false otherwise.
   bool isEmpty() const;

 protected:
   /// Construct a SimplexBase with the specified number of variables and fixed
   /// columns.
   ///
   /// For example, Simplex uses two fixed columns: the denominator and the
   /// constant term, whereas LexSimplex has an extra fixed column for the
   /// so-called big M parameter. For more information see the documentation for
   /// LexSimplex.
-  SimplexBase(unsigned nVar, bool mustUseBigM);
+  SimplexBase(unsigned nVar, bool mustUseBigM, unsigned symbolOffset,
+              unsigned nSymbol);
 
   enum class Orientation { Row, Column };
 
   /// An Unknown is either a variable or a constraint. It is always associated
   /// with either a row or column. Whether it's a row or a column is specified
   /// by the orientation and pos identifies the specific row or column it is
   /// associated with. If the unknown is restricted, then it has a
   /// non-negativity constraint associated with it, i.e., its sample value must
   /// always be non-negative and if it cannot be made non-negative without
   /// violating other constraints, the tableau is empty.
   struct Unknown {
-    Unknown(Orientation oOrientation, bool oRestricted, unsigned oPos)
-        : pos(oPos), orientation(oOrientation), restricted(oRestricted) {}
+    Unknown(Orientation oOrientation, bool oRestricted, unsigned oPos,
+            bool oIsSymbol = false)
+        : pos(oPos), orientation(oOrientation), restricted(oRestricted),
+          isSymbol(oIsSymbol) {}
     unsigned pos;
     Orientation orientation;
     bool restricted : 1;
+    bool isSymbol : 1;
 
     void print(raw_ostream &os) const {
       os << (orientation == Orientation::Row ? "r" : "c");
       os << pos;
       if (restricted)
         os << " [>=0]";
     }
   };
protected:
   /// The number of columns in the tableau, including the common denominator
   /// and the constant column.
   unsigned nCol;
 
   /// The number of redundant rows in the tableau. These are the first
   /// nRedundant rows.
   unsigned nRedundant;
 
+  /// The number of parameters. This must be consistent with the number of
+  /// Unknowns in `var` below that have `isSymbol` set to true.
+  unsigned nSymbol;
+
   /// The matrix representing the tableau.
   Matrix tableau;
 
   /// This is true if the tableau has been detected to be empty, false
   /// otherwise.
   bool empty;
 
   /// Holds a log of operations, used for rolling back to a previous state.

 /// optimization. The implementation of this class is based on the paper
 /// "Parametric Integer Programming" by Paul Feautrier.
 ///
 /// This does not directly support negative-valued variables, so it uses the big
 /// M parameter trick to make all the variables non-negative. Basically we
 /// introduce an artifical variable M that is considered to have a value of
 /// +infinity and instead of the variables x, y, z, we internally use variables
 /// M + x, M + y, M + z, which are now guaranteed to be non-negative. See the
-/// documentation for Simplex for more details. The whole algorithm is performed
-/// without having to fix a "big enough" value of the big M parameter; it is
-/// just considered to be infinite throughout and it never appears in the final
-/// outputs. We will deal with sample values throughout that may in general be
-/// some linear expression involving M like pM + q or aM + b. We can compare
-/// these with each other. They have a total order:
-/// aM + b < pM + q iff a < p or (a == p and b < q).
+/// documentation for SimplexBase for more details. M is also considered to be
+/// an integer that is divisible by everything.
+///
+/// The whole algorithm is performed with M treated as a symbol;
+/// it is just considered to be infinite throughout and it never appears in the
+/// final outputs. We will deal with sample values throughout that may in
+/// general be some affine expression involving M, like pM + q or aM + b. We can
+/// compare these with each other. They have a total order:
+///
+/// aM + b < pM + q iff  a < p or (a == p and b < q).
 /// In particular, aM + b < 0 iff a < 0 or (a == 0 and b < 0).
 ///
+/// When performing symbolic optimization, sample values will be affine
+/// expressions in M and the symbols. For example, we could have sample values
+/// aM + bS + c and pM + qS + r, where S is a symbol. Now we have
+/// aM + bS + c < pM + qS + r iff (a < p) or (a == p and bS + c < qS + r).
+/// bS + c < qS + r can be always true, always false, or neither,
+/// depending on the set of values S can take. The symbols are always stored
+/// in columns [3, 3 + nSymbols). For more details, see the
+/// documentation for SymbolicLexSimplex.
+///
 /// Initially all the constraints to be added are added as rows, with no attempt
 /// to keep the tableau consistent. Pivots are only performed when some query
 /// is made, such as a call to getRationalLexMin. Care is taken to always
 /// maintain a lexicopositive basis transform, explained below.
 ///
-/// Let the variables be x = (x_1, ... x_n). Let the basis unknowns at a
-/// particular point be  y = (y_1, ... y_n). We know that x = A*y + b for some
-/// n x n matrix A and n x 1 column vector b. We want every column in A to be
-/// lexicopositive, i.e., have at least one non-zero element, with the first
-/// such element being positive. This property is preserved throughout the
-/// operation of LexSimplex. Note that on construction, the basis transform A is
-/// the indentity matrix and so every column is lexicopositive. Note that for
-/// LexSimplex, for the tableau to be consistent we must have non-negative
-/// sample values not only for the constraints but also for the variables.
-/// So if the tableau is consistent then x >= 0 and y >= 0, by which we mean
-/// every element in these vectors is non-negative. (note that this is a
-/// different concept from lexicopositivity!)
-///
-/// When we arrive at a basis such the basis transform is lexicopositive and the
-/// tableau is consistent, the sample point is the lexiographically minimum
-/// point in the polytope. We will show that A*y is zero or lexicopositive when
-/// y >= 0. Adding a lexicopositive vector to b will make it lexicographically
-/// bigger, so A*y + b is lexicographically bigger than b for any y >= 0 except
-/// y = 0. This shows that no point lexicographically smaller than x = b can be
-/// obtained. Since we already know that x = b is valid point in the space, this
-/// shows that x = b is the lexicographic minimum.
-///
-/// Proof that A*y is lexicopositive or zero when y > 0. Recall that every
-/// column of A is lexicopositive. Begin by considering A_1, the first row of A.
-/// If this row is all zeros, then (A*y)_1 = (A_1)*y = 0; proceed to the next
-/// row. If we run out of rows, A*y is zero and we are done; otherwise, we
-/// encounter some row A_i that has a non-zero element. Every column is
-/// lexicopositive and so has some positive element before any negative elements
-/// occur, so the element in this row for any column, if non-zero, must be
-/// positive. Consider (A*y)_i = (A_i)*y. All the elements in both vectors are
-/// non-negative, so if this is non-zero then it must be positive. Then the
-/// first non-zero element of A*y is positive so A*y is lexicopositive.
-///
-/// Otherwise, if (A_i)*y is zero, then for every column j that had a non-zero
-/// element in A_i, y_j is zero. Thus these columns have no contribution to A*y
-/// and we can completely ignore these columns of A. We now continue downwards,
-/// looking for rows of A that have a non-zero element other than in the ignored
-/// columns. If we find one, say A_k, once again these elements must be positive
-/// since they are the first non-zero element in each of these columns, so if
-/// (A_k)*y is not zero then we have that A*y is lexicopositive and if not we
-/// ignore more columns; eventually if all these dot products become zero then
-/// A*y is zero and we are done.
+/// Let the variables be x = (x_1, ... x_n).
+/// Let the symbols be   s = (s_1, ... s_m). Let the basis unknowns at a
+/// particular point be  y = (y_1, ... y_n). We know that x = A*y + T*s + b for
+/// some n x n matrix A, n x m matrix s, and n x 1 column vector b. We want
+/// every column in A to be lexicopositive, i.e., have at least one non-zero
+/// element, with the first such element being positive. This property is
+/// preserved throughout the operation of LexSimplexBase. Note that on
+/// construction, the basis transform A is the identity matrix and so every
+/// column is lexicopositive. Note that for LexSimplexBase, for the tableau to
+/// be consistent we must have non-negative sample values not only for the
+/// constraints but also for the variables. So if the tableau is consistent then
+/// x >= 0 and y >= 0, by which we mean every element in these vectors is
+/// non-negative. (note that this is a different concept from lexicopositivity!)
 class LexSimplexBase : public SimplexBase {
 public:
   ~LexSimplexBase() override = default;
 
   /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
   /// is the current number of variables, then the corresponding inequality is
   /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0.
   ///
   /// This just adds the inequality to the tableau and does not try to create a
   /// consistent tableau configuration.
   void addInequality(ArrayRef<int64_t> coeffs) final;
 
   /// Get a snapshot of the current state. This is used for rolling back.
   unsigned getSnapshot() { return SimplexBase::getSnapshotBasis(); }
 
 protected:
-  LexSimplexBase(unsigned nVar) : SimplexBase(nVar, /*mustUseBigM=*/true) {}
+  LexSimplexBase(unsigned nVar, unsigned symbolOffset, unsigned nSymbol)
+      : SimplexBase(nVar, /*mustUseBigM=*/true, symbolOffset, nSymbol) {}
   explicit LexSimplexBase(const IntegerRelation &constraints)
-      : LexSimplexBase(constraints.getNumIds()) {
+      : LexSimplexBase(constraints.getNumIds(),
+                       constraints.getIdKindOffset(IdKind::Symbol),
+                       constraints.getNumSymbolIds()) {
     intersectIntegerRelation(constraints);
   }
 
+  /// Add new symbolic variables to the end of the list of variables.
+  void appendSymbol();
+
   /// Try to move the specified row to column orientation while preserving the
-  /// lexicopositivity of the basis transform. If this is not possible, return
-  /// failure. This only occurs when the constraints have no solution; the
-  /// tableau will be marked empty in such a case.
+  /// lexicopositivity of the basis transform. The row must have a negative
+  /// sample value. If this is not possible, return failure. This only occurs
+  /// when the constraints have no solution; the tableau will be marked empty in
+  /// such a case.
   LogicalResult moveRowUnknownToColumn(unsigned row);
 
-  /// Given a row that has a non-integer sample value, add an inequality such
-  /// that this fractional sample value is cut away from the polytope. The added
-  /// inequality will be such that no integer points are removed.
+  /// Given a row that has a non-integer sample value, add an inequality to cut
+  /// away this fractional sample value from the polytope without removing any
+  /// integer points. The integer lexmin, if one existed, remains the same on
+  /// return.
+  ///
+  /// This assumes that the symbolic part of the sample is integral,
+  /// i.e., if the symbolic sample is (c + aM + b_1*s_1 + ... b_n*s_n)/d,
+  /// where s_1, ... s_n are symbols, this assumes that
+  /// (b_1*s_1 + ... + b_n*s_n)/s is integral.
   ///
-  /// Returns whether the cut constraint could be enforced, i.e. failure if the
-  /// cut made the polytope empty, and success if it didn't. Failure status
-  /// indicates that the polytope didn't have any integer points.
+  /// Return failure if the tableau became empty, and success if it didn't.
+  /// Failure status indicates that the polytope was integer empty.
   LogicalResult addCut(unsigned row);
 
   /// Undo the addition of the last constraint. This is only called while
   /// rolling back.
   void undoLastConstraint() final;
 
   /// Given two potential pivot columns for a row, return the one that results
-  /// in the lexicographically smallest sample vector.
+  /// in the lexicographically smallest sample vector. The row's sample value
+  /// must be negative. If symbols are involved, the sample value must be
+  /// negative for all possible assignments to the symbols.
   unsigned getLexMinPivotColumn(unsigned row, unsigned colA,
                                 unsigned colB) const;
 };
 
+/// A class for lexicographic optimization without any symbols. This also
+/// provides support for integer-exact redundancy and separateness checks.
 class LexSimplex : public LexSimplexBase {
 public:
-  explicit LexSimplex(unsigned nVar) : LexSimplexBase(nVar) {}
+  explicit LexSimplex(unsigned nVar)
+      : LexSimplexBase(nVar, /*symbolOffset=*/0, /*nSymbol=*/0) {}
   explicit LexSimplex(const IntegerRelation &constraints)
       : LexSimplexBase(constraints) {
     assert(constraints.getNumSymbolIds() == 0 &&
            "LexSimplex does not support symbols!");
   }
 
   /// Return the lexicographically minimum rational solution to the constraints.
   MaybeOptimum<SmallVector<Fraction, 8>> findRationalLexMin();
private:
   /// coordinates. Returns an empty optimum when the tableau is empty.
   ///
   /// Returns an unbounded optimum when the big M parameter is used and a
   /// variable has a non-zero big M coefficient, meaning its value is infinite
   /// or unbounded.
   MaybeOptimum<SmallVector<Fraction, 8>> getRationalSample() const;
 
   /// Make the tableau configuration consistent.
-  void restoreRationalConsistency();
+  LogicalResult restoreRationalConsistency();
 
   /// Return whether the specified row is violated;
   bool rowIsViolated(unsigned row) const;
 
   /// Get a constraint row that is violated, if one exists.
   /// Otherwise, return an empty optional.
   Optional<unsigned> maybeGetViolatedRow() const;
 
   /// Get a row corresponding to a var that has a non-integral sample value, if
   /// one exists. Otherwise, return an empty optional.
   Optional<unsigned> maybeGetNonIntegralVarRow() const;
+};
 
-  /// Given two potential pivot columns for a row, return the one that results
-  /// in the lexicographically smallest sample vector.
-  unsigned getLexMinPivotColumn(unsigned row, unsigned colA,
-                                unsigned colB) const;
+/// Represents the result of a symbolic lexicographic minimization computation.
+struct SymbolicLexMin {
+  SymbolicLexMin(unsigned nSymbols, unsigned nNonSymbols)
+      : lexmin(PresburgerSpace::getSetSpace(nSymbols), nNonSymbols),
+        unboundedDomain(
+            PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(nSymbols))) {}
+
+  /// This maps assignments of symbols to the corresponding lexmin.
+  /// Takes no value when no integer sample exists for the assignment or if the
+  /// lexmin is unbounded.
+  PWMAFunction lexmin;
+  /// Contains all assignments to the symbols that made the lexmin unbounded.
+  /// Note that the symbols of the input set to the symbolic lexmin are dims
+  /// of this PrebsurgerSet.
+  PresburgerSet unboundedDomain;
+};
+
+/// A class to perform symbolic lexicographic optimization,
+/// i.e., to find, for every assignment to the symbols the specified
+/// `symbolDomain`, the lexicographically minimum value integer value attained
+/// by the non-symbol variables.
+///
+/// The input is a set parametrized by some symbols, i.e., the constant terms
+/// of the constraints in the set are affine expressions in the symbols, and
+/// every assignment to the symbols defines a non-symbolic set.
+///
+/// Accordingly, the sample values of the rows in our tableau will be affine
+/// expressions in the symbols, and every assignment to the symbols will define
+/// a non-symbolic LexSimplex. We then run the algorithm of
+/// LexSimplex::findIntegerLexMin simultaneously for every value of the symbols
+/// in the domain.
+///
+/// Often, the pivot to be performed is the same for all values of the symbols,
+/// in which case we just do it. For example, if the symbolic sample of a row is
+/// negative for all values in the symbol domain, the row needs to be pivoted
+/// irrespective of the precise value of the symbols. To answer queries like
+/// "Is this symbolic sample always negative in the symbol domain?", we maintain
+/// a `LexSimplex domainSimplex` correponding to the symbol domain.
+///
+/// In other cases, it may be that the symbolic sample is violated at some
+/// values in the symbol domain and not violated at others. In this case,
+/// the pivot to be performed does depend on the value of the symbols. We
+/// handle this by splitting the symbol domain. We run the algorithm for the
+/// case where the row isn't violated, and then come back and run the case
+/// where it is.
+class SymbolicLexSimplex : public LexSimplexBase {
+public:
+  /// `constraints` is the set for which the symbolic lexmin will be computed.
+  /// `symbolDomain` is the set of values of the symbols for which the lexmin
+  /// will be computed. `symbolDomain` should have a dim id for every symbol in
+  /// `constraints`, and no other ids.
+  SymbolicLexSimplex(const IntegerPolyhedron &constraints,
+                     const IntegerPolyhedron &symbolDomain)
+      : LexSimplexBase(constraints), domainPoly(symbolDomain),
+        domainSimplex(symbolDomain) {
+    assert(domainPoly.getNumIds() == constraints.getNumSymbolIds());
+    assert(domainPoly.getNumDimIds() == constraints.getNumSymbolIds());
+  }
+
+  /// The lexmin will be stored as a function `lexmin` from symbols to
+  /// non-symbols in the result.
+  ///
+  /// For some values of the symbols, the lexmin may be unbounded.
+  /// These parts of the symbol domain will be stored in `unboundedDomain`.
+  SymbolicLexMin computeSymbolicIntegerLexMin();
+
+private:
+  /// Perform all pivots that do not require branching.
+  ///
+  /// Return failure if the tableau became empty, indicating that the polytope
+  /// is always integer empty in the current symbol domain.
+  /// Return success otherwise.
+  LogicalResult doNonBranchingPivots();
+
+  /// Get a row that is always violated in the current domain, if one exists.
+  Optional<unsigned> maybeGetAlwaysViolatedRow();
+
+  /// Get a row corresponding to a variable with non-integral sample value, if
+  /// one exists.
+  Optional<unsigned> maybeGetNonIntegralVarRow();
+
+  /// Given a row that has a non-integer sample value, cut away this fractional
+  /// sample value witahout removing any integer points, i.e., the integer
+  /// lexmin, if it exists, remains the same after a call to this function. This
+  /// may add constraints or local variables to the tableau, as well as to the
+  /// domain.
+  ///
+  /// Returns whether the cut constraint could be enforced, i.e. failure if the
+  /// cut made the polytope empty, and success if it didn't. Failure status
+  /// indicates that the polytope is always integer empty in the symbol domain
+  /// at the time of the call. (This function may modify the symbol domain, but
+  /// failure statu indicates that the polytope was empty for all symbol values
+  /// in the initial domain.)
+  LogicalResult addSymbolicCut(unsigned row);
+
+  /// Get the numerator of the symbolic sample of the specific row.
+  /// This is an affine expression in the symbols with integer coefficients.
+  /// The last element is the constant term. This ignores the big M coefficient.
+  SmallVector<int64_t, 8> getSymbolicSampleNumerator(unsigned row) const;
+
+  /// Return whether all the coefficients of the symbolic sample are integers.
+  ///
+  /// This does not consult the domain to check if the specified expression
+  /// is always integral despite coefficients being fractional.
+  bool isSymbolicSampleIntegral(unsigned row) const;
+
+  /// Record a lexmin. The tableau must be consistent with all variables
+  /// having symbolic samples with integer coefficients.
+  void recordOutput(SymbolicLexMin &result) const;
+
+  /// The symbol domain.
+  IntegerPolyhedron domainPoly;
+  /// Simplex corresponding to the symbol domain.
+  LexSimplex domainSimplex;
 };
 
 /// The Simplex class uses the Normal pivot rule and supports integer emptiness
 /// checks as well as detecting redundancies.
 ///
 /// The Simplex class supports redundancy checking via detectRedundant and
 /// isMarkedRedundant. A redundant constraint is one which is never violated as
 /// long as the other constraints are not violated, i.e., removing a redundant
 /// constraint does not change the set of solutions to the constraints. As a
 /// heuristic, constraints that have been marked redundant can be ignored for
 /// most operations. Therefore, these constraints are kept in rows 0 to
 /// nRedundant - 1, where nRedundant is a member variable that tracks the number
 /// of constraints that have been marked redundant.
 ///
 /// Finding an integer sample is done with the Generalized Basis Reduction
 /// algorithm. See the documentation for findIntegerSample and reduceBasis.
 class Simplex : public SimplexBase {
 public:
   enum class Direction { Up, Down };
 
   Simplex() = delete;
-  explicit Simplex(unsigned nVar) : SimplexBase(nVar, /*mustUseBigM=*/false) {}
+  explicit Simplex(unsigned nVar)
+      : SimplexBase(nVar, /*mustUseBigM=*/false, /*symbolOffset=*/0,
+                    /*nSymbol=*/0) {}
   explicit Simplex(const IntegerRelation &constraints)
       : Simplex(constraints.getNumIds()) {
     intersectIntegerRelation(constraints);
   }
   ~Simplex() override = default;
 
   /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
   /// is the current number of variables, then the corresponding inequality is

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

 //===- IntegerRelation.cpp - MLIR IntegerRelation 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 relation over integer tuples. A relation is
 // represented as a constraint system over a space of tuples of integer valued
 // varaiables supporting symbolic identifiers and existential quantification.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/LinearTransform.h"
+#include "mlir/Analysis/Presburger/PWMAFunction.h"
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "presburger"
void IntegerRelation::truncate(const CountsSnapshot &counts) {
   truncateIdKind(IdKind::Domain, counts);
   truncateIdKind(IdKind::Range, counts);
   truncateIdKind(IdKind::Symbol, counts);
   truncateIdKind(IdKind::Local, counts);
   removeInequalityRange(counts.getNumIneqs(), getNumInequalities());
   removeEqualityRange(counts.getNumEqs(), getNumEqualities());
 }
 
+SymbolicLexMin IntegerPolyhedron::findSymbolicIntegerLexMin() const {
+  // Compute the symbolic lexmin of the dims and locals, with the symbols being
+  // the actual symbols of this set.
+  SymbolicLexMin result =
+      SymbolicLexSimplex(
+          *this, PresburgerSpace::getSetSpace(/*numDims=*/getNumSymbolIds()))
+          .computeSymbolicIntegerLexMin();
+
+  // We want to return only the lexmin over the dims, so strip the locals from
+  // the computed lexmin.
+  result.lexmin.truncateOutput(result.lexmin.getNumOutputs() -
+                               getNumLocalIds());
+  return result;
+}
+
 unsigned IntegerRelation::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
   unsigned insertPos = PresburgerSpace::insertId(kind, pos, num);
   inequalities.insertColumns(insertPos, num);
   equalities.insertColumns(insertPos, num);
   return insertPos;
 }

mlir/lib/Analysis/Presburger/Matrix.cpp

void Matrix::reserveRows(unsigned rows) {
   data.reserve(rows * nReservedColumns);
 }
 
 unsigned Matrix::appendExtraRow() {
   resizeVertically(nRows + 1);
   return nRows - 1;
 }
 
+unsigned Matrix::appendExtraRow(ArrayRef<int64_t> elems) {
+  assert(elems.size() == nColumns && "elems must match row length!");
+  unsigned row = appendExtraRow();
+  for (unsigned col = 0; col < nColumns; ++col)
+    at(row, col) = elems[col];
+  return row;
+}
+
 void Matrix::resizeHorizontally(unsigned newNColumns) {
   if (newNColumns < nColumns)
     removeColumns(newNColumns, nColumns - newNColumns);
   if (newNColumns > nColumns)
     insertColumns(nColumns, newNColumns - nColumns);
 }
 
 void Matrix::resize(unsigned newNRows, unsigned newNColumns) {

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

 
 void MultiAffineFunction::eliminateRedundantLocalId(unsigned posA,
                                                     unsigned posB) {
   unsigned localOffset = getIdKindOffset(IdKind::Local);
   output.addToColumn(localOffset + posB, localOffset + posA, /*scale=*/1);
   IntegerPolyhedron::eliminateRedundantLocalId(posA, posB);
 }
 
+void MultiAffineFunction::truncateOutput(unsigned count) {
+  assert(count <= output.getNumRows());
+  output.resizeVertically(count);
+}
+
+void PWMAFunction::truncateOutput(unsigned count) {
+  assert(count <= numOutputs);
+  for (MultiAffineFunction &piece : pieces)
+    piece.truncateOutput(count);
+  numOutputs = count;
+}
+
 bool MultiAffineFunction::isEqualWhereDomainsOverlap(
     MultiAffineFunction other) const {
   if (!isSpaceCompatible(other))
     return false;
 
   // `commonFunc` has the same output as `this`.
   MultiAffineFunction commonFunc = *this;
   // After this merge, `commonFunc` and `other` have the same local ids; they

mlir/lib/Analysis/Presburger/Simplex.cpp

 
 using namespace mlir;
 using namespace presburger;
 
 using Direction = Simplex::Direction;
 
 const int nullIndex = std::numeric_limits<int>::max();
 
-SimplexBase::SimplexBase(unsigned nVar, bool mustUseBigM)
+SimplexBase::SimplexBase(unsigned nVar, bool mustUseBigM, unsigned symbolOffset,
+                         unsigned nSymbol)
     : usingBigM(mustUseBigM), nRow(0), nCol(getNumFixedCols() + nVar),
-      nRedundant(0), tableau(0, nCol), empty(false) {
+      nRedundant(0), nSymbol(nSymbol), tableau(0, nCol), empty(false) {
+  assert(symbolOffset + nSymbol <= nVar);
+
   colUnknown.insert(colUnknown.begin(), getNumFixedCols(), nullIndex);
   for (unsigned i = 0; i < nVar; ++i) {
     var.emplace_back(Orientation::Column, /*restricted=*/false,
                      /*pos=*/getNumFixedCols() + i);
     colUnknown.push_back(i);
   }
+
+  // Move the symbols to be in columns [3, 3 + nSymbol).
+  for (unsigned i = 0; i < nSymbol; ++i) {
+    var[symbolOffset + i].isSymbol = true;
+    swapColumns(var[symbolOffset + i].pos, getNumFixedCols() + i);
+  }
 }
 
 const Simplex::Unknown &SimplexBase::unknownFromIndex(int index) const {
   assert(index != nullIndex && "nullIndex passed to unknownFromIndex");
   return index >= 0 ? var[index] : con[~index];
 }
 
 const Simplex::Unknown &SimplexBase::unknownFromColumn(unsigned col) const {
if (usingBigM) {
     //
     // we internally use the variables
     //
     // M, M + x, M + y, M + z, ...
     //
     // where M is the big M parameter. As such, when the user tries to add
     // a row ax + by + cz + d, we express it in terms of our internal variables
     // as -(a + b + c)M + a(M + x) + b(M + y) + c(M + z) + d.
+    //
+    // Symbols don't use the big M parameter since they do not get lex
+    // optimized.
     int64_t bigMCoeff = 0;
     for (unsigned i = 0; i < coeffs.size() - 1; ++i)
+      if (!var[i].isSymbol)
         bigMCoeff -= coeffs[i];
     // The coefficient to the big M parameter is stored in column 2.
     tableau(nRow - 1, 2) = bigMCoeff;
   }
 
   // Process each given variable coefficient.
   for (unsigned i = 0; i < var.size(); ++i) {
     unsigned pos = var[i].pos;
     if (coeffs[i] == 0)
bool signMatchesDirection(int64_t elem, Direction direction) {
   return direction == Direction::Up ? elem > 0 : elem < 0;
 }
 
 Direction flippedDirection(Direction direction) {
   return direction == Direction::Up ? Direction::Down : Simplex::Direction::Up;
 }
 } // namespace
 
+/// We simply make the tableau consistent while maintaining a lexicopositive
+/// basis transform, and then return the sample value. If the tableau becomes
+/// empty, we return empty.
+///
+/// Let the variables be x = (x_1, ... x_n).
+/// Let the basis unknowns be y = (y_1, ... y_n).
+/// We have that x = A*y + b for some n x n matrix A and n x 1 column vector b.
+///
+/// As we will show below, A*y is either zero or lexicopositive.
+/// Adding a lexicopositive vector to b will make it lexicographically
+/// greater, so A*y + b is always equal to or lexicographically greater than b.
+/// Thus, since we can attain x = b, that is the lexicographic minimum.
+///
+/// We have that that every column in A is lexicopositive, i.e., has at least
+/// one non-zero element, with the first such element being positive. Since for
+/// the tableau to be consistent we must have non-negative sample values not
+/// only for the constraints but also for the variables, we also have x >= 0 and
+/// y >= 0, by which we mean every element in these vectors is non-negative.
+///
+/// Proof that if every column in A is lexicopositive, and y >= 0, then
+/// A*y is zero or lexicopositive. Begin by considering A_1, the first row of A.
+/// If this row is all zeros, then (A*y)_1 = (A_1)*y = 0; proceed to the next
+/// row. If we run out of rows, A*y is zero and we are done; otherwise, we
+/// encounter some row A_i that has a non-zero element. Every column is
+/// lexicopositive and so has some positive element before any negative elements
+/// occur, so the element in this row for any column, if non-zero, must be
+/// positive. Consider (A*y)_i = (A_i)*y. All the elements in both vectors are
+/// non-negative, so if this is non-zero then it must be positive. Then the
+/// first non-zero element of A*y is positive so A*y is lexicopositive.
+///
+/// Otherwise, if (A_i)*y is zero, then for every column j that had a non-zero
+/// element in A_i, y_j is zero. Thus these columns have no contribution to A*y
+/// and we can completely ignore these columns of A. We now continue downwards,
+/// looking for rows of A that have a non-zero element other than in the ignored
+/// columns. If we find one, say A_k, once again these elements must be positive
+/// since they are the first non-zero element in each of these columns, so if
+/// (A_k)*y is not zero then we have that A*y is lexicopositive and if not we
+/// add these to the set of ignored columns and continue to the next row. If we
+/// run out of rows, then A*y is zero and we are done.
 MaybeOptimum<SmallVector<Fraction, 8>> LexSimplex::findRationalLexMin() {
-  restoreRationalConsistency();
+  if (restoreRationalConsistency().failed())
+    return OptimumKind::Empty;
   return getRationalSample();
 }
 
+/// Given a row that has a non-integer sample value, add an inequality such
+/// that this fractional sample value is cut away from the polytope. The added
+/// inequality will be such that no integer points are removed. i.e., the
+/// integer lexmin, if it exists, is the same with and without this constraint.
+///
+/// Let the row be
+/// (c + coeffM*M + a_1*s_1 + ... + a_m*s_m + b_1*y_1 + ... + b_n*y_n)/d,
+/// where s_1, ... s_m are the symbols and
+///       y_1, ... y_n are the other basis unknowns.
+///
+/// For this to be an integer, we want
+/// coeffM*M + a_1*s_1 + ... + a_m*s_m + b_1*y_1 + ... + b_n*y_n = -c (mod d)
+/// Note that this constraint must always hold, independent of the basis,
+/// becuse the row unknown's value always equals this expression, even if *we*
+/// later compute the sample value from a different expression based on a
+/// different basis.
+///
+/// Let us assume that M has a factor of d in it. Imposing this constraint on M
+/// does not in any way hinder us from finding a value of M that is big enough.
+/// Moreover, this function is only called when the symbolic part of the sample,
+/// a_1*s_1 + ... + a_m*s_m, is known to be an integer.
+///
+/// Also, we can safely reduce the coefficients modulo d, so we have:
+///
+/// (b_1%d)y_1 + ... + (b_n%d)y_n = (-c%d) + k*d for some integer `k`
+///
+/// Note that all coefficient modulos here are non-negative. Also, all the
+/// unknowns are non-negative here as both constraints and variables are
+/// non-negative in LexSimplexBase. (We used the big M trick to make the
+/// variables non-negative). Therefore, the LHS here is non-negative.
+/// Since 0 <= (-c%d) < d, k is the quotient of dividing the LHS by d and
+/// is therefore non-negative as well.
+///
+/// So we have
+/// ((b_1%d)y_1 + ... + (b_n%d)y_n - (-c%d))/d >= 0.
+///
+/// The constraint is violated when added (it would be useless otherwise)
+/// so we immediately try to move it to a column.
 LogicalResult LexSimplexBase::addCut(unsigned row) {
-  int64_t denom = tableau(row, 0);
+  int64_t d = tableau(row, 0);
   addZeroRow(/*makeRestricted=*/true);
-  tableau(nRow - 1, 0) = denom;
-  tableau(nRow - 1, 1) = -mod(-tableau(row, 1), denom);
-  tableau(nRow - 1, 2) = 0; // M has all factors in it.
-  for (unsigned col = 3; col < nCol; ++col)
-    tableau(nRow - 1, col) = mod(tableau(row, col), denom);
+  tableau(nRow - 1, 0) = d;
+  tableau(nRow - 1, 1) = -mod(-tableau(row, 1), d); // -c%d.
+  tableau(nRow - 1, 2) = 0;
+  for (unsigned col = 3 + nSymbol; col < nCol; ++col)
+    tableau(nRow - 1, col) = mod(tableau(row, col), d); // b_i%d.
   return moveRowUnknownToColumn(nRow - 1);
 }
 
 Optional<unsigned> LexSimplex::maybeGetNonIntegralVarRow() const {
   for (const Unknown &u : var) {
     if (u.orientation == Orientation::Column)
       continue;
     // If the sample value is of the form (a/d)M + b/d, we need b to be
-    // divisible by d. We assume M is very large and contains all possible
+    // divisible by d. We assume M contains all possible
     // factors and is divisible by everything.
     unsigned row = u.pos;
     if (tableau(row, 1) % tableau(row, 0) != 0)
       return row;
   }
   return {};
 }
 
 MaybeOptimum<SmallVector<int64_t, 8>> LexSimplex::findIntegerLexMin() {
-  while (!empty) {
-    restoreRationalConsistency();
-    if (empty)
-      return OptimumKind::Empty;
+  // We first try to make the tableau consistent.
+  if (restoreRationalConsistency().failed())
 
-    if (Optional<unsigned> maybeRow = maybeGetNonIntegralVarRow()) {
-      // Failure occurs when the polytope is integer empty.
-      if (failed(addCut(*maybeRow)))
-        return OptimumKind::Empty;
-      continue;
-    }
+  // Then, if the sample value is integral, we are done.
+  while (Optional<unsigned> maybeRow = maybeGetNonIntegralVarRow()) {
+    // Otherwise, for the variable whose row has a non-integral sample value,
+    // we add a cut, a constraint that remove this rational point
+    // while preserving all integer points, thus keeping the lexmin the same.
+    // We then again try to make the tableau with the new constraint
+    // consistent. This continues until the tableau becomes empty, in which
+    // case there is no integer point, or until there are no variables with
+    // non-integral sample values.
+    //
+    // Failure indicates that the tableau became empty, which occurs when the
+    // polytope is integer empty.
+    if (addCut(*maybeRow).failed())
+      return OptimumKind::Empty;
+    if (restoreRationalConsistency().failed())
 
-    MaybeOptimum<SmallVector<Fraction, 8>> sample = getRationalSample();
-    assert(!sample.isEmpty() && "If we reached here the sample should exist!");
-    if (sample.isUnbounded())
-      return OptimumKind::Unbounded;
-    return llvm::to_vector<8>(
-        llvm::map_range(*sample, std::mem_fn(&Fraction::getAsInteger)));
-  }
 
-  // Polytope is integer empty.
-  return OptimumKind::Empty;
-}
-
 bool LexSimplex::isSeparateInequality(ArrayRef<int64_t> coeffs) {
   SimplexRollbackScopeExit scopeExit(*this);
   addInequality(coeffs);
   return findIntegerLexMin().isEmpty();
 }
 
 bool LexSimplex::isRedundantInequality(ArrayRef<int64_t> coeffs) {
   return isSeparateInequality(getComplementIneq(coeffs));
 }
+
+SmallVector<int64_t, 8>
+SymbolicLexSimplex::getSymbolicSampleNumerator(unsigned row) const {
+  SmallVector<int64_t, 8> sample;
+  sample.reserve(nSymbol + 1);
+  for (unsigned col = 3; col < 3 + nSymbol; ++col)
+    sample.push_back(tableau(row, col));
+  sample.push_back(tableau(row, 1));
+  return sample;
+}
+
+void LexSimplexBase::appendSymbol() {
+  appendVariable();
+  swapColumns(3 + nSymbol, nCol - 1);
+  var.back().isSymbol = true;
+  nSymbol++;
+}
+
+static bool isRangeDivisibleBy(ArrayRef<int64_t> range, int64_t divisor) {
+  assert(divisor > 0 && "divisor must be positive!");
+  return llvm::all_of(range, [divisor](int64_t x) { return x % divisor == 0; });
+}
+
+bool SymbolicLexSimplex::isSymbolicSampleIntegral(unsigned row) const {
+  int64_t denom = tableau(row, 0);
+  return tableau(row, 1) % denom == 0 &&
+         isRangeDivisibleBy(tableau.getRow(row).slice(3, nSymbol), denom);
+}
+
+/// This proceeds similarly to LexSimplex::addCut(). We are given a row that has
+/// a symbolic sample value with fractional coefficients.
+///
+/// Let the row be
+/// (c + coeffM*M + sum_i a_i*s_i + sum_j b_j*y_j)/d,
+/// where s_1, ... s_m are the symbols and
+///       y_1, ... y_n are the other basis unknowns.
+///
+/// As in LexSimplex::addCut, for this to be an integer, we want
+///
+/// coeffM*M + sum_j b_j*y_j = -c + sum_i (-a_i*s_i) (mod d)
+///
+/// This time, a_1*s_1 + ... + a_m*s_m may not be an integer. We find that
+///
+/// sum_i (b_i%d)y_i = ((-c%d) + sum_i (-a_i%d)s_i)%d + k*d for some integer k
+///
+/// where we take a modulo of the whole symbolic expression on the right to
+/// bring it into the range [0, d - 1]. Therefore, as in LexSimplex::addCut,
+/// k is the quotient on dividing the LHS by d, and since LHS >= 0, we have
+/// k >= 0 as well. We realize the modulo of the symbolic expression by adding a
+/// division variable
+///
+/// q = ((-c%d) + sum_i (-a_i%d)s_i)/d
+///
+/// to the symbol domain, so the equality becomes
+///
+/// sum_i (b_i%d)y_i = (-c%d) + sum_i (-a_i%d)s_i - q*d + k*d for some integer k
+///
+/// So the cut is
+/// (sum_i (b_i%d)y_i - (-c%d) - sum_i (-a_i%d)s_i + q*d)/d >= 0
+/// This constraint is violated when added so we immediately try to move it to a
+/// column.
+LogicalResult SymbolicLexSimplex::addSymbolicCut(unsigned row) {
+  int64_t d = tableau(row, 0);
+
+  // Add the division variable `q` described above to the symbol domain.
+  // q = ((-c%d) + sum_i (-a_i%d)s_i)/d.
+  SmallVector<int64_t, 8> domainDivCoeffs;
+  domainDivCoeffs.reserve(nSymbol + 1);
+  for (unsigned col = 3; col < 3 + nSymbol; ++col)
+    domainDivCoeffs.push_back(mod(-tableau(row, col), d)); // (-a_i%d)s_i
+  domainDivCoeffs.push_back(mod(-tableau(row, 1), d));     // -c%d.
+
+  domainSimplex.addDivisionVariable(domainDivCoeffs, d);
+  domainPoly.addLocalFloorDiv(domainDivCoeffs, d);
+
+  // Update `this` to account for the additional symbol we just added.
+  appendSymbol();
+
+  // Add the cut (sum_i (b_i%d)y_i - (-c%d) + sum_i -(-a_i%d)s_i + q*d)/d >= 0.
+  addZeroRow(/*makeRestricted=*/true);
+  tableau(nRow - 1, 0) = d;
+  tableau(nRow - 1, 2) = 0;
+
+  tableau(nRow - 1, 1) = -mod(-tableau(row, 1), d); // -(-c%d).
+  for (unsigned col = 3; col < 3 + nSymbol - 1; ++col)
+    tableau(nRow - 1, col) = -mod(-tableau(row, col), d); // -(-a_i%d)s_i.
+  tableau(nRow - 1, 3 + nSymbol - 1) = d;                 // q*d.
+
+  for (unsigned col = 3 + nSymbol; col < nCol; ++col)
+    tableau(nRow - 1, col) = mod(tableau(row, col), d); // (b_i%d)y_i.
+  return moveRowUnknownToColumn(nRow - 1);
+}
+
+void SymbolicLexSimplex::recordOutput(SymbolicLexMin &result) const {
+  Matrix output(0, domainPoly.getNumIds() + 1);
+  output.reserveRows(result.lexmin.getNumOutputs());
+  for (const Unknown &u : var) {
+    if (u.isSymbol)
+      continue;
+
+    if (u.orientation == Orientation::Column) {
+      // M + u has a sample value of zero so u has a sample value of -M, i.e,
+      // unbounded.
+      result.unboundedDomain.unionInPlace(domainPoly);
+      return;
+    }
+
+    int64_t denom = tableau(u.pos, 0);
+    if (tableau(u.pos, 2) < denom) {
+      // M + u has a sample value of fM + something, where f < 1, so
+      // u = (f - 1)M + something, which has a negative coefficient for M,
+      // and so is unbounded.
+      result.unboundedDomain.unionInPlace(domainPoly);
+      return;
+    }
+    assert(tableau(u.pos, 2) == denom &&
+           "Coefficient of M should not be greater than 1!");
+
+    SmallVector<int64_t, 8> sample = getSymbolicSampleNumerator(u.pos);
+    for (int64_t &elem : sample) {
+      assert(elem % denom == 0 && "coefficients must be integral!");
+      elem /= denom;
+    }
+    output.appendExtraRow(sample);
+  }
+  result.lexmin.addPiece(domainPoly, output);
+}
+
+Optional<unsigned> SymbolicLexSimplex::maybeGetAlwaysViolatedRow() {
+  // First look for rows that are clearly violated just from the big M
+  // coefficient, without needing to perform any simplex queries on the domain.
+  for (unsigned row = 0; row < nRow; ++row)
+    if (tableau(row, 2) < 0)
+      return row;
+
+  for (unsigned row = 0; row < nRow; ++row) {
+    if (tableau(row, 2) > 0)
+      continue;
+    if (domainSimplex.isSeparateInequality(getSymbolicSampleNumerator(row))) {
+      // Sample numerator always takes negative values in the symbol domain.
+      return row;
+    }
+  }
+  return {};
+}
+
+Optional<unsigned> SymbolicLexSimplex::maybeGetNonIntegralVarRow() {
+  for (const Unknown &u : var) {
+    if (u.orientation == Orientation::Column)
+      continue;
+    assert(!u.isSymbol && "Symbol should not be in row orientation!");
+    if (!isSymbolicSampleIntegral(u.pos))
+      return u.pos;
+  }
+  return {};
+}
+
+/// The non-branching pivots are just the ones moving the rows
+/// that are always violated in the symbol domain.
+LogicalResult SymbolicLexSimplex::doNonBranchingPivots() {
+  while (Optional<unsigned> row = maybeGetAlwaysViolatedRow())
+    if (moveRowUnknownToColumn(*row).failed())
+      return failure();
+  return success();
+}
+
+SymbolicLexMin SymbolicLexSimplex::computeSymbolicIntegerLexMin() {
+  SymbolicLexMin result(nSymbol, var.size() - nSymbol);
+
+  /// The algorithm is more naturally expressed recursively, but we implement
+  /// it iteratively here to avoid potential issues with stack overflows in the
+  /// compiler. We explicitly maintain the stack frames in a vector.
+  ///
+  /// To "recurse", we store the current "stack frame", i.e., state variables
+  /// that we will need when we "return", into `stack`, increment `level`, and
+  /// `continue`. To "tail recurse", we just `continue`.
+  /// To "return", we decrement `level` and `continue`.
+  ///
+  /// When there is no stack frame for the current `level`, this indicates that
+  /// we have just "recursed" or "tail recursed". When there does exist one,
+  /// this indicates that we have just "returned" from recursing. There is only
+  /// one point at which non-tail calls occur so we always "return" there.
+  unsigned level = 1;
+  struct StackFrame {
+    int splitIndex;
+    unsigned snapshot;
+    unsigned domainSnapshot;
+    IntegerRelation::CountsSnapshot domainPolyCounts;
+  };
+  SmallVector<StackFrame, 8> stack;
+
+  while (level > 0) {
+    assert(level >= stack.size());
+    if (level > stack.size()) {
+      if (empty || domainSimplex.findIntegerLexMin().isEmpty()) {
+        // No integer points; return.
+        --level;
+        continue;
+      }
+
+      if (doNonBranchingPivots().failed()) {
+        // Could not find pivots for violated constraints; return.
+        --level;
+        continue;
+      }
+
+      unsigned splitRow;
+      SmallVector<int64_t, 8> symbolicSample;
+      for (splitRow = 0; splitRow < nRow; ++splitRow) {
+        if (tableau(splitRow, 2) > 0)
+          continue;
+        assert(tableau(splitRow, 2) == 0 &&
+               "Non-branching pivots should have been handled already!");
+
+        symbolicSample = getSymbolicSampleNumerator(splitRow);
+        if (domainSimplex.isRedundantInequality(symbolicSample))
+          continue;
+
+        // It's neither redundant nor separate, so it takes both positive and
+        // negative values, and hence constitutes a row for which we need to
+        // split the domain and separately run each case.
+        assert(!domainSimplex.isSeparateInequality(symbolicSample) &&
+               "Non-branching pivots should have been handled already!");
+        break;
+      }
+
+      if (splitRow < nRow) {
+        unsigned domainSnapshot = domainSimplex.getSnapshot();
+        IntegerRelation::CountsSnapshot domainPolyCounts =
+            domainPoly.getCounts();
+
+        // First, we consider the part of the domain where the row is not
+        // violated. We don't have to do any pivots for the row in this case,
+        // but we record the additional constraint that defines this part of
+        // the domain.
+        domainSimplex.addInequality(symbolicSample);
+        domainPoly.addInequality(symbolicSample);
+
+        // Recurse.
+        //
+        // On return, the basis as a set is preserved but not the internal
+        // ordering within rows or columns. Thus, we take note of the index of
+        // the Unknown that caused the split, which may be in a different
+        // row when we come back from recursing. We will need this to recurse
+        // on the other part of the split domain, where the row is violated.
+        //
+        // Note that we have to capture the index above and not a reference to
+        // the Unknown itself, since the array it lives in might get
+        // reallocated.
+        int splitIndex = rowUnknown[splitRow];
+        unsigned snapshot = getSnapshot();
+        stack.push_back(
+            {splitIndex, snapshot, domainSnapshot, domainPolyCounts});
+        ++level;
+        continue;
+      }
+
+      // The tableau is rationally consistent for the current domain.
+      // Now we look for non-integral sample values and add cuts for them.
+      if (Optional<unsigned> row = maybeGetNonIntegralVarRow()) {
+        if (addSymbolicCut(*row).failed()) {
+          // No integral points; return.
+          --level;
+          continue;
+        }
+
+        // Rerun this level with the added cut constraint (tail recurse).
+        continue;
+      }
+
+      // Record output and return.
+      recordOutput(result);
+      --level;
+      continue;
+    }
+
+    if (level == stack.size()) {
+      // We have "returned" from "recursing".
+      const StackFrame &frame = stack.back();
+      domainPoly.truncate(frame.domainPolyCounts);
+      domainSimplex.rollback(frame.domainSnapshot);
+      rollback(frame.snapshot);
+      const Unknown &u = unknownFromIndex(frame.splitIndex);
+
+      // Drop the frame. We don't need it anymore.
+      stack.pop_back();
+
+      // Now we consider the part of the domain where the unknown `splitIndex`
+      // was negative.
+      assert(u.orientation == Orientation::Row &&
+             "The split row should have been returned to row orientation!");
+      SmallVector<int64_t, 8> splitIneq =
+          getComplementIneq(getSymbolicSampleNumerator(u.pos));
+      if (moveRowUnknownToColumn(u.pos).failed()) {
+        // The unknown can't be made non-negative; return.
+        --level;
+        continue;
+      }
+
+      // The unknown can be made negative; recurse with the corresponding domain
+      // constraints.
+      domainSimplex.addInequality(splitIneq);
+      domainPoly.addInequality(splitIneq);
+
+      // We are now taking care of the second half of the domain and we don't
+      // need to do anything else here after returning, so it's a tail recurse.
+      continue;
+    }
+  }
+
+  return result;
+}
+
 bool LexSimplex::rowIsViolated(unsigned row) const {
   if (tableau(row, 2) < 0)
     return true;
   if (tableau(row, 2) == 0 && tableau(row, 1) < 0)
     return true;
   return false;
 }
 
 Optional<unsigned> LexSimplex::maybeGetViolatedRow() const {
   for (unsigned row = 0; row < nRow; ++row)
     if (rowIsViolated(row))
       return row;
   return {};
 }
 
-// We simply look for violated rows and keep trying to move them to column
-// orientation, which always succeeds unless the constraints have no solution
-// in which case we just give up and return.
-void LexSimplex::restoreRationalConsistency() {
-  while (Optional<unsigned> maybeViolatedRow = maybeGetViolatedRow()) {
-    LogicalResult status = moveRowUnknownToColumn(*maybeViolatedRow);
-    if (failed(status))
-      return;
-  }
+/// We simply look for violated rows and keep trying to move them to column
+/// orientation, which always succeeds unless the constraints have no solution
+/// in which case we just give up and return.
+LogicalResult LexSimplex::restoreRationalConsistency() {
+  if (empty)
+    return failure();
+  while (Optional<unsigned> maybeViolatedRow = maybeGetViolatedRow())
+    if (moveRowUnknownToColumn(*maybeViolatedRow).failed())
+      return failure();
+  return success();
 }
 
 // Move the row unknown to column orientation while preserving lexicopositivity
-// of the basis transform.
+// of the basis transform. The sample value of the row must be negative.
 //
 // We only consider pivots where the pivot element is positive. Suppose no such
 // pivot exists, i.e., some violated row has no positive coefficient for any
 // basis unknown. The row can be represented as (s + c_1*u_1 + ... + c_n*u_n)/d,
 // where d is the denominator, s is the sample value and the c_i are the basis
 // coefficients. Since any feasible assignment of the basis satisfies u_i >= 0
 // for all i, and we have s < 0 as well as c_i < 0 for all i, any feasible
 // assignment would violate this row and therefore the constraints have no

 // This means that for some p, for all t < p,
 // C(t,k) = 0 => B(t,k) = B(t,j) * B(i,k) / B(i,j) and
 // C(t,k) < 0 => B(p,k) < B(t,j) * B(i,k) / B(i,j),
 // which is in contradiction to the fact that B.col(j) / B(i,j) must be
 // lexicographically smaller than B.col(k) / B(i,k), since it lexicographically
 // minimizes the change in sample value.
 LogicalResult LexSimplexBase::moveRowUnknownToColumn(unsigned row) {
   Optional<unsigned> maybeColumn;
-  for (unsigned col = 3; col < nCol; ++col) {
+  for (unsigned col = 3 + nSymbol; col < nCol; ++col) {
     if (tableau(row, col) <= 0)
       continue;
     maybeColumn =
         !maybeColumn ? col : getLexMinPivotColumn(row, *maybeColumn, col);
   }
 
   if (!maybeColumn) {
     markEmpty();
     return failure();
   }
 
   pivot(row, *maybeColumn);
   return success();
 }
 
 unsigned LexSimplexBase::getLexMinPivotColumn(unsigned row, unsigned colA,
                                               unsigned colB) const {
+  // First, let's consider the non-symbolic case.
   // A pivot causes the following change. (in the diagram the matrix elements
   // are shown as rationals and there is no common denominator used)
   //
   //            pivot col    big M col      const col
   // pivot row     a            p               b
   // other row     c            q               d
   //                        |
   //                        v
   //
   //            pivot col    big M col      const col
   // pivot row     1/a         -p/a           -b/a
   // other row     c/a        q - pc/a       d - bc/a
   //
   // Let the sample value of the pivot row be s = pM + b before the pivot. Since
   // the pivot row represents a violated constraint we know that s < 0.
   //
   // If the variable is a non-pivot column, its sample value is zero before and
   // after the pivot.
   //
   // If the variable is the pivot column, then its sample value goes from 0 to
   // (-p/a)M + (-b/a), i.e. 0 to -(pM + b)/a. Thus the change in the sample
   // value is -s/a.
   //
-  // If the variable is the pivot row, it sampel value goes from s to 0, for a
+  // If the variable is the pivot row, its sample value goes from s to 0, for a
   // change of -s.
   //
   // If the variable is a non-pivot row, its sample value changes from
   // qM + d to qM + d + (-pc/a)M + (-bc/a). Thus the change in sample value
   // is -(pM + b)(c/a) = -sc/a.
   //
   // Thus the change in sample value is either 0, -s/a, -s, or -sc/a. Here -s is
   // fixed for all calls to this function since the row and tableau are fixed.
   // The callee just wants to compare the return values with the return value of
   // other invocations of the same function. So the -s is common for all
   // comparisons involved and can be ignored, since -s is strictly positive.
   //
   // Thus we take away this common factor and just return 0, 1/a, 1, or c/a as
-  // appropriate. This allows us to run the entire algorithm without ever having
-  // to fix a value of M.
+  // appropriate. This allows us to run the entire algorithm treating M
+  // symbolically, as the pivot to be performed does not depend on the value
+  // of M, so long as the sample value s is negative. Note that this is not
+  // because of any special feature of M; by the same argument, we ignore the
+  // symbols too. The caller ensure that the sample value s is negative for
+  // all possible values of the symbols.
   auto getSampleChangeCoeffForVar = [this, row](unsigned col,
                                                 const Unknown &u) -> Fraction {
     int64_t a = tableau(row, col);
     if (u.orientation == Orientation::Column) {
       // Pivot column case.
       if (u.pos == col)
         return {1, a};
 

 /// pivot row     a/p          b/p     ->   pivot row      p/a         -b/a
 /// other row     c/q          d/q          other row     cp/aq    (da - bc)/aq
 ///
 /// The pivot row transform is accomplished be swapping a with the pivot row's
 /// common denominator and negating the pivot row except for the pivot column
 /// element.
 void SimplexBase::pivot(unsigned pivotRow, unsigned pivotCol) {
   assert(pivotCol >= getNumFixedCols() && "Refusing to pivot invalid column");
+  assert(!unknownFromColumn(pivotCol).isSymbol);
 
   swapRowWithCol(pivotRow, pivotCol);
   std::swap(tableau(pivotRow, 0), tableau(pivotRow, pivotCol));
   // We need to negate the whole pivot row except for the pivot column.
   if (tableau(pivotRow, 0) < 0) {
     // If the denominator is negative, we negate the row by simply negating the
     // denominator.
     tableau(pivotRow, 0) = -tableau(pivotRow, 0);
if (entry == UndoLogEntry::RemoveLastConstraint) {
     // was added after this variable was added, so the addition of such
     // constraints should already have been rolled back by the time we get to
     // rolling back the addition of the variable. Therefore, no constraint
     // currently has a component along the variable, so the variable itself must
     // be part of the basis.
     assert(var.back().orientation == Orientation::Column &&
            "Variable to be removed must be in column orientation!");
 
+    if (var.back().isSymbol)
+      nSymbol--;
+
     // Move this variable to the last column and remove the column from the
     // tableau.
     swapColumns(var.back().pos, nCol - 1);
     tableau.resizeHorizontally(nCol - 1);
     var.pop_back();
     colUnknown.pop_back();
     nCol--;
   } else if (entry == UndoLogEntry::UnmarkEmpty) {

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 "./Utils.h"
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
+#include "mlir/Analysis/Presburger/PWMAFunction.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 #include <numeric>
 
 using namespace mlir;
expectIntegerLexMin(parsePoly("(x, y, z) : (2*x + 13 >= 0, 4*y - 3*x - 2  >= "
                                 "0, 11*z + 5*y - 3*x + 7 >= 0)"),
                       {-6, -4, 0});
   // Similar to above but no lower bound on z.
   expectNoIntegerLexMin(OptimumKind::Unbounded,
                         parsePoly("(x, y, z) : (2*x + 13 >= 0, 4*y - 3*x - 2  "
                                   ">= 0, -11*z + 5*y - 3*x + 7 >= 0)"));
 }
 
+void expectSymbolicIntegerLexMin(
+    StringRef polyStr,
+    ArrayRef<std::pair<StringRef, SmallVector<SmallVector<int64_t, 8>, 8>>>
+        expectedLexminRepr,
+    ArrayRef<StringRef> expectedUnboundedDomainRepr) {
+  IntegerPolyhedron poly = parsePoly(polyStr);
+
+  ASSERT_NE(poly.getNumDimIds(), 0u);
+  ASSERT_NE(poly.getNumSymbolIds(), 0u);
+
+  PWMAFunction expectedLexmin =
+      parsePWMAF(/*numInputs=*/poly.getNumSymbolIds(),
+                 /*numOutputs=*/poly.getNumDimIds(), expectedLexminRepr);
+
+  PresburgerSet expectedUnboundedDomain = parsePresburgerSetFromPolyStrings(
+      poly.getNumSymbolIds(), expectedUnboundedDomainRepr);
+
+  SymbolicLexMin result = poly.findSymbolicIntegerLexMin();
+
+  EXPECT_TRUE(result.lexmin.isEqual(expectedLexmin));
+  if (!result.lexmin.isEqual(expectedLexmin)) {
+    llvm::errs() << "got:\n";
+    result.lexmin.dump();
+    llvm::errs() << "expected:\n";
+    expectedLexmin.dump();
+  }
+
+  EXPECT_TRUE(result.unboundedDomain.isEqual(expectedUnboundedDomain));
+  if (!result.unboundedDomain.isEqual(expectedUnboundedDomain))
+    result.unboundedDomain.dump();
+}
+
+void expectSymbolicIntegerLexMin(
+    StringRef polyStr,
+    ArrayRef<std::pair<StringRef, SmallVector<SmallVector<int64_t, 8>, 8>>>
+        result) {
+  expectSymbolicIntegerLexMin(polyStr, result, {});
+}
+
+TEST(IntegerPolyhedronTest, findSymbolicIntegerLexMin) {
+  expectSymbolicIntegerLexMin("(x)[a] : (x - a >= 0)",
+                              {
+                                  {"(a) : ()", {{1, 0}}}, // a
+                              });
+
+  expectSymbolicIntegerLexMin(
+      "(x)[a, b] : (x - a >= 0, x - b >= 0)",
+      {
+          {"(a, b) : (a - b >= 0)", {{1, 0, 0}}},     // a
+          {"(a, b) : (b - a - 1 >= 0)", {{0, 1, 0}}}, // b
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(x)[a, b, c] : (x -a >= 0, x - b >= 0, x - c >= 0)",
+      {
+          {"(a, b, c) : (a - b >= 0, a - c >= 0)", {{1, 0, 0, 0}}},         // a
+          {"(a, b, c) : (b - a - 1 >= 0, b - c >= 0)", {{0, 1, 0, 0}}},     // b
+          {"(a, b, c) : (c - a - 1 >= 0, c - b - 1 >= 0)", {{0, 0, 1, 0}}}, // c
+      });
+
+  expectSymbolicIntegerLexMin("(x, y)[a] : (x - a >= 0, x + y >= 0)",
+                              {
+                                  {"(a) : ()", {{1, 0}, {-1, 0}}}, // (a, -a)
+                              });
+
+  expectSymbolicIntegerLexMin(
+      "(x, y)[a] : (x - a >= 0, x + y >= 0, y >= 0)",
+      {
+          {"(a) : (a >= 0)", {{1, 0}, {0, 0}}},       // (a, 0)
+          {"(a) : (-a - 1 >= 0)", {{1, 0}, {-1, 0}}}, // (a, -a)
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(x, y)[a, b, c] : (x - a >= 0, y - b >= 0, c - x - y >= 0)",
+      {
+          {"(a, b, c) : (c - a - b >= 0)",
+           {{1, 0, 0, 0}, {0, 1, 0, 0}}}, // (a, b)
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(x, y, z)[a, b, c] : (c - z >= 0, b - y >= 0, x + y + z - a == 0)",
+      {
+          {"(a, b, c) : ()",
+           {{1, -1, -1, 0}, {0, 1, 0, 0}, {0, 0, 1, 0}}}, // (a - b - c, b, c)
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(x)[a, b] : (a >= 0, b >= 0, x >= 0, a + b + x - 1 >= 0)",
+      {
+          {"(a, b) : (a >= 0, b >= 0, a + b - 1 >= 0)", {{0, 0, 0}}}, // 0
+          {"(a, b) : (a == 0, b == 0)", {{0, 0, 1}}},                 // 1
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(x)[a, b] : (1 - a >= 0, a >= 0, 1 - b >= 0, b >= 0, 1 - x >= 0, x >= "
+      "0, a + b + x - 1 >= 0)",
+      {
+          {"(a, b) : (1 - a >= 0, a >= 0, 1 - b >= 0, b >= 0, a + b - 1 >= 0)",
+           {{0, 0, 0}}},                              // 0
+          {"(a, b) : (a == 0, b == 0)", {{0, 0, 1}}}, // 1
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(x, y, z)[a, b] : (x - a == 0, y - b == 0, x >= 0, y >= 0, z >= 0, x + "
+      "y + z - 1 >= 0)",
+      {
+          {"(a, b) : (a >= 0, b >= 0, 1 - a - b >= 0)",
+           {{1, 0, 0}, {0, 1, 0}, {-1, -1, 1}}}, // (a, b, 1 - a - b)
+          {"(a, b) : (a >= 0, b >= 0, a + b - 2 >= 0)",
+           {{1, 0, 0}, {0, 1, 0}, {0, 0, 0}}}, // (a, b, 0)
+      });
+
+  expectSymbolicIntegerLexMin("(x)[a, b] : (x - a == 0, x - b >= 0)",
+                              {
+                                  {"(a, b) : (a - b >= 0)", {{1, 0, 0}}}, // a
+                              });
+
+  expectSymbolicIntegerLexMin(
+      "(q)[a] : (a - 1 - 3*q == 0, q >= 0)",
+      {
+          {"(a) : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
+           {{0, 1, 0}}}, // a floordiv 3
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(r, q)[a] : (a - r - 3*q == 0, q >= 0, 1 - r >= 0, r >= 0)",
+      {
+          {"(a) : (a - 0 - 3*(a floordiv 3) == 0, a >= 0)",
+           {{0, 0, 0}, {0, 1, 0}}}, // (0, a floordiv 3)
+          {"(a) : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
+           {{0, 0, 1}, {0, 1, 0}}}, // (1 a floordiv 3)
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(r, q)[a] : (a - r - 3*q == 0, q >= 0, 2 - r >= 0, r - 1 >= 0)",
+      {
+          {"(a) : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
+           {{0, 0, 1}, {0, 1, 0}}}, // (1, a floordiv 3)
+          {"(a) : (a - 2 - 3*(a floordiv 3) == 0, a >= 0)",
+           {{0, 0, 2}, {0, 1, 0}}}, // (2, a floordiv 3)
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(r, q)[a] : (a - r - 3*q == 0, q >= 0, r >= 0)",
+      {
+          {"(a) : (a - 3*(a floordiv 3) == 0, a >= 0)",
+           {{0, 0, 0}, {0, 1, 0}}}, // (0, a floordiv 3)
+          {"(a) : (a - 1 - 3*(a floordiv 3) == 0, a >= 0)",
+           {{0, 0, 1}, {0, 1, 0}}}, // (1, a floordiv 3)
+          {"(a) : (a - 2 - 3*(a floordiv 3) == 0, a >= 0)",
+           {{0, 0, 2}, {0, 1, 0}}}, // (2, a floordiv 3)
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(x, y, z, w)[g] : ("
+      // x, y, z, w are boolean variables.
+      "1 - x >= 0, x >= 0, 1 - y >= 0, y >= 0,"
+      "1 - z >= 0, z >= 0, 1 - w >= 0, w >= 0,"
+      // We have some constraints on them:
+      "x + y + z - 1 >= 0,"             // x or y or z
+      "x + y + w - 1 >= 0,"             // x or y or w
+      "1 - x + 1 - y + 1 - w - 1 >= 0," // ~x or ~y or ~w
+      // What's the lexmin solution using exactly g true vars?
+      "g - x - y - z - w == 0)",
+      {
+          {"(g) : (g - 1 == 0)",
+           {{0, 0}, {0, 1}, {0, 0}, {0, 0}}}, // (0, 1, 0, 0)
+          {"(g) : (g - 2 == 0)",
+           {{0, 0}, {0, 0}, {0, 1}, {0, 1}}}, // (0, 0, 1, 1)
+          {"(g) : (g - 3 == 0)",
+           {{0, 0}, {0, 1}, {0, 1}, {0, 1}}}, // (0, 1, 1, 1)
+      });
+
+  // Bezout's lemma: if a, b are constants,
+  // the set of values that ax + by can take is all multiples of gcd(a, b).
+  expectSymbolicIntegerLexMin(
+      // If (x, y) is a solution for a given [a, r], then so is (x - 5, y + 2).
+      // So the lexmin is unbounded if it exists.
+      "(x, y)[a, r] : (a >= 0, r - a + 14*x + 35*y == 0)", {},
+      // According to Bezout's lemma, 14x + 35y can take on all multiples
+      // of 7 and no other values. So the solution exists iff r - a is a
+      // multiple of 7.
+      {"(a, r) : (a >= 0, r - a - 7*((r - a) floordiv 7) == 0)"});
+
+  // The lexmins are unbounded.
+  expectSymbolicIntegerLexMin("(x, y)[a] : (9*x - 4*y - 2*a >= 0)", {},
+                              {"(a) : ()"});
+
+  // Test cases adapted from isl.
+  expectSymbolicIntegerLexMin(
+      // a = 2b - 2(c - b), c - b >= 0.
+      // So b is minimized when c = b.
+      "(b, c)[a] : (a - 4*b + 2*c == 0, c - b >= 0)",
+      {
+          {"(a) : (a - 2*(a floordiv 2) == 0)",
+           {{0, 1, 0}, {0, 1, 0}}}, // (a floordiv 2, a floordiv 2)
+      });
+
+  expectSymbolicIntegerLexMin(
+      // 0 <= b <= 255, 1 <= a - 512b <= 509,
+      // b + 8 >= 1 + 16*(b + 8 floordiv 16) // i.e. b % 16 != 8
+      "(b)[a] : (255 - b >= 0, b >= 0, a - 512*b - 1 >= 0, 512*b -a + 509 >= "
+      "0, b + 7 - 16*((8 + b) floordiv 16) >= 0)",
+      {
+          {"(a) : (255 - (a floordiv 512) >= 0, a >= 0, a - 512*(a floordiv "
+           "512) - 1 >= 0, 512*(a floordiv 512) - a + 509 >= 0, (a floordiv "
+           "512) + 7 - 16*((8 + (a floordiv 512)) floordiv 16) >= 0)",
+           {{0, 1, 0, 0}}}, // (a floordiv 2, a floordiv 2)
+      });
+
+  expectSymbolicIntegerLexMin(
+      "(a, b)[K, N, x, y] : (N - K - 2 >= 0, K + 4 - N >= 0, x - 4 >= 0, x + 6 "
+      "- 2*N >= 0, K+N - x - 1 >= 0, a - N + 1 >= 0, K+N-1-a >= 0,a + 6 - b - "
+      "N >= 0, 2*N - 4 - a >= 0,"
+      "2*N - 3*K + a - b >= 0, 4*N - K + 1 - 3*b >= 0, b - N >= 0, a - x - 1 "
+      ">= 0)",
+      {{
+          "(K, N, x, y) : (x + 6 - 2*N >= 0, 2*N - 5 - x >= 0, x + 1 -3*K + N "
+          ">= 0, N + K - 2 - x >= 0, x - 4 >= 0)",
+          {{0, 0, 1, 0, 1}, {0, 1, 0, 0, 0}} // (1 + x, N)
+      }});
+}
+
 static void
 expectComputedVolumeIsValidOverapprox(const IntegerPolyhedron &poly,
                                       Optional<uint64_t> trueVolume,
                                       Optional<uint64_t> resultBound) {
   expectComputedVolumeIsValidOverapprox(poly.computeVolume(), trueVolume,
                                         resultBound);
 }