From 61b6d275cd832436133d3ccd7f607fa734a2263f Mon Sep 17 00:00:00 2001 From: Yue Huang <30948580+AdUhTkJm@users.noreply.github.com> Date: Mon, 16 Mar 2026 15:52:50 +0000 Subject: [PATCH] [MLIR][Presburger] Add support for Smith normal form (#185328) FPL already has support for computing Hermite normal form for integer matrices. Here we add support to computing Smith normal form. This is a preparation for Barvinok's algorithm. Given a polyhedron $P = \{ x | Ax + b = 0, Cx + d \leq 0 \}$, we must find a particular solution $x_0$ of $Ax + b = 0$ in order to project lower-dimensional polyhedra into full-dimensional ones. This requires the Smith normal form of the integer matrix $A$. The implementation here follows the algorithm in [wikipedia](https://en.wikipedia.org/wiki/Smith_normal_form#Algorithm). --- .../include/mlir/Analysis/Presburger/Matrix.h | 12 ++ mlir/lib/Analysis/Presburger/Matrix.cpp | 163 ++++++++++++++++++ .../Analysis/Presburger/MatrixTest.cpp | 51 ++++++ 3 files changed, 226 insertions(+) diff --git a/mlir/include/mlir/Analysis/Presburger/Matrix.h b/mlir/include/mlir/Analysis/Presburger/Matrix.h index 15069fa2a55f..4592efda29f7 100644 --- a/mlir/include/mlir/Analysis/Presburger/Matrix.h +++ b/mlir/include/mlir/Analysis/Presburger/Matrix.h @@ -233,6 +233,9 @@ public: /// 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 postMultiply(const Matrix &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 @@ public: /// pivot. std::pair 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 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 83a2c280c3d4..6ea04543146b 100644 --- a/mlir/lib/Analysis/Presburger/Matrix.cpp +++ b/mlir/lib/Analysis/Presburger/Matrix.cpp @@ -284,6 +284,24 @@ void Matrix::moveColumns(unsigned srcPos, unsigned num, unsigned dstPos) { } } +template +Matrix Matrix::postMultiply(const Matrix &other) const { + assert(getNumColumns() == other.getNumRows()); + unsigned n = getNumRows(); + unsigned m = other.getNumRows(); + unsigned p = other.getNumColumns(); + Matrix 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 void Matrix::addToRow(unsigned sourceRow, unsigned targetRow, const T &scale) { @@ -535,6 +553,151 @@ std::pair IntMatrix::computeHermiteNormalForm() const { return {h, u}; } +// In the submatrix `mat(from:, from:)`, the function finds the position (row, +// col) of the element with smallest non-zero absolute value. When all elements +// in the submatrix are zero, returns std::nullopt. +static std::optional> +findNonZeroMinInSubmatrix(const IntMatrix &mat, unsigned from) { + unsigned numRows = mat.getNumRows(); + unsigned numCols = mat.getNumColumns(); + unsigned minRow = from, minCol = from; + + std::optional minVal; + for (unsigned r = from; r < numRows; r++) { + for (unsigned c = from; c < numCols; c++) { + DynamicAPInt val = llvm::abs(mat(r, c)); + if (val == 0 || (minVal && val >= *minVal)) + continue; + + minVal = val; + minRow = r; + minCol = c; + } + } + + if (!minVal) + return std::nullopt; + + return std::make_pair(minRow, minCol); +} + +// Finds the first row in submatrix `mat(from:, from:)` that contains an element +// `d` such that `d` is not a multiple of `divisor`. When there is no such row, +// returns std::nullopt. +static std::optional findNonMultipleRow(const IntMatrix &mat, + unsigned from, + const DynamicAPInt &divisor) { + unsigned numRows = mat.getNumRows(); + unsigned numCols = mat.getNumColumns(); + for (unsigned row = from; row < numRows; ++row) { + for (unsigned col = from; col < numCols; ++col) { + if (mat(row, col) % divisor != 0) + return row; + } + } + return std::nullopt; +} + +std::tuple +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 changed; + do { + changed = 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. + auto pivotPos = findNonZeroMinInSubmatrix(d, i); + if (!pivotPos) + break; + auto [pvtRow, pvtCol] = *pivotPos; + + // The remaining submatrix is zero. + if (d(pvtRow, pvtCol) == 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); + changed = true; + } + } + } + // 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); + changed = true; + } + } + } + + if (auto row = findNonMultipleRow(d, i + 1, pivot)) { + // 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); + changed = true; + } + } while (changed); + } + + 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 e2c2a9bcb7d2..d40760fa8d89 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);