[Mlir-commits] [mlir] [MLIR][Presburger] Implement computation of generating function for unimodular cones (PR #77235)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Mon Jan 8 11:29:24 PST 2024


https://github.com/Abhinav271828 updated https://github.com/llvm/llvm-project/pull/77235

>From 26fb9722414ea7e43cc554b5d2c0d96699b0071e Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 7 Jan 2024 16:39:37 +0530
Subject: [PATCH 1/5] Initial commit

---
 .../mlir/Analysis/Presburger/Barvinok.h       |  6 ++
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 68 +++++++++++++++++++
 .../Analysis/Presburger/BarvinokTest.cpp      | 36 ++++++++++
 3 files changed, 110 insertions(+)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 15e805860db237..93b29e2d718e59 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -24,6 +24,7 @@
 #ifndef MLIR_ANALYSIS_PRESBURGER_BARVINOK_H
 #define MLIR_ANALYSIS_PRESBURGER_BARVINOK_H
 
+#include "mlir/Analysis/Presburger/GeneratingFunction.h"
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
 #include <optional>
@@ -77,6 +78,11 @@ ConeV getDual(ConeH cone);
 /// The returned cone is pointed at the origin.
 ConeH getDual(ConeV cone);
 
+/// Compute the generating function for a unimodular cone.
+/// It assert-fails if the input cone is not unimodular.
+GeneratingFunction unimodularConeGeneratingFunction(ParamPoint vertex, int sign,
+                                                    ConeH cone);
+
 } // namespace detail
 } // namespace presburger
 } // namespace mlir
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 9152b66968a1f5..f0cabc36e537a9 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -63,3 +63,71 @@ MPInt mlir::presburger::detail::getIndex(ConeV cone) {
 
   return cone.determinant();
 }
+
+/// Compute the generating function for a unimodular cone.
+GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
+    ParamPoint vertex, int sign, ConeH cone) {
+  // `cone` is assumed to be unimodular.
+  assert(getIndex(getDual(cone)) == 1 && "input cone is not unimodular!");
+
+  unsigned numVar = cone.getNumVars();
+  unsigned numIneq = cone.getNumInequalities();
+
+  // Thus its ray matrix, U, is the inverse of the
+  // transpose of its inequality matrix, `cone`.
+  FracMatrix transp(numVar, numIneq);
+  for (unsigned i = 0; i < numVar; i++)
+    for (unsigned j = 0; j < numIneq; j++)
+      transp(j, i) = Fraction(cone.atIneq(i, j), 1);
+
+  FracMatrix generators(numVar, numIneq);
+  transp.determinant(&generators); // This is the U-matrix.
+
+  // The denominators of the generating function
+  // are given by the generators of the cone, i.e.,
+  // the rows of the matrix U.
+  std::vector<Point> denominator(numIneq);
+  ArrayRef<Fraction> row;
+  for (unsigned i = 0; i < numVar; i++) {
+    row = generators.getRow(i);
+    denominator[i] = Point(row);
+  }
+
+  // The vertex is v : [d, n+1].
+  // We need to find affine functions of parameters λi(p)
+  // such that v = Σ λi(p)*ui.
+  // The λi are given by the columns of Λ = v^T @ U^{-1} = v^T @ transp.
+  // Then the numerator will be Σ -floor(-λi(p))*u_i.
+  // Thus we store the numerator as the affine function -Λ,
+  // since the generators are already stored in the denominator.
+  // Note that the outer -1 will have to be accounted for, as it is not stored.
+  // See end for an example.
+
+  unsigned numColumns = vertex.getNumColumns();
+  unsigned numRows = vertex.getNumRows();
+  ParamPoint numerator(numColumns, numRows);
+  SmallVector<Fraction> ithCol(numRows);
+  for (unsigned i = 0; i < numColumns; i++) {
+    for (unsigned j = 0; j < numRows; j++)
+      ithCol[j] = vertex(j, i);
+    numerator.setRow(i, transp.preMultiplyWithRow(ithCol));
+    numerator.negateRow(i);
+  }
+
+  return GeneratingFunction(numColumns - 1, SmallVector<int>(1, sign),
+                            std::vector({numerator}),
+                            std::vector({denominator}));
+
+  // Suppose the vertex is given by the matrix [ 2  2   0], with 2 params
+  //                                           [-1 -1/2 1]
+  // and the cone has H-representation [0  -1] => U-matrix [ 2 -1]
+  //                                   [-1 -2]             [-1  0]
+  // Therefore Λ will be given by [ 1    0 ] and the negation of this will be
+  // stored as the numerator.
+  //                              [ 1/2 -1 ]
+  //                              [ -1  -2 ]
+
+  // Algebraically, the numerator exponent is
+  // [ -2 ⌊ -N - M/2 +1 ⌋ + 1 ⌊ 0 +M +2 ⌋ ] -> first  COLUMN of U is [2, -1]
+  // [  1 ⌊ -N - M/2 +1 ⌋ + 0 ⌊ 0 +M +2 ⌋ ] -> second COLUMN of U is [-1, 0]
+}
diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index b88baa6c6b48a4..2936d95c802e9c 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -46,3 +46,39 @@ TEST(BarvinokTest, getIndex) {
       4, 4, {{4, 2, 5, 1}, {4, 1, 3, 6}, {8, 2, 5, 6}, {5, 2, 5, 7}});
   EXPECT_EQ(getIndex(cone), cone.determinant());
 }
+
+// The following cones and vertices are randomly generated
+// (s.t. the cones are unimodular) and the generating functions
+// are computed. We check that the results contain the correct
+// matrices.
+TEST(BarvinokTest, unimodularConeGeneratingFunction) {
+  ConeH cone = defineHRep(2);
+  cone.addInequality({0, -1, 0});
+  cone.addInequality({-1, -2, 0});
+
+  ParamPoint vertex =
+      makeFracMatrix(2, 3, {{2, 2, 0}, {-1, -Fraction(1, 2), 1}});
+
+  GeneratingFunction gf = unimodularConeGeneratingFunction(vertex, 1, cone);
+
+  EXPECT_EQ_REPR_GENERATINGFUNCTION(
+      gf, GeneratingFunction(
+              2, {1},
+              {makeFracMatrix(3, 2, {{-1, 0}, {-Fraction(1, 2), 1}, {1, 2}})},
+              {{{2, -1}, {-1, 0}}}));
+
+  cone = defineHRep(3);
+  cone.addInequality({7, 1, 6, 0});
+  cone.addInequality({9, 1, 7, 0});
+  cone.addInequality({8, -1, 1, 0});
+
+  vertex = makeFracMatrix(3, 2, {{5, 2}, {6, 2}, {7, 1}});
+
+  gf = unimodularConeGeneratingFunction(vertex, 1, cone);
+
+  EXPECT_EQ_REPR_GENERATINGFUNCTION(
+      gf,
+      GeneratingFunction(
+          1, {1}, {makeFracMatrix(2, 3, {{-83, -100, -41}, {-22, -27, -15}})},
+          {{{8, 47, -17}, {-7, -41, 15}, {1, 5, -2}}}));
+}

>From 7a0acf6e325e7142787e9eea651c3e29a08c745b Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 7 Jan 2024 23:14:33 +0530
Subject: [PATCH 2/5] Improve doc and fix loops

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index f0cabc36e537a9..0efb2ac8003e86 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -65,6 +65,13 @@ MPInt mlir::presburger::detail::getIndex(ConeV cone) {
 }
 
 /// Compute the generating function for a unimodular cone.
+/// This consists of a single term of the form
+/// x^num / prod_j (1 - x^den_j)
+///
+/// den_j is defined as the set of generators of the cone.
+/// num is computed by expressing the vertex as a weighted
+/// sum of the generators, and then taking the floor of the
+/// coefficients.
 GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
     ParamPoint vertex, int sign, ConeH cone) {
   // `cone` is assumed to be unimodular.
@@ -76,8 +83,8 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
   // Thus its ray matrix, U, is the inverse of the
   // transpose of its inequality matrix, `cone`.
   FracMatrix transp(numVar, numIneq);
-  for (unsigned i = 0; i < numVar; i++)
-    for (unsigned j = 0; j < numIneq; j++)
+  for (unsigned i = 0; i < numVar; ++i)
+    for (unsigned j = 0; j < numIneq; ++j)
       transp(j, i) = Fraction(cone.atIneq(i, j), 1);
 
   FracMatrix generators(numVar, numIneq);
@@ -88,7 +95,7 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
   // the rows of the matrix U.
   std::vector<Point> denominator(numIneq);
   ArrayRef<Fraction> row;
-  for (unsigned i = 0; i < numVar; i++) {
+  for (unsigned i = 0; i < numVar; ++i) {
     row = generators.getRow(i);
     denominator[i] = Point(row);
   }
@@ -107,8 +114,8 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
   unsigned numRows = vertex.getNumRows();
   ParamPoint numerator(numColumns, numRows);
   SmallVector<Fraction> ithCol(numRows);
-  for (unsigned i = 0; i < numColumns; i++) {
-    for (unsigned j = 0; j < numRows; j++)
+  for (unsigned i = 0; i < numColumns; ++i) {
+    for (unsigned j = 0; j < numRows; ++j)
       ithCol[j] = vertex(j, i);
     numerator.setRow(i, transp.preMultiplyWithRow(ithCol));
     numerator.negateRow(i);

>From 3353e7f416f290e746c505bb7009a5789db9373d Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 8 Jan 2024 22:27:49 +0530
Subject: [PATCH 3/5] Fix documentation

---
 .../mlir/Analysis/Presburger/Barvinok.h       |  2 +-
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 25 ++++++++++---------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 93b29e2d718e59..213af636e5964d 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -79,7 +79,7 @@ ConeV getDual(ConeH cone);
 ConeH getDual(ConeV cone);
 
 /// Compute the generating function for a unimodular cone.
-/// It assert-fails if the input cone is not unimodular.
+/// The input cone must be unimodular; it assert-fails otherwise.
 GeneratingFunction unimodularConeGeneratingFunction(ParamPoint vertex, int sign,
                                                     ConeH cone);
 
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 0efb2ac8003e86..f91744f754f96e 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -74,7 +74,7 @@ MPInt mlir::presburger::detail::getIndex(ConeV cone) {
 /// coefficients.
 GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
     ParamPoint vertex, int sign, ConeH cone) {
-  // `cone` is assumed to be unimodular.
+  // `cone` must be unimodular.
   assert(getIndex(getDual(cone)) == 1 && "input cone is not unimodular!");
 
   unsigned numVar = cone.getNumVars();
@@ -88,11 +88,11 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
       transp(j, i) = Fraction(cone.atIneq(i, j), 1);
 
   FracMatrix generators(numVar, numIneq);
-  transp.determinant(&generators); // This is the U-matrix.
+  transp.determinant(/*inverse=*/&generators); // This is the U-matrix.
 
-  // The denominators of the generating function
-  // are given by the generators of the cone, i.e.,
-  // the rows of the matrix U.
+  // The powers in the denominator of the generating
+  // function are given by the generators of the cone,
+  // i.e., the rows of the matrix U.
   std::vector<Point> denominator(numIneq);
   ArrayRef<Fraction> row;
   for (unsigned i = 0; i < numVar; ++i) {
@@ -100,13 +100,14 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
     denominator[i] = Point(row);
   }
 
-  // The vertex is v : [d, n+1].
-  // We need to find affine functions of parameters λi(p)
-  // such that v = Σ λi(p)*ui.
-  // The λi are given by the columns of Λ = v^T @ U^{-1} = v^T @ transp.
-  // Then the numerator will be Σ -floor(-λi(p))*u_i.
-  // Thus we store the numerator as the affine function -Λ,
-  // since the generators are already stored in the denominator.
+  // The vertex is v : d x (n+1)
+  // We need to find affine functions of parameters λ_i(p)
+  // such that v = Σ λ_i(p)*u_i.
+  // The λi are given by the columns of Λ = v^T U^{-1} = v^T transp.
+  // Then the exponent in the numerator will be
+  // Σ -floor(-λ_i(p))*u_i.
+  // Thus we store the (exponent of the) numerator as the affine function -Λ,
+  // since the generators are already stored as the exponent of the denominator.
   // Note that the outer -1 will have to be accounted for, as it is not stored.
   // See end for an example.
 

>From 2ad8dcf560ecde87e5f477ae3a2b89b0ad60fe69 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 9 Jan 2024 00:51:42 +0530
Subject: [PATCH 4/5] Use range-style loops

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index f91744f754f96e..434719f66c5fc7 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/Barvinok.h"
+#include "llvm/ADT/Sequence.h"
 
 using namespace mlir;
 using namespace presburger;
@@ -24,7 +25,7 @@ ConeV mlir::presburger::detail::getDual(ConeH cone) {
   // is represented as a row [a1, ..., an, b]
   // and that b = 0.
 
-  for (unsigned i = 0; i < numIneq; ++i) {
+  for (auto i : llvm::seq<int>(0, numIneq)) {
     assert(cone.atIneq(i, numVar) == 0 &&
            "H-representation of cone is not centred at the origin!");
     for (unsigned j = 0; j < numVar; ++j) {
@@ -83,8 +84,8 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
   // Thus its ray matrix, U, is the inverse of the
   // transpose of its inequality matrix, `cone`.
   FracMatrix transp(numVar, numIneq);
-  for (unsigned i = 0; i < numVar; ++i)
-    for (unsigned j = 0; j < numIneq; ++j)
+  for (auto i : llvm::seq<int>(0, numVar))
+    for (auto j : llvm::seq<int>(0, numIneq))
       transp(j, i) = Fraction(cone.atIneq(i, j), 1);
 
   FracMatrix generators(numVar, numIneq);
@@ -95,7 +96,7 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
   // i.e., the rows of the matrix U.
   std::vector<Point> denominator(numIneq);
   ArrayRef<Fraction> row;
-  for (unsigned i = 0; i < numVar; ++i) {
+  for (auto i : llvm::seq<int>(0, numVar)) {
     row = generators.getRow(i);
     denominator[i] = Point(row);
   }
@@ -115,8 +116,8 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
   unsigned numRows = vertex.getNumRows();
   ParamPoint numerator(numColumns, numRows);
   SmallVector<Fraction> ithCol(numRows);
-  for (unsigned i = 0; i < numColumns; ++i) {
-    for (unsigned j = 0; j < numRows; ++j)
+  for (auto i : llvm::seq<int>(0, numColumns)) {
+    for (auto j : llvm::seq<int>(0, numRows))
       ithCol[j] = vertex(j, i);
     numerator.setRow(i, transp.preMultiplyWithRow(ithCol));
     numerator.negateRow(i);

>From feaad7957eb07fac3da57776523914d81e545dec Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 9 Jan 2024 00:59:10 +0530
Subject: [PATCH 5/5] Define getInequalities for IntegerRelation, transpose for
 Matrix<T>, and use in ucgf

---
 .../include/mlir/Analysis/Presburger/IntegerRelation.h |  2 ++
 mlir/include/mlir/Analysis/Presburger/Matrix.h         |  3 +++
 mlir/lib/Analysis/Presburger/Barvinok.cpp              |  8 ++++----
 mlir/lib/Analysis/Presburger/Matrix.cpp                | 10 ++++++++++
 4 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index 4c6b810f92e95a..41aeda07ac7912 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -221,6 +221,8 @@ class IntegerRelation {
     return getInt64Vec(inequalities.getRow(idx));
   }
 
+  inline IntMatrix getInequalities() const { return inequalities; }
+
   /// Get the number of vars of the specified kind.
   unsigned getNumVarKind(VarKind kind) const {
     return space.getNumVarKind(kind);
diff --git a/mlir/include/mlir/Analysis/Presburger/Matrix.h b/mlir/include/mlir/Analysis/Presburger/Matrix.h
index 347e2e0489786f..5502528a087b25 100644
--- a/mlir/include/mlir/Analysis/Presburger/Matrix.h
+++ b/mlir/include/mlir/Analysis/Presburger/Matrix.h
@@ -181,6 +181,9 @@ class Matrix {
   /// `elems` must be equal to the number of columns.
   unsigned appendExtraRow(ArrayRef<T> elems);
 
+  // Transpose the matrix without modifying it.
+  Matrix<T> transpose();
+
   /// Print the matrix.
   void print(raw_ostream &os) const;
   void dump() const;
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 434719f66c5fc7..d771854717225e 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -83,10 +83,10 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
 
   // Thus its ray matrix, U, is the inverse of the
   // transpose of its inequality matrix, `cone`.
-  FracMatrix transp(numVar, numIneq);
-  for (auto i : llvm::seq<int>(0, numVar))
-    for (auto j : llvm::seq<int>(0, numIneq))
-      transp(j, i) = Fraction(cone.atIneq(i, j), 1);
+  // The last column of the inequality matrix is null,
+  // so we remove it to obtain a square matrix.
+  FracMatrix transp = FracMatrix(cone.getInequalities()).transpose();
+  transp.removeRow(numVar);
 
   FracMatrix generators(numVar, numIneq);
   transp.determinant(/*inverse=*/&generators); // This is the U-matrix.
diff --git a/mlir/lib/Analysis/Presburger/Matrix.cpp b/mlir/lib/Analysis/Presburger/Matrix.cpp
index b68a7b7004bba9..c53bef467e632b 100644
--- a/mlir/lib/Analysis/Presburger/Matrix.cpp
+++ b/mlir/lib/Analysis/Presburger/Matrix.cpp
@@ -62,6 +62,16 @@ unsigned Matrix<T>::appendExtraRow(ArrayRef<T> elems) {
   return row;
 }
 
+template <typename T>
+Matrix<T> Matrix<T>::transpose() {
+  Matrix<T> transp(nColumns, nRows);
+  for (unsigned row = 0; row < nRows; ++row)
+    for (unsigned col = 0; col < nColumns; ++col)
+      transp(col, row) = at(row, col);
+
+  return transp;
+}
+
 template <typename T>
 void Matrix<T>::resizeHorizontally(unsigned newNColumns) {
   if (newNColumns < nColumns)



More information about the Mlir-commits mailing list