[Mlir-commits] [mlir] [MLIR][Presburger] (PR #70843)

[llvmlistbot at llvm.org]

https://github.com/Abhinav271828 created https://github.com/llvm/llvm-project/pull/70843

Implement Gram-Schmidt orthogonalisation for the FracMatrix class.
This requires dotProduct, which has been added as a util.

>From 060e3c8f7f76e083723e4c303ec396447e6a1df3 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 31 Oct 2023 17:55:43 +0000
Subject: [PATCH] Add gram-schmidt and tests

---
 .../include/mlir/Analysis/Presburger/Matrix.h |  5 ++++
 mlir/include/mlir/Analysis/Presburger/Utils.h |  5 ++++
 mlir/lib/Analysis/Presburger/Matrix.cpp       | 23 +++++++++++++++++++
 mlir/lib/Analysis/Presburger/Utils.cpp        |  8 +++++++
 .../Analysis/Presburger/MatrixTest.cpp        | 15 ++++++++++++
 5 files changed, 56 insertions(+)

diff --git a/mlir/include/mlir/Analysis/Presburger/Matrix.h b/mlir/include/mlir/Analysis/Presburger/Matrix.h
index 4d9f13832e0692a..b591b0b4fdad167 100644
--- a/mlir/include/mlir/Analysis/Presburger/Matrix.h
+++ b/mlir/include/mlir/Analysis/Presburger/Matrix.h
@@ -265,6 +265,11 @@ class FracMatrix : public Matrix<Fraction> {
   // does not exist, which happens iff det = 0.
   // Assert-fails if the matrix is not square.
   Fraction determinant(FracMatrix *inverse = nullptr) const;
+
+  // Computes the Gram-Schmidt orthogonalisation
+  // of the matrix (cubic time).
+  FracMatrix gramSchmidt() const;
+
 };
 
 } // namespace presburger
diff --git a/mlir/include/mlir/Analysis/Presburger/Utils.h b/mlir/include/mlir/Analysis/Presburger/Utils.h
index a451ae8bf55723e..639683cea0d9f2e 100644
--- a/mlir/include/mlir/Analysis/Presburger/Utils.h
+++ b/mlir/include/mlir/Analysis/Presburger/Utils.h
@@ -276,6 +276,11 @@ SmallVector<MPInt, 8> getNegatedCoeffs(ArrayRef<MPInt> coeffs);
 /// a_1 x_1 + ... + a_n x_ + c < 0, i.e., -a_1 x_1 - ... - a_n x_ - c - 1 >= 0,
 /// since all the variables are constrained to be integers.
 SmallVector<MPInt, 8> getComplementIneq(ArrayRef<MPInt> ineq);
+
+/// Compute the dot product of two vectors.
+/// Assumes that the vectors have the same sizes.
+Fraction dotProduct(MutableArrayRef<Fraction> a, MutableArrayRef<Fraction> b);
+
 } // namespace presburger
 } // namespace mlir
 
diff --git a/mlir/lib/Analysis/Presburger/Matrix.cpp b/mlir/lib/Analysis/Presburger/Matrix.cpp
index ae97e456d9820cf..7669e461f079472 100644
--- a/mlir/lib/Analysis/Presburger/Matrix.cpp
+++ b/mlir/lib/Analysis/Presburger/Matrix.cpp
@@ -548,4 +548,27 @@ Fraction FracMatrix::determinant(FracMatrix *inverse) const {
     determinant *= m.at(i, i);
 
   return determinant;
+}
+
+FracMatrix FracMatrix::gramSchmidt() const {
+
+    // Create a copy of the argument to store
+    // the orthogonalised version.
+    FracMatrix orth(*this);
+    Fraction projectionScale;
+
+    // For each vector (row) in the matrix,
+    // subtract its unit projection along
+    // each of the previous vectors.
+    // This ensures that it has no component
+    // in the direction of any of the
+    // previous vectors.
+    for (unsigned i = 1; i < getNumRows(); i++) {
+        for (unsigned j = 0; j < i; j++) {
+            projectionScale = dotProduct(orth.getRow(i), orth.getRow(j)) /
+                              dotProduct(orth.getRow(j), orth.getRow(j));
+            orth.addToRow(j, i, -projectionScale);
+        }
+    }
+    return orth;
 }
\ No newline at end of file
diff --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp
index 9aef2f5de109357..30f01f545819101 100644
--- a/mlir/lib/Analysis/Presburger/Utils.cpp
+++ b/mlir/lib/Analysis/Presburger/Utils.cpp
@@ -520,3 +520,11 @@ SmallVector<int64_t, 8> presburger::getInt64Vec(ArrayRef<MPInt> range) {
   std::transform(range.begin(), range.end(), result.begin(), int64FromMPInt);
   return result;
 }
+
+Fraction presburger::dotProduct(MutableArrayRef<Fraction> a, MutableArrayRef<Fraction> b) {
+  assert(a.size() == b.size() && "Dot product of two unequal vectors!");
+  Fraction sum = 0;
+  for (unsigned i = 0; i < a.size(); i++)
+    sum += a[i] * b[i];
+  return sum;
+}
\ No newline at end of file
diff --git a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
index d05b05e004c5c5f..b7cc8e07a54e1ea 100644
--- a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
@@ -310,3 +310,18 @@ TEST(MatrixTest, intInverse) {
 
   EXPECT_EQ(det, 0);
 }
+
+TEST(MatrixTest, gramSchmidt) {
+  FracMatrix mat = makeFracMatrix(3, 5, {{Fraction(3, 1), Fraction(4, 1), Fraction(5, 1), Fraction(12, 1), Fraction(19, 1)},
+                                         {Fraction(4, 1), Fraction(5, 1), Fraction(6, 1), Fraction(13, 1), Fraction(20, 1)},
+                                         {Fraction(7, 1), Fraction(8, 1), Fraction(9, 1), Fraction(16, 1), Fraction(24, 1)}});
+
+  FracMatrix gramSchmidt = makeFracMatrix(3, 5,
+         {{Fraction(3, 1),     Fraction(4, 1),     Fraction(5, 1),    Fraction(12, 1),     Fraction(19, 1)},
+          {Fraction(142, 185), Fraction(383, 555), Fraction(68, 111), Fraction(13, 185),   Fraction(-262, 555)},
+          {Fraction(53, 463),  Fraction(27, 463),  Fraction(1, 463),  Fraction(-181, 463), Fraction(100, 463)}});
+
+  FracMatrix gs = mat.gramSchmidt();
+
+  EXPECT_EQ_FRAC_MATRIX(gs, gramSchmidt);
+}
\ No newline at end of file