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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Wed Feb 14 21:33:35 PST 2024


Author: Abhinav271828
Date: 2024-02-15T11:03:32+05:30
New Revision: 562790f371f230d8f67a1a8fb4b54e02e8d1e31f

URL: https://github.com/llvm/llvm-project/commit/562790f371f230d8f67a1a8fb4b54e02e8d1e31f
DIFF: https://github.com/llvm/llvm-project/commit/562790f371f230d8f67a1a8fb4b54e02e8d1e31f.diff

LOG: [MLIR][Presburger] Implement vertex enumeration and chamber decomposition for polytope generating function computation. (#78987)

We implement a function to compute the generating function corresponding
to a full-dimensional parametric polytope whose tangent cones are all
unimodular.
We fix a bug in unimodGenFunc to check the absolute value of the index.
We also implement Matrix<T>::negateMatrix() and Matrix<T>::scaleRow for
convenience.

Added: 
    

Modified: 
    mlir/include/mlir/Analysis/Presburger/Barvinok.h
    mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
    mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
    mlir/include/mlir/Analysis/Presburger/Matrix.h
    mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h
    mlir/include/mlir/Analysis/Presburger/Simplex.h
    mlir/include/mlir/Analysis/Presburger/Utils.h
    mlir/lib/Analysis/Presburger/Barvinok.cpp
    mlir/lib/Analysis/Presburger/IntegerRelation.cpp
    mlir/lib/Analysis/Presburger/Matrix.cpp
    mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
    mlir/lib/Analysis/Presburger/Simplex.cpp
    mlir/lib/Analysis/Presburger/Utils.cpp
    mlir/unittests/Analysis/Presburger/BarvinokTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index b70ec33b693235..f730a07393331a 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -27,7 +27,9 @@
 #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 <bitset>
 #include <optional>
 
 namespace mlir {
@@ -47,16 +49,22 @@ using PolyhedronV = IntMatrix;
 using ConeH = PolyhedronH;
 using ConeV = PolyhedronV;
 
-inline ConeH defineHRep(int numVars) {
+inline PolyhedronH defineHRep(int numVars, int numSymbols = 0) {
   // We don't distinguish between domain and range variables, so
   // we set the number of domain variables as 0 and the number of
   // range variables as the number of actual variables.
-  // There are no symbols (we don't work with parametric cones) and no local
-  // (existentially quantified) variables.
+  //
+  // numSymbols is the number of parameters.
+  //
+  // There are no local (existentially quantified) variables.
+  //
+  // The number of symbols is the number of parameters. By default, we consider
+  // nonparametric polyhedra.
+  //
   // Once the cone is defined, we use `addInequality()` to set inequalities.
-  return ConeH(PresburgerSpace::getSetSpace(/*numDims=*/numVars,
-                                            /*numSymbols=*/0,
-                                            /*numLocals=*/0));
+  return PolyhedronH(PresburgerSpace::getSetSpace(/*numDims=*/numVars,
+                                                  /*numSymbols=*/numSymbols,
+                                                  /*numLocals=*/0));
 }
 
 /// Get the index of a cone, i.e., the volume of the parallelepiped
@@ -81,8 +89,38 @@ ConeH getDual(ConeV cone);
 
 /// Compute the generating function for a unimodular cone.
 /// The input cone must be unimodular; it assert-fails otherwise.
-GeneratingFunction unimodularConeGeneratingFunction(ParamPoint vertex, int sign,
-                                                    ConeH cone);
+GeneratingFunction computeUnimodularConeGeneratingFunction(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 null.
+std::optional<ParamPoint> solveParametricEquations(FracMatrix equations);
+
+/// Given a list of possibly intersecting regions (PresburgerSet) and the
+/// generating functions active in each region, produce a pairwise disjoint
+/// list of regions (chambers) and identify the generating function of the
+/// polytope in each chamber.
+///
+/// "Disjoint" here means that the intersection of two chambers is no full-
+/// dimensional.
+///
+/// The returned list partitions the universe into parts depending on which
+/// subset of GFs is active there, and gives the sum of active GFs for each
+/// part.
+std::vector<std::pair<PresburgerSet, GeneratingFunction>>
+computeChamberDecomposition(
+    unsigned numSymbols, ArrayRef<std::pair<PresburgerSet, GeneratingFunction>>
+                             regionsAndGeneratingFunctions);
+
+/// Compute the generating function corresponding to a polytope.
+///
+/// All tangent cones of the polytope must be unimodular.
+std::vector<std::pair<PresburgerSet, GeneratingFunction>>
+computePolytopeGeneratingFunction(const 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/include/mlir/Analysis/Presburger/GeneratingFunction.h b/mlir/include/mlir/Analysis/Presburger/GeneratingFunction.h
index c38eab6efd0fc1..db5b6b6a959186 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+(const GeneratingFunction &gf) const {
     assert(numParam == gf.getNumParams() &&
            "two generating functions with 
diff erent 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..27dc382c1d5dbe 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -711,6 +711,17 @@ 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.,
+  // no equality holds for the relation.
+  //
+  // If there are no variables, it always returns true.
+  // If there is at least one variable and the relation is empty, it returns
+  // false.
+  bool isFullDim();
+
   void print(raw_ostream &os) const;
   void dump() const;
 
@@ -871,6 +882,26 @@ class IntegerPolyhedron : public IntegerRelation {
                           /*numReservedEqualities=*/0,
                           /*numReservedCols=*/space.getNumVars() + 1, space) {}
 
+  /// Constructs a relation with the specified number of dimensions and symbols
+  /// and adds the given inequalities.
+  explicit IntegerPolyhedron(const PresburgerSpace &space,
+                             IntMatrix inequalities)
+      : IntegerPolyhedron(space) {
+    for (unsigned i = 0, e = inequalities.getNumRows(); i < e; i++)
+      addInequality(inequalities.getRow(i));
+  }
+
+  /// Constructs a relation with the specified number of dimensions and symbols
+  /// and adds the given inequalities, after normalizing row-wise to integer
+  /// values.
+  explicit IntegerPolyhedron(const PresburgerSpace &space,
+                             FracMatrix inequalities)
+      : IntegerPolyhedron(space) {
+    IntMatrix ineqsNormalized = inequalities.normalizeRows();
+    for (unsigned i = 0, e = inequalities.getNumRows(); i < e; i++)
+      addInequality(ineqsNormalized.getRow(i));
+  }
+
   /// Construct a set from an IntegerRelation. The relation should have
   /// no domain vars.
   explicit IntegerPolyhedron(const IntegerRelation &rel)

diff  --git a/mlir/include/mlir/Analysis/Presburger/Matrix.h b/mlir/include/mlir/Analysis/Presburger/Matrix.h
index 0d4a593a95b1c9..4484ebc747e61c 100644
--- a/mlir/include/mlir/Analysis/Presburger/Matrix.h
+++ b/mlir/include/mlir/Analysis/Presburger/Matrix.h
@@ -20,6 +20,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include <bitset>
 #include <cassert>
 
 namespace mlir {
@@ -73,6 +74,8 @@ class Matrix {
 
   T operator()(unsigned row, unsigned column) const { return at(row, column); }
 
+  bool operator==(const Matrix<T> &m) const;
+
   /// Swap the given columns.
   void swapColumns(unsigned column, unsigned otherColumn);
 
@@ -142,6 +145,9 @@ class Matrix {
   /// Add `scale` multiples of the rowVec row to the specified row.
   void addToRow(unsigned row, ArrayRef<T> rowVec, const T &scale);
 
+  /// Multiply the specified row by a factor of `scale`.
+  void scaleRow(unsigned row, const T &scale);
+
   /// Add `scale` multiples of the source column to the target column.
   void addToColumn(unsigned sourceColumn, unsigned targetColumn,
                    const T &scale);
@@ -156,6 +162,9 @@ class Matrix {
   /// Negate the specified row.
   void negateRow(unsigned row);
 
+  /// Negate the entire matrix.
+  void negateMatrix();
+
   /// The given vector is interpreted as a row vector v. Post-multiply v with
   /// this matrix, say M, and return vM.
   SmallVector<T, 8> preMultiplyWithRow(ArrayRef<T> rowVec) const;
@@ -184,6 +193,19 @@ class Matrix {
   // Transpose the matrix without modifying it.
   Matrix<T> transpose() const;
 
+  // Copy the cells in the intersection of
+  // the rows between `fromRows` and `toRows` and
+  // the columns between `fromColumns` and `toColumns`, both inclusive.
+  Matrix<T> getSubMatrix(unsigned fromRow, unsigned toRow, unsigned fromColumn,
+                         unsigned toColumn) const;
+
+  /// Split the rows of a matrix into two matrices according to which bits are
+  /// 1 and which are 0 in a given bitset.
+  ///
+  /// The first matrix returned has the rows corresponding to 1 and the second
+  /// corresponding to 2.
+  std::pair<Matrix<T>, Matrix<T>> splitByBitset(ArrayRef<int> indicator);
+
   /// Print the matrix.
   void print(raw_ostream &os) const;
   void dump() const;
@@ -297,6 +319,10 @@ class FracMatrix : public Matrix<Fraction> {
   // paper](https://www.cs.cmu.edu/~avrim/451f11/lectures/lect1129_LLL.pdf)
   // calls `y`, usually 3/4.
   void LLL(Fraction delta);
+
+  // Multiply each row of the matrix by the LCM of the denominators, thereby
+  // converting it to an integer matrix.
+  IntMatrix normalizeRows() const;
 };
 
 } // namespace presburger

diff  --git a/mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h b/mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h
index c6b00eca90733a..9634df6d58a1a1 100644
--- a/mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h
@@ -217,6 +217,10 @@ class PresburgerRelation {
   /// redundencies.
   PresburgerRelation simplify() const;
 
+  /// Return whether the given PresburgerRelation is full-dimensional. By full-
+  /// dimensional we mean that it is not flat along any dimension.
+  bool isFullDim() const;
+
   /// Print the set's internal state.
   void print(raw_ostream &os) const;
   void dump() const;

diff  --git a/mlir/include/mlir/Analysis/Presburger/Simplex.h b/mlir/include/mlir/Analysis/Presburger/Simplex.h
index 9482f69b31cd66..7ee74c150867c1 100644
--- a/mlir/include/mlir/Analysis/Presburger/Simplex.h
+++ b/mlir/include/mlir/Analysis/Presburger/Simplex.h
@@ -771,6 +771,12 @@ class Simplex : public SimplexBase {
   std::pair<MaybeOptimum<MPInt>, MaybeOptimum<MPInt>>
   computeIntegerBounds(ArrayRef<MPInt> coeffs);
 
+  /// Check if the simplex takes only one rational value along the
+  /// direction of `coeffs`.
+  ///
+  /// `this` must be nonempty.
+  bool isFlatAlong(ArrayRef<MPInt> coeffs);
+
   /// Returns true if the polytope is unbounded, i.e., extends to infinity in
   /// some direction. Otherwise, returns false.
   bool isUnbounded();

diff  --git a/mlir/include/mlir/Analysis/Presburger/Utils.h b/mlir/include/mlir/Analysis/Presburger/Utils.h
index e6d29e4ef6d062..38262a65f97542 100644
--- a/mlir/include/mlir/Analysis/Presburger/Utils.h
+++ b/mlir/include/mlir/Analysis/Presburger/Utils.h
@@ -286,6 +286,8 @@ Fraction dotProduct(ArrayRef<Fraction> a, ArrayRef<Fraction> b);
 std::vector<Fraction> multiplyPolynomials(ArrayRef<Fraction> a,
                                           ArrayRef<Fraction> b);
 
+bool isRangeZero(ArrayRef<Fraction> arr);
+
 } // namespace presburger
 } // namespace mlir
 

diff  --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index d2752ded6b43f5..b6d1f99df8ba55 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;
@@ -76,7 +77,8 @@ MPInt mlir::presburger::detail::getIndex(ConeV cone) {
 /// num is computed by expressing the vertex as a weighted
 /// sum of the generators, and then taking the floor of the
 /// coefficients.
-GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
+GeneratingFunction
+mlir::presburger::detail::computeUnimodularConeGeneratingFunction(
     ParamPoint vertex, int sign, ConeH cone) {
   // Consider a cone with H-representation [0  -1].
   //                                       [-1 -2]
@@ -84,7 +86,7 @@ GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
   //                                       [-1 -1/2 1]
 
   // `cone` must be unimodular.
-  assert(getIndex(getDual(cone)) == 1 && "input cone is not unimodular!");
+  assert(abs(getIndex(getDual(cone))) == 1 && "input cone is not unimodular!");
 
   unsigned numVar = cone.getNumVars();
   unsigned numIneq = cone.getNumInequalities();
@@ -147,6 +149,304 @@ 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::solveParametricEquations(FracMatrix equations) {
+  // equations is a d x (d + p + 1) matrix.
+  // Each row represents an equation.
+  unsigned d = equations.getNumRows();
+  unsigned numCols = equations.getNumColumns();
+
+  // If the determinant is zero, there is no unique solution.
+  // Thus we return null.
+  if (FracMatrix(equations.getSubMatrix(/*fromRow=*/0, /*toRow=*/d - 1,
+                                        /*fromColumn=*/0,
+                                        /*toColumn=*/d - 1))
+          .determinant() == 0)
+    return std::nullopt;
+
+  // Perform row operations to make each column all zeros except for the
+  // diagonal element, which is made to be one.
+  for (unsigned i = 0; i < d; ++i) {
+    // First ensure that the diagonal element is nonzero, by swapping
+    // it with a row that is non-zero at column i.
+    if (equations(i, i) != 0)
+      continue;
+    for (unsigned j = i + 1; j < d; ++j) {
+      if (equations(j, i) == 0)
+        continue;
+      equations.swapRows(j, i);
+      break;
+    }
+
+    Fraction diagElement = equations(i, i);
+
+    // Apply row operations to make all elements except the diagonal to zero.
+    for (unsigned j = 0; j < d; ++j) {
+      if (i == j)
+        continue;
+      if (equations(j, i) == 0)
+        continue;
+      // Apply row operations to make element (j, i) zero by subtracting the
+      // ith row, appropriately scaled.
+      Fraction currentElement = equations(j, i);
+      equations.addToRow(/*sourceRow=*/i, /*targetRow=*/j,
+                         /*scale=*/-currentElement / diagElement);
+    }
+  }
+
+  // Rescale diagonal elements to 1.
+  for (unsigned i = 0; i < d; ++i)
+    equations.scaleRow(i, 1 / equations(i, i));
+
+  // Now we have reduced the equations to the form
+  // x_i + b_1' m_1 + ... + b_p' m_p + c' = 0
+  // i.e. each variable appears exactly once in the system, and has coefficient
+  // one.
+  //
+  // Thus we have
+  // x_i = - b_1' m_1 - ... - b_p' m_p - c
+  // and so we return the negation of the last p + 1 columns of the matrix.
+  //
+  // We copy these columns and return them.
+  ParamPoint vertex =
+      equations.getSubMatrix(/*fromRow=*/0, /*toRow=*/d - 1,
+                             /*fromColumn=*/d, /*toColumn=*/numCols - 1);
+  vertex.negateMatrix();
+  return vertex;
+}
+
+/// This is an implementation of the Clauss-Loechner algorithm for chamber
+/// decomposition.
+///
+/// We maintain a list of pairwise disjoint chambers and the generating
+/// functions corresponding to each one. We iterate over the list of regions,
+/// each time adding the current region's generating function to the chambers
+/// where it is active and separating the chambers where it is not.
+///
+/// Given the region each generating function is active in, for each subset of
+/// generating functions the region that (the sum of) precisely this subset is
+/// in, is the intersection of the regions that these are active in,
+/// intersected with the complements of the remaining regions.
+std::vector<std::pair<PresburgerSet, GeneratingFunction>>
+mlir::presburger::detail::computeChamberDecomposition(
+    unsigned numSymbols, ArrayRef<std::pair<PresburgerSet, GeneratingFunction>>
+                             regionsAndGeneratingFunctions) {
+  assert(!regionsAndGeneratingFunctions.empty() &&
+         "there must be at least one chamber!");
+  // We maintain a list of regions and their associated generating function
+  // initialized with the universe and the empty generating function.
+  std::vector<std::pair<PresburgerSet, GeneratingFunction>> chambers = {
+      {PresburgerSet::getUniverse(PresburgerSpace::getSetSpace(numSymbols)),
+       GeneratingFunction(numSymbols, {}, {}, {})}};
+
+  // We iterate over the region list.
+  //
+  // For each activity region R_j (corresponding to the generating function
+  // gf_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 the generating function is
+  // gf_i + gf_j.
+  // 2. the 
diff erence R_i - R_j, where the generating function is gf_i.
+  //
+  // At each step, we define a new chamber list after considering gf_j,
+  // replacing and appending chambers as discussed above.
+  //
+  // The loop has the invariant that the union over all the chambers gives the
+  // universe at every step.
+  for (const auto &[region, generatingFunction] :
+       regionsAndGeneratingFunctions) {
+    std::vector<std::pair<PresburgerSet, GeneratingFunction>> newChambers;
+
+    for (const auto &[currentRegion, currentGeneratingFunction] : chambers) {
+      PresburgerSet intersection = currentRegion.intersect(region);
+
+      // If the intersection is not full-dimensional, we do not modify
+      // the chamber list.
+      if (!intersection.isFullDim()) {
+        newChambers.emplace_back(currentRegion, currentGeneratingFunction);
+        continue;
+      }
+
+      // If it is, we add the intersection and the 
diff erence as chambers.
+      newChambers.emplace_back(intersection,
+                               currentGeneratingFunction + generatingFunction);
+      newChambers.emplace_back(currentRegion.subtract(region),
+                               currentGeneratingFunction);
+    }
+    chambers = std::move(newChambers);
+  }
+
+  return chambers;
+}
+
+/// For a polytope expressed as a set of n inequalities, compute the generating
+/// function corresponding to the lattice points included in the polytope. This
+/// algorithm has three main steps:
+/// 1. Enumerate the vertices, by iterating over subsets of inequalities and
+///    checking for satisfiability. For each d-subset of inequalities (where d
+///    is the number of variables), we solve to obtain the vertex in terms of
+///    the parameters, and then check for the region in parameter space where
+///    this vertex satisfies the remaining (n - d) inequalities.
+/// 2. For each vertex, identify the tangent cone and compute the generating
+///    function corresponding to it. The generating function depends on the
+///    parametric expression of the vertex and the (non-parametric) generators
+///    of the tangent cone.
+/// 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<PresburgerSet, GeneratingFunction>>
+mlir::presburger::detail::computePolytopeGeneratingFunction(
+    const PolyhedronH &poly) {
+  unsigned numVars = poly.getNumRangeVars();
+  unsigned numSymbols = poly.getNumSymbolVars();
+  unsigned numIneqs = poly.getNumInequalities();
+
+  // We store a list of the computed vertices.
+  std::vector<ParamPoint> vertices;
+  // For each vertex, we store the corresponding active region and the
+  // generating functions of the tangent cone, in order.
+  std::vector<std::pair<PresburgerSet, GeneratingFunction>>
+      regionsAndGeneratingFunctions;
+
+  // We iterate over all subsets of inequalities with cardinality numVars,
+  // using permutations of numVars 1's and (numIneqs - numVars) 0's.
+  //
+  // For a given permutation, we consider a subset which contains
+  // the i'th inequality if the i'th bit in the bitset is 1.
+  //
+  // We start with the permutation that takes the last numVars inequalities.
+  SmallVector<int> indicator(numIneqs);
+  for (unsigned i = numIneqs - numVars; i < numIneqs; ++i)
+    indicator[i] = 1;
+
+  do {
+    // Collect the inequalities corresponding to the bits which are set
+    // and the remaining ones.
+    auto [subset, remainder] = poly.getInequalities().splitByBitset(indicator);
+    // 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.
+    FracMatrix a2(numIneqs - numVars, numVars);
+    FracMatrix b2c2(numIneqs - numVars, numSymbols + 1);
+    a2 = FracMatrix(
+        remainder.getSubMatrix(0, numIneqs - numVars - 1, 0, numVars - 1));
+    b2c2 = FracMatrix(remainder.getSubMatrix(0, numIneqs - numVars - 1, numVars,
+                                             numVars + numSymbols));
+
+    // Find the vertex, if any, corresponding to the current subset of
+    // inequalities.
+    std::optional<ParamPoint> vertex =
+        solveParametricEquations(FracMatrix(subset)); // d x (p+1)
+
+    if (!vertex)
+      continue;
+    if (std::find(vertices.begin(), vertices.end(), vertex) != vertices.end())
+      continue;
+    // If this subset corresponds to a vertex that has not been considered,
+    // store it.
+    vertices.push_back(*vertex);
+
+    // If a vertex is formed by the intersection of more than d facets, we
+    // assume that any d-subset of these facets can be solved to obtain its
+    // expression. This assumption is valid because, if the vertex has two
+    // distinct parametric expressions, then a nontrivial equality among the
+    // parameters holds, which is a contradiction as we know the parameter
+    // space to be full-dimensional.
+
+    // 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., into the
+    // inequalities [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].
+    FracMatrix activeRegion(numIneqs - numVars, numSymbols + 1);
+    for (unsigned i = 0; i < numIneqs - numVars; i++) {
+      activeRegion.setRow(i, vertex->preMultiplyWithRow(a2.getRow(i)));
+      activeRegion.addToRow(i, b2c2.getRow(i), 1);
+    }
+
+    // We convert the representation of the active region to an integers-only
+    // form so as to store it as a PresburgerSet.
+    IntegerPolyhedron activeRegionRel(
+        PresburgerSpace::getRelationSpace(0, numSymbols, 0, 0), activeRegion);
+
+    // Now, we compute the generating function at this vertex.
+    // We collect the inequalities corresponding to each vertex to compute
+    // the tangent cone at that vertex.
+    //
+    // We only need the coefficients of the variables (NOT the parameters)
+    // as the generating function only depends on these.
+    // We translate the cones to be pointed at the origin by making the
+    // constant terms zero.
+    ConeH tangentCone = defineHRep(numVars);
+    for (unsigned j = 0, e = subset.getNumRows(); j < e; ++j) {
+      SmallVector<MPInt> ineq(numVars + 1);
+      for (unsigned k = 0; k < numVars; ++k)
+        ineq[k] = subset(j, k);
+      tangentCone.addInequality(ineq);
+    }
+    // We assume that the tangent cone is unimodular, so there is no need
+    // to decompose it.
+    //
+    // In the general case, the unimodular decomposition may have several
+    // cones.
+    GeneratingFunction vertexGf(numSymbols, {}, {}, {});
+    SmallVector<std::pair<int, ConeH>, 4> unimodCones = {{1, tangentCone}};
+    for (std::pair<int, ConeH> signedCone : unimodCones) {
+      auto [sign, cone] = signedCone;
+      vertexGf = vertexGf +
+                 computeUnimodularConeGeneratingFunction(*vertex, sign, cone);
+    }
+    // We store the vertex we computed with the generating function of its
+    // tangent cone.
+    regionsAndGeneratingFunctions.emplace_back(PresburgerSet(activeRegionRel),
+                                               vertexGf);
+  } while (std::next_permutation(indicator.begin(), indicator.end()));
+
+  // 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 "facets" or "edges", but their intersection can only have
+  // up to numVars - 1 dimensions.
+  //
+  // In each chamber, we sum up the generating functions of the active vertices
+  // to find the generating function of the polytope.
+  return computeChamberDecomposition(numSymbols, regionsAndGeneratingFunctions);
+}
+
 /// 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..2ac271e2e05531 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -26,6 +26,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -2498,6 +2499,31 @@ void IntegerRelation::printSpace(raw_ostream &os) const {
   os << getNumConstraints() << " constraints\n";
 }
 
+void IntegerRelation::removeTrivialEqualities() {
+  for (int i = getNumEqualities() - 1; i >= 0; --i)
+    if (rangeIsZero(getEquality(i)))
+      removeEquality(i);
+}
+
+bool IntegerRelation::isFullDim() {
+  if (getNumVars() == 0)
+    return true;
+  if (isEmpty())
+    return false;
+
+  // If there is a non-trivial equality, the space cannot be full-dimensional.
+  removeTrivialEqualities();
+  if (getNumEqualities() > 0)
+    return false;
+
+  // The polytope is full-dimensional iff it is not flat along any of the
+  // inequality directions.
+  Simplex simplex(*this);
+  return llvm::none_of(llvm::seq<int>(getNumInequalities()), [&](int i) {
+    return simplex.isFlatAlong(getInequality(i));
+  });
+}
+
 void IntegerRelation::print(raw_ostream &os) const {
   assert(hasConsistentState());
   printSpace(os);

diff  --git a/mlir/lib/Analysis/Presburger/Matrix.cpp b/mlir/lib/Analysis/Presburger/Matrix.cpp
index bd7f7f58a932f3..4cb6e6b16bc878 100644
--- a/mlir/lib/Analysis/Presburger/Matrix.cpp
+++ b/mlir/lib/Analysis/Presburger/Matrix.cpp
@@ -29,6 +29,22 @@ Matrix<T>::Matrix(unsigned rows, unsigned columns, unsigned reservedRows,
   data.reserve(std::max(nRows, reservedRows) * nReservedColumns);
 }
 
+/// We cannot use the default implementation of operator== as it compares
+/// fields like `reservedColumns` etc., which are not part of the data.
+template <typename T>
+bool Matrix<T>::operator==(const Matrix<T> &m) const {
+  if (nRows != m.getNumRows())
+    return false;
+  if (nColumns != m.getNumColumns())
+    return false;
+
+  for (unsigned i = 0; i < nRows; i++)
+    if (getRow(i) != m.getRow(i))
+      return false;
+
+  return true;
+}
+
 template <typename T>
 Matrix<T> Matrix<T>::identity(unsigned dimension) {
   Matrix matrix(dimension, dimension);
@@ -295,6 +311,12 @@ void Matrix<T>::addToRow(unsigned row, ArrayRef<T> rowVec, const T &scale) {
     at(row, col) += scale * rowVec[col];
 }
 
+template <typename T>
+void Matrix<T>::scaleRow(unsigned row, const T &scale) {
+  for (unsigned col = 0; col < nColumns; ++col)
+    at(row, col) *= scale;
+}
+
 template <typename T>
 void Matrix<T>::addToColumn(unsigned sourceColumn, unsigned targetColumn,
                             const T &scale) {
@@ -316,6 +338,12 @@ void Matrix<T>::negateRow(unsigned row) {
     at(row, column) = -at(row, column);
 }
 
+template <typename T>
+void Matrix<T>::negateMatrix() {
+  for (unsigned row = 0; row < nRows; ++row)
+    negateRow(row);
+}
+
 template <typename T>
 SmallVector<T, 8> Matrix<T>::preMultiplyWithRow(ArrayRef<T> rowVec) const {
   assert(rowVec.size() == getNumRows() && "Invalid row vector dimension!");
@@ -354,6 +382,22 @@ static void modEntryColumnOperation(Matrix<MPInt> &m, unsigned row,
   otherMatrix.addToColumn(sourceCol, targetCol, ratio);
 }
 
+template <typename T>
+Matrix<T> Matrix<T>::getSubMatrix(unsigned fromRow, unsigned toRow,
+                                  unsigned fromColumn,
+                                  unsigned toColumn) const {
+  assert(fromRow <= toRow && "end of row range must be after beginning!");
+  assert(toRow < nRows && "end of row range out of bounds!");
+  assert(fromColumn <= toColumn &&
+         "end of column range must be after beginning!");
+  assert(toColumn < nColumns && "end of column range out of bounds!");
+  Matrix<T> subMatrix(toRow - fromRow + 1, toColumn - fromColumn + 1);
+  for (unsigned i = fromRow; i <= toRow; ++i)
+    for (unsigned j = fromColumn; j <= toColumn; ++j)
+      subMatrix(i - fromRow, j - fromColumn) = at(i, j);
+  return subMatrix;
+}
+
 template <typename T>
 void Matrix<T>::print(raw_ostream &os) const {
   for (unsigned row = 0; row < nRows; ++row) {
@@ -363,6 +407,21 @@ void Matrix<T>::print(raw_ostream &os) const {
   }
 }
 
+/// We iterate over the `indicator` bitset, checking each bit. If a bit is 1,
+/// we append it to one matrix, and if it is zero, we append it to the other.
+template <typename T>
+std::pair<Matrix<T>, Matrix<T>>
+Matrix<T>::splitByBitset(ArrayRef<int> indicator) {
+  Matrix<T> rowsForOne(0, nColumns), rowsForZero(0, nColumns);
+  for (unsigned i = 0; i < nRows; i++) {
+    if (indicator[i] == 1)
+      rowsForOne.appendExtraRow(getRow(i));
+    else
+      rowsForZero.appendExtraRow(getRow(i));
+  }
+  return {rowsForOne, rowsForZero};
+}
+
 template <typename T>
 void Matrix<T>::dump() const {
   print(llvm::errs());
@@ -697,3 +756,20 @@ void FracMatrix::LLL(Fraction delta) {
     }
   }
 }
+
+IntMatrix FracMatrix::normalizeRows() const {
+  unsigned numRows = getNumRows();
+  unsigned numColumns = getNumColumns();
+  IntMatrix normalized(numRows, numColumns);
+
+  MPInt lcmDenoms = MPInt(1);
+  for (unsigned i = 0; i < numRows; i++) {
+    // For a row, first compute the LCM of the denominators.
+    for (unsigned j = 0; j < numColumns; j++)
+      lcmDenoms = lcm(lcmDenoms, at(i, j).den);
+    // Then, multiply by it throughout and convert to integers.
+    for (unsigned j = 0; j < numColumns; j++)
+      normalized(i, j) = (at(i, j) * lcmDenoms).getAsInteger();
+  }
+  return normalized;
+}

diff  --git a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp b/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
index 787fc1c659a12e..3af6baae0e7001 100644
--- a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
@@ -1041,6 +1041,12 @@ PresburgerRelation PresburgerRelation::simplify() const {
   return result;
 }
 
+bool PresburgerRelation::isFullDim() const {
+  return llvm::any_of(getAllDisjuncts(), [&](IntegerRelation disjunct) {
+    return disjunct.isFullDim();
+  });
+}
+
 void PresburgerRelation::print(raw_ostream &os) const {
   os << "Number of Disjuncts: " << getNumDisjuncts() << "\n";
   for (const IntegerRelation &disjunct : disjuncts) {

diff  --git a/mlir/lib/Analysis/Presburger/Simplex.cpp b/mlir/lib/Analysis/Presburger/Simplex.cpp
index 42bbc3363d5830..1969cce93ad2e0 100644
--- a/mlir/lib/Analysis/Presburger/Simplex.cpp
+++ b/mlir/lib/Analysis/Presburger/Simplex.cpp
@@ -2104,6 +2104,19 @@ Simplex::computeIntegerBounds(ArrayRef<MPInt> coeffs) {
   return {minRoundedUp, maxRoundedDown};
 }
 
+bool Simplex::isFlatAlong(ArrayRef<MPInt> coeffs) {
+  assert(!isEmpty() && "cannot check for flatness of empty simplex!");
+  auto upOpt = computeOptimum(Simplex::Direction::Up, coeffs);
+  auto downOpt = computeOptimum(Simplex::Direction::Down, coeffs);
+
+  if (!upOpt.isBounded())
+    return false;
+  if (!downOpt.isBounded())
+    return false;
+
+  return *upOpt == *downOpt;
+}
+
 void SimplexBase::print(raw_ostream &os) const {
   os << "rows = " << getNumRows() << ", columns = " << getNumColumns() << "\n";
   if (empty)

diff  --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp
index a8d860885ef106..f717a4de5d7283 100644
--- a/mlir/lib/Analysis/Presburger/Utils.cpp
+++ b/mlir/lib/Analysis/Presburger/Utils.cpp
@@ -564,4 +564,8 @@ std::vector<Fraction> presburger::multiplyPolynomials(ArrayRef<Fraction> a,
     convolution.push_back(sum);
   }
   return convolution;
-}
\ No newline at end of file
+}
+
+bool presburger::isRangeZero(ArrayRef<Fraction> arr) {
+  return llvm::all_of(arr, [&](Fraction f) { return f == 0; });
+}

diff  --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index 919aaa7a428593..5e279b542fdf95 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -1,5 +1,6 @@
 #include "mlir/Analysis/Presburger/Barvinok.h"
 #include "./Utils.h"
+#include "Parser.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
@@ -59,7 +60,8 @@ TEST(BarvinokTest, unimodularConeGeneratingFunction) {
   ParamPoint vertex =
       makeFracMatrix(2, 3, {{2, 2, 0}, {-1, -Fraction(1, 2), 1}});
 
-  GeneratingFunction gf = unimodularConeGeneratingFunction(vertex, 1, cone);
+  GeneratingFunction gf =
+      computeUnimodularConeGeneratingFunction(vertex, 1, cone);
 
   EXPECT_EQ_REPR_GENERATINGFUNCTION(
       gf, GeneratingFunction(
@@ -74,7 +76,7 @@ TEST(BarvinokTest, unimodularConeGeneratingFunction) {
 
   vertex = makeFracMatrix(3, 2, {{5, 2}, {6, 2}, {7, 1}});
 
-  gf = unimodularConeGeneratingFunction(vertex, 1, cone);
+  gf = computeUnimodularConeGeneratingFunction(vertex, 1, cone);
 
   EXPECT_EQ_REPR_GENERATINGFUNCTION(
       gf,
@@ -125,7 +127,7 @@ TEST(BarvinokTest, getCoefficientInRationalFunction) {
   EXPECT_EQ(coeff.getConstantTerm(), Fraction(55, 64));
 }
 
-TEST(BarvinokTest, computeNumTerms) {
+TEST(BarvinokTest, computeNumTermsCone) {
   // The following test is taken from
   // Verdoolaege, Sven, et al. "Counting integer points in parametric
   // polytopes using Barvinok's rational functions." Algorithmica 48 (2007):
@@ -233,4 +235,69 @@ TEST(BarvinokTest, computeNumTerms) {
     for (unsigned j = 0; j < 2; j++)
       for (unsigned k = 0; k < 2; k++)
         EXPECT_EQ(count[i][j][k], 1);
-}
\ No newline at end of file
+}
+
+/// We define some simple polyhedra with unimodular tangent cones and verify
+/// that the returned generating functions correspond to those calculated by
+/// hand.
+TEST(BarvinokTest, computeNumTermsPolytope) {
+  // A cube of side 1.
+  PolyhedronH poly =
+      parseRelationFromSet("(x, y, z) : (x >= 0, y >= 0, z >= 0, -x + 1 >= 0, "
+                           "-y + 1 >= 0, -z + 1 >= 0)",
+                           0);
+
+  std::vector<std::pair<PresburgerSet, GeneratingFunction>> count =
+      computePolytopeGeneratingFunction(poly);
+  // There is only one chamber, as it is non-parametric.
+  EXPECT_EQ(count.size(), 9u);
+
+  GeneratingFunction gf = count[0].second;
+  EXPECT_EQ_REPR_GENERATINGFUNCTION(
+      gf,
+      GeneratingFunction(
+          0, {1, 1, 1, 1, 1, 1, 1, 1},
+          {makeFracMatrix(1, 3, {{1, 1, 1}}), makeFracMatrix(1, 3, {{0, 1, 1}}),
+           makeFracMatrix(1, 3, {{0, 1, 1}}), makeFracMatrix(1, 3, {{0, 0, 1}}),
+           makeFracMatrix(1, 3, {{0, 1, 1}}), makeFracMatrix(1, 3, {{0, 0, 1}}),
+           makeFracMatrix(1, 3, {{0, 0, 1}}),
+           makeFracMatrix(1, 3, {{0, 0, 0}})},
+          {{{-1, 0, 0}, {0, -1, 0}, {0, 0, -1}},
+           {{0, 0, 1}, {-1, 0, 0}, {0, -1, 0}},
+           {{0, 1, 0}, {-1, 0, 0}, {0, 0, -1}},
+           {{0, 1, 0}, {0, 0, 1}, {-1, 0, 0}},
+           {{1, 0, 0}, {0, -1, 0}, {0, 0, -1}},
+           {{1, 0, 0}, {0, 0, 1}, {0, -1, 0}},
+           {{1, 0, 0}, {0, 1, 0}, {0, 0, -1}},
+           {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}}}));
+
+  // A right-angled triangle with side p.
+  poly =
+      parseRelationFromSet("(x, y)[N] : (x >= 0, y >= 0, -x - y + N >= 0)", 0);
+
+  count = computePolytopeGeneratingFunction(poly);
+  // There is only one chamber: p ≥ 0
+  EXPECT_EQ(count.size(), 4u);
+
+  gf = count[0].second;
+  EXPECT_EQ_REPR_GENERATINGFUNCTION(
+      gf, GeneratingFunction(
+              1, {1, 1, 1},
+              {makeFracMatrix(2, 2, {{0, 1}, {0, 0}}),
+               makeFracMatrix(2, 2, {{0, 1}, {0, 0}}),
+               makeFracMatrix(2, 2, {{0, 0}, {0, 0}})},
+              {{{-1, 1}, {-1, 0}}, {{1, -1}, {0, -1}}, {{1, 0}, {0, 1}}}));
+
+  // Cartesian product of a cube with side M and a right triangle with side N.
+  poly = parseRelationFromSet(
+      "(x, y, z, w, a)[M, N] : (x >= 0, y >= 0, z >= 0, -x + M >= 0, -y + M >= "
+      "0, -z + M >= 0, w >= 0, a >= 0, -w - a + N >= 0)",
+      0);
+
+  count = computePolytopeGeneratingFunction(poly);
+
+  EXPECT_EQ(count.size(), 25u);
+
+  gf = count[0].second;
+  EXPECT_EQ(gf.getNumerators().size(), 24u);
+}


        


More information about the Mlir-commits mailing list