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

source code of mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp