[Mlir-commits] [mlir] [MLIR][Presburger] Generating functions and quasi-polynomials for Barvinok's algorithm (PR #75702)

[Arjun P]

================
@@ -0,0 +1,117 @@
+//===- QuasiPolynomial.cpp - Quasipolynomial Class --------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/Presburger/QuasiPolynomial.h"
+#include "mlir/Analysis/Presburger/Fraction.h"
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
+#include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Analysis/Presburger/PresburgerSpace.h"
+#include "mlir/Analysis/Presburger/Utils.h"
+#include "mlir/Support/LogicalResult.h"
+#include <optional>
+
+using namespace mlir;
+using namespace presburger;
+
+QuasiPolynomial::QuasiPolynomial(
+    unsigned numParam, SmallVector<Fraction> coeffs,
+    std::vector<std::vector<SmallVector<Fraction>>> aff)
+    : numParam(numParam), coefficients(coeffs), affine(aff) {
+  // Find the first term which involves some affine function.
+  for (const std::vector<SmallVector<Fraction>> &term : affine) {
+    if (term.size() == 0)
+      continue;
+    // The number of elements in each affine function is
+    // one more than the number of parameters.
+    for (const SmallVector<Fraction> &aff : term) {
+      assert(aff.size() - 1 == numParam &&
+             "dimensionality of affine functions does not match number of "
+             "parameters!");
+    }
+  }
+}
+
+QuasiPolynomial QuasiPolynomial::operator+(const QuasiPolynomial &x) const {
+  assert(numParam == x.getNumParams() &&
+         "two quasi-polynomials with different numbers of parameters cannot "
+         "be added!");
+  SmallVector<Fraction> sumCoeffs = coefficients;
+  sumCoeffs.append(x.coefficients);
+  std::vector<std::vector<SmallVector<Fraction>>> sumAff = affine;
+  sumAff.insert(sumAff.end(), x.affine.begin(), x.affine.end());
+  return QuasiPolynomial(numParam, sumCoeffs, sumAff);
+}
+
+QuasiPolynomial QuasiPolynomial::operator-(const QuasiPolynomial &x) const {
+  assert(numParam == x.getNumParams() &&
+         "two quasi-polynomials with different numbers of parameters cannot "
+         "be subtracted!");
+  QuasiPolynomial qp(numParam, x.coefficients, x.affine);
+  for (Fraction &coeff : qp.coefficients)
+    coeff = -coeff;
+  return *this + qp;
+}
+
+QuasiPolynomial QuasiPolynomial::operator*(const QuasiPolynomial &x) const {
+  assert(numParam == x.getNumParams() &&
+         "two quasi-polynomials with different numbers of "
+         "parameters cannot be multiplied!");
+
+  SmallVector<Fraction> coeffs;
+  coeffs.reserve(coefficients.size() * x.coefficients.size());
+  for (const Fraction &coeff : coefficients) {
+    for (const Fraction &xcoeff : x.coefficients) {
+      coeffs.push_back(coeff * xcoeff);
+    }
+  }
+
+  std::vector<SmallVector<Fraction>> product;
+  std::vector<std::vector<SmallVector<Fraction>>> aff;
+  aff.reserve(affine.size() * x.affine.size());
+  for (const std::vector<SmallVector<Fraction>> &term : affine) {
----------------
Superty wrote:

...but not in this case.

https://github.com/llvm/llvm-project/pull/75702