Skip to content

Commit 1022feb

Browse files
Abhinav271828Abhinav271828
authored
[MLIR][Presburger] Generating functions and quasi-polynomials for Barvinok's algorithm (#75702)
Define basic types and classes for Barvinok's algorithm, including polyhedra, generating functions and quasi-polynomials. The class definitions include methods for arithmetic manipulation, printing, logical relations, etc.
1 parent 01bf29b commit 1022feb

File tree

7 files changed

+481
-0
lines changed

7 files changed

+481
-0
lines changed

‎mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
//===- QuasiPolynomial.h - QuasiPolynomial Class ----------------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// Definition of the QuasiPolynomial class for Barvinok's algorithm,
10+
// which represents a single-valued function on a set of parameters.
11+
// It is an expression of the form
12+
// f(x) = \sum_i c_i * \prod_j ⌊g_{ij}(x)⌋
13+
// where c_i \in Q and
14+
// g_{ij} : Q^d -> Q are affine functionals over d parameters.
15+
//
16+
//===----------------------------------------------------------------------===//
17+
18+
#ifndef MLIR_ANALYSIS_PRESBURGER_QUASIPOLYNOMIAL_H
19+
#define MLIR_ANALYSIS_PRESBURGER_QUASIPOLYNOMIAL_H
20+
21+
#include "mlir/Analysis/Presburger/Fraction.h"
22+
#include "mlir/Analysis/Presburger/PresburgerSpace.h"
23+
24+
namespace mlir {
25+
namespace presburger {
26+
27+
// A class to describe quasi-polynomials.
28+
// A quasipolynomial consists of a set of terms.
29+
// The ith term is a constant `coefficients[i]`, multiplied
30+
// by the product of a set of affine functions on n parameters.
31+
// Represents functions f : Q^n -> Q of the form
32+
//
33+
// f(x) = \sum_i c_i * \prod_j ⌊g_{ij}(x)⌋
34+
//
35+
// where c_i \in Q and
36+
// g_{ij} : Q^n -> Q are affine functionals.
37+
class QuasiPolynomial : public PresburgerSpace {
38+
public:
39+
QuasiPolynomial(unsigned numVars, SmallVector<Fraction> coeffs = {},
40+
std::vector<std::vector<SmallVector<Fraction>>> aff = {});
41+
42+
// Find the number of inputs (numDomain) to the polynomial.
43+
// numSymbols is set to zero.
44+
unsigned getNumInputs() const {
45+
return getNumDomainVars() + getNumSymbolVars();
46+
}
47+
48+
const SmallVector<Fraction> &getCoefficients() const { return coefficients; }
49+
50+
const std::vector<std::vector<SmallVector<Fraction>>> &getAffine() const {
51+
return affine;
52+
}
53+
54+
// Arithmetic operations.
55+
QuasiPolynomial operator+(const QuasiPolynomial &x) const;
56+
QuasiPolynomial operator-(const QuasiPolynomial &x) const;
57+
QuasiPolynomial operator*(const QuasiPolynomial &x) const;
58+
QuasiPolynomial operator/(const Fraction x) const;
59+
60+
// Removes terms which evaluate to zero from the expression.
61+
QuasiPolynomial simplify();
62+
63+
private:
64+
SmallVector<Fraction> coefficients;
65+
std::vector<std::vector<SmallVector<Fraction>>> affine;
66+
};
67+
68+
} // namespace presburger
69+
} // namespace mlir
70+
71+
#endif // MLIR_ANALYSIS_PRESBURGER_QUASIPOLYNOMIAL_H

‎mlir/lib/Analysis/Presburger/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ add_mlir_library(MLIRPresburger
66
PresburgerRelation.cpp
77
PresburgerSpace.cpp
88
PWMAFunction.cpp
9+
QuasiPolynomial.cpp
910
Simplex.cpp
1011
SlowMPInt.cpp
1112
Utils.cpp

‎mlir/lib/Analysis/Presburger/GeneratingFunction.h

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
//===- GeneratingFunction.h - Generating Functions over Q^d -----*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// Definition of the GeneratingFunction class for Barvinok's algorithm,
10+
// which represents a function over Q^n, parameterized by d parameters.
11+
//
12+
//===----------------------------------------------------------------------===//
13+
14+
#ifndef MLIR_ANALYSIS_PRESBURGER_GENERATINGFUNCTION_H
15+
#define MLIR_ANALYSIS_PRESBURGER_GENERATINGFUNCTION_H
16+
17+
#include "mlir/Analysis/Presburger/Fraction.h"
18+
#include "mlir/Analysis/Presburger/Matrix.h"
19+
20+
namespace mlir {
21+
namespace presburger {
22+
23+
// A parametric point is a vector, each of whose elements
24+
// is an affine function of n parameters. Each row
25+
// in the matrix represents the affine function and
26+
// has n+1 elements.
27+
using ParamPoint = FracMatrix;
28+
29+
// A point is simply a vector.
30+
using Point = SmallVector<Fraction>;
31+
32+
// A class to describe the type of generating function
33+
// used to enumerate the integer points in a polytope.
34+
// Consists of a set of terms, where the ith term has
35+
// * a sign, ±1, stored in `signs[i]`
36+
// * a numerator, of the form x^{n},
37+
// where n, stored in `numerators[i]`,
38+
// is a parametric point.
39+
// * a denominator, of the form (1 - x^{d1})...(1 - x^{dn}),
40+
// where each dj, stored in `denominators[i][j]`,
41+
// is a vector.
42+
//
43+
// Represents functions f_p : Q^n -> Q of the form
44+
//
45+
// f_p(x) = \sum_i s_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
46+
//
47+
// where s_i is ±1,
48+
// n_i \in Q^d -> Q^n is an n-vector of affine functions on d parameters, and
49+
// g_{ij} \in Q^n are vectors.
50+
class GeneratingFunction {
51+
public:
52+
GeneratingFunction(unsigned numParam, SmallVector<int, 8> signs,
53+
std::vector<ParamPoint> nums,
54+
std::vector<std::vector<Point>> dens)
55+
: numParam(numParam), signs(signs), numerators(nums), denominators(dens) {
56+
for (const ParamPoint &term : numerators)
57+
assert(term.getNumColumns() == numParam + 1 &&
58+
"dimensionality of numerator exponents does not match number of "
59+
"parameters!");
60+
}
61+
62+
unsigned getNumParams() { return numParam; }
63+
64+
SmallVector<int> getSigns() { return signs; }
65+
66+
std::vector<ParamPoint> getNumerators() { return numerators; }
67+
68+
std::vector<std::vector<Point>> getDenominators() { return denominators; }
69+
70+
GeneratingFunction operator+(const GeneratingFunction &gf) const {
71+
assert(numParam == gf.getNumParams() &&
72+
"two generating functions with different numbers of parameters "
73+
"cannot be added!");
74+
SmallVector<int> sumSigns = signs;
75+
sumSigns.append(gf.signs);
76+
77+
std::vector<ParamPoint> sumNumerators = numerators;
78+
sumNumerators.insert(sumNumerators.end(), gf.numerators.begin(),
79+
gf.numerators.end());
80+
81+
std::vector<std::vector<Point>> sumDenominators = denominators;
82+
sumDenominators.insert(sumDenominators.end(), gf.denominators.begin(),
83+
gf.denominators.end());
84+
return GeneratingFunction(sumSigns, sumNumerators, sumDenominators);
85+
}
86+
87+
llvm::raw_ostream &print(llvm::raw_ostream &os) const {
88+
for (unsigned i = 0, e = signs.size(); i < e; i++) {
89+
if (i == 0) {
90+
if (signs[i] == -1)
91+
os << "- ";
92+
} else {
93+
if (signs[i] == 1)
94+
os << " + ";
95+
else
96+
os << " - ";
97+
}
98+
99+
os << "x^[";
100+
unsigned r = numerators[i].getNumRows();
101+
for (unsigned j = 0; j < r - 1; j++) {
102+
os << "[";
103+
for (unsigned k = 0, c = numerators[i].getNumColumns(); k < c - 1; k++)
104+
os << numerators[i].at(j, k) << ",";
105+
os << numerators[i].getRow(j).back() << "],";
106+
}
107+
os << "[";
108+
for (unsigned k = 0, c = numerators[i].getNumColumns(); k < c - 1; k++)
109+
os << numerators[i].at(r - 1, k) << ",";
110+
os << numerators[i].getRow(r - 1).back() << "]]/";
111+
112+
for (const Point &den : denominators[i]) {
113+
os << "(x^[";
114+
for (unsigned j = 0, e = den.size(); j < e - 1; j++)
115+
os << den[j] << ",";
116+
os << den.back() << "])";
117+
}
118+
}
119+
return os;
120+
}
121+
122+
private:
123+
unsigned numParam;
124+
SmallVector<int, 8> signs;
125+
std::vector<ParamPoint> numerators;
126+
std::vector<std::vector<Point>> denominators;
127+
};
128+
129+
} // namespace presburger
130+
} // namespace mlir
131+
132+
#endif // MLIR_ANALYSIS_PRESBURGER_GENERATINGFUNCTION_H

‎mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
//===- QuasiPolynomial.cpp - Quasipolynomial Class --------------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "mlir/Analysis/Presburger/QuasiPolynomial.h"
10+
#include "mlir/Analysis/Presburger/Fraction.h"
11+
#include "mlir/Analysis/Presburger/PresburgerSpace.h"
12+
#include "mlir/Analysis/Presburger/Utils.h"
13+
14+
using namespace mlir;
15+
using namespace presburger;
16+
17+
QuasiPolynomial::QuasiPolynomial(
18+
unsigned numVars, SmallVector<Fraction> coeffs,
19+
std::vector<std::vector<SmallVector<Fraction>>> aff)
20+
: PresburgerSpace(/*numDomain=*/numVars, /*numRange=*/1, /*numSymbols=*/0,
21+
/*numLocals=*/0),
22+
coefficients(coeffs), affine(aff) {
23+
// For each term which involves at least one affine function,
24+
for (const std::vector<SmallVector<Fraction>> &term : affine) {
25+
if (term.size() == 0)
26+
continue;
27+
// the number of elements in each affine function is
28+
// one more than the number of symbols.
29+
for (const SmallVector<Fraction> &aff : term) {
30+
assert(aff.size() == getNumInputs() + 1 &&
31+
"dimensionality of affine functions does not match number of "
32+
"symbols!");
33+
}
34+
}
35+
}
36+
37+
QuasiPolynomial QuasiPolynomial::operator+(const QuasiPolynomial &x) const {
38+
assert(getNumInputs() == x.getNumInputs() &&
39+
"two quasi-polynomials with different numbers of symbols cannot "
40+
"be added!");
41+
SmallVector<Fraction> sumCoeffs = coefficients;
42+
sumCoeffs.append(x.coefficients);
43+
std::vector<std::vector<SmallVector<Fraction>>> sumAff = affine;
44+
sumAff.insert(sumAff.end(), x.affine.begin(), x.affine.end());
45+
return QuasiPolynomial(getNumInputs(), sumCoeffs, sumAff);
46+
}
47+
48+
QuasiPolynomial QuasiPolynomial::operator-(const QuasiPolynomial &x) const {
49+
assert(getNumInputs() == x.getNumInputs() &&
50+
"two quasi-polynomials with different numbers of symbols cannot "
51+
"be subtracted!");
52+
QuasiPolynomial qp(getNumInputs(), x.coefficients, x.affine);
53+
for (Fraction &coeff : qp.coefficients)
54+
coeff = -coeff;
55+
return *this + qp;
56+
}
57+
58+
QuasiPolynomial QuasiPolynomial::operator*(const QuasiPolynomial &x) const {
59+
assert(getNumInputs() == x.getNumInputs() &&
60+
"two quasi-polynomials with different numbers of "
61+
"symbols cannot be multiplied!");
62+
63+
SmallVector<Fraction> coeffs;
64+
coeffs.reserve(coefficients.size() * x.coefficients.size());
65+
for (const Fraction &coeff : coefficients)
66+
for (const Fraction &xcoeff : x.coefficients)
67+
coeffs.push_back(coeff * xcoeff);
68+
69+
std::vector<SmallVector<Fraction>> product;
70+
std::vector<std::vector<SmallVector<Fraction>>> aff;
71+
aff.reserve(affine.size() * x.affine.size());
72+
for (const std::vector<SmallVector<Fraction>> &term : affine) {
73+
for (const std::vector<SmallVector<Fraction>> &xterm : x.affine) {
74+
product.clear();
75+
product.insert(product.end(), term.begin(), term.end());
76+
product.insert(product.end(), xterm.begin(), xterm.end());
77+
aff.push_back(product);
78+
}
79+
}
80+
81+
return QuasiPolynomial(getNumInputs(), coeffs, aff);
82+
}
83+
84+
QuasiPolynomial QuasiPolynomial::operator/(const Fraction x) const {
85+
assert(x != 0 && "division by zero!");
86+
QuasiPolynomial qp(*this);
87+
for (Fraction &coeff : qp.coefficients)
88+
coeff /= x;
89+
return qp;
90+
}
91+
92+
// Removes terms which evaluate to zero from the expression.
93+
QuasiPolynomial QuasiPolynomial::simplify() {
94+
SmallVector<Fraction> newCoeffs({});
95+
std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
96+
for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
97+
// A term is zero if its coefficient is zero, or
98+
if (coefficients[i] == Fraction(0, 1))
99+
continue;
100+
bool product_is_zero =
101+
// if any of the affine functions in the product
102+
llvm::any_of(affine[i], [](const SmallVector<Fraction> &affine_ij) {
103+
// has all its coefficients as zero.
104+
return llvm::all_of(affine_ij,
105+
[](const Fraction &f) { return f == 0; });
106+
});
107+
if (product_is_zero)
108+
continue;
109+
newCoeffs.push_back(coefficients[i]);
110+
newAffine.push_back(affine[i]);
111+
}
112+
return QuasiPolynomial(getNumInputs(), newCoeffs, newAffine);
113+
}

‎mlir/unittests/Analysis/Presburger/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ add_mlir_unittest(MLIRPresburgerTests
1111
PresburgerRelationTest.cpp
1212
PresburgerSpaceTest.cpp
1313
PWMAFunctionTest.cpp
14+
QuasiPolynomialTest.cpp
1415
SimplexTest.cpp
1516
UtilsTest.cpp
1617
)

0 commit comments

Comments
 (0)