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
13using namespace mlir;
14using namespace presburger;
15
16QuasiPolynomial::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.
39QuasiPolynomial::QuasiPolynomial(unsigned numVars, const Fraction &constant)
40 : PresburgerSpace(/*numDomain=*/numVars, /*numRange=*/1, /*numSymbols=*/0,
41 /*numLocals=*/0),
42 coefficients({constant}), affine({{}}) {}
43
44QuasiPolynomial 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
55QuasiPolynomial 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
65QuasiPolynomial 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
91QuasiPolynomial 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.
102QuasiPolynomial 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
146QuasiPolynomial 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
167Fraction 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

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