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

[llvmlistbot at llvm.org]

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

>From b65a078edbbd8d711b1bb9ec12b938f47c550e9a Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sat, 16 Dec 2023 19:37:00 +0530
Subject: [PATCH 01/38] Definitions for gf and qp classes

---
 .../mlir/Analysis/Presburger/Barvinok.h       | 229 ++++++++++++++++++
 1 file changed, 229 insertions(+)
 create mode 100644 mlir/include/mlir/Analysis/Presburger/Barvinok.h

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
new file mode 100644
index 00000000000000..3ee9d5dc452d0e
--- /dev/null
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -0,0 +1,229 @@
+//===- Barvinok.h - Barvinok's Algorithm -----------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Functions and classes for Barvinok's algorithm in MLIR.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_BARVINOK_H
+#define MLIR_ANALYSIS_PRESBURGER_BARVINOK_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>
+
+namespace mlir {
+namespace presburger {
+
+// The H (inequality) representation of both general
+// polyhedra and cones specifically is an integer relation.
+using PolyhedronH = IntegerRelation;
+using ConeH = PolyhedronH;
+
+// The V (generator) representation of both general
+// polyhedra and cones specifically is simply a matrix
+// whose rows are the generators.
+using PolyhedronV = Matrix<MPInt>;
+using ConeV = PolyhedronV;
+
+// A parametric point is a vector, each of whose elements
+// is an affine function of n parameters. Each row
+// in the matrix represents the affine function and
+// has n+1 elements.
+using ParamPoint = Matrix<Fraction>;
+
+// A point is simply a vector.
+using Point = SmallVector<Fraction>;
+
+// A class to describe the type of generating function
+// used to enumerate the integer points in a polytope.
+// Consists of a set of terms, where the ith term has
+// * a sign, ±1, stored in `signs[i]`
+// * a numerator, of the form x^{n},
+//      where n, stored in `numerators[i]`,
+//      is a parametric point (a vertex).
+// * a denominator, of the form (1 - x^{d1})...(1 - x^{dn}),
+//      where each dj, stored in `denominators[i][j]`,
+//      is a vector (a generator).
+class GeneratingFunction {
+public:
+  GeneratingFunction(SmallVector<int, 16> s, std::vector<ParamPoint> nums,
+                     std::vector<std::vector<Point>> dens)
+      : signs(s), numerators(nums), denominators(dens){};
+
+  unsigned getNumParams() {
+    for (auto term : numerators)
+      return (term.getNumColumns() - 1);
+    return -1;
+  }
+
+  bool operator==(const GeneratingFunction &gf) const {
+    if (signs != gf.signs || numerators != gf.numerators ||
+        denominators != gf.denominators)
+      return false;
+    return true;
+  }
+
+  GeneratingFunction operator+(const GeneratingFunction &gf) {
+    bool sameNumParams = (getNumParams() == -1) || (gf.getNumParams() == -1) ||
+                         (getNumParams() == gf.getNumParams());
+    assert(
+        sameNumParams &&
+        "two generators with different numbers of parameters cannot be added!");
+    signs.insert(signs.end(), gf.signs.begin(), gf.signs.end());
+    numerators.insert(numerators.end(), gf.numerators.begin(),
+                      gf.numerators.end());
+    denominators.insert(denominators.end(), gf.denominators.begin(),
+                        gf.denominators.end());
+    return *this;
+  }
+
+  llvm::raw_ostream &print(llvm::raw_ostream &os) const {
+    for (unsigned i = 0; i < signs.size(); i++) {
+      if (signs[i] == 1)
+        os << "+";
+      else
+        os << "-";
+
+      os << "*(x^[";
+      for (unsigned j = 0; j < numerators[i].size() - 1; j++)
+        os << numerators[i][j] << ",";
+      os << numerators[i][numerators[i].size() - 1] << "])/";
+
+      for (Point den : denominators[i]) {
+        os << "(x^[";
+        for (unsigned j = 0; j < den.size(); j++)
+          os << den[j] << ",";
+        os << den[den.size() - 1] << "])";
+      }
+    }
+    return os;
+  }
+
+  SmallVector<int, 16> signs;
+  std::vector<ParamPoint> numerators;
+  std::vector<std::vector<Point>> denominators;
+};
+
+// A class to describe the quasi-polynomials obtained by
+// substituting the unit vector in the type of generating
+// function described above.
+// Consists of a set of terms.
+// The ith term is a constant `coefficients[i]`, multiplied
+// by the product of a set of affine functions on n parameters.
+class QuasiPolynomial {
+public:
+  QuasiPolynomial(SmallVector<Fraction> coeffs = {},
+                  std::vector<std::vector<SmallVector<Fraction>>> aff = {})
+      : coefficients(coeffs), affine(aff){};
+
+  QuasiPolynomial(Fraction cons) : coefficients({cons}), affine({{}}){};
+
+  QuasiPolynomial(QuasiPolynomial const &) = default;
+
+  SmallVector<Fraction> coefficients;
+  std::vector<std::vector<SmallVector<Fraction>>> affine;
+
+  // Find the number of parameters involved in the polynomial
+  // from the dimensionality of the affine functions.
+  unsigned getNumParams() {
+    // Find the first term which involves some affine function.
+    for (auto term : affine) {
+      if (term.size() == 0)
+        continue;
+      // The number of elements in the affine function is
+      // one more than the number of parameters.
+      return (term[0].size() - 1);
+    }
+    // The polynomial can be treated as having any number
+    // of parameters.
+    return -1;
+  }
+
+  QuasiPolynomial operator+(const QuasiPolynomial &x) {
+    bool sameNumParams = (getNumParams() == -1) || (x.getNumParams() == -1) ||
+                         (getNumParams() == x.getNumParams());
+    assert(sameNumParams && "two quasi-polynomials with different numbers of "
+                            "parameters cannot be added!");
+    coefficients.append(x.coefficients);
+    affine.insert(affine.end(), x.affine.begin(), x.affine.end());
+    return *this;
+  }
+
+  QuasiPolynomial operator-(const QuasiPolynomial &x) {
+    bool sameNumParams = (getNumParams() == -1) || (x.getNumParams() == -1) ||
+                         (getNumParams() == x.getNumParams());
+    assert(sameNumParams && "two quasi-polynomials with different numbers of "
+                            "parameters cannot be subtracted!");
+    QuasiPolynomial qp(x.coefficients, x.affine);
+    for (unsigned i = 0; i < x.coefficients.size(); i++)
+      qp.coefficients[i] = -qp.coefficients[i];
+    return (*this + qp);
+  }
+
+  QuasiPolynomial operator*(const QuasiPolynomial &x) {
+    bool sameNumParams = (getNumParams() == -1) || (x.getNumParams() == -1) ||
+                         (getNumParams() == x.getNumParams());
+    assert(sameNumParams && "two quasi-polynomials with different numbers of "
+                            "parameters cannot be multiplied!");
+    QuasiPolynomial qp();
+    std::vector<SmallVector<Fraction>> product;
+    for (unsigned i = 0; i < coefficients.size(); i++) {
+      for (unsigned j = 0; j < x.coefficients.size(); j++) {
+        qp.coefficients.append({coefficients[i] * x.coefficients[j]});
+
+        product.clear();
+        product.insert(product.end(), affine[i].begin(), affine[i].end());
+        product.insert(product.end(), x.affine[j].begin(), x.affine[j].end());
+
+        qp.affine.push_back(product);
+      }
+    }
+
+    return qp;
+  }
+
+  QuasiPolynomial operator/(Fraction x) {
+    assert(x != 0 && "division by zero!");
+    for (unsigned i = 0; i < coefficients.size(); i++)
+      coefficients[i] = coefficients[i] / x;
+    return *this;
+  };
+
+  // Removes terms which evaluate to zero from the expression.
+  QuasiPolynomial reduce() {
+    SmallVector<Fraction> newCoeffs({});
+    std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
+    bool prodIsNonz, sumIsNonz;
+    for (unsigned i = 0; i < coefficients.size(); i++) {
+      prodIsNonz = true;
+      for (unsigned j = 0; j < affine[i].size(); j++) {
+        sumIsNonz = false;
+        for (unsigned k = 0; k < affine[i][j].size(); k++)
+          if (affine[i][j][k] != Fraction(0, 1))
+            sumIsNonz = true;
+        if (sumIsNonz == false)
+          prodIsNonz = false;
+      }
+      if (prodIsNonz == true && coefficients[i] != Fraction(0, 1)) {
+        newCoeffs.append({coefficients[i]});
+        newAffine.push_back({affine[i]});
+      }
+    }
+    return QuasiPolynomial(newCoeffs, newAffine);
+  }
+};
+
+} // namespace presburger
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGER_BARVINOK_H
\ No newline at end of file

>From 6101acb97250f3455bf3aaf4cf3b88165b2c740b Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sat, 16 Dec 2023 21:02:51 +0530
Subject: [PATCH 02/38] Comments

---
 mlir/include/mlir/Analysis/Presburger/Barvinok.h | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 3ee9d5dc452d0e..d753f97d3bcf74 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -60,9 +60,15 @@ class GeneratingFunction {
                      std::vector<std::vector<Point>> dens)
       : signs(s), numerators(nums), denominators(dens){};
 
+  // Find the number of parameters involved in the function
+  // from the dimensionality of the affine functions.
   unsigned getNumParams() {
     for (auto term : numerators)
+      // The number of elements in the affine function is
+      // one more than the number of parameters.
       return (term.getNumColumns() - 1);
+    // The polynomial can be treated as having any number
+    // of parameters.
     return -1;
   }
 

>From 5d425bd2c83576f3ba389aed3aa3291375dab185 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 01:27:14 +0530
Subject: [PATCH 03/38] Miscellaneous reformatting

---
 .../mlir/Analysis/Presburger/Barvinok.h       | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index d753f97d3bcf74..bbd4baf3dd8cd8 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -56,9 +56,9 @@ using Point = SmallVector<Fraction>;
 //      is a vector (a generator).
 class GeneratingFunction {
 public:
-  GeneratingFunction(SmallVector<int, 16> s, std::vector<ParamPoint> nums,
+  GeneratingFunction(SmallVector<int, 16> signs, std::vector<ParamPoint> nums,
                      std::vector<std::vector<Point>> dens)
-      : signs(s), numerators(nums), denominators(dens){};
+      : signs(signs), numerators(nums), denominators(dens) {}
 
   // Find the number of parameters involved in the function
   // from the dimensionality of the affine functions.
@@ -96,18 +96,18 @@ class GeneratingFunction {
   llvm::raw_ostream &print(llvm::raw_ostream &os) const {
     for (unsigned i = 0; i < signs.size(); i++) {
       if (signs[i] == 1)
-        os << "+";
+        os << " + ";
       else
-        os << "-";
+        os << " - ";
 
-      os << "*(x^[";
+      os << "(x^[";
       for (unsigned j = 0; j < numerators[i].size() - 1; j++)
         os << numerators[i][j] << ",";
-      os << numerators[i][numerators[i].size() - 1] << "])/";
+      os << numerators[i].back() << "])/";
 
       for (Point den : denominators[i]) {
         os << "(x^[";
-        for (unsigned j = 0; j < den.size(); j++)
+        for (unsigned j = 0; j < den.size() - 1; j++)
           os << den[j] << ",";
         os << den[den.size() - 1] << "])";
       }
@@ -115,7 +115,7 @@ class GeneratingFunction {
     return os;
   }
 
-  SmallVector<int, 16> signs;
+  SmallVector<int, 8> signs;
   std::vector<ParamPoint> numerators;
   std::vector<std::vector<Point>> denominators;
 };
@@ -130,9 +130,9 @@ class QuasiPolynomial {
 public:
   QuasiPolynomial(SmallVector<Fraction> coeffs = {},
                   std::vector<std::vector<SmallVector<Fraction>>> aff = {})
-      : coefficients(coeffs), affine(aff){};
+      : coefficients(coeffs), affine(aff) {}
 
-  QuasiPolynomial(Fraction cons) : coefficients({cons}), affine({{}}){};
+  QuasiPolynomial(Fraction cons) : coefficients({cons}), affine({{}}) {}
 
   QuasiPolynomial(QuasiPolynomial const &) = default;
 

>From 883e99356b4875c83eb79bba4bf591d5f5a95444 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 01:31:44 +0530
Subject: [PATCH 04/38] Remove superfluous == operator

---
 mlir/include/mlir/Analysis/Presburger/Barvinok.h | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index bbd4baf3dd8cd8..d99ec62bd93a41 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -72,13 +72,6 @@ class GeneratingFunction {
     return -1;
   }
 
-  bool operator==(const GeneratingFunction &gf) const {
-    if (signs != gf.signs || numerators != gf.numerators ||
-        denominators != gf.denominators)
-      return false;
-    return true;
-  }
-
   GeneratingFunction operator+(const GeneratingFunction &gf) {
     bool sameNumParams = (getNumParams() == -1) || (gf.getNumParams() == -1) ||
                          (getNumParams() == gf.getNumParams());

>From 2812e75ff47666d56053c1396b9157d52031ed88 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 01:40:04 +0530
Subject: [PATCH 05/38] Use SmallVector::append() instead of insert()

---
 mlir/include/mlir/Analysis/Presburger/Barvinok.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index d99ec62bd93a41..b7dd528f92ffd3 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -78,7 +78,7 @@ class GeneratingFunction {
     assert(
         sameNumParams &&
         "two generators with different numbers of parameters cannot be added!");
-    signs.insert(signs.end(), gf.signs.begin(), gf.signs.end());
+    signs.append(gf.signs);
     numerators.insert(numerators.end(), gf.numerators.begin(),
                       gf.numerators.end());
     denominators.insert(denominators.end(), gf.denominators.begin(),

>From 625b1b1bbe0c95addf9af0a00d2bbb7aa470d8a7 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 01:51:01 +0530
Subject: [PATCH 06/38] Refactor numParam as private attribute

---
 .../mlir/Analysis/Presburger/Barvinok.h       | 87 +++++++++----------
 1 file changed, 43 insertions(+), 44 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index b7dd528f92ffd3..64e5d98b20adbc 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -56,28 +56,24 @@ using Point = SmallVector<Fraction>;
 //      is a vector (a generator).
 class GeneratingFunction {
 public:
-  GeneratingFunction(SmallVector<int, 16> signs, std::vector<ParamPoint> nums,
+  GeneratingFunction(unsigned numParam, SmallVector<int, 8> signs,
+                     std::vector<ParamPoint> nums,
                      std::vector<std::vector<Point>> dens)
-      : signs(signs), numerators(nums), denominators(dens) {}
+      : numParam(numParam), signs(signs), numerators(nums), denominators(dens) {
+    for (auto term : numerators)
+      assert(term.getNumColumns() - 1 == numParam &&
+             "dimensionality of numerator exponents does not match number of "
+             "parameters!");
+  }
 
   // Find the number of parameters involved in the function
   // from the dimensionality of the affine functions.
-  unsigned getNumParams() {
-    for (auto term : numerators)
-      // The number of elements in the affine function is
-      // one more than the number of parameters.
-      return (term.getNumColumns() - 1);
-    // The polynomial can be treated as having any number
-    // of parameters.
-    return -1;
-  }
+  unsigned getNumParams() { return numParam; }
 
   GeneratingFunction operator+(const GeneratingFunction &gf) {
-    bool sameNumParams = (getNumParams() == -1) || (gf.getNumParams() == -1) ||
-                         (getNumParams() == gf.getNumParams());
-    assert(
-        sameNumParams &&
-        "two generators with different numbers of parameters cannot be added!");
+    assert(numParam == gf.getNumParams() &&
+           "two generating functions with different numbers of parameters "
+           "cannot be added!");
     signs.append(gf.signs);
     numerators.insert(numerators.end(), gf.numerators.begin(),
                       gf.numerators.end());
@@ -111,6 +107,9 @@ class GeneratingFunction {
   SmallVector<int, 8> signs;
   std::vector<ParamPoint> numerators;
   std::vector<std::vector<Point>> denominators;
+
+private:
+  unsigned numParam;
 };
 
 // A class to describe the quasi-polynomials obtained by
@@ -121,38 +120,32 @@ class GeneratingFunction {
 // by the product of a set of affine functions on n parameters.
 class QuasiPolynomial {
 public:
-  QuasiPolynomial(SmallVector<Fraction> coeffs = {},
+  QuasiPolynomial(unsigned numParam, SmallVector<Fraction> coeffs = {},
                   std::vector<std::vector<SmallVector<Fraction>>> aff = {})
-      : coefficients(coeffs), affine(aff) {}
+      : numParam(numParam), coefficients(coeffs), affine(aff) {
+    // Find the first term which involves some affine function.
+    for (auto term : affine)
+      if (term.size() == 0)
+        continue;
+    // The number of elements in the affine function is
+    // one more than the number of parameters.
+    assert(term[0].size() - 1 == numParam &&
+           "dimensionality of affine functions does not match number of "
+           "parameters!")
+  }
 
   QuasiPolynomial(Fraction cons) : coefficients({cons}), affine({{}}) {}
 
   QuasiPolynomial(QuasiPolynomial const &) = default;
 
-  SmallVector<Fraction> coefficients;
-  std::vector<std::vector<SmallVector<Fraction>>> affine;
-
   // Find the number of parameters involved in the polynomial
   // from the dimensionality of the affine functions.
-  unsigned getNumParams() {
-    // Find the first term which involves some affine function.
-    for (auto term : affine) {
-      if (term.size() == 0)
-        continue;
-      // The number of elements in the affine function is
-      // one more than the number of parameters.
-      return (term[0].size() - 1);
-    }
-    // The polynomial can be treated as having any number
-    // of parameters.
-    return -1;
-  }
+  unsigned getNumParams() { return numParam; }
 
   QuasiPolynomial operator+(const QuasiPolynomial &x) {
-    bool sameNumParams = (getNumParams() == -1) || (x.getNumParams() == -1) ||
-                         (getNumParams() == x.getNumParams());
-    assert(sameNumParams && "two quasi-polynomials with different numbers of "
-                            "parameters cannot be added!");
+    assert(numParam == x.getNumParams() &&
+           "two quasi-polynomials with different numbers of parameters cannot "
+           "be added!");
     coefficients.append(x.coefficients);
     affine.insert(affine.end(), x.affine.begin(), x.affine.end());
     return *this;
@@ -161,8 +154,9 @@ class QuasiPolynomial {
   QuasiPolynomial operator-(const QuasiPolynomial &x) {
     bool sameNumParams = (getNumParams() == -1) || (x.getNumParams() == -1) ||
                          (getNumParams() == x.getNumParams());
-    assert(sameNumParams && "two quasi-polynomials with different numbers of "
-                            "parameters cannot be subtracted!");
+    assert(numParam == x.getNumParams() &&
+           "two quasi-polynomials with different numbers of parameters cannot "
+           "be subtracted!");
     QuasiPolynomial qp(x.coefficients, x.affine);
     for (unsigned i = 0; i < x.coefficients.size(); i++)
       qp.coefficients[i] = -qp.coefficients[i];
@@ -170,10 +164,9 @@ class QuasiPolynomial {
   }
 
   QuasiPolynomial operator*(const QuasiPolynomial &x) {
-    bool sameNumParams = (getNumParams() == -1) || (x.getNumParams() == -1) ||
-                         (getNumParams() == x.getNumParams());
-    assert(sameNumParams && "two quasi-polynomials with different numbers of "
-                            "parameters cannot be multiplied!");
+    assert(numParam = x.getNumParams() &&
+                      "two quasi-polynomials with different numbers of "
+                      "parameters cannot be multiplied!");
     QuasiPolynomial qp();
     std::vector<SmallVector<Fraction>> product;
     for (unsigned i = 0; i < coefficients.size(); i++) {
@@ -220,6 +213,12 @@ class QuasiPolynomial {
     }
     return QuasiPolynomial(newCoeffs, newAffine);
   }
+
+  SmallVector<Fraction> coefficients;
+  std::vector<std::vector<SmallVector<Fraction>>> affine;
+
+private:
+  unsigned numParam;
 };
 
 } // namespace presburger

>From 94b5996de245280e1e9cd60673d46f793b995416 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 12:49:32 +0530
Subject: [PATCH 07/38] Formatting fixes

---
 .../mlir/Analysis/Presburger/Barvinok.h       | 68 +++++++++++--------
 1 file changed, 39 insertions(+), 29 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 64e5d98b20adbc..80ab8fb4df1ef7 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -1,4 +1,4 @@
-//===- Barvinok.h - Barvinok's Algorithm -----------*- C++ -*-===//
+//===- Barvinok.h - Barvinok's Algorithm ------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -54,13 +54,21 @@ using Point = SmallVector<Fraction>;
 // * a denominator, of the form (1 - x^{d1})...(1 - x^{dn}),
 //      where each dj, stored in `denominators[i][j]`,
 //      is a vector (a generator).
+//
+// Represents functions f : Q^n -> Q of the form
+//
+// f(x) = \sum_i s_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
+//
+// where s_i is ±1,
+// n_i \in (Q^d -> Q)^n is an n-vector of affine functions on d parameters, and
+// g_{ij} \in Q^n are vectors.
 class GeneratingFunction {
 public:
   GeneratingFunction(unsigned numParam, SmallVector<int, 8> signs,
                      std::vector<ParamPoint> nums,
                      std::vector<std::vector<Point>> dens)
       : numParam(numParam), signs(signs), numerators(nums), denominators(dens) {
-    for (auto term : numerators)
+    for (const auto &term : numerators)
       assert(term.getNumColumns() - 1 == numParam &&
              "dimensionality of numerator exponents does not match number of "
              "parameters!");
@@ -83,20 +91,20 @@ class GeneratingFunction {
   }
 
   llvm::raw_ostream &print(llvm::raw_ostream &os) const {
-    for (unsigned i = 0; i < signs.size(); i++) {
+    for (unsigned i = 0, e = signs.size(); i < e; i++) {
       if (signs[i] == 1)
         os << " + ";
       else
         os << " - ";
 
-      os << "(x^[";
-      for (unsigned j = 0; j < numerators[i].size() - 1; j++)
+      os << "x^[";
+      for (unsigned j = 0, e = numerators[i].size(); j < e - 1; j++)
         os << numerators[i][j] << ",";
-      os << numerators[i].back() << "])/";
+      os << numerators[i].back() << "]/";
 
       for (Point den : denominators[i]) {
         os << "(x^[";
-        for (unsigned j = 0; j < den.size() - 1; j++)
+        for (unsigned j = 0, e = den.size(); j < e - 1; j++)
           os << den[j] << ",";
         os << den[den.size() - 1] << "])";
       }
@@ -118,28 +126,30 @@ class GeneratingFunction {
 // Consists of a set of terms.
 // The ith term is a constant `coefficients[i]`, multiplied
 // by the product of a set of affine functions on n parameters.
+// Represents functions f : Q^n -> Q of the form
+//
+// f(x) = \sum_i c_i * \prod_j ⌊g_{ij}(x)⌋
+//
+// where c_i \in Q and
+// g_{ij} : Q^n -> Q are affine functionals.
 class QuasiPolynomial {
 public:
   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 (auto term : affine)
+    for (auto term : affine) {
       if (term.size() == 0)
         continue;
-    // The number of elements in the affine function is
-    // one more than the number of parameters.
-    assert(term[0].size() - 1 == numParam &&
-           "dimensionality of affine functions does not match number of "
-           "parameters!")
+      // The number of elements in the affine function is
+      // one more than the number of parameters.
+      assert(term[0].size() - 1 == numParam &&
+             "dimensionality of affine functions does not match number of "
+             "parameters!")
+    }
   }
 
-  QuasiPolynomial(Fraction cons) : coefficients({cons}), affine({{}}) {}
-
-  QuasiPolynomial(QuasiPolynomial const &) = default;
-
-  // Find the number of parameters involved in the polynomial
-  // from the dimensionality of the affine functions.
+  // Find the number of parameters involved in the polynomial.
   unsigned getNumParams() { return numParam; }
 
   QuasiPolynomial operator+(const QuasiPolynomial &x) {
@@ -158,7 +168,7 @@ class QuasiPolynomial {
            "two quasi-polynomials with different numbers of parameters cannot "
            "be subtracted!");
     QuasiPolynomial qp(x.coefficients, x.affine);
-    for (unsigned i = 0; i < x.coefficients.size(); i++)
+    for (unsigned i = 0, e = x.coefficients.size(); i < e; i++)
       qp.coefficients[i] = -qp.coefficients[i];
     return (*this + qp);
   }
@@ -169,8 +179,8 @@ class QuasiPolynomial {
                       "parameters cannot be multiplied!");
     QuasiPolynomial qp();
     std::vector<SmallVector<Fraction>> product;
-    for (unsigned i = 0; i < coefficients.size(); i++) {
-      for (unsigned j = 0; j < x.coefficients.size(); j++) {
+    for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
+      for (unsigned j = 0, e = x.coefficients.size(); j < e; j++) {
         qp.coefficients.append({coefficients[i] * x.coefficients[j]});
 
         product.clear();
@@ -186,27 +196,27 @@ class QuasiPolynomial {
 
   QuasiPolynomial operator/(Fraction x) {
     assert(x != 0 && "division by zero!");
-    for (unsigned i = 0; i < coefficients.size(); i++)
-      coefficients[i] = coefficients[i] / x;
+    for (Fraction &coeff : coefficients)
+      coeff /= x;
     return *this;
   };
 
   // Removes terms which evaluate to zero from the expression.
-  QuasiPolynomial reduce() {
+  QuasiPolynomial simplify() {
     SmallVector<Fraction> newCoeffs({});
     std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
     bool prodIsNonz, sumIsNonz;
-    for (unsigned i = 0; i < coefficients.size(); i++) {
+    for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
       prodIsNonz = true;
-      for (unsigned j = 0; j < affine[i].size(); j++) {
+      for (unsigned j = 0, e = affine[i].size(); j < e; j++) {
         sumIsNonz = false;
-        for (unsigned k = 0; k < affine[i][j].size(); k++)
+        for (unsigned k = 0, e = affine[i][j].size(); k < e; k++)
           if (affine[i][j][k] != Fraction(0, 1))
             sumIsNonz = true;
         if (sumIsNonz == false)
           prodIsNonz = false;
       }
-      if (prodIsNonz == true && coefficients[i] != Fraction(0, 1)) {
+      if (prodIsNonz && coefficients[i] != Fraction(0, 1)) {
         newCoeffs.append({coefficients[i]});
         newAffine.push_back({affine[i]});
       }

>From f005fdc6cdac0b7b859066e5d8d3cd6c47b9a312 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 12:51:58 +0530
Subject: [PATCH 08/38] Add type instead of auto

---
 mlir/include/mlir/Analysis/Presburger/Barvinok.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 80ab8fb4df1ef7..682099ed84b029 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -68,7 +68,7 @@ class GeneratingFunction {
                      std::vector<ParamPoint> nums,
                      std::vector<std::vector<Point>> dens)
       : numParam(numParam), signs(signs), numerators(nums), denominators(dens) {
-    for (const auto &term : numerators)
+    for (const ParamPoint &term : numerators)
       assert(term.getNumColumns() - 1 == numParam &&
              "dimensionality of numerator exponents does not match number of "
              "parameters!");
@@ -138,7 +138,7 @@ class QuasiPolynomial {
                   std::vector<std::vector<SmallVector<Fraction>>> aff = {})
       : numParam(numParam), coefficients(coeffs), affine(aff) {
     // Find the first term which involves some affine function.
-    for (auto term : affine) {
+    for (const std::vector<SmallVector<Fraction>> &term : affine) {
       if (term.size() == 0)
         continue;
       // The number of elements in the affine function is

>From 34da1b2e557c503cb6207d2d3d039adea1e85674 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 21:21:13 +0530
Subject: [PATCH 09/38] Pass numParam

---
 mlir/include/mlir/Analysis/Presburger/Barvinok.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 682099ed84b029..f9c06ef56157cc 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -177,7 +177,7 @@ class QuasiPolynomial {
     assert(numParam = x.getNumParams() &&
                       "two quasi-polynomials with different numbers of "
                       "parameters cannot be multiplied!");
-    QuasiPolynomial qp();
+    QuasiPolynomial qp(numParam);
     std::vector<SmallVector<Fraction>> product;
     for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
       for (unsigned j = 0, e = x.coefficients.size(); j < e; j++) {

>From 07e5be7b0c5d943ba67895340ecc07a737055fa6 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 21:28:29 +0530
Subject: [PATCH 10/38] Use range-based loop

---
 mlir/include/mlir/Analysis/Presburger/Barvinok.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index f9c06ef56157cc..e8f2c20662f2c3 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -168,8 +168,8 @@ class QuasiPolynomial {
            "two quasi-polynomials with different numbers of parameters cannot "
            "be subtracted!");
     QuasiPolynomial qp(x.coefficients, x.affine);
-    for (unsigned i = 0, e = x.coefficients.size(); i < e; i++)
-      qp.coefficients[i] = -qp.coefficients[i];
+    for (Fraction &coeff : qp.coefficients)
+      coeff = -coeff;
     return (*this + qp);
   }
 

>From f255ad13abd7ee2fb15c4aebe5bf5abd2b33c2b8 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 21:29:48 +0530
Subject: [PATCH 11/38] Formatting

---
 .../mlir/Analysis/Presburger/Barvinok.h        | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index e8f2c20662f2c3..a5ddf5e28b2360 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -162,15 +162,13 @@ class QuasiPolynomial {
   }
 
   QuasiPolynomial operator-(const QuasiPolynomial &x) {
-    bool sameNumParams = (getNumParams() == -1) || (x.getNumParams() == -1) ||
-                         (getNumParams() == x.getNumParams());
     assert(numParam == x.getNumParams() &&
            "two quasi-polynomials with different numbers of parameters cannot "
            "be subtracted!");
     QuasiPolynomial qp(x.coefficients, x.affine);
     for (Fraction &coeff : qp.coefficients)
       coeff = -coeff;
-    return (*this + qp);
+    return *this + qp;
   }
 
   QuasiPolynomial operator*(const QuasiPolynomial &x) {
@@ -205,18 +203,18 @@ class QuasiPolynomial {
   QuasiPolynomial simplify() {
     SmallVector<Fraction> newCoeffs({});
     std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
-    bool prodIsNonz, sumIsNonz;
+    bool prodIsNonzero, sumIsNonzero;
     for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
-      prodIsNonz = true;
+      prodIsNonzero = true;
       for (unsigned j = 0, e = affine[i].size(); j < e; j++) {
-        sumIsNonz = false;
+        sumIsNonzero = false;
         for (unsigned k = 0, e = affine[i][j].size(); k < e; k++)
           if (affine[i][j][k] != Fraction(0, 1))
-            sumIsNonz = true;
-        if (sumIsNonz == false)
-          prodIsNonz = false;
+            sumIsNonzero = true;
+        if (sumIsNonzero == false)
+          prodIsNonzero = false;
       }
-      if (prodIsNonz && coefficients[i] != Fraction(0, 1)) {
+      if (prodIsNonzero && coefficients[i] != Fraction(0, 1)) {
         newCoeffs.append({coefficients[i]});
         newAffine.push_back({affine[i]});
       }

>From ef899d9ab25ffb489c5c36a8ce970eaee187b9b3 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Sun, 17 Dec 2023 23:36:27 +0530
Subject: [PATCH 12/38] Use llvm::any_of and llvm::all_of

---
 .../mlir/Analysis/Presburger/Barvinok.h       | 25 ++++++++-----------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index a5ddf5e28b2360..7f24882457512d 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -203,21 +203,18 @@ class QuasiPolynomial {
   QuasiPolynomial simplify() {
     SmallVector<Fraction> newCoeffs({});
     std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
-    bool prodIsNonzero, sumIsNonzero;
     for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
-      prodIsNonzero = true;
-      for (unsigned j = 0, e = affine[i].size(); j < e; j++) {
-        sumIsNonzero = false;
-        for (unsigned k = 0, e = affine[i][j].size(); k < e; k++)
-          if (affine[i][j][k] != Fraction(0, 1))
-            sumIsNonzero = true;
-        if (sumIsNonzero == false)
-          prodIsNonzero = false;
-      }
-      if (prodIsNonzero && coefficients[i] != Fraction(0, 1)) {
-        newCoeffs.append({coefficients[i]});
-        newAffine.push_back({affine[i]});
-      }
+      // A term is zero if its coefficient is zero, or
+      if (coefficients[i] == Fraction(0, 1) ||
+              // if any of the affine functions in the product
+              llvm::any_of(affine[i], [](SmallVector<Fraction> affine_ij) {
+                // has all its coefficients as zero.
+                return llvm::all_of(affine_ij,
+                                    [](Fraction f) { return f == 0; });
+              });)
+        continue;
+      newCoeffs.append({coefficients[i]});
+      newAffine.push_back({affine[i]});
     }
     return QuasiPolynomial(newCoeffs, newAffine);
   }

>From 9bc0cb5d2598f9feba119feab7e2e1be7b449532 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 18 Dec 2023 15:41:12 +0530
Subject: [PATCH 13/38] Make data private

---
 mlir/include/mlir/Analysis/Presburger/Barvinok.h | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 7f24882457512d..1980b5d634b6b2 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -78,6 +78,12 @@ class GeneratingFunction {
   // from the dimensionality of the affine functions.
   unsigned getNumParams() { return numParam; }
 
+  SmallVector<int, 8> getSigns() { return signs; }
+
+  std::vector<ParamPoint> getNumerators() { return numerators; }
+
+  std::vector<std::vector<Point>> getDenominators() { return denominators; }
+
   GeneratingFunction operator+(const GeneratingFunction &gf) {
     assert(numParam == gf.getNumParams() &&
            "two generating functions with different numbers of parameters "
@@ -112,12 +118,11 @@ class GeneratingFunction {
     return os;
   }
 
+private:
+  unsigned numParam;
   SmallVector<int, 8> signs;
   std::vector<ParamPoint> numerators;
   std::vector<std::vector<Point>> denominators;
-
-private:
-  unsigned numParam;
 };
 
 // A class to describe the quasi-polynomials obtained by
@@ -219,11 +224,10 @@ class QuasiPolynomial {
     return QuasiPolynomial(newCoeffs, newAffine);
   }
 
-  SmallVector<Fraction> coefficients;
-  std::vector<std::vector<SmallVector<Fraction>>> affine;
-
 private:
   unsigned numParam;
+  SmallVector<Fraction> coefficients;
+  std::vector<std::vector<SmallVector<Fraction>>> affine;
 };
 
 } // namespace presburger

>From d1e6b249951e836131acf5963dd12c260617faed Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 18 Dec 2023 16:05:42 +0530
Subject: [PATCH 14/38] Define getters

---
 .../include/mlir/Analysis/Presburger/Barvinok.h | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 1980b5d634b6b2..37ba5cc040e871 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -24,17 +24,6 @@
 namespace mlir {
 namespace presburger {
 
-// The H (inequality) representation of both general
-// polyhedra and cones specifically is an integer relation.
-using PolyhedronH = IntegerRelation;
-using ConeH = PolyhedronH;
-
-// The V (generator) representation of both general
-// polyhedra and cones specifically is simply a matrix
-// whose rows are the generators.
-using PolyhedronV = Matrix<MPInt>;
-using ConeV = PolyhedronV;
-
 // A parametric point is a vector, each of whose elements
 // is an affine function of n parameters. Each row
 // in the matrix represents the affine function and
@@ -78,7 +67,7 @@ class GeneratingFunction {
   // from the dimensionality of the affine functions.
   unsigned getNumParams() { return numParam; }
 
-  SmallVector<int, 8> getSigns() { return signs; }
+  SmallVector<int> getSigns() { return signs; }
 
   std::vector<ParamPoint> getNumerators() { return numerators; }
 
@@ -157,6 +146,10 @@ class QuasiPolynomial {
   // Find the number of parameters involved in the polynomial.
   unsigned getNumParams() { return numParam; }
 
+  SmallVector<Fraction> getCoefficients() { return coefficients; }
+
+  std::vector<std::vector<SmallVector<Fraction>>> getAffine() { return affine; }
+
   QuasiPolynomial operator+(const QuasiPolynomial &x) {
     assert(numParam == x.getNumParams() &&
            "two quasi-polynomials with different numbers of parameters cannot "

>From c04f6d9f5a7aed53740a4ab776d236d4bc1b26df Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 18 Dec 2023 19:22:56 +0530
Subject: [PATCH 15/38] Refactor headers

---
 .../mlir/Analysis/Presburger/Barvinok.h       | 229 ------------------
 .../Analysis/Presburger/QuasiPolynomial.h     | 103 ++++++++
 mlir/lib/Analysis/Presburger/CMakeLists.txt   |   1 +
 .../Analysis/Presburger/GeneratingFunction.h  | 116 +++++++++
 .../Analysis/Presburger/QuasiPolynomial.cpp   |  65 +++++
 5 files changed, 285 insertions(+), 229 deletions(-)
 delete mode 100644 mlir/include/mlir/Analysis/Presburger/Barvinok.h
 create mode 100644 mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
 create mode 100644 mlir/lib/Analysis/Presburger/GeneratingFunction.h
 create mode 100644 mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
deleted file mode 100644
index 37ba5cc040e871..00000000000000
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ /dev/null
@@ -1,229 +0,0 @@
-//===- Barvinok.h - Barvinok's Algorithm ------------------------*- 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
-//
-//===----------------------------------------------------------------------===//
-//
-// Functions and classes for Barvinok's algorithm in MLIR.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_ANALYSIS_PRESBURGER_BARVINOK_H
-#define MLIR_ANALYSIS_PRESBURGER_BARVINOK_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>
-
-namespace mlir {
-namespace presburger {
-
-// A parametric point is a vector, each of whose elements
-// is an affine function of n parameters. Each row
-// in the matrix represents the affine function and
-// has n+1 elements.
-using ParamPoint = Matrix<Fraction>;
-
-// A point is simply a vector.
-using Point = SmallVector<Fraction>;
-
-// A class to describe the type of generating function
-// used to enumerate the integer points in a polytope.
-// Consists of a set of terms, where the ith term has
-// * a sign, ±1, stored in `signs[i]`
-// * a numerator, of the form x^{n},
-//      where n, stored in `numerators[i]`,
-//      is a parametric point (a vertex).
-// * a denominator, of the form (1 - x^{d1})...(1 - x^{dn}),
-//      where each dj, stored in `denominators[i][j]`,
-//      is a vector (a generator).
-//
-// Represents functions f : Q^n -> Q of the form
-//
-// f(x) = \sum_i s_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
-//
-// where s_i is ±1,
-// n_i \in (Q^d -> Q)^n is an n-vector of affine functions on d parameters, and
-// g_{ij} \in Q^n are vectors.
-class GeneratingFunction {
-public:
-  GeneratingFunction(unsigned numParam, SmallVector<int, 8> signs,
-                     std::vector<ParamPoint> nums,
-                     std::vector<std::vector<Point>> dens)
-      : numParam(numParam), signs(signs), numerators(nums), denominators(dens) {
-    for (const ParamPoint &term : numerators)
-      assert(term.getNumColumns() - 1 == numParam &&
-             "dimensionality of numerator exponents does not match number of "
-             "parameters!");
-  }
-
-  // Find the number of parameters involved in the function
-  // from the dimensionality of the affine functions.
-  unsigned getNumParams() { return numParam; }
-
-  SmallVector<int> getSigns() { return signs; }
-
-  std::vector<ParamPoint> getNumerators() { return numerators; }
-
-  std::vector<std::vector<Point>> getDenominators() { return denominators; }
-
-  GeneratingFunction operator+(const GeneratingFunction &gf) {
-    assert(numParam == gf.getNumParams() &&
-           "two generating functions with different numbers of parameters "
-           "cannot be added!");
-    signs.append(gf.signs);
-    numerators.insert(numerators.end(), gf.numerators.begin(),
-                      gf.numerators.end());
-    denominators.insert(denominators.end(), gf.denominators.begin(),
-                        gf.denominators.end());
-    return *this;
-  }
-
-  llvm::raw_ostream &print(llvm::raw_ostream &os) const {
-    for (unsigned i = 0, e = signs.size(); i < e; i++) {
-      if (signs[i] == 1)
-        os << " + ";
-      else
-        os << " - ";
-
-      os << "x^[";
-      for (unsigned j = 0, e = numerators[i].size(); j < e - 1; j++)
-        os << numerators[i][j] << ",";
-      os << numerators[i].back() << "]/";
-
-      for (Point den : denominators[i]) {
-        os << "(x^[";
-        for (unsigned j = 0, e = den.size(); j < e - 1; j++)
-          os << den[j] << ",";
-        os << den[den.size() - 1] << "])";
-      }
-    }
-    return os;
-  }
-
-private:
-  unsigned numParam;
-  SmallVector<int, 8> signs;
-  std::vector<ParamPoint> numerators;
-  std::vector<std::vector<Point>> denominators;
-};
-
-// A class to describe the quasi-polynomials obtained by
-// substituting the unit vector in the type of generating
-// function described above.
-// Consists of a set of terms.
-// The ith term is a constant `coefficients[i]`, multiplied
-// by the product of a set of affine functions on n parameters.
-// Represents functions f : Q^n -> Q of the form
-//
-// f(x) = \sum_i c_i * \prod_j ⌊g_{ij}(x)⌋
-//
-// where c_i \in Q and
-// g_{ij} : Q^n -> Q are affine functionals.
-class QuasiPolynomial {
-public:
-  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 the affine function is
-      // one more than the number of parameters.
-      assert(term[0].size() - 1 == numParam &&
-             "dimensionality of affine functions does not match number of "
-             "parameters!")
-    }
-  }
-
-  // Find the number of parameters involved in the polynomial.
-  unsigned getNumParams() { return numParam; }
-
-  SmallVector<Fraction> getCoefficients() { return coefficients; }
-
-  std::vector<std::vector<SmallVector<Fraction>>> getAffine() { return affine; }
-
-  QuasiPolynomial operator+(const QuasiPolynomial &x) {
-    assert(numParam == x.getNumParams() &&
-           "two quasi-polynomials with different numbers of parameters cannot "
-           "be added!");
-    coefficients.append(x.coefficients);
-    affine.insert(affine.end(), x.affine.begin(), x.affine.end());
-    return *this;
-  }
-
-  QuasiPolynomial operator-(const QuasiPolynomial &x) {
-    assert(numParam == x.getNumParams() &&
-           "two quasi-polynomials with different numbers of parameters cannot "
-           "be subtracted!");
-    QuasiPolynomial qp(x.coefficients, x.affine);
-    for (Fraction &coeff : qp.coefficients)
-      coeff = -coeff;
-    return *this + qp;
-  }
-
-  QuasiPolynomial operator*(const QuasiPolynomial &x) {
-    assert(numParam = x.getNumParams() &&
-                      "two quasi-polynomials with different numbers of "
-                      "parameters cannot be multiplied!");
-    QuasiPolynomial qp(numParam);
-    std::vector<SmallVector<Fraction>> product;
-    for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
-      for (unsigned j = 0, e = x.coefficients.size(); j < e; j++) {
-        qp.coefficients.append({coefficients[i] * x.coefficients[j]});
-
-        product.clear();
-        product.insert(product.end(), affine[i].begin(), affine[i].end());
-        product.insert(product.end(), x.affine[j].begin(), x.affine[j].end());
-
-        qp.affine.push_back(product);
-      }
-    }
-
-    return qp;
-  }
-
-  QuasiPolynomial operator/(Fraction x) {
-    assert(x != 0 && "division by zero!");
-    for (Fraction &coeff : coefficients)
-      coeff /= x;
-    return *this;
-  };
-
-  // Removes terms which evaluate to zero from the expression.
-  QuasiPolynomial simplify() {
-    SmallVector<Fraction> newCoeffs({});
-    std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
-    for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
-      // A term is zero if its coefficient is zero, or
-      if (coefficients[i] == Fraction(0, 1) ||
-              // if any of the affine functions in the product
-              llvm::any_of(affine[i], [](SmallVector<Fraction> affine_ij) {
-                // has all its coefficients as zero.
-                return llvm::all_of(affine_ij,
-                                    [](Fraction f) { return f == 0; });
-              });)
-        continue;
-      newCoeffs.append({coefficients[i]});
-      newAffine.push_back({affine[i]});
-    }
-    return QuasiPolynomial(newCoeffs, newAffine);
-  }
-
-private:
-  unsigned numParam;
-  SmallVector<Fraction> coefficients;
-  std::vector<std::vector<SmallVector<Fraction>>> affine;
-};
-
-} // namespace presburger
-} // namespace mlir
-
-#endif // MLIR_ANALYSIS_PRESBURGER_BARVINOK_H
\ No newline at end of file
diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
new file mode 100644
index 00000000000000..6c457a8d94bead
--- /dev/null
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -0,0 +1,103 @@
+//===- QuasiPolynomial.h - 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Definition of the QuasiPolynomial class for Barvinok's algorithm,
+// which represents a single-valued function on a set of parameters.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_QUASIPOLYNOMIAL_H
+#define MLIR_ANALYSIS_PRESBURGER_QUASIPOLYNOMIAL_H
+
+#include "mlir/Analysis/Presburger/Fraction.h"
+#include "mlir/Analysis/Presburger/Matrix.h"
+
+namespace mlir {
+namespace presburger {
+
+// A class to describe the quasi-polynomials obtained by
+// substituting the unit vector in the type of generating
+// function described above.
+// Consists of a set of terms.
+// The ith term is a constant `coefficients[i]`, multiplied
+// by the product of a set of affine functions on n parameters.
+// Represents functions f : Q^n -> Q of the form
+//
+// f(x) = \sum_i c_i * \prod_j ⌊g_{ij}(x)⌋
+//
+// where c_i \in Q and
+// g_{ij} : Q^n -> Q are affine functionals.
+class QuasiPolynomial {
+public:
+  QuasiPolynomial(unsigned numParam, SmallVector<Fraction> coeffs = {},
+                  std::vector<std::vector<SmallVector<Fraction>>> aff = {});
+
+  unsigned getNumParams();
+  SmallVector<Fraction> getCoefficients();
+  std::vector<std::vector<SmallVector<Fraction>>> getAffine();
+
+  QuasiPolynomial operator+(QuasiPolynomial &x) {
+    assert(numParam == x.getNumParams() &&
+           "two quasi-polynomials with different numbers of parameters cannot "
+           "be added!");
+    coefficients.append(x.coefficients);
+    affine.insert(affine.end(), x.affine.begin(), x.affine.end());
+    return *this;
+  }
+
+  QuasiPolynomial operator-(QuasiPolynomial &x) {
+    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 operator*(QuasiPolynomial &x) {
+    assert(numParam = x.getNumParams() &&
+                      "two quasi-polynomials with different numbers of "
+                      "parameters cannot be multiplied!");
+    QuasiPolynomial qp(numParam);
+    std::vector<SmallVector<Fraction>> product;
+    for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
+      for (unsigned j = 0, e = x.coefficients.size(); j < e; j++) {
+        qp.coefficients.append({coefficients[i] * x.coefficients[j]});
+
+        product.clear();
+        product.insert(product.end(), affine[i].begin(), affine[i].end());
+        product.insert(product.end(), x.affine[j].begin(), x.affine[j].end());
+
+        qp.affine.push_back(product);
+      }
+    }
+
+    return qp;
+  }
+
+  QuasiPolynomial operator/(Fraction x) {
+    assert(x != 0 && "division by zero!");
+    for (Fraction &coeff : coefficients)
+      coeff /= x;
+    return *this;
+  };
+
+  // Removes terms which evaluate to zero from the expression.
+  QuasiPolynomial simplify();
+
+private:
+  unsigned numParam;
+  SmallVector<Fraction> coefficients;
+  std::vector<std::vector<SmallVector<Fraction>>> affine;
+};
+
+} // namespace presburger
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGER_QUASIPOLYNOMIAL_H
\ No newline at end of file
diff --git a/mlir/lib/Analysis/Presburger/CMakeLists.txt b/mlir/lib/Analysis/Presburger/CMakeLists.txt
index 22f1a4cac44055..e77e1623dae175 100644
--- a/mlir/lib/Analysis/Presburger/CMakeLists.txt
+++ b/mlir/lib/Analysis/Presburger/CMakeLists.txt
@@ -6,6 +6,7 @@ add_mlir_library(MLIRPresburger
   PresburgerRelation.cpp
   PresburgerSpace.cpp
   PWMAFunction.cpp
+  QuasiPolynomial.cpp
   Simplex.cpp
   SlowMPInt.cpp
   Utils.cpp
diff --git a/mlir/lib/Analysis/Presburger/GeneratingFunction.h b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
new file mode 100644
index 00000000000000..11832e822dc021
--- /dev/null
+++ b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
@@ -0,0 +1,116 @@
+//===- GeneratingFunction.h - Rational Polynomials over Q^d -----*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Definition of the GeneratingFunction class for Barvinok's algorithm,
+// which represents a function over Q^n, parameterized by d parameters.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_GENERATINGFUNCTION_H
+#define MLIR_ANALYSIS_PRESBURGER_GENERATINGFUNCTION_H
+
+#include "mlir/Analysis/Presburger/Fraction.h"
+#include "mlir/Analysis/Presburger/Matrix.h"
+
+namespace mlir {
+namespace presburger {
+
+// A parametric point is a vector, each of whose elements
+// is an affine function of n parameters. Each row
+// in the matrix represents the affine function and
+// has n+1 elements.
+using ParamPoint = Matrix<Fraction>;
+
+// A point is simply a vector.
+using Point = SmallVector<Fraction>;
+
+// A class to describe the type of generating function
+// used to enumerate the integer points in a polytope.
+// Consists of a set of terms, where the ith term has
+// * a sign, ±1, stored in `signs[i]`
+// * a numerator, of the form x^{n},
+//      where n, stored in `numerators[i]`,
+//      is a parametric point (a vertex).
+// * a denominator, of the form (1 - x^{d1})...(1 - x^{dn}),
+//      where each dj, stored in `denominators[i][j]`,
+//      is a vector (a generator).
+//
+// Represents functions f : Q^n -> Q of the form
+//
+// f(x) = \sum_i s_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
+//
+// where s_i is ±1,
+// n_i \in (Q^d -> Q)^n is an n-vector of affine functions on d parameters, and
+// g_{ij} \in Q^n are vectors.
+class GeneratingFunction {
+public:
+  GeneratingFunction(unsigned numParam, SmallVector<int, 8> signs,
+                     std::vector<ParamPoint> nums,
+                     std::vector<std::vector<Point>> dens)
+      : numParam(numParam), signs(signs), numerators(nums), denominators(dens) {
+    for (const ParamPoint &term : numerators)
+      assert(term.getNumColumns() - 1 == numParam &&
+             "dimensionality of numerator exponents does not match number of "
+             "parameters!");
+  }
+
+  // Find the number of parameters involved in the function
+  // from the dimensionality of the affine functions.
+  unsigned getNumParams() { return numParam; }
+
+  SmallVector<int> getSigns() { return signs; }
+
+  std::vector<ParamPoint> getNumerators() { return numerators; }
+
+  std::vector<std::vector<Point>> getDenominators() { return denominators; }
+
+  GeneratingFunction operator+(const GeneratingFunction &gf) {
+    assert(numParam == gf.getNumParams() &&
+           "two generating functions with different numbers of parameters "
+           "cannot be added!");
+    signs.append(gf.signs);
+    numerators.insert(numerators.end(), gf.numerators.begin(),
+                      gf.numerators.end());
+    denominators.insert(denominators.end(), gf.denominators.begin(),
+                        gf.denominators.end());
+    return *this;
+  }
+
+  llvm::raw_ostream &print(llvm::raw_ostream &os) const {
+    for (unsigned i = 0, e = signs.size(); i < e; i++) {
+      if (signs[i] == 1)
+        os << " + ";
+      else
+        os << " - ";
+
+      os << "x^[";
+      for (unsigned j = 0, e = numerators[i].size(); j < e - 1; j++)
+        os << numerators[i][j] << ",";
+      os << numerators[i].back() << "]/";
+
+      for (Point den : denominators[i]) {
+        os << "(x^[";
+        for (unsigned j = 0, e = den.size(); j < e - 1; j++)
+          os << den[j] << ",";
+        os << den[den.size() - 1] << "])";
+      }
+    }
+    return os;
+  }
+
+private:
+  unsigned numParam;
+  SmallVector<int, 8> signs;
+  std::vector<ParamPoint> numerators;
+  std::vector<std::vector<Point>> denominators;
+};
+
+} // namespace presburger
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_PRESBURGER_GENERATINGFUNCTION_H
\ No newline at end of file
diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
new file mode 100644
index 00000000000000..14cfdcad0fb255
--- /dev/null
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -0,0 +1,65 @@
+//===- 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 the affine function is
+    // one more than the number of parameters.
+    assert(term[0].size() - 1 == numParam &&
+           "dimensionality of affine functions does not match number of "
+           "parameters!");
+  }
+}
+
+// Find the number of parameters involved in the polynomial.
+unsigned QuasiPolynomial::getNumParams() { return numParam; }
+
+SmallVector<Fraction> QuasiPolynomial::getCoefficients() {
+  return coefficients;
+}
+
+std::vector<std::vector<SmallVector<Fraction>>> QuasiPolynomial::getAffine() {
+  return affine;
+}
+
+// Removes terms which evaluate to zero from the expression.
+QuasiPolynomial QuasiPolynomial::simplify() {
+  SmallVector<Fraction> newCoeffs({});
+  std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
+  for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
+    // A term is zero if its coefficient is zero, or
+    if (coefficients[i] == Fraction(0, 1) ||
+        // if any of the affine functions in the product
+        llvm::any_of(affine[i], [](SmallVector<Fraction> affine_ij) {
+          // has all its coefficients as zero.
+          return llvm::all_of(affine_ij, [](Fraction f) { return f == 0; });
+        }))
+      continue;
+    newCoeffs.append({coefficients[i]});
+    newAffine.push_back({affine[i]});
+  }
+  return QuasiPolynomial(numParam, newCoeffs, newAffine);
+}
\ No newline at end of file

>From ebccab232e6eb48f3b800fa5f9578826a2310041 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 18:47:10 +0530
Subject: [PATCH 16/38] Define equality of QPs

---
 .../Analysis/Presburger/CMakeLists.txt        |  1 +
 mlir/unittests/Analysis/Presburger/Utils.h    | 20 +++++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/mlir/unittests/Analysis/Presburger/CMakeLists.txt b/mlir/unittests/Analysis/Presburger/CMakeLists.txt
index b6ce273e35a0ee..e37133354e53ca 100644
--- a/mlir/unittests/Analysis/Presburger/CMakeLists.txt
+++ b/mlir/unittests/Analysis/Presburger/CMakeLists.txt
@@ -11,6 +11,7 @@ add_mlir_unittest(MLIRPresburgerTests
   PresburgerRelationTest.cpp
   PresburgerSpaceTest.cpp
   PWMAFunctionTest.cpp
+  QuasiPolynomialTest.cpp
   SimplexTest.cpp
   UtilsTest.cpp
 )
diff --git a/mlir/unittests/Analysis/Presburger/Utils.h b/mlir/unittests/Analysis/Presburger/Utils.h
index 544577375dd1d1..b3b8d5d9aeb007 100644
--- a/mlir/unittests/Analysis/Presburger/Utils.h
+++ b/mlir/unittests/Analysis/Presburger/Utils.h
@@ -16,6 +16,7 @@
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
 #include "mlir/Analysis/Presburger/PWMAFunction.h"
+#include "mlir/Analysis/Presburger/QuasiPolynomial.h"
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/IR/MLIRContext.h"
@@ -71,6 +72,25 @@ inline void EXPECT_EQ_FRAC_MATRIX(FracMatrix a, FracMatrix b) {
       EXPECT_EQ(a(row, col), b(row, col));
 }
 
+inline void EXPECT_EQ_QUASIPOLYNOMIAL(QuasiPolynomial a, QuasiPolynomial b) {
+  EXPECT_EQ(a.getNumParams(), b.getNumParams());
+
+  SmallVector<Fraction> aCoeffs = a.getCoefficients(), bCoeffs = b.getCoefficients();
+  EXPECT_EQ(aCoeffs.size(), bCoeffs.size());
+  for (unsigned i = 0, e = aCoeffs.size(); i < e; i++)
+    EXPECT_EQ(aCoeffs[i], bCoeffs[i]);
+  
+  std::vector<std::vector<SmallVector<Fraction>>> aAff = a.getAffine(), bAff = b.getAffine();
+  EXPECT_EQ(aAff.size(), bAff.size());
+  for (unsigned i = 0, e = aAff.size(); i < e; i++) {
+    EXPECT_EQ(aAff[i].size(), bAff[i].size());
+    for (unsigned j = 0, f = aAff[i].size(); j < f; j++) {
+      for (unsigned k = 0, g = a.getNumParams(); k <= g; k++)
+        EXPECT_EQ(aAff[i][j][k], bAff[i][j][k]);
+    }
+  }
+}
+
 /// lhs and rhs represent non-negative integers or positive infinity. The
 /// infinity case corresponds to when the Optional is empty.
 inline bool infinityOrUInt64LE(std::optional<MPInt> lhs,

>From e89e96b58f11e445ac2615a7a468e6fdbfaba900 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 18:53:52 +0530
Subject: [PATCH 17/38] Formatting

---
 mlir/unittests/Analysis/Presburger/Utils.h | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlir/unittests/Analysis/Presburger/Utils.h b/mlir/unittests/Analysis/Presburger/Utils.h
index b3b8d5d9aeb007..b0d37f27989103 100644
--- a/mlir/unittests/Analysis/Presburger/Utils.h
+++ b/mlir/unittests/Analysis/Presburger/Utils.h
@@ -75,12 +75,14 @@ inline void EXPECT_EQ_FRAC_MATRIX(FracMatrix a, FracMatrix b) {
 inline void EXPECT_EQ_QUASIPOLYNOMIAL(QuasiPolynomial a, QuasiPolynomial b) {
   EXPECT_EQ(a.getNumParams(), b.getNumParams());
 
-  SmallVector<Fraction> aCoeffs = a.getCoefficients(), bCoeffs = b.getCoefficients();
+  SmallVector<Fraction> aCoeffs = a.getCoefficients(),
+                        bCoeffs = b.getCoefficients();
   EXPECT_EQ(aCoeffs.size(), bCoeffs.size());
   for (unsigned i = 0, e = aCoeffs.size(); i < e; i++)
     EXPECT_EQ(aCoeffs[i], bCoeffs[i]);
   
-  std::vector<std::vector<SmallVector<Fraction>>> aAff = a.getAffine(), bAff = b.getAffine();
+  std::vector<std::vector<SmallVector<Fraction>>> aAff = a.getAffine(),
+                                                  bAff = b.getAffine();
   EXPECT_EQ(aAff.size(), bAff.size());
   for (unsigned i = 0, e = aAff.size(); i < e; i++) {
     EXPECT_EQ(aAff[i].size(), bAff[i].size());

>From dcc0bf415c7f7e7162c62306373ae4d9033ae07c Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 19:14:54 +0530
Subject: [PATCH 18/38] Fix arith operators

---
 .../mlir/Analysis/Presburger/QuasiPolynomial.h    | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index 6c457a8d94bead..6f05216ff63199 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -45,9 +45,13 @@ class QuasiPolynomial {
     assert(numParam == x.getNumParams() &&
            "two quasi-polynomials with different numbers of parameters cannot "
            "be added!");
-    coefficients.append(x.coefficients);
-    affine.insert(affine.end(), x.affine.begin(), x.affine.end());
-    return *this;
+    SmallVector<Fraction> sumCoeffs;
+    sumCoeffs.append(coefficients);
+    sumCoeffs.append(x.coefficients);
+    std::vector<std::vector<SmallVector<Fraction>>> sumAff;
+    sumAff.insert(sumAff.end(), affine.begin(), affine.end());
+    sumAff.insert(sumAff.end(), x.affine.begin(), x.affine.end());
+    return QuasiPolynomial(numParam, sumCoeffs, sumAff);
   }
 
   QuasiPolynomial operator-(QuasiPolynomial &x) {
@@ -83,9 +87,10 @@ class QuasiPolynomial {
 
   QuasiPolynomial operator/(Fraction x) {
     assert(x != 0 && "division by zero!");
-    for (Fraction &coeff : coefficients)
+    QuasiPolynomial qp(*this);
+    for (Fraction &coeff : qp.coefficients)
       coeff /= x;
-    return *this;
+    return qp;
   };
 
   // Removes terms which evaluate to zero from the expression.

>From 7b5a469303f3a449972ff60ded85e5a560a8cff3 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 19:31:55 +0530
Subject: [PATCH 19/38] Formatting

---
 mlir/unittests/Analysis/Presburger/Utils.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/unittests/Analysis/Presburger/Utils.h b/mlir/unittests/Analysis/Presburger/Utils.h
index b0d37f27989103..6a62c56afe098d 100644
--- a/mlir/unittests/Analysis/Presburger/Utils.h
+++ b/mlir/unittests/Analysis/Presburger/Utils.h
@@ -80,7 +80,7 @@ inline void EXPECT_EQ_QUASIPOLYNOMIAL(QuasiPolynomial a, QuasiPolynomial b) {
   EXPECT_EQ(aCoeffs.size(), bCoeffs.size());
   for (unsigned i = 0, e = aCoeffs.size(); i < e; i++)
     EXPECT_EQ(aCoeffs[i], bCoeffs[i]);
-  
+
   std::vector<std::vector<SmallVector<Fraction>>> aAff = a.getAffine(),
                                                   bAff = b.getAffine();
   EXPECT_EQ(aAff.size(), bAff.size());

>From deca816d0e058dc8b114adff65a43b1bad342207 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 19:39:46 +0530
Subject: [PATCH 20/38] Formatting

---
 mlir/unittests/Analysis/Presburger/Utils.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/unittests/Analysis/Presburger/Utils.h b/mlir/unittests/Analysis/Presburger/Utils.h
index 6a62c56afe098d..223f0472704768 100644
--- a/mlir/unittests/Analysis/Presburger/Utils.h
+++ b/mlir/unittests/Analysis/Presburger/Utils.h
@@ -16,8 +16,8 @@
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
 #include "mlir/Analysis/Presburger/PWMAFunction.h"
-#include "mlir/Analysis/Presburger/QuasiPolynomial.h"
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
+#include "mlir/Analysis/Presburger/QuasiPolynomial.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Support/LLVM.h"

>From fb8a772de3e6350c44e108187ae8fb18b74f1eb1 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 19:50:09 +0530
Subject: [PATCH 21/38] Fix multiplication operator

---
 .../Analysis/Presburger/QuasiPolynomial.h     | 31 +++++++++++--------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index 6f05216ff63199..8df5c5f5ab58cb 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -65,24 +65,29 @@ class QuasiPolynomial {
   }
 
   QuasiPolynomial operator*(QuasiPolynomial &x) {
-    assert(numParam = x.getNumParams() &&
-                      "two quasi-polynomials with different numbers of "
-                      "parameters cannot be multiplied!");
-    QuasiPolynomial qp(numParam);
-    std::vector<SmallVector<Fraction>> product;
-    for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
-      for (unsigned j = 0, e = x.coefficients.size(); j < e; j++) {
-        qp.coefficients.append({coefficients[i] * x.coefficients[j]});
+    assert(numParam == x.getNumParams() &&
+           "two quasi-polynomials with different numbers of "
+           "parameters cannot be multiplied!");
 
-        product.clear();
-        product.insert(product.end(), affine[i].begin(), affine[i].end());
-        product.insert(product.end(), x.affine[j].begin(), x.affine[j].end());
+    SmallVector<Fraction> coeffs;
+    for (const Fraction &coeff : coefficients) {
+      for (const Fraction &xcoeff : x.coefficients) {
+        coeffs.append({coeff * xcoeff});
+      }
+    }
 
-        qp.affine.push_back(product);
+    std::vector<SmallVector<Fraction>> product;
+    std::vector<std::vector<SmallVector<Fraction>>> aff;
+    for (const std::vector<SmallVector<Fraction>> &term : affine) {
+      for (const std::vector<SmallVector<Fraction>> &xterm : x.affine) {
+        product.clear();
+        product.insert(product.end(), term.begin(), term.end());
+        product.insert(product.end(), xterm.begin(), xterm.end());
+        aff.push_back(product);
       }
     }
 
-    return qp;
+    return QuasiPolynomial(numParam, coeffs, aff);
   }
 
   QuasiPolynomial operator/(Fraction x) {

>From bb52a3515d4fc1b003397dd12c0344f2a8563dfe Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 19:50:20 +0530
Subject: [PATCH 22/38] Arithmetic tests for QP

---
 .../Presburger/QuasiPolynomialTest.cpp        | 98 +++++++++++++++++++
 1 file changed, 98 insertions(+)
 create mode 100644 mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp

diff --git a/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
new file mode 100644
index 00000000000000..e1b6218253416e
--- /dev/null
+++ b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
@@ -0,0 +1,98 @@
+//===- MatrixTest.cpp - Tests for QuasiPolynomial -------------------------===//
+//
+// 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 "./Utils.h"
+#include "mlir/Analysis/Presburger/Fraction.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+using namespace mlir;
+using namespace presburger;
+
+TEST(QuasiPolynomialTest, arith) {
+  QuasiPolynomial qp1(
+      3, {Fraction(1, 3), Fraction(1, 1), Fraction(1, 2)},
+      {{{Fraction(1, 1), Fraction(-1, 2), Fraction(4, 5), Fraction(0, 1)},
+        {Fraction(2, 3), Fraction(3, 4), Fraction(-1, 1), Fraction(5, 7)}},
+       {{Fraction(1, 2), Fraction(1, 1), Fraction(4, 5), Fraction(1, 1)}},
+       {{Fraction(-3, 2), Fraction(1, 1), Fraction(5, 6), Fraction(7, 5)},
+        {Fraction(1, 4), Fraction(2, 1), Fraction(6, 5), Fraction(-9, 8)},
+        {Fraction(3, 2), Fraction(2, 5), Fraction(-7, 4), Fraction(0, 1)}}});
+  QuasiPolynomial qp2(
+      3, {Fraction(1, 1), Fraction(2, 1)},
+      {{{Fraction(1, 2), Fraction(0, 1), Fraction(-1, 3), Fraction(5, 3)},
+        {Fraction(2, 1), Fraction(5, 4), Fraction(9, 7), Fraction(-1, 5)}},
+       {{Fraction(1, 3), Fraction(-2, 3), Fraction(1, 1), Fraction(0, 1)}}});
+
+  QuasiPolynomial sum = qp1 + qp2;
+  EXPECT_EQ_QUASIPOLYNOMIAL(
+      sum,
+      QuasiPolynomial(
+          3,
+          {Fraction(1, 3), Fraction(1, 1), Fraction(1, 2), Fraction(1, 1),
+           Fraction(2, 1)},
+          {{{Fraction(1, 1), Fraction(-1, 2), Fraction(4, 5), Fraction(0, 1)},
+            {Fraction(2, 3), Fraction(3, 4), Fraction(-1, 1), Fraction(5, 7)}},
+           {{Fraction(1, 2), Fraction(1, 1), Fraction(4, 5), Fraction(1, 1)}},
+           {{Fraction(-3, 2), Fraction(1, 1), Fraction(5, 6), Fraction(7, 5)},
+            {Fraction(1, 4), Fraction(2, 1), Fraction(6, 5), Fraction(-9, 8)},
+            {Fraction(3, 2), Fraction(2, 5), Fraction(-7, 4), Fraction(0, 1)}},
+           {{Fraction(1, 2), Fraction(0, 1), Fraction(-1, 3), Fraction(5, 3)},
+            {Fraction(2, 1), Fraction(5, 4), Fraction(9, 7), Fraction(-1, 5)}},
+           {{Fraction(1, 3), Fraction(-2, 3), Fraction(1, 1),
+             Fraction(0, 1)}}}));
+
+  QuasiPolynomial diff = qp1 - qp2;
+  EXPECT_EQ_QUASIPOLYNOMIAL(
+      diff,
+      QuasiPolynomial(
+          3,
+          {Fraction(1, 3), Fraction(1, 1), Fraction(1, 2), Fraction(-1, 1),
+           Fraction(-2, 1)},
+          {{{Fraction(1, 1), Fraction(-1, 2), Fraction(4, 5), Fraction(0, 1)},
+            {Fraction(2, 3), Fraction(3, 4), Fraction(-1, 1), Fraction(5, 7)}},
+           {{Fraction(1, 2), Fraction(1, 1), Fraction(4, 5), Fraction(1, 1)}},
+           {{Fraction(-3, 2), Fraction(1, 1), Fraction(5, 6), Fraction(7, 5)},
+            {Fraction(1, 4), Fraction(2, 1), Fraction(6, 5), Fraction(-9, 8)},
+            {Fraction(3, 2), Fraction(2, 5), Fraction(-7, 4), Fraction(0, 1)}},
+           {{Fraction(1, 2), Fraction(0, 1), Fraction(-1, 3), Fraction(5, 3)},
+            {Fraction(2, 1), Fraction(5, 4), Fraction(9, 7), Fraction(-1, 5)}},
+           {{Fraction(1, 3), Fraction(-2, 3), Fraction(1, 1),
+             Fraction(0, 1)}}}));
+
+  QuasiPolynomial prod = qp1 * qp2;
+  EXPECT_EQ_QUASIPOLYNOMIAL(
+      prod,
+      QuasiPolynomial(
+          3,
+          {Fraction(1, 3), Fraction(2, 3), Fraction(1, 1), Fraction(2, 1),
+           Fraction(1, 2), Fraction(1, 1)},
+          {{{Fraction(1, 1), Fraction(-1, 2), Fraction(4, 5), Fraction(0, 1)},
+            {Fraction(2, 3), Fraction(3, 4), Fraction(-1, 1), Fraction(5, 7)},
+            {Fraction(1, 2), Fraction(0, 1), Fraction(-1, 3), Fraction(5, 3)},
+            {Fraction(2, 1), Fraction(5, 4), Fraction(9, 7), Fraction(-1, 5)}},
+           {{Fraction(1, 1), Fraction(-1, 2), Fraction(4, 5), Fraction(0, 1)},
+            {Fraction(2, 3), Fraction(3, 4), Fraction(-1, 1), Fraction(5, 7)},
+            {Fraction(1, 3), Fraction(-2, 3), Fraction(1, 1), Fraction(0, 1)}},
+           {{Fraction(1, 2), Fraction(1, 1), Fraction(4, 5), Fraction(1, 1)},
+            {Fraction(1, 2), Fraction(0, 1), Fraction(-1, 3), Fraction(5, 3)},
+            {Fraction(2, 1), Fraction(5, 4), Fraction(9, 7), Fraction(-1, 5)}},
+           {{Fraction(1, 2), Fraction(1, 1), Fraction(4, 5), Fraction(1, 1)},
+            {Fraction(1, 3), Fraction(-2, 3), Fraction(1, 1), Fraction(0, 1)}},
+           {{Fraction(-3, 2), Fraction(1, 1), Fraction(5, 6), Fraction(7, 5)},
+            {Fraction(1, 4), Fraction(2, 1), Fraction(6, 5), Fraction(-9, 8)},
+            {Fraction(3, 2), Fraction(2, 5), Fraction(-7, 4), Fraction(0, 1)},
+            {Fraction(1, 2), Fraction(0, 1), Fraction(-1, 3), Fraction(5, 3)},
+            {Fraction(2, 1), Fraction(5, 4), Fraction(9, 7), Fraction(-1, 5)}},
+           {{Fraction(-3, 2), Fraction(1, 1), Fraction(5, 6), Fraction(7, 5)},
+            {Fraction(1, 4), Fraction(2, 1), Fraction(6, 5), Fraction(-9, 8)},
+            {Fraction(3, 2), Fraction(2, 5), Fraction(-7, 4), Fraction(0, 1)},
+            {Fraction(1, 3), Fraction(-2, 3), Fraction(1, 1),
+             Fraction(0, 1)}}}));
+}
\ No newline at end of file

>From 8803dbfc15e64d82540568de96d3b4a5849559d8 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 19:53:20 +0530
Subject: [PATCH 23/38] Test for division

---
 .../Analysis/Presburger/QuasiPolynomialTest.cpp      | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
index e1b6218253416e..f3e4c496349975 100644
--- a/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
@@ -95,4 +95,16 @@ TEST(QuasiPolynomialTest, arith) {
             {Fraction(3, 2), Fraction(2, 5), Fraction(-7, 4), Fraction(0, 1)},
             {Fraction(1, 3), Fraction(-2, 3), Fraction(1, 1),
              Fraction(0, 1)}}}));
+
+  QuasiPolynomial quot = qp1 / 2;
+  EXPECT_EQ_QUASIPOLYNOMIAL(
+      quot,
+      QuasiPolynomial(
+          3, {Fraction(1, 6), Fraction(1, 2), Fraction(1, 4)},
+          {{{Fraction(1, 1), Fraction(-1, 2), Fraction(4, 5), Fraction(0, 1)},
+            {Fraction(2, 3), Fraction(3, 4), Fraction(-1, 1), Fraction(5, 7)}},
+           {{Fraction(1, 2), Fraction(1, 1), Fraction(4, 5), Fraction(1, 1)}},
+           {{Fraction(-3, 2), Fraction(1, 1), Fraction(5, 6), Fraction(7, 5)},
+            {Fraction(3, 2), Fraction(2, 5), Fraction(-7, 4),
+             Fraction(0, 1)}}}));
 }
\ No newline at end of file

>From 5666b5c74926e2dd589fa43770a97167050bf768 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 20:04:41 +0530
Subject: [PATCH 24/38] Fix test and add simplify test

---
 .../Presburger/QuasiPolynomialTest.cpp        | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
index f3e4c496349975..120c683b6b86e2 100644
--- a/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
@@ -105,6 +105,26 @@ TEST(QuasiPolynomialTest, arith) {
             {Fraction(2, 3), Fraction(3, 4), Fraction(-1, 1), Fraction(5, 7)}},
            {{Fraction(1, 2), Fraction(1, 1), Fraction(4, 5), Fraction(1, 1)}},
            {{Fraction(-3, 2), Fraction(1, 1), Fraction(5, 6), Fraction(7, 5)},
+            {Fraction(1, 4), Fraction(2, 1), Fraction(6, 5), Fraction(-9, 8)},
             {Fraction(3, 2), Fraction(2, 5), Fraction(-7, 4),
              Fraction(0, 1)}}}));
+}
+
+TEST(QuasiPolynomialTest, simplify) {
+  QuasiPolynomial qp(2,
+                     {Fraction(2, 3), Fraction(0, 1), Fraction(1, 1),
+                      Fraction(1, 2), Fraction(0, 1)},
+                     {{{Fraction(1, 1), Fraction(3, 4), Fraction(5, 3)},
+                       {Fraction(2, 1), Fraction(0, 1), Fraction(0, 1)}},
+                      {{Fraction(1, 3), Fraction(8, 5), Fraction(2, 5)}},
+                      {{Fraction(2, 7), Fraction(9, 5), Fraction(0, 1)},
+                       {Fraction(0, 1), Fraction(0, 1), Fraction(0, 1)}},
+                      {{Fraction(1, 1), Fraction(4, 5), Fraction(6, 5)}},
+                      {{Fraction(1, 3), Fraction(4, 3), Fraction(7, 8)}}});
+  EXPECT_EQ_QUASIPOLYNOMIAL(
+      qp.simplify(),
+      QuasiPolynomial(2, {Fraction(2, 3), Fraction(1, 2)},
+                      {{{Fraction(1, 1), Fraction(3, 4), Fraction(5, 3)},
+                        {Fraction(2, 1), Fraction(0, 1), Fraction(0, 1)}},
+                       {{Fraction(1, 1), Fraction(4, 5), Fraction(6, 5)}}}));
 }
\ No newline at end of file

>From 47cbb9492153e4d99e473726ba3ef1716e3762a9 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 20:19:50 +0530
Subject: [PATCH 25/38] Fix GF addition

---
 mlir/lib/Analysis/Presburger/GeneratingFunction.h | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/GeneratingFunction.h b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
index 11832e822dc021..6b93212f7df2c1 100644
--- a/mlir/lib/Analysis/Presburger/GeneratingFunction.h
+++ b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
@@ -73,12 +73,16 @@ class GeneratingFunction {
     assert(numParam == gf.getNumParams() &&
            "two generating functions with different numbers of parameters "
            "cannot be added!");
-    signs.append(gf.signs);
-    numerators.insert(numerators.end(), gf.numerators.begin(),
-                      gf.numerators.end());
-    denominators.insert(denominators.end(), gf.denominators.begin(),
+    SmallVector<int> sumSigns(signs);
+    sumSigns.append(gf.signs);
+
+    std::vector<ParamPoint> sumNumerators(numerators);
+    sumNumerators.insert(sumNumerators.end(), gf.numerators.begin(), gf.numerators.end());
+
+    std::vector<std::vector<Point>> sumDenominators(denominators);
+    sumDenominators.insert(sumDenominators.end(), gf.denominators.begin(),
                         gf.denominators.end());
-    return *this;
+    return GeneratingFunction(sumSigns, sumNumerators, sumDenominators);
   }
 
   llvm::raw_ostream &print(llvm::raw_ostream &os) const {

>From 419c410fa4d2e84b96a17a9e4adc7b93030b8e9c Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Tue, 19 Dec 2023 20:23:24 +0530
Subject: [PATCH 26/38] Formatting

---
 mlir/lib/Analysis/Presburger/GeneratingFunction.h | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/GeneratingFunction.h b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
index 6b93212f7df2c1..f8a30c223f1685 100644
--- a/mlir/lib/Analysis/Presburger/GeneratingFunction.h
+++ b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
@@ -77,11 +77,12 @@ class GeneratingFunction {
     sumSigns.append(gf.signs);
 
     std::vector<ParamPoint> sumNumerators(numerators);
-    sumNumerators.insert(sumNumerators.end(), gf.numerators.begin(), gf.numerators.end());
+    sumNumerators.insert(sumNumerators.end(), gf.numerators.begin(),
+                         gf.numerators.end());
 
     std::vector<std::vector<Point>> sumDenominators(denominators);
     sumDenominators.insert(sumDenominators.end(), gf.denominators.begin(),
-                        gf.denominators.end());
+                           gf.denominators.end());
     return GeneratingFunction(sumSigns, sumNumerators, sumDenominators);
   }
 

>From c46cb868dd7fabb560d56c514515eec686a324bc Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 20 Dec 2023 17:54:12 +0530
Subject: [PATCH 27/38] Fix comments

---
 .../mlir/Analysis/Presburger/QuasiPolynomial.h     | 12 +++++++-----
 mlir/lib/Analysis/Presburger/GeneratingFunction.h  | 14 ++++++--------
 mlir/unittests/Analysis/Presburger/Utils.h         |  2 ++
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index 8df5c5f5ab58cb..3e44bb689c5bbd 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -1,4 +1,4 @@
-//===- QuasiPolynomial.h - Quasipolynomial Class ----------------*- C++ -*-===//
+//===- QuasiPolynomial.h - 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.
@@ -8,6 +8,10 @@
 //
 // Definition of the QuasiPolynomial class for Barvinok's algorithm,
 // which represents a single-valued function on a set of parameters.
+// It is an expression of the form
+// f(x) = \sum_i c_i * \prod_j ⌊g_{ij}(x)⌋
+// where c_i \in Q and
+// g_{ij} : Q^d -> Q are affine functionals over d parameters.
 //
 //===----------------------------------------------------------------------===//
 
@@ -20,10 +24,8 @@
 namespace mlir {
 namespace presburger {
 
-// A class to describe the quasi-polynomials obtained by
-// substituting the unit vector in the type of generating
-// function described above.
-// Consists of a set of terms.
+// A class to describe quasi-polynomials.
+// A quasipolynomial consists of a set of terms.
 // The ith term is a constant `coefficients[i]`, multiplied
 // by the product of a set of affine functions on n parameters.
 // Represents functions f : Q^n -> Q of the form
diff --git a/mlir/lib/Analysis/Presburger/GeneratingFunction.h b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
index f8a30c223f1685..264c185db6865f 100644
--- a/mlir/lib/Analysis/Presburger/GeneratingFunction.h
+++ b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
@@ -1,4 +1,4 @@
-//===- GeneratingFunction.h - Rational Polynomials over Q^d -----*- C++ -*-===//
+//===- GeneratingFunction.h - Generating Functions over Q^d -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -35,17 +35,17 @@ using Point = SmallVector<Fraction>;
 // * a sign, ±1, stored in `signs[i]`
 // * a numerator, of the form x^{n},
 //      where n, stored in `numerators[i]`,
-//      is a parametric point (a vertex).
+//      is a parametric point.
 // * a denominator, of the form (1 - x^{d1})...(1 - x^{dn}),
 //      where each dj, stored in `denominators[i][j]`,
-//      is a vector (a generator).
+//      is a vector.
 //
-// Represents functions f : Q^n -> Q of the form
+// Represents functions f_p : Q^n -> Q of the form
 //
-// f(x) = \sum_i s_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
+// f_p(x) = \sum_i s_i * (x^n_i(p)) / (\prod_j (1 - x^d_{ij})
 //
 // where s_i is ±1,
-// n_i \in (Q^d -> Q)^n is an n-vector of affine functions on d parameters, and
+// n_i \in Q^d -> Q^n is an n-vector of affine functions on d parameters, and
 // g_{ij} \in Q^n are vectors.
 class GeneratingFunction {
 public:
@@ -59,8 +59,6 @@ class GeneratingFunction {
              "parameters!");
   }
 
-  // Find the number of parameters involved in the function
-  // from the dimensionality of the affine functions.
   unsigned getNumParams() { return numParam; }
 
   SmallVector<int> getSigns() { return signs; }
diff --git a/mlir/unittests/Analysis/Presburger/Utils.h b/mlir/unittests/Analysis/Presburger/Utils.h
index 223f0472704768..b6d4344e4c8cc0 100644
--- a/mlir/unittests/Analysis/Presburger/Utils.h
+++ b/mlir/unittests/Analysis/Presburger/Utils.h
@@ -72,6 +72,8 @@ inline void EXPECT_EQ_FRAC_MATRIX(FracMatrix a, FracMatrix b) {
       EXPECT_EQ(a(row, col), b(row, col));
 }
 
+// Check the coefficients (in order) of two quasipolynomials.
+// Note that this is not a true equality check.
 inline void EXPECT_EQ_QUASIPOLYNOMIAL(QuasiPolynomial a, QuasiPolynomial b) {
   EXPECT_EQ(a.getNumParams(), b.getNumParams());
 

>From 5bbd0f21208619e1ed88512eb70869172a700746 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 20 Dec 2023 17:59:47 +0530
Subject: [PATCH 28/38] Use const

---
 mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h | 8 ++++----
 mlir/lib/Analysis/Presburger/GeneratingFunction.h       | 2 +-
 mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp        | 5 +++--
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index 3e44bb689c5bbd..0a13df5c9ca520 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -43,7 +43,7 @@ class QuasiPolynomial {
   SmallVector<Fraction> getCoefficients();
   std::vector<std::vector<SmallVector<Fraction>>> getAffine();
 
-  QuasiPolynomial operator+(QuasiPolynomial &x) {
+  QuasiPolynomial operator+(const QuasiPolynomial &x) const {
     assert(numParam == x.getNumParams() &&
            "two quasi-polynomials with different numbers of parameters cannot "
            "be added!");
@@ -56,7 +56,7 @@ class QuasiPolynomial {
     return QuasiPolynomial(numParam, sumCoeffs, sumAff);
   }
 
-  QuasiPolynomial operator-(QuasiPolynomial &x) {
+  QuasiPolynomial operator-(const QuasiPolynomial &x) const {
     assert(numParam == x.getNumParams() &&
            "two quasi-polynomials with different numbers of parameters cannot "
            "be subtracted!");
@@ -66,7 +66,7 @@ class QuasiPolynomial {
     return *this + qp;
   }
 
-  QuasiPolynomial operator*(QuasiPolynomial &x) {
+  QuasiPolynomial operator*(const QuasiPolynomial &x) const {
     assert(numParam == x.getNumParams() &&
            "two quasi-polynomials with different numbers of "
            "parameters cannot be multiplied!");
@@ -92,7 +92,7 @@ class QuasiPolynomial {
     return QuasiPolynomial(numParam, coeffs, aff);
   }
 
-  QuasiPolynomial operator/(Fraction x) {
+  QuasiPolynomial operator/(const Fraction x) const {
     assert(x != 0 && "division by zero!");
     QuasiPolynomial qp(*this);
     for (Fraction &coeff : qp.coefficients)
diff --git a/mlir/lib/Analysis/Presburger/GeneratingFunction.h b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
index 264c185db6865f..630a2bad26ec49 100644
--- a/mlir/lib/Analysis/Presburger/GeneratingFunction.h
+++ b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
@@ -67,7 +67,7 @@ class GeneratingFunction {
 
   std::vector<std::vector<Point>> getDenominators() { return denominators; }
 
-  GeneratingFunction operator+(const GeneratingFunction &gf) {
+  GeneratingFunction operator+(const GeneratingFunction &gf) const {
     assert(numParam == gf.getNumParams() &&
            "two generating functions with different numbers of parameters "
            "cannot be added!");
diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index 14cfdcad0fb255..5524be392e184e 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -53,9 +53,10 @@ QuasiPolynomial QuasiPolynomial::simplify() {
     // A term is zero if its coefficient is zero, or
     if (coefficients[i] == Fraction(0, 1) ||
         // if any of the affine functions in the product
-        llvm::any_of(affine[i], [](SmallVector<Fraction> affine_ij) {
+        llvm::any_of(affine[i], [](const SmallVector<Fraction> &affine_ij) {
           // has all its coefficients as zero.
-          return llvm::all_of(affine_ij, [](Fraction f) { return f == 0; });
+          return llvm::all_of(affine_ij,
+                              [](const Fraction &f) { return f == 0; });
         }))
       continue;
     newCoeffs.append({coefficients[i]});

>From dee486f68adbaa99900c63f4249aaad2b5bfdea8 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 20 Dec 2023 18:12:10 +0530
Subject: [PATCH 29/38] Fix vector definitions

---
 .../mlir/Analysis/Presburger/QuasiPolynomial.h   | 16 ++++++++--------
 .../lib/Analysis/Presburger/GeneratingFunction.h |  6 +++---
 mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp | 11 ++++++-----
 3 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index 0a13df5c9ca520..52a5bd5d4611ba 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -39,19 +39,17 @@ class QuasiPolynomial {
   QuasiPolynomial(unsigned numParam, SmallVector<Fraction> coeffs = {},
                   std::vector<std::vector<SmallVector<Fraction>>> aff = {});
 
-  unsigned getNumParams();
-  SmallVector<Fraction> getCoefficients();
-  std::vector<std::vector<SmallVector<Fraction>>> getAffine();
+  unsigned getNumParams() const;
+  SmallVector<Fraction> getCoefficients() const;
+  std::vector<std::vector<SmallVector<Fraction>>> getAffine() const;
 
   QuasiPolynomial operator+(const QuasiPolynomial &x) const {
     assert(numParam == x.getNumParams() &&
            "two quasi-polynomials with different numbers of parameters cannot "
            "be added!");
-    SmallVector<Fraction> sumCoeffs;
-    sumCoeffs.append(coefficients);
+    SmallVector<Fraction> sumCoeffs = coefficients;
     sumCoeffs.append(x.coefficients);
-    std::vector<std::vector<SmallVector<Fraction>>> sumAff;
-    sumAff.insert(sumAff.end(), affine.begin(), affine.end());
+    std::vector<std::vector<SmallVector<Fraction>>> sumAff = affine;
     sumAff.insert(sumAff.end(), x.affine.begin(), x.affine.end());
     return QuasiPolynomial(numParam, sumCoeffs, sumAff);
   }
@@ -72,14 +70,16 @@ class QuasiPolynomial {
            "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.append({coeff * xcoeff});
+        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) {
       for (const std::vector<SmallVector<Fraction>> &xterm : x.affine) {
         product.clear();
diff --git a/mlir/lib/Analysis/Presburger/GeneratingFunction.h b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
index 630a2bad26ec49..19fd3b429d79da 100644
--- a/mlir/lib/Analysis/Presburger/GeneratingFunction.h
+++ b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
@@ -71,14 +71,14 @@ class GeneratingFunction {
     assert(numParam == gf.getNumParams() &&
            "two generating functions with different numbers of parameters "
            "cannot be added!");
-    SmallVector<int> sumSigns(signs);
+    SmallVector<int> sumSigns = signs;
     sumSigns.append(gf.signs);
 
-    std::vector<ParamPoint> sumNumerators(numerators);
+    std::vector<ParamPoint> sumNumerators = numerators;
     sumNumerators.insert(sumNumerators.end(), gf.numerators.begin(),
                          gf.numerators.end());
 
-    std::vector<std::vector<Point>> sumDenominators(denominators);
+    std::vector<std::vector<Point>> sumDenominators = denominators;
     sumDenominators.insert(sumDenominators.end(), gf.denominators.begin(),
                            gf.denominators.end());
     return GeneratingFunction(sumSigns, sumNumerators, sumDenominators);
diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index 5524be392e184e..e9aa852b1de343 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -35,13 +35,14 @@ QuasiPolynomial::QuasiPolynomial(
 }
 
 // Find the number of parameters involved in the polynomial.
-unsigned QuasiPolynomial::getNumParams() { return numParam; }
+unsigned QuasiPolynomial::getNumParams() const { return numParam; }
 
-SmallVector<Fraction> QuasiPolynomial::getCoefficients() {
+SmallVector<Fraction> QuasiPolynomial::getCoefficients() const {
   return coefficients;
 }
 
-std::vector<std::vector<SmallVector<Fraction>>> QuasiPolynomial::getAffine() {
+std::vector<std::vector<SmallVector<Fraction>>>
+QuasiPolynomial::getAffine() const {
   return affine;
 }
 
@@ -59,8 +60,8 @@ QuasiPolynomial QuasiPolynomial::simplify() {
                               [](const Fraction &f) { return f == 0; });
         }))
       continue;
-    newCoeffs.append({coefficients[i]});
-    newAffine.push_back({affine[i]});
+    newCoeffs.push_back(coefficients[i]);
+    newAffine.push_back(affine[i]);
   }
   return QuasiPolynomial(numParam, newCoeffs, newAffine);
 }
\ No newline at end of file

>From ab1d13064e86ffe236a0a666f21f26c9564b6afc Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 20 Dec 2023 18:19:00 +0530
Subject: [PATCH 30/38] Comments on tests

---
 .../Analysis/Presburger/QuasiPolynomialTest.cpp        | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
index 120c683b6b86e2..4f86a1238c6924 100644
--- a/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
@@ -15,6 +15,11 @@
 using namespace mlir;
 using namespace presburger;
 
+// Test the arithmetic operations on QuasiPolynomials;
+// addition, subtraction, multiplication, and division
+// by a constant.
+// Two QPs of 3 parameters each were generated randomly
+// and their sum, difference, and product computed by hand.
 TEST(QuasiPolynomialTest, arith) {
   QuasiPolynomial qp1(
       3, {Fraction(1, 3), Fraction(1, 1), Fraction(1, 2)},
@@ -110,6 +115,11 @@ TEST(QuasiPolynomialTest, arith) {
              Fraction(0, 1)}}}));
 }
 
+// Test the simplify() operation on QPs, which removes terms that
+// are identically zero. A random QP was generated and terms were
+// changed to account for each condition in simplify() – 
+// the term coefficient being zero, or all the coefficients in some
+// affine term in the product being zero.
 TEST(QuasiPolynomialTest, simplify) {
   QuasiPolynomial qp(2,
                      {Fraction(2, 3), Fraction(0, 1), Fraction(1, 1),

>From 1a3575e3d32a9fe380b3395b44e48dc1b74e989e Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 20 Dec 2023 18:20:30 +0530
Subject: [PATCH 31/38] Fix name of test

---
 mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
index 4f86a1238c6924..b8ffa7648646b8 100644
--- a/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp
@@ -20,7 +20,7 @@ using namespace presburger;
 // by a constant.
 // Two QPs of 3 parameters each were generated randomly
 // and their sum, difference, and product computed by hand.
-TEST(QuasiPolynomialTest, arith) {
+TEST(QuasiPolynomialTest, arithmetic) {
   QuasiPolynomial qp1(
       3, {Fraction(1, 3), Fraction(1, 1), Fraction(1, 2)},
       {{{Fraction(1, 1), Fraction(-1, 2), Fraction(4, 5), Fraction(0, 1)},

>From 712404be88769545ee336d20e8cf48bd2701cd8c Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 20 Dec 2023 18:22:49 +0530
Subject: [PATCH 32/38] Break down condition

---
 mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index e9aa852b1de343..974af1858b5713 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -52,13 +52,16 @@ QuasiPolynomial QuasiPolynomial::simplify() {
   std::vector<std::vector<SmallVector<Fraction>>> newAffine({});
   for (unsigned i = 0, e = coefficients.size(); i < e; i++) {
     // A term is zero if its coefficient is zero, or
-    if (coefficients[i] == Fraction(0, 1) ||
+    if (coefficients[i] == Fraction(0, 1))
+      continue;
+    bool product_is_zero =
         // if any of the affine functions in the product
         llvm::any_of(affine[i], [](const SmallVector<Fraction> &affine_ij) {
           // has all its coefficients as zero.
           return llvm::all_of(affine_ij,
                               [](const Fraction &f) { return f == 0; });
-        }))
+        });
+    if (product_is_zero)
       continue;
     newCoeffs.push_back(coefficients[i]);
     newAffine.push_back(affine[i]);

>From 22b8429030fa0304405d42db721e49c6d5d2d30f Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 20 Dec 2023 18:25:00 +0530
Subject: [PATCH 33/38] Check all terms in assert

---
 mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index 974af1858b5713..6e96928bef1d00 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -26,11 +26,13 @@ QuasiPolynomial::QuasiPolynomial(
   for (const std::vector<SmallVector<Fraction>> &term : affine) {
     if (term.size() == 0)
       continue;
-    // The number of elements in the affine function is
+    // The number of elements in each affine function is
     // one more than the number of parameters.
-    assert(term[0].size() - 1 == numParam &&
-           "dimensionality of affine functions does not match number of "
-           "parameters!");
+    for (const SmallVector<Fraction> &aff : term) {
+      assert(aff.size() - 1 == numParam &&
+             "dimensionality of affine functions does not match number of "
+             "parameters!");
+    }
   }
 }
 

>From e6e05e12dbbe76183ed884e63001ef0f1eb48e80 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Wed, 20 Dec 2023 18:29:11 +0530
Subject: [PATCH 34/38] Move getters to header and operators to cpp

---
 .../Analysis/Presburger/QuasiPolynomial.h     | 66 +++----------------
 .../Analysis/Presburger/QuasiPolynomial.cpp   | 59 +++++++++++++++--
 2 files changed, 62 insertions(+), 63 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index 52a5bd5d4611ba..5ccee451c0b11f 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -39,66 +39,20 @@ class QuasiPolynomial {
   QuasiPolynomial(unsigned numParam, SmallVector<Fraction> coeffs = {},
                   std::vector<std::vector<SmallVector<Fraction>>> aff = {});
 
-  unsigned getNumParams() const;
-  SmallVector<Fraction> getCoefficients() const;
-  std::vector<std::vector<SmallVector<Fraction>>> getAffine() const;
+  // Find the number of parameters involved in the polynomial.
+  unsigned getNumParams() const { return numParam; }
 
-  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 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 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) {
-      for (const std::vector<SmallVector<Fraction>> &xterm : x.affine) {
-        product.clear();
-        product.insert(product.end(), term.begin(), term.end());
-        product.insert(product.end(), xterm.begin(), xterm.end());
-        aff.push_back(product);
-      }
-    }
+  SmallVector<Fraction> getCoefficients() const { return coefficients; }
 
-    return QuasiPolynomial(numParam, coeffs, aff);
+  std::vector<std::vector<SmallVector<Fraction>>> getAffine() const {
+    return affine;
   }
 
-  QuasiPolynomial operator/(const Fraction x) const {
-    assert(x != 0 && "division by zero!");
-    QuasiPolynomial qp(*this);
-    for (Fraction &coeff : qp.coefficients)
-      coeff /= x;
-    return qp;
-  };
+  // Arithmetic operations.
+  QuasiPolynomial operator+(const QuasiPolynomial &x) const;
+  QuasiPolynomial operator-(const QuasiPolynomial &x) const;
+  QuasiPolynomial operator*(const QuasiPolynomial &x) const;
+  QuasiPolynomial operator/(const Fraction x) const;
 
   // Removes terms which evaluate to zero from the expression.
   QuasiPolynomial simplify();
diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index 6e96928bef1d00..d56377cc119cbd 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -36,16 +36,61 @@ QuasiPolynomial::QuasiPolynomial(
   }
 }
 
-// Find the number of parameters involved in the polynomial.
-unsigned QuasiPolynomial::getNumParams() const { return numParam; }
+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) {
+    for (const std::vector<SmallVector<Fraction>> &xterm : x.affine) {
+      product.clear();
+      product.insert(product.end(), term.begin(), term.end());
+      product.insert(product.end(), xterm.begin(), xterm.end());
+      aff.push_back(product);
+    }
+  }
 
-SmallVector<Fraction> QuasiPolynomial::getCoefficients() const {
-  return coefficients;
+  return QuasiPolynomial(numParam, coeffs, aff);
 }
 
-std::vector<std::vector<SmallVector<Fraction>>>
-QuasiPolynomial::getAffine() const {
-  return affine;
+QuasiPolynomial QuasiPolynomial::operator/(const Fraction x) const {
+  assert(x != 0 && "division by zero!");
+  QuasiPolynomial qp(*this);
+  for (Fraction &coeff : qp.coefficients)
+    coeff /= x;
+  return qp;
 }
 
 // Removes terms which evaluate to zero from the expression.

>From 94ec54b734c9d7f6c0df883d2e5779d2cbf26fe5 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Fri, 22 Dec 2023 10:35:00 +0530
Subject: [PATCH 35/38] Fix QP initialization

---
 mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index d56377cc119cbd..ff703f32ae58cc 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -22,14 +22,14 @@ 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 each term which involves at least one affine function,
   for (const std::vector<SmallVector<Fraction>> &term : affine) {
     if (term.size() == 0)
       continue;
-    // The number of elements in each affine function is
+    // 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 &&
+      assert(aff.size() == numParam + 1 &&
              "dimensionality of affine functions does not match number of "
              "parameters!");
     }

>From 59f697aa7eef5218d3c9b08180de967ec105452f Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Fri, 22 Dec 2023 10:37:18 +0530
Subject: [PATCH 36/38] Return const refs in getters

---
 mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
index 5ccee451c0b11f..40c3c4529c8eb9 100644
--- a/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
+++ b/mlir/include/mlir/Analysis/Presburger/QuasiPolynomial.h
@@ -40,11 +40,11 @@ class QuasiPolynomial {
                   std::vector<std::vector<SmallVector<Fraction>>> aff = {});
 
   // Find the number of parameters involved in the polynomial.
-  unsigned getNumParams() const { return numParam; }
+  const unsigned &getNumParams() const { return numParam; }
 
-  SmallVector<Fraction> getCoefficients() const { return coefficients; }
+  const SmallVector<Fraction> &getCoefficients() const { return coefficients; }
 
-  std::vector<std::vector<SmallVector<Fraction>>> getAffine() const {
+  const std::vector<std::vector<SmallVector<Fraction>>> &getAffine() const {
     return affine;
   }
 

>From 3eb698bb207ed48765d44c42bda9ba0a3462a89b Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Fri, 22 Dec 2023 10:42:54 +0530
Subject: [PATCH 37/38] Fix loop braces

---
 mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp | 6 ++----
 mlir/unittests/Analysis/Presburger/Utils.h       | 3 +--
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
index ff703f32ae58cc..e54803beb1571e 100644
--- a/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
+++ b/mlir/lib/Analysis/Presburger/QuasiPolynomial.cpp
@@ -64,11 +64,9 @@ QuasiPolynomial QuasiPolynomial::operator*(const QuasiPolynomial &x) const {
 
   SmallVector<Fraction> coeffs;
   coeffs.reserve(coefficients.size() * x.coefficients.size());
-  for (const Fraction &coeff : coefficients) {
-    for (const Fraction &xcoeff : x.coefficients) {
+  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;
diff --git a/mlir/unittests/Analysis/Presburger/Utils.h b/mlir/unittests/Analysis/Presburger/Utils.h
index b6d4344e4c8cc0..923db6a207ee04 100644
--- a/mlir/unittests/Analysis/Presburger/Utils.h
+++ b/mlir/unittests/Analysis/Presburger/Utils.h
@@ -88,10 +88,9 @@ inline void EXPECT_EQ_QUASIPOLYNOMIAL(QuasiPolynomial a, QuasiPolynomial b) {
   EXPECT_EQ(aAff.size(), bAff.size());
   for (unsigned i = 0, e = aAff.size(); i < e; i++) {
     EXPECT_EQ(aAff[i].size(), bAff[i].size());
-    for (unsigned j = 0, f = aAff[i].size(); j < f; j++) {
+    for (unsigned j = 0, f = aAff[i].size(); j < f; j++)
       for (unsigned k = 0, g = a.getNumParams(); k <= g; k++)
         EXPECT_EQ(aAff[i][j][k], bAff[i][j][k]);
-    }
   }
 }
 

>From 1bada4e12148aed46081fb9b3c3cdacf779efaf8 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Fri, 22 Dec 2023 11:07:30 +0530
Subject: [PATCH 38/38] Shift -1 to +1 on unsigned int

---
 mlir/lib/Analysis/Presburger/GeneratingFunction.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Analysis/Presburger/GeneratingFunction.h b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
index 19fd3b429d79da..33bb97c59fe2c2 100644
--- a/mlir/lib/Analysis/Presburger/GeneratingFunction.h
+++ b/mlir/lib/Analysis/Presburger/GeneratingFunction.h
@@ -54,7 +54,7 @@ class GeneratingFunction {
                      std::vector<std::vector<Point>> dens)
       : numParam(numParam), signs(signs), numerators(nums), denominators(dens) {
     for (const ParamPoint &term : numerators)
-      assert(term.getNumColumns() - 1 == numParam &&
+      assert(term.getNumColumns() == numParam + 1 &&
              "dimensionality of numerator exponents does not match number of "
              "parameters!");
   }