[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

154

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

154

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/Simplex.h

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/Simplex.cpp

unsigned SimplexBase::addRow(ArrayRef<int64_t> coeffs, bool makeRestricted) {
   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;
+  normalizeRange(tableau.getRow(row));
 }
 
 namespace {
 bool signMatchesDirection(int64_t elem, Direction direction) {
   assert(elem != 0 && "elem should not be 0");
   return direction == Direction::Up ? elem > 0 : elem < 0;
 }
 
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.

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::greatestCommonDivisor(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;
 }