[Mlir-commits] [mlir] [MLIR][Presburger] Add support for Smith normal form (PR #185328)

Yue Huang llvmlistbot at llvm.org
Sun Mar 15 04:43:47 PDT 2026


https://github.com/AdUhTkJm updated https://github.com/llvm/llvm-project/pull/185328

>From 279bce8412b13005ca70c6e01e720b90de4197c5 Mon Sep 17 00:00:00 2001
From: Yue Huang <yh548 at cam.ac.uk>
Date: Sun, 8 Mar 2026 22:27:55 +0000
Subject: [PATCH] [MLIR][Presburger] Add support for Smith normal form

---
 .../include/mlir/Analysis/Presburger/Matrix.h |  12 ++
 mlir/lib/Analysis/Presburger/Matrix.cpp       | 139 ++++++++++++++++++
 .../Analysis/Presburger/MatrixTest.cpp        |  51 +++++++
 3 files changed, 202 insertions(+)

diff --git a/mlir/include/mlir/Analysis/Presburger/Matrix.h b/mlir/include/mlir/Analysis/Presburger/Matrix.h
index 15069fa2a55f1..4592efda29f70 100644
--- a/mlir/include/mlir/Analysis/Presburger/Matrix.h
+++ b/mlir/include/mlir/Analysis/Presburger/Matrix.h
@@ -233,6 +233,9 @@ class Matrix {
   /// The left shift operation (i.e. dstPos < srcPos) works in a similar way.
   void moveColumns(unsigned srcPos, unsigned num, unsigned dstPos);
 
+  /// Returns the matrix right-multiplied with `other`.
+  Matrix<T> postMultiply(const Matrix<T> &other) const;
+
 protected:
   /// The current number of rows, columns, and reserved columns. The underlying
   /// data vector is viewed as an nRows x nReservedColumns matrix, of which the
@@ -274,6 +277,15 @@ class IntMatrix : public Matrix<DynamicAPInt> {
   ///    pivot.
   std::pair<IntMatrix, IntMatrix> computeHermiteNormalForm() const;
 
+  /// Given the current matrix M, returns the matrices U, D, V such that
+  /// UMV = D, where D is called the Smith Normal Form (SNF).
+  /// The matrices have the following properties:
+  ///   - U, V are unimodular. In other words, det(U), det(V) are 1 or -1;
+  ///     their inverses also contain integer entries.
+  ///   - D is diagonal.
+  ///   - For all i, the diagonal element D_{i, i} divides D_{i + 1, i + 1}.
+  std::tuple<IntMatrix, IntMatrix, IntMatrix> computeSmithNormalForm() const;
+
   /// Divide the first `nCols` of the specified row by their GCD.
   /// Returns the GCD of the first `nCols` of the specified row.
   DynamicAPInt normalizeRow(unsigned row, unsigned nCols);
diff --git a/mlir/lib/Analysis/Presburger/Matrix.cpp b/mlir/lib/Analysis/Presburger/Matrix.cpp
index 83a2c280c3d4e..58cde7c167aba 100644
--- a/mlir/lib/Analysis/Presburger/Matrix.cpp
+++ b/mlir/lib/Analysis/Presburger/Matrix.cpp
@@ -284,6 +284,24 @@ void Matrix<T>::moveColumns(unsigned srcPos, unsigned num, unsigned dstPos) {
   }
 }
 
+template <typename T>
+Matrix<T> Matrix<T>::postMultiply(const Matrix<T> &other) const {
+  assert(getNumColumns() == other.getNumRows());
+  unsigned n = getNumRows();
+  unsigned m = other.getNumRows();
+  unsigned p = other.getNumColumns();
+  Matrix<T> result(n, p);
+
+  for (unsigned i = 0; i < n; i++) {
+    for (unsigned j = 0; j < m; j++) {
+      for (unsigned k = 0; k < p; k++) {
+        result.at(i, k) += at(i, j) * other.at(j, k);
+      }
+    }
+  }
+  return result;
+}
+
 template <typename T>
 void Matrix<T>::addToRow(unsigned sourceRow, unsigned targetRow,
                          const T &scale) {
@@ -535,6 +553,127 @@ std::pair<IntMatrix, IntMatrix> IntMatrix::computeHermiteNormalForm() const {
   return {h, u};
 }
 
+std::tuple<IntMatrix, IntMatrix, IntMatrix>
+IntMatrix::computeSmithNormalForm() const {
+  IntMatrix d = *this;
+  // We put D into diagonal form by applying row and columns operations to it.
+  // The matrix U records row operations applied in the process, and V records
+  // column operations.
+  IntMatrix u = IntMatrix::identity(d.getNumRows());
+  IntMatrix v = IntMatrix::identity(d.getNumColumns());
+
+  unsigned numRows = d.getNumRows();
+  unsigned numCols = d.getNumColumns();
+  for (unsigned i = 0, e = std::min(numRows, numCols); i < e; i++) {
+    // We first put D into diagonal form, and then ensure the divisibility
+    // condition. The latter step is better illustrated with an example:
+    //
+    // [6 0 ] ---(1)--> [6 10] ---(2)--> [2 0 ]
+    // [0 10]           [0 10]           [0 10]
+    //
+    // (1) adds the element violating the divisibility constraint to the same
+    // column in row i;
+    // (2) does an elimination of the column.
+    //
+    // There can be many elements that violate the constraint, hence the loop.
+    bool fixDivisibility;
+    do {
+      fixDivisibility = false;
+
+      // Find the entry in the submatrix d(i:, i:) with the smallest non-zero
+      // absolute value.
+      // The element is the pivot, and we record its current row and column.
+      unsigned pvtRow = i, pvtCol = i;
+      DynamicAPInt minVal(0);
+      for (unsigned r = i; r < numRows; r++) {
+        for (unsigned c = i; c < numCols; c++) {
+          DynamicAPInt val = llvm::abs(d(r, c));
+          if (val == 0 || (minVal != 0 && val >= minVal))
+            continue;
+
+          minVal = val;
+          pvtRow = r;
+          pvtCol = c;
+        }
+      }
+
+      // The remaining submatrix is zero.
+      if (minVal == 0)
+        break;
+
+      // Bring pivot to d(i, i). Record the operation in u, v respectively.
+      if (pvtRow != i) {
+        d.swapRows(pvtRow, i);
+        u.swapRows(pvtRow, i);
+      }
+      if (pvtCol != i) {
+        d.swapColumns(pvtCol, i);
+        v.swapColumns(pvtCol, i);
+      }
+
+      // Ensure the pivot is positive.
+      if (d(i, i) < 0) {
+        d.negateRow(i);
+        u.negateRow(i);
+      }
+
+      DynamicAPInt pivot = d(i, i);
+
+      // Clear other entries in row i and column i with Euclid's algorithm.
+      for (unsigned r = i + 1; r < numRows; ++r) {
+        while (d(r, i) != 0) {
+          DynamicAPInt quotient = d(r, i) / d(i, i);
+          d.addToRow(i, r, -quotient);
+          u.addToRow(i, r, -quotient);
+
+          if (d(r, i) != 0) {
+            d.swapRows(r, i);
+            u.swapRows(r, i);
+          }
+        }
+      }
+      // Similar to the rows operations, this time it works on columns.
+      for (unsigned c = i + 1; c < numCols; ++c) {
+        while (d(i, c) != 0) {
+          DynamicAPInt quotient = d(i, c) / d(i, i);
+          d.addToColumn(i, c, -quotient);
+          v.addToColumn(i, c, -quotient);
+
+          if (d(i, c) != 0) {
+            d.swapColumns(c, i);
+            v.swapColumns(c, i);
+          }
+        }
+      }
+
+      const auto &fixRowDivisibility = [&](int row) {
+        for (unsigned col = i + 1; col < numCols; ++col) {
+          if (d(row, col) % pivot == 0)
+            continue;
+
+          // Add the row (r) to row i. This brings d(r, c) into the i-th row,
+          // creating a new value at d(i, c) that will be used to reduce the
+          // pivot size.
+          d.addToRow(row, i, 1);
+          u.addToRow(row, i, 1);
+          fixDivisibility = true;
+
+          // We must break and restart the inner process to find the new,
+          // smaller pivot in the submatrix d(i:, i:).
+          return true;
+        }
+        return false;
+      };
+      for (unsigned r = i + 1; r < numRows; ++r) {
+        if (fixRowDivisibility(r))
+          break;
+      }
+    } while (fixDivisibility);
+  }
+
+  return {u, d, v};
+}
+
 DynamicAPInt IntMatrix::normalizeRow(unsigned row, unsigned cols) {
   return normalizeRange(getRow(row).slice(0, cols));
 }
diff --git a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
index e2c2a9bcb7d26..d40760fa8d894 100644
--- a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
@@ -258,6 +258,57 @@ TEST(MatrixTest, computeHermiteNormalForm) {
   }
 }
 
+static void checkSmithNormalForm(const IntMatrix &mat) {
+  auto [u, d, v] = mat.computeSmithNormalForm();
+
+  // Check u and v are unimodular.
+  EXPECT_EQ(llvm::abs(u.determinant()), 1);
+  EXPECT_EQ(llvm::abs(v.determinant()), 1);
+
+  // Check u @ mat @ v = d (@ is matrix multiplication).
+  EXPECT_EQ(u.postMultiply(mat).postMultiply(v), d);
+
+  // Check d is diagonal, i.e. non-diagonal elements are zero.
+  for (unsigned i = 0, e = d.getNumRows(); i < e; i++) {
+    for (unsigned j = 0, f = d.getNumColumns(); j < f; j++) {
+      if (i != j)
+        EXPECT_EQ(d(i, j), 0);
+    }
+  }
+
+  // Check d(i, i) divides d(i + 1, i + 1).
+  unsigned end = std::min(d.getNumRows(), d.getNumColumns()) - 1;
+  unsigned i = 0;
+  for (; i < end; i++) {
+    if (d(i, i) == 0)
+      break;
+
+    EXPECT_EQ(d(i + 1, i + 1) % d(i, i), 0);
+  }
+  for (; i < end; i++)
+    EXPECT_EQ(d(i, i), 0);
+}
+
+TEST(MatrixTest, computeSmithNormalForm) {
+  {
+    IntMatrix mat =
+        makeIntMatrix(4, 3, {{2, 5, 8}, {3, 6, 9}, {4, 7, 1}, {5, 8, 2}});
+    checkSmithNormalForm(mat);
+  }
+
+  {
+    // Smith normal form of this matrix has trailing zeroes on the diagonal.
+    IntMatrix mat = makeIntMatrix(2, 3, {{6, 4, 2}, {3, 2, 0}});
+    checkSmithNormalForm(mat);
+  }
+
+  {
+    // 1x1 edge case.
+    IntMatrix mat = makeIntMatrix(1, 1, {{9}});
+    checkSmithNormalForm(mat);
+  }
+}
+
 TEST(MatrixTest, inverse) {
   IntMatrix mat1 = makeIntMatrix(2, 2, {{2, 1}, {7, 0}});
   EXPECT_EQ(mat1.determinant(), -7);



More information about the Mlir-commits mailing list