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

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.

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D122985

Author: Superty

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

@@ -536,6 +536,7 @@ class IntegerRelation : public PresburgerSpace {
   Matrix inequalities;
 };
 
+struct SymbolicLexMin;
 /// An IntegerPolyhedron is a PresburgerSpace subject to affine
 /// constraints. Affine constraints can be inequalities or equalities in the
 /// form:
@@ -593,6 +594,28 @@ class IntegerPolyhedron : public IntegerRelation {
   /// 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

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

@@ -151,6 +151,9 @@ class Matrix {
 
   /// 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;

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

@@ -106,6 +106,11 @@ class MultiAffineFunction : protected IntegerPolyhedron {
   /// 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;
 
@@ -165,6 +170,11 @@ class PWMAFunction : public PresburgerSpace {
   /// 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;
 

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

@@ -14,6 +14,7 @@
 
 #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"
@@ -145,6 +146,21 @@ void IntegerRelation::truncate(const CountsSnapshot &counts) {
   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));
 

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

@@ -66,6 +66,14 @@ unsigned Matrix::appendExtraRow() {
   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);

mlir/lib/Analysis/Presburger/Matrix.cpp

@@ -114,6 +114,18 @@ void MultiAffineFunction::eliminateRedundantLocalId(unsigned posA,
   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))