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(x.coefficients); 50 std::vector<std::vector<SmallVector<Fraction>>> sumAff = affine; 51 sumAff.insert(sumAff.end(), x.affine.begin(), x.affine.end()); 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(coefficients.size() * x.coefficients.size()); 72 for (const Fraction &coeff : coefficients) 73 for (const Fraction &xcoeff : x.coefficients) 74 coeffs.emplace_back(coeff * xcoeff); 75 76 std::vector<SmallVector<Fraction>> product; 77 std::vector<std::vector<SmallVector<Fraction>>> aff; 78 aff.reserve(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 product.insert(product.end(), term.begin(), term.end()); 83 product.insert(product.end(), xterm.begin(), xterm.end()); 84 aff.emplace_back(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(affine[i], [](const SmallVector<Fraction> &affine_ij) { 118 // has all its coefficients as zero. 119 return llvm::all_of(affine_ij, 120 [](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 term.slice(0, numParam), [](const Fraction &c) { return c == 0; }); 134 if (allCoeffsZero) 135 newCoeff *= term[numParam]; 136 else 137 newAffineTerm.emplace_back(term); 138 } 139 140 newCoeffs.emplace_back(newCoeff); 141 newAffine.emplace_back(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(coefficients[i]); 161 newAffine.emplace_back(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