[Mlir-commits] [mlir] [MLIR][Presburger] Implement function to evaluate the number of terms in a generating function. (PR #78078)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Sun Jan 14 10:38:20 PST 2024


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

>From 1ec4bd325ab80f6d9f1dfbe1db1100ce9fb22219 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 09:39:25 +0530
Subject: [PATCH 01/14] Initial commit

---
 .../mlir/Analysis/Presburger/Barvinok.h       |   4 +
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 197 ++++++++++++++++++
 2 files changed, 201 insertions(+)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index edee19f0e1a535..2e2273dab4bc9d 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -99,6 +99,10 @@ QuasiPolynomial getCoefficientInRationalFunction(unsigned power,
                                                  ArrayRef<QuasiPolynomial> num,
                                                  ArrayRef<Fraction> den);
 
+/// Substitute the generating function with the unit vector
+/// to find the number of terms.
+QuasiPolynomial substituteWithUnitVector(GeneratingFunction);
+
 } // namespace detail
 } // namespace presburger
 } // namespace mlir
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 4ba4462af0317f..0e9d3be7a3289f 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -245,3 +245,200 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
   }
   return coefficients[power].simplify();
 }
+
+// We have a generating function of the form
+// f_p(x) = \sum_i sign_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
+//
+// where sign_i is ±1,
+// n_i \in Q^p -> Q^d is a d-vector of affine functions on p parameters, and
+// d_{ij} \in Q^d are vectors.
+//
+// We need to find the number of terms of the form x^t in the expansion of
+// this function, for which we substitute x = (1, ..., 1).
+// However, direct substitution leads to an undefined answer due to the
+// form of the denominator.
+//
+// We therefore use the following procedure instead:
+// Find a vector μ that is not orthogonal to any of the d_{ij}.
+// Substitute x_i = (s+1)^μ_i. As μ_i is not orthogonal to d_{ij},
+// we never have (1 - (s+1)^0) = 0 in any of the terms in denominator.
+// We then find the constant term in this function, i.e., we evaluate it
+// at s = 0, which is equivalent to x = (1, ..., 1).
+//
+// Now, we have a function of the form
+// f_p(s) = \sum_i sign_i * (s+1)^n_i / (\prod_j (1 - (s+1)^d'_{ij}))
+// in which we need to find the constant term.
+// For the i'th term, we first convert all the d'_{ij} to their
+// absolute values by multiplying and dividing by (s+1)^(-d'_{ij}) if it is
+// negative. We change the sign accordingly.
+// Then, we replace each (1 - (s+1)^(d'_{ij})) with
+// (-s)(\sum_{0 ≤ k < d'_{ij}} (s+1)^k).
+// Thus the term overall has now the form
+// sign'_i * (s+1)^n'_i / (s^r * \prod_j (\sum_k (s+1)^k)).
+// This means that
+// the numerator is a polynomial in s, with coefficients as quasipolynomials,
+// and the denominator is polynomial in s, with fractional coefficients.
+// We need to find the constant term in the expansion of this term,
+// which is the same as finding the coefficient of s^r in
+// sign'_i * (s+1)^n'_i / (\prod_j (\sum_k (s+1)^k)),
+// for which we use the `getCoefficientInRationalFunction()` function.
+//
+// Verdoolaege, Sven, et al. "Counting integer points in parametric polytopes
+// using Barvinok's rational functions." Algorithmica 48 (2007): 37-66.
+QuasiPolynomial
+mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
+  std::vector<Point> allDenominators;
+  for (std::vector<Point> den : gf.getDenominators())
+    allDenominators.insert(allDenominators.end(), den.begin(), den.end());
+  Point mu = getNonOrthogonalVector(allDenominators);
+
+  unsigned num_params = gf.getNumParams();
+  unsigned num_dims = mu.size();
+  unsigned num_terms = gf.getDenominators().size();
+
+  std::vector<Fraction> dens;
+
+  std::vector<QuasiPolynomial> numeratorCoefficients;
+  std::vector<Fraction> singleTermDenCoefficients, denominatorCoefficients;
+  std::vector<std::vector<Fraction>> eachTermDenCoefficients;
+  std::vector<Fraction> convolution;
+
+  QuasiPolynomial totalTerm(num_params, 0);
+  for (unsigned i = 0; i < num_terms; i++) {
+    int sign = gf.getSigns()[i];
+    ParamPoint v = gf.getNumerators()[i];
+    std::vector<Point> ds = gf.getDenominators()[i];
+
+    // Substitute x_i = (s+1)^μ_i
+    // Then the exponent in the numerator becomes
+    // - (μ • u_1) * (floor(first col of v))
+    // - (μ • u_2) * (floor(second col of v)) - ...
+    // - (μ • u_d) * (floor(d'th col of v))
+    // So we store the negation of the  dot produts.
+
+    // We have d terms, each of whose coefficient is the negative dot product,
+    SmallVector<Fraction> coefficients;
+    coefficients.reserve(num_dims);
+    for (Point d : ds)
+      coefficients.push_back(-dotProduct(mu, d));
+
+    // and whose affine fn is a single floor expression, given by the
+    // corresponding column of v.
+    std::vector<std::vector<SmallVector<Fraction>>> affine(num_dims);
+    for (unsigned j = 0; j < num_dims; j++)
+      SmallVector<Fraction> jthCol(v.transpose().getRow(j));
+
+    QuasiPolynomial num(num_params, coefficients, affine);
+    num = num.simplify();
+
+    // Now the numerator is (s+1)^num.
+
+    dens.clear();
+    // Similarly, each term in the denominator has exponent
+    // given by the dot product of μ with u_i.
+    for (Point d : ds)
+      dens.push_back(dotProduct(d, mu));
+    // This term in the denominator is
+    // (1 - (s+1)^dens.back())
+
+    // We track the number of exponents that are negative in the
+    // denominator, and convert them to their absolute values
+    // (see lines 361-71).
+    unsigned numNegExps = 0;
+    Fraction sumNegExps(0, 1);
+    for (unsigned j = 0; j < dens.size(); j++) {
+      if (dens[j] < Fraction(0, 1)) {
+        numNegExps += 1;
+        sumNegExps = sumNegExps + dens[j];
+      }
+      // All exponents will be made positive; then reduce
+      // (1 - (s+1)^x)
+      // to
+      // (-s)*(Σ_{x-1} (s+1)^j) because x > 0
+      dens[j] = abs(dens[j]) - 1;
+    }
+
+    // If we have (1 - (s+1)^-c) in the denominator,
+    // multiply and divide by (s+1)^c.
+    // We convert all negative-exponent terms at once; therefore
+    // we multiply and divide by (s+1)^sumNegExps.
+    // Then we get
+    // -(1 - (s+1)^c) in the denominator,
+    // increase the numerator by c, and
+    // flip the sign.
+    if (numNegExps % 2 == 1)
+      sign = -sign;
+    num = num - QuasiPolynomial(num_params, sumNegExps);
+
+    // Take all the (-s) out, from line 328.
+    unsigned r = dens.size();
+    if (r % 2 == 1)
+      sign = -sign;
+
+    // Now the expression is
+    // (s+1)^num /
+    // (s^r * Π_(0 ≤ i < r) (Σ_{0 ≤ j ≤ dens[i]} (s+1)^j))
+
+    // Letting P(s) = (s+1)^num and Q(s) = Π_r (...),
+    // we need to find the coefficient of s^r in
+    // P(s)/Q(s).
+
+    // First, the coefficients of P(s), which are binomial coefficients.
+    // We need r+1 of these.
+    numeratorCoefficients.clear();
+    numeratorCoefficients.push_back(
+        QuasiPolynomial(num_params, 1)); // Coeff of s^0
+    for (unsigned j = 1; j <= r; j++)
+      numeratorCoefficients.push_back(
+          (numeratorCoefficients[j - 1] *
+           (num - QuasiPolynomial(num_params, j - 1)) / Fraction(j, 1))
+              .simplify());
+    // Coeff of s^j
+
+    // Then the coefficients of each individual term in Q(s),
+    // which are (di+1) C (k+1) for 0 ≤ k ≤ di
+    eachTermDenCoefficients.clear();
+    for (Fraction den : dens) {
+      singleTermDenCoefficients.clear();
+      singleTermDenCoefficients.push_back(den + 1);
+      for (unsigned j = 1; j <= den; j++)
+        singleTermDenCoefficients.push_back(singleTermDenCoefficients[j - 1] *
+                                            (den - (j - 1)) / (j + 1));
+
+      eachTermDenCoefficients.push_back(singleTermDenCoefficients);
+    }
+
+    // Now we find the coefficients in Q(s) itself
+    // by taking the convolution of the coefficients
+    // of all the terms.
+    denominatorCoefficients.clear();
+    denominatorCoefficients = eachTermDenCoefficients[0];
+    for (unsigned j = 1; j < eachTermDenCoefficients.size(); j++) {
+      // The length of the convolution is the maximum of the lengths
+      // of the two sequences. We pad the shorter one with zeroes.
+      unsigned convlen = std::max(denominatorCoefficients.size(),
+                                  eachTermDenCoefficients[j].size());
+      for (unsigned k = denominatorCoefficients.size(); k < convlen; k++)
+        denominatorCoefficients.push_back(0);
+      for (unsigned k = eachTermDenCoefficients[j].size(); k < convlen; k++)
+        eachTermDenCoefficients[j].push_back(0);
+
+      convolution.clear();
+      for (unsigned k = 0; k < convlen; k++) {
+        Fraction sum(0, 1);
+        for (unsigned l = 0; l <= k; l++)
+          sum = sum +
+                denominatorCoefficients[l] * eachTermDenCoefficients[j][k - l];
+        convolution.push_back(sum);
+      }
+      denominatorCoefficients = convolution;
+    }
+
+    totalTerm =
+        totalTerm + getCoefficientInRationalFunction(r, numeratorCoefficients,
+                                                     denominatorCoefficients) *
+                        QuasiPolynomial(num_params, sign);
+  }
+
+  return totalTerm.simplify();
+}
\ No newline at end of file

>From f286b7b216f98a139b78d7339086dd8c650d336f Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 18:17:25 +0530
Subject: [PATCH 02/14] Bug fix

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp     |  9 +--
 .../Analysis/Presburger/BarvinokTest.cpp      | 56 +++++++++++++++++++
 2 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 0e9d3be7a3289f..785f61b578242f 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -324,9 +324,10 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
 
     // and whose affine fn is a single floor expression, given by the
     // corresponding column of v.
-    std::vector<std::vector<SmallVector<Fraction>>> affine(num_dims);
+    std::vector<std::vector<SmallVector<Fraction>>> affine;
+    affine.reserve(num_dims);
     for (unsigned j = 0; j < num_dims; j++)
-      SmallVector<Fraction> jthCol(v.transpose().getRow(j));
+      affine.push_back({SmallVector<Fraction>(v.transpose().getRow(j))});
 
     QuasiPolynomial num(num_params, coefficients, affine);
     num = num.simplify();
@@ -343,7 +344,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
 
     // We track the number of exponents that are negative in the
     // denominator, and convert them to their absolute values
-    // (see lines 361-71).
+    // (see lines 362-72).
     unsigned numNegExps = 0;
     Fraction sumNegExps(0, 1);
     for (unsigned j = 0; j < dens.size(); j++) {
@@ -370,7 +371,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
       sign = -sign;
     num = num - QuasiPolynomial(num_params, sumNegExps);
 
-    // Take all the (-s) out, from line 328.
+    // Take all the (-s) out, from line 359.
     unsigned r = dens.size();
     if (r % 2 == 1)
       sign = -sign;
diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index e304f81de21f0b..8badf8983a2a39 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -124,3 +124,59 @@ TEST(BarvinokTest, getCoefficientInRationalFunction) {
   coeff = getCoefficientInRationalFunction(3, numerator, denominator);
   EXPECT_EQ(coeff.getConstantTerm(), Fraction(55, 64));
 }
+
+// The following test is taken from
+// 
+TEST(BarvinokTest, substituteWithUnitVector) {
+  GeneratingFunction gf(
+      1, {1, 1, 1},
+      {makeFracMatrix(2, 2, {{0, Fraction(1, 2)}, {0, 0}}),
+       makeFracMatrix(2, 2, {{0, Fraction(1, 2)}, {0, 0}}),
+       makeFracMatrix(2, 2, {{0, 0}, {0, 0}})},
+      {{{-1, 1}, {-1, 0}},
+       {{1, -1}, {0, -1}},
+       {{1, 0}, {0, 1}}});
+
+  QuasiPolynomial numPoints = substituteWithUnitVector(gf);
+  EXPECT_EQ(
+      numPoints.getCoefficients(),
+      SmallVector<Fraction>(
+          {Fraction(-1, 2), Fraction(-1, 2), Fraction(1, 2), Fraction(-1, 2),
+           Fraction(-1, 2), Fraction(1, 2), Fraction(1, 1), Fraction(3, 2),
+           Fraction(-1, 2), Fraction(3, 2), Fraction(9, 4), Fraction(-3, 4),
+           Fraction(-1, 2), Fraction(-3, 4), Fraction(1, 8), Fraction(1, 8)}));
+  EXPECT_EQ(
+      numPoints.getAffine(),
+      std::vector<std::vector<Point>>(
+          {{{Fraction(1, 2), Fraction(0, 1)}, {Fraction(1, 2), Fraction(0, 1)}},
+           {{Fraction(1, 2), Fraction(0, 1)}},
+           {{Fraction(1, 2), Fraction(0, 1)}},
+           {{Fraction(1, 2), Fraction(0, 1)}},
+           {},
+           {},
+           {{Fraction(1, 2), Fraction(0, 1)}, {Fraction(1, 2), Fraction(0, 1)}},
+           {{Fraction(1, 2), Fraction(0, 1)}},
+           {{Fraction(1, 2), Fraction(0, 1)}},
+           {{Fraction(1, 2), Fraction(0, 1)}},
+           {},
+           {},
+           {{Fraction(1, 2), Fraction(0, 1)}},
+           {},
+           {},
+           {}}));
+
+  // We can gather the like terms because we know there's only
+  // either ⌊p/2⌋^2, ⌊p/2⌋, or constants.
+  Fraction pSquaredCoeff = 0, pCoeff = 0, constantTerm = 0;
+  for (unsigned i = 0; i < numPoints.getCoefficients().size(); i++)
+    if (numPoints.getAffine()[i].size() == 2)
+      pSquaredCoeff = pSquaredCoeff + numPoints.getCoefficients()[i];
+    else if (numPoints.getAffine()[i].size() == 1)
+      pCoeff = pCoeff + numPoints.getCoefficients()[i];
+    else
+      constantTerm = constantTerm + numPoints.getCoefficients()[i];
+
+  EXPECT_EQ(pSquaredCoeff, Fraction(1, 2));
+  EXPECT_EQ(pCoeff, Fraction(3, 2));
+  EXPECT_EQ(constantTerm, Fraction(1, 1));
+}
\ No newline at end of file

>From 19ade35598076669bccebfbc184b7956a1cd9fec Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 18:24:09 +0530
Subject: [PATCH 03/14] Add test

---
 .../Analysis/Presburger/BarvinokTest.cpp      | 54 +++++++------------
 1 file changed, 19 insertions(+), 35 deletions(-)

diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index 8badf8983a2a39..2842fc14eb8298 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -126,57 +126,41 @@ TEST(BarvinokTest, getCoefficientInRationalFunction) {
 }
 
 // The following test is taken from
-// 
+//
 TEST(BarvinokTest, substituteWithUnitVector) {
   GeneratingFunction gf(
       1, {1, 1, 1},
       {makeFracMatrix(2, 2, {{0, Fraction(1, 2)}, {0, 0}}),
        makeFracMatrix(2, 2, {{0, Fraction(1, 2)}, {0, 0}}),
        makeFracMatrix(2, 2, {{0, 0}, {0, 0}})},
-      {{{-1, 1}, {-1, 0}},
-       {{1, -1}, {0, -1}},
-       {{1, 0}, {0, 1}}});
+      {{{-1, 1}, {-1, 0}}, {{1, -1}, {0, -1}}, {{1, 0}, {0, 1}}});
 
   QuasiPolynomial numPoints = substituteWithUnitVector(gf);
-  EXPECT_EQ(
-      numPoints.getCoefficients(),
-      SmallVector<Fraction>(
-          {Fraction(-1, 2), Fraction(-1, 2), Fraction(1, 2), Fraction(-1, 2),
-           Fraction(-1, 2), Fraction(1, 2), Fraction(1, 1), Fraction(3, 2),
-           Fraction(-1, 2), Fraction(3, 2), Fraction(9, 4), Fraction(-3, 4),
-           Fraction(-1, 2), Fraction(-3, 4), Fraction(1, 8), Fraction(1, 8)}));
-  EXPECT_EQ(
-      numPoints.getAffine(),
-      std::vector<std::vector<Point>>(
-          {{{Fraction(1, 2), Fraction(0, 1)}, {Fraction(1, 2), Fraction(0, 1)}},
-           {{Fraction(1, 2), Fraction(0, 1)}},
-           {{Fraction(1, 2), Fraction(0, 1)}},
-           {{Fraction(1, 2), Fraction(0, 1)}},
-           {},
-           {},
-           {{Fraction(1, 2), Fraction(0, 1)}, {Fraction(1, 2), Fraction(0, 1)}},
-           {{Fraction(1, 2), Fraction(0, 1)}},
-           {{Fraction(1, 2), Fraction(0, 1)}},
-           {{Fraction(1, 2), Fraction(0, 1)}},
-           {},
-           {},
-           {{Fraction(1, 2), Fraction(0, 1)}},
-           {},
-           {},
-           {}}));
-
-  // We can gather the like terms because we know there's only
+
+  // First, we make sure that all the affine functions are of the form ⌊p/2⌋.
+  for (const std::vector<SmallVector<Fraction>> &term : numPoints.getAffine()) {
+    for (const SmallVector<Fraction> &aff : term) {
+      EXPECT_EQ(aff.size(), 2u);
+      EXPECT_EQ(aff[0], Fraction(1, 2));
+      EXPECT_EQ(aff[1], Fraction(0, 1));
+    }
+  }
+
+  // Now, we can gather the like terms because we know there's only
   // either ⌊p/2⌋^2, ⌊p/2⌋, or constants.
+  // The total coefficient of ⌊p/2⌋^2 is the sum of coefficients of all
+  // terms with 2 affine functions, and
+  // the coefficient of total ⌊p/2⌋ is the sum of coefficients of all
+  // terms with 1 affine function,
   Fraction pSquaredCoeff = 0, pCoeff = 0, constantTerm = 0;
   for (unsigned i = 0; i < numPoints.getCoefficients().size(); i++)
     if (numPoints.getAffine()[i].size() == 2)
       pSquaredCoeff = pSquaredCoeff + numPoints.getCoefficients()[i];
     else if (numPoints.getAffine()[i].size() == 1)
       pCoeff = pCoeff + numPoints.getCoefficients()[i];
-    else
-      constantTerm = constantTerm + numPoints.getCoefficients()[i];
 
+  // We expect the answer to be (1/2)⌊p/2⌋^2 + (3/2)⌊p/2⌋ + 1.
   EXPECT_EQ(pSquaredCoeff, Fraction(1, 2));
   EXPECT_EQ(pCoeff, Fraction(3, 2));
-  EXPECT_EQ(constantTerm, Fraction(1, 1));
+  EXPECT_EQ(numPoints.getConstantTerm(), Fraction(1, 1));
 }
\ No newline at end of file

>From 39dd6a8a1c87ac33d6b6e31a92374c7a2409892c Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 18:27:59 +0530
Subject: [PATCH 04/14] Minor fixes

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 785f61b578242f..e580e40f6034cc 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -288,7 +288,7 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
 QuasiPolynomial
 mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
   std::vector<Point> allDenominators;
-  for (std::vector<Point> den : gf.getDenominators())
+  for (ArrayRef<Point> den : gf.getDenominators())
     allDenominators.insert(allDenominators.end(), den.begin(), den.end());
   Point mu = getNonOrthogonalVector(allDenominators);
 
@@ -304,7 +304,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
   std::vector<Fraction> convolution;
 
   QuasiPolynomial totalTerm(num_params, 0);
-  for (unsigned i = 0; i < num_terms; i++) {
+  for (unsigned i = 0; i < num_terms; ++i) {
     int sign = gf.getSigns()[i];
     ParamPoint v = gf.getNumerators()[i];
     std::vector<Point> ds = gf.getDenominators()[i];
@@ -326,7 +326,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     // corresponding column of v.
     std::vector<std::vector<SmallVector<Fraction>>> affine;
     affine.reserve(num_dims);
-    for (unsigned j = 0; j < num_dims; j++)
+    for (unsigned j = 0; j < num_dims; ++j)
       affine.push_back({SmallVector<Fraction>(v.transpose().getRow(j))});
 
     QuasiPolynomial num(num_params, coefficients, affine);
@@ -347,10 +347,10 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     // (see lines 362-72).
     unsigned numNegExps = 0;
     Fraction sumNegExps(0, 1);
-    for (unsigned j = 0; j < dens.size(); j++) {
-      if (dens[j] < Fraction(0, 1)) {
+    for (unsigned j = 0, e = dens.size(); j < e; ++j) {
+      if (dens[j] < 0) {
         numNegExps += 1;
-        sumNegExps = sumNegExps + dens[j];
+        sumNegExps += dens[j];
       }
       // All exponents will be made positive; then reduce
       // (1 - (s+1)^x)
@@ -389,7 +389,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     numeratorCoefficients.clear();
     numeratorCoefficients.push_back(
         QuasiPolynomial(num_params, 1)); // Coeff of s^0
-    for (unsigned j = 1; j <= r; j++)
+    for (unsigned j = 1; j <= r; ++j)
       numeratorCoefficients.push_back(
           (numeratorCoefficients[j - 1] *
            (num - QuasiPolynomial(num_params, j - 1)) / Fraction(j, 1))
@@ -402,7 +402,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     for (Fraction den : dens) {
       singleTermDenCoefficients.clear();
       singleTermDenCoefficients.push_back(den + 1);
-      for (unsigned j = 1; j <= den; j++)
+      for (unsigned j = 1; j <= den; ++j)
         singleTermDenCoefficients.push_back(singleTermDenCoefficients[j - 1] *
                                             (den - (j - 1)) / (j + 1));
 
@@ -414,20 +414,20 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     // of all the terms.
     denominatorCoefficients.clear();
     denominatorCoefficients = eachTermDenCoefficients[0];
-    for (unsigned j = 1; j < eachTermDenCoefficients.size(); j++) {
+    for (unsigned j = 1, e = eachTermDenCoefficients.size(); j < e; ++j) {
       // The length of the convolution is the maximum of the lengths
       // of the two sequences. We pad the shorter one with zeroes.
       unsigned convlen = std::max(denominatorCoefficients.size(),
                                   eachTermDenCoefficients[j].size());
-      for (unsigned k = denominatorCoefficients.size(); k < convlen; k++)
+      for (unsigned k = denominatorCoefficients.size(); k < convlen; ++k)
         denominatorCoefficients.push_back(0);
-      for (unsigned k = eachTermDenCoefficients[j].size(); k < convlen; k++)
+      for (unsigned k = eachTermDenCoefficients[j].size(); k < convlen; ++k)
         eachTermDenCoefficients[j].push_back(0);
 
       convolution.clear();
-      for (unsigned k = 0; k < convlen; k++) {
+      for (unsigned k = 0; k < convlen; ++k) {
         Fraction sum(0, 1);
-        for (unsigned l = 0; l <= k; l++)
+        for (unsigned l = 0; l <= k; ++l)
           sum = sum +
                 denominatorCoefficients[l] * eachTermDenCoefficients[j][k - l];
         convolution.push_back(sum);

>From 64a55bc0a8991270f3a7e7ff56c33de6fba3de8b Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 18:33:17 +0530
Subject: [PATCH 05/14] Use const

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index e580e40f6034cc..180c3f047c6010 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -319,7 +319,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     // We have d terms, each of whose coefficient is the negative dot product,
     SmallVector<Fraction> coefficients;
     coefficients.reserve(num_dims);
-    for (Point d : ds)
+    for (const Point &d : ds)
       coefficients.push_back(-dotProduct(mu, d));
 
     // and whose affine fn is a single floor expression, given by the
@@ -337,7 +337,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     dens.clear();
     // Similarly, each term in the denominator has exponent
     // given by the dot product of μ with u_i.
-    for (Point d : ds)
+    for (const Point &d : ds)
       dens.push_back(dotProduct(d, mu));
     // This term in the denominator is
     // (1 - (s+1)^dens.back())
@@ -399,7 +399,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     // Then the coefficients of each individual term in Q(s),
     // which are (di+1) C (k+1) for 0 ≤ k ≤ di
     eachTermDenCoefficients.clear();
-    for (Fraction den : dens) {
+    for (const Fraction &den : dens) {
       singleTermDenCoefficients.clear();
       singleTermDenCoefficients.push_back(den + 1);
       for (unsigned j = 1; j <= den; ++j)

>From b3331b2c736d1a99fe360c90e17fd98ce2ac08b0 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 18:41:09 +0530
Subject: [PATCH 06/14] Push down declarations

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp | 51 +++++++++++------------
 1 file changed, 25 insertions(+), 26 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 180c3f047c6010..8b3a91acdc83a9 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -292,19 +292,12 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     allDenominators.insert(allDenominators.end(), den.begin(), den.end());
   Point mu = getNonOrthogonalVector(allDenominators);
 
-  unsigned num_params = gf.getNumParams();
-  unsigned num_dims = mu.size();
-  unsigned num_terms = gf.getDenominators().size();
+  unsigned numParams = gf.getNumParams();
+  unsigned numDims = mu.size();
+  unsigned numTerms = gf.getDenominators().size();
 
-  std::vector<Fraction> dens;
-
-  std::vector<QuasiPolynomial> numeratorCoefficients;
-  std::vector<Fraction> singleTermDenCoefficients, denominatorCoefficients;
-  std::vector<std::vector<Fraction>> eachTermDenCoefficients;
-  std::vector<Fraction> convolution;
-
-  QuasiPolynomial totalTerm(num_params, 0);
-  for (unsigned i = 0; i < num_terms; ++i) {
+  QuasiPolynomial totalTerm(numParams, 0);
+  for (unsigned i = 0; i < numTerms; ++i) {
     int sign = gf.getSigns()[i];
     ParamPoint v = gf.getNumerators()[i];
     std::vector<Point> ds = gf.getDenominators()[i];
@@ -318,23 +311,24 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
 
     // We have d terms, each of whose coefficient is the negative dot product,
     SmallVector<Fraction> coefficients;
-    coefficients.reserve(num_dims);
+    coefficients.reserve(numDims);
     for (const Point &d : ds)
       coefficients.push_back(-dotProduct(mu, d));
 
     // and whose affine fn is a single floor expression, given by the
     // corresponding column of v.
     std::vector<std::vector<SmallVector<Fraction>>> affine;
-    affine.reserve(num_dims);
-    for (unsigned j = 0; j < num_dims; ++j)
+    affine.reserve(numDims);
+    for (unsigned j = 0; j < numDims; ++j)
       affine.push_back({SmallVector<Fraction>(v.transpose().getRow(j))});
 
-    QuasiPolynomial num(num_params, coefficients, affine);
+    QuasiPolynomial num(numParams, coefficients, affine);
     num = num.simplify();
 
     // Now the numerator is (s+1)^num.
 
-    dens.clear();
+    std::vector<Fraction> dens;
+    dens.reserve(ds.size());
     // Similarly, each term in the denominator has exponent
     // given by the dot product of μ with u_i.
     for (const Point &d : ds)
@@ -344,7 +338,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
 
     // We track the number of exponents that are negative in the
     // denominator, and convert them to their absolute values
-    // (see lines 362-72).
+    // (see lines 356-66).
     unsigned numNegExps = 0;
     Fraction sumNegExps(0, 1);
     for (unsigned j = 0, e = dens.size(); j < e; ++j) {
@@ -369,9 +363,9 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     // flip the sign.
     if (numNegExps % 2 == 1)
       sign = -sign;
-    num = num - QuasiPolynomial(num_params, sumNegExps);
+    num = num - QuasiPolynomial(numParams, sumNegExps);
 
-    // Take all the (-s) out, from line 359.
+    // Take all the (-s) out, from line 353.
     unsigned r = dens.size();
     if (r % 2 == 1)
       sign = -sign;
@@ -386,19 +380,22 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
 
     // First, the coefficients of P(s), which are binomial coefficients.
     // We need r+1 of these.
-    numeratorCoefficients.clear();
+    std::vector<QuasiPolynomial> numeratorCoefficients;
+    numeratorCoefficients.reserve(r + 1);
     numeratorCoefficients.push_back(
-        QuasiPolynomial(num_params, 1)); // Coeff of s^0
+        QuasiPolynomial(numParams, 1)); // Coeff of s^0
     for (unsigned j = 1; j <= r; ++j)
       numeratorCoefficients.push_back(
           (numeratorCoefficients[j - 1] *
-           (num - QuasiPolynomial(num_params, j - 1)) / Fraction(j, 1))
+           (num - QuasiPolynomial(numParams, j - 1)) / Fraction(j, 1))
               .simplify());
     // Coeff of s^j
 
     // Then the coefficients of each individual term in Q(s),
     // which are (di+1) C (k+1) for 0 ≤ k ≤ di
-    eachTermDenCoefficients.clear();
+    std::vector<std::vector<Fraction>> eachTermDenCoefficients;
+    std::vector<Fraction> singleTermDenCoefficients;
+    eachTermDenCoefficients.reserve(r);
     for (const Fraction &den : dens) {
       singleTermDenCoefficients.clear();
       singleTermDenCoefficients.push_back(den + 1);
@@ -412,7 +409,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     // Now we find the coefficients in Q(s) itself
     // by taking the convolution of the coefficients
     // of all the terms.
-    denominatorCoefficients.clear();
+    std::vector<Fraction> denominatorCoefficients;
     denominatorCoefficients = eachTermDenCoefficients[0];
     for (unsigned j = 1, e = eachTermDenCoefficients.size(); j < e; ++j) {
       // The length of the convolution is the maximum of the lengths
@@ -424,6 +421,8 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
       for (unsigned k = eachTermDenCoefficients[j].size(); k < convlen; ++k)
         eachTermDenCoefficients[j].push_back(0);
 
+      std::vector<Fraction> convolution;
+      convolution.reserve(convlen);
       convolution.clear();
       for (unsigned k = 0; k < convlen; ++k) {
         Fraction sum(0, 1);
@@ -438,7 +437,7 @@ mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
     totalTerm =
         totalTerm + getCoefficientInRationalFunction(r, numeratorCoefficients,
                                                      denominatorCoefficients) *
-                        QuasiPolynomial(num_params, sign);
+                        QuasiPolynomial(numParams, sign);
   }
 
   return totalTerm.simplify();

>From f78ef764817c10bbbf2375556dd1fb502d2686e6 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 23:02:15 +0530
Subject: [PATCH 07/14] Mark getters as const

---
 .../mlir/Analysis/Presburger/GeneratingFunction.h      | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h b/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
index eaf0449fe8a118..c38eab6efd0fc1 100644
--- a/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
+++ b/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
@@ -62,13 +62,15 @@ class GeneratingFunction {
 #endif // NDEBUG
   }
 
-  unsigned getNumParams() { return numParam; }
+  unsigned getNumParams() const { return numParam; }
 
-  SmallVector<int> getSigns() { return signs; }
+  SmallVector<int> getSigns() const { return signs; }
 
-  std::vector<ParamPoint> getNumerators() { return numerators; }
+  std::vector<ParamPoint> getNumerators() const { return numerators; }
 
-  std::vector<std::vector<Point>> getDenominators() { return denominators; }
+  std::vector<std::vector<Point>> getDenominators() const {
+    return denominators;
+  }
 
   GeneratingFunction operator+(GeneratingFunction &gf) const {
     assert(numParam == gf.getNumParams() &&

>From 06ca87a6de82a4d36ba54508bc979cc0eb5ef5de Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 23:02:41 +0530
Subject: [PATCH 08/14] Fix doc

---
 .../mlir/Analysis/Presburger/Barvinok.h       |  6 +-
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 95 +++++++++++--------
 2 files changed, 58 insertions(+), 43 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 2e2273dab4bc9d..d72332d220d2b5 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -99,9 +99,9 @@ QuasiPolynomial getCoefficientInRationalFunction(unsigned power,
                                                  ArrayRef<QuasiPolynomial> num,
                                                  ArrayRef<Fraction> den);
 
-/// Substitute the generating function with the unit vector
-/// to find the number of terms.
-QuasiPolynomial substituteWithUnitVector(GeneratingFunction);
+/// Find the number of terms in the generating function corresponding to
+/// a polytope.
+QuasiPolynomial computeNumTerms(const GeneratingFunction& gf);
 
 } // namespace detail
 } // namespace presburger
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 8b3a91acdc83a9..9eac53a96e9e0b 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -246,47 +246,62 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
   return coefficients[power].simplify();
 }
 
-// We have a generating function of the form
-// f_p(x) = \sum_i sign_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
-//
-// where sign_i is ±1,
-// n_i \in Q^p -> Q^d is a d-vector of affine functions on p parameters, and
-// d_{ij} \in Q^d are vectors.
-//
-// We need to find the number of terms of the form x^t in the expansion of
-// this function, for which we substitute x = (1, ..., 1).
-// However, direct substitution leads to an undefined answer due to the
-// form of the denominator.
-//
-// We therefore use the following procedure instead:
-// Find a vector μ that is not orthogonal to any of the d_{ij}.
-// Substitute x_i = (s+1)^μ_i. As μ_i is not orthogonal to d_{ij},
-// we never have (1 - (s+1)^0) = 0 in any of the terms in denominator.
-// We then find the constant term in this function, i.e., we evaluate it
-// at s = 0, which is equivalent to x = (1, ..., 1).
-//
-// Now, we have a function of the form
-// f_p(s) = \sum_i sign_i * (s+1)^n_i / (\prod_j (1 - (s+1)^d'_{ij}))
-// in which we need to find the constant term.
-// For the i'th term, we first convert all the d'_{ij} to their
-// absolute values by multiplying and dividing by (s+1)^(-d'_{ij}) if it is
-// negative. We change the sign accordingly.
-// Then, we replace each (1 - (s+1)^(d'_{ij})) with
-// (-s)(\sum_{0 ≤ k < d'_{ij}} (s+1)^k).
-// Thus the term overall has now the form
-// sign'_i * (s+1)^n'_i / (s^r * \prod_j (\sum_k (s+1)^k)).
-// This means that
-// the numerator is a polynomial in s, with coefficients as quasipolynomials,
-// and the denominator is polynomial in s, with fractional coefficients.
-// We need to find the constant term in the expansion of this term,
-// which is the same as finding the coefficient of s^r in
-// sign'_i * (s+1)^n'_i / (\prod_j (\sum_k (s+1)^k)),
-// for which we use the `getCoefficientInRationalFunction()` function.
-//
-// Verdoolaege, Sven, et al. "Counting integer points in parametric polytopes
-// using Barvinok's rational functions." Algorithmica 48 (2007): 37-66.
+/// We have a generating function of the form
+/// f_p(x) = \sum_i sign_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
+///
+/// where sign_i is ±1,
+/// n_i \in Q^p -> Q^d is a d-vector of affine functions on p parameters, and
+/// d_{ij} \in Q^d are vectors.
+///
+/// We need to find the number of terms of the form x^t in the expansion of
+/// this function, for which we substitute x = (1, ..., 1).
+/// However, direct substitution causes the denominator to become zero.
+///
+/// We therefore use the following procedure instead:
+/// 1. Substitute x_i = (s+1)^μ_i for some vector μ. This makes the generating
+/// a function of a scalar s.
+/// 2. Write each term in this function as P(s)/Q(s), where P and Q are
+/// polynomials. P has coefficients as quasipolynomials in d parameters, while
+/// Q has coefficients as scalars.
+/// 3. Find the constant term in the expansion of each term P(s)/Q(s). This is
+/// equivalent to substituting s = 0, which by step 1's substitution is
+/// equivalent to letting x = (1, ..., 1).
+/// In this step, we cancel the factor with root zero from the numerator and
+/// denominator, thus preventing the denominator from becoming zero.
+/// Step (1) We need to find a μ_i such that we can substitute x_i =
+/// (s+1)^μ_i. After this substitution, the exponent of (s+1) in the
+/// denominator is (μ_i • d_{ij}) in each term. Clearly, this cannot become
+/// zero. Thus we find a vector μ that is not orthogonal to any of the
+/// d_{ij} and substitute x accordingly.
+///
+/// Step (2) We need to express the terms in the function as quotients of
+/// polynomials. Each term is now of the form
+/// sign_i * (s+1)^n_i / (\prod_j (1 - (s+1)^d'_{ij}))
+/// For the i'th term, we first convert all the d'_{ij} to their
+/// absolute values by multiplying and dividing by (s+1)^(-d'_{ij}) if it is
+/// negative. We change the sign accordingly.
+/// Then, we replace each (1 - (s+1)^(d'_{ij})) with
+/// (-s)(\sum_{0 ≤ k < d'_{ij}} (s+1)^k).
+/// Thus the term overall has now the form
+/// sign'_i * (s+1)^n'_i / (s^r * \prod_j (\sum_k (s+1)^k)).
+/// This means that
+/// the numerator is a polynomial in s, with coefficients as
+/// quasipolynomials (given by binomial coefficients), and the denominator
+/// is polynomial in s, with fractional coefficients (given by taking the
+/// convolution over all j).
+///
+/// Step (3) We need to find the constant term in the expansion of each
+/// term. Since each term has s^r as a factor in the denominator, we avoid
+/// substituting s = 0 directly; instead, we find the coefficient of s^r in
+/// sign'_i * (s+1)^n'_i / (\prod_j (\sum_k (s+1)^k)),
+/// for which we use the `getCoefficientInRationalFunction()` function.
+///
+/// Verdoolaege, Sven, et al. "Counting integer points in parametric
+/// polytopes using Barvinok's rational functions." Algorithmica 48 (2007):
+/// 37-66.
+
 QuasiPolynomial
-mlir::presburger::detail::substituteWithUnitVector(GeneratingFunction gf) {
+mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
   std::vector<Point> allDenominators;
   for (ArrayRef<Point> den : gf.getDenominators())
     allDenominators.insert(allDenominators.end(), den.begin(), den.end());

>From 92a73fceb9b547917c4a2f0f6dcb68c1aab2794c Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 23:32:06 +0530
Subject: [PATCH 09/14] Abstract convolution

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 63 ++++++++++---------
 .../Analysis/Presburger/BarvinokTest.cpp      |  4 +-
 2 files changed, 34 insertions(+), 33 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 9eac53a96e9e0b..ca3edd68c82483 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -246,6 +246,28 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
   return coefficients[power].simplify();
 }
 
+static std::vector<Fraction> convolution(std::vector<Fraction> a,
+                                         std::vector<Fraction> b) {
+  // The length of the convolution is the maximum of the lengths
+  // of the two sequences. We pad the shorter one with zeroes.
+  unsigned convlen = std::max(a.size(), b.size());
+  for (unsigned k = a.size(); k < convlen; ++k)
+    a.push_back(0);
+  for (unsigned k = b.size(); k < convlen; ++k)
+    b.push_back(0);
+
+  std::vector<Fraction> convolution;
+  convolution.reserve(convlen);
+  convolution.clear();
+  for (unsigned k = 0; k < convlen; ++k) {
+    Fraction sum(0, 1);
+    for (unsigned l = 0; l <= k; ++l)
+      sum = sum + a[l] * b[k - l];
+    convolution.push_back(sum);
+  }
+  return convolution;
+}
+
 /// We have a generating function of the form
 /// f_p(x) = \sum_i sign_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
 ///
@@ -254,7 +276,7 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
 /// d_{ij} \in Q^d are vectors.
 ///
 /// We need to find the number of terms of the form x^t in the expansion of
-/// this function, for which we substitute x = (1, ..., 1).
+/// this function.
 /// However, direct substitution causes the denominator to become zero.
 ///
 /// We therefore use the following procedure instead:
@@ -264,10 +286,8 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
 /// polynomials. P has coefficients as quasipolynomials in d parameters, while
 /// Q has coefficients as scalars.
 /// 3. Find the constant term in the expansion of each term P(s)/Q(s). This is
-/// equivalent to substituting s = 0, which by step 1's substitution is
-/// equivalent to letting x = (1, ..., 1).
-/// In this step, we cancel the factor with root zero from the numerator and
-/// denominator, thus preventing the denominator from becoming zero.
+/// equivalent to substituting s = 0.
+///
 /// Step (1) We need to find a μ_i such that we can substitute x_i =
 /// (s+1)^μ_i. After this substitution, the exponent of (s+1) in the
 /// denominator is (μ_i • d_{ij}) in each term. Clearly, this cannot become
@@ -279,7 +299,8 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
 /// sign_i * (s+1)^n_i / (\prod_j (1 - (s+1)^d'_{ij}))
 /// For the i'th term, we first convert all the d'_{ij} to their
 /// absolute values by multiplying and dividing by (s+1)^(-d'_{ij}) if it is
-/// negative. We change the sign accordingly.
+/// negative. We change the sign accordingly to keep the denominator in the
+/// same form.
 /// Then, we replace each (1 - (s+1)^(d'_{ij})) with
 /// (-s)(\sum_{0 ≤ k < d'_{ij}} (s+1)^k).
 /// Thus the term overall has now the form
@@ -299,7 +320,6 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
 /// Verdoolaege, Sven, et al. "Counting integer points in parametric
 /// polytopes using Barvinok's rational functions." Algorithmica 48 (2007):
 /// 37-66.
-
 QuasiPolynomial
 mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
   std::vector<Point> allDenominators;
@@ -314,7 +334,6 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
   QuasiPolynomial totalTerm(numParams, 0);
   for (unsigned i = 0; i < numTerms; ++i) {
     int sign = gf.getSigns()[i];
-    ParamPoint v = gf.getNumerators()[i];
     std::vector<Point> ds = gf.getDenominators()[i];
 
     // Substitute x_i = (s+1)^μ_i
@@ -332,10 +351,11 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
 
     // and whose affine fn is a single floor expression, given by the
     // corresponding column of v.
+    ParamPoint vTranspose = gf.getNumerators()[i].transpose();
     std::vector<std::vector<SmallVector<Fraction>>> affine;
     affine.reserve(numDims);
     for (unsigned j = 0; j < numDims; ++j)
-      affine.push_back({SmallVector<Fraction>(v.transpose().getRow(j))});
+      affine.push_back({SmallVector<Fraction>(vTranspose.getRow(j))});
 
     QuasiPolynomial num(numParams, coefficients, affine);
     num = num.simplify();
@@ -426,28 +446,9 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     // of all the terms.
     std::vector<Fraction> denominatorCoefficients;
     denominatorCoefficients = eachTermDenCoefficients[0];
-    for (unsigned j = 1, e = eachTermDenCoefficients.size(); j < e; ++j) {
-      // The length of the convolution is the maximum of the lengths
-      // of the two sequences. We pad the shorter one with zeroes.
-      unsigned convlen = std::max(denominatorCoefficients.size(),
-                                  eachTermDenCoefficients[j].size());
-      for (unsigned k = denominatorCoefficients.size(); k < convlen; ++k)
-        denominatorCoefficients.push_back(0);
-      for (unsigned k = eachTermDenCoefficients[j].size(); k < convlen; ++k)
-        eachTermDenCoefficients[j].push_back(0);
-
-      std::vector<Fraction> convolution;
-      convolution.reserve(convlen);
-      convolution.clear();
-      for (unsigned k = 0; k < convlen; ++k) {
-        Fraction sum(0, 1);
-        for (unsigned l = 0; l <= k; ++l)
-          sum = sum +
-                denominatorCoefficients[l] * eachTermDenCoefficients[j][k - l];
-        convolution.push_back(sum);
-      }
-      denominatorCoefficients = convolution;
-    }
+    for (unsigned j = 1, e = eachTermDenCoefficients.size(); j < e; ++j)
+      denominatorCoefficients =
+          convolution(denominatorCoefficients, eachTermDenCoefficients[j]);
 
     totalTerm =
         totalTerm + getCoefficientInRationalFunction(r, numeratorCoefficients,
diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index 2842fc14eb8298..82af731b54f71a 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -127,7 +127,7 @@ TEST(BarvinokTest, getCoefficientInRationalFunction) {
 
 // The following test is taken from
 //
-TEST(BarvinokTest, substituteWithUnitVector) {
+TEST(BarvinokTest, computeNumTerms) {
   GeneratingFunction gf(
       1, {1, 1, 1},
       {makeFracMatrix(2, 2, {{0, Fraction(1, 2)}, {0, 0}}),
@@ -135,7 +135,7 @@ TEST(BarvinokTest, substituteWithUnitVector) {
        makeFracMatrix(2, 2, {{0, 0}, {0, 0}})},
       {{{-1, 1}, {-1, 0}}, {{1, -1}, {0, -1}}, {{1, 0}, {0, 1}}});
 
-  QuasiPolynomial numPoints = substituteWithUnitVector(gf);
+  QuasiPolynomial numPoints = computeNumTerms(gf);
 
   // First, we make sure that all the affine functions are of the form ⌊p/2⌋.
   for (const std::vector<SmallVector<Fraction>> &term : numPoints.getAffine()) {

>From aad6ca25ac2ac56091dce748a0e980841397b7a4 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 23:38:25 +0530
Subject: [PATCH 10/14] Abstract out mu substitution

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp | 59 +++++++++++++----------
 1 file changed, 34 insertions(+), 25 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index ca3edd68c82483..92e942bc42430c 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -268,6 +268,38 @@ static std::vector<Fraction> convolution(std::vector<Fraction> a,
   return convolution;
 }
 
+/// Substitute x_i = (s+1)^μ_i in one term of a generating function,
+/// returning a quasipolynomial which represents the exponent of the
+/// numerator of the result.
+QuasiPolynomial substituteMuInTerm(unsigned numParams, ParamPoint v,
+                                   std::vector<Point> ds, Point mu) {
+  unsigned numDims = mu.size();
+  // First, the exponent in the numerator becomes
+  // - (μ • u_1) * (floor(first col of v))
+  // - (μ • u_2) * (floor(second col of v)) - ...
+  // - (μ • u_d) * (floor(d'th col of v))
+  // So we store the negation of the  dot produts.
+
+  // We have d terms, each of whose coefficient is the negative dot product,
+  SmallVector<Fraction> coefficients;
+  coefficients.reserve(numDims);
+  for (const Point &d : ds)
+    coefficients.push_back(-dotProduct(mu, d));
+
+  // Then, the affine fn is a single floor expression, given by the
+  // corresponding column of v.
+  ParamPoint vTranspose = v.transpose();
+  std::vector<std::vector<SmallVector<Fraction>>> affine;
+  affine.reserve(numDims);
+  for (unsigned j = 0; j < numDims; ++j)
+    affine.push_back({SmallVector<Fraction>(vTranspose.getRow(j))});
+
+  QuasiPolynomial num(numParams, coefficients, affine);
+  num = num.simplify();
+
+  return num;
+}
+
 /// We have a generating function of the form
 /// f_p(x) = \sum_i sign_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
 ///
@@ -328,7 +360,6 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
   Point mu = getNonOrthogonalVector(allDenominators);
 
   unsigned numParams = gf.getNumParams();
-  unsigned numDims = mu.size();
   unsigned numTerms = gf.getDenominators().size();
 
   QuasiPolynomial totalTerm(numParams, 0);
@@ -336,30 +367,8 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     int sign = gf.getSigns()[i];
     std::vector<Point> ds = gf.getDenominators()[i];
 
-    // Substitute x_i = (s+1)^μ_i
-    // Then the exponent in the numerator becomes
-    // - (μ • u_1) * (floor(first col of v))
-    // - (μ • u_2) * (floor(second col of v)) - ...
-    // - (μ • u_d) * (floor(d'th col of v))
-    // So we store the negation of the  dot produts.
-
-    // We have d terms, each of whose coefficient is the negative dot product,
-    SmallVector<Fraction> coefficients;
-    coefficients.reserve(numDims);
-    for (const Point &d : ds)
-      coefficients.push_back(-dotProduct(mu, d));
-
-    // and whose affine fn is a single floor expression, given by the
-    // corresponding column of v.
-    ParamPoint vTranspose = gf.getNumerators()[i].transpose();
-    std::vector<std::vector<SmallVector<Fraction>>> affine;
-    affine.reserve(numDims);
-    for (unsigned j = 0; j < numDims; ++j)
-      affine.push_back({SmallVector<Fraction>(vTranspose.getRow(j))});
-
-    QuasiPolynomial num(numParams, coefficients, affine);
-    num = num.simplify();
-
+    QuasiPolynomial num =
+        substituteMuInTerm(numParams, gf.getNumerators()[i], ds, mu);
     // Now the numerator is (s+1)^num.
 
     std::vector<Fraction> dens;

>From 520576574b29719a6104ed2f1c6039dd32e05727 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 23:54:20 +0530
Subject: [PATCH 11/14] Fix doc

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 92e942bc42430c..de78ab8e2493be 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -304,7 +304,8 @@ QuasiPolynomial substituteMuInTerm(unsigned numParams, ParamPoint v,
 /// f_p(x) = \sum_i sign_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
 ///
 /// where sign_i is ±1,
-/// n_i \in Q^p -> Q^d is a d-vector of affine functions on p parameters, and
+/// n_i \in Q^p -> Q^d is the sum of the vectors d_{ij}, weighted by the
+/// floors of d affine functions.
 /// d_{ij} \in Q^d are vectors.
 ///
 /// We need to find the number of terms of the form x^t in the expansion of

>From fc0dba7ffe164b8c03ac279b4affb32604c7e849 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 14 Jan 2024 23:58:22 +0530
Subject: [PATCH 12/14] Fix doc

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

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index d72332d220d2b5..50d6742019161c 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -101,7 +101,7 @@ QuasiPolynomial getCoefficientInRationalFunction(unsigned power,
 
 /// Find the number of terms in the generating function corresponding to
 /// a polytope.
-QuasiPolynomial computeNumTerms(const GeneratingFunction& gf);
+QuasiPolynomial computeNumTerms(const GeneratingFunction &gf);
 
 } // namespace detail
 } // namespace presburger
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index de78ab8e2493be..809dfdca6376d9 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -314,7 +314,7 @@ QuasiPolynomial substituteMuInTerm(unsigned numParams, ParamPoint v,
 ///
 /// We therefore use the following procedure instead:
 /// 1. Substitute x_i = (s+1)^μ_i for some vector μ. This makes the generating
-/// a function of a scalar s.
+/// function a function of a scalar s.
 /// 2. Write each term in this function as P(s)/Q(s), where P and Q are
 /// polynomials. P has coefficients as quasipolynomials in d parameters, while
 /// Q has coefficients as scalars.

>From d3e758e8a4fc261db7e5bf8b9007ed6f44430fd5 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 00:01:31 +0530
Subject: [PATCH 13/14] Fix doc

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 809dfdca6376d9..497987676d4d3b 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -324,12 +324,12 @@ QuasiPolynomial substituteMuInTerm(unsigned numParams, ParamPoint v,
 /// Step (1) We need to find a μ_i such that we can substitute x_i =
 /// (s+1)^μ_i. After this substitution, the exponent of (s+1) in the
 /// denominator is (μ_i • d_{ij}) in each term. Clearly, this cannot become
-/// zero. Thus we find a vector μ that is not orthogonal to any of the
+/// zero. Hence we find a vector μ that is not orthogonal to any of the
 /// d_{ij} and substitute x accordingly.
 ///
 /// Step (2) We need to express the terms in the function as quotients of
 /// polynomials. Each term is now of the form
-/// sign_i * (s+1)^n_i / (\prod_j (1 - (s+1)^d'_{ij}))
+/// sign_i * (s+1)^n'_i / (\prod_j (1 - (s+1)^d'_{ij}))
 /// For the i'th term, we first convert all the d'_{ij} to their
 /// absolute values by multiplying and dividing by (s+1)^(-d'_{ij}) if it is
 /// negative. We change the sign accordingly to keep the denominator in the

>From ac6b0388aa8d43c7b2426fc023a05f1efb64dcf3 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 00:08:00 +0530
Subject: [PATCH 14/14] Fix doc

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 497987676d4d3b..cde2ce3dcd91c2 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -330,7 +330,8 @@ QuasiPolynomial substituteMuInTerm(unsigned numParams, ParamPoint v,
 /// Step (2) We need to express the terms in the function as quotients of
 /// polynomials. Each term is now of the form
 /// sign_i * (s+1)^n'_i / (\prod_j (1 - (s+1)^d'_{ij}))
-/// For the i'th term, we first convert all the d'_{ij} to their
+/// For the i'th term, we first normalize the denominator to have only
+/// positive exponents. We convert all the d'_{ij} to their
 /// absolute values by multiplying and dividing by (s+1)^(-d'_{ij}) if it is
 /// negative. We change the sign accordingly to keep the denominator in the
 /// same form.
@@ -366,7 +367,7 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
   QuasiPolynomial totalTerm(numParams, 0);
   for (unsigned i = 0; i < numTerms; ++i) {
     int sign = gf.getSigns()[i];
-    std::vector<Point> ds = gf.getDenominators()[i];
+    const std::vector<Point> &ds = gf.getDenominators()[i];
 
     QuasiPolynomial num =
         substituteMuInTerm(numParams, gf.getNumerators()[i], ds, mu);
@@ -382,8 +383,7 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     // (1 - (s+1)^dens.back())
 
     // We track the number of exponents that are negative in the
-    // denominator, and convert them to their absolute values
-    // (see lines 356-66).
+    // denominator, and convert them to their absolute values.
     unsigned numNegExps = 0;
     Fraction sumNegExps(0, 1);
     for (unsigned j = 0, e = dens.size(); j < e; ++j) {
@@ -410,7 +410,7 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
       sign = -sign;
     num = num - QuasiPolynomial(numParams, sumNegExps);
 
-    // Take all the (-s) out, from line 353.
+    // Take all the (-s) out.
     unsigned r = dens.size();
     if (r % 2 == 1)
       sign = -sign;



More information about the Mlir-commits mailing list