[Mlir-commits] [mlir] [MLIR][Presburger] Helper functions to compute the constant term of a generating function (PR #77819)

[Arjun P]

================
@@ -144,3 +145,102 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
                             std::vector({numerator}),
                             std::vector({denominator}));
 }
+
+/// We use an iterative procedure to find a vector not orthogonal
+/// to a given set, ignoring the null vectors.
+/// Let the inputs be {x_1, ..., x_k}, all vectors of length n.
+///
+/// In the following,
+/// vs[:i] means the elements of vs up to and including the i'th one,
+/// <vs, us> means the dot product of vs and us,
+/// vs ++ [v] means the vector vs with the new element v appended to it.
+///
+/// We proceed iteratively; for steps d = 0, ... n-1, we construct a vector
+/// which is not orthogonal to any of {x_1[:d], ..., x_n[:d]}, ignoring
+/// the null vectors.
+/// At step d = 0, we let vs = [1]. Clearly this is not orthogonal to
+/// any vector in the set {x_1[0], ..., x_n[0]}, except the null ones,
+/// which we ignore.
+/// At step d > 0 , we need a number v
+/// s.t. <x_i[:d], vs++[v]> != 0 for all i.
+/// => <x_i[:d-1], vs> + x_i[d]*v != 0
+/// => v != - <x_i[:d-1], vs> / x_i[d]
+/// We compute this value for all x_i, and then
+/// set v to be the maximum element of this set plus one. Thus
+/// v is outside the set as desired, and we append it to vs
+/// to obtain the result of the d'th step.
+Point mlir::presburger::detail::getNonOrthogonalVector(
+    ArrayRef<Point> vectors) {
+  unsigned dim = vectors[0].size();
+  for (const Point &vector : vectors)
+    assert(vector.size() == dim && "all vectors need to be the same size!");
+
+  SmallVector<Fraction> newPoint = {Fraction(1, 1)};
+  std::vector<Fraction> lowerDimDotProducts;
+  Fraction dotP;
+  Fraction maxDisallowedValue = -Fraction(1, 0),
+           disallowedValue = Fraction(0, 1);
+  Fraction newValue;
+
+  for (unsigned d = 1; d < dim; ++d) {
+
+    // Compute the disallowed values  - <x_i[:d-1], vs> / x_i[d] for each i.
+    maxDisallowedValue = -Fraction(1, 0);
+    for (const Point &vector : vectors) {
+      if (vector[d] == 0)
+        continue;
+      dotP = dotProduct(ArrayRef(vector).slice(0, d), newPoint);
+      disallowedValue =
+          -dotProduct(ArrayRef(vector).slice(0, d), newPoint) / vector[d];
+
+      // Find the biggest such value
+      maxDisallowedValue = std::max(maxDisallowedValue, disallowedValue);
+    }
+
+    newValue = Fraction(ceil(maxDisallowedValue + Fraction(1, 1)), 1);
+    newPoint.push_back(newValue);
+  }
+  return newPoint;
+}
+
+/// We use the following recursive formula to find the coefficient of
+/// s^power in the rational function given by P(s)/Q(s).
+///
+/// Let P[i] denote the coefficient of s^i in the polynomial P(s).
+/// (P/Q)[r] =
+/// if (r == 0) then
+///   P[0]/Q[0]
+/// else
+///   (P[r] - {Σ_{i=1}^r (P/Q)[r-i] * Q[i])}/(Q[0])
+/// We therefore recursively call `getCoefficientInRationalFunction` on
+/// all i \in [0, power).
+///
+/// https://math.ucdavis.edu/~deloera/researchsummary/
+/// barvinokalgorithm-latte1.pdf, p. 1285
+QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
+    unsigned power, ArrayRef<QuasiPolynomial> num, ArrayRef<Fraction> den) {
+  assert(den.size() != 0 &&
+         "division by empty denominator in rational function!");
+
+  unsigned numParam = num[0].getNumInputs();
+  for (const QuasiPolynomial &qp : num)
+    assert(numParam == qp.getNumInputs() &&
+           "the quasipolynomials should all belong to the same space!");
+
+  std::vector<QuasiPolynomial> coefficients;
+  coefficients.reserve(power + 1);
+
+  coefficients.push_back(num[0] / den[0]);
+  for (unsigned i = 1; i <= power; ++i) {
+    coefficients.push_back(i < num.size() ? num[i]
+                                          : QuasiPolynomial(numParam, 0));
+
+    unsigned limit = std::min<unsigned long>(i, den.size() - 1);
+    for (unsigned j = 1; j <= limit; ++j)
+      coefficients[i] = coefficients[i] -
----------------
Superty wrote:

you can do that later

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