[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 12:05:25 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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;

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

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index cde2ce3dcd91c2..0419e590d3481f 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -367,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];
-    const std::vector<Point> &ds = gf.getDenominators()[i];
+    std::vector<Point> ds = gf.getDenominators()[i];
 
     QuasiPolynomial num =
         substituteMuInTerm(numParams, gf.getNumerators()[i], ds, mu);

>From fd7879a583f6fc1f4749ac30cf999fa27dca796e Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 00:42:32 +0530
Subject: [PATCH 16/23] Reorg

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp | 120 ++++++++++++----------
 1 file changed, 65 insertions(+), 55 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 0419e590d3481f..7fdb5d44c0c6b7 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -269,16 +269,20 @@ static std::vector<Fraction> convolution(std::vector<Fraction> a,
 }
 
 /// 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) {
+/// returning
+/// a quasipolynomial which represents the exponent of the numerator
+/// of the result, and
+/// a vector which represents the exponents of the denominator of the
+/// result.
+std::pair<QuasiPolynomial, std::vector<Fraction>>
+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.
+  // So we store the negation of the dot products.
 
   // We have d terms, each of whose coefficient is the negative dot product,
   SmallVector<Fraction> coefficients;
@@ -297,7 +301,52 @@ QuasiPolynomial substituteMuInTerm(unsigned numParams, ParamPoint v,
   QuasiPolynomial num(numParams, coefficients, affine);
   num = num.simplify();
 
-  return num;
+  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)
+    dens.push_back(dotProduct(d, mu));
+  // This term in the denominator is
+  // (1 - (s+1)^dens.back())
+
+  return std::make_pair(num, dens);
+}
+
+void normalizeDenominatorExponents(int &sign, QuasiPolynomial &num,
+                                   std::vector<Fraction> &dens) {
+  // We track the number of exponents that are negative in the
+  // 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) {
+    if (dens[j] < 0) {
+      numNegExps += 1;
+      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.getNumInputs(), sumNegExps);
+
+  // Take all the (-s) out.
+  unsigned r = dens.size();
+  if (r % 2 == 1)
+    sign = -sign;
 }
 
 /// We have a generating function of the form
@@ -362,58 +411,18 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
   Point mu = getNonOrthogonalVector(allDenominators);
 
   unsigned numParams = gf.getNumParams();
-  unsigned numTerms = gf.getDenominators().size();
 
+  const std::vector<std::vector<Point>> &ds = gf.getDenominators();
   QuasiPolynomial totalTerm(numParams, 0);
-  for (unsigned i = 0; i < numTerms; ++i) {
+  for (unsigned i = 0, e = ds.size(); i < e; ++i) {
     int sign = gf.getSigns()[i];
-    std::vector<Point> ds = gf.getDenominators()[i];
-
-    QuasiPolynomial num =
-        substituteMuInTerm(numParams, gf.getNumerators()[i], ds, mu);
-    // Now the numerator is (s+1)^num.
-
-    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)
-      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.
-    unsigned numNegExps = 0;
-    Fraction sumNegExps(0, 1);
-    for (unsigned j = 0, e = dens.size(); j < e; ++j) {
-      if (dens[j] < 0) {
-        numNegExps += 1;
-        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(numParams, sumNegExps);
-
-    // Take all the (-s) out.
-    unsigned r = dens.size();
-    if (r % 2 == 1)
-      sign = -sign;
+    auto [numExp, dens] =
+        substituteMuInTerm(numParams, gf.getNumerators()[i], ds[i], mu);
+    // Now the numerator is (s+1)^numExp
+    // and the denominator is \prod_j (1 - (s+1)^dens[j]).
+
+    normalizeDenominatorExponents(sign, numExp, dens);
 
     // Now the expression is
     // (s+1)^num /
@@ -423,6 +432,7 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     // we need to find the coefficient of s^r in
     // P(s)/Q(s).
 
+    unsigned r = dens.size();
     // First, the coefficients of P(s), which are binomial coefficients.
     // We need r+1 of these.
     std::vector<QuasiPolynomial> numeratorCoefficients;
@@ -432,7 +442,7 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     for (unsigned j = 1; j <= r; ++j)
       numeratorCoefficients.push_back(
           (numeratorCoefficients[j - 1] *
-           (num - QuasiPolynomial(numParams, j - 1)) / Fraction(j, 1))
+           (numExp - QuasiPolynomial(numParams, j - 1)) / Fraction(j, 1))
               .simplify());
     // Coeff of s^j
 

>From 819ec57a60e99ced09624eb5d652c63763cb5d3f Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 00:43:09 +0530
Subject: [PATCH 17/23] Typo

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 7fdb5d44c0c6b7..85bdfda4a79abd 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -290,7 +290,7 @@ substituteMuInTerm(unsigned numParams, ParamPoint v, std::vector<Point> ds,
   for (const Point &d : ds)
     coefficients.push_back(-dotProduct(mu, d));
 
-  // Then, the affine fn is a single floor expression, given by the
+  // Then, the affine function is a single floor expression, given by the
   // corresponding column of v.
   ParamPoint vTranspose = v.transpose();
   std::vector<std::vector<SmallVector<Fraction>>> affine;

>From eeb020c739eb5aff0f77bbc39985bbd6003ef54c Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 00:54:33 +0530
Subject: [PATCH 18/23] Reorg doc

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp | 74 +++++++++++------------
 1 file changed, 36 insertions(+), 38 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 85bdfda4a79abd..a1cfdd8114c13a 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -313,6 +313,9 @@ substituteMuInTerm(unsigned numParams, ParamPoint v, std::vector<Point> ds,
   return std::make_pair(num, dens);
 }
 
+/// Normalize all denominator exponents `dens` to their absolute values
+/// by multiplying and dividing by the inverses.
+/// Also, take the factors (-s) out of each term of the product.
 void normalizeDenominatorExponents(int &sign, QuasiPolynomial &num,
                                    std::vector<Fraction> &dens) {
   // We track the number of exponents that are negative in the
@@ -370,41 +373,16 @@ void normalizeDenominatorExponents(int &sign, QuasiPolynomial &num,
 /// 3. Find the constant term in the expansion of each term P(s)/Q(s). This is
 /// 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
-/// 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}))
-/// 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.
-/// 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::computeNumTerms(const GeneratingFunction &gf) {
+  // Step (1) We need to find a μ 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. Hence we find a vector μ that is not orthogonal to any of the
+  // d_{ij} and substitute x accordingly.
   std::vector<Point> allDenominators;
   for (ArrayRef<Point> den : gf.getDenominators())
     allDenominators.insert(allDenominators.end(), den.begin(), den.end());
@@ -417,22 +395,42 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
   for (unsigned i = 0, e = ds.size(); i < e; ++i) {
     int sign = gf.getSigns()[i];
 
+    // Compute the new numerator and denominator after substituting μ.
     auto [numExp, dens] =
         substituteMuInTerm(numParams, gf.getNumerators()[i], ds[i], mu);
     // Now the numerator is (s+1)^numExp
     // and the denominator is \prod_j (1 - (s+1)^dens[j]).
 
+    // 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)^numExp / (\prod_j (1 - (s+1)^dens[j]))
+    // For the i'th term, we first normalize the denominator to have only
+    // positive exponents. We convert all the dens[j] to their
+    // absolute values by multiplying and dividing by (s+1)^(-dens[j]) if it is
+    // negative. We change the sign accordingly to keep the denominator in the
+    // same form.
+    // Then, we replace each (1 - (s+1)^(dens[j])) with
+    // (-s)(\sum_{0 ≤ k < dens[j]} (s+1)^k).
     normalizeDenominatorExponents(sign, numExp, dens);
 
-    // 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).
-
+    // Thus the term overall has now the form
+    // sign'_i * (s+1)^numExp /
+    // (s^r * \prod_j (\sum_{0 ≤ k < dens[j]} (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)^numExp / (\prod_j (\sum_k (s+1)^k)),
+    // Letting P(s) = (s+1)^numExp and Q(s) = \prod_j (...),
+    // we need to find the coefficient of s^r in P(s)/Q(s),
+    // for which we use the `getCoefficientInRationalFunction()` function.
     unsigned r = dens.size();
+
     // First, the coefficients of P(s), which are binomial coefficients.
     // We need r+1 of these.
     std::vector<QuasiPolynomial> numeratorCoefficients;

>From 0bda7729d154a2f4232f9926d4905b5007bdbd96 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 01:24:05 +0530
Subject: [PATCH 19/23] Abstract out binomial coefficient computation

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp | 61 +++++++++++++++--------
 1 file changed, 40 insertions(+), 21 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index a1cfdd8114c13a..cf42b9494c8bc1 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -352,6 +352,34 @@ void normalizeDenominatorExponents(int &sign, QuasiPolynomial &num,
     sign = -sign;
 }
 
+/// Compute the binomial coefficients nCi for 0 ≤ i ≤ r,
+/// where n is a QuasiPolynomial.
+std::vector<QuasiPolynomial> getBinomialCoefficients(QuasiPolynomial n,
+                                                     unsigned r) {
+  unsigned numParams = n.getNumInputs();
+  std::vector<QuasiPolynomial> coefficients;
+  coefficients.reserve(r + 1);
+  coefficients.push_back(QuasiPolynomial(numParams, 1));
+  for (unsigned j = 1; j <= r; ++j)
+    // We use the recursive formula for binomial coefficients here and below.
+    coefficients.push_back(
+        (coefficients[j - 1] * (n - QuasiPolynomial(numParams, j - 1)) /
+         Fraction(j, 1))
+            .simplify());
+  return coefficients;
+}
+
+/// Compute the binomial coefficients nCi for 0 ≤ i ≤ r,
+/// where n is a QuasiPolynomial.
+std::vector<Fraction> getBinomialCoefficients(Fraction n, Fraction r) {
+  std::vector<Fraction> coefficients;
+  coefficients.reserve((int64_t)floor(r));
+  coefficients.push_back(1);
+  for (unsigned j = 1; j <= r; ++j)
+    coefficients.push_back(coefficients[j - 1] * (n - (j - 1)) / (j));
+  return coefficients;
+}
+
 /// 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})
 ///
@@ -409,7 +437,8 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     // absolute values by multiplying and dividing by (s+1)^(-dens[j]) if it is
     // negative. We change the sign accordingly to keep the denominator in the
     // same form.
-    // Then, we replace each (1 - (s+1)^(dens[j])) with
+    // Then, using the formula for geometric series, we replace each (1 -
+    // (s+1)^(dens[j])) with
     // (-s)(\sum_{0 ≤ k < dens[j]} (s+1)^k).
     normalizeDenominatorExponents(sign, numExp, dens);
 
@@ -431,32 +460,22 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     // for which we use the `getCoefficientInRationalFunction()` function.
     unsigned r = dens.size();
 
-    // First, the coefficients of P(s), which are binomial coefficients.
-    // We need r+1 of these.
-    std::vector<QuasiPolynomial> numeratorCoefficients;
-    numeratorCoefficients.reserve(r + 1);
-    numeratorCoefficients.push_back(
-        QuasiPolynomial(numParams, 1)); // Coeff of s^0
-    for (unsigned j = 1; j <= r; ++j)
-      numeratorCoefficients.push_back(
-          (numeratorCoefficients[j - 1] *
-           (numExp - QuasiPolynomial(numParams, j - 1)) / Fraction(j, 1))
-              .simplify());
-    // Coeff of s^j
+    // First, we compute the coefficients of P(s),
+    // which are binomial coefficients.
+    // We only need the first r+1 of these, as higher-order terms do not
+    // contribute to the coefficient of s^r.
+    std::vector<QuasiPolynomial> numeratorCoefficients =
+        getBinomialCoefficients(numExp, r);
 
     // Then the coefficients of each individual term in Q(s),
-    // which are (di+1) C (k+1) for 0 ≤ k ≤ di
+    // which are (di+1) C (k+1) for 0 ≤ k ≤ di.
     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);
-      for (unsigned j = 1; j <= den; ++j)
-        singleTermDenCoefficients.push_back(singleTermDenCoefficients[j - 1] *
-                                            (den - (j - 1)) / (j + 1));
-
-      eachTermDenCoefficients.push_back(singleTermDenCoefficients);
+      singleTermDenCoefficients = getBinomialCoefficients(den + 1, den + 1);
+      eachTermDenCoefficients.push_back(
+          ArrayRef<Fraction>(singleTermDenCoefficients).slice(1));
     }
 
     // Now we find the coefficients in Q(s) itself

>From e811dfda57ef3e389d9e869a5425c240cb2a8600 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 01:29:23 +0530
Subject: [PATCH 20/23] Fix conv

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index cf42b9494c8bc1..f8b3bd6226ed19 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -250,19 +250,16 @@ 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);
+  unsigned len = std::max(a.size(), b.size());
+  a.resize(len);
+  b.resize(len);
 
   std::vector<Fraction> convolution;
-  convolution.reserve(convlen);
-  convolution.clear();
-  for (unsigned k = 0; k < convlen; ++k) {
+  convolution.reserve(len);
+  for (unsigned k = 0; k < len; ++k) {
     Fraction sum(0, 1);
     for (unsigned l = 0; l <= k; ++l)
-      sum = sum + a[l] * b[k - l];
+      sum += a[l] * b[k - l];
     convolution.push_back(sum);
   }
   return convolution;

>From 3b40fb5243694f162b9ca20e7c715a1bf55b8af7 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 01:31:47 +0530
Subject: [PATCH 21/23] Fix conv

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index f8b3bd6226ed19..cbc5d30a29dab1 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -250,7 +250,7 @@ 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 len = std::max(a.size(), b.size());
+  unsigned len = a.size() + b.size();
   a.resize(len);
   b.resize(len);
 

>From 0ef21593f729cbce801d3682ae1ce710bd47501f Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 01:33:50 +0530
Subject: [PATCH 22/23] Fix doc

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index cbc5d30a29dab1..f9cda3f34c9291 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -271,6 +271,9 @@ static std::vector<Fraction> convolution(std::vector<Fraction> a,
 /// of the result, and
 /// a vector which represents the exponents of the denominator of the
 /// result.
+/// v represents the affine functions whose floors are multiplied by
+/// by the generators, and
+/// ds represents the list of generators.
 std::pair<QuasiPolynomial, std::vector<Fraction>>
 substituteMuInTerm(unsigned numParams, ParamPoint v, std::vector<Point> ds,
                    Point mu) {
@@ -302,12 +305,13 @@ substituteMuInTerm(unsigned numParams, ParamPoint v, std::vector<Point> ds,
   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)
+  for (const Point &d : ds) {
+    // This term in the denominator is
+    // (1 - (s+1)^dens.back())
     dens.push_back(dotProduct(d, mu));
-  // This term in the denominator is
-  // (1 - (s+1)^dens.back())
+  }
 
-  return std::make_pair(num, dens);
+  return {num, dens};
 }
 
 /// Normalize all denominator exponents `dens` to their absolute values

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

---
 mlir/unittests/Analysis/Presburger/BarvinokTest.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index 82af731b54f71a..b0948334740fe6 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -126,7 +126,9 @@ TEST(BarvinokTest, getCoefficientInRationalFunction) {
 }
 
 // The following test is taken from
-//
+/// Verdoolaege, Sven, et al. "Counting integer points in parametric
+/// polytopes using Barvinok's rational functions." Algorithmica 48 (2007):
+/// 37-66.
 TEST(BarvinokTest, computeNumTerms) {
   GeneratingFunction gf(
       1, {1, 1, 1},



More information about the Mlir-commits mailing list