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

[Arjun P]

================
@@ -245,3 +245,236 @@ 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;
+}
+
+/// 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, 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 products.
+
+  // 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 function 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();
+
+  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
+/// 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 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
+/// this function.
+/// 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
+/// 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.
+/// 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) {
+  std::vector<Point> allDenominators;
+  for (ArrayRef<Point> den : gf.getDenominators())
+    allDenominators.insert(allDenominators.end(), den.begin(), den.end());
+  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) {
+    int sign = gf.getSigns()[i];
+
+    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 /
+    // (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).
+
+    unsigned r = dens.size();
+    // First, the coefficients of P(s), which are binomial coefficients.
+    // We need r+1 of these.
----------------
Superty wrote:

We only need the first r+1 of these because...

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