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 | |
13 | using namespace mlir; |
14 | using namespace presburger; |
15 | |
16 | QuasiPolynomial::QuasiPolynomial( |
17 | unsigned numVars, ArrayRef<Fraction> coeffs, |
18 | ArrayRef<std::vector<SmallVector<Fraction>>> aff) |
19 | : PresburgerSpace(/*numDomain=*/numVars, /*numRange=*/1, /*numSymbols=*/0, |
20 | /*numLocals=*/0), |
21 | coefficients(coeffs), affine(aff) { |
22 | #ifndef NDEBUG |
23 | // For each term which involves at least one affine function, |
24 | for (const std::vector<SmallVector<Fraction>> &term : affine) { |
25 | if (term.empty()) |
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 | #endif // NDEBUG |
36 | } |
37 | |
38 | /// Define a quasipolynomial which is a single constant. |
39 | QuasiPolynomial::QuasiPolynomial(unsigned numVars, const Fraction &constant) |
40 | : PresburgerSpace(/*numDomain=*/numVars, /*numRange=*/1, /*numSymbols=*/0, |
41 | /*numLocals=*/0), |
42 | coefficients({constant}), affine({{}}) {} |
43 | |
44 | QuasiPolynomial QuasiPolynomial::operator+(const QuasiPolynomial &x) const { |
45 | assert(getNumInputs() == x.getNumInputs() && |
46 | "two quasi-polynomials with different numbers of symbols cannot " |
47 | "be added!" ); |
48 | SmallVector<Fraction> sumCoeffs = coefficients; |
49 | sumCoeffs.append(RHS: x.coefficients); |
50 | std::vector<std::vector<SmallVector<Fraction>>> sumAff = affine; |
51 | llvm::append_range(C&: sumAff, R: x.affine); |
52 | return QuasiPolynomial(getNumInputs(), sumCoeffs, sumAff); |
53 | } |
54 | |
55 | QuasiPolynomial QuasiPolynomial::operator-(const QuasiPolynomial &x) const { |
56 | assert(getNumInputs() == x.getNumInputs() && |
57 | "two quasi-polynomials with different numbers of symbols cannot " |
58 | "be subtracted!" ); |
59 | QuasiPolynomial qp(getNumInputs(), x.coefficients, x.affine); |
60 | for (Fraction &coeff : qp.coefficients) |
61 | coeff = -coeff; |
62 | return *this + qp; |
63 | } |
64 | |
65 | QuasiPolynomial QuasiPolynomial::operator*(const QuasiPolynomial &x) const { |
66 | assert(getNumInputs() == x.getNumInputs() && |
67 | "two quasi-polynomials with different numbers of " |
68 | "symbols cannot be multiplied!" ); |
69 | |
70 | SmallVector<Fraction> coeffs; |
71 | coeffs.reserve(N: coefficients.size() * x.coefficients.size()); |
72 | for (const Fraction &coeff : coefficients) |
73 | for (const Fraction &xcoeff : x.coefficients) |
74 | coeffs.emplace_back(Args: coeff * xcoeff); |
75 | |
76 | std::vector<SmallVector<Fraction>> product; |
77 | std::vector<std::vector<SmallVector<Fraction>>> aff; |
78 | aff.reserve(n: affine.size() * x.affine.size()); |
79 | for (const std::vector<SmallVector<Fraction>> &term : affine) { |
80 | for (const std::vector<SmallVector<Fraction>> &xterm : x.affine) { |
81 | product.clear(); |
82 | llvm::append_range(C&: product, R: term); |
83 | llvm::append_range(C&: product, R: xterm); |
84 | aff.emplace_back(args&: product); |
85 | } |
86 | } |
87 | |
88 | return QuasiPolynomial(getNumInputs(), coeffs, aff); |
89 | } |
90 | |
91 | QuasiPolynomial QuasiPolynomial::operator/(const Fraction &x) const { |
92 | assert(x != 0 && "division by zero!" ); |
93 | QuasiPolynomial qp(*this); |
94 | for (Fraction &coeff : qp.coefficients) |
95 | coeff /= x; |
96 | return qp; |
97 | } |
98 | |
99 | // Removes terms which evaluate to zero from the expression and |
100 | // integrate affine functions which are constants into the |
101 | // coefficients. |
102 | QuasiPolynomial QuasiPolynomial::simplify() { |
103 | Fraction newCoeff = 0; |
104 | SmallVector<Fraction> newCoeffs({}); |
105 | |
106 | std::vector<SmallVector<Fraction>> newAffineTerm({}); |
107 | std::vector<std::vector<SmallVector<Fraction>>> newAffine({}); |
108 | |
109 | unsigned numParam = getNumInputs(); |
110 | |
111 | for (unsigned i = 0, e = coefficients.size(); i < e; i++) { |
112 | // A term is zero if its coefficient is zero, or |
113 | if (coefficients[i] == Fraction(0, 1)) |
114 | continue; |
115 | bool product_is_zero = |
116 | // if any of the affine functions in the product |
117 | llvm::any_of(Range&: affine[i], P: [](const SmallVector<Fraction> &affine_ij) { |
118 | // has all its coefficients as zero. |
119 | return llvm::all_of(Range: affine_ij, |
120 | P: [](const Fraction &f) { return f == 0; }); |
121 | }); |
122 | if (product_is_zero) |
123 | continue; |
124 | |
125 | // Now, we know the term is nonzero. |
126 | |
127 | // We now eliminate the affine functions which are constant |
128 | // by merging them into the coefficients. |
129 | newAffineTerm = {}; |
130 | newCoeff = coefficients[i]; |
131 | for (ArrayRef<Fraction> term : affine[i]) { |
132 | bool allCoeffsZero = llvm::all_of( |
133 | Range: term.slice(N: 0, M: numParam), P: [](const Fraction &c) { return c == 0; }); |
134 | if (allCoeffsZero) |
135 | newCoeff *= term[numParam]; |
136 | else |
137 | newAffineTerm.emplace_back(args&: term); |
138 | } |
139 | |
140 | newCoeffs.emplace_back(Args&: newCoeff); |
141 | newAffine.emplace_back(args&: newAffineTerm); |
142 | } |
143 | return QuasiPolynomial(getNumInputs(), newCoeffs, newAffine); |
144 | } |
145 | |
146 | QuasiPolynomial QuasiPolynomial::collectTerms() { |
147 | SmallVector<Fraction> newCoeffs({}); |
148 | std::vector<std::vector<SmallVector<Fraction>>> newAffine({}); |
149 | |
150 | for (unsigned i = 0, e = affine.size(); i < e; i++) { |
151 | bool alreadyPresent = false; |
152 | for (unsigned j = 0, f = newAffine.size(); j < f; j++) { |
153 | if (affine[i] == newAffine[j]) { |
154 | newCoeffs[j] += coefficients[i]; |
155 | alreadyPresent = true; |
156 | } |
157 | } |
158 | if (alreadyPresent) |
159 | continue; |
160 | newCoeffs.emplace_back(Args&: coefficients[i]); |
161 | newAffine.emplace_back(args&: affine[i]); |
162 | } |
163 | |
164 | return QuasiPolynomial(getNumInputs(), newCoeffs, newAffine); |
165 | } |
166 | |
167 | Fraction QuasiPolynomial::getConstantTerm() { |
168 | Fraction constTerm = 0; |
169 | for (unsigned i = 0, e = coefficients.size(); i < e; ++i) |
170 | if (affine[i].empty()) |
171 | constTerm += coefficients[i]; |
172 | return constTerm; |
173 | } |
174 | |