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

[Arjun P]

================
@@ -245,3 +245,253 @@ 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);
----------------
Superty wrote:

you can do return {num, dens}

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