[MLIR][Presburger] Implement matrix inverse

[Abhinav271828]

Shift the determinant() function from LinearTransform to Matrix.
Implement a FracMatrix class, inheriting from Matrix, for inverses.
Implement inverse for FracMatrix and intInverse for IntMatrix.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[llvmbot]

@llvm/pr-subscribers-mlir-presburger

Changes

A new class, FracMatrix, has been created for rational matrices.
It inherits from Matrix<Fraction> with no new data attributes and one new method, inverse().
Gaussian elimination is used in the inverse implementation.
A basic test for inverses has also been added.

---

Full diff: https://github.com/llvm/llvm-project/pull/67382.diff

5 Files Affected:

• (modified) mlir/include/mlir/Analysis/Presburger/Fraction.h (+1-1)
• (modified) mlir/include/mlir/Analysis/Presburger/Matrix.h (+26)
• (modified) mlir/lib/Analysis/Presburger/Matrix.cpp (+66)
• (modified) mlir/unittests/Analysis/Presburger/MatrixTest.cpp (+11)
• (modified) mlir/unittests/Analysis/Presburger/Utils.h (+2-2)

diff --git a/mlir/include/mlir/Analysis/Presburger/Fraction.h b/mlir/include/mlir/Analysis/Presburger/Fraction.h
index 74127a900d53ed2..f633192a870c8d4 100644
--- a/mlir/include/mlir/Analysis/Presburger/Fraction.h
+++ b/mlir/include/mlir/Analysis/Presburger/Fraction.h
@@ -102,7 +102,7 @@ inline bool operator>=(const Fraction &x, const Fraction &y) {
 inline Fraction reduce(const Fraction &f) {
   if (f == Fraction(0))
     return Fraction(0, 1);
-  MPInt g = gcd(f.num, f.den);
+  MPInt g = gcd(abs(f.num), abs(f.den));
   return Fraction(f.num / g, f.den / g);
 }
 
diff --git a/mlir/include/mlir/Analysis/Presburger/Matrix.h b/mlir/include/mlir/Analysis/Presburger/Matrix.h
index f05a4fb2ffbb70d..b9048334c39f1f3 100644
--- a/mlir/include/mlir/Analysis/Presburger/Matrix.h
+++ b/mlir/include/mlir/Analysis/Presburger/Matrix.h
@@ -244,6 +244,32 @@ class IntMatrix : public Matrix<MPInt>
 
 };
 
+// An inherited class for rational matrices, with no new data attributes.
+// This is only used for the matrix-related method which apply only
+// to fractions (inverse).
+class FracMatrix : public Matrix<Fraction>
+{
+public:
+  FracMatrix(unsigned rows, unsigned columns, unsigned reservedRows = 0,
+            unsigned reservedColumns = 0) :
+    Matrix<Fraction>(rows, columns, reservedRows, reservedColumns) {};
+
+  FracMatrix(Matrix<Fraction> m) :
+    Matrix<Fraction>(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 FracMatrix identity(unsigned dimension);
+
+  // Return the inverse of the matrix, leaving the calling object unmodified.
+  FracMatrix inverse();
+
+};
+
 } // namespace presburger
 } // namespace mlir
 
diff --git a/mlir/lib/Analysis/Presburger/Matrix.cpp b/mlir/lib/Analysis/Presburger/Matrix.cpp
index f0bcb09fb28f7b1..c0fbc688be2e1b2 100644
--- a/mlir/lib/Analysis/Presburger/Matrix.cpp
+++ b/mlir/lib/Analysis/Presburger/Matrix.cpp
@@ -390,4 +390,70 @@ MPInt IntMatrix::normalizeRow(unsigned row, unsigned cols) {
 
 MPInt IntMatrix::normalizeRow(unsigned row) {
   return normalizeRow(row, getNumColumns());
+}
+
+
+FracMatrix FracMatrix::identity(unsigned dimension) {
+  FracMatrix matrix(dimension, dimension);
+  for (unsigned i = 0; i < dimension; ++i)
+    matrix(i, i) = 1;
+  return matrix;
+}
+
+FracMatrix FracMatrix::inverse()
+{
+    // We use Gaussian elimination on the rows of [M | I]
+    // to find the integer inverse. We proceed left-to-right,
+    // top-to-bottom. M is assumed to be a dim x dim matrix.
+
+    unsigned dim = getNumRows();
+
+    // Construct the augmented matrix [M | I]
+    FracMatrix augmented(dim, dim + dim);
+    for (unsigned i = 0; i < dim; i++)
+    {
+        augmented.fillRow(i, 0);
+        for (unsigned j = 0; j < dim; j++)
+            augmented(i, j) = at(i, j);
+        augmented(i, dim+i).num = 1;
+        augmented(i, dim+i).den = 1;
+    }
+    Fraction a, b;
+    for (unsigned i = 0; i < dim; i++)
+    {
+        if (augmented(i, i) == Fraction(0, 1))
+            for (unsigned j = i+1; j < dim; j++)
+                if (augmented(j, i) != Fraction(0, 1))
+                {
+                    augmented.addToRow(i, augmented.getRow(j), Fraction(1, 1));
+                    break;
+                }
+
+        b = augmented(i, i);
+        for (unsigned j = 0; j < dim; j++)
+        {
+            if (i == j || augmented(j, i) == 0) continue;
+            a = augmented(j, i);
+            // Rj -> Rj - (b/a)Ri
+            augmented.addToRow(j, augmented.getRow(i), - a / b);
+            // Now (Rj)i = 0
+        }
+    }
+    
+    // Now only diagonal elements are nonzero, but they are
+    // not necessarily 1.
+    for (unsigned i = 0; i < dim; i++)
+    {
+        a = augmented(i, i);
+        for (unsigned j = dim; j < dim + dim; j++)
+            augmented(i, j) = augmented(i, j) / a;
+    }
+
+    // Copy the right half of the augmented matrix.
+    FracMatrix inverse(dim, dim);
+    for (unsigned i = 0; i < dim; i++)
+        for (unsigned j = 0; j < dim; j++)
+            inverse(i, j) = augmented(i, j+dim);
+
+    return inverse;
 }
\ No newline at end of file
diff --git a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
index 6b23cedabf624ec..d07318617e1a249 100644
--- a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
@@ -248,3 +248,14 @@ TEST(MatrixTest, computeHermiteNormalForm) {
     checkHermiteNormalForm(mat, hermiteForm);
   }
 }
+
+TEST(MatrixTest, inverse) {
+    FracMatrix mat = makeFracMatrix(2, 2, {{Fraction(2, 1), Fraction(1, 1)}, {Fraction(7, 1), Fraction(0, 1)}});
+    FracMatrix inverse = makeFracMatrix(2, 2, {{Fraction(0, 1), Fraction(1, 7)}, {Fraction(1, 1), Fraction(-2, 7)}});
+
+    FracMatrix inv = mat.inverse();
+
+    for (unsigned row = 0; row < 2; row++)
+      for (unsigned col = 0; col < 2; col++)
+        EXPECT_EQ(inv(row, col), inverse(row, col));
+}
diff --git a/mlir/unittests/Analysis/Presburger/Utils.h b/mlir/unittests/Analysis/Presburger/Utils.h
index ef4a67d0b8c004f..c652daa583a882c 100644
--- a/mlir/unittests/Analysis/Presburger/Utils.h
+++ b/mlir/unittests/Analysis/Presburger/Utils.h
@@ -40,9 +40,9 @@ inline IntMatrix makeIntMatrix(unsigned numRow, unsigned numColumns,
   return results;
 }
 
-inline Matrix<Fraction> makeFracMatrix(unsigned numRow, unsigned numColumns,
+inline FracMatrix makeFracMatrix(unsigned numRow, unsigned numColumns,
                          ArrayRef<SmallVector<Fraction, 8>> matrix) {
-  Matrix<Fraction> results(numRow, numColumns);
+  FracMatrix results(numRow, numColumns);
   assert(matrix.size() == numRow);
   for (unsigned i = 0; i < numRow; ++i) {
     assert(matrix[i].size() == numColumns &&