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

[Arjun Pitchanathan]

================
@@ -535,6 +553,129 @@ 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 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.
+      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);
+            changed = true;
----------------
superty wrote:

it's also unfortuntae that this is almost entirely duplicated with the below loop nest except for the row/column swap. feel free to brainstorm ways to deduplicate this, but it's fine if there's no easy way.

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