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

[Arjun P]

================
@@ -124,3 +124,139 @@ TEST(BarvinokTest, getCoefficientInRationalFunction) {
   coeff = getCoefficientInRationalFunction(3, numerator, denominator);
   EXPECT_EQ(coeff.getConstantTerm(), Fraction(55, 64));
 }
+
+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}}),
+       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 = 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()) {
+    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];
+
+  // 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(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) {
----------------
Superty wrote:

This is error-prone/difficult to maintain. Please use a better way.

You can just compute the total exponent for each of M, N, P by iterating through the affines. Then you can index into some count[2][2][2] array.

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