[MLIR][Presburger] Use Matrix utilities for IntegerPolyhedron

[Groverkss]

Summary

This patch replaces various functions over inequalities/equalities in
IntegerPolyhedron with Matrix functions already implementing them or refactors
them to a Matrix function.

[Groverkss]

• Replace swapping with Matrix functions

[arjunp]

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

126

nCols would be clearer

126

The meaning of normalize is not clear in this context. Maybe just write, divide the first nCols cols of row by its gcd.

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

836–837

The order of these is switched now, please update the comment accordingly

838

Can just call it gcd now, there is no other gcd left

mlir/lib/Analysis/Presburger/Matrix.cpp

216–222

Why not just declare gcd as an int64_t and use llvm::greatestCommonDivisor?

[Groverkss]

• Address Arjun's comments

[arjunp]

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

127

Please update this comment too

[Groverkss]

• Fix comment for normalizeRow

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

public:
   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);
 
+  /// 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);
+  /// 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);
+
   /// 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>
   postMultiplyWithColumn(ArrayRef<int64_t> colVec) const;

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

 
 void IntegerPolyhedron::swapId(unsigned posA, unsigned posB) {
   assert(posA < getNumIds() && "invalid position A");
   assert(posB < getNumIds() && "invalid position B");
 
   if (posA == posB)
     return;
 
-  for (unsigned r = 0, e = getNumInequalities(); r < e; r++)
-    std::swap(atIneq(r, posA), atIneq(r, posB));
-  for (unsigned r = 0, e = getNumEqualities(); r < e; r++)
-    std::swap(atEq(r, posA), atEq(r, posB));
+  inequalities.swapColumns(posA, posB);
+  equalities.swapColumns(posA, posB);
 }
 
 void IntegerPolyhedron::clearConstraints() {
   equalities.resizeVertically(0);
   inequalities.resizeVertically(0);
 }
 
 /// Gather all lower and upper bounds of the identifier at `pos`, and
void IntegerPolyhedron::setAndEliminate(unsigned pos,
                                         ArrayRef<int64_t> values) {
   if (values.empty())
     return;
   assert(pos + values.size() <= getNumIds() &&
          "invalid position or too many values");
   // Setting x_j = p in sum_i a_i x_i + c is equivalent to adding p*a_j to the
   // constant term and removing the id x_j. We do this for all the ids
   // pos, pos + 1, ... pos + values.size() - 1.
-  for (unsigned r = 0, e = getNumInequalities(); r < e; r++)
-    for (unsigned i = 0, numVals = values.size(); i < numVals; ++i)
-      atIneq(r, getNumCols() - 1) += atIneq(r, pos + i) * values[i];
-  for (unsigned r = 0, e = getNumEqualities(); r < e; r++)
-    for (unsigned i = 0, numVals = values.size(); i < numVals; ++i)
-      atEq(r, getNumCols() - 1) += atEq(r, pos + i) * values[i];
+  unsigned constantColPos = getNumCols() - 1;
+    inequalities.addToColumn(i + pos, constantColPos, values[i]);
+    equalities.addToColumn(i + pos, constantColPos, values[i]);
   removeIdRange(pos, pos + values.size());
 }
 
 void IntegerPolyhedron::clearAndCopyFrom(const IntegerPolyhedron &other) {
   *this = other;
 }
 
 // Searches for a constraint with a non-zero coefficient at `colIdx` in
bool IntegerPolyhedron::findConstraintWithNonZeroAt(unsigned colIdx, bool isEq,
   for (*rowIdx = 0; *rowIdx < e; ++(*rowIdx)) {
     if (at(*rowIdx) != 0) {
       return true;
     }
   }
   return false;
 }
 
-// Normalizes the coefficient values across all columns in `rowIdx` by their
-// GCD in equality or inequality constraints as specified by `isEq`.
-template <bool isEq>
-static void normalizeConstraintByGCD(IntegerPolyhedron *constraints,
-                                     unsigned rowIdx) {
-  auto at = [&](unsigned colIdx) -> int64_t {
-    return isEq ? constraints->atEq(rowIdx, colIdx)
-                : constraints->atIneq(rowIdx, colIdx);
-  };
-  uint64_t gcd = std::abs(at(0));
-  for (unsigned j = 1, e = constraints->getNumCols(); j < e; ++j) {
-    gcd = llvm::GreatestCommonDivisor64(gcd, std::abs(at(j)));
-  }
-  if (gcd > 0 && gcd != 1) {
-    for (unsigned j = 0, e = constraints->getNumCols(); j < e; ++j) {
-      int64_t v = at(j) / static_cast<int64_t>(gcd);
-      isEq ? constraints->atEq(rowIdx, j) = v
-           : constraints->atIneq(rowIdx, j) = v;
-    }
-  }
-}
-
 void IntegerPolyhedron::normalizeConstraintsByGCD() {
-  for (unsigned i = 0, e = getNumEqualities(); i < e; ++i) {
-    normalizeConstraintByGCD</*isEq=*/true>(this, i);
-  }
-  for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
-    normalizeConstraintByGCD</*isEq=*/false>(this, i);
-  }
+  for (unsigned i = 0, e = getNumEqualities(); i < e; ++i)
+    equalities.normalizeRow(i);
+  for (unsigned i = 0, e = getNumInequalities(); i < e; ++i)
+    inequalities.normalizeRow(i);
 }
 
 bool IntegerPolyhedron::hasInvalidConstraint() const {
   assert(hasConsistentState());
   auto check = [&](bool isEq) -> bool {
     unsigned numCols = getNumCols();
     unsigned numRows = isEq ? getNumEqualities() : getNumInequalities();
     for (unsigned i = 0, e = numRows; i < e; ++i) {

 /// be reduced to the preceding multiple of the GCD of the coefficients, i.e.,
 ///  64*i - 100 >= 0  =>  64*i - 128 >= 0 (since 'i' is an integer). This is a
 /// 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 IntegerPolyhedron::gcdTightenInequalities() {
   unsigned numCols = getNumCols();
   for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
-    uint64_t gcd = std::abs(atIneq(i, 0));
-    for (unsigned j = 1; j < numCols - 1; ++j) {
-      gcd = llvm::GreatestCommonDivisor64(gcd, std::abs(atIneq(i, j)));
-    }
-    if (gcd > 0 && gcd != 1) {
-      int64_t gcdI = static_cast<int64_t>(gcd);
-      // Tighten the constant term and normalize the constraint by the GCD.
-      atIneq(i, numCols - 1) = mlir::floorDiv(atIneq(i, numCols - 1), gcdI);
-      for (unsigned j = 0, e = numCols - 1; j < e; ++j)
-        atIneq(i, j) /= gcdI;
-    }
+    // Normalize the constraint and tighten the constant term by the GCD.
+    uint64_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 IntegerPolyhedron::gaussianEliminateIds(unsigned posStart,
                                                  unsigned posLimit) {
   // Return if identifier positions to eliminate are out of range.
if (!findConstraintWithNonZeroAt(pivotCol, /*isEq=*/true, &pivotRow)) {
       }
       break;
     }
 
     // Eliminate identifier at 'pivotCol' from each equality row.
     for (unsigned i = 0, e = getNumEqualities(); i < e; ++i) {
       eliminateFromConstraint(this, i, pivotRow, pivotCol, posStart,
                               /*isEq=*/true);
-      normalizeConstraintByGCD</*isEq=*/true>(this, i);
+      equalities.normalizeRow(i);
     }
 
     // Eliminate identifier at 'pivotCol' from each inequality row.
     for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
       eliminateFromConstraint(this, i, pivotRow, pivotCol, posStart,
                               /*isEq=*/false);
-      normalizeConstraintByGCD</*isEq=*/false>(this, i);
+      inequalities.normalizeRow(i);
     }
     removeEquality(pivotRow);
     gcdTightenInequalities();
   }
   // Update position limit based on number eliminated.
   posLimit = pivotCol;
   // Remove eliminated columns from all constraints.
   removeIdRange(posStart, posLimit);
   return posLimit - posStart;
 }
 
-// Fills an inequality row with the value 'val'.
-static inline void fillInequality(IntegerPolyhedron *cst, unsigned r,
-                                  int64_t val) {
-  for (unsigned c = 0, f = cst->getNumCols(); c < f; c++) {
-    cst->atIneq(r, c) = val;
-  }
-}
-
-// Negates an inequality.
-static inline void negateInequality(IntegerPolyhedron *cst, unsigned r) {
-  for (unsigned c = 0, f = cst->getNumCols(); c < f; c++) {
-    cst->atIneq(r, c) = -cst->atIneq(r, c);
-  }
-}
-
 // A more complex check to eliminate redundant inequalities. Uses FourierMotzkin
 // to check if a constraint is redundant.
 void IntegerPolyhedron::removeRedundantInequalities() {
   SmallVector<bool, 32> redun(getNumInequalities(), false);
   // To check if an inequality is redundant, we replace the inequality by its
   // complement (for eg., i - 1 >= 0 by i <= 0), and check if the resulting
   // system is empty. If it is, the inequality is redundant.
   IntegerPolyhedron tmpCst(*this);
   for (unsigned r = 0, e = getNumInequalities(); r < e; r++) {
     // Change the inequality to its complement.
-    negateInequality(&tmpCst, r);
+    tmpCst.inequalities.negateRow(r);
     tmpCst.atIneq(r, tmpCst.getNumCols() - 1)--;
     if (tmpCst.isEmpty()) {
       redun[r] = true;
       // Zero fill the redundant inequality.
-      fillInequality(this, r, /*val=*/0);
-      fillInequality(&tmpCst, r, /*val=*/0);
+      inequalities.fillRow(r, /*value=*/0);
+      tmpCst.inequalities.fillRow(r, /*value=*/0);
     } else {
       // Reverse the change (to avoid recreating tmpCst each time).
       tmpCst.atIneq(r, tmpCst.getNumCols() - 1)++;
-      negateInequality(&tmpCst, r);
+      tmpCst.inequalities.negateRow(r);
     }
   }
 
-  // Scan to get rid of all rows marked redundant, in-place.
-  auto copyRow = [&](unsigned src, unsigned dest) {
-    if (src == dest)
-      return;
-    for (unsigned c = 0, e = getNumCols(); c < e; c++) {
-      atIneq(dest, c) = atIneq(src, c);
-    }
-  };
   unsigned pos = 0;
-  for (unsigned r = 0, e = getNumInequalities(); r < e; r++) {
+  for (unsigned r = 0, e = getNumInequalities(); r < e; ++r) {
     if (!redun[r])
-      copyRow(r, pos++);
+      inequalities.copyRow(r, pos++);
   }
   inequalities.resizeVertically(pos);
 }
 
 // A more complex check to eliminate redundant inequalities and equalities. Uses
 // Simplex to check if a constraint is redundant.
 void IntegerPolyhedron::removeRedundantConstraints() {
   // First, we run gcdTightenInequalities. This allows us to catch some
   // constraints which are not redundant when considering rational solutions
   // but are redundant in terms of integer solutions.
   gcdTightenInequalities();
   Simplex simplex(*this);
   simplex.detectRedundant();
 
-  auto copyInequality = [&](unsigned src, unsigned dest) {
-    if (src == dest)
-      return;
-    for (unsigned c = 0, e = getNumCols(); c < e; c++)
-      atIneq(dest, c) = atIneq(src, c);
-  };
   unsigned pos = 0;
   unsigned numIneqs = getNumInequalities();
   // Scan to get rid of all inequalities marked redundant, in-place. In Simplex,
   // the first constraints added are the inequalities.
   for (unsigned r = 0; r < numIneqs; r++) {
     if (!simplex.isMarkedRedundant(r))
-      copyInequality(r, pos++);
+      inequalities.copyRow(r, pos++);
   }
   inequalities.resizeVertically(pos);
 
   // Scan to get rid of all equalities marked redundant, in-place. In Simplex,
   // after the inequalities, a pair of constraints for each equality is added.
   // An equality is redundant if both the inequalities in its pair are
   // redundant.
-  auto copyEquality = [&](unsigned src, unsigned dest) {
-    if (src == dest)
-      return;
-    for (unsigned c = 0, e = getNumCols(); c < e; c++)
-      atEq(dest, c) = atEq(src, c);
-  };
   pos = 0;
   for (unsigned r = 0, e = getNumEqualities(); r < e; r++) {
     if (!(simplex.isMarkedRedundant(numIneqs + 2 * r) &&
           simplex.isMarkedRedundant(numIneqs + 2 * r + 1)))
-      copyEquality(r, pos++);
+      equalities.copyRow(r, pos++);
   }
   equalities.resizeVertically(pos);
 }
 
 Optional<uint64_t> IntegerPolyhedron::computeVolume() const {
   assert(getNumSymbolIds() == 0 && "Symbols are not yet supported!");
 
   Simplex simplex(*this);

 /// variable and `affine-expr` is an affine expression not containing `e`.
 /// If an equality satisfies this form, the local variable is replaced in
 /// each constraint and then removed. The equality used to replace this local
 /// variable is also removed.
 void IntegerPolyhedron::removeRedundantLocalVars() {
   // Normalize the equality constraints to reduce coefficients of local
   // variables to 1 wherever possible.
   for (unsigned i = 0, e = getNumEqualities(); i < e; ++i)
-    normalizeConstraintByGCD</*isEq=*/true>(this, i);
+    equalities.normalizeRow(i);
 
   while (true) {
     unsigned i, e, j, f;
     for (i = 0, e = getNumEqualities(); i < e; ++i) {
       // Find a local variable to eliminate using ith equality.
       for (j = getNumDimAndSymbolIds(), f = getNumIds(); j < f; ++j)
         if (std::abs(atEq(i, j)) == 1)
           break;
 
       // Local variable can be eliminated using ith equality.
       if (j < f)
         break;
     }
 
     // No equality can be used to eliminate a local variable.
     if (i == e)
       break;
 
     // Use the ith equality to simplify other equalities. If any changes
     // are made to an equality constraint, it is normalized by GCD.
     for (unsigned k = 0, t = getNumEqualities(); k < t; ++k) {
       if (atEq(k, j) != 0) {
         eliminateFromConstraint(this, k, i, j, j, /*isEq=*/true);
-        normalizeConstraintByGCD</*isEq=*/true>(this, k);
+        equalities.normalizeRow(k);
       }
     }
 
     // Use the ith equality to simplify inequalities.
     for (unsigned k = 0, t = getNumInequalities(); k < t; ++k)
       eliminateFromConstraint(this, k, i, j, j, /*isEq=*/false);
 
     // Remove the ith equality and the found local variable.

mlir/lib/Analysis/Presburger/Matrix.cpp

for (unsigned row = 0, e = getNumRows(); row < e; ++row)
     at(row, targetColumn) += scale * at(row, sourceColumn);
 }
 
 void Matrix::negateColumn(unsigned column) {
   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;
+}
+
+uint64_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);
   for (unsigned col = 0, e = getNumColumns(); col < e; ++col)
     for (unsigned i = 0, e = getNumRows(); i < e; ++i)
       result[col] += rowVec[i] * at(i, col);