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

[Arjun P]

================
@@ -147,6 +149,307 @@ 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) {
+    Fraction diagElement = equations(i, i);
+    for (unsigned j = 0; j < numCols; ++j)
+      equations(i, j) = equations(i, j) / diagElement;
+  }
+
+  // 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);
+  for (unsigned i = 0; i < d; ++i)
+    vertex.negateRow(i);
+
+  return vertex;
+}
+
+/// This is an implementation of the Clauss-Loechner algorithm for chamber
+/// decomposition.
+/// We maintain a list of pairwise disjoint chambers and their vertex-sets;
+/// we iterate over the vertex list, each time appending the vertex to the
+/// chambers where it is active and separating the chambers where it is not.
+///
+/// Given the region each vertex is active in, for each subset of vertices,
+/// the region that 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, std::vector<unsigned>>>
+mlir::presburger::detail::computeChamberDecomposition(
+    ArrayRef<PresburgerSet> activeRegions, ArrayRef<ParamPoint> vertices) {
+  // We maintain a list of regions and their associated vertex sets,
+  // initialized with the universe and the empty list of vertices.
+  std::vector<std::pair<PresburgerSet, std::vector<unsigned>>> chambers = {
+      {PresburgerSet::getUniverse(activeRegions[0].getSpace()),
+       std::vector<unsigned>({})}};
+  // 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]}.
+
+  // 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.
+
+  // At each step, we define a new chamber list after considering vertex v_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 (modulo lower-dimensional spaces).
+  for (unsigned j = 0, e = vertices.size(); j < e; j++) {
+    std::vector<std::pair<PresburgerSet, std::vector<unsigned>>> newChambers;
+
+    for (auto [currentRegion, currentVertices] : chambers) {
----------------
Superty wrote:

this is copying. use cost auto &[...] to avoid the copy and take const refs

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