[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 21 10:44:41 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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/44] 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},

>From 98ae71e2649a31ba619f773b7b81a2e148c65a5e Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 15 Jan 2024 01:36:28 +0530
Subject: [PATCH 24/44] Fix doc

---
 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 f9cda3f34c9291..969c327ab60fa8 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -469,7 +469,7 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
         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 (dens[i]+1) C (k+1) for 0 ≤ k ≤ dens[i].
     std::vector<std::vector<Fraction>> eachTermDenCoefficients;
     std::vector<Fraction> singleTermDenCoefficients;
     eachTermDenCoefficients.reserve(r);

>From 3ab2f4894098918a730a3fe7704a6c33419a1fd6 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 17 Jan 2024 02:55:11 +0530
Subject: [PATCH 25/44] Improve simplify()

---
 .../Analysis/Presburger/QuasiPolynomial.h     |  3 +-
 .../Analysis/Presburger/QuasiPolynomial.cpp   | 30 +++++++++++++++++--
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index d03446f50264fa..c5db820a20c5c8 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -59,7 +59,8 @@ class QuasiPolynomial : public PresburgerSpace {
   QuasiPolynomial operator*(const QuasiPolynomial &x) const;
   QuasiPolynomial operator/(const Fraction x) const;
 
-  // Removes terms which evaluate to zero from the expression.
+  // Removes terms which evaluate to zero and affine functions
+  // which are constant from the expression.
   QuasiPolynomial simplify();
 
   Fraction getConstantTerm();
diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index 385d4c354be18a..af0a2e56bedb8c 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -97,10 +97,18 @@ QuasiPolynomial QuasiPolynomial::operator/(const Fraction x) const {
   return qp;
 }
 
-// Removes terms which evaluate to zero from the expression.
+// Removes terms which evaluate to zero from the expression and
+// integrate affine functions which are constants into the
+// coefficients.
 QuasiPolynomial QuasiPolynomial::simplify() {
+  Fraction newCoeff = 0;
   SmallVector<Fraction> newCoeffs({});
+
+  std::vector<SmallVector<Fraction>> newAffineTerm({});
   std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
+
+  unsigned numParam = getNumInputs();
+
   for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
     // A term is zero if its coefficient is zero, or
     if (coefficients[i] == Fraction(0, 1))
@@ -114,8 +122,24 @@ QuasiPolynomial QuasiPolynomial::simplify() {
         });
     if (product_is_zero)
       continue;
-    newCoeffs.push_back(coefficients[i]);
-    newAffine.push_back(affine[i]);
+
+    // Now, we know the term is nonzero.
+
+    // We now eliminate the affine functions which are constant
+    // by integrating them into the coefficients.
+    newAffineTerm = {};
+    newCoeff = coefficients[i];
+    for (ArrayRef<Fraction> term : affine[i]) {
+      bool allCoeffsZero = llvm::all_of(
+          term.slice(0, numParam), [](const Fraction c) { return c == 0; });
+      if (allCoeffsZero)
+        newCoeff *= term[numParam];
+      else
+        newAffineTerm.push_back(SmallVector<Fraction>(term));
+    }
+
+    newCoeffs.push_back(newCoeff);
+    newAffine.push_back(newAffineTerm);
   }
   return QuasiPolynomial(getNumInputs(), newCoeffs, newAffine);
 }

>From 8b2c962137f4e2d2a3230e62b9203f85cc443443 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 17 Jan 2024 03:50:08 +0530
Subject: [PATCH 26/44] Implement collectTerms()

---
 .../Analysis/Presburger/QuasiPolynomial.h     |  3 +++
 .../Analysis/Presburger/QuasiPolynomial.cpp   | 21 +++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index c5db820a20c5c8..e5b4aae36db31e 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -63,6 +63,9 @@ class QuasiPolynomial : public PresburgerSpace {
   // which are constant from the expression.
   QuasiPolynomial simplify();
 
+  // Group together like terms in the expression.
+  QuasiPolynomial collectTerms();
+
   Fraction getConstantTerm();
 
 private:
diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index af0a2e56bedb8c..36f37771a63dfc 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -144,6 +144,27 @@ QuasiPolynomial QuasiPolynomial::simplify() {
   return QuasiPolynomial(getNumInputs(), newCoeffs, newAffine);
 }
 
+QuasiPolynomial QuasiPolynomial::collectTerms() {
+  SmallVector<Fraction> newCoeffs({});
+  std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
+
+  for (unsigned i = 0, e = affine.size(); i < e; i++) {
+    bool alreadyPresent = false;
+    for (unsigned j = 0, f = newAffine.size(); j < f; j++) {
+      if (affine[i] == newAffine[j]) {
+        newCoeffs[j] += coefficients[i];
+        alreadyPresent = true;
+      }
+    }
+    if (alreadyPresent)
+      continue;
+    newCoeffs.push_back(coefficients[i]);
+    newAffine.push_back(affine[i]);
+  }
+
+  return QuasiPolynomial(getNumInputs(), newCoeffs, newAffine);
+}
+
 Fraction QuasiPolynomial::getConstantTerm() {
   Fraction constTerm = 0;
   for (unsigned i = 0, e = coefficients.size(); i < e; ++i)

>From aa10c1f161c67731e2341eb38dd17cd37c029253 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 17 Jan 2024 03:50:49 +0530
Subject: [PATCH 27/44] Add test for computeNumTerms()

---
 .../Analysis/Presburger/BarvinokTest.cpp      | 104 +++++++++++++++++-
 1 file changed, 99 insertions(+), 5 deletions(-)

diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index b0948334740fe6..e60ff9432ce865 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -125,11 +125,13 @@ TEST(BarvinokTest, getCoefficientInRationalFunction) {
   EXPECT_EQ(coeff.getConstantTerm(), Fraction(55, 64));
 }
 
-// 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) {
+  // 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.
+  // It represents a right-angled triangle with right angle at the origin,
+  // with height and base lengths (p/2).
   GeneratingFunction gf(
       1, {1, 1, 1},
       {makeFracMatrix(2, 2, {{0, Fraction(1, 2)}, {0, 0}}),
@@ -137,7 +139,7 @@ TEST(BarvinokTest, computeNumTerms) {
        makeFracMatrix(2, 2, {{0, 0}, {0, 0}})},
       {{{-1, 1}, {-1, 0}}, {{1, -1}, {0, -1}}, {{1, 0}, {0, 1}}});
 
-  QuasiPolynomial numPoints = computeNumTerms(gf);
+  QuasiPolynomial numPoints = computeNumTerms(gf).collectTerms();
 
   // First, we make sure that all the affine functions are of the form ⌊p/2⌋.
   for (const std::vector<SmallVector<Fraction>> &term : numPoints.getAffine()) {
@@ -165,4 +167,96 @@ TEST(BarvinokTest, computeNumTerms) {
   EXPECT_EQ(pSquaredCoeff, Fraction(1, 2));
   EXPECT_EQ(pCoeff, Fraction(3, 2));
   EXPECT_EQ(numPoints.getConstantTerm(), Fraction(1, 1));
+
+  // The following generating function corresponds to a cuboid
+  // with length (x-axis) M, width (y-axis) N, and height (z-axis)
+  // P.
+  // There are eight terms.
+  gf = GeneratingFunction(
+      3, {1, 1, 1, 1, 1, 1, 1, 1},
+      {makeFracMatrix(4, 3, {{0, 0, 0}, {0, 0, 0}, {0, 0, 0}, {0, 0, 0}}),
+       makeFracMatrix(4, 3, {{1, 0, 0}, {0, 0, 0}, {0, 0, 0}, {0, 0, 0}}),
+       makeFracMatrix(4, 3, {{0, 0, 0}, {0, 1, 0}, {0, 0, 0}, {0, 0, 0}}),
+       makeFracMatrix(4, 3, {{0, 0, 0}, {0, 0, 0}, {0, 0, 1}, {0, 0, 0}}),
+       makeFracMatrix(4, 3, {{1, 0, 0}, {0, 1, 0}, {0, 0, 0}, {0, 0, 0}}),
+       makeFracMatrix(4, 3, {{1, 0, 0}, {0, 0, 0}, {0, 0, 1}, {0, 0, 0}}),
+       makeFracMatrix(4, 3, {{0, 0, 0}, {0, 1, 0}, {0, 0, 1}, {0, 0, 0}}),
+       makeFracMatrix(4, 3, {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}, {0, 0, 0}})},
+      {{{1, 0, 0}, {0, 1, 0}, {0, 0, 1}},
+       {{-1, 0, 0}, {0, 1, 0}, {0, 0, 1}},
+       {{1, 0, 0}, {0, -1, 0}, {0, 0, 1}},
+       {{1, 0, 0}, {0, 1, 0}, {0, 0, -1}},
+       {{-1, 0, 0}, {0, -1, 0}, {0, 0, 1}},
+       {{-1, 0, 0}, {0, 1, 0}, {0, 0, -1}},
+       {{1, 0, 0}, {0, -1, 0}, {0, 0, -1}},
+       {{-1, 0, 0}, {0, -1, 0}, {0, 0, -1}}});
+
+  numPoints = computeNumTerms(gf);
+  numPoints = numPoints.collectTerms().simplify();
+
+  // First, we make sure all the affine functions are either
+  // ⌊M⌋, ⌊N⌋, or ⌊P⌋.
+  for (const std::vector<SmallVector<Fraction>> &term : numPoints.getAffine())
+    for (const SmallVector<Fraction> &aff : term)
+      EXPECT_EQ(aff[0] + aff[1] + aff[2] + aff[3], 1);
+
+  Fraction m = 0, n = 0, p = 0, mn = 0, np = 0, pm = 0, mnp = 0;
+
+  for (unsigned i = 0, e = numPoints.getAffine().size(); i < e; i++) {
+    if (numPoints.getAffine()[i].size() == 1u) {
+      if (numPoints.getAffine()[i][0][0] == 1)
+        m += numPoints.getCoefficients()[i];
+      if (numPoints.getAffine()[i][0][1] == 1)
+        n += numPoints.getCoefficients()[i];
+      if (numPoints.getAffine()[i][0][2] == 1)
+        p += numPoints.getCoefficients()[i];
+    } else if (numPoints.getAffine()[i].size() == 2u) {
+      if ((numPoints.getAffine()[i][0][0] == 1 &&
+           numPoints.getAffine()[i][1][1] == 1) ||
+          (numPoints.getAffine()[i][0][1] == 1 &&
+           numPoints.getAffine()[i][1][0] == 1))
+        mn += numPoints.getCoefficients()[i];
+      if ((numPoints.getAffine()[i][0][1] == 1 &&
+           numPoints.getAffine()[i][1][2] == 1) ||
+          (numPoints.getAffine()[i][0][2] == 1 &&
+           numPoints.getAffine()[i][1][1] == 1))
+        np += numPoints.getCoefficients()[i];
+      if ((numPoints.getAffine()[i][0][2] == 1 &&
+           numPoints.getAffine()[i][1][0] == 1) ||
+          (numPoints.getAffine()[i][0][0] == 1 &&
+           numPoints.getAffine()[i][1][2] == 1))
+        pm += numPoints.getCoefficients()[i];
+    } else if (numPoints.getAffine()[i].size() == 3u) {
+      if ((numPoints.getAffine()[i][0][0] == 1 &&
+           numPoints.getAffine()[i][1][1] == 1 &&
+           numPoints.getAffine()[i][2][2] == 1) ||
+          (numPoints.getAffine()[i][0][0] == 1 &&
+           numPoints.getAffine()[i][1][2] == 1 &&
+           numPoints.getAffine()[i][2][1] == 1) ||
+          (numPoints.getAffine()[i][0][1] == 1 &&
+           numPoints.getAffine()[i][1][0] == 1 &&
+           numPoints.getAffine()[i][2][2] == 1) ||
+          (numPoints.getAffine()[i][0][1] == 1 &&
+           numPoints.getAffine()[i][1][2] == 1 &&
+           numPoints.getAffine()[i][2][0] == 1) ||
+          (numPoints.getAffine()[i][0][2] == 1 &&
+           numPoints.getAffine()[i][1][1] == 1 &&
+           numPoints.getAffine()[i][2][0] == 1) ||
+          (numPoints.getAffine()[i][0][2] == 1 &&
+           numPoints.getAffine()[i][1][0] == 1 &&
+           numPoints.getAffine()[i][2][1] == 1))
+        mnp += numPoints.getCoefficients()[i];
+    }
+  }
+
+  // We expect the answer to be
+  // (⌊M⌋ + 1)(⌊N⌋ + 1)(⌊P⌋ + 1) =
+  // ⌊M⌋⌊N⌋⌊P⌋ + ⌊M⌋⌊N⌋ + ⌊N⌋⌊P⌋ + ⌊M⌋⌊P⌋ + ⌊M⌋ + ⌊N⌋ + ⌊P⌋ + 1.
+  EXPECT_EQ(mn, 1);
+  EXPECT_EQ(np, 1);
+  EXPECT_EQ(pm, 1);
+  EXPECT_EQ(m, 1);
+  EXPECT_EQ(n, 1);
+  EXPECT_EQ(p, 1);
+  EXPECT_EQ(numPoints.getConstantTerm(), 1);
 }
\ No newline at end of file

>From 5ebc95da4b78f4c32952dda9dd48c3be8646b018 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 17 Jan 2024 14:34:25 +0530
Subject: [PATCH 28/44] Fix doc

---
 mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h | 5 +++--
 mlir/lib/Analysis/Presburger/Barvinok.cpp               | 2 +-
 mlir/unittests/Analysis/Presburger/BarvinokTest.cpp     | 2 +-
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index e5b4aae36db31e..aeac19e827b44f 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -59,8 +59,9 @@ class QuasiPolynomial : public PresburgerSpace {
   QuasiPolynomial operator*(const QuasiPolynomial &x) const;
   QuasiPolynomial operator/(const Fraction x) const;
 
-  // Removes terms which evaluate to zero and affine functions
-  // which are constant from the expression.
+  // Removes terms which evaluate to zero from the expression
+  // and folds affine functions which are constant into the
+  // constant coefficients.
   QuasiPolynomial simplify();
 
   // Group together like terms in the expression.
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 969c327ab60fa8..5327f4581e8baa 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -248,7 +248,7 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
 
 static std::vector<Fraction> convolution(std::vector<Fraction> a,
                                          std::vector<Fraction> b) {
-  // The length of the convolution is the maximum of the lengths
+  // The length of the convolution is the sum of the lengths
   // of the two sequences. We pad the shorter one with zeroes.
   unsigned len = a.size() + b.size();
   a.resize(len);
diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index e60ff9432ce865..ca628a6022da2f 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -195,7 +195,7 @@ TEST(BarvinokTest, computeNumTerms) {
   numPoints = numPoints.collectTerms().simplify();
 
   // First, we make sure all the affine functions are either
-  // ⌊M⌋, ⌊N⌋, or ⌊P⌋.
+  // M, N, or P.
   for (const std::vector<SmallVector<Fraction>> &term : numPoints.getAffine())
     for (const SmallVector<Fraction> &aff : term)
       EXPECT_EQ(aff[0] + aff[1] + aff[2] + aff[3], 1);

>From e6fae1d877cf88670e4d78e702eb91160bbe32f8 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 17 Jan 2024 14:35:20 +0530
Subject: [PATCH 29/44] 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 5327f4581e8baa..6596c0353711b2 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -265,7 +265,7 @@ static std::vector<Fraction> convolution(std::vector<Fraction> a,
   return convolution;
 }
 
-/// Substitute x_i = (s+1)^μ_i in one term of a generating function,
+/// Substitute x_i = t^μ_i in one term of a generating function,
 /// returning
 /// a quasipolynomial which represents the exponent of the numerator
 /// of the result, and
@@ -307,7 +307,7 @@ substituteMuInTerm(unsigned numParams, ParamPoint v, std::vector<Point> ds,
   // given by the dot product of μ with u_i.
   for (const Point &d : ds) {
     // This term in the denominator is
-    // (1 - (s+1)^dens.back())
+    // (1 - t^dens.back())
     dens.push_back(dotProduct(d, mu));
   }
 

>From b387fbfd1e40d0826ccb49607fb97472b234e9a2 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 17 Jan 2024 14:44:06 +0530
Subject: [PATCH 30/44] Add assert

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 6596c0353711b2..5974b0ec7e1f41 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -278,6 +278,10 @@ std::pair<QuasiPolynomial, std::vector<Fraction>>
 substituteMuInTerm(unsigned numParams, ParamPoint v, std::vector<Point> ds,
                    Point mu) {
   unsigned numDims = mu.size();
+  for (const Point &d : ds)
+    assert(d.size() == numDims &&
+           "μ has to have the same number of dimensions as the generators!");
+
   // First, the exponent in the numerator becomes
   // - (μ • u_1) * (floor(first col of v))
   // - (μ • u_2) * (floor(second col of v)) - ...

>From afefb1bc55c3676a4e8e5d1e272eb8def2b4a13d Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Fri, 19 Jan 2024 22:13:37 +0530
Subject: [PATCH 31/44] Fix test

---
 .../Analysis/Presburger/BarvinokTest.cpp      | 69 +++++--------------
 1 file changed, 18 insertions(+), 51 deletions(-)

diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index ca628a6022da2f..ffe3d7daad4ea5 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -202,61 +202,28 @@ TEST(BarvinokTest, computeNumTerms) {
 
   Fraction m = 0, n = 0, p = 0, mn = 0, np = 0, pm = 0, mnp = 0;
 
-  for (unsigned i = 0, e = numPoints.getAffine().size(); i < e; i++) {
-    if (numPoints.getAffine()[i].size() == 1u) {
-      if (numPoints.getAffine()[i][0][0] == 1)
-        m += numPoints.getCoefficients()[i];
-      if (numPoints.getAffine()[i][0][1] == 1)
-        n += numPoints.getCoefficients()[i];
-      if (numPoints.getAffine()[i][0][2] == 1)
-        p += numPoints.getCoefficients()[i];
-    } else if (numPoints.getAffine()[i].size() == 2u) {
-      if ((numPoints.getAffine()[i][0][0] == 1 &&
-           numPoints.getAffine()[i][1][1] == 1) ||
-          (numPoints.getAffine()[i][0][1] == 1 &&
-           numPoints.getAffine()[i][1][0] == 1))
-        mn += numPoints.getCoefficients()[i];
-      if ((numPoints.getAffine()[i][0][1] == 1 &&
-           numPoints.getAffine()[i][1][2] == 1) ||
-          (numPoints.getAffine()[i][0][2] == 1 &&
-           numPoints.getAffine()[i][1][1] == 1))
-        np += numPoints.getCoefficients()[i];
-      if ((numPoints.getAffine()[i][0][2] == 1 &&
-           numPoints.getAffine()[i][1][0] == 1) ||
-          (numPoints.getAffine()[i][0][0] == 1 &&
-           numPoints.getAffine()[i][1][2] == 1))
-        pm += numPoints.getCoefficients()[i];
-    } else if (numPoints.getAffine()[i].size() == 3u) {
-      if ((numPoints.getAffine()[i][0][0] == 1 &&
-           numPoints.getAffine()[i][1][1] == 1 &&
-           numPoints.getAffine()[i][2][2] == 1) ||
-          (numPoints.getAffine()[i][0][0] == 1 &&
-           numPoints.getAffine()[i][1][2] == 1 &&
-           numPoints.getAffine()[i][2][1] == 1) ||
-          (numPoints.getAffine()[i][0][1] == 1 &&
-           numPoints.getAffine()[i][1][0] == 1 &&
-           numPoints.getAffine()[i][2][2] == 1) ||
-          (numPoints.getAffine()[i][0][1] == 1 &&
-           numPoints.getAffine()[i][1][2] == 1 &&
-           numPoints.getAffine()[i][2][0] == 1) ||
-          (numPoints.getAffine()[i][0][2] == 1 &&
-           numPoints.getAffine()[i][1][1] == 1 &&
-           numPoints.getAffine()[i][2][0] == 1) ||
-          (numPoints.getAffine()[i][0][2] == 1 &&
-           numPoints.getAffine()[i][1][0] == 1 &&
-           numPoints.getAffine()[i][2][1] == 1))
-        mnp += numPoints.getCoefficients()[i];
+  // We store the coefficients of M, N and P in this array.
+  Fraction count[2][2][2];
+  unsigned mIndex, nIndex, pIndex;
+  for (unsigned i = 0, e = numPoints.getCoefficients().size(); i < e; i++) {
+    mIndex = nIndex = pIndex = 0;
+    for (const SmallVector<Fraction> &aff : numPoints.getAffine()[i]) {
+      if (aff[0] == 1)
+        mIndex = 1;
+      if (aff[1] == 1)
+        nIndex = 1;
+      if (aff[2] == 1)
+        pIndex = 1;
+      EXPECT_EQ(aff[3], 0);
     }
+    count[mIndex][nIndex][pIndex] += numPoints.getCoefficients()[i];
   }
 
   // We expect the answer to be
   // (⌊M⌋ + 1)(⌊N⌋ + 1)(⌊P⌋ + 1) =
   // ⌊M⌋⌊N⌋⌊P⌋ + ⌊M⌋⌊N⌋ + ⌊N⌋⌊P⌋ + ⌊M⌋⌊P⌋ + ⌊M⌋ + ⌊N⌋ + ⌊P⌋ + 1.
-  EXPECT_EQ(mn, 1);
-  EXPECT_EQ(np, 1);
-  EXPECT_EQ(pm, 1);
-  EXPECT_EQ(m, 1);
-  EXPECT_EQ(n, 1);
-  EXPECT_EQ(p, 1);
-  EXPECT_EQ(numPoints.getConstantTerm(), 1);
+  for (unsigned i = 0; i < 2; i++)
+    for (unsigned j = 0; j < 2; j++)
+      for (unsigned k = 0; k < 2; k++)
+        EXPECT_EQ(count[i][j][k], 1);
 }
\ No newline at end of file

>From ca40b83f43fbcd0f34045389a5de844e27ee9845 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Fri, 19 Jan 2024 22:55:18 +0530
Subject: [PATCH 32/44] Add accessors

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 5974b0ec7e1f41..050269bb1a59d5 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -246,20 +246,32 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
   return coefficients[power].simplify();
 }
 
-static std::vector<Fraction> convolution(std::vector<Fraction> a,
-                                         std::vector<Fraction> b) {
+static std::vector<Fraction> convolution(ArrayRef<Fraction> a,
+                                         ArrayRef<Fraction> b) {
   // The length of the convolution is the sum of the lengths
   // of the two sequences. We pad the shorter one with zeroes.
   unsigned len = a.size() + b.size();
-  a.resize(len);
-  b.resize(len);
+
+  // We define accessors to avoid out-of-bounds errors.
+  std::function<Fraction(unsigned i)> aGetItem = [a](unsigned i) -> Fraction {
+    if (i < a.size())
+      return a[i];
+    else
+      return 0;
+  };
+  std::function<Fraction(unsigned i)> bGetItem = [b](unsigned i) -> Fraction {
+    if (i < b.size())
+      return b[i];
+    else
+      return 0;
+  };
 
   std::vector<Fraction> convolution;
   convolution.reserve(len);
   for (unsigned k = 0; k < len; ++k) {
     Fraction sum(0, 1);
     for (unsigned l = 0; l <= k; ++l)
-      sum += a[l] * b[k - l];
+      sum += aGetItem(l) * bGetItem(k - l);
     convolution.push_back(sum);
   }
   return convolution;

>From 64bf7386e980ebc2cf1331185af7580c6cd7fe51 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Fri, 19 Jan 2024 22:55:40 +0530
Subject: [PATCH 33/44] Fix doc

---
 mlir/include/mlir/Analysis/Presburger/Barvinok.h | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 50d6742019161c..055606fcb9e7b7 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -99,8 +99,10 @@ QuasiPolynomial getCoefficientInRationalFunction(unsigned power,
                                                  ArrayRef<QuasiPolynomial> num,
                                                  ArrayRef<Fraction> den);
 
-/// Find the number of terms in the generating function corresponding to
-/// a polytope.
+/// Find the number of terms in a generating function, which is
+/// a quasipolynomial in the parameter space of the input function.
+/// The generating function must be such that for all values of the
+/// parameters, the number of terms is finite; else it returns 0.
 QuasiPolynomial computeNumTerms(const GeneratingFunction &gf);
 
 } // namespace detail

>From 0c261a10e5301d09c2f85dd22699b17b54e20e21 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sat, 20 Jan 2024 22:25:28 +0530
Subject: [PATCH 34/44] Move convolution

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

diff --git a/mlir/include/mlir/Analysis/Presburger/Utils.h b/mlir/include/mlir/Analysis/Presburger/Utils.h
index 20af0bfcd62ba6..0f74eca0da940f 100644
--- a/mlir/include/mlir/Analysis/Presburger/Utils.h
+++ b/mlir/include/mlir/Analysis/Presburger/Utils.h
@@ -281,6 +281,8 @@ SmallVector<MPInt, 8> getComplementIneq(ArrayRef<MPInt> ineq);
 /// The vectors must have the same sizes.
 Fraction dotProduct(ArrayRef<Fraction> a, ArrayRef<Fraction> b);
 
+std::vector<Fraction> convolution(ArrayRef<Fraction> a, ArrayRef<Fraction> b);
+
 } // namespace presburger
 } // namespace mlir
 
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 050269bb1a59d5..f49bb575af3cc7 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/Barvinok.h"
+#include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/Sequence.h"
 #include <algorithm>
 
@@ -246,37 +247,6 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
   return coefficients[power].simplify();
 }
 
-static std::vector<Fraction> convolution(ArrayRef<Fraction> a,
-                                         ArrayRef<Fraction> b) {
-  // The length of the convolution is the sum of the lengths
-  // of the two sequences. We pad the shorter one with zeroes.
-  unsigned len = a.size() + b.size();
-
-  // We define accessors to avoid out-of-bounds errors.
-  std::function<Fraction(unsigned i)> aGetItem = [a](unsigned i) -> Fraction {
-    if (i < a.size())
-      return a[i];
-    else
-      return 0;
-  };
-  std::function<Fraction(unsigned i)> bGetItem = [b](unsigned i) -> Fraction {
-    if (i < b.size())
-      return b[i];
-    else
-      return 0;
-  };
-
-  std::vector<Fraction> convolution;
-  convolution.reserve(len);
-  for (unsigned k = 0; k < len; ++k) {
-    Fraction sum(0, 1);
-    for (unsigned l = 0; l <= k; ++l)
-      sum += aGetItem(l) * bGetItem(k - l);
-    convolution.push_back(sum);
-  }
-  return convolution;
-}
-
 /// Substitute x_i = t^μ_i in one term of a generating function,
 /// returning
 /// a quasipolynomial which represents the exponent of the numerator
@@ -402,12 +372,13 @@ std::vector<Fraction> getBinomialCoefficients(Fraction n, Fraction r) {
 ///
 /// where sign_i is ±1,
 /// n_i \in Q^p -> Q^d is the sum of the vectors d_{ij}, weighted by the
-/// floors of d affine functions.
+/// floors of d affine functions on p parameters.
 /// 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.
-/// However, direct substitution causes the denominator to become zero.
+/// However, direct substitution (x = (1, ..., 1)) 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
diff --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp
index 5e267d2045bc1c..5d5927f1450b16 100644
--- a/mlir/lib/Analysis/Presburger/Utils.cpp
+++ b/mlir/lib/Analysis/Presburger/Utils.cpp
@@ -537,4 +537,35 @@ Fraction presburger::dotProduct(ArrayRef<Fraction> a, ArrayRef<Fraction> b) {
   for (unsigned i = 0, e = a.size(); i < e; i++)
     sum += a[i] * b[i];
   return sum;
+}
+
+std::vector<Fraction> presburger::convolution(ArrayRef<Fraction> a,
+                                              ArrayRef<Fraction> b) {
+  // The length of the convolution is the sum of the lengths
+  // of the two sequences. We pad the shorter one with zeroes.
+  unsigned len = a.size() + b.size();
+
+  // We define accessors to avoid out-of-bounds errors.
+  std::function<Fraction(unsigned i)> aGetItem = [a](unsigned i) -> Fraction {
+    if (i < a.size())
+      return a[i];
+    else
+      return 0;
+  };
+  std::function<Fraction(unsigned i)> bGetItem = [b](unsigned i) -> Fraction {
+    if (i < b.size())
+      return b[i];
+    else
+      return 0;
+  };
+
+  std::vector<Fraction> convolution;
+  convolution.reserve(len);
+  for (unsigned k = 0; k < len; ++k) {
+    Fraction sum(0, 1);
+    for (unsigned l = 0; l <= k; ++l)
+      sum += aGetItem(l) * bGetItem(k - l);
+    convolution.push_back(sum);
+  }
+  return convolution;
 }
\ No newline at end of file

>From b6094454541ce279efee6b7a136bf57526494fef Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 21 Jan 2024 23:25:31 +0530
Subject: [PATCH 35/44] Refactor accessors

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

diff --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp
index 5d5927f1450b16..51e5cac614d923 100644
--- a/mlir/lib/Analysis/Presburger/Utils.cpp
+++ b/mlir/lib/Analysis/Presburger/Utils.cpp
@@ -546,15 +546,10 @@ std::vector<Fraction> presburger::convolution(ArrayRef<Fraction> a,
   unsigned len = a.size() + b.size();
 
   // We define accessors to avoid out-of-bounds errors.
-  std::function<Fraction(unsigned i)> aGetItem = [a](unsigned i) -> Fraction {
-    if (i < a.size())
-      return a[i];
-    else
-      return 0;
-  };
-  std::function<Fraction(unsigned i)> bGetItem = [b](unsigned i) -> Fraction {
-    if (i < b.size())
-      return b[i];
+  std::function<Fraction(ArrayRef<Fraction>, unsigned)> getItem =
+      [](ArrayRef<Fraction> arr, unsigned i) -> Fraction {
+    if (i < arr.size())
+      return arr[i];
     else
       return 0;
   };
@@ -564,7 +559,7 @@ std::vector<Fraction> presburger::convolution(ArrayRef<Fraction> a,
   for (unsigned k = 0; k < len; ++k) {
     Fraction sum(0, 1);
     for (unsigned l = 0; l <= k; ++l)
-      sum += aGetItem(l) * bGetItem(k - l);
+      sum += getItem(a, l) * getItem(b, k - l);
     convolution.push_back(sum);
   }
   return convolution;

>From 02c2ce3dba94739b24165e575e8dd43580517dba Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 21 Jan 2024 23:30:58 +0530
Subject: [PATCH 36/44] Fix docs

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp          | 14 +++++---------
 mlir/lib/Analysis/Presburger/Utils.cpp             |  3 +--
 .../unittests/Analysis/Presburger/BarvinokTest.cpp |  3 +--
 3 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index f49bb575af3cc7..444fdc362ba958 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -247,15 +247,12 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
   return coefficients[power].simplify();
 }
 
-/// Substitute x_i = t^μ_i in one term of a generating function,
-/// returning
+/// Substitute x_i = t^μ_i in one term of a generating function, 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.
-/// v represents the affine functions whose floors are multiplied by
-/// by the generators, and
-/// ds represents the list of generators.
+/// 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 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) {
@@ -405,7 +402,6 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
   Point mu = getNonOrthogonalVector(allDenominators);
 
   unsigned numParams = gf.getNumParams();
-
   const std::vector<std::vector<Point>> &ds = gf.getDenominators();
   QuasiPolynomial totalTerm(numParams, 0);
   for (unsigned i = 0, e = ds.size(); i < e; ++i) {
diff --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp
index 51e5cac614d923..48bc9137b6f9bf 100644
--- a/mlir/lib/Analysis/Presburger/Utils.cpp
+++ b/mlir/lib/Analysis/Presburger/Utils.cpp
@@ -546,8 +546,7 @@ std::vector<Fraction> presburger::convolution(ArrayRef<Fraction> a,
   unsigned len = a.size() + b.size();
 
   // We define accessors to avoid out-of-bounds errors.
-  std::function<Fraction(ArrayRef<Fraction>, unsigned)> getItem =
-      [](ArrayRef<Fraction> arr, unsigned i) -> Fraction {
+  auto getItem = [](ArrayRef<Fraction> arr, unsigned i) -> Fraction {
     if (i < arr.size())
       return arr[i];
     else
diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index ffe3d7daad4ea5..dced217a6f1589 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -169,8 +169,7 @@ TEST(BarvinokTest, computeNumTerms) {
   EXPECT_EQ(numPoints.getConstantTerm(), Fraction(1, 1));
 
   // The following generating function corresponds to a cuboid
-  // with length (x-axis) M, width (y-axis) N, and height (z-axis)
-  // P.
+  // with length M (x-axis), width N (y-axis), and height P (z-axis).
   // There are eight terms.
   gf = GeneratingFunction(
       3, {1, 1, 1, 1, 1, 1, 1, 1},

>From 65baf078bf5057b8063e81b2288d21f3a7d64b8f Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 21 Jan 2024 23:36:14 +0530
Subject: [PATCH 37/44] Minor fixes

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 444fdc362ba958..be98a46776b8d9 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -432,7 +432,7 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     // 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
+    // is polynomial in s, with integral coefficients (given by taking the
     // convolution over all j).
 
     // Step (3) We need to find the constant term in the expansion of each
diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index dced217a6f1589..df0933f0ddcc74 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -157,11 +157,12 @@ TEST(BarvinokTest, computeNumTerms) {
   // 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;
+  SmallVector<Fraction> coefficients = numPoints.getCoefficients();
   for (unsigned i = 0; i < numPoints.getCoefficients().size(); i++)
     if (numPoints.getAffine()[i].size() == 2)
-      pSquaredCoeff = pSquaredCoeff + numPoints.getCoefficients()[i];
+      pSquaredCoeff = pSquaredCoeff + coefficients[i];
     else if (numPoints.getAffine()[i].size() == 1)
-      pCoeff = pCoeff + numPoints.getCoefficients()[i];
+      pCoeff = pCoeff + coefficients[i];
 
   // We expect the answer to be (1/2)⌊p/2⌋^2 + (3/2)⌊p/2⌋ + 1.
   EXPECT_EQ(pSquaredCoeff, Fraction(1, 2));
@@ -199,13 +200,11 @@ TEST(BarvinokTest, computeNumTerms) {
     for (const SmallVector<Fraction> &aff : term)
       EXPECT_EQ(aff[0] + aff[1] + aff[2] + aff[3], 1);
 
-  Fraction m = 0, n = 0, p = 0, mn = 0, np = 0, pm = 0, mnp = 0;
-
   // We store the coefficients of M, N and P in this array.
   Fraction count[2][2][2];
-  unsigned mIndex, nIndex, pIndex;
-  for (unsigned i = 0, e = numPoints.getCoefficients().size(); i < e; i++) {
-    mIndex = nIndex = pIndex = 0;
+  coefficients = numPoints.getCoefficients();
+  for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
+    unsigned mIndex = 0, nIndex = 0, pIndex = 0;
     for (const SmallVector<Fraction> &aff : numPoints.getAffine()[i]) {
       if (aff[0] == 1)
         mIndex = 1;
@@ -215,7 +214,7 @@ TEST(BarvinokTest, computeNumTerms) {
         pIndex = 1;
       EXPECT_EQ(aff[3], 0);
     }
-    count[mIndex][nIndex][pIndex] += numPoints.getCoefficients()[i];
+    count[mIndex][nIndex][pIndex] += coefficients[i];
   }
 
   // We expect the answer to be

>From 5289ffdf3e004f64b47760b9726ec9ab4163f505 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 21 Jan 2024 23:43:36 +0530
Subject: [PATCH 38/44] FIx normalize

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index be98a46776b8d9..562602c192a69a 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -311,11 +311,6 @@ void normalizeDenominatorExponents(int &sign, QuasiPolynomial &num,
       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,
@@ -329,11 +324,6 @@ void normalizeDenominatorExponents(int &sign, QuasiPolynomial &num,
   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;
 }
 
 /// Compute the binomial coefficients nCi for 0 ≤ i ≤ r,
@@ -421,10 +411,21 @@ 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.
+    normalizeDenominatorExponents(sign, numExp, dens);
+
     // 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);
+    for (unsigned j = 0, e = dens.size(); j < e; ++j)
+      dens[j] = abs(dens[j]) - 1;
+    // Note that at this point, the semantics of `dens[j]` changes to mean
+    // the limit of summation of (s+1)^t terms, and not the actual exponent.
+
+    // Since the -s are taken out, the sign changes if there is an odd number
+    // of such terms.
+    unsigned r = dens.size();
+    if (dens.size() % 2 == 1)
+      sign = -sign;
 
     // Thus the term overall has now the form
     // sign'_i * (s+1)^numExp /
@@ -442,7 +443,6 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     // 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, we compute the coefficients of P(s),
     // which are binomial coefficients.

>From ff5f967a2dbe083a5c2aac8248acf1ad474a0497 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 21 Jan 2024 23:44:38 +0530
Subject: [PATCH 39/44] Fix doc

---
 mlir/lib/Analysis/Presburger/Barvinok.cpp        | 4 +---
 mlir/unittests/Analysis/Presburger/UtilsTest.cpp | 9 +++++++++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 562602c192a69a..b56a5bd8072bc1 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -408,9 +408,7 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     // 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.
+    // absolute values and change the sign and exponent in the numerator.
     normalizeDenominatorExponents(sign, numExp, dens);
 
     // Then, using the formula for geometric series, we replace each (1 -
diff --git a/mlir/unittests/Analysis/Presburger/UtilsTest.cpp b/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
index 7c1646aa841162..ba286d9af3ab5c 100644
--- a/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
@@ -66,3 +66,12 @@ TEST(UtilsTest, DivisionReprNormalizeTest) {
   checkEqual(a, b);
   checkEqual(c, d);
 }
+
+TEST(UtilsTest, convolution) {
+  ArrayRef<Fraction> x({1, 2, 3, 4});
+  ArrayRef<Fraction> y({7, 3, 1, 6});
+
+  std::vector<Fraction> conv = convolution(x, y);
+
+  EXPECT_EQ(conv, std::vector<Fraction>({7, 17, 28, 45, 27, 22, 24}));
+}
\ No newline at end of file

>From b919d7c94bcf01719bceaefa065e30d79186a202 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 21 Jan 2024 23:48:49 +0530
Subject: [PATCH 40/44] Fix doc

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

diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index b56a5bd8072bc1..364467e55b0fc9 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -397,7 +397,8 @@ 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 μ.
+    // Compute the new exponents of (s+1) for the numerator and the
+    // denominator after substituting μ.
     auto [numExp, dens] =
         substituteMuInTerm(numParams, gf.getNumerators()[i], ds[i], mu);
     // Now the numerator is (s+1)^numExp
@@ -442,8 +443,8 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     // we need to find the coefficient of s^r in P(s)/Q(s),
     // for which we use the `getCoefficientInRationalFunction()` function.
 
-    // First, we compute the coefficients of P(s),
-    // which are binomial coefficients.
+    // 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 =

>From 66bc0d24a6a090d06be87d17738bceee27caff47 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 21 Jan 2024 23:51:22 +0530
Subject: [PATCH 41/44] Rename conv

---
 mlir/include/mlir/Analysis/Presburger/Utils.h    | 5 ++++-
 mlir/lib/Analysis/Presburger/Barvinok.cpp        | 4 ++--
 mlir/lib/Analysis/Presburger/Utils.cpp           | 6 ++++--
 mlir/unittests/Analysis/Presburger/UtilsTest.cpp | 2 +-
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Utils.h b/mlir/include/mlir/Analysis/Presburger/Utils.h
index 0f74eca0da940f..e6d29e4ef6d062 100644
--- a/mlir/include/mlir/Analysis/Presburger/Utils.h
+++ b/mlir/include/mlir/Analysis/Presburger/Utils.h
@@ -281,7 +281,10 @@ SmallVector<MPInt, 8> getComplementIneq(ArrayRef<MPInt> ineq);
 /// The vectors must have the same sizes.
 Fraction dotProduct(ArrayRef<Fraction> a, ArrayRef<Fraction> b);
 
-std::vector<Fraction> convolution(ArrayRef<Fraction> a, ArrayRef<Fraction> b);
+/// Find the product of two polynomials, each given by an array of
+/// coefficients.
+std::vector<Fraction> multiplyPolynomials(ArrayRef<Fraction> a,
+                                          ArrayRef<Fraction> b);
 
 } // namespace presburger
 } // namespace mlir
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 364467e55b0fc9..90eaae2e87f930 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -467,8 +467,8 @@ mlir::presburger::detail::computeNumTerms(const GeneratingFunction &gf) {
     std::vector<Fraction> denominatorCoefficients;
     denominatorCoefficients = eachTermDenCoefficients[0];
     for (unsigned j = 1, e = eachTermDenCoefficients.size(); j < e; ++j)
-      denominatorCoefficients =
-          convolution(denominatorCoefficients, eachTermDenCoefficients[j]);
+      denominatorCoefficients = multiplyPolynomials(denominatorCoefficients,
+                                                    eachTermDenCoefficients[j]);
 
     totalTerm =
         totalTerm + getCoefficientInRationalFunction(r, numeratorCoefficients,
diff --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp
index 48bc9137b6f9bf..923e003b7220e6 100644
--- a/mlir/lib/Analysis/Presburger/Utils.cpp
+++ b/mlir/lib/Analysis/Presburger/Utils.cpp
@@ -539,8 +539,10 @@ Fraction presburger::dotProduct(ArrayRef<Fraction> a, ArrayRef<Fraction> b) {
   return sum;
 }
 
-std::vector<Fraction> presburger::convolution(ArrayRef<Fraction> a,
-                                              ArrayRef<Fraction> b) {
+/// Find the product of two polynomials, each given by an array of
+/// coefficients, by taking the convolution.
+std::vector<Fraction> presburger::multiplyPolynomials(ArrayRef<Fraction> a,
+                                                      ArrayRef<Fraction> b) {
   // The length of the convolution is the sum of the lengths
   // of the two sequences. We pad the shorter one with zeroes.
   unsigned len = a.size() + b.size();
diff --git a/mlir/unittests/Analysis/Presburger/UtilsTest.cpp b/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
index ba286d9af3ab5c..ea116a6f51e0df 100644
--- a/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
@@ -71,7 +71,7 @@ TEST(UtilsTest, convolution) {
   ArrayRef<Fraction> x({1, 2, 3, 4});
   ArrayRef<Fraction> y({7, 3, 1, 6});
 
-  std::vector<Fraction> conv = convolution(x, y);
+  std::vector<Fraction> conv = multiplyPolynomials(x, y);
 
   EXPECT_EQ(conv, std::vector<Fraction>({7, 17, 28, 45, 27, 22, 24}));
 }
\ No newline at end of file

>From 44d9fc5a2ad92a573970ec29257801352bc74b89 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 21 Jan 2024 23:54:03 +0530
Subject: [PATCH 42/44] Fix test

---
 .../Analysis/Presburger/BarvinokTest.cpp          | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index df0933f0ddcc74..919aaa7a428593 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -195,10 +195,19 @@ TEST(BarvinokTest, computeNumTerms) {
   numPoints = numPoints.collectTerms().simplify();
 
   // First, we make sure all the affine functions are either
-  // M, N, or P.
-  for (const std::vector<SmallVector<Fraction>> &term : numPoints.getAffine())
-    for (const SmallVector<Fraction> &aff : term)
+  // M, N, P, or 1.
+  for (const std::vector<SmallVector<Fraction>> &term : numPoints.getAffine()) {
+    for (const SmallVector<Fraction> &aff : term) {
+      // First, ensure that the coefficients are all nonnegative integers.
+      for (const Fraction &c : aff) {
+        EXPECT_TRUE(c >= 0);
+        EXPECT_EQ(c, c.getAsInteger());
+      }
+      // Now, if the coefficients add up to 1, we can be sure the term is
+      // either M, N, P, or 1.
       EXPECT_EQ(aff[0] + aff[1] + aff[2] + aff[3], 1);
+    }
+  }
 
   // We store the coefficients of M, N and P in this array.
   Fraction count[2][2][2];

>From 8c7e15bdefc13077bb66c35d061538812964a2fc Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 22 Jan 2024 00:09:29 +0530
Subject: [PATCH 43/44] Remove extra element from conv

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

diff --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp
index 923e003b7220e6..250cb6284654e9 100644
--- a/mlir/lib/Analysis/Presburger/Utils.cpp
+++ b/mlir/lib/Analysis/Presburger/Utils.cpp
@@ -545,7 +545,7 @@ std::vector<Fraction> presburger::multiplyPolynomials(ArrayRef<Fraction> a,
                                                       ArrayRef<Fraction> b) {
   // The length of the convolution is the sum of the lengths
   // of the two sequences. We pad the shorter one with zeroes.
-  unsigned len = a.size() + b.size();
+  unsigned len = a.size() + b.size() - 1;
 
   // We define accessors to avoid out-of-bounds errors.
   auto getItem = [](ArrayRef<Fraction> arr, unsigned i) -> Fraction {

>From f809eae2ff9754c9489bc7e283fc336927018c36 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 22 Jan 2024 00:14:10 +0530
Subject: [PATCH 44/44] Conv test

---
 .../Analysis/Presburger/UtilsTest.cpp          | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/mlir/unittests/Analysis/Presburger/UtilsTest.cpp b/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
index ea116a6f51e0df..f09a1a760ce609 100644
--- a/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
@@ -68,10 +68,20 @@ TEST(UtilsTest, DivisionReprNormalizeTest) {
 }
 
 TEST(UtilsTest, convolution) {
-  ArrayRef<Fraction> x({1, 2, 3, 4});
-  ArrayRef<Fraction> y({7, 3, 1, 6});
+  std::vector<Fraction> aVals({1, 2, 3, 4});
+  std::vector<Fraction> bVals({7, 3, 1, 6});
+  ArrayRef<Fraction> a(aVals);
+  ArrayRef<Fraction> b(bVals);
 
-  std::vector<Fraction> conv = multiplyPolynomials(x, y);
+  std::vector<Fraction> conv = multiplyPolynomials(a, b);
 
   EXPECT_EQ(conv, std::vector<Fraction>({7, 17, 28, 45, 27, 22, 24}));
-}
\ No newline at end of file
+
+  aVals = {3, 6, 0, 2, 5};
+  bVals = {2, 0, 6};
+  a = aVals;
+  b = bVals;
+
+  conv = multiplyPolynomials(a, b);
+  EXPECT_EQ(conv, std::vector<Fraction>({6, 12, 18, 40, 10, 12, 30}));
+}



More information about the Mlir-commits mailing list