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

[Arjun P]

================
@@ -147,6 +148,319 @@ 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) {
+  // equations is a d x (d + p + 1) matrix.
+  // Each row represents an equation.
+
+  unsigned numEqs = equations.getNumRows();
+  unsigned numCols = equations.getNumColumns();
+
+  // First, we check that the system has a solution, and return
+  // null if not.
+  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)
+    return std::nullopt;
+
+  // We work with rational numbers.
+  FracMatrix equationsF(equations);
+
+  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 < numEqs; ++j) {
+        if (equationsF(j, i) != 0) {
+          equationsF.swapRows(j, i);
+          break;
+        }
+      }
+    }
+
+    Fraction b = equationsF(i, i);
+
+    // Set all elements except the diagonal to zero.
+    for (unsigned j = 0; j < numEqs; ++j) {
+      if (equationsF(j, i) == 0 || j == i)
+        continue;
+      // 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);
+    }
+  }
+
+  // Rescale diagonal elements to 1.
+  for (unsigned i = 0; i < numEqs; ++i) {
+    Fraction a = equationsF(i, i);
+    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(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 numVars = poly.getNumRangeVars();
+  unsigned numParams = poly.getNumSymbolVars();
+  unsigned numIneqs = poly.getNumInequalities();
+
+  // 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({});
+
+  // 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;
+
+  FracMatrix a2(numIneqs - numVars, numVars);
+  FracMatrix b2c2(numIneqs - numVars, numParams + 1);
+
+  // We iterate over all subsets of inequalities with cardinality numVars,
+  // using bitsets to enumerate.
----------------
Superty wrote:

if you want to do this this way, "using bitsets to enumerate" needs a bit more elaboration.

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