174 lines · cpp
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-exception6//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 NDEBUG23 // 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 is28 // 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 // NDEBUG36}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(x.coefficients);50 std::vector<std::vector<SmallVector<Fraction>>> sumAff = affine;51 llvm::append_range(sumAff, 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(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 llvm::append_range(product, term);83 llvm::append_range(product, xterm);84 aff.emplace_back(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 and100// integrate affine functions which are constants into the101// 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, or113 if (coefficients[i] == Fraction(0, 1))114 continue;115 bool productIsZero =116 // if any of the affine functions in the product117 llvm::any_of(affine[i], [](const SmallVector<Fraction> &affineIj) {118 // has all its coefficients as zero.119 return llvm::all_of(affineIj,120 [](const Fraction &f) { return f == 0; });121 });122 if (productIsZero)123 continue;124 125 // Now, we know the term is nonzero.126 127 // We now eliminate the affine functions which are constant128 // 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 else137 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 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(coefficients[i]);161 newAffine.emplace_back(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