[MLIR][Presburger] Add hermite normal form computation to Matrix

[Groverkss]

Summary

This patch adds hermite normal form computation to Matrix. Part of this algorithm
lived in LinearTransform, being used for compuing column echelon form. This
patch moves the implementation to Matrix::hermiteNormalForm and generalises it
to compute the hermite normal form.

[kaitingwang]

Thanks @Groverkss ! Any performance numbers to share?

[Groverkss]

In D133510#3779041, @kaitingwang wrote:

Thanks @Groverkss ! Any performance numbers to share?

I would assume that the computation works fast and should not bottleneck anything assuming that the number of columns is not too huge (which is generally true for polyhedra).

I still tried to run some experiments:

For a 10x10 matrix with almost 4 non-zero elements per row (generally polyhedra 2-3 non-zero elements per row in inequality/equality matrix):

{
  Matrix mat = makeMatrix(10, 10,
                          {
                              {1, 0, 0, 0, 1, 0, 0, 0, 0, 0},
                              {0, 1, 0, 0, 0, 0, 0, 0, 0, 1},
                              {0, 0, 0, 3, 0, 0, 4, 0, 0, -1},
                              {0, 0, 1, 0, 0, -1, 0, 0, 0, 0},
                              {0, 0, 0, -1, 0, 0, 0, 0, 0, -1},
                              {0, 0, 0, 1, 0, 0, 0, 0, 2, 10},
                              {0, 0, 0, 0, 3, 0, 0, 2, 0, 0},
                              {0, 1, 0, 0, 0, 0, 5, 0, 0, 0},
                              {0, 0, 0, 1, 0, 1, 0, 0, 1, 1},
                              {0, 0, 1, 0, 0, 1, 0, 1, 0, 0},
                          });
  auto [h, u] = mat.hermiteNormalForm();
}

Running this computation 100 times takes 2ms in a debug build. I would expect it to be faster in a release build.

I also tried the computation for dense 10x10 matrices, but while it runs fast, it is incorrect due to overflows (which should be fixed by https://reviews.llvm.org/D129510)

Are there any specific numbers that you are looking for? Or are these fine?

[arjunp]

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

166

168

use std::pair.

168

mlir/lib/Analysis/Presburger/LinearTransform.cpp

20–21

24

mlir/lib/Analysis/Presburger/Matrix.cpp

348

350–364

It would be better to just support negative targetCol in modEntryColumnOperation and use that

mlir/unittests/Analysis/Presburger/MatrixTest.cpp

201

229–239

Please write a function to do this

[Groverkss]

Address Arjun's comments

[Groverkss]

mlir/lib/Analysis/Presburger/Matrix.cpp

348

This is not required because the scale is always assigned a value before use.

[Groverkss]

Address more comments.

[arjunp]

mlir/lib/Analysis/Presburger/Matrix.cpp

264–271

ratio should just be floorDiv(m(row, targetCol), m(row, sourceCol)) right? (with scale =-ratio below)

[Groverkss]

Address Arjun's comment

[arjunp]

mlir/lib/Analysis/Presburger/Matrix.cpp

263

Maybe add some documentation explaining this.

[Groverkss]

Rebase

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

public:
 
   // Returns a linear transform T such that MT is M in column echelon form.
   // Also returns the number of non-zero columns in MT.
   //
   // Specifically, T is such that in every column the first non-zero row is
   // strictly below that of the previous column, and all columns which have only
   // zeros are at the end.
   static std::pair<unsigned, LinearTransform>
-  makeTransformToColumnEchelon(Matrix m);
+  makeTransformToColumnEchelon(const Matrix &m);
 
   // Returns an IntegerRelation having a constraint vector vT for every
   // constraint vector v in rel, where T is this transform.
   IntegerRelation applyTo(const IntegerRelation &rel) const;
 
   // The given vector is interpreted as a row vector v. Post-multiply v with
   // this transform, say T, and return vT.
   SmallVector<int64_t, 8> preMultiplyWithRow(ArrayRef<int64_t> rowVec) const {

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

public:
   /// 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;
 
+  /// 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.
   ///
   /// Due to the representation of the matrix, resizing vertically (adding rows)
   /// is less expensive than increasing the number of columns beyond
   /// nReservedColumns.
   void resize(unsigned newNRows, unsigned newNColumns);

mlir/lib/Analysis/Presburger/LinearTransform.cpp

 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 
 using namespace mlir;
 using namespace presburger;
 
 LinearTransform::LinearTransform(Matrix &&oMatrix) : matrix(oMatrix) {}
 LinearTransform::LinearTransform(const Matrix &oMatrix) : matrix(oMatrix) {}
 
-// Set M(row, targetCol) to its remainder on division by M(row, sourceCol)
-// by subtracting from column targetCol an appropriate integer multiple of
-// sourceCol. This brings M(row, targetCol) to the range [0, M(row, sourceCol)).
-// Apply the same column operation to otherMatrix, with the same integer
-// multiple.
-static void modEntryColumnOperation(Matrix &m, unsigned row, unsigned sourceCol,
-                                    unsigned targetCol, Matrix &otherMatrix) {
-  assert(m(row, sourceCol) != 0 && "Cannot divide by zero!");
-  assert((m(row, sourceCol) > 0 && m(row, targetCol) > 0) &&
-         "Operands must be positive!");
-  int64_t ratio = m(row, targetCol) / m(row, sourceCol);
-  m.addToColumn(sourceCol, targetCol, -ratio);
-  otherMatrix.addToColumn(sourceCol, targetCol, -ratio);
-}
-
 std::pair<unsigned, LinearTransform>
-LinearTransform::makeTransformToColumnEchelon(Matrix m) {
-  // We start with an identity result matrix and perform operations on m
-  // until m is in column echelon form. We apply the same sequence of operations
-  // on resultMatrix to obtain a transform that takes m to column echelon
-  // form.
-  Matrix resultMatrix = Matrix::identity(m.getNumColumns());
-
-  unsigned echelonCol = 0;
-  // Invariant: in all rows above row, all columns from echelonCol onwards
-  // are all zero elements. In an iteration, if the curent row has any non-zero
-  // elements echelonCol onwards, we bring one to echelonCol and use it to
-  // make all elements echelonCol + 1 onwards zero.
-  for (unsigned row = 0; row < m.getNumRows(); ++row) {
-    // Search row for a non-empty entry, starting at echelonCol.
-    unsigned nonZeroCol = echelonCol;
-    for (unsigned e = m.getNumColumns(); nonZeroCol < e; ++nonZeroCol) {
-      if (m(row, nonZeroCol) == 0)
-        continue;
+LinearTransform::makeTransformToColumnEchelon(const Matrix &m) {
+  // Compute the hermite normal form of m. This, is by definition, is in column
+  // echelon form.
+  auto [h, u] = m.computeHermiteNormalForm();
+
+  // Since the matrix is in column ecehlon form, a zero column means the rest of
+  // the columns are zero. Thus, once we find a zero column, we can stop.
+  unsigned col, e;
+  for (col = 0, e = m.getNumColumns(); col < e; ++col) {
+    bool zeroCol = true;
+    for (unsigned row = 0, f = m.getNumRows(); row < f; ++row) {
+      if (h(row, col) != 0) {
+        zeroCol = false;
         break;
       }
-
-    // Continue to the next row with the same echelonCol if this row is all
-    // zeros from echelonCol onwards.
-    if (nonZeroCol == m.getNumColumns())
-      continue;
-
-    // Bring the non-zero column to echelonCol. This doesn't affect rows
-    // above since they are all zero at these columns.
-    if (nonZeroCol != echelonCol) {
-      m.swapColumns(nonZeroCol, echelonCol);
-      resultMatrix.swapColumns(nonZeroCol, echelonCol);
-    }
-
-    // Make m(row, echelonCol) non-negative.
-    if (m(row, echelonCol) < 0) {
-      m.negateColumn(echelonCol);
-      resultMatrix.negateColumn(echelonCol);
-    }
-
-    // Make all the entries in row after echelonCol zero.
-    for (unsigned i = echelonCol + 1, e = m.getNumColumns(); i < e; ++i) {
-      // We make m(row, i) non-negative, and then apply the Euclidean GCD
-      // algorithm to (row, i) and (row, echelonCol). At the end, one of them
-      // has value equal to the gcd of the two entries, and the other is zero.
-
-      if (m(row, i) < 0) {
-        m.negateColumn(i);
-        resultMatrix.negateColumn(i);
-      }
-
-      unsigned targetCol = i, sourceCol = echelonCol;
-      // At every step, we set m(row, targetCol) %= m(row, sourceCol), and
-      // swap the indices sourceCol and targetCol. (not the columns themselves)
-      // This modulo is implemented as a subtraction
-      // m(row, targetCol) -= quotient * m(row, sourceCol),
-      // where quotient = floor(m(row, targetCol) / m(row, sourceCol)),
-      // which brings m(row, targetCol) to the range [0, m(row, sourceCol)).
-      //
-      // We are only allowed column operations; we perform the above
-      // for every row, i.e., the above subtraction is done as a column
-      // operation. This does not affect any rows above us since they are
-      // guaranteed to be zero at these columns.
-      while (m(row, targetCol) != 0 && m(row, sourceCol) != 0) {
-        modEntryColumnOperation(m, row, sourceCol, targetCol, resultMatrix);
-        std::swap(targetCol, sourceCol);
-      }
-
-      // One of (row, echelonCol) and (row, i) is zero and the other is the gcd.
-      // Make it so that (row, echelonCol) holds the non-zero value.
-      if (m(row, echelonCol) == 0) {
-        m.swapColumns(i, echelonCol);
-        resultMatrix.swapColumns(i, echelonCol);
-      }
     }
 
-    ++echelonCol;
+    if (zeroCol)
+      break;
   }
 
-  return {echelonCol, LinearTransform(std::move(resultMatrix))};
+  return {col, LinearTransform(std::move(u))};
 }
 
 IntegerRelation LinearTransform::applyTo(const IntegerRelation &rel) const {
   IntegerRelation result(rel.getSpace());
 
   for (unsigned i = 0, e = rel.getNumEqualities(); i < e; ++i) {
     ArrayRef<int64_t> eq = rel.getEquality(i);
 

mlir/lib/Analysis/Presburger/Matrix.cpp

Matrix::postMultiplyWithColumn(ArrayRef<int64_t> colVec) const {
 
   SmallVector<int64_t, 8> result(getNumRows(), 0);
   for (unsigned row = 0, e = getNumRows(); row < e; row++)
     for (unsigned i = 0, e = getNumColumns(); i < e; i++)
       result[row] += at(row, i) * colVec[i];
   return result;
 }
 
+/// Set M(row, targetCol) to its remainder on division by M(row, sourceCol)
+/// by subtracting from column targetCol an appropriate integer multiple of
+/// sourceCol. This brings M(row, targetCol) to the range [0, M(row,
+/// sourceCol)). Apply the same column operation to otherMatrix, with the same
+/// integer multiple.
+static void modEntryColumnOperation(Matrix &m, unsigned row, unsigned sourceCol,
+                                    unsigned targetCol, Matrix &otherMatrix) {
+  assert(m(row, sourceCol) != 0 && "Cannot divide by zero!");
+  assert(m(row, sourceCol) > 0 && "Source must be positive!");
+  int64_t ratio = -floorDiv(m(row, targetCol), m(row, sourceCol));
+  m.addToColumn(sourceCol, targetCol, ratio);
+  otherMatrix.addToColumn(sourceCol, targetCol, ratio);
+}
+
+std::pair<Matrix, Matrix> Matrix::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 h = *this;
+  Matrix u = Matrix::identity(h.getNumColumns());
+
+  unsigned echelonCol = 0;
+  // Invariant: in all rows above row, all columns from echelonCol onwards
+  // are all zero elements. In an iteration, if the curent row has any non-zero
+  // elements echelonCol onwards, we bring one to echelonCol and use it to
+  // make all elements echelonCol + 1 onwards zero.
+  for (unsigned row = 0; row < h.getNumRows(); ++row) {
+    // Search row for a non-empty entry, starting at echelonCol.
+    unsigned nonZeroCol = echelonCol;
+    for (unsigned e = h.getNumColumns(); nonZeroCol < e; ++nonZeroCol) {
+      if (h(row, nonZeroCol) == 0)
+        continue;
+      break;
+    }
+
+    // Continue to the next row with the same echelonCol if this row is all
+    // zeros from echelonCol onwards.
+    if (nonZeroCol == h.getNumColumns())
+      continue;
+
+    // Bring the non-zero column to echelonCol. This doesn't affect rows
+    // above since they are all zero at these columns.
+    if (nonZeroCol != echelonCol) {
+      h.swapColumns(nonZeroCol, echelonCol);
+      u.swapColumns(nonZeroCol, echelonCol);
+    }
+
+    // Make h(row, echelonCol) non-negative.
+    if (h(row, echelonCol) < 0) {
+      h.negateColumn(echelonCol);
+      u.negateColumn(echelonCol);
+    }
+
+    // Make all the entries in row after echelonCol zero.
+    for (unsigned i = echelonCol + 1, e = h.getNumColumns(); i < e; ++i) {
+      // We make h(row, i) non-negative, and then apply the Euclidean GCD
+      // algorithm to (row, i) and (row, echelonCol). At the end, one of them
+      // has value equal to the gcd of the two entries, and the other is zero.
+
+      if (h(row, i) < 0) {
+        h.negateColumn(i);
+        u.negateColumn(i);
+      }
+
+      unsigned targetCol = i, sourceCol = echelonCol;
+      // At every step, we set h(row, targetCol) %= h(row, sourceCol), and
+      // swap the indices sourceCol and targetCol. (not the columns themselves)
+      // This modulo is implemented as a subtraction
+      // h(row, targetCol) -= quotient * h(row, sourceCol),
+      // where quotient = floor(h(row, targetCol) / h(row, sourceCol)),
+      // which brings h(row, targetCol) to the range [0, h(row, sourceCol)).
+      //
+      // We are only allowed column operations; we perform the above
+      // for every row, i.e., the above subtraction is done as a column
+      // operation. This does not affect any rows above us since they are
+      // guaranteed to be zero at these columns.
+      while (h(row, targetCol) != 0 && h(row, sourceCol) != 0) {
+        modEntryColumnOperation(h, row, sourceCol, targetCol, u);
+        std::swap(targetCol, sourceCol);
+      }
+
+      // One of (row, echelonCol) and (row, i) is zero and the other is the gcd.
+      // Make it so that (row, echelonCol) holds the non-zero value.
+      if (h(row, echelonCol) == 0) {
+        h.swapColumns(i, echelonCol);
+        u.swapColumns(i, echelonCol);
+      }
+    }
+
+    // Make all entries before echelonCol non-negative and strictly smaller
+    // than the pivot entry.
+    for (unsigned i = 0; i < echelonCol; ++i)
+      modEntryColumnOperation(h, row, echelonCol, i, u);
+
+    ++echelonCol;
+  }
+
+  return {h, u};
+}
+
 void Matrix::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';
   }
 }
 
 void Matrix::dump() const { print(llvm::errs()); }
 
 bool Matrix::hasConsistentState() const {
   if (data.size() != nRows * nReservedColumns)
     return false;
   if (nColumns > nReservedColumns)
     return false;
   for (unsigned r = 0; r < nRows; ++r)
     for (unsigned c = nColumns; c < nReservedColumns; ++c)
       if (data[r * nReservedColumns + c] != 0)
         return false;
   return true;
 }

mlir/unittests/Analysis/Presburger/MatrixTest.cpp

 //===- MatrixTest.cpp - Tests for Matrix ----------------------------------===//
 //
 // 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 "./Utils.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 using namespace mlir;
 using namespace presburger;
 
 TEST(MatrixTest, ReadWrite) {
   Matrix mat(5, 5);
TEST(MatrixTest, resize) {
   mat.resize(7, 7);
   ASSERT_TRUE(mat.hasConsistentState());
   EXPECT_EQ(mat.getNumRows(), 7u);
   EXPECT_EQ(mat.getNumColumns(), 7u);
   for (unsigned row = 0; row < 7; ++row)
     for (unsigned col = 0; col < 7; ++col)
       EXPECT_EQ(mat(row, col), row >= 3 || col >= 3 ? 0 : int(10 * row + col));
 }
+
+static void checkHermiteNormalForm(const Matrix &mat,
+                                   const Matrix &hermiteForm) {
+  auto [h, u] = mat.computeHermiteNormalForm();
+
+  for (unsigned row = 0; row < mat.getNumRows(); row++)
+    for (unsigned col = 0; col < mat.getNumColumns(); col++)
+      EXPECT_EQ(h(row, col), hermiteForm(row, col));
+}
+
+TEST(MatrixTest, computeHermiteNormalForm) {
+  // TODO: Add a check to test the original statement of hermite normal form
+  // instead of using a precomputed result.
+
+  {
+    // Hermite form of a unimodular matrix is the identity matrix.
+    Matrix mat = makeMatrix(3, 3, {{2, 3, 6}, {3, 2, 3}, {17, 11, 16}});
+    Matrix hermiteForm = makeMatrix(3, 3, {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}});
+    checkHermiteNormalForm(mat, hermiteForm);
+  }
+
+  {
+    // Hermite form of a unimodular is the identity matrix.
+    Matrix mat = makeMatrix(
+        4, 4,
+        {{-6, -1, -19, -20}, {0, 1, 0, 0}, {-5, 0, -15, -16}, {6, 0, 18, 19}});
+    Matrix hermiteForm = makeMatrix(
+        4, 4, {{1, 0, 0, 0}, {0, 1, 0, 0}, {0, 0, 1, 0}, {0, 0, 0, 1}});
+    checkHermiteNormalForm(mat, hermiteForm);
+  }
+
+  {
+    Matrix mat = makeMatrix(
+        4, 4, {{3, 3, 1, 4}, {0, 1, 0, 0}, {0, 0, 19, 16}, {0, 0, 0, 3}});
+    Matrix hermiteForm = makeMatrix(
+        4, 4, {{1, 0, 0, 0}, {0, 1, 0, 0}, {1, 0, 3, 0}, {18, 0, 54, 57}});
+    checkHermiteNormalForm(mat, hermiteForm);
+  }
+
+  {
+    Matrix mat = makeMatrix(
+        4, 4, {{3, 3, 1, 4}, {0, 1, 0, 0}, {0, 0, 19, 16}, {0, 0, 0, 3}});
+    Matrix hermiteForm = makeMatrix(
+        4, 4, {{1, 0, 0, 0}, {0, 1, 0, 0}, {1, 0, 3, 0}, {18, 0, 54, 57}});
+    checkHermiteNormalForm(mat, hermiteForm);
+  }
+
+  {
+    Matrix mat =
+        makeMatrix(3, 5, {{0, 2, 0, 7, 1}, {-1, 0, 0, -3, 0}, {0, 4, 1, 0, 8}});
+    Matrix hermiteForm =
+        makeMatrix(3, 5, {{1, 0, 0, 0, 0}, {0, 1, 0, 0, 0}, {0, 0, 1, 0, 0}});
+    checkHermiteNormalForm(mat, hermiteForm);
+  }
+}