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

[Arjun P]

================
@@ -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 (ArrayRef<Point> den : gf.getDenominators())
+    allDenominators.insert(allDenominators.end(), den.begin(), den.end());
+  Point mu = getNonOrthogonalVector(allDenominators);
+
+  unsigned numParams = gf.getNumParams();
+  unsigned numDims = mu.size();
+  unsigned numTerms = gf.getDenominators().size();
+
+  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
+    // 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.
+    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))});
+
+    QuasiPolynomial num(numParams, coefficients, affine);
+    num = num.simplify();
+
+    // 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
+    // (see lines 356-66).
+    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, from line 353.
+    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.
+    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] *
+           (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
+    std::vector<std::vector<Fraction>> eachTermDenCoefficients;
----------------
Superty wrote:

pull this convolution into another free functions... it helps to break up independent subparts into separate functions. it's harder to read through a long function body where there may be arbitrary interactions between parts

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