[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.
+  // The largest possible bitset that corresponds to such a subset can be
+  // written as numVar 1's followed by (numIneqs - numVars) 0's.
+  unsigned upperBound = ((1ul << numVars) - 1ul) << (numIneqs - numVars);
+  for (std::bitset<16> indicator(((1ul << numVars) - 1ul)
+                                 << (numIneqs - numVars));
+       indicator.to_ulong() <= upperBound;
----------------
Superty wrote:

Please abstract the body of the loop into multiple other functions. Functions are not only for code reuse, but also as a method of organization. Once you do this, you should be able to get the high level steps by looking at the loop body. And then if you want to go into the details of any step, you can by going to that function definitino. Right now, it's impossible to understand what's doing on at a glance because of the sprawling function body.

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