[Mlir-commits] [mlir] [MLIR][Presburger][WIP] Implement vertex enumeration and chamber decomposition for polytope generating function computation. (PR #78987)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Mon Jan 22 06:57:01 PST 2024


https://github.com/Abhinav271828 created https://github.com/llvm/llvm-project/pull/78987

We implement a function to compute the generating function corresponding to a full-dimensional parametric polytope whose tangent cones are all unimodular.

>From be5fa415840591c7e001c18cbe0c3f2e0c9aa274 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 22 Jan 2024 14:56:23 +0530
Subject: [PATCH 1/6] Initial commit

---
 .../mlir/Analysis/Presburger/Barvinok.h       |   5 +
 .../Analysis/Presburger/GeneratingFunction.h  |   2 +-
 .../Analysis/Presburger/IntegerRelation.h     |   4 +
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 217 ++++++++++++++++++
 .../Analysis/Presburger/IntegerRelation.cpp   |  36 +++
 5 files changed, 263 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index b70ec33b693235..6feb7e8f09233f 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -27,6 +27,7 @@
 #include "mlir/Analysis/Presburger/GeneratingFunction.h"
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Analysis/Presburger/PresburgerRelation.h"
 #include "mlir/Analysis/Presburger/QuasiPolynomial.h"
 #include <optional>
 
@@ -84,6 +85,10 @@ ConeH getDual(ConeV cone);
 GeneratingFunction unimodularConeGeneratingFunction(ParamPoint vertex, int sign,
                                                     ConeH cone);
 
+std::optional<ParamPoint> findVertex(Matrix<MPInt> equations);
+
+std::vector<std::pair<PresburgerRelation, GeneratingFunction>> polytopeGeneratingFunction(PolyhedronH poly);
+
 /// Find a vector that is not orthogonal to any of the given vectors,
 /// i.e., has nonzero dot product with those of the given vectors
 /// that are not null.
diff --git a/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h b/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
index c38eab6efd0fc1..35f15319b4568e 100644
--- a/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
+++ b/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
@@ -72,7 +72,7 @@ class GeneratingFunction {
     return denominators;
   }
 
-  GeneratingFunction operator+(GeneratingFunction &gf) const {
+  GeneratingFunction operator+(GeneratingFunction gf) const {
     assert(numParam == gf.getNumParams() &&
            "two generating functions with different numbers of parameters "
            "cannot be added!");
diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index c476a022a48272..e7b03c44d33e38 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -711,6 +711,10 @@ class IntegerRelation {
   /// return `this \ set`.
   PresburgerRelation subtract(const PresburgerRelation &set) const;
 
+  void removeTrivialEqualities();
+
+  bool isFullDim();
+
   void print(raw_ostream &os) const;
   void dump() const;
 
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index e0fd0dd8caa4d3..c8cec45a21dc4b 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -10,6 +10,7 @@
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/Sequence.h"
 #include <algorithm>
+#include <bitset>
 
 using namespace mlir;
 using namespace presburger;
@@ -147,6 +148,222 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
                             std::vector({denominator}));
 }
 
+std::optional<ParamPoint> mlir::presburger::detail::findVertex(Matrix<MPInt> equations)
+{
+    // `equalities` is a d x (d + p + 1) matrix.
+
+    unsigned r = equations.getNumRows();
+    unsigned c = equations.getNumColumns();
+
+    IntMatrix coeffs(r, r);
+    for (unsigned i = 0; i < r; i++)
+        for (unsigned j = 0; j < r; j++)
+            coeffs(i, j) = equations(i, j);
+    
+    if (coeffs.determinant() == MPInt(0))
+        return std::nullopt;
+
+    Matrix<Fraction> equationsF(r, c);
+    for (unsigned i = 0; i < r; i++)
+        for (unsigned j = 0; j < c; j++)
+            equationsF(i, j) = Fraction(equations(i, j), 1);
+    
+    Fraction a, b;
+    for (unsigned i = 0; i < r; i++)
+    {
+        if (equationsF(i, i) == Fraction(0, 1))
+            for (unsigned j = i+1; j < r; j++)
+                if (equationsF(j, i) != 0)
+                {
+                    equationsF.addToRow(i, equationsF.getRow(j), Fraction(1, 1));
+                    break;
+                }
+        b = equationsF(i, i);
+
+        for (unsigned j = 0; j < r; j++)
+        {
+            if (equationsF(j, i) == 0 || j == i) continue;
+            a = equationsF(j, i);
+            equationsF.addToRow(j, equationsF.getRow(i), - a / b);
+        }
+    }
+
+    for (unsigned i = 0; i < r; i++)
+    {
+        a = equationsF(i, i);
+        for (unsigned j = 0; j < c; j++)
+            equationsF(i, j) = equationsF(i, j) / a;
+    }
+
+    ParamPoint vertex(r, c-r); // d x p+1
+    for (unsigned i = 0; i < r; i++)
+        for (unsigned j = 0; j < c-r; j++)
+            vertex(i, j) = -equationsF(i, r+j);
+    
+    return vertex;
+}
+
+std::vector<std::pair<PresburgerRelation, GeneratingFunction>> mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly)
+{
+    unsigned d = poly.getNumRangeVars(); 
+    unsigned p = poly.getNumSymbolVars();
+    unsigned n = poly.getNumInequalities();
+
+    SmallVector<std::pair<int, ConeH>, 4> unimodCones;
+    GeneratingFunction chamberGf(p, {}, {}, {});
+    std::vector<std::pair<PresburgerRelation, GeneratingFunction>> gf({});
+    ConeH tgtCone = defineHRep(d);
+
+    Matrix<MPInt> subset(d, d+p+1);
+    std::vector<Matrix<MPInt>> subsets; // Stores the inequality subsets corresponding to each vertex.
+    Matrix<Fraction> remaining(n-d, d+p+1);
+
+    std::optional<ParamPoint> vertex;
+    std::vector<ParamPoint> vertices;
+
+    Matrix<Fraction> a2(n-d, d);
+    Matrix<Fraction> b2c2(n-d, p+1);
+
+    Matrix<Fraction> activeRegion(n-d, p+1);
+    Matrix<MPInt> activeRegionNorm(n-d, p+1);
+    MPInt lcmDenoms;
+    IntegerRelation activeRegionRel(PresburgerSpace::getRelationSpace(0, p, 0, 0));
+    // The active region will be defined as activeRegionCoeffs @ p + activeRegionConstant ≥ 0.
+    // The active region is a polyhedron in parameter space.
+    std::vector<PresburgerRelation> activeRegions;
+    
+
+    for (std::bitset<16> indicator(((1ul << d)-1ul) << (n-d));
+        indicator.to_ulong() <= ((1ul << d)-1ul) << (n-d);              // d 1's followed by n-d 0's
+        indicator = std::bitset<16>(indicator.to_ulong() - 1))
+    {
+        if (indicator.count() != d)
+            continue;
+
+        subset = Matrix<MPInt>(d, d+p+1);
+        remaining = Matrix<Fraction>(n-d, d+p+1);
+        unsigned j1 = 0, j2 = 0;
+        for (unsigned i = 0; i < n; i++)
+            if (indicator.test(i))
+                subset.setRow(j1++, poly.getInequality(i));
+                // [A1 | B1 | c1]
+            else
+            {
+                for (unsigned k = 0; k < d; k++)
+                    a2(j2, k) = Fraction(poly.atIneq(i, k), 1);
+                for (unsigned k = d; k < d+p+1; k++)
+                    b2c2(j2, k-d) = Fraction(poly.atIneq(i, k), 1);
+                j2++;
+                // [A2 | B2 | c2]
+            }
+
+        vertex = findVertex(subset); // d x (p+1)
+
+        if (vertex == std::nullopt) continue;
+        vertices.push_back(*vertex);
+        subsets.push_back(subset);
+
+        // Region is given by (A2 @ X + B2) p + (A2 @ y + c2) ≥ 0
+        // This is equivt to A2 @ [X | y] + [B2 | c2]
+        // We premultiply [X | y] with each row of A2 and add each row of [B2 | c2].
+        for (unsigned i = 0; i < n-d; i++)
+        {
+            activeRegion.setRow(i, (*vertex).preMultiplyWithRow(a2.getRow(i)));
+            activeRegion.addToRow(i, b2c2.getRow(i), Fraction(1, 1));
+        }
+
+        activeRegionNorm = Matrix<MPInt>(n-d, p+1);
+        activeRegionRel = IntegerRelation(PresburgerSpace::getRelationSpace(0, p, 0, 0));
+        lcmDenoms = 1;
+        for (unsigned i = 0; i < n-d; i++)
+        {
+            for (unsigned j = 0; j < p+1; j++)
+                lcmDenoms = lcm(lcmDenoms, activeRegion(i, j).den);
+            for (unsigned j = 0; j < p+1; j++)
+                activeRegionNorm(i, j) = (activeRegion(i, j) * Fraction(lcmDenoms, 1)).getAsInteger();
+
+            activeRegionRel.addInequality(activeRegionNorm.getRow(i));
+        }
+
+        activeRegions.push_back(PresburgerRelation(activeRegionRel));
+    }
+
+    // Clauss-Loechner chamber decomposition
+    std::vector<std::pair<PresburgerRelation, std::vector<unsigned>>> chambers = 
+        {std::make_pair(activeRegions[0], std::vector({0u}))};
+    std::vector<std::pair<PresburgerRelation, std::vector<unsigned>>> newChambers;
+    for (unsigned j = 1; j < vertices.size(); j++)
+    {
+        newChambers.clear();
+        PresburgerRelation r_j = activeRegions[j];
+        ParamPoint v_j = vertices[j];
+        for (unsigned i = 0; i < chambers.size(); i++)
+        {
+            auto [r_i, v_i] = chambers[i];
+
+            PresburgerRelation intersection = r_i.intersect(r_j);
+            bool isFullDim = false;
+            for (auto disjunct : intersection.getAllDisjuncts())
+                if (disjunct.isFullDim())
+                {
+                    isFullDim = true;
+                    break;
+                }
+            isFullDim = (p == 0) || isFullDim;
+            if (!isFullDim) newChambers.push_back(chambers[i]);
+            else
+            {
+                PresburgerRelation subtraction = r_i.subtract(r_j);
+                newChambers.push_back(std::make_pair(subtraction, v_i));
+
+                v_i.push_back(j);
+                newChambers.push_back(std::make_pair(intersection, v_i));
+            }
+
+        }
+        for (auto chamber : newChambers)
+            r_j = r_j.subtract(chamber.first);
+
+        newChambers.push_back(std::make_pair(r_j, std::vector({j})));
+
+        chambers.clear();
+        for (auto chamber : newChambers)
+        {
+            bool empty = true;
+            for (auto disjunct : chamber.first.getAllDisjuncts())
+                if (!disjunct.isEmpty())
+                {
+                    empty = false;
+                    break;
+                }
+            if (!empty)
+                chambers.push_back(chamber);
+        }
+    }
+
+    SmallVector<MPInt> ineq(d+1);
+    for (auto chamber : chambers)
+    {
+        chamberGf = GeneratingFunction(p, {}, {}, {});
+        for (unsigned i : chamber.second)
+        {
+            tgtCone = defineHRep(d);
+            for (unsigned j = 0; j < d; j++)
+            {
+                for (unsigned k = 0; k < d; k++)
+                    ineq[k] = subsets[i](j, k);
+                ineq[d] = subsets[i](j, d+p);
+                tgtCone.addInequality(ineq);
+            }
+            unimodCones = {std::make_pair(1, tgtCone)};
+            for (auto signedCone : unimodCones)
+                chamberGf = chamberGf + unimodularConeGeneratingFunction(vertices[i], signedCone.first, signedCone.second);
+        }
+        gf.push_back(std::make_pair(chamber.first, chamberGf));
+    }
+    return gf;
+}
+
 /// We use an iterative procedure to find a vector not orthogonal
 /// to a given set, ignoring the null vectors.
 /// Let the inputs be {x_1, ..., x_k}, all vectors of length n.
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 7d2a63d17676f5..578b95571fff24 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -2498,6 +2498,42 @@ void IntegerRelation::printSpace(raw_ostream &os) const {
   os << getNumConstraints() << " constraints\n";
 }
 
+void IntegerRelation::removeTrivialEqualities() {
+    bool flag;
+    for (unsigned i = 0; i < getNumEqualities(); i++)
+    {
+        flag = true;
+        for (unsigned j = 0; j < getNumVars()+1; j++)
+            if (atEq(i, j) != 0)
+                flag = false;
+        if (flag)
+            removeEquality(i);
+    }
+}
+
+bool IntegerRelation::isFullDim() {
+    removeTrivialEqualities(); // to implement: remove equalities that are `0 = 0`
+
+    if (getNumEqualities() > 0)
+      return true;
+
+    Simplex simplex(*this);
+    for (unsigned i = 0; i < getNumInequalities(); i++)
+    {
+      auto ineq = inequalities.getRow(i);
+      auto upOpt = simplex.computeOptimum(Simplex::Direction::Up, ineq);
+      auto downOpt = simplex.computeOptimum(Simplex::Direction::Down, ineq);
+      if (upOpt.getKind() == OptimumKind::Unbounded ||
+          downOpt.getKind() == OptimumKind::Unbounded)
+            continue;
+      if (upOpt.getKind() == OptimumKind::Bounded &&
+          downOpt.getKind() == OptimumKind::Bounded &&
+          *upOpt == *downOpt) // ineq == 0 holds for this
+        return false;
+    }
+    return true;
+  }
+
 void IntegerRelation::print(raw_ostream &os) const {
   assert(hasConsistentState());
   printSpace(os);

>From 3a147d0ee7542720acea7349d15968a1414cf017 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 22 Jan 2024 14:57:29 +0530
Subject: [PATCH 2/6] Formatting

---
 .../mlir/Analysis/Presburger/Barvinok.h       |   3 +-
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 390 +++++++++---------
 .../Analysis/Presburger/IntegerRelation.cpp   |  54 ++-
 3 files changed, 219 insertions(+), 228 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 6feb7e8f09233f..1fafa1e709d5e3 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -87,7 +87,8 @@ GeneratingFunction unimodularConeGeneratingFunction(ParamPoint vertex, int sign,
 
 std::optional<ParamPoint> findVertex(Matrix<MPInt> equations);
 
-std::vector<std::pair<PresburgerRelation, GeneratingFunction>> polytopeGeneratingFunction(PolyhedronH poly);
+std::vector<std::pair<PresburgerRelation, GeneratingFunction>>
+polytopeGeneratingFunction(PolyhedronH poly);
 
 /// Find a vector that is not orthogonal to any of the given vectors,
 /// i.e., has nonzero dot product with those of the given vectors
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index c8cec45a21dc4b..f9f75a8a0dcef3 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -148,220 +148,212 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
                             std::vector({denominator}));
 }
 
-std::optional<ParamPoint> mlir::presburger::detail::findVertex(Matrix<MPInt> equations)
-{
-    // `equalities` is a d x (d + p + 1) matrix.
-
-    unsigned r = equations.getNumRows();
-    unsigned c = equations.getNumColumns();
-
-    IntMatrix coeffs(r, r);
-    for (unsigned i = 0; i < r; i++)
-        for (unsigned j = 0; j < r; j++)
-            coeffs(i, j) = equations(i, j);
-    
-    if (coeffs.determinant() == MPInt(0))
-        return std::nullopt;
-
-    Matrix<Fraction> equationsF(r, c);
-    for (unsigned i = 0; i < r; i++)
-        for (unsigned j = 0; j < c; j++)
-            equationsF(i, j) = Fraction(equations(i, j), 1);
-    
-    Fraction a, b;
-    for (unsigned i = 0; i < r; i++)
-    {
-        if (equationsF(i, i) == Fraction(0, 1))
-            for (unsigned j = i+1; j < r; j++)
-                if (equationsF(j, i) != 0)
-                {
-                    equationsF.addToRow(i, equationsF.getRow(j), Fraction(1, 1));
-                    break;
-                }
-        b = equationsF(i, i);
-
-        for (unsigned j = 0; j < r; j++)
-        {
-            if (equationsF(j, i) == 0 || j == i) continue;
-            a = equationsF(j, i);
-            equationsF.addToRow(j, equationsF.getRow(i), - a / b);
+std::optional<ParamPoint>
+mlir::presburger::detail::findVertex(Matrix<MPInt> equations) {
+  // `equalities` is a d x (d + p + 1) matrix.
+
+  unsigned r = equations.getNumRows();
+  unsigned c = equations.getNumColumns();
+
+  IntMatrix coeffs(r, r);
+  for (unsigned i = 0; i < r; i++)
+    for (unsigned j = 0; j < r; j++)
+      coeffs(i, j) = equations(i, j);
+
+  if (coeffs.determinant() == MPInt(0))
+    return std::nullopt;
+
+  Matrix<Fraction> equationsF(r, c);
+  for (unsigned i = 0; i < r; i++)
+    for (unsigned j = 0; j < c; j++)
+      equationsF(i, j) = Fraction(equations(i, j), 1);
+
+  Fraction a, b;
+  for (unsigned i = 0; i < r; i++) {
+    if (equationsF(i, i) == Fraction(0, 1))
+      for (unsigned j = i + 1; j < r; j++)
+        if (equationsF(j, i) != 0) {
+          equationsF.addToRow(i, equationsF.getRow(j), Fraction(1, 1));
+          break;
         }
-    }
+    b = equationsF(i, i);
 
-    for (unsigned i = 0; i < r; i++)
-    {
-        a = equationsF(i, i);
-        for (unsigned j = 0; j < c; j++)
-            equationsF(i, j) = equationsF(i, j) / a;
+    for (unsigned j = 0; j < r; j++) {
+      if (equationsF(j, i) == 0 || j == i)
+        continue;
+      a = equationsF(j, i);
+      equationsF.addToRow(j, equationsF.getRow(i), -a / b);
     }
+  }
 
-    ParamPoint vertex(r, c-r); // d x p+1
-    for (unsigned i = 0; i < r; i++)
-        for (unsigned j = 0; j < c-r; j++)
-            vertex(i, j) = -equationsF(i, r+j);
-    
-    return vertex;
-}
+  for (unsigned i = 0; i < r; i++) {
+    a = equationsF(i, i);
+    for (unsigned j = 0; j < c; j++)
+      equationsF(i, j) = equationsF(i, j) / a;
+  }
 
-std::vector<std::pair<PresburgerRelation, GeneratingFunction>> mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly)
-{
-    unsigned d = poly.getNumRangeVars(); 
-    unsigned p = poly.getNumSymbolVars();
-    unsigned n = poly.getNumInequalities();
-
-    SmallVector<std::pair<int, ConeH>, 4> unimodCones;
-    GeneratingFunction chamberGf(p, {}, {}, {});
-    std::vector<std::pair<PresburgerRelation, GeneratingFunction>> gf({});
-    ConeH tgtCone = defineHRep(d);
-
-    Matrix<MPInt> subset(d, d+p+1);
-    std::vector<Matrix<MPInt>> subsets; // Stores the inequality subsets corresponding to each vertex.
-    Matrix<Fraction> remaining(n-d, d+p+1);
-
-    std::optional<ParamPoint> vertex;
-    std::vector<ParamPoint> vertices;
-
-    Matrix<Fraction> a2(n-d, d);
-    Matrix<Fraction> b2c2(n-d, p+1);
-
-    Matrix<Fraction> activeRegion(n-d, p+1);
-    Matrix<MPInt> activeRegionNorm(n-d, p+1);
-    MPInt lcmDenoms;
-    IntegerRelation activeRegionRel(PresburgerSpace::getRelationSpace(0, p, 0, 0));
-    // The active region will be defined as activeRegionCoeffs @ p + activeRegionConstant ≥ 0.
-    // The active region is a polyhedron in parameter space.
-    std::vector<PresburgerRelation> activeRegions;
-    
-
-    for (std::bitset<16> indicator(((1ul << d)-1ul) << (n-d));
-        indicator.to_ulong() <= ((1ul << d)-1ul) << (n-d);              // d 1's followed by n-d 0's
-        indicator = std::bitset<16>(indicator.to_ulong() - 1))
-    {
-        if (indicator.count() != d)
-            continue;
-
-        subset = Matrix<MPInt>(d, d+p+1);
-        remaining = Matrix<Fraction>(n-d, d+p+1);
-        unsigned j1 = 0, j2 = 0;
-        for (unsigned i = 0; i < n; i++)
-            if (indicator.test(i))
-                subset.setRow(j1++, poly.getInequality(i));
-                // [A1 | B1 | c1]
-            else
-            {
-                for (unsigned k = 0; k < d; k++)
-                    a2(j2, k) = Fraction(poly.atIneq(i, k), 1);
-                for (unsigned k = d; k < d+p+1; k++)
-                    b2c2(j2, k-d) = Fraction(poly.atIneq(i, k), 1);
-                j2++;
-                // [A2 | B2 | c2]
-            }
-
-        vertex = findVertex(subset); // d x (p+1)
-
-        if (vertex == std::nullopt) continue;
-        vertices.push_back(*vertex);
-        subsets.push_back(subset);
-
-        // Region is given by (A2 @ X + B2) p + (A2 @ y + c2) ≥ 0
-        // This is equivt to A2 @ [X | y] + [B2 | c2]
-        // We premultiply [X | y] with each row of A2 and add each row of [B2 | c2].
-        for (unsigned i = 0; i < n-d; i++)
-        {
-            activeRegion.setRow(i, (*vertex).preMultiplyWithRow(a2.getRow(i)));
-            activeRegion.addToRow(i, b2c2.getRow(i), Fraction(1, 1));
-        }
+  ParamPoint vertex(r, c - r); // d x p+1
+  for (unsigned i = 0; i < r; i++)
+    for (unsigned j = 0; j < c - r; j++)
+      vertex(i, j) = -equationsF(i, r + j);
 
-        activeRegionNorm = Matrix<MPInt>(n-d, p+1);
-        activeRegionRel = IntegerRelation(PresburgerSpace::getRelationSpace(0, p, 0, 0));
-        lcmDenoms = 1;
-        for (unsigned i = 0; i < n-d; i++)
-        {
-            for (unsigned j = 0; j < p+1; j++)
-                lcmDenoms = lcm(lcmDenoms, activeRegion(i, j).den);
-            for (unsigned j = 0; j < p+1; j++)
-                activeRegionNorm(i, j) = (activeRegion(i, j) * Fraction(lcmDenoms, 1)).getAsInteger();
-
-            activeRegionRel.addInequality(activeRegionNorm.getRow(i));
-        }
+  return vertex;
+}
 
-        activeRegions.push_back(PresburgerRelation(activeRegionRel));
+std::vector<std::pair<PresburgerRelation, GeneratingFunction>>
+mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly) {
+  unsigned d = poly.getNumRangeVars();
+  unsigned p = poly.getNumSymbolVars();
+  unsigned n = poly.getNumInequalities();
+
+  SmallVector<std::pair<int, ConeH>, 4> unimodCones;
+  GeneratingFunction chamberGf(p, {}, {}, {});
+  std::vector<std::pair<PresburgerRelation, GeneratingFunction>> gf({});
+  ConeH tgtCone = defineHRep(d);
+
+  Matrix<MPInt> subset(d, d + p + 1);
+  std::vector<Matrix<MPInt>>
+      subsets; // Stores the inequality subsets corresponding to each vertex.
+  Matrix<Fraction> remaining(n - d, d + p + 1);
+
+  std::optional<ParamPoint> vertex;
+  std::vector<ParamPoint> vertices;
+
+  Matrix<Fraction> a2(n - d, d);
+  Matrix<Fraction> b2c2(n - d, p + 1);
+
+  Matrix<Fraction> activeRegion(n - d, p + 1);
+  Matrix<MPInt> activeRegionNorm(n - d, p + 1);
+  MPInt lcmDenoms;
+  IntegerRelation activeRegionRel(
+      PresburgerSpace::getRelationSpace(0, p, 0, 0));
+  // The active region will be defined as activeRegionCoeffs @ p +
+  // activeRegionConstant ≥ 0. The active region is a polyhedron in parameter
+  // space.
+  std::vector<PresburgerRelation> activeRegions;
+
+  for (std::bitset<16> indicator(((1ul << d) - 1ul) << (n - d));
+       indicator.to_ulong() <= ((1ul << d) - 1ul)
+                                   << (n - d); // d 1's followed by n-d 0's
+       indicator = std::bitset<16>(indicator.to_ulong() - 1)) {
+    if (indicator.count() != d)
+      continue;
+
+    subset = Matrix<MPInt>(d, d + p + 1);
+    remaining = Matrix<Fraction>(n - d, d + p + 1);
+    unsigned j1 = 0, j2 = 0;
+    for (unsigned i = 0; i < n; i++)
+      if (indicator.test(i))
+        subset.setRow(j1++, poly.getInequality(i));
+      // [A1 | B1 | c1]
+      else {
+        for (unsigned k = 0; k < d; k++)
+          a2(j2, k) = Fraction(poly.atIneq(i, k), 1);
+        for (unsigned k = d; k < d + p + 1; k++)
+          b2c2(j2, k - d) = Fraction(poly.atIneq(i, k), 1);
+        j2++;
+        // [A2 | B2 | c2]
+      }
+
+    vertex = findVertex(subset); // d x (p+1)
+
+    if (vertex == std::nullopt)
+      continue;
+    vertices.push_back(*vertex);
+    subsets.push_back(subset);
+
+    // Region is given by (A2 @ X + B2) p + (A2 @ y + c2) ≥ 0
+    // This is equivt to A2 @ [X | y] + [B2 | c2]
+    // We premultiply [X | y] with each row of A2 and add each row of [B2 | c2].
+    for (unsigned i = 0; i < n - d; i++) {
+      activeRegion.setRow(i, (*vertex).preMultiplyWithRow(a2.getRow(i)));
+      activeRegion.addToRow(i, b2c2.getRow(i), Fraction(1, 1));
     }
 
-    // Clauss-Loechner chamber decomposition
-    std::vector<std::pair<PresburgerRelation, std::vector<unsigned>>> chambers = 
-        {std::make_pair(activeRegions[0], std::vector({0u}))};
-    std::vector<std::pair<PresburgerRelation, std::vector<unsigned>>> newChambers;
-    for (unsigned j = 1; j < vertices.size(); j++)
-    {
-        newChambers.clear();
-        PresburgerRelation r_j = activeRegions[j];
-        ParamPoint v_j = vertices[j];
-        for (unsigned i = 0; i < chambers.size(); i++)
-        {
-            auto [r_i, v_i] = chambers[i];
-
-            PresburgerRelation intersection = r_i.intersect(r_j);
-            bool isFullDim = false;
-            for (auto disjunct : intersection.getAllDisjuncts())
-                if (disjunct.isFullDim())
-                {
-                    isFullDim = true;
-                    break;
-                }
-            isFullDim = (p == 0) || isFullDim;
-            if (!isFullDim) newChambers.push_back(chambers[i]);
-            else
-            {
-                PresburgerRelation subtraction = r_i.subtract(r_j);
-                newChambers.push_back(std::make_pair(subtraction, v_i));
-
-                v_i.push_back(j);
-                newChambers.push_back(std::make_pair(intersection, v_i));
-            }
+    activeRegionNorm = Matrix<MPInt>(n - d, p + 1);
+    activeRegionRel =
+        IntegerRelation(PresburgerSpace::getRelationSpace(0, p, 0, 0));
+    lcmDenoms = 1;
+    for (unsigned i = 0; i < n - d; i++) {
+      for (unsigned j = 0; j < p + 1; j++)
+        lcmDenoms = lcm(lcmDenoms, activeRegion(i, j).den);
+      for (unsigned j = 0; j < p + 1; j++)
+        activeRegionNorm(i, j) =
+            (activeRegion(i, j) * Fraction(lcmDenoms, 1)).getAsInteger();
+
+      activeRegionRel.addInequality(activeRegionNorm.getRow(i));
+    }
+
+    activeRegions.push_back(PresburgerRelation(activeRegionRel));
+  }
 
+  // Clauss-Loechner chamber decomposition
+  std::vector<std::pair<PresburgerRelation, std::vector<unsigned>>> chambers = {
+      std::make_pair(activeRegions[0], std::vector({0u}))};
+  std::vector<std::pair<PresburgerRelation, std::vector<unsigned>>> newChambers;
+  for (unsigned j = 1; j < vertices.size(); j++) {
+    newChambers.clear();
+    PresburgerRelation r_j = activeRegions[j];
+    ParamPoint v_j = vertices[j];
+    for (unsigned i = 0; i < chambers.size(); i++) {
+      auto [r_i, v_i] = chambers[i];
+
+      PresburgerRelation intersection = r_i.intersect(r_j);
+      bool isFullDim = false;
+      for (auto disjunct : intersection.getAllDisjuncts())
+        if (disjunct.isFullDim()) {
+          isFullDim = true;
+          break;
         }
-        for (auto chamber : newChambers)
-            r_j = r_j.subtract(chamber.first);
-
-        newChambers.push_back(std::make_pair(r_j, std::vector({j})));
-
-        chambers.clear();
-        for (auto chamber : newChambers)
-        {
-            bool empty = true;
-            for (auto disjunct : chamber.first.getAllDisjuncts())
-                if (!disjunct.isEmpty())
-                {
-                    empty = false;
-                    break;
-                }
-            if (!empty)
-                chambers.push_back(chamber);
+      isFullDim = (p == 0) || isFullDim;
+      if (!isFullDim)
+        newChambers.push_back(chambers[i]);
+      else {
+        PresburgerRelation subtraction = r_i.subtract(r_j);
+        newChambers.push_back(std::make_pair(subtraction, v_i));
+
+        v_i.push_back(j);
+        newChambers.push_back(std::make_pair(intersection, v_i));
+      }
+    }
+    for (auto chamber : newChambers)
+      r_j = r_j.subtract(chamber.first);
+
+    newChambers.push_back(std::make_pair(r_j, std::vector({j})));
+
+    chambers.clear();
+    for (auto chamber : newChambers) {
+      bool empty = true;
+      for (auto disjunct : chamber.first.getAllDisjuncts())
+        if (!disjunct.isEmpty()) {
+          empty = false;
+          break;
         }
+      if (!empty)
+        chambers.push_back(chamber);
     }
+  }
 
-    SmallVector<MPInt> ineq(d+1);
-    for (auto chamber : chambers)
-    {
-        chamberGf = GeneratingFunction(p, {}, {}, {});
-        for (unsigned i : chamber.second)
-        {
-            tgtCone = defineHRep(d);
-            for (unsigned j = 0; j < d; j++)
-            {
-                for (unsigned k = 0; k < d; k++)
-                    ineq[k] = subsets[i](j, k);
-                ineq[d] = subsets[i](j, d+p);
-                tgtCone.addInequality(ineq);
-            }
-            unimodCones = {std::make_pair(1, tgtCone)};
-            for (auto signedCone : unimodCones)
-                chamberGf = chamberGf + unimodularConeGeneratingFunction(vertices[i], signedCone.first, signedCone.second);
-        }
-        gf.push_back(std::make_pair(chamber.first, chamberGf));
+  SmallVector<MPInt> ineq(d + 1);
+  for (auto chamber : chambers) {
+    chamberGf = GeneratingFunction(p, {}, {}, {});
+    for (unsigned i : chamber.second) {
+      tgtCone = defineHRep(d);
+      for (unsigned j = 0; j < d; j++) {
+        for (unsigned k = 0; k < d; k++)
+          ineq[k] = subsets[i](j, k);
+        ineq[d] = subsets[i](j, d + p);
+        tgtCone.addInequality(ineq);
+      }
+      unimodCones = {std::make_pair(1, tgtCone)};
+      for (auto signedCone : unimodCones)
+        chamberGf =
+            chamberGf + unimodularConeGeneratingFunction(
+                            vertices[i], signedCone.first, signedCone.second);
     }
-    return gf;
+    gf.push_back(std::make_pair(chamber.first, chamberGf));
+  }
+  return gf;
 }
 
 /// We use an iterative procedure to find a vector not orthogonal
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 578b95571fff24..edb95b4072e1d8 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -2499,40 +2499,38 @@ void IntegerRelation::printSpace(raw_ostream &os) const {
 }
 
 void IntegerRelation::removeTrivialEqualities() {
-    bool flag;
-    for (unsigned i = 0; i < getNumEqualities(); i++)
-    {
-        flag = true;
-        for (unsigned j = 0; j < getNumVars()+1; j++)
-            if (atEq(i, j) != 0)
-                flag = false;
-        if (flag)
-            removeEquality(i);
-    }
+  bool flag;
+  for (unsigned i = 0; i < getNumEqualities(); i++) {
+    flag = true;
+    for (unsigned j = 0; j < getNumVars() + 1; j++)
+      if (atEq(i, j) != 0)
+        flag = false;
+    if (flag)
+      removeEquality(i);
+  }
 }
 
 bool IntegerRelation::isFullDim() {
-    removeTrivialEqualities(); // to implement: remove equalities that are `0 = 0`
-
-    if (getNumEqualities() > 0)
-      return true;
+  removeTrivialEqualities(); // to implement: remove equalities that are `0 = 0`
 
-    Simplex simplex(*this);
-    for (unsigned i = 0; i < getNumInequalities(); i++)
-    {
-      auto ineq = inequalities.getRow(i);
-      auto upOpt = simplex.computeOptimum(Simplex::Direction::Up, ineq);
-      auto downOpt = simplex.computeOptimum(Simplex::Direction::Down, ineq);
-      if (upOpt.getKind() == OptimumKind::Unbounded ||
-          downOpt.getKind() == OptimumKind::Unbounded)
-            continue;
-      if (upOpt.getKind() == OptimumKind::Bounded &&
-          downOpt.getKind() == OptimumKind::Bounded &&
-          *upOpt == *downOpt) // ineq == 0 holds for this
-        return false;
-    }
+  if (getNumEqualities() > 0)
     return true;
+
+  Simplex simplex(*this);
+  for (unsigned i = 0; i < getNumInequalities(); i++) {
+    auto ineq = inequalities.getRow(i);
+    auto upOpt = simplex.computeOptimum(Simplex::Direction::Up, ineq);
+    auto downOpt = simplex.computeOptimum(Simplex::Direction::Down, ineq);
+    if (upOpt.getKind() == OptimumKind::Unbounded ||
+        downOpt.getKind() == OptimumKind::Unbounded)
+      continue;
+    if (upOpt.getKind() == OptimumKind::Bounded &&
+        downOpt.getKind() == OptimumKind::Bounded &&
+        *upOpt == *downOpt) // ineq == 0 holds for this
+      return false;
   }
+  return true;
+}
 
 void IntegerRelation::print(raw_ostream &os) const {
   assert(hasConsistentState());

>From 294dc0ae9147463248b9b40f04d4a2529ff09ed4 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 22 Jan 2024 15:07:38 +0530
Subject: [PATCH 3/6] Use intmatrix and fracmatrix

---
 .../mlir/Analysis/Presburger/Barvinok.h       |  2 +-
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 24 +++++++++----------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 1fafa1e709d5e3..320848348552c9 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -85,7 +85,7 @@ ConeH getDual(ConeV cone);
 GeneratingFunction unimodularConeGeneratingFunction(ParamPoint vertex, int sign,
                                                     ConeH cone);
 
-std::optional<ParamPoint> findVertex(Matrix<MPInt> equations);
+std::optional<ParamPoint> findVertex(IntMatrix equations);
 
 std::vector<std::pair<PresburgerRelation, GeneratingFunction>>
 polytopeGeneratingFunction(PolyhedronH poly);
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index f9f75a8a0dcef3..3753f715c7d925 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -149,7 +149,7 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
 }
 
 std::optional<ParamPoint>
-mlir::presburger::detail::findVertex(Matrix<MPInt> equations) {
+mlir::presburger::detail::findVertex(IntMatrix equations) {
   // `equalities` is a d x (d + p + 1) matrix.
 
   unsigned r = equations.getNumRows();
@@ -163,7 +163,7 @@ mlir::presburger::detail::findVertex(Matrix<MPInt> equations) {
   if (coeffs.determinant() == MPInt(0))
     return std::nullopt;
 
-  Matrix<Fraction> equationsF(r, c);
+  FracMatrix equationsF(r, c);
   for (unsigned i = 0; i < r; i++)
     for (unsigned j = 0; j < c; j++)
       equationsF(i, j) = Fraction(equations(i, j), 1);
@@ -211,19 +211,19 @@ mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly) {
   std::vector<std::pair<PresburgerRelation, GeneratingFunction>> gf({});
   ConeH tgtCone = defineHRep(d);
 
-  Matrix<MPInt> subset(d, d + p + 1);
-  std::vector<Matrix<MPInt>>
+  IntMatrix subset(d, d + p + 1);
+  std::vector<IntMatrix>
       subsets; // Stores the inequality subsets corresponding to each vertex.
-  Matrix<Fraction> remaining(n - d, d + p + 1);
+  FracMatrix remaining(n - d, d + p + 1);
 
   std::optional<ParamPoint> vertex;
   std::vector<ParamPoint> vertices;
 
-  Matrix<Fraction> a2(n - d, d);
-  Matrix<Fraction> b2c2(n - d, p + 1);
+  FracMatrix a2(n - d, d);
+  FracMatrix b2c2(n - d, p + 1);
 
-  Matrix<Fraction> activeRegion(n - d, p + 1);
-  Matrix<MPInt> activeRegionNorm(n - d, p + 1);
+  FracMatrix activeRegion(n - d, p + 1);
+  IntMatrix activeRegionNorm(n - d, p + 1);
   MPInt lcmDenoms;
   IntegerRelation activeRegionRel(
       PresburgerSpace::getRelationSpace(0, p, 0, 0));
@@ -239,8 +239,8 @@ mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly) {
     if (indicator.count() != d)
       continue;
 
-    subset = Matrix<MPInt>(d, d + p + 1);
-    remaining = Matrix<Fraction>(n - d, d + p + 1);
+    subset = IntMatrix(d, d + p + 1);
+    remaining = FracMatrix(n - d, d + p + 1);
     unsigned j1 = 0, j2 = 0;
     for (unsigned i = 0; i < n; i++)
       if (indicator.test(i))
@@ -270,7 +270,7 @@ mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly) {
       activeRegion.addToRow(i, b2c2.getRow(i), Fraction(1, 1));
     }
 
-    activeRegionNorm = Matrix<MPInt>(n - d, p + 1);
+    activeRegionNorm = IntMatrix(n - d, p + 1);
     activeRegionRel =
         IntegerRelation(PresburgerSpace::getRelationSpace(0, p, 0, 0));
     lcmDenoms = 1;

>From 35e9b7665fcb8b5445d39aa98274e3e63904e368 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 22 Jan 2024 19:11:47 +0530
Subject: [PATCH 4/6] Doc

---
 .../mlir/Analysis/Presburger/Barvinok.h       |  4 ++
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 71 ++++++++++++-------
 2 files changed, 50 insertions(+), 25 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 320848348552c9..9c940231d5e147 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -85,6 +85,10 @@ ConeH getDual(ConeV cone);
 GeneratingFunction unimodularConeGeneratingFunction(ParamPoint vertex, int sign,
                                                     ConeH cone);
 
+/// Find the solution of a set of equations that express affine constraints
+/// between a set of variables and a set of parameters. The solution expresses
+/// each variable as an affine function of the parameters.
+/// If there is no solution, return std::null.
 std::optional<ParamPoint> findVertex(IntMatrix equations);
 
 std::vector<std::pair<PresburgerRelation, GeneratingFunction>>
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index 3753f715c7d925..ff202984bd3362 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -148,54 +148,75 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
                             std::vector({denominator}));
 }
 
+/// We use Gaussian elimination to find the solution to a set of d equations
+/// of the form
+/// a_1 x_1 + ... + a_d x_d + b_1 m_1 + ... + b_p m_p + c = 0
+/// where x_i are variables,
+/// m_i are parameters and
+/// a_i, b_i, c are rational coefficients.
+/// The solution expresses each x_i as an affine function of the m_i, and is
+/// therefore represented as a matrix of size d x (p+1).
+/// If there is no solution, we return null.
 std::optional<ParamPoint>
 mlir::presburger::detail::findVertex(IntMatrix equations) {
-  // `equalities` is a d x (d + p + 1) matrix.
+  // equations is a d x (d + p + 1) matrix.
+  // Each row represents an equation.
 
-  unsigned r = equations.getNumRows();
+  unsigned d = equations.getNumRows();
   unsigned c = equations.getNumColumns();
 
-  IntMatrix coeffs(r, r);
-  for (unsigned i = 0; i < r; i++)
-    for (unsigned j = 0; j < r; j++)
+  // First, we check that the system has a solution, and return
+  // null if not.
+  IntMatrix coeffs(d, d);
+  for (unsigned i = 0; i < d; i++)
+    for (unsigned j = 0; j < d; j++)
       coeffs(i, j) = equations(i, j);
 
-  if (coeffs.determinant() == MPInt(0))
+  if (coeffs.determinant() == 0)
     return std::nullopt;
 
-  FracMatrix equationsF(r, c);
-  for (unsigned i = 0; i < r; i++)
-    for (unsigned j = 0; j < c; j++)
-      equationsF(i, j) = Fraction(equations(i, j), 1);
+  // We work with rational numbers.
+  FracMatrix equationsF(equations);
 
-  Fraction a, b;
-  for (unsigned i = 0; i < r; i++) {
-    if (equationsF(i, i) == Fraction(0, 1))
-      for (unsigned j = i + 1; j < r; j++)
+  for (unsigned i = 0; i < d; ++i) {
+    // First ensure that the diagonal element is nonzero, by swapping
+    // it with a nonzero row.
+    if (equationsF(i, i) == 0) {
+      for (unsigned j = i + 1; j < d; ++j) {
         if (equationsF(j, i) != 0) {
-          equationsF.addToRow(i, equationsF.getRow(j), Fraction(1, 1));
+          equationsF.swapRows(j, i);
           break;
         }
-    b = equationsF(i, i);
+      }
+    }
+
+    Fraction b = equationsF(i, i);
 
-    for (unsigned j = 0; j < r; j++) {
+    // Set all elements except the diagonal to zero.
+    for (unsigned j = 0; j < d; ++j) {
       if (equationsF(j, i) == 0 || j == i)
         continue;
-      a = equationsF(j, i);
+      // Set element (j, i) to zero
+      // by subtracting the ith row,
+      // appropriately scaled.
+      Fraction a = equationsF(j, i);
       equationsF.addToRow(j, equationsF.getRow(i), -a / b);
     }
   }
 
-  for (unsigned i = 0; i < r; i++) {
-    a = equationsF(i, i);
-    for (unsigned j = 0; j < c; j++)
+  // Rescale diagonal elements to 1.
+  for (unsigned i = 0; i < d; ++i) {
+    Fraction a = equationsF(i, i);
+    for (unsigned j = 0; j < c; ++j)
       equationsF(i, j) = equationsF(i, j) / a;
   }
 
-  ParamPoint vertex(r, c - r); // d x p+1
-  for (unsigned i = 0; i < r; i++)
-    for (unsigned j = 0; j < c - r; j++)
-      vertex(i, j) = -equationsF(i, r + j);
+  // We copy the last p+1 columns of the matrix as the values of x_i.
+  // We shift the parameter terms to the RHS, and so flip their sign.
+  ParamPoint vertex(d, c - d);
+  for (unsigned i = 0; i < d; ++i)
+    for (unsigned j = 0; j < c - d; ++j)
+      vertex(i, j) = -equationsF(i, d + j);
 
   return vertex;
 }

>From 15e20fe3d2fe845efcb6b3f32165117980ce46dd Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 22 Jan 2024 20:19:08 +0530
Subject: [PATCH 5/6] Doc

---
 .../mlir/Analysis/Presburger/Barvinok.h       |   2 +
 mlir/lib/Analysis/Presburger/Barvinok.cpp     | 228 ++++++++++++------
 2 files changed, 157 insertions(+), 73 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index 9c940231d5e147..a5bc49d79408c6 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -91,6 +91,8 @@ GeneratingFunction unimodularConeGeneratingFunction(ParamPoint vertex, int sign,
 /// If there is no solution, return std::null.
 std::optional<ParamPoint> findVertex(IntMatrix equations);
 
+/// Compute the generating function corresponding to a polytope.
+/// All tangent cones of the polytope must be unimodular.
 std::vector<std::pair<PresburgerRelation, GeneratingFunction>>
 polytopeGeneratingFunction(PolyhedronH poly);
 
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index ff202984bd3362..3afc9ad8638ccd 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -162,14 +162,14 @@ mlir::presburger::detail::findVertex(IntMatrix equations) {
   // equations is a d x (d + p + 1) matrix.
   // Each row represents an equation.
 
-  unsigned d = equations.getNumRows();
-  unsigned c = equations.getNumColumns();
+  unsigned numEqs = equations.getNumRows();
+  unsigned numCols = equations.getNumColumns();
 
   // First, we check that the system has a solution, and return
   // null if not.
-  IntMatrix coeffs(d, d);
-  for (unsigned i = 0; i < d; i++)
-    for (unsigned j = 0; j < d; j++)
+  IntMatrix coeffs(numEqs, numEqs);
+  for (unsigned i = 0; i < numEqs; i++)
+    for (unsigned j = 0; j < numEqs; j++)
       coeffs(i, j) = equations(i, j);
 
   if (coeffs.determinant() == 0)
@@ -178,11 +178,11 @@ mlir::presburger::detail::findVertex(IntMatrix equations) {
   // We work with rational numbers.
   FracMatrix equationsF(equations);
 
-  for (unsigned i = 0; i < d; ++i) {
+  for (unsigned i = 0; i < numEqs; ++i) {
     // First ensure that the diagonal element is nonzero, by swapping
     // it with a nonzero row.
     if (equationsF(i, i) == 0) {
-      for (unsigned j = i + 1; j < d; ++j) {
+      for (unsigned j = i + 1; j < numEqs; ++j) {
         if (equationsF(j, i) != 0) {
           equationsF.swapRows(j, i);
           break;
@@ -193,7 +193,7 @@ mlir::presburger::detail::findVertex(IntMatrix equations) {
     Fraction b = equationsF(i, i);
 
     // Set all elements except the diagonal to zero.
-    for (unsigned j = 0; j < d; ++j) {
+    for (unsigned j = 0; j < numEqs; ++j) {
       if (equationsF(j, i) == 0 || j == i)
         continue;
       // Set element (j, i) to zero
@@ -205,102 +205,140 @@ mlir::presburger::detail::findVertex(IntMatrix equations) {
   }
 
   // Rescale diagonal elements to 1.
-  for (unsigned i = 0; i < d; ++i) {
+  for (unsigned i = 0; i < numEqs; ++i) {
     Fraction a = equationsF(i, i);
-    for (unsigned j = 0; j < c; ++j)
+    for (unsigned j = 0; j < numCols; ++j)
       equationsF(i, j) = equationsF(i, j) / a;
   }
 
   // We copy the last p+1 columns of the matrix as the values of x_i.
   // We shift the parameter terms to the RHS, and so flip their sign.
-  ParamPoint vertex(d, c - d);
-  for (unsigned i = 0; i < d; ++i)
-    for (unsigned j = 0; j < c - d; ++j)
-      vertex(i, j) = -equationsF(i, d + j);
+  ParamPoint vertex(numEqs, numCols - numEqs);
+  for (unsigned i = 0; i < numEqs; ++i)
+    for (unsigned j = 0; j < numCols - numEqs; ++j)
+      vertex(i, j) = -equationsF(i, numEqs + j);
 
   return vertex;
 }
 
+/// For a polytope expressed as a set of inequalities, compute the generating
+/// function corresponding to the number of lattice points present. This
+/// algorithm has three main steps:
+/// 1. Enumerate the vertices, by iterating over subsets of inequalities and
+///    checking for solubility.
+/// 2. For each vertex, identify the tangent cone and compute the generating
+///    function corresponding to it. The sum of these GFs is the GF of the
+///    polytope.
+/// 3. [Clauss-Loechner decomposition] Identify the regions in parameter space
+///    (chambers) where each vertex is active, and accordingly compute the
+///    GF of the polytope in each chamber.
+///
+/// Verdoolaege, Sven, et al. "Counting integer points in parametric
+/// polytopes using Barvinok's rational functions." Algorithmica 48 (2007):
+/// 37-66.
 std::vector<std::pair<PresburgerRelation, GeneratingFunction>>
 mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly) {
-  unsigned d = poly.getNumRangeVars();
-  unsigned p = poly.getNumSymbolVars();
-  unsigned n = poly.getNumInequalities();
+  unsigned numVars = poly.getNumRangeVars();
+  unsigned numParams = poly.getNumSymbolVars();
+  unsigned numIneqs = poly.getNumInequalities();
 
-  SmallVector<std::pair<int, ConeH>, 4> unimodCones;
-  GeneratingFunction chamberGf(p, {}, {}, {});
+  // The generating function of the polytope is computed as a set of generating
+  // functions, each one associated with a region in parameter space (chamber).
   std::vector<std::pair<PresburgerRelation, GeneratingFunction>> gf({});
-  ConeH tgtCone = defineHRep(d);
-
-  IntMatrix subset(d, d + p + 1);
-  std::vector<IntMatrix>
-      subsets; // Stores the inequality subsets corresponding to each vertex.
-  FracMatrix remaining(n - d, d + p + 1);
 
-  std::optional<ParamPoint> vertex;
-  std::vector<ParamPoint> vertices;
-
-  FracMatrix a2(n - d, d);
-  FracMatrix b2c2(n - d, p + 1);
-
-  FracMatrix activeRegion(n - d, p + 1);
-  IntMatrix activeRegionNorm(n - d, p + 1);
-  MPInt lcmDenoms;
-  IntegerRelation activeRegionRel(
-      PresburgerSpace::getRelationSpace(0, p, 0, 0));
   // The active region will be defined as activeRegionCoeffs @ p +
   // activeRegionConstant ≥ 0. The active region is a polyhedron in parameter
   // space.
+  FracMatrix activeRegion(numIneqs - numVars, numParams + 1);
+
+  // These vectors store lists of
+  // subsets of inequalities,
+  // the vertices corresponding to them, and
+  // the active regions of the vertices, in order.
+  std::vector<IntMatrix> subsets;
+  std::vector<ParamPoint> vertices;
   std::vector<PresburgerRelation> activeRegions;
 
-  for (std::bitset<16> indicator(((1ul << d) - 1ul) << (n - d));
-       indicator.to_ulong() <= ((1ul << d) - 1ul)
-                                   << (n - d); // d 1's followed by n-d 0's
+  FracMatrix a2(numIneqs - numVars, numVars);
+  FracMatrix b2c2(numIneqs - numVars, numParams + 1);
+
+  // We iterate over all subsets of inequalities with cardinality numVars,
+  // using bitsets up to 2^numIneqs to enumerate.
+  for (std::bitset<16> indicator(((1ul << numVars) - 1ul)
+                                 << (numIneqs - numVars));
+       indicator.to_ulong() <=
+       ((1ul << numVars) - 1ul)
+           << (numIneqs - numVars); // d 1's followed by n-numVars 0's
        indicator = std::bitset<16>(indicator.to_ulong() - 1)) {
-    if (indicator.count() != d)
+
+    if (indicator.count() != numVars)
       continue;
 
-    subset = IntMatrix(d, d + p + 1);
-    remaining = FracMatrix(n - d, d + p + 1);
+    // Collect the inequalities corresponding to the bits which are set.
+    IntMatrix subset(numVars, numVars + numParams + 1);
     unsigned j1 = 0, j2 = 0;
-    for (unsigned i = 0; i < n; i++)
+    for (unsigned i = 0; i < numIneqs; i++)
       if (indicator.test(i))
         subset.setRow(j1++, poly.getInequality(i));
-      // [A1 | B1 | c1]
+
       else {
-        for (unsigned k = 0; k < d; k++)
-          a2(j2, k) = Fraction(poly.atIneq(i, k), 1);
-        for (unsigned k = d; k < d + p + 1; k++)
-          b2c2(j2, k - d) = Fraction(poly.atIneq(i, k), 1);
+        // All other inequalities are stored in a2 and b2c2.
+        // These are column-wise splits of the inequalities;
+        // a2 stores the coefficients of the variables, and
+        // b2c2 stores the coefficients of the parameters and the constant term.
+        for (unsigned k = 0; k < numVars; k++)
+          a2(j2, k) = poly.atIneq(i, k);
+        for (unsigned k = numVars; k < numVars + numParams + 1; k++)
+          b2c2(j2, k - numVars) = poly.atIneq(i, k);
         j2++;
-        // [A2 | B2 | c2]
       }
 
-    vertex = findVertex(subset); // d x (p+1)
+    // Find the vertex, if any, corresponding to the current subset of
+    // inequalities.
+    std::optional<ParamPoint> vertex = findVertex(subset); // d x (p+1)
 
     if (vertex == std::nullopt)
       continue;
+    // If this subset corresponds to a vertex, store it.
     vertices.push_back(*vertex);
     subsets.push_back(subset);
 
-    // Region is given by (A2 @ X + B2) p + (A2 @ y + c2) ≥ 0
-    // This is equivt to A2 @ [X | y] + [B2 | c2]
-    // We premultiply [X | y] with each row of A2 and add each row of [B2 | c2].
-    for (unsigned i = 0; i < n - d; i++) {
+    // Let the current vertex be [X | y], where
+    // X represents the coefficients of the parameters and
+    // y represents the constant term.
+
+    // The region (in parameter space) where this vertex is active is given
+    // by substituting the vertex into the *remaining* inequalities of the
+    // polytope (those which were not collected into `subset`), i.e.,
+    // [A2 | B2 | c2].
+    // Thus, the coefficients of the parameters after substitution become
+    // (A2 • X + B2)
+    // and the constant terms become
+    // (A2 • y + c2).
+    // The region is therefore given by
+    // (A2 • X + B2) p + (A2 • y + c2) ≥ 0
+    // This is equivalent to A2 • [X | y] + [B2 | c2]
+    // Thus we premultiply [X | y] with each row of A2
+    // and add each row of [B2 | c2].
+    for (unsigned i = 0; i < numIneqs - numVars; i++) {
       activeRegion.setRow(i, (*vertex).preMultiplyWithRow(a2.getRow(i)));
-      activeRegion.addToRow(i, b2c2.getRow(i), Fraction(1, 1));
+      activeRegion.addToRow(i, b2c2.getRow(i), 1);
     }
 
-    activeRegionNorm = IntMatrix(n - d, p + 1);
-    activeRegionRel =
-        IntegerRelation(PresburgerSpace::getRelationSpace(0, p, 0, 0));
-    lcmDenoms = 1;
-    for (unsigned i = 0; i < n - d; i++) {
-      for (unsigned j = 0; j < p + 1; j++)
+    // We convert the representation of the active region to an integers-only
+    // form so as to store it as an PresburgerRelation.
+    // We do this by taking the LCM of the denominators of all the coefficients
+    // and multiplying by it throughout.
+    IntMatrix activeRegionNorm = IntMatrix(numIneqs - numVars, numParams + 1);
+    IntegerRelation activeRegionRel =
+        IntegerRelation(PresburgerSpace::getRelationSpace(0, numParams, 0, 0));
+    MPInt lcmDenoms = MPInt(1);
+    for (unsigned i = 0; i < numIneqs - numVars; i++) {
+      for (unsigned j = 0; j < numParams + 1; j++)
         lcmDenoms = lcm(lcmDenoms, activeRegion(i, j).den);
-      for (unsigned j = 0; j < p + 1; j++)
+      for (unsigned j = 0; j < numParams + 1; j++)
         activeRegionNorm(i, j) =
-            (activeRegion(i, j) * Fraction(lcmDenoms, 1)).getAsInteger();
+            (activeRegion(i, j) * lcmDenoms).getAsInteger();
 
       activeRegionRel.addInequality(activeRegionNorm.getRow(i));
     }
@@ -308,17 +346,45 @@ mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly) {
     activeRegions.push_back(PresburgerRelation(activeRegionRel));
   }
 
-  // Clauss-Loechner chamber decomposition
+  // Now, we use Clauss-Loechner decomposition to identify regions in parameter
+  // space where each vertex is active. These regions (chambers) have the
+  // property that no two of them have a full-dimensional intersection, i.e.,
+  // they may share "faces" or "edges", but their intersection can only have
+  // up to numVars-1 dimensions.
+
+  // We maintain a list of regions and their associated vertex sets,
+  // initialized with the first vertex and its corresponding activity region.
   std::vector<std::pair<PresburgerRelation, std::vector<unsigned>>> chambers = {
       std::make_pair(activeRegions[0], std::vector({0u}))};
+  // Note that instead of storing lists of actual vertices, we store lists
+  // of indices. Thus the set {2, 3, 4} represents the vertex set
+  // {vertices[2], vertices[3], vertices[4]}.
+
   std::vector<std::pair<PresburgerRelation, std::vector<unsigned>>> newChambers;
+
+  // We iterate over the vertex set.
+  // For each vertex v_j and its activity region R_j,
+  // we examine all the current chambers R_i.
+  // If R_j has a full-dimensional intersection with an existing chamber R_i,
+  // then that chamber is replaced by two new ones:
+  // 1. the intersection R_i \cap R_j, where v_j is active;
+  // 2. the difference R_i - R_j, where v_j is inactive.
+  // Once we have examined all R_i, we add a final chamber
+  // R_j - (union of all existing chambers),
+  // in which only v_j is active.
   for (unsigned j = 1; j < vertices.size(); j++) {
     newChambers.clear();
+
     PresburgerRelation r_j = activeRegions[j];
     ParamPoint v_j = vertices[j];
+
     for (unsigned i = 0; i < chambers.size(); i++) {
       auto [r_i, v_i] = chambers[i];
 
+      // First, we check if the intersection of R_j and R_i.
+      // It is a disjoint union of convex regions in the parameter space,
+      // and so we know that it is full-dimensional if any of the disjuncts
+      // is full-dimensional.
       PresburgerRelation intersection = r_i.intersect(r_j);
       bool isFullDim = false;
       for (auto disjunct : intersection.getAllDisjuncts())
@@ -326,10 +392,14 @@ mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly) {
           isFullDim = true;
           break;
         }
-      isFullDim = (p == 0) || isFullDim;
+      isFullDim = (numParams == 0) || isFullDim;
+
+      // If the intersection is not full-dimensional, we do not modify
+      // the chamber list.
       if (!isFullDim)
         newChambers.push_back(chambers[i]);
       else {
+        // If it is, we add the intersection and the difference as new chambers.
         PresburgerRelation subtraction = r_i.subtract(r_j);
         newChambers.push_back(std::make_pair(subtraction, v_i));
 
@@ -337,11 +407,14 @@ mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly) {
         newChambers.push_back(std::make_pair(intersection, v_i));
       }
     }
+
+    // Finally we compute the chamber where only v_j is active by subtracting
+    // all existing chambers from R_j.
     for (auto chamber : newChambers)
       r_j = r_j.subtract(chamber.first);
-
     newChambers.push_back(std::make_pair(r_j, std::vector({j})));
 
+    // We filter `chambers` to remove empty regions.
     chambers.clear();
     for (auto chamber : newChambers) {
       bool empty = true;
@@ -355,18 +428,27 @@ mlir::presburger::detail::polytopeGeneratingFunction(PolyhedronH poly) {
     }
   }
 
-  SmallVector<MPInt> ineq(d + 1);
+  // Now, we compute the generating function. For each chamber, we iterate over
+  // the vertices active in it, and compute the generating function for each
+  // of them. The sum of these generating functions is the GF corresponding to
+  // the entire polytope.
+  SmallVector<MPInt> ineq(numVars + 1);
   for (auto chamber : chambers) {
-    chamberGf = GeneratingFunction(p, {}, {}, {});
+    GeneratingFunction chamberGf(numParams, {}, {}, {});
     for (unsigned i : chamber.second) {
-      tgtCone = defineHRep(d);
-      for (unsigned j = 0; j < d; j++) {
-        for (unsigned k = 0; k < d; k++)
+      // We collect the inequalities corresponding to each vertex.
+      // We only need the coefficients of the variables (NOT the parameters)
+      // as the generating function only depends on these.
+      ConeH tgtCone = defineHRep(numVars);
+      for (unsigned j = 0; j < numVars; j++) {
+        for (unsigned k = 0; k < numVars; k++)
           ineq[k] = subsets[i](j, k);
-        ineq[d] = subsets[i](j, d + p);
+        ineq[numVars] = subsets[i](j, numVars + numParams);
         tgtCone.addInequality(ineq);
       }
-      unimodCones = {std::make_pair(1, tgtCone)};
+      // We assume that the tangent cone is unimodular.
+      SmallVector<std::pair<int, ConeH>, 4> unimodCones = {
+          std::make_pair(1, tgtCone)};
       for (auto signedCone : unimodCones)
         chamberGf =
             chamberGf + unimodularConeGeneratingFunction(

>From 6176e0cc4008946f1ada6bea336d36dbee850ee5 Mon Sep 17 00:00:00 2001
From: Abhinav271828 <abhinav.m at research.iiit.ac.in>
Date: Mon, 22 Jan 2024 20:24:54 +0530
Subject: [PATCH 6/6] Doc

---
 .../mlir/Analysis/Presburger/IntegerRelation.h   |  3 +++
 mlir/lib/Analysis/Presburger/IntegerRelation.cpp | 16 ++++++++++++----
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index e7b03c44d33e38..141dc1f8fa97ed 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -711,8 +711,11 @@ class IntegerRelation {
   /// return `this \ set`.
   PresburgerRelation subtract(const PresburgerRelation &set) const;
 
+  // Remove equalities which have only zero coefficients.
   void removeTrivialEqualities();
 
+  // Verify whether the relation is full-dimensional, i.e.,
+  // has the same number of dimensions as the number of variables.
   bool isFullDim();
 
   void print(raw_ostream &os) const;
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index edb95b4072e1d8..b05075b2892ace 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -2511,24 +2511,32 @@ void IntegerRelation::removeTrivialEqualities() {
 }
 
 bool IntegerRelation::isFullDim() {
-  removeTrivialEqualities(); // to implement: remove equalities that are `0 = 0`
+  removeTrivialEqualities();
 
+  // If there is a non-trivial equality, the space cannot be full-dimensional.
   if (getNumEqualities() > 0)
-    return true;
+    return false;
 
+  // If along the direction of any of the inequalities, the upper and lower
+  // optima are the same, then the region is not full-dimensional.
   Simplex simplex(*this);
   for (unsigned i = 0; i < getNumInequalities(); i++) {
     auto ineq = inequalities.getRow(i);
     auto upOpt = simplex.computeOptimum(Simplex::Direction::Up, ineq);
     auto downOpt = simplex.computeOptimum(Simplex::Direction::Down, ineq);
+
     if (upOpt.getKind() == OptimumKind::Unbounded ||
         downOpt.getKind() == OptimumKind::Unbounded)
       continue;
+
+    // Check if the upper and lower optima are equal.
     if (upOpt.getKind() == OptimumKind::Bounded &&
-        downOpt.getKind() == OptimumKind::Bounded &&
-        *upOpt == *downOpt) // ineq == 0 holds for this
+        downOpt.getKind() == OptimumKind::Bounded && *upOpt == *downOpt)
       return false;
   }
+  // If none of the inequalities were such that the upper and lower optima
+  // along their direction were equal, then we conclude that the region is full
+  // dimensional.
   return true;
 }
 



More information about the Mlir-commits mailing list