[MLIR][Presburger][Simplex] symbolic lexmin: add some normalization heuristics

[arjunp]

Summary

Normalize some of the division and inequality expressions used,
which can improve performance. Also deduplicate some of the
normalization functionality throughout the Presburger library.

[Groverkss]

mlir/lib/Analysis/Presburger/Simplex.cpp

150–151

Why not just use Matrix::normalizeRow?

https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Analysis/Presburger/Matrix.h#L127

[arjunp]

mlir/lib/Analysis/Presburger/Simplex.cpp

150–151

Good point. I removed SimplexBase::normalizeRow and replaced it with Matrix::normalizeRow. I also replaced Matrix::normalizeRow's implementation with just a call through to normalizeRange

[arjunp]

Address comments

[Groverkss]

LGTM. The Debian build failure seems unrelated.

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

public:
   /// Negate the specified column.
   void negateColumn(unsigned column);
 
   /// Negate the specified row.
   void negateRow(unsigned row);
 
   /// Divide the first `nCols` of the specified row by their GCD.
   /// Returns the GCD of the first `nCols` of the specified row.
-  uint64_t normalizeRow(unsigned row, unsigned nCols);
+  int64_t normalizeRow(unsigned row, unsigned nCols);
   /// Divide the columns of the specified row by their GCD.
   /// Returns the GCD of the columns of the specified row.
-  uint64_t normalizeRow(unsigned row);
+  int64_t normalizeRow(unsigned row);
 
   /// The given vector is interpreted as a row vector v. Post-multiply v with
   /// this matrix, say M, and return vM.
   SmallVector<int64_t, 8> preMultiplyWithRow(ArrayRef<int64_t> rowVec) const;
 
   /// The given vector is interpreted as a column vector v. Pre-multiply v with
   /// this matrix, say M, and return Mv.
   SmallVector<int64_t, 8>

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

protected:
 
   /// Add a new row to the tableau and the associated data structures.
   /// 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);
 
   /// Enum to denote operations that need to be undone during rollback.
   enum class UndoLogEntry {
     RemoveLastConstraint,
     RemoveLastVariable,
private:
   /// in the initial domain.)
   LogicalResult addSymbolicCut(unsigned row);
 
   /// Get the numerator of the symbolic sample of the specific row.
   /// This is an affine expression in the symbols with integer coefficients.
   /// The last element is the constant term. This ignores the big M coefficient.
   SmallVector<int64_t, 8> getSymbolicSampleNumerator(unsigned row) const;
 
+  /// Get an affine inequality in the symbols with integer coefficients that
+  /// holds iff the symbolic sample of the specified row is non-negative.
+  SmallVector<int64_t, 8> getSymbolicSampleIneq(unsigned row) const;
+
   /// Return whether all the coefficients of the symbolic sample are integers.
   ///
   /// This does not consult the domain to check if the specified expression
   /// is always integral despite coefficients being fractional.
   bool isSymbolicSampleIntegral(unsigned row) const;
 
   /// Record a lexmin. The tableau must be consistent with all variables
   /// having symbolic samples with integer coefficients.

mlir/include/mlir/Analysis/Presburger/Utils.h

 /// it's each instance is replaced by `i^th` division. If it returns `false`,
 /// the divisions are not merged. `merge` can also do side effects, For example
 /// it can merge the local identifiers in IntegerRelation.
 void removeDuplicateDivs(
     std::vector<SmallVector<int64_t, 8>> &divs,
     SmallVectorImpl<unsigned> &denoms, unsigned localOffset,
     llvm::function_ref<bool(unsigned i, unsigned j)> merge);
 
+/// Compute the gcd of the range.
+int64_t gcdRange(ArrayRef<int64_t> range);
+
+/// Divide the range by its gcd and return the gcd.
+int64_t normalizeRange(MutableArrayRef<int64_t> range);
+
+/// Normalize the given (numerator, denominator) pair by dividing out the
+/// common factors between them. The numerator here is an affine expression
+/// with integer coefficients.
+void normalizeDiv(MutableArrayRef<int64_t> num, int64_t &denom);
+
 /// Return `coeffs` with all the elements negated.
 SmallVector<int64_t, 8> getNegatedCoeffs(ArrayRef<int64_t> coeffs);
 
 /// Return the complement of the given inequality.
 ///
 /// The complement of a_1 x_1 + ... + a_n x_ + c >= 0 is
 /// a_1 x_1 + ... + a_n x_ + c < 0, i.e., -a_1 x_1 - ... - a_n x_ - c - 1 >= 0,
 /// since all the variables are constrained to be integers.
 SmallVector<int64_t, 8> getComplementIneq(ArrayRef<int64_t> ineq);
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_UTILS_H

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

 /// fast method - linear in the number of coefficients.
 // Example on how this affects practical cases: consider the scenario:
 // 64*i >= 100, j = 64*i; without a tightening, elimination of i would yield
 // j >= 100 instead of the tighter (exact) j >= 128.
 void IntegerRelation::gcdTightenInequalities() {
   unsigned numCols = getNumCols();
   for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
     // Normalize the constraint and tighten the constant term by the GCD.
-    uint64_t gcd = inequalities.normalizeRow(i, getNumCols() - 1);
+    int64_t gcd = inequalities.normalizeRow(i, getNumCols() - 1);
     if (gcd > 1)
       atIneq(i, numCols - 1) = mlir::floorDiv(atIneq(i, numCols - 1), gcd);
   }
 }
 
 // Eliminates all identifier variables in column range [posStart, posLimit).
 // Returns the number of variables eliminated.
 unsigned IntegerRelation::gaussianEliminateIds(unsigned posStart,

mlir/lib/Analysis/Presburger/Matrix.cpp

 //===- Matrix.cpp - MLIR Matrix Class -------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/Support/MathExtras.h"
 
 using namespace mlir;
 using namespace presburger;
 
 Matrix::Matrix(unsigned rows, unsigned columns, unsigned reservedRows,
                unsigned reservedColumns)
     : nRows(rows), nColumns(columns),
for (unsigned row = 0, e = getNumRows(); row < e; ++row)
     at(row, column) = -at(row, column);
 }
 
 void Matrix::negateRow(unsigned row) {
   for (unsigned column = 0, e = getNumColumns(); column < e; ++column)
     at(row, column) = -at(row, column);
 }
 
-uint64_t Matrix::normalizeRow(unsigned row, unsigned cols) {
-  if (cols == 0)
-    return 0;
-
-  int64_t gcd = std::abs(at(row, 0));
-  for (unsigned j = 1, e = cols; j < e; ++j)
-    gcd = llvm::GreatestCommonDivisor64(gcd, std::abs(at(row, j)));
-
-  if (gcd > 1)
-    for (unsigned j = 0, e = cols; j < e; ++j)
-      at(row, j) /= gcd;
-
-  return gcd;
+int64_t Matrix::normalizeRow(unsigned row, unsigned cols) {
+  return normalizeRange(getRow(row).slice(0, cols));
 }
 
-uint64_t Matrix::normalizeRow(unsigned row) {
+int64_t Matrix::normalizeRow(unsigned row) {
   return normalizeRow(row, getNumColumns());
 }
 
 SmallVector<int64_t, 8>
 Matrix::preMultiplyWithRow(ArrayRef<int64_t> rowVec) const {
   assert(rowVec.size() == getNumRows() && "Invalid row vector dimension!");
 
   SmallVector<int64_t, 8> result(getNumColumns(), 0);

mlir/lib/Analysis/Presburger/Simplex.cpp

for (unsigned i = 0; i < var.size(); ++i) {
     int64_t nRowCoeff = lcm / tableau(nRow - 1, 0);
     int64_t idxRowCoeff = coeffs[i] * (lcm / tableau(pos, 0));
     tableau(nRow - 1, 0) = lcm;
     for (unsigned col = 1; col < nCol; ++col)
       tableau(nRow - 1, col) =
           nRowCoeff * tableau(nRow - 1, col) + idxRowCoeff * tableau(pos, col);
   }
 
-  normalizeRow(nRow - 1);
+  tableau.normalizeRow(nRow - 1);
   // Push to undo log along with the index of the new constraint.
   return con.size() - 1;
 }
 
-/// Normalize the row by removing factors that are common between the
-/// denominator and all the numerator coefficients.
-void SimplexBase::normalizeRow(unsigned row) {
-  int64_t gcd = 0;
-  for (unsigned col = 0; col < nCol; ++col) {
-    gcd = llvm::greatestCommonDivisor(gcd, std::abs(tableau(row, col)));
-    // If the gcd becomes 1 then the row is already normalized.
-    if (gcd == 1)
-      return;
-  }
-
-  // Note that the gcd can never become zero since the first element of the row,
-  // the denominator, is non-zero.
-  assert(gcd != 0);
-  for (unsigned col = 0; col < nCol; ++col)
-    tableau(row, col) /= gcd;
-}
-
 namespace {
 bool signMatchesDirection(int64_t elem, Direction direction) {
   assert(elem != 0 && "elem should not be 0");
   return direction == Direction::Up ? elem > 0 : elem < 0;
 }
 
 Direction flippedDirection(Direction direction) {
   return direction == Direction::Up ? Direction::Down : Simplex::Direction::Up;
 }
SymbolicLexSimplex::getSymbolicSampleNumerator(unsigned row) const {
   SmallVector<int64_t, 8> sample;
   sample.reserve(nSymbol + 1);
   for (unsigned col = 3; col < 3 + nSymbol; ++col)
     sample.push_back(tableau(row, col));
   sample.push_back(tableau(row, 1));
   return sample;
 }
 
+SmallVector<int64_t, 8>
+SymbolicLexSimplex::getSymbolicSampleIneq(unsigned row) const {
+  SmallVector<int64_t, 8> sample = getSymbolicSampleNumerator(row);
+  // The inequality is equivalent to the GCD-normalized one.
+  normalizeRange(sample);
+  return sample;
+}
+
 void LexSimplexBase::appendSymbol() {
   appendVariable();
   swapColumns(3 + nSymbol, nCol - 1);
   var.back().isSymbol = true;
   nSymbol++;
 }
 
 static bool isRangeDivisibleBy(ArrayRef<int64_t> range, int64_t divisor) {

 /// (sum_i (b_i%d)y_i - (-c%d) - sum_i (-a_i%d)s_i + q*d)/d >= 0
 /// This constraint is violated when added so we immediately try to move it to a
 /// column.
 LogicalResult SymbolicLexSimplex::addSymbolicCut(unsigned row) {
   int64_t d = tableau(row, 0);
 
   // Add the division variable `q` described above to the symbol domain.
   // q = ((-c%d) + sum_i (-a_i%d)s_i)/d.
-  SmallVector<int64_t, 8> domainDivCoeffs;
-  domainDivCoeffs.reserve(nSymbol + 1);
+  SmallVector<int64_t, 8> divCoeffs;
+  divCoeffs.reserve(nSymbol + 1);
+  int64_t divDenom = d;
   for (unsigned col = 3; col < 3 + nSymbol; ++col)
-    domainDivCoeffs.push_back(mod(-tableau(row, col), d)); // (-a_i%d)s_i
-  domainDivCoeffs.push_back(mod(-tableau(row, 1), d));     // -c%d.
+    divCoeffs.push_back(mod(-tableau(row, col), divDenom)); // (-a_i%d)s_i
+  divCoeffs.push_back(mod(-tableau(row, 1), divDenom));     // -c%d.
 
-  domainSimplex.addDivisionVariable(domainDivCoeffs, d);
-  domainPoly.addLocalFloorDiv(domainDivCoeffs, d);
+  normalizeDiv(divCoeffs, divDenom);
+  domainSimplex.addDivisionVariable(divCoeffs, divDenom);
+  domainPoly.addLocalFloorDiv(divCoeffs, divDenom);
 
   // Update `this` to account for the additional symbol we just added.
   appendSymbol();
 
   // Add the cut (sum_i (b_i%d)y_i - (-c%d) + sum_i -(-a_i%d)s_i + q*d)/d >= 0.
   addZeroRow(/*makeRestricted=*/true);
   tableau(nRow - 1, 0) = d;
   tableau(nRow - 1, 2) = 0;
Optional<unsigned> SymbolicLexSimplex::maybeGetAlwaysViolatedRow() {
   // coefficient, without needing to perform any simplex queries on the domain.
   for (unsigned row = 0; row < nRow; ++row)
     if (tableau(row, 2) < 0)
       return row;
 
   for (unsigned row = 0; row < nRow; ++row) {
     if (tableau(row, 2) > 0)
       continue;
-    if (domainSimplex.isSeparateInequality(getSymbolicSampleNumerator(row))) {
+    if (domainSimplex.isSeparateInequality(getSymbolicSampleIneq(row))) {
       // Sample numerator always takes negative values in the symbol domain.
       return row;
     }
   }
   return {};
 }
 
 Optional<unsigned> SymbolicLexSimplex::maybeGetNonIntegralVarRow() {
if (level > stack.size()) {
       unsigned splitRow;
       SmallVector<int64_t, 8> symbolicSample;
       for (splitRow = 0; splitRow < nRow; ++splitRow) {
         if (tableau(splitRow, 2) > 0)
           continue;
         assert(tableau(splitRow, 2) == 0 &&
                "Non-branching pivots should have been handled already!");
 
-        symbolicSample = getSymbolicSampleNumerator(splitRow);
+        symbolicSample = getSymbolicSampleIneq(splitRow);
         if (domainSimplex.isRedundantInequality(symbolicSample))
           continue;
 
         // It's neither redundant nor separate, so it takes both positive and
         // negative values, and hence constitutes a row for which we need to
         // split the domain and separately run each case.
         assert(!domainSimplex.isSeparateInequality(symbolicSample) &&
                "Non-branching pivots should have been handled already!");
if (level == stack.size()) {
       // Drop the frame. We don't need it anymore.
       stack.pop_back();
 
       // Now we consider the part of the domain where the unknown `splitIndex`
       // was negative.
       assert(u.orientation == Orientation::Row &&
              "The split row should have been returned to row orientation!");
       SmallVector<int64_t, 8> splitIneq =
-          getComplementIneq(getSymbolicSampleNumerator(u.pos));
+          getComplementIneq(getSymbolicSampleIneq(u.pos));
+      normalizeRange(splitIneq);
       if (moveRowUnknownToColumn(u.pos).failed()) {
         // The unknown can't be made non-negative; return.
         --level;
         continue;
       }
 
       // The unknown can be made negative; recurse with the corresponding domain
       // constraints.
if (tableau(pivotRow, 0) < 0) {
     tableau(pivotRow, pivotCol) = -tableau(pivotRow, pivotCol);
   } else {
     for (unsigned col = 1; col < nCol; ++col) {
       if (col == pivotCol)
         continue;
       tableau(pivotRow, col) = -tableau(pivotRow, col);
     }
   }
-  normalizeRow(pivotRow);
+  tableau.normalizeRow(pivotRow);
 
   for (unsigned row = 0; row < nRow; ++row) {
     if (row == pivotRow)
       continue;
     if (tableau(row, pivotCol) == 0) // Nothing to do.
       continue;
     tableau(row, 0) *= tableau(pivotRow, 0);
     for (unsigned j = 1; j < nCol; ++j) {
       if (j == pivotCol)
         continue;
       // Add rather than subtract because the pivot row has been negated.
       tableau(row, j) = tableau(row, j) * tableau(pivotRow, 0) +
                         tableau(row, pivotCol) * tableau(pivotRow, j);
     }
     tableau(row, pivotCol) *= tableau(pivotRow, pivotCol);
-    normalizeRow(row);
+    tableau.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 Simplex::restoreRow(Unknown &u) {
   assert(u.orientation == Orientation::Row &&

mlir/lib/Analysis/Presburger/Utils.cpp

for (unsigned j = i + 1; j < divs.size(); ++j) {
       divs.erase(divs.begin() + j);
       denoms.erase(denoms.begin() + j);
       // Since `j` can never be zero, we do not need to worry about overflows.
       --j;
     }
   }
 }
 
+int64_t presburger::gcdRange(ArrayRef<int64_t> range) {
+  int64_t gcd = 0;
+  for (int64_t elem : range) {
+    gcd = llvm::GreatestCommonDivisor64(gcd, std::abs(elem));
+    if (gcd == 1)
+      return gcd;
+  }
+  return gcd;
+}
+
+int64_t presburger::normalizeRange(MutableArrayRef<int64_t> range) {
+  int64_t gcd = gcdRange(range);
+  if (gcd == 0 || gcd == 1)
+    return gcd;
+  for (int64_t &elem : range)
+    elem /= gcd;
+  return gcd;
+}
+
+void presburger::normalizeDiv(MutableArrayRef<int64_t> num, int64_t &denom) {
+  int64_t gcd = llvm::greatestCommonDivisor(gcdRange(num), denom);
+  for (int64_t &coeff : num)
+    coeff /= gcd;
+  denom /= gcd;
+}
+
 SmallVector<int64_t, 8> presburger::getNegatedCoeffs(ArrayRef<int64_t> coeffs) {
   SmallVector<int64_t, 8> negatedCoeffs;
   negatedCoeffs.reserve(coeffs.size());
   for (int64_t coeff : coeffs)
     negatedCoeffs.emplace_back(-coeff);
   return negatedCoeffs;
 }
 
 SmallVector<int64_t, 8> presburger::getComplementIneq(ArrayRef<int64_t> ineq) {
   SmallVector<int64_t, 8> coeffs;
   coeffs.reserve(ineq.size());
   for (int64_t coeff : ineq)
     coeffs.emplace_back(-coeff);
   --coeffs.back();
   return coeffs;
 }