[MLIR] Introduce LexSimplex to support lexicographic optimization

[arjunp]

Summary

This patch introduces a class LexSimplex that can currently be used to find the
lexicographically minimal rational point in an IntegerPolyhedron. This is a
series of patches leading to computing the lexicographically minimal integer
lattice point as well parametric lexicographic minimization.

[Groverkss]

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

203

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

82

83

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

983

Is there any reason for not adding tests with symbols/local variables? If not, could you add one test to just check if they are supported.

[arjunp]

• clang-format
• Address kunwar's comments

[arjunp]

• small update to doc

[Groverkss]

LGTM.

[arjunp]

I'll wait for a couple of days in case @bondhugula has some comments.

[arjunp]

Rebase

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/Fraction.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 
 /// An integer polyhedron is the set of solutions to a list of affine
 /// constraints over n integer-valued variables/identifiers. Affine constraints
 /// can be inequalities or equalities in the form:

 /// example, consider the set: (x) : (exists q : 1 <= x <= 7, x = 2q). An
 /// assignment to symbolic and dimension variables is valid if there exists some
 /// assignment to the local variable `q` satisfying these constraints. For this
 /// example, the set is equivalent to {2, 4, 6}. Mathematically, existential
 /// quantification can be thought of as the result of projection. In this
 /// example, `q` is existentially quantified. This can be thought of as the
 /// result of projecting out `q` from the previous example, i.e. we obtained {2,
 /// 4, 6} by projecting out the second dimension from {(2, 1), (4, 2), (6, 2)}.
-///
 class IntegerPolyhedron {
 public:
   /// All derived classes of IntegerPolyhedron.
   enum class Kind {
     FlatAffineConstraints,
     FlatAffineValueConstraints,
     IntegerPolyhedron
   };
public:
 
   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);
 
+  /// Get the lexicographically minimum rational point satisfying the
+  /// constraints. Returns an empty optional if the polyhedron is empty or if
+  /// the lexmin is unbounded. Symbols are not supported and will result in
+  /// assert-failure.
+  Optional<SmallVector<Fraction, 8>> getRationalLexMin() const;
+
   /// Swap the posA^th identifier with the posB^th identifier.
   virtual void swapId(unsigned posA, unsigned posB);
 
   /// Removes all equalities and inequalities.
   void clearConstraints();
 
   /// Sets the `values.size()` identifiers starting at `po`s to the specified
   /// values and removes them.

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

public:
   /// rows that were at positions pos + count - 1 or later will be pushed to the
   /// right. The rows to be deleted must be valid rows: pos + count - 1 must be
   /// at most nRows - 1.
   void removeRows(unsigned pos, unsigned count);
   void removeRow(unsigned pos);
 
   void copyRow(unsigned sourceRow, unsigned targetRow);
 
+  void fillRow(unsigned row, int64_t value);
+
   /// Add `scale` multiples of the source row to the target row.
   void addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale);
 
   /// Add `scale` multiples of the source column to the target column.
   void addToColumn(unsigned sourceColumn, unsigned targetColumn, int64_t scale);
 
   /// Negate the specified column.
   void negateColumn(unsigned column);

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

 //===- Simplex.h - MLIR Simplex 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Functionality to perform analysis on an IntegerPolyhedron. In particular,
-// support for performing emptiness checks and redundancy checks.
+// support for performing emptiness checks, redundancy checks and obtaining the
+// lexicographically minimum rational element in a set.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H
 #define MLIR_ANALYSIS_PRESBURGER_SIMPLEX_H
 
 #include "mlir/Analysis/Presburger/Fraction.h"
 #include "mlir/Analysis/Presburger/IntegerPolyhedron.h"

 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. Currently, this only works for bounded
-/// sets. 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. Simplex can also be
-/// constructed from an IntegerPolyhedron object.
+/// 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 IntegerPolyhedron. Both these classes can be
+/// constructed from an IntegerPolyhedron, and both inherit common
+/// functionality from SimplexBase.
 ///
-/// The implementation of the Simplex and SimplexBase classes, other than the
-/// functionality for sampling, is based on the paper
+/// 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
 /// two-dimensional set defined by 1 + 2x + 3y >= 0 and 2x - 3y >= 0. Here,
-/// x, y, are variables while 1 + 2x + 3y >= 0, 2x - 3y >= 0 are
-/// constraints. Unknowns are either variables or constraints, i.e., x, y,
-/// 1 + 2x + 3y >= 0, 2x - 3y >= 0 are all unknowns.
+/// x, y, are variables while 1 + 2x + 3y >= 0, 2x - 3y >= 0 are constraints.
+/// Unknowns are either variables or constraints, i.e., x, y, 1 + 2x + 3y >= 0,
+/// 2x - 3y >= 0 are all unknowns.
 ///
 /// The implementation involves a matrix called a tableau, which can be thought
 /// of as a 2D matrix of rational numbers having number of rows equal to the
 /// number of constraints and number of columns equal to one plus the number of
 /// variables. In our implementation, instead of storing rational numbers, we
 /// store a common denominator for each row, so it is in fact a matrix of
 /// integers with number of rows equal to number of constraints and number of
 /// columns equal to _two_ plus the number of variables. For example, instead of
 /// storing a row of three rationals [1/2, 2/3, 3], we would store [6, 3, 4, 18]
 /// 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. The second
-/// column represents the constant, explained in more detail below. An unknown
-/// associated with a row or column is said to be in row or column position
-/// respectively.
+/// 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.
+///
+/// LexSimplex 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.
 ///
 /// 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

 /// by eliminating the old column unknown in favour of the new column unknown.
 /// If we had pivoted the column for x with the row for 2x - 3y >= 0,
 /// the new row for x would be [2, 1, 3] since x = (1*(2x - 3y) + 3*y)/2.
 /// See the documentation for the pivot member function for details.
 ///
 /// The association of unknowns to rows and columns is called the _tableau
 /// configuration_. The _sample value_ of an unknown in a particular tableau
 /// configuration is its value if all the column unknowns were set to zero.
-/// Concretely, for unknowns in column position the sample value is zero and
-/// for unknowns in row position the sample value is the constant term divided
-/// by the common denominator.
+/// Concretely, for unknowns in column position the sample value is zero; when
+/// the big M parameter is not used, for unknowns in row position the sample
+/// value is the constant term divided by the common denominator. When the big M
+/// parameter is used, if d is the denominator, p is the big M coefficient, and
+/// c is the constant term, then the sample value is (p*M + c)/d. Since M is
+/// considered to be positive infinity, this is positive (negative) infinity
+/// when p is positive or negative, and c/d when p is zero.
 ///
 /// The tableau configuration is called _consistent_ if the sample value of all
 /// restricted unknowns is non-negative. Initially there are no constraints, and
 /// the tableau is consistent. When a new constraint is added, its sample value
 /// in the current tableau configuration may be negative. In that case, we try
 /// to find a series of pivots to bring us to a consistent tableau
 /// configuration, i.e. we try to make the new constraint's sample value
 /// non-negative without making that of any other constraints negative. (See
 /// findPivot and findPivotRow for details.) If this is not possible, then the
 /// set of constraints is mutually contradictory and the tableau is marked
 /// _empty_, which means the set of constraints has no solution.
 ///
-/// 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.
-///
-/// This Simplex class also supports taking snapshots of the current state
+/// 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.
-///
-/// Finding an integer sample is done with the Generalized Basis Reduction
-/// algorithm. See the documentation for findIntegerSample and reduceBasis.
-///
-/// The SimplexBase class contains some basic functionality. In the future,
-/// lexicographic optimization will be supported by a class inheriting from
-/// SimplexBase. Functionality that will not supported by that class is placed
-/// in `Simplex`.
-
+/// 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.
 class SimplexBase {
 public:
-  enum class Direction { Up, Down };
-
   SimplexBase() = delete;
-  explicit SimplexBase(unsigned nVar);
-  explicit SimplexBase(const IntegerPolyhedron &constraints);
+  virtual ~SimplexBase() = default;
+
+  /// 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);
 
   /// Returns true if the tableau is empty (has conflicting constraints),
   /// false otherwise.
   bool isEmpty() const;
 
   /// 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.
-  void addInequality(ArrayRef<int64_t> coeffs);
+  virtual void addInequality(ArrayRef<int64_t> coeffs) = 0;
 
   /// Returns the number of variables in the tableau.
   unsigned getNumVariables() const;
 
   /// Returns the number of constraints in the tableau.
   unsigned getNumConstraints() const;
 
   /// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
   /// is the current number of variables, then the corresponding equality is
   /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0.
   void addEquality(ArrayRef<int64_t> coeffs);
 
   /// Add new variables to the end of the list of variables.
   void appendVariable(unsigned count = 1);
 
   /// Mark the tableau as being empty.
   void markEmpty();
 
   /// Get a snapshot of the current state. This is used for rolling back.
+  /// The same basis will not necessarily be restored on rolling back.
+  /// The snapshot only captures the set of variables and constraints present
+  /// in the Simplex.
   unsigned getSnapshot() const;
 
+  /// Get a snapshot of the current state including the basis. When rolling
+  /// back, the exact basis will be restored.
+  unsigned getSnapshotBasis();
+
   /// Rollback to a snapshot. This invalidates all later snapshots.
   void rollback(unsigned snapshot);
 
   /// Add all the constraints from the given IntegerPolyhedron.
   void intersectIntegerPolyhedron(const IntegerPolyhedron &poly);
 
-  /// Returns a rational sample point. Returns an empty optional if Simplex is
-  /// empty.
+  /// Returns the current sample point, which may contain non-integer (rational)
+  /// coordinates. Returns an empty optional when the tableau is empty.
+  ///
+  /// Also returns empty when the big M parameter is used and a variable
+  /// has a non-zero big M coefficient, meaning its value is infinite or
+  /// unbounded.
   Optional<SmallVector<Fraction, 8>> getRationalSample() const;
 
-  /// Returns the current sample point if it is integral. Otherwise, returns
-  /// None.
-  Optional<SmallVector<int64_t, 8>> getSamplePointIfIntegral() const;
-
   /// Print the tableau's internal state.
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
   enum class Orientation { Row, Column };
 
   /// An Unknown is either a variable or a constraint. It is always associated
void print(raw_ostream &os) const {
         os << " [>=0]";
     }
   };
 
   struct Pivot {
     unsigned row, column;
   };
 
-  /// Find a pivot to change the sample value of row in the specified
-  /// direction. The returned pivot row will be row if and only
-  /// if the unknown is unbounded in the specified direction.
+  /// Return any row that this column can be pivoted with, ignoring tableau
+  /// consistency.
   ///
-  /// Returns a (row, col) pair denoting a pivot, or an empty Optional if
-  /// no valid pivot exists.
-  Optional<Pivot> findPivot(int row, Direction direction) const;
+  /// Returns an empty optional if no pivot is possible, which happens only when
+  /// the column unknown is a variable and no constraint has a non-zero
+  /// coefficient for it.
+  Optional<unsigned> findAnyPivotRow(unsigned col);
 
   /// Swap the row with the column in the tableau's data structures but not the
   /// tableau itself. This is used by pivot.
   void swapRowWithCol(unsigned row, unsigned col);
 
   /// Pivot the row with the column.
   void pivot(unsigned row, unsigned col);
   void pivot(Pivot pair);
 
   /// Returns the unknown associated with index.
   const Unknown &unknownFromIndex(int index) const;
   /// Returns the unknown associated with col.
   const Unknown &unknownFromColumn(unsigned col) const;
   /// Returns the unknown associated with row.
   const Unknown &unknownFromRow(unsigned row) const;
   /// Returns the unknown associated with index.
   Unknown &unknownFromIndex(int index);
   /// Returns the unknown associated with col.
   Unknown &unknownFromColumn(unsigned col);
   /// Returns the unknown associated with row.
   Unknown &unknownFromRow(unsigned row);
 
   /// Add a new row to the tableau and the associated data structures.
-  unsigned addRow(ArrayRef<int64_t> coeffs);
+  /// The new row is considered to be a constraint; the new Unknown lives in
+  /// con.
+  ///
+  /// Returns the index of the new Unknown in con.
+  unsigned addRow(ArrayRef<int64_t> coeffs, bool makeRestricted = false);
 
   /// Normalize the given row by removing common factors between the numerator
   /// and the denominator.
   void normalizeRow(unsigned row);
 
   /// Swap the two rows/columns in the tableau and associated data structures.
   void swapRows(unsigned i, unsigned j);
   void swapColumns(unsigned i, unsigned j);
 
-  /// Restore the unknown to a non-negative sample value.
-  ///
-  /// Returns success if the unknown was successfully restored to a non-negative
-  /// sample value, failure otherwise.
-  LogicalResult restoreRow(Unknown &u);
-
   /// Enum to denote operations that need to be undone during rollback.
   enum class UndoLogEntry {
     RemoveLastConstraint,
     RemoveLastVariable,
     UnmarkEmpty,
-    UnmarkLastRedundant
+    UnmarkLastRedundant,
+    RestoreBasis
   };
 
+  /// Undo the addition of the last constraint. This will only be called from
+  /// undo, when rolling back.
+  virtual void undoLastConstraint() = 0;
+
+  /// Remove the last constraint, which must be in row orientation.
+  void removeLastConstraintRowOrientation();
+
   /// Undo the operation represented by the log entry.
   void undo(UndoLogEntry entry);
 
-  /// Find a row that can be used to pivot the column in the specified
-  /// direction. If skipRow is not null, then this row is excluded
-  /// from consideration. The returned pivot will maintain all constraints
-  /// except the column itself and skipRow, if it is set. (if these unknowns
-  /// are restricted).
-  ///
-  /// Returns the row to pivot to, or an empty Optional if the column
-  /// is unbounded in the specified direction.
-  Optional<unsigned> findPivotRow(Optional<unsigned> skipRow,
-                                  Direction direction, unsigned col) const;
+  /// Return the number of fixed columns, as described in the constructor above,
+  /// this is the number of columns beyond those for the variables in var.
+  unsigned getNumFixedCols() const { return usingBigM ? 3u : 2u; }
+
+  /// Stores whether or not a big M column is present in the tableau.
+  const bool usingBigM;
 
   /// The number of rows in the tableau.
   unsigned nRow;
 
   /// 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 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.
   SmallVector<UndoLogEntry, 8> undoLog;
 
+  /// Holds a vector of bases. The ith saved basis is the basis that should be
+  /// restored when processing the ith occurrance of UndoLogEntry::RestoreBasis
+  /// in undoLog. This is used by getSnapshotBasis.
+  SmallVector<SmallVector<int, 8>, 8> savedBases;
+
   /// These hold the indexes of the unknown at a given row or column position.
   /// We keep these as signed integers since that makes it convenient to check
   /// if an index corresponds to a variable or a constraint by checking the
   /// sign.
   ///
   /// colUnknown is padded with two null indexes at the front since the first
   /// two columns don't correspond to any unknowns.
   SmallVector<int, 8> rowUnknown, colUnknown;
 
   /// These hold information about each unknown.
   SmallVector<Unknown, 8> con, var;
 };
 
+/// Simplex class using the lexicographic pivot rule. Used for lexicographic
+/// 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).
+/// In particular, aM + b < 0 iff a < 0 or (a == 0 and b < 0).
+///
+/// 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.
+class LexSimplex : public SimplexBase {
+public:
+  explicit LexSimplex(unsigned nVar)
+      : SimplexBase(nVar, /*mustUseBigM=*/true) {}
+  explicit LexSimplex(const IntegerPolyhedron &constraints)
+      : LexSimplex(constraints.getNumIds()) {
+    intersectIntegerPolyhedron(constraints);
+  }
+  ~LexSimplex() 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 {
+    addRow(coeffs, /*makeRestricted=*/true);
+  }
+
+  /// Get a snapshot of the current state. This is used for rolling back.
+  unsigned getSnapshot() { return SimplexBase::getSnapshotBasis(); }
+
+  /// Return the lexicographically minimum rational solution to the constraints.
+  Optional<SmallVector<Fraction, 8>> getRationalLexMin();
+
+protected:
+  /// Undo the addition of the last constraint. This is only called while
+  /// rolling back.
+  void undoLastConstraint() final;
+
+  /// Make the tableau configuration consistent.
+  void 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;
+
+  /// 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;
+
+  /// 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.
+  LogicalResult moveRowUnknownToColumn(unsigned row);
+};
+
+/// 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) {}
+  explicit Simplex(unsigned nVar) : SimplexBase(nVar, /*mustUseBigM=*/false) {}
   explicit Simplex(const IntegerPolyhedron &constraints)
-      : SimplexBase(constraints) {}
+      : Simplex(constraints.getNumIds()) {
+    intersectIntegerPolyhedron(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
+  /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} >= 0.
+  ///
+  /// This also tries to restore the tableau configuration to a consistent
+  /// state and marks the Simplex empty if this is not possible.
+  void addInequality(ArrayRef<int64_t> coeffs) final;
 
   /// Compute the maximum or minimum value of the given row, depending on
   /// direction. The specified row is never pivoted. On return, the row may
   /// have a negative sample value if the direction is down.
   ///
   /// Returns a Fraction denoting the optimum, or a null value if no optimum
   /// exists, i.e., if the expression is unbounded in this direction.
   Optional<Fraction> computeRowOptimum(Direction direction, unsigned row);
public:
 
   /// Check if the specified equality already holds in the polytope.
   bool isRedundantEquality(ArrayRef<int64_t> coeffs);
 
   /// Returns true if this Simplex's polytope is a rational subset of `poly`.
   /// Otherwise, returns false.
   bool isRationalSubsetOf(const IntegerPolyhedron &poly);
 
+  /// Returns the current sample point if it is integral. Otherwise, returns
+  /// None.
+  Optional<SmallVector<int64_t, 8>> getSamplePointIfIntegral() const;
+
 private:
   friend class GBRSimplex;
 
+  /// Restore the unknown to a non-negative sample value.
+  ///
+  /// Returns success if the unknown was successfully restored to a non-negative
+  /// sample value, failure otherwise.
+  LogicalResult restoreRow(Unknown &u);
+
+  /// Find a pivot to change the sample value of row in the specified
+  /// direction while preserving tableau consistency, except that if the
+  /// direction is down then the pivot may make the specified row take a
+  /// negative value. The returned pivot row will be row if and only if the
+  /// unknown is unbounded in the specified direction.
+  ///
+  /// Returns a (row, col) pair denoting a pivot, or an empty Optional if
+  /// no valid pivot exists.
+  Optional<Pivot> findPivot(int row, Direction direction) const;
+
+  /// Find a row that can be used to pivot the column in the specified
+  /// direction. If skipRow is not null, then this row is excluded
+  /// from consideration. The returned pivot will maintain all constraints
+  /// except the column itself and skipRow, if it is set. (if these unknowns
+  /// are restricted).
+  ///
+  /// Returns the row to pivot to, or an empty Optional if the column
+  /// is unbounded in the specified direction.
+  Optional<unsigned> findPivotRow(Optional<unsigned> skipRow,
+                                  Direction direction, unsigned col) const;
+
+  /// Undo the addition of the last constraint while preserving tableau
+  /// consistency.
+  void undoLastConstraint() final;
+
   /// Compute the maximum or minimum of the specified Unknown, depending on
   /// direction. The specified unknown may be pivoted. If the unknown is
   /// restricted, it will have a non-negative sample value on return.
   /// Should not be called if the Simplex is empty.
   ///
   /// Returns a Fraction denoting the optimum, or a null value if no optimum
   /// exists, i.e., if the expression is unbounded in this direction.
   Optional<Fraction> computeOptimum(Direction direction, Unknown &u);

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

void IntegerPolyhedron::append(const IntegerPolyhedron &other) {
   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));
   }
 }
 
+Optional<SmallVector<Fraction, 8>>
+IntegerPolyhedron::getRationalLexMin() const {
+  assert(numSymbols == 0 && "Symbols are not supported!");
+  Optional<SmallVector<Fraction, 8>> maybeLexMin =
+      LexSimplex(*this).getRationalLexMin();
+
+  if (!maybeLexMin)
+    return {};
+
+  // The Simplex returns the lexmin over all the variables including locals. But
+  // locals are not actually part of the space and should not be returned in the
+  // result. Since the locals are placed last in the list of identifiers, they
+  // will be minimized last in the lexmin. So simply truncating out the locals
+  // from the end of the answer gives the desired lexmin over the dimensions.
+  assert(maybeLexMin->size() == getNumIds() &&
+         "Incorrect number of vars in lexMin!");
+  maybeLexMin->resize(getNumDimAndSymbolIds());
+  return maybeLexMin;
+}
+
 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);
 }
 

mlir/lib/Analysis/Presburger/Matrix.cpp

 
 void Matrix::copyRow(unsigned sourceRow, unsigned targetRow) {
   if (sourceRow == targetRow)
     return;
   for (unsigned c = 0; c < nColumns; ++c)
     at(targetRow, c) = at(sourceRow, c);
 }
 
+void Matrix::fillRow(unsigned row, int64_t value) {
+  for (unsigned col = 0; col < nColumns; ++col)
+    at(row, col) = value;
+}
+
 void Matrix::addToRow(unsigned sourceRow, unsigned targetRow, int64_t scale) {
   if (scale == 0)
     return;
   for (unsigned col = 0; col < nColumns; ++col)
     at(targetRow, col) += scale * at(sourceRow, col);
 }
 
 void Matrix::addToColumn(unsigned sourceColumn, unsigned targetColumn,

mlir/lib/Analysis/Presburger/Simplex.cpp

 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/Optional.h"
 
 namespace mlir {
 using Direction = Simplex::Direction;
 
 const int nullIndex = std::numeric_limits<int>::max();
 
-/// Construct a Simplex object with `nVar` variables.
-SimplexBase::SimplexBase(unsigned nVar)
-    : nRow(0), nCol(2), nRedundant(0), tableau(0, 2 + nVar), empty(false) {
-  colUnknown.push_back(nullIndex);
-  colUnknown.push_back(nullIndex);
+SimplexBase::SimplexBase(unsigned nVar, bool mustUseBigM)
+    : usingBigM(mustUseBigM), nRow(0), nCol(getNumFixedCols() + nVar),
+      nRedundant(0), tableau(0, nCol), empty(false) {
+  colUnknown.insert(colUnknown.begin(), getNumFixedCols(), nullIndex);
   for (unsigned i = 0; i < nVar; ++i) {
-    var.emplace_back(Orientation::Column, /*restricted=*/false, /*pos=*/nCol);
+    var.emplace_back(Orientation::Column, /*restricted=*/false,
+                     /*pos=*/getNumFixedCols() + i);
     colUnknown.push_back(i);
-    nCol++;
   }
 }
 
-SimplexBase::SimplexBase(const IntegerPolyhedron &constraints)
-    : SimplexBase(constraints.getNumIds()) {
-  for (unsigned i = 0, numIneqs = constraints.getNumInequalities();
-       i < numIneqs; ++i)
-    addInequality(constraints.getInequality(i));
-  for (unsigned i = 0, numEqs = constraints.getNumEqualities(); i < numEqs; ++i)
-    addEquality(constraints.getEquality(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 {
   assert(col < nCol && "Invalid column");
   return unknownFromIndex(colUnknown[col]);

 Simplex::Unknown &SimplexBase::unknownFromRow(unsigned row) {
   assert(row < nRow && "Invalid row");
   return unknownFromIndex(rowUnknown[row]);
 }
 
 /// Add a new row to the tableau corresponding to the given constant term and
 /// list of coefficients. The coefficients are specified as a vector of
 /// (variable index, coefficient) pairs.
-unsigned SimplexBase::addRow(ArrayRef<int64_t> coeffs) {
-  assert(coeffs.size() == 1 + var.size() &&
+unsigned SimplexBase::addRow(ArrayRef<int64_t> coeffs, bool makeRestricted) {
+  assert(coeffs.size() == var.size() + 1 &&
          "Incorrect number of coefficients!");
 
   ++nRow;
   // If the tableau is not big enough to accomodate the extra row, we extend it.
   if (nRow >= tableau.getNumRows())
     tableau.resizeVertically(nRow);
   rowUnknown.push_back(~con.size());
-  con.emplace_back(Orientation::Row, false, nRow - 1);
+  con.emplace_back(Orientation::Row, makeRestricted, nRow - 1);
+
+  // Zero out the new row.
+  tableau.fillRow(nRow - 1, 0);
 
   tableau(nRow - 1, 0) = 1;
   tableau(nRow - 1, 1) = coeffs.back();
-  for (unsigned col = 2; col < nCol; ++col)
-    tableau(nRow - 1, col) = 0;
+  if (usingBigM) {
+    // When the lexicographic pivot rule is used, instead of the variables
+    //
+    // x, y, z ...
+    //
+    // 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.
+    int64_t bigMCoeff = 0;
+    for (unsigned i = 0; i < coeffs.size() - 1; ++i)
+      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)
       continue;
 
     if (var[i].orientation == Orientation::Column) {
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
 
+Optional<SmallVector<Fraction, 8>> LexSimplex::getRationalLexMin() {
+  restoreRationalConsistency();
+  return getRationalSample();
+}
+
+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;
+  }
+}
+
+// Move the row unknown to column orientation while preserving lexicopositivity
+// of the basis transform.
+//
+// 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
+// solution.
+//
+// We can preserve lexicopositivity by picking the pivot column with positive
+// pivot element that makes the lexicographically smallest change to the sample
+// point.
+//
+// Proof. Let
+// x = (x_1, ... x_n) be the variables,
+// z = (z_1, ... z_m) be the constraints,
+// y = (y_1, ... y_n) be the current basis, and
+// define w = (x_1, ... x_n, z_1, ... z_m) = B*y + s.
+// B is basically the simplex tableau of our implementation except that instead
+// of only describing the transform to get back the non-basis unknowns, it
+// defines the values of all the unknowns in terms of the basis unknowns.
+// Similarly, s is the column for the sample value.
+//
+// Our goal is to show that each column in B, restricted to the first n
+// rows, is lexicopositive after the pivot if it is so before. This is
+// equivalent to saying the columns in the whole matrix are lexicopositive;
+// there must be some non-zero element in every column in the first n rows since
+// the n variables cannot be spanned without using all the n basis unknowns.
+//
+// Consider a pivot where z_i replaces y_j in the basis. Recall the pivot
+// transform for the tableau derived for SimplexBase::pivot:
+//
+//            pivot col    other col                   pivot col    other col
+// pivot row     a             b       ->   pivot row     1/a         -b/a
+// other row     c             d            other row     c/a        d - bc/a
+//
+// Similarly, a pivot results in B changing to B' and c to c'; the difference
+// between the tableau and these matrices B and B' is that there is no special
+// case for the pivot row, since it continues to represent the same unknown. The
+// same formula applies for all rows:
+//
+// B'.col(j) = B.col(j) / B(i,j)
+// B'.col(k) = B.col(k) - B(i,k) * B.col(j) / B(i,j) for k != j
+// and similarly, s' = s - s_i * B.col(j) / B(i,j).
+//
+// Since the row is violated, we have s_i < 0, so the change in sample value
+// when pivoting with column a is lexicographically smaller than that when
+// pivoting with column b iff B.col(a) / B(i, a) is lexicographically smaller
+// than B.col(b) / B(i, b).
+//
+// Since B(i, j) > 0, column j remains lexicopositive.
+//
+// For the other columns, suppose C.col(k) is not lexicopositive.
+// 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 LexSimplex::moveRowUnknownToColumn(unsigned row) {
+  Optional<unsigned> maybeColumn;
+  for (unsigned col = 3; 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 LexSimplex::getLexMinPivotColumn(unsigned row, unsigned colA,
+                                          unsigned colB) const {
+  // 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
+  // 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.
+  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};
+
+      // Non-pivot column case.
+      return {0, 1};
+    }
+
+    // Pivot row case.
+    if (u.pos == row)
+      return {1, 1};
+
+    // Non-pivot row case.
+    int64_t c = tableau(u.pos, col);
+    return {c, a};
+  };
+
+  for (const Unknown &u : var) {
+    Fraction changeA = getSampleChangeCoeffForVar(colA, u);
+    Fraction changeB = getSampleChangeCoeffForVar(colB, u);
+    if (changeA < changeB)
+      return colA;
+    if (changeA > changeB)
+      return colB;
+  }
+
+  // If we reached here, both result in exactly the same changes, so it
+  // doesn't matter which we return.
+  return colA;
+}
+
 /// Find a pivot to change the sample value of the row in the specified
 /// direction. The returned pivot row will involve `row` if and only if the
 /// unknown is unbounded in the specified direction.
 ///
 /// To increase (resp. decrease) the value of a row, we need to find a live
 /// column with a non-zero coefficient. If the coefficient is positive, we need
 /// to increase (decrease) the value of the column, and if the coefficient is
 /// negative, we need to decrease (increase) the value of the column. Also,
 /// we cannot decrease the sample value of restricted columns.
 ///
 /// If multiple columns are valid, we break ties by considering a lexicographic
 /// ordering where we prefer unknowns with lower index.
-Optional<SimplexBase::Pivot> SimplexBase::findPivot(int row,
-                                                    Direction direction) const {
+Optional<SimplexBase::Pivot> Simplex::findPivot(int row,
   Optional<unsigned> col;
   for (unsigned j = 2; j < nCol; ++j) {
     int64_t elem = tableau(row, j);
     if (elem == 0)
       continue;
 
     if (unknownFromColumn(j).restricted &&
         !signMatchesDirection(elem, direction))

 ///            pivot col    other col                    pivot col   other col
 /// 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 >= 2 && "Refusing to pivot invalid column");
+  assert(pivotCol >= getNumFixedCols() && "Refusing to pivot invalid column");
 
   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);
for (unsigned row = 0; row < nRow; ++row) {
     tableau(row, pivotCol) *= tableau(pivotRow, pivotCol);
     normalizeRow(row);
   }
 }
 
 /// Perform pivots until the unknown has a non-negative sample value or until
 /// no more upward pivots can be performed. Return success if we were able to
 /// bring the row to a non-negative sample value, and failure otherwise.
-LogicalResult SimplexBase::restoreRow(Unknown &u) {
+LogicalResult Simplex::restoreRow(Unknown &u) {
   assert(u.orientation == Orientation::Row &&
          "unknown should be in row position");
 
   while (tableau(u.pos, 1) < 0) {
     Optional<Pivot> maybePivot = findPivot(u.pos, Direction::Up);
     if (!maybePivot)
       break;
 

 /// change in sample value (in the specified direction). (note that c is
 /// non-negative here since the row is restricted and the tableau is consistent)
 ///
 /// We iterate through the rows and pick the row which imposes the most
 /// stringent bound, since pivoting with a row changes the row's sample value to
 /// 0 and hence saturates the bound it imposes. We break ties between rows that
 /// impose the same bound by considering a lexicographic ordering where we
 /// prefer unknowns with lower index value.
-Optional<unsigned> SimplexBase::findPivotRow(Optional<unsigned> skipRow,
-                                             Direction direction,
-                                             unsigned col) const {
+Optional<unsigned> Simplex::findPivotRow(Optional<unsigned> skipRow,
   Optional<unsigned> retRow;
   // Initialize these to zero in order to silence a warning about retElem and
   // retConst being used uninitialized in the initialization of `diff` below. In
   // reality, these are always initialized when that line is reached since these
   // are set whenever retRow is set.
   int64_t retElem = 0, retConst = 0;
   for (unsigned row = nRedundant; row < nRow; ++row) {
     if (skipRow && row == *skipRow)

 
 /// 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.
 ///
 /// We add the inequality and mark it as restricted. We then try to make its
 /// sample value non-negative. If this is not possible, the tableau has become
 /// empty and we mark it as such.
-void SimplexBase::addInequality(ArrayRef<int64_t> coeffs) {
-  unsigned conIndex = addRow(coeffs);
-  Unknown &u = con[conIndex];
-  u.restricted = true;
-  LogicalResult result = restoreRow(u);
+void Simplex::addInequality(ArrayRef<int64_t> coeffs) {
+  unsigned conIndex = addRow(coeffs, /*makeRestricted=*/true);
+  LogicalResult result = restoreRow(con[conIndex]);
   if (failed(result))
     markEmpty();
 }
 
 /// Add an equality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
 /// is the current number of variables, then the corresponding equality is
 /// c_n + c_0*x_0 + c_1*x_1 + ... + c_{n-1}*x_{n-1} == 0.
 ///

 
 unsigned SimplexBase::getNumVariables() const { return var.size(); }
 unsigned SimplexBase::getNumConstraints() const { return con.size(); }
 
 /// Return a snapshot of the current state. This is just the current size of the
 /// undo log.
 unsigned SimplexBase::getSnapshot() const { return undoLog.size(); }
 
-void SimplexBase::undo(UndoLogEntry entry) {
-  if (entry == UndoLogEntry::RemoveLastConstraint) {
-    Unknown &constraint = con.back();
-    if (constraint.orientation == Orientation::Column) {
-      unsigned column = constraint.pos;
-      Optional<unsigned> row;
+unsigned SimplexBase::getSnapshotBasis() {
+  SmallVector<int, 8> basis;
+  for (int index : colUnknown) {
+    if (index != nullIndex)
+      basis.push_back(index);
+  }
+  savedBases.push_back(std::move(basis));
+
+  undoLog.emplace_back(UndoLogEntry::RestoreBasis);
+  return undoLog.size() - 1;
+}
+
+void SimplexBase::removeLastConstraintRowOrientation() {
+  assert(con.back().orientation == Orientation::Row);
 
-      // Try to find any pivot row for this column that preserves tableau
-      // consistency (except possibly the column itself, which is going to be
-      // deallocated anyway).
-      //
-      // If no pivot row is found in either direction, then the unknown is
-      // unbounded in both directions and we are free to
-      // perform any pivot at all. To do this, we just need to find any row with
-      // a non-zero coefficient for the column.
-      if (Optional<unsigned> maybeRow =
-              findPivotRow({}, Direction::Up, column)) {
-        row = *maybeRow;
-      } else if (Optional<unsigned> maybeRow =
-                     findPivotRow({}, Direction::Down, column)) {
-        row = *maybeRow;
-      } else {
-        // The loop doesn't find a pivot row only if the column has zero
-        // coefficients for every row. But the unknown is a constraint,
-        // so it was added initially as a row. Such a row could never have been
-        // pivoted to a column. So a pivot row will always be found.
-        for (unsigned i = nRedundant; i < nRow; ++i) {
-          if (tableau(i, column) != 0) {
-            row = i;
-            break;
-          }
-        }
-      }
-      assert(row.hasValue() && "No pivot row found!");
-      pivot(*row, column);
-    }
+  // Move this unknown to the last row and remove the last row from the
+  // tableau.
+  swapRows(con.back().pos, nRow - 1);
+  // It is not strictly necessary to shrink the tableau, but for now we
+  // maintain the invariant that the tableau has exactly nRow rows.
+  tableau.resizeVertically(nRow - 1);
+  nRow--;
+  rowUnknown.pop_back();
+  con.pop_back();
+}
+
+// This doesn't find a pivot row only if the column has zero
+// coefficients for every row.
+//
+// If the unknown is a constraint, this can't happen, since it was added
+// initially as a row. Such a row could never have been pivoted to a column. So
+// a pivot row will always be found if we have a constraint.
+//
+// If we have a variable, then the column has zero coefficients for every row
+// iff no constraints have been added with a non-zero coefficient for this row.
+Optional<unsigned> SimplexBase::findAnyPivotRow(unsigned col) {
+  for (unsigned row = nRedundant; row < nRow; ++row)
+    if (tableau(row, col) != 0)
+      return row;
+  return {};
+}
+
+// It's not valid to remove the constraint by deleting the column since this
+// would result in an invalid basis.
+void Simplex::undoLastConstraint() {
+  if (con.back().orientation == Orientation::Column) {
+    // We try to find any pivot row for this column that preserves tableau
+    // unbounded in both directions and we are free to perform any pivot at
+    // all. To do this, we just need to find any row with a non-zero
+    // coefficient for the column. findAnyPivotRow will always be able to
+    // find such a row for a constraint.
+    unsigned column = con.back().pos;
+    if (Optional<unsigned> maybeRow = findPivotRow({}, Direction::Up, column)) {
+      pivot(*maybeRow, column);
+      pivot(*maybeRow, column);
+      Optional<unsigned> row = findAnyPivotRow(column);
+      assert(row.hasValue() && "Pivot should always exist for a constraint!");
+      pivot(*row, column);
+  removeLastConstraintRowOrientation();
+
+// It's not valid to remove the constraint by deleting the column since this
+// would result in an invalid basis.
+void LexSimplex::undoLastConstraint() {
+  if (con.back().orientation == Orientation::Column) {
+    // When removing the last constraint during a rollback, we just need to find
+    // any pivot at all, i.e., any row with non-zero coefficient for the
+    // column, because when rolling back a lexicographic simplex, we always
+    // end by restoring the exact basis that was present at the time of the
+    // snapshot, so what pivots we perform while undoing doesn't matter as
+    // long as we get the unknown to row orientation and remove it.
+    unsigned column = con.back().pos;
+    Optional<unsigned> row = findAnyPivotRow(column);
+    assert(row.hasValue() && "Pivot should always exist for a constraint!");
+  removeLastConstraintRowOrientation();
+}
 
-    // Move this unknown to the last row and remove the last row from the
-    // tableau.
-    swapRows(constraint.pos, nRow - 1);
-    // It is not strictly necessary to shrink the tableau, but for now we
-    // maintain the invariant that the tableau has exactly nRow rows.
-    tableau.resizeVertically(nRow - 1);
-    nRow--;
-    rowUnknown.pop_back();
-    con.pop_back();
+void SimplexBase::undo(UndoLogEntry entry) {
+  if (entry == UndoLogEntry::RemoveLastConstraint) {
+    // Simplex and LexSimplex handle this differently, so we call out to a
+    // virtual function to handle this.
+    undoLastConstraint();
   } else if (entry == UndoLogEntry::RemoveLastVariable) {
     // Whenever we are rolling back the addition of a variable, it is guaranteed
     // that the variable will be in column position.
     //
     // We can see this as follows: any constraint that depends on this variable
     // 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!");
 
     // 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) {
     empty = false;
   } else if (entry == UndoLogEntry::UnmarkLastRedundant) {
     nRedundant--;
+  } else if (entry == UndoLogEntry::RestoreBasis) {
+    assert(!savedBases.empty() && "No bases saved!");
+
+    SmallVector<int, 8> basis = std::move(savedBases.back());
+    savedBases.pop_back();
+
+    for (int index : basis) {
+      Unknown &u = unknownFromIndex(index);
+      if (u.orientation == Orientation::Column)
+        continue;
+      for (unsigned col = getNumFixedCols(); col < nCol; col++) {
+        assert(colUnknown[col] != nullIndex &&
+               "Column should not be a fixed column!");
+        if (std::find(basis.begin(), basis.end(), colUnknown[col]) !=
+            basis.end())
+          continue;
+        if (tableau(u.pos, col) == 0)
+          continue;
+        pivot(u.pos, col);
+        break;
+      }
+
+      assert(u.orientation == Orientation::Column && "No pivot found!");
+    }
   }
 }
 
 /// Rollback to the specified snapshot.
 ///
 /// We undo all the log entries until the log size when the snapshot was taken
 /// is reached.
 void SimplexBase::rollback(unsigned snapshot) {
Optional<SmallVector<Fraction, 8>> SimplexBase::getRationalSample() const {
   SmallVector<Fraction, 8> sample;
   sample.reserve(var.size());
   // Push the sample value for each variable into the vector.
   for (const Unknown &u : var) {
     if (u.orientation == Orientation::Column) {
       // If the variable is in column position, its sample value is zero.
       sample.emplace_back(0, 1);
     } else {
-      // If the variable is in row position, its sample value is the entry in
-      // the constant column divided by the entry in the common denominator
-      // column.
-      sample.emplace_back(tableau(u.pos, 1), tableau(u.pos, 0));
+      int64_t denom = tableau(u.pos, 0);
+
+      // When the big M parameter is being used, each variable x is represented
+      // as M + x, so its sample value is finite only if it is of the form
+      // 1*M + c. If the coefficient of M is not one then the sample value is
+      // infinite, and we return an empty optional.
+      if (usingBigM)
+        if (tableau(u.pos, 2) != denom)
+          return {};
+
+      // Otherwise, If the variable is in row position, its sample value is the
+      // entry in the constant column divided by the denominator.
+      sample.emplace_back(tableau(u.pos, 1), denom);
     }
   }
   return sample;
 }
 
-Optional<SmallVector<int64_t, 8>>
-SimplexBase::getSamplePointIfIntegral() const {
+Optional<SmallVector<int64_t, 8>> Simplex::getSamplePointIfIntegral() const {
   // If the tableau is empty, no sample point exists.
   if (empty)
     return {};
 
   // The value will always exist since the Simplex is non-empty.
   SmallVector<Fraction, 8> rationalSample = *getRationalSample();
   SmallVector<int64_t, 8> integerSample;
   integerSample.reserve(var.size());

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

TEST(IntegerPolyhedronTest, mergeDivisionsConstants) {
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(poly1.getNumLocalIds(), 2u);
     EXPECT_EQ(poly2.getNumLocalIds(), 2u);
   }
 }
 
+void expectRationalLexMin(const IntegerPolyhedron &poly,
+                          ArrayRef<Fraction> min) {
+  auto lexMin = poly.getRationalLexMin();
+  ASSERT_TRUE(lexMin.hasValue());
+  EXPECT_EQ(ArrayRef<Fraction>(*lexMin), min);
+}
+
+void expectNoRationalLexMin(const IntegerPolyhedron &poly) {
+  EXPECT_FALSE(poly.getRationalLexMin().hasValue());
+}
+
+TEST(IntegerPolyhedronTest, getRationalLexMin) {
+  MLIRContext context;
+  expectRationalLexMin(
+      parsePoly("(x, y, z) : (x + 10 >= 0, y + 40 >= 0, z + 30 >= 0)",
+                &context),
+      {{-10, 1}, {-40, 1}, {-30, 1}});
+  expectRationalLexMin(
+      parsePoly(
+          "(x, y, z) : (2*x + 7 >= 0, 3*y - 5 >= 0, 8*z + 10 >= 0, 9*z >= 0)",
+          &context),
+      {{-7, 2}, {5, 3}, {0, 1}});
+  expectRationalLexMin(
+      parsePoly(
+          "(x, y) : (3*x + 2*y + 10 >= 0, -3*y + 10 >= 0, 4*x - 7*y - 10 >= 0)",
+          &context),
+      {{-50, 29}, {-70, 29}});
+
+  // Test with some locals. This is basically x >= 11, 0 <= x - 2e <= 1.
+  // It'll just choose x = 11, e = 5.5 since it's rational lexmin.
+  expectRationalLexMin(
+      parsePoly(
+          "(x, y) : (x - 2*(x floordiv 2) == 0, y - 2*x >= 0, x - 11 >= 0)",
+          &context),
+      {{11, 1}, {22, 1}});
+
+  expectRationalLexMin(parsePoly("(x, y) : (3*x + 2*y + 10 >= 0,"
+                                 "-4*x + 7*y + 10 >= 0, -3*y + 10 >= 0)",
+                                 &context),
+                       {{-50, 9}, {10, 3}});
+
+  // Cartesian product of above with itself.
+  expectRationalLexMin(
+      parsePoly("(x, y, z, w) : (3*x + 2*y + 10 >= 0, -4*x + 7*y + 10 >= 0,"
+                "-3*y + 10 >= 0, 3*z + 2*w + 10 >= 0, -4*z + 7*w + 10 >= 0,"
+                "-3*w + 10 >= 0)",
+                &context),
+      {{-50, 9}, {10, 3}, {-50, 9}, {10, 3}});
+
+  // Same as above but for the constraints on z and w, we express "10" in terms
+  // of x and y. We know that x and y still have to take the values
+  // -50/9 and 10/3 since their constraints are the same and their values are
+  // minimized first. Accordingly, the values -9x - 12y,  -9x - 0y - 10,
+  // and -9x - 15y + 10 are all equal to 10.
+  expectRationalLexMin(
+      parsePoly(
+          "(x, y, z, w) : (3*x + 2*y + 10 >= 0, -4*x + 7*y + 10 >= 0, "
+          "-3*y + 10 >= 0, 3*z + 2*w - 9*x - 12*y >= 0,"
+          "-4*z + 7*w + - 9*x - 9*y - 10 >= 0, -3*w - 9*x - 15*y + 10 >= 0)",
+          &context),
+      {{-50, 9}, {10, 3}, {-50, 9}, {10, 3}});
+
+  // Same as above with one constraint removed, making the lexmin unbounded.
+  expectNoRationalLexMin(
+      parsePoly("(x, y, z, w) : (3*x + 2*y + 10 >= 0, -4*x + 7*y + 10 >= 0,"
+                "-3*y + 10 >= 0, 3*z + 2*w - 9*x - 12*y >= 0,"
+                "-4*z + 7*w + - 9*x - 9*y - 10>= 0)",
+                &context));
+
+  // Again, the lexmin is unbounded.
+  expectNoRationalLexMin(
+      parsePoly("(x, y, z) : (2*x + 5*y + 8*z - 10 >= 0,"
+                "2*x + 10*y + 8*z - 10 >= 0, 2*x + 5*y + 10*z - 10 >= 0)",
+                &context));
+
+  // The set is empty.
+  expectNoRationalLexMin(parsePoly("(x) : (2*x >= 0, -x - 1 >= 0)", &context));
+}
+
 } // namespace mlir