[MLIR][Presburger] Implement findSymbolicIntegerLexMax for IntegerRelation

[iambrj]

Summary

This patch implements findSymbolicIntegerLexMax for IntegerRelation

[Groverkss]

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

672–693

nit: I think you can just share the documentation of lexmin and lexmax functions. Just easier to maintain.

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

532

538

542–544

Please change the documentation to reflect lexmin -> lexopt

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

257

Can you write a doc comment explaining the flipping stratergy we use?

mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp

141

I think a few more tests would be nice.

[iambrj]

Address comments

[iambrj]

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

672–693

I am unsure how to do this, do you suggest the lexmax function documentation be moved to lexmin or just scrapped altogether?

[iambrj]

Delete commented code in unittest

[Groverkss]

LGTM.

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

672–693

Just add a "Same as findSymbolicIntegerLexMin but produces lexmax instead of lexmin" comment and it should be fine.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

259–260

"With the sign of each dimension variable in range flipped" is better.

[iambrj]

Address comments

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

 
 namespace mlir {
 namespace presburger {
 
 class IntegerRelation;
 class IntegerPolyhedron;
 class PresburgerSet;
 class PresburgerRelation;
-struct SymbolicLexMin;
+struct SymbolicLexOpt;
 
 /// The type of bound: equal, lower bound or upper bound.
 enum class BoundType { EQ, LB, UB };
 
 /// An IntegerRelation represents the set of points from a PresburgerSpace that
 /// satisfy a list of affine constraints. Affine constraints can be inequalities
 /// or equalities in the form:
 ///
public:
   ///
   /// (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.
+  /// This function is stored in the `lexopt` 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
+  /// `SymbolicLexOpt` stores these parts of the symbolic domain in a separate
   /// `PresburgerSet`, `unboundedDomain`.
-  SymbolicLexMin findSymbolicIntegerLexMin() const;
+  SymbolicLexOpt findSymbolicIntegerLexMin() const;
+
+  /// Compute the symbolic integer lexmax of the relation.
+  ///
+  /// This finds, for every assignment to the symbols and domain,
+  /// the lexicographically maximum value attained by the range.
+  ///
+  /// For example, the symbolic lexmax of the set
+  ///
+  /// (x, y)[a, b, c] : (a <= x, x <= b, x <= c)
+  ///
+  /// can be written as
+  ///
+  /// x = b if b <= c, a <= b
+  /// x = c if c <  b, a <= c
+  ///
+  /// This function is stored in the `lexopt` 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 min(b, c) < a.
+  ///
+  /// For some values of the symbols, the lexmax may be unbounded.
+  /// `SymbolicLexOpt` stores these parts of the symbolic domain in a separate
+  /// `PresburgerSet`, `unboundedDomain`.
+  SymbolicLexOpt findSymbolicIntegerLexMax() const;
 
   /// Return the set difference of this set and the given set, i.e.,
   /// return `this \ set`.
   PresburgerRelation subtract(const PresburgerRelation &set) const;
 
   void print(raw_ostream &os) const;
   void dump() const;
 

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

 /// 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
-/// lexicographic minimum of an IntegerRelation. The SymbolicLexMin class
+/// lexicographic minimum of an IntegerRelation. The SymbolicLexOpt 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.
private:
   /// Otherwise, return an empty optional.
   std::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.
   std::optional<unsigned> maybeGetNonIntegralVarRow() const;
 };
 
 /// Represents the result of a symbolic lexicographic minimization computation.
-struct SymbolicLexMin {
-  SymbolicLexMin(const PresburgerSpace &space)
-      : lexmin(space),
+struct SymbolicLexOpt {
+  SymbolicLexOpt(const PresburgerSpace &space)
+      : lexopt(space),
         unboundedDomain(PresburgerSet::getEmpty(space.getDomainSpace())) {}
 
   /// 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;
+  PWMAFunction lexopt;
   /// 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.
 ///

 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 var for every symbol in
   /// `constraints`, and no other vars. `isSymbol` specifies which vars of
   /// `constraints` should be considered as symbols.
   ///
-  /// The resulting SymbolicLexMin's space will be compatible with that of
+  /// The resulting SymbolicLexOpt's space will be compatible with that of
   /// symbolDomain.
   SymbolicLexSimplex(const IntegerRelation &constraints,
                      const IntegerPolyhedron &symbolDomain,
                      const llvm::SmallBitVector &isSymbol)
       : LexSimplexBase(constraints, isSymbol), domainPoly(symbolDomain),
         domainSimplex(symbolDomain) {
     // TODO consider supporting this case. It amounts
     // to just returning the input constraints.
public:
   /// 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`.
   ///
   /// The spaces of the sets in the result are compatible with the symbolDomain
   /// passed in the SymbolicLexSimplex constructor.
-  SymbolicLexMin computeSymbolicIntegerLexMin();
+  SymbolicLexOpt 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();
private:
   /// 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;
+  void recordOutput(SymbolicLexOpt &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

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

 }
 
 PresburgerRelation IntegerRelation::computeReprWithOnlyDivLocals() const {
   // If there are no locals, we're done.
   if (getNumLocalVars() == 0)
     return PresburgerRelation(*this);
 
   // Move all the non-div locals to the end, as the current API to
-  // SymbolicLexMin requires these to form a contiguous range.
+  // SymbolicLexOpt requires these to form a contiguous range.
   //
   // Take a copy so we can perform mutations.
   IntegerRelation copy = *this;
   std::vector<MaybeLocalRepr> reprs(getNumLocalVars());
   copy.getLocalReprs(&reprs);
 
   // Iterate through all the locals. The last `numNonDivLocals` are the locals
   // that have been scanned already and do not have division representations.
PresburgerRelation IntegerRelation::computeReprWithOnlyDivLocals() const {
   // lexicographically minimal valid assignment of "non-symbols", when a
   // satisfying assignment exists. It separately returns the set of assignments
   // to the "symbols" such that a satisfying assignment to the "non-symbols"
   // exists but the lexmin is unbounded. We basically want to find the set of
   // values of the "symbols" such that an assignment to the "non-symbols"
   // exists, which is the union of the domain of the returned lexmin function
   // and the returned set of assignments to the "symbols" that makes the lexmin
   // unbounded.
-  SymbolicLexMin lexminResult =
+  SymbolicLexOpt lexminResult =
       SymbolicLexSimplex(copy, /*symbolOffset*/ 0,
                          IntegerPolyhedron(PresburgerSpace::getSetSpace(
                              /*numDims=*/copy.getNumVars() - numNonDivLocals)))
           .computeSymbolicIntegerLexMin();
   PresburgerRelation result =
-      lexminResult.lexmin.getDomain().unionSet(lexminResult.unboundedDomain);
+      lexminResult.lexopt.getDomain().unionSet(lexminResult.unboundedDomain);
 
   // The result set might lie in the wrong space -- all its ids are dims.
   // Set it to the desired space and return.
   PresburgerSpace space = getSpace();
   space.removeVarRange(VarKind::Local, 0, getNumLocalVars());
   result.setSpace(space);
   return result;
 }
 
-SymbolicLexMin IntegerRelation::findSymbolicIntegerLexMin() const {
+SymbolicLexOpt IntegerRelation::findSymbolicIntegerLexMin() const {
   // Symbol and Domain vars will be used as symbols for symbolic lexmin.
   // In other words, for every value of the symbols and domain, return the
   // lexmin value of the (range, locals).
   llvm::SmallBitVector isSymbol(getNumVars(), false);
   isSymbol.set(getVarKindOffset(VarKind::Symbol),
                getVarKindEnd(VarKind::Symbol));
   isSymbol.set(getVarKindOffset(VarKind::Domain),
                getVarKindEnd(VarKind::Domain));
   // Compute the symbolic lexmin of the dims and locals, with the symbols being
   // the actual symbols of this set.
   // The resultant space of lexmin is the space of the relation itself.
-  SymbolicLexMin result =
+  SymbolicLexOpt result =
       SymbolicLexSimplex(*this,
                          IntegerPolyhedron(PresburgerSpace::getSetSpace(
                              /*numDims=*/getNumDomainVars(),
                              /*numSymbols=*/getNumSymbolVars())),
                          isSymbol)
           .computeSymbolicIntegerLexMin();
 
   // We want to return only the lexmin over the dims, so strip the locals from
   // the computed lexmin.
-  result.lexmin.removeOutputs(result.lexmin.getNumOutputs() - getNumLocalVars(),
-                              result.lexmin.getNumOutputs());
+  result.lexopt.removeOutputs(result.lexopt.getNumOutputs() - getNumLocalVars(),
+                              result.lexopt.getNumOutputs());
   return result;
 }
 
+SymbolicLexOpt IntegerRelation::findSymbolicIntegerLexMax() const {
+  IntegerRelation flippedRel = *this;
+  // Flip range sign by flipping the sign of range variables in all constraints.
+  for (unsigned j = getNumDomainVars(),
+                b = getNumDomainVars() + getNumRangeVars();
+       j < b; j++) {
+    for (unsigned i = 0, a = getNumEqualities(); i < a; i++)
+      flippedRel.atEq(i, j) = -1 * atEq(i, j);
+    for (unsigned i = 0, a = getNumInequalities(); i < a; i++)
+      flippedRel.atIneq(i, j) = -1 * atIneq(i, j);
+  }
+  // Compute negated lexmax by computing lexmin.
+  SymbolicLexOpt flippedSymbolicIntegerLexMax =
+                     flippedRel.findSymbolicIntegerLexMin(),
+                 symbolicIntegerLexMax(
+                     flippedSymbolicIntegerLexMax.lexopt.getSpace());
+  // Get lexmax by again flipping range sign in the PWMA constraints.
+  for (auto &flippedPiece :
+       flippedSymbolicIntegerLexMax.lexopt.getAllPieces()) {
+    Matrix mat = flippedPiece.output.getOutputMatrix();
+    for (unsigned i = 0, e = mat.getNumRows(); i < e; i++)
+      mat.negateRow(i);
+    MultiAffineFunction maf(flippedPiece.output.getSpace(), mat);
+    PWMAFunction::Piece piece = {flippedPiece.domain, maf};
+    symbolicIntegerLexMax.lexopt.addPiece(piece);
+  }
+  symbolicIntegerLexMax.unboundedDomain =
+      flippedSymbolicIntegerLexMax.unboundedDomain;
+  return symbolicIntegerLexMax;
+}
+
 PresburgerRelation
 IntegerRelation::subtract(const PresburgerRelation &set) const {
   return PresburgerRelation(*this).subtract(set);
 }
 
 unsigned IntegerRelation::insertVar(VarKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumVarKind(kind));
 

mlir/lib/Analysis/Presburger/Simplex.cpp

for (unsigned col = 3; col < 3 + nSymbol - 1; ++col)
     tableau(cutRow, col) = -mod(-tableau(row, col), d); // -(-a_i%d)s_i.
   tableau(cutRow, 3 + nSymbol - 1) = d;                 // q*d.
 
   for (unsigned col = 3 + nSymbol, e = getNumColumns(); col < e; ++col)
     tableau(cutRow, col) = mod(tableau(row, col), d); // (b_i%d)y_i.
   return moveRowUnknownToColumn(cutRow);
 }
 
-void SymbolicLexSimplex::recordOutput(SymbolicLexMin &result) const {
+void SymbolicLexSimplex::recordOutput(SymbolicLexOpt &result) const {
   Matrix output(0, domainPoly.getNumVars() + 1);
-  output.reserveRows(result.lexmin.getNumOutputs());
+  output.reserveRows(result.lexopt.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);
for (const Unknown &u : var) {
     for (MPInt &elem : sample) {
       assert(elem % denom == 0 && "coefficients must be integral!");
       elem /= denom;
     }
     output.appendExtraRow(sample);
   }
 
   // Store the output in a MultiAffineFunction and add it the result.
-  PresburgerSpace funcSpace = result.lexmin.getSpace();
+  PresburgerSpace funcSpace = result.lexopt.getSpace();
   funcSpace.insertVar(VarKind::Local, 0, domainPoly.getNumLocalVars());
 
-  result.lexmin.addPiece(
+  result.lexopt.addPiece(
       {PresburgerSet(domainPoly),
        MultiAffineFunction(funcSpace, output, domainPoly.getLocalReprs())});
 }
 
 std::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, e = getNumRows(); row < e; ++row)

 /// that are always violated in the symbol domain.
 LogicalResult SymbolicLexSimplex::doNonBranchingPivots() {
   while (std::optional<unsigned> row = maybeGetAlwaysViolatedRow())
     if (moveRowUnknownToColumn(*row).failed())
       return failure();
   return success();
 }
 
-SymbolicLexMin SymbolicLexSimplex::computeSymbolicIntegerLexMin() {
-  SymbolicLexMin result(PresburgerSpace::getRelationSpace(
+SymbolicLexOpt SymbolicLexSimplex::computeSymbolicIntegerLexMin() {
+  SymbolicLexOpt result(PresburgerSpace::getRelationSpace(
       /*numDomain=*/domainPoly.getNumDimVars(),
       /*numRange=*/var.size() - nSymbol,
       /*numSymbols=*/domainPoly.getNumSymbolVars()));
 
   /// 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.
   ///

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

void expectSymbolicIntegerLexMin(
     StringRef polyStr,
     ArrayRef<std::pair<StringRef, StringRef>> expectedLexminRepr,
     ArrayRef<StringRef> expectedUnboundedDomainRepr) {
   IntegerPolyhedron poly = parseIntegerPolyhedron(polyStr);
 
   ASSERT_NE(poly.getNumDimVars(), 0u);
   ASSERT_NE(poly.getNumSymbolVars(), 0u);
 
-  SymbolicLexMin result = poly.findSymbolicIntegerLexMin();
+  SymbolicLexOpt result = poly.findSymbolicIntegerLexMin();
 
   if (expectedLexminRepr.empty()) {
-    EXPECT_TRUE(result.lexmin.getDomain().isIntegerEmpty());
+    EXPECT_TRUE(result.lexopt.getDomain().isIntegerEmpty());
   } else {
     PWMAFunction expectedLexmin = parsePWMAF(expectedLexminRepr);
-    EXPECT_TRUE(result.lexmin.isEqual(expectedLexmin));
+    EXPECT_TRUE(result.lexopt.isEqual(expectedLexmin));
   }
 
   if (expectedUnboundedDomainRepr.empty()) {
     EXPECT_TRUE(result.unboundedDomain.isIntegerEmpty());
   } else {
     PresburgerSet expectedUnboundedDomain =
         parsePresburgerSet(expectedUnboundedDomainRepr);
     EXPECT_TRUE(result.unboundedDomain.isEqual(expectedUnboundedDomain));

mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp

TEST(IntegerRelationTest, applyDomainAndRange) {
 
     map1.applyDomain(map2);
 
     EXPECT_TRUE(map1.isEqual(map3));
   }
 }
 
 TEST(IntegerRelationTest, symbolicLexmin) {
-  SymbolicLexMin lexmin =
+  SymbolicLexOpt lexmin =
       parseRelationFromSet("(a, x)[b] : (x - a >= 0, x - b >= 0)", 1)
           .findSymbolicIntegerLexMin();
 
   PWMAFunction expectedLexmin = parsePWMAF({
       {"(a)[b] : (a - b >= 0)", "(a)[b] -> (a)"},     // a
       {"(a)[b] : (b - a - 1 >= 0)", "(a)[b] -> (b)"}, // b
   });
   EXPECT_TRUE(lexmin.unboundedDomain.isIntegerEmpty());
-  EXPECT_TRUE(lexmin.lexmin.isEqual(expectedLexmin));
+  EXPECT_TRUE(lexmin.lexopt.isEqual(expectedLexmin));
+}
+
+TEST(IntegerRelationTest, symbolicLexmax) {
+  SymbolicLexOpt lexmax =
+      parseRelationFromSet("(a, x)[b] : (a - x >= 0, b - x >= 0)", 1)
+          .findSymbolicIntegerLexMax();
+
+  PWMAFunction expectedLexmax = parsePWMAF({
+      {"(a)[b] : (a - b >= 0)", "(a)[b] -> (b)"},
+      {"(a)[b] : (b - a - 1 >= 0)", "(a)[b] -> (a)"},
+  });
+  EXPECT_TRUE(lexmax.unboundedDomain.isIntegerEmpty());
+  EXPECT_TRUE(lexmax.lexopt.isEqual(expectedLexmax));
 }