xref: /llvm-project/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h (revision 266a5a9cb9daa96c1eeaebc18e10f5a37d638734)
1 //===- QuasiPolynomial.h - 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 // Definition of the QuasiPolynomial class for Barvinok's algorithm,
10 // which represents a single-valued function on a set of parameters.
11 // It is an expression of the form
12 // f(x) = \sum_i c_i * \prod_j ⌊g_{ij}(x)⌋
13 // where c_i \in Q and
14 // g_{ij} : Q^d -> Q are affine functionals over d parameters.
15 //
16 //===----------------------------------------------------------------------===//
17 
18 #ifndef MLIR_ANALYSIS_PRESBURGER_QUASIPOLYNOMIAL_H
19 #define MLIR_ANALYSIS_PRESBURGER_QUASIPOLYNOMIAL_H
20 
21 #include "mlir/Analysis/Presburger/Fraction.h"
22 #include "mlir/Analysis/Presburger/PresburgerSpace.h"
23 
24 namespace mlir {
25 namespace presburger {
26 
27 // A class to describe quasi-polynomials.
28 // A quasipolynomial consists of a set of terms.
29 // The ith term is a constant `coefficients[i]`, multiplied
30 // by the product of a set of affine functions on n parameters.
31 // Represents functions f : Q^n -> Q of the form
32 //
33 // f(x) = \sum_i c_i * \prod_j ⌊g_{ij}(x)⌋
34 //
35 // where c_i \in Q and
36 // g_{ij} : Q^n -> Q are affine functionals.
37 class QuasiPolynomial : public PresburgerSpace {
38 public:
39   QuasiPolynomial(unsigned numVars, ArrayRef<Fraction> coeffs = {},
40                   ArrayRef<std::vector<SmallVector<Fraction>>> aff = {});
41 
42   QuasiPolynomial(unsigned numVars, const Fraction &constant);
43 
44   // Find the number of inputs (numDomain) to the polynomial.
45   // numSymbols is set to zero.
46   unsigned getNumInputs() const {
47     return getNumDomainVars() + getNumSymbolVars();
48   }
49 
50   const SmallVector<Fraction> &getCoefficients() const { return coefficients; }
51 
52   const std::vector<std::vector<SmallVector<Fraction>>> &getAffine() const {
53     return affine;
54   }
55 
56   // Arithmetic operations.
57   QuasiPolynomial operator+(const QuasiPolynomial &x) const;
58   QuasiPolynomial operator-(const QuasiPolynomial &x) const;
59   QuasiPolynomial operator*(const QuasiPolynomial &x) const;
60   QuasiPolynomial operator/(const Fraction &x) const;
61 
62   // Removes terms which evaluate to zero from the expression
63   // and folds affine functions which are constant into the
64   // constant coefficients.
65   QuasiPolynomial simplify();
66 
67   // Group together like terms in the expression.
68   QuasiPolynomial collectTerms();
69 
70   Fraction getConstantTerm();
71 
72 private:
73   SmallVector<Fraction> coefficients;
74   std::vector<std::vector<SmallVector<Fraction>>> affine;
75 };
76 
77 } // namespace presburger
78 } // namespace mlir
79 
80 #endif // MLIR_ANALYSIS_PRESBURGER_QUASIPOLYNOMIAL_H
81