[Mlir-commits] [mlir] 68a5261 - [MLIR][Presburger] Implement function to evaluate the number of terms in a generating function. (#78078)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jan 22 00:52:06 PST 2024
Author: Abhinav271828
Date: 2024-01-22T14:22:01+05:30
New Revision: 68a5261d260e95148755c6725f56957cb8fd23a3
URL: https://github.com/llvm/llvm-project/commit/68a5261d260e95148755c6725f56957cb8fd23a3
DIFF: https://github.com/llvm/llvm-project/commit/68a5261d260e95148755c6725f56957cb8fd23a3.diff
LOG: [MLIR][Presburger] Implement function to evaluate the number of terms in a generating function. (#78078)
We implement `computeNumTerms()`, which counts the number of terms in a
generating function by substituting the unit vector in it.
This is the main function in Barvinok's algorithm – the number of points
in a polytope is given by the number of terms in the generating function
corresponding to it.
We also modify the GeneratingFunction class to have `const` getters and
improve the simplification of QuasiPolynomials.
Added:
Modified:
mlir/include/mlir/Analysis/Presburger/Barvinok.h
mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
mlir/include/mlir/Analysis/Presburger/Utils.h
mlir/lib/Analysis/Presburger/Barvinok.cpp
mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
mlir/lib/Analysis/Presburger/Utils.cpp
mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
mlir/unittests/Analysis/Presburger/UtilsTest.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index edee19f0e1a535..b70ec33b693235 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -99,6 +99,12 @@ QuasiPolynomial getCoefficientInRationalFunction(unsigned power,
ArrayRef<QuasiPolynomial> num,
ArrayRef<Fraction> den);
+/// Find the number of terms in a generating function, as
+/// a quasipolynomial in the parameter space of the input function.
+/// The generating function must be such that for all values of the
+/// parameters, the number of terms is finite.
+QuasiPolynomial computeNumTerms(const GeneratingFunction &gf);
+
} // namespace detail
} // namespace presburger
} // namespace mlir
diff --git a/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h b/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
index eaf0449fe8a118..c38eab6efd0fc1 100644
--- a/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
+++ b/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
@@ -62,13 +62,15 @@ class GeneratingFunction {
#endif // NDEBUG
}
- unsigned getNumParams() { return numParam; }
+ unsigned getNumParams() const { return numParam; }
- SmallVector<int> getSigns() { return signs; }
+ SmallVector<int> getSigns() const { return signs; }
- std::vector<ParamPoint> getNumerators() { return numerators; }
+ std::vector<ParamPoint> getNumerators() const { return numerators; }
- std::vector<std::vector<Point>> getDenominators() { return denominators; }
+ std::vector<std::vector<Point>> getDenominators() const {
+ return denominators;
+ }
GeneratingFunction operator+(GeneratingFunction &gf) const {
assert(numParam == gf.getNumParams() &&
diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index d03446f50264fa..aeac19e827b44f 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -59,9 +59,14 @@ class QuasiPolynomial : public PresburgerSpace {
QuasiPolynomial operator*(const QuasiPolynomial &x) const;
QuasiPolynomial operator/(const Fraction x) const;
- // Removes terms which evaluate to zero from the expression.
+ // Removes terms which evaluate to zero from the expression
+ // and folds affine functions which are constant into the
+ // constant coefficients.
QuasiPolynomial simplify();
+ // Group together like terms in the expression.
+ QuasiPolynomial collectTerms();
+
Fraction getConstantTerm();
private:
diff --git a/mlir/include/mlir/Analysis/Presburger/Utils.h b/mlir/include/mlir/Analysis/Presburger/Utils.h
index 20af0bfcd62ba6..e6d29e4ef6d062 100644
--- a/mlir/include/mlir/Analysis/Presburger/Utils.h
+++ b/mlir/include/mlir/Analysis/Presburger/Utils.h
@@ -281,6 +281,11 @@ SmallVector<MPInt, 8> getComplementIneq(ArrayRef<MPInt> ineq);
/// The vectors must have the same sizes.
Fraction dotProduct(ArrayRef<Fraction> a, ArrayRef<Fraction> b);
+/// Find the product of two polynomials, each given by an array of
+/// coefficients.
+std::vector<Fraction> multiplyPolynomials(ArrayRef<Fraction> a,
+ ArrayRef<Fraction> b);
+
} // namespace presburger
} // namespace mlir
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 4ba4462af0317f..e0fd0dd8caa4d3 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Analysis/Presburger/Barvinok.h"
+#include "mlir/Analysis/Presburger/Utils.h"
#include "llvm/ADT/Sequence.h"
#include <algorithm>
@@ -245,3 +246,241 @@ QuasiPolynomial mlir::presburger::detail::getCoefficientInRationalFunction(
}
return coefficients[power].simplify();
}
+
+/// Substitute x_i = t^μ_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.
+/// If the returned value is {num, dens}, it represents the function
+/// t^num / \prod_j (1 - t^dens[j]).
+/// v represents the affine functions whose floors are multiplied by the
+/// generators, and ds represents the list of generators.
+std::pair<QuasiPolynomial, std::vector<Fraction>>
+substituteMuInTerm(unsigned numParams, ParamPoint v, std::vector<Point> ds,
+ Point mu) {
+ unsigned numDims = mu.size();
+ for (const Point &d : ds)
+ assert(d.size() == numDims &&
+ "μ has to have the same number of dimensions as the generators!");
+
+ // 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) {
+ // This term in the denominator is
+ // (1 - t^dens.back())
+ dens.push_back(dotProduct(d, mu));
+ }
+
+ return {num, dens};
+}
+
+/// Normalize all denominator exponents `dens` to their absolute values
+/// by multiplying and dividing by the inverses, in a function of the form
+/// sign * t^num / prod_j (1 - t^dens[j]).
+/// Here, sign = ± 1,
+/// num is a QuasiPolynomial, and
+/// each dens[j] is a Fraction.
+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];
+ }
+ }
+
+ // If we have (1 - t^-c) in the denominator, for positive c,
+ // multiply and divide by t^c.
+ // We convert all negative-exponent terms at once; therefore
+ // we multiply and divide by t^sumNegExps.
+ // Then we get
+ // -(1 - t^c) in the denominator,
+ // increase the numerator by c, and
+ // flip the sign of the function.
+ if (numNegExps % 2 == 1)
+ sign = -sign;
+ num = num - QuasiPolynomial(num.getNumInputs(), sumNegExps);
+}
+
+/// Compute the binomial coefficients nCi for 0 ≤ i ≤ r,
+/// where n is a QuasiPolynomial.
+std::vector<QuasiPolynomial> getBinomialCoefficients(QuasiPolynomial n,
+ unsigned r) {
+ unsigned numParams = n.getNumInputs();
+ std::vector<QuasiPolynomial> coefficients;
+ coefficients.reserve(r + 1);
+ coefficients.push_back(QuasiPolynomial(numParams, 1));
+ for (unsigned j = 1; j <= r; ++j)
+ // We use the recursive formula for binomial coefficients here and below.
+ coefficients.push_back(
+ (coefficients[j - 1] * (n - QuasiPolynomial(numParams, j - 1)) /
+ Fraction(j, 1))
+ .simplify());
+ return coefficients;
+}
+
+/// Compute the binomial coefficients nCi for 0 ≤ i ≤ r,
+/// where n is a QuasiPolynomial.
+std::vector<Fraction> getBinomialCoefficients(Fraction n, Fraction r) {
+ std::vector<Fraction> coefficients;
+ coefficients.reserve((int64_t)floor(r));
+ coefficients.push_back(1);
+ for (unsigned j = 1; j <= r; ++j)
+ coefficients.push_back(coefficients[j - 1] * (n - (j - 1)) / (j));
+ return coefficients;
+}
+
+/// 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 on p parameters.
+/// 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 (x = (1, ..., 1)) 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.
+///
+/// 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) {
+ // Step (1) We need to find a μ 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.
+ 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];
+
+ // Compute the new exponents of (s+1) for the numerator and the
+ // denominator after substituting μ.
+ 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]).
+
+ // 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)^numExp / (\prod_j (1 - (s+1)^dens[j]))
+ // For the i'th term, we first normalize the denominator to have only
+ // positive exponents. We convert all the dens[j] to their
+ // absolute values and change the sign and exponent in the numerator.
+ normalizeDenominatorExponents(sign, numExp, dens);
+
+ // Then, using the formula for geometric series, we replace each (1 -
+ // (s+1)^(dens[j])) with
+ // (-s)(\sum_{0 ≤ k < dens[j]} (s+1)^k).
+ for (unsigned j = 0, e = dens.size(); j < e; ++j)
+ dens[j] = abs(dens[j]) - 1;
+ // Note that at this point, the semantics of `dens[j]` changes to mean
+ // a term (\sum_{0 ≤ k ≤ dens[j]} (s+1)^k). The denominator is, as before,
+ // a product of these terms.
+
+ // Since the -s are taken out, the sign changes if there is an odd number
+ // of such terms.
+ unsigned r = dens.size();
+ if (dens.size() % 2 == 1)
+ sign = -sign;
+
+ // Thus the term overall now has the form
+ // sign'_i * (s+1)^numExp /
+ // (s^r * \prod_j (\sum_{0 ≤ k < dens[j]} (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 a polynomial in s, with integral 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)^numExp / (\prod_j (\sum_k (s+1)^k)),
+ // Letting P(s) = (s+1)^numExp and Q(s) = \prod_j (...),
+ // we need to find the coefficient of s^r in P(s)/Q(s),
+ // for which we use the `getCoefficientInRationalFunction()` function.
+
+ // First, we compute the coefficients of P(s), which are binomial
+ // coefficients.
+ // We only need the first r+1 of these, as higher-order terms do not
+ // contribute to the coefficient of s^r.
+ std::vector<QuasiPolynomial> numeratorCoefficients =
+ getBinomialCoefficients(numExp, r);
+
+ // Then we compute the coefficients of each individual term in Q(s),
+ // which are (dens[i]+1) C (k+1) for 0 ≤ k ≤ dens[i].
+ std::vector<std::vector<Fraction>> eachTermDenCoefficients;
+ std::vector<Fraction> singleTermDenCoefficients;
+ eachTermDenCoefficients.reserve(r);
+ for (const Fraction &den : dens) {
+ singleTermDenCoefficients = getBinomialCoefficients(den + 1, den + 1);
+ eachTermDenCoefficients.push_back(
+ ArrayRef<Fraction>(singleTermDenCoefficients).slice(1));
+ }
+
+ // Now we find the coefficients in Q(s) itself
+ // by taking the convolution of the coefficients
+ // of all the terms.
+ std::vector<Fraction> denominatorCoefficients;
+ denominatorCoefficients = eachTermDenCoefficients[0];
+ for (unsigned j = 1, e = eachTermDenCoefficients.size(); j < e; ++j)
+ denominatorCoefficients = multiplyPolynomials(denominatorCoefficients,
+ eachTermDenCoefficients[j]);
+
+ totalTerm =
+ totalTerm + getCoefficientInRationalFunction(r, numeratorCoefficients,
+ denominatorCoefficients) *
+ QuasiPolynomial(numParams, sign);
+ }
+
+ return totalTerm.simplify();
+}
\ No newline at end of file
diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index 385d4c354be18a..4fd4886d22536d 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -97,10 +97,18 @@ QuasiPolynomial QuasiPolynomial::operator/(const Fraction x) const {
return qp;
}
-// Removes terms which evaluate to zero from the expression.
+// Removes terms which evaluate to zero from the expression and
+// integrate affine functions which are constants into the
+// coefficients.
QuasiPolynomial QuasiPolynomial::simplify() {
+ Fraction newCoeff = 0;
SmallVector<Fraction> newCoeffs({});
+
+ std::vector<SmallVector<Fraction>> newAffineTerm({});
std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
+
+ unsigned numParam = getNumInputs();
+
for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
// A term is zero if its coefficient is zero, or
if (coefficients[i] == Fraction(0, 1))
@@ -114,9 +122,46 @@ QuasiPolynomial QuasiPolynomial::simplify() {
});
if (product_is_zero)
continue;
+
+ // Now, we know the term is nonzero.
+
+ // We now eliminate the affine functions which are constant
+ // by merging them into the coefficients.
+ newAffineTerm = {};
+ newCoeff = coefficients[i];
+ for (ArrayRef<Fraction> term : affine[i]) {
+ bool allCoeffsZero = llvm::all_of(
+ term.slice(0, numParam), [](const Fraction c) { return c == 0; });
+ if (allCoeffsZero)
+ newCoeff *= term[numParam];
+ else
+ newAffineTerm.push_back(SmallVector<Fraction>(term));
+ }
+
+ newCoeffs.push_back(newCoeff);
+ newAffine.push_back(newAffineTerm);
+ }
+ return QuasiPolynomial(getNumInputs(), newCoeffs, newAffine);
+}
+
+QuasiPolynomial QuasiPolynomial::collectTerms() {
+ SmallVector<Fraction> newCoeffs({});
+ std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
+
+ for (unsigned i = 0, e = affine.size(); i < e; i++) {
+ bool alreadyPresent = false;
+ for (unsigned j = 0, f = newAffine.size(); j < f; j++) {
+ if (affine[i] == newAffine[j]) {
+ newCoeffs[j] += coefficients[i];
+ alreadyPresent = true;
+ }
+ }
+ if (alreadyPresent)
+ continue;
newCoeffs.push_back(coefficients[i]);
newAffine.push_back(affine[i]);
}
+
return QuasiPolynomial(getNumInputs(), newCoeffs, newAffine);
}
diff --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp
index 5e267d2045bc1c..a8d860885ef106 100644
--- a/mlir/lib/Analysis/Presburger/Utils.cpp
+++ b/mlir/lib/Analysis/Presburger/Utils.cpp
@@ -537,4 +537,31 @@ Fraction presburger::dotProduct(ArrayRef<Fraction> a, ArrayRef<Fraction> b) {
for (unsigned i = 0, e = a.size(); i < e; i++)
sum += a[i] * b[i];
return sum;
+}
+
+/// Find the product of two polynomials, each given by an array of
+/// coefficients, by taking the convolution.
+std::vector<Fraction> presburger::multiplyPolynomials(ArrayRef<Fraction> a,
+ ArrayRef<Fraction> b) {
+ // The length of the convolution is the sum of the lengths
+ // of the two sequences. We pad the shorter one with zeroes.
+ unsigned len = a.size() + b.size() - 1;
+
+ // We define accessors to avoid out-of-bounds errors.
+ auto getCoeff = [](ArrayRef<Fraction> arr, unsigned i) -> Fraction {
+ if (i < arr.size())
+ return arr[i];
+ else
+ return 0;
+ };
+
+ std::vector<Fraction> convolution;
+ convolution.reserve(len);
+ for (unsigned k = 0; k < len; ++k) {
+ Fraction sum(0, 1);
+ for (unsigned l = 0; l <= k; ++l)
+ sum += getCoeff(a, l) * getCoeff(b, k - l);
+ convolution.push_back(sum);
+ }
+ return convolution;
}
\ No newline at end of file
diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index e304f81de21f0b..919aaa7a428593 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -124,3 +124,113 @@ 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;
+ SmallVector<Fraction> coefficients = numPoints.getCoefficients();
+ for (unsigned i = 0; i < numPoints.getCoefficients().size(); i++)
+ if (numPoints.getAffine()[i].size() == 2)
+ pSquaredCoeff = pSquaredCoeff + coefficients[i];
+ else if (numPoints.getAffine()[i].size() == 1)
+ pCoeff = pCoeff + coefficients[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 M (x-axis), width N (y-axis), and height P (z-axis).
+ // 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, P, or 1.
+ for (const std::vector<SmallVector<Fraction>> &term : numPoints.getAffine()) {
+ for (const SmallVector<Fraction> &aff : term) {
+ // First, ensure that the coefficients are all nonnegative integers.
+ for (const Fraction &c : aff) {
+ EXPECT_TRUE(c >= 0);
+ EXPECT_EQ(c, c.getAsInteger());
+ }
+ // Now, if the coefficients add up to 1, we can be sure the term is
+ // either M, N, P, or 1.
+ EXPECT_EQ(aff[0] + aff[1] + aff[2] + aff[3], 1);
+ }
+ }
+
+ // We store the coefficients of M, N and P in this array.
+ Fraction count[2][2][2];
+ coefficients = numPoints.getCoefficients();
+ for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
+ unsigned mIndex = 0, nIndex = 0, pIndex = 0;
+ for (const SmallVector<Fraction> &aff : numPoints.getAffine()[i]) {
+ if (aff[0] == 1)
+ mIndex = 1;
+ if (aff[1] == 1)
+ nIndex = 1;
+ if (aff[2] == 1)
+ pIndex = 1;
+ EXPECT_EQ(aff[3], 0);
+ }
+ count[mIndex][nIndex][pIndex] += coefficients[i];
+ }
+
+ // We expect the answer to be
+ // (⌊M⌋ + 1)(⌊N⌋ + 1)(⌊P⌋ + 1) =
+ // ⌊M⌋⌊N⌋⌊P⌋ + ⌊M⌋⌊N⌋ + ⌊N⌋⌊P⌋ + ⌊M⌋⌊P⌋ + ⌊M⌋ + ⌊N⌋ + ⌊P⌋ + 1.
+ for (unsigned i = 0; i < 2; i++)
+ for (unsigned j = 0; j < 2; j++)
+ for (unsigned k = 0; k < 2; k++)
+ EXPECT_EQ(count[i][j][k], 1);
+}
\ No newline at end of file
diff --git a/mlir/unittests/Analysis/Presburger/UtilsTest.cpp b/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
index 7c1646aa841162..f09a1a760ce609 100644
--- a/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/UtilsTest.cpp
@@ -66,3 +66,22 @@ TEST(UtilsTest, DivisionReprNormalizeTest) {
checkEqual(a, b);
checkEqual(c, d);
}
+
+TEST(UtilsTest, convolution) {
+ std::vector<Fraction> aVals({1, 2, 3, 4});
+ std::vector<Fraction> bVals({7, 3, 1, 6});
+ ArrayRef<Fraction> a(aVals);
+ ArrayRef<Fraction> b(bVals);
+
+ std::vector<Fraction> conv = multiplyPolynomials(a, b);
+
+ EXPECT_EQ(conv, std::vector<Fraction>({7, 17, 28, 45, 27, 22, 24}));
+
+ aVals = {3, 6, 0, 2, 5};
+ bVals = {2, 0, 6};
+ a = aVals;
+ b = bVals;
+
+ conv = multiplyPolynomials(a, b);
+ EXPECT_EQ(conv, std::vector<Fraction>({6, 12, 18, 40, 10, 12, 30}));
+}
More information about the Mlir-commits
mailing list