[Mlir-commits] [mlir] [MLIR][Presburger] Add LLL basis reduction (PR #75565)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-presburger

@llvm/pr-subscribers-mlir

Author: None (Abhinav271828)

<details>
<summary>Changes</summary>

Add a method for LLL basis reduction to the FracMatrix class.
This needs an abs() method for Fractions, which is added to Fraction.h.

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


4 Files Affected:

- (modified) mlir/include/mlir/Analysis/Presburger/Fraction.h (+2) 
- (modified) mlir/include/mlir/Analysis/Presburger/Matrix.h (+4) 
- (modified) mlir/lib/Analysis/Presburger/Matrix.cpp (+36-1) 
- (modified) mlir/unittests/Analysis/Presburger/MatrixTest.cpp (+48-1) 


``````````diff
diff --git a/mlir/include/mlir/Analysis/Presburger/Fraction.h b/mlir/include/mlir/Analysis/Presburger/Fraction.h
index afcbed84c66bc3..e95056ae5fc961 100644
--- a/mlir/include/mlir/Analysis/Presburger/Fraction.h
+++ b/mlir/include/mlir/Analysis/Presburger/Fraction.h
@@ -101,6 +101,8 @@ inline bool operator>=(const Fraction &x, const Fraction &y) {
   return compare(x, y) >= 0;
 }
 
+inline Fraction abs(const Fraction &f) { return Fraction(abs(f.num), f.den); }
+
 inline Fraction reduce(const Fraction &f) {
   if (f == Fraction(0))
     return Fraction(0, 1);
diff --git a/mlir/include/mlir/Analysis/Presburger/Matrix.h b/mlir/include/mlir/Analysis/Presburger/Matrix.h
index 89fad85c0c3374..fca3164bda6278 100644
--- a/mlir/include/mlir/Analysis/Presburger/Matrix.h
+++ b/mlir/include/mlir/Analysis/Presburger/Matrix.h
@@ -270,6 +270,10 @@ class FracMatrix : public Matrix<Fraction> {
   // of the rows of matrix (cubic time).
   // The rows of the matrix must be linearly independent.
   FracMatrix gramSchmidt() const;
+
+  // Run LLL basis reduction on the matrix, modifying it in-place.
+  // The parameter is delta.
+  void LLL(Fraction delta);
 };
 
 } // namespace presburger
diff --git a/mlir/lib/Analysis/Presburger/Matrix.cpp b/mlir/lib/Analysis/Presburger/Matrix.cpp
index 1fcc6d072b44b7..ae04c9f9149a5a 100644
--- a/mlir/lib/Analysis/Presburger/Matrix.cpp
+++ b/mlir/lib/Analysis/Presburger/Matrix.cpp
@@ -576,4 +576,39 @@ FracMatrix FracMatrix::gramSchmidt() const {
     }
   }
   return orth;
-}
\ No newline at end of file
+}
+
+void FracMatrix::LLL(Fraction delta) {
+  MPInt nearest;
+  Fraction mu;
+
+  // `bStar` holds the Gram-Schmidt orthogonalisation
+  // of the matrix at all times. It is recomputed every
+  // time the matrix is modified during the algorithm.
+  FracMatrix bStar = gramSchmidt();
+
+  unsigned k = 1;
+  while (k < getNumRows()) {
+    for (unsigned j = k - 1; j < k; j--) {
+      mu = dotProduct(getRow(k), bStar.getRow(j)) /
+           dotProduct(bStar.getRow(j), bStar.getRow(j));
+      if (abs(mu) > Fraction(1, 2)) {
+        nearest = floor(mu + Fraction(1, 2));
+        addToRow(k, getRow(j), -Fraction(nearest, 1));
+        bStar = gramSchmidt();
+      }
+    }
+    mu = dotProduct(getRow(k), bStar.getRow(k - 1)) /
+         dotProduct(bStar.getRow(k - 1), bStar.getRow(k - 1));
+    if (dotProduct(bStar.getRow(k), bStar.getRow(k)) >
+        (delta - mu * mu) *
+            dotProduct(bStar.getRow(k - 1), bStar.getRow(k - 1)))
+      k += 1;
+    else {
+      swapRows(k, k - 1);
+      bStar = gramSchmidt();
+      k = k > 1 ? k - 1 : 1;
+    }
+  }
+  return;
+}
diff --git a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
index 508d4fa369c14c..4d7d0531e5ee84 100644
--- a/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/MatrixTest.cpp
@@ -377,4 +377,51 @@ TEST(MatrixTest, gramSchmidt) {
   gs = mat.gramSchmidt();
 
   EXPECT_EQ_FRAC_MATRIX(gs, FracMatrix::identity(10));
-}
\ No newline at end of file
+}
+
+TEST(MatrixTest, LLL) {
+  FracMatrix mat =
+      makeFracMatrix(3, 3,
+                     {{Fraction(1, 1), Fraction(1, 1), Fraction(1, 1)},
+                      {Fraction(-1, 1), Fraction(0, 1), Fraction(2, 1)},
+                      {Fraction(3, 1), Fraction(5, 1), Fraction(6, 1)}});
+  mat.LLL(Fraction(3, 4));
+
+  FracMatrix LLL =
+      makeFracMatrix(3, 3,
+                     {{Fraction(0, 1), Fraction(1, 1), Fraction(0, 1)},
+                      {Fraction(1, 1), Fraction(0, 1), Fraction(1, 1)},
+                      {Fraction(-1, 1), Fraction(0, 1), Fraction(2, 1)}});
+
+  for (unsigned row = 0; row < 3; row++)
+    for (unsigned col = 0; col < 3; col++)
+      EXPECT_EQ(mat(row, col), LLL(row, col));
+
+  mat = makeFracMatrix(
+      2, 2,
+      {{Fraction(12, 1), Fraction(2, 1)}, {Fraction(13, 1), Fraction(4, 1)}});
+  LLL = makeFracMatrix(
+      2, 2,
+      {{Fraction(1, 1), Fraction(2, 1)}, {Fraction(9, 1), Fraction(-4, 1)}});
+
+  mat.LLL(Fraction(3, 4));
+
+  for (unsigned row = 0; row < 2; row++)
+    for (unsigned col = 0; col < 2; col++)
+      EXPECT_EQ(mat(row, col), LLL(row, col));
+
+  mat = makeFracMatrix(3, 3,
+                       {{Fraction(1, 1), Fraction(0, 1), Fraction(2, 1)},
+                        {Fraction(0, 1), Fraction(1, 3), -Fraction(5, 3)},
+                        {Fraction(0, 1), Fraction(0, 1), Fraction(1, 1)}});
+  LLL = makeFracMatrix(3, 3,
+                       {{Fraction(0, 1), Fraction(1, 3), Fraction(1, 3)},
+                        {Fraction(0, 1), Fraction(1, 3), -Fraction(2, 3)},
+                        {Fraction(1, 1), Fraction(0, 1), Fraction(0, 1)}});
+
+  mat.LLL(Fraction(3, 4));
+
+  for (unsigned row = 0; row < 3; row++)
+    for (unsigned col = 0; col < 3; col++)
+      EXPECT_EQ(mat(row, col), LLL(row, col));
+}

``````````

</details>


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