[MLIR][Presburger] Template Matrix to allow MPInt and Fraction and separate out IntMatrix

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

@@ -366,7 +366,7 @@ class IntegerRelation {
   /// bounded. The span of the returned vectors is guaranteed to contain all
   /// such vectors. The returned vectors are NOT guaranteed to be linearly
   /// independent. This function should not be called on empty sets.
-  Matrix<MPInt> getBoundedDirections() const;
+  IntMatrix getBoundedDirections() const;
 
   /// Find an integer sample point satisfying the constraints using a
   /// branch and bound algorithm with generalized basis reduction, with some
@@ -792,10 +792,10 @@ class IntegerRelation {
   PresburgerSpace space;
 
   /// Coefficients of affine equalities (in == 0 form).
-  Matrix<MPInt> equalities;
+  IntMatrix equalities;
 
   /// Coefficients of affine inequalities (in >= 0 form).
-  Matrix<MPInt> inequalities;
+  IntMatrix inequalities;
 };
 
 /// An IntegerPolyhedron represents the set of points from a PresburgerSpace

mlir/include/mlir/Analysis/Presburger/LinearTransform.h

@@ -22,8 +22,8 @@ namespace presburger {
 
 class LinearTransform {
 public:
-  explicit LinearTransform(Matrix<MPInt> &&oMatrix);
-  explicit LinearTransform(const Matrix<MPInt> &oMatrix);
+  explicit LinearTransform(IntMatrix &&oMatrix);
+  explicit LinearTransform(const IntMatrix &oMatrix);
 
   // Returns a linear transform T such that MT is M in column echelon form.
   // Also returns the number of non-zero columns in MT.
@@ -32,7 +32,7 @@ class LinearTransform {
   // strictly below that of the previous column, and all columns which have only
   // zeros are at the end.
   static std::pair<unsigned, LinearTransform>
-  makeTransformToColumnEchelon(const Matrix<MPInt> &m);
+  makeTransformToColumnEchelon(const IntMatrix &m);
 
   // Returns an IntegerRelation having a constraint vector vT for every
   // constraint vector v in rel, where T is this transform.
@@ -55,7 +55,7 @@ class LinearTransform {
   MPInt determinant();
 
 private:
-  Matrix<MPInt> matrix;
+  IntMatrix matrix;
 };
 
 } // namespace presburger

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

@@ -157,13 +157,6 @@ static_assert(std::is_same_v<T,MPInt> || std::is_same_v<T,Fraction>, "T must be
   /// 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.
-  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.
-  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<T, 8> preMultiplyWithRow(ArrayRef<T> rowVec) const;
@@ -172,18 +165,6 @@ static_assert(std::is_same_v<T,MPInt> || std::is_same_v<T,Fraction>, "T must be
   /// this matrix, say M, and return Mv.
   SmallVector<T, 8> postMultiplyWithColumn(ArrayRef<T> colVec) const;
 
-  /// Given the current matrix M, returns the matrices H, U such that H is the
-  /// column hermite normal form of M, i.e. H = M * U, where U is unimodular and
-  /// the matrix H has the following restrictions:
-  ///  - H is lower triangular.
-  ///  - The leading coefficient (the first non-zero entry from the top, called
-  ///    the pivot) of a non-zero column is always strictly below of the leading
-  ///    coefficient of the column before it; moreover, it is positive.
-  ///  - The elements to the right of the pivots are zero and the elements to
-  ///    the left of the pivots are nonnegative and strictly smaller than the
-  ///    pivot.
-  std::pair<Matrix, Matrix> computeHermiteNormalForm() const;
-
   /// Resize the matrix to the specified dimensions. If a dimension is smaller,
   /// the values are truncated; if it is bigger, the new values are initialized
   /// to zero.
@@ -221,7 +202,49 @@ static_assert(std::is_same_v<T,MPInt> || std::is_same_v<T,Fraction>, "T must be
   SmallVector<T, 16> data;
 };
 
+// An inherited class for integer matrices, with no new data attributes.
+// This is only used for the matrix-related methods which apply only
+// to integers (hermite normal form computation and row normalisation).
+class IntMatrix : public Matrix<MPInt>
+{
+public:
+  IntMatrix(unsigned rows, unsigned columns, unsigned reservedRows = 0,
+            unsigned reservedColumns = 0) :
+    Matrix<MPInt>(rows, columns, reservedRows, reservedColumns) {};
+
+  IntMatrix(Matrix<MPInt> m) :
+    Matrix<MPInt>(m.getNumRows(), m.getNumColumns(), m.getNumReservedRows(), m.getNumReservedColumns())
+  {
+    for (unsigned i = 0; i < m.getNumRows(); i++)
+      for (unsigned j = 0; j < m.getNumColumns(); j++)
+        at(i, j) = m(i, j);
+  };
+
+  /// Return the identity matrix of the specified dimension.
+  static IntMatrix identity(unsigned dimension);
+
+  /// Given the current matrix M, returns the matrices H, U such that H is the
+  /// column hermite normal form of M, i.e. H = M * U, where U is unimodular and
+  /// the matrix H has the following restrictions:
+  ///  - H is lower triangular.
+  ///  - The leading coefficient (the first non-zero entry from the top, called
+  ///    the pivot) of a non-zero column is always strictly below of the leading
+  ///    coefficient of the column before it; moreover, it is positive.
+  ///  - The elements to the right of the pivots are zero and the elements to
+  ///    the left of the pivots are nonnegative and strictly smaller than the
+  ///    pivot.
+  std::pair<IntMatrix, IntMatrix> computeHermiteNormalForm() const;
+
+  /// Divide the first `nCols` of the specified row by their GCD.
+  /// Returns the GCD of the first `nCols` of the specified row.
+  MPInt 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.
+  MPInt normalizeRow(unsigned row);
+
+};
+
 } // namespace presburger
 } // namespace mlir
 
-#endif // MLIR_ANALYSIS_PRESBURGER_MATRIX_H
+#endif // MLIR_ANALYSIS_PRESBURGER_MATRIX_H

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

@@ -40,13 +40,13 @@ enum class OrderingKind { EQ, NE, LT, LE, GT, GE };
 /// value of the function at a specified point.
 class MultiAffineFunction {
 public:
-  MultiAffineFunction(const PresburgerSpace &space, const Matrix<MPInt> &output)
+  MultiAffineFunction(const PresburgerSpace &space, const IntMatrix &output)
       : space(space), output(output),
         divs(space.getNumVars() - space.getNumRangeVars()) {
     assertIsConsistent();
   }
 
-  MultiAffineFunction(const PresburgerSpace &space, const Matrix<MPInt> &output,
+  MultiAffineFunction(const PresburgerSpace &space, const IntMatrix &output,
                       const DivisionRepr &divs)
       : space(space), output(output), divs(divs) {
     assertIsConsistent();
@@ -65,7 +65,7 @@ class MultiAffineFunction {
   PresburgerSpace getOutputSpace() const { return space.getRangeSpace(); }
 
   /// Get a matrix with each row representing row^th output expression.
-  const Matrix<MPInt> &getOutputMatrix() const { return output; }
+  const IntMatrix &getOutputMatrix() const { return output; }
   /// Get the `i^th` output expression.
   ArrayRef<MPInt> getOutputExpr(unsigned i) const { return output.getRow(i); }
 
@@ -124,7 +124,7 @@ class MultiAffineFunction {
   /// The function's output is a tuple of integers, with the ith element of the
   /// tuple defined by the affine expression given by the ith row of this output
   /// matrix.
-  Matrix<MPInt> output;
+  IntMatrix output;
 
   /// Storage for division representation for each local variable in space.
   DivisionRepr divs;

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

@@ -338,7 +338,7 @@ class SimplexBase {
   unsigned nSymbol;
 
   /// The matrix representing the tableau.
-  Matrix<MPInt> tableau;
+  IntMatrix tableau;
 
   /// This is true if the tableau has been detected to be empty, false
   /// otherwise.
@@ -861,7 +861,7 @@ class Simplex : public SimplexBase {
 
   /// Reduce the given basis, starting at the specified level, using general
   /// basis reduction.
-  void reduceBasis(Matrix<MPInt> &basis, unsigned level);
+  void reduceBasis(IntMatrix &basis, unsigned level);
 };
 
 /// Takes a snapshot of the simplex state on construction and rolls back to the

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

@@ -182,7 +182,7 @@ class DivisionRepr {
   /// Each row of the Matrix represents a single division dividend. The
   /// `i^th` row represents the dividend of the variable at `divOffset + i`
   /// in the constraint system (and the `i^th` local variable).
-  Matrix<MPInt> dividends;
+  IntMatrix dividends;
 
   /// Denominators of each division. If a denominator of a division is `0`, the
   /// division variable is considered to not have a division representation.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

@@ -304,7 +304,7 @@ SymbolicLexOpt IntegerRelation::findSymbolicIntegerLexMax() const {
   // Get lexmax by flipping range sign in the PWMA constraints.
   for (auto &flippedPiece :
        flippedSymbolicIntegerLexMax.lexopt.getAllPieces()) {
-    Matrix<MPInt> mat = flippedPiece.output.getOutputMatrix();
+    IntMatrix mat = flippedPiece.output.getOutputMatrix();
     for (unsigned i = 0, e = mat.getNumRows(); i < e; i++)
       mat.negateRow(i);
     MultiAffineFunction maf(flippedPiece.output.getSpace(), mat);
@@ -738,7 +738,7 @@ bool IntegerRelation::isEmptyByGCDTest() const {
 //
 // It is sufficient to check the perpendiculars of the constraints, as the set
 // of perpendiculars which are bounded must span all bounded directions.
-Matrix<MPInt> IntegerRelation::getBoundedDirections() const {
+IntMatrix IntegerRelation::getBoundedDirections() const {
   // Note that it is necessary to add the equalities too (which the constructor
   // does) even though we don't need to check if they are bounded; whether an
   // inequality is bounded or not depends on what other constraints, including
@@ -759,7 +759,7 @@ Matrix<MPInt> IntegerRelation::getBoundedDirections() const {
   // The direction vector is given by the coefficients and does not include the
   // constant term, so the matrix has one fewer column.
   unsigned dirsNumCols = getNumCols() - 1;
-  Matrix<MPInt> dirs(boundedIneqs.size() + getNumEqualities(), dirsNumCols);
+  IntMatrix dirs(boundedIneqs.size() + getNumEqualities(), dirsNumCols);
 
   // Copy the bounded inequalities.
   unsigned row = 0;
@@ -845,7 +845,7 @@ IntegerRelation::findIntegerSample() const {
   // m is a matrix containing, in each row, a vector in which S is
   // bounded, such that the linear span of all these dimensions contains all
   // bounded dimensions in S.
-  Matrix<MPInt> m = getBoundedDirections();
+  IntMatrix m = getBoundedDirections();
   // In column echelon form, each row of m occupies only the first rank(m)
   // columns and has zeros on the other columns. The transform T that brings S
   // to column echelon form is unimodular as well, so this is a suitable

mlir/lib/Analysis/Presburger/LinearTransform.cpp

@@ -12,11 +12,11 @@
 using namespace mlir;
 using namespace presburger;
 
-LinearTransform::LinearTransform(Matrix<MPInt> &&oMatrix) : matrix(oMatrix) {}
-LinearTransform::LinearTransform(const Matrix<MPInt> &oMatrix) : matrix(oMatrix) {}
+LinearTransform::LinearTransform(IntMatrix &&oMatrix) : matrix(oMatrix) {}
+LinearTransform::LinearTransform(const IntMatrix &oMatrix) : matrix(oMatrix) {}
 
 std::pair<unsigned, LinearTransform>
-LinearTransform::makeTransformToColumnEchelon(const Matrix<MPInt> &m) {
+LinearTransform::makeTransformToColumnEchelon(const IntMatrix &m) {
   // Compute the hermite normal form of m. This, is by definition, is in column
   // echelon form.
   auto [h, u] = m.computeHermiteNormalForm();

mlir/lib/Analysis/Presburger/Matrix.cpp

@@ -223,14 +223,6 @@ template <typename T> void Matrix<T>::negateRow(unsigned row) {
     at(row, column) = -at(row, column);
 }
 
-template <> MPInt Matrix<MPInt>::normalizeRow(unsigned row, unsigned cols) {
-  return normalizeRange(getRow(row).slice(0, cols));
-}
-
-template <> MPInt Matrix<MPInt>::normalizeRow(unsigned row) {
-  return normalizeRow(row, getNumColumns());
-}
-
 template <typename T> SmallVector<T, 8> Matrix<T>::preMultiplyWithRow(ArrayRef<T> rowVec) const {
   assert(rowVec.size() == getNumRows() && "Invalid row vector dimension!");
 
@@ -267,12 +259,53 @@ static void modEntryColumnOperation(Matrix<MPInt> &m, unsigned row, unsigned sou
   otherMatrix.addToColumn(sourceCol, targetCol, ratio);
 }
 
-template <> std::pair<Matrix<MPInt>, Matrix<MPInt>> Matrix<MPInt>::computeHermiteNormalForm() const {
+template <typename T> void Matrix<T>::print(raw_ostream &os) const {
+  for (unsigned row = 0; row < nRows; ++row) {
+    for (unsigned column = 0; column < nColumns; ++column)
+      os << at(row, column) << ' ';
+    os << '\n';
+  }
+}
+
+template <typename T> void Matrix<T>::dump() const { print(llvm::errs()); }
+
+template <typename T> bool Matrix<T>::hasConsistentState() const {
+  if (data.size() != nRows * nReservedColumns)
+    return false;
+  if (nColumns > nReservedColumns)
+    return false;
+#ifdef EXPENSIVE_CHECKS
+  for (unsigned r = 0; r < nRows; ++r)
+    for (unsigned c = nColumns; c < nReservedColumns; ++c)
+      if (data[r * nReservedColumns + c] != 0)
+        return false;
+#endif
+  return true;
+}
+
+namespace mlir
+{
+namespace presburger
+{
+template class Matrix<MPInt>;
+template class Matrix<Fraction>;
+}
+}
+
+IntMatrix IntMatrix::identity(unsigned dimension) {
+  IntMatrix matrix(dimension, dimension);
+  for (unsigned i = 0; i < dimension; ++i)
+    matrix(i, i) = 1;
+  return matrix;
+}
+
+
+std::pair<IntMatrix, IntMatrix> IntMatrix::computeHermiteNormalForm() const {
   // We start with u as an identity matrix and perform operations on h until h
   // is in hermite normal form. We apply the same sequence of operations on u to
   // obtain a transform that takes h to hermite normal form.
-  Matrix<MPInt> h = *this;
-  Matrix<MPInt> u = Matrix<MPInt>::identity(h.getNumColumns());
+  IntMatrix h = *this;
+  IntMatrix u = IntMatrix::identity(h.getNumColumns());
 
   unsigned echelonCol = 0;
   // Invariant: in all rows above row, all columns from echelonCol onwards
@@ -353,35 +386,10 @@ template <> std::pair<Matrix<MPInt>, Matrix<MPInt>> Matrix<MPInt>::computeHermit
   return {h, u};
 }
 
-template <typename T> void Matrix<T>::print(raw_ostream &os) const {
-  for (unsigned row = 0; row < nRows; ++row) {
-    for (unsigned column = 0; column < nColumns; ++column)
-      os << at(row, column) << ' ';
-    os << '\n';
-  }
-}
-
-template <typename T> void Matrix<T>::dump() const { print(llvm::errs()); }
-
-template <typename T> bool Matrix<T>::hasConsistentState() const {
-  if (data.size() != nRows * nReservedColumns)
-    return false;
-  if (nColumns > nReservedColumns)
-    return false;
-#ifdef EXPENSIVE_CHECKS
-  for (unsigned r = 0; r < nRows; ++r)
-    for (unsigned c = nColumns; c < nReservedColumns; ++c)
-      if (data[r * nReservedColumns + c] != 0)
-        return false;
-#endif
-  return true;
+MPInt IntMatrix::normalizeRow(unsigned row, unsigned cols) {
+  return normalizeRange(getRow(row).slice(0, cols));
 }
 
-namespace mlir
-{
-namespace presburger
-{
-template class Matrix<MPInt>;
-template class Matrix<Fraction>;
-}
+MPInt IntMatrix::normalizeRow(unsigned row) {
+  return normalizeRow(row, getNumColumns());
 }

mlir/lib/Analysis/Presburger/Simplex.cpp

@@ -436,7 +436,7 @@ LogicalResult SymbolicLexSimplex::addSymbolicCut(unsigned row) {
 }
 
 void SymbolicLexSimplex::recordOutput(SymbolicLexOpt &result) const {
-  Matrix<MPInt> output(0, domainPoly.getNumVars() + 1);
+  IntMatrix output(0, domainPoly.getNumVars() + 1);
   output.reserveRows(result.lexopt.getNumOutputs());
   for (const Unknown &u : var) {
     if (u.isSymbol)
@@ -1801,7 +1801,7 @@ class presburger::GBRSimplex {
 ///
 /// When incrementing i, no cached f values get invalidated. However, the cached
 /// duals do get invalidated as the duals for the higher levels are different.
-void Simplex::reduceBasis(Matrix<MPInt> &basis, unsigned level) {
+void Simplex::reduceBasis(IntMatrix &basis, unsigned level) {
   const Fraction epsilon(3, 4);
 
   if (level == basis.getNumRows() - 1)
@@ -1975,7 +1975,7 @@ std::optional<SmallVector<MPInt, 8>> Simplex::findIntegerSample() {
     return {};
 
   unsigned nDims = var.size();
-  Matrix<MPInt> basis = Matrix<MPInt>::identity(nDims);
+  IntMatrix basis = IntMatrix::identity(nDims);
 
   unsigned level = 0;
   // The snapshot just before constraining a direction to a value at each level.