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

[llvmlistbot at llvm.org]

================
@@ -147,6 +148,324 @@ 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;
+
+  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,
+                         -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 creating a new chamber if necessary.
+std::vector<std::pair<PresburgerRelation, std::vector<unsigned>>>
+mlir::presburger::detail::computeChamberDecomposition(
+    std::vector<PresburgerRelation> activeRegions,
+    std::vector<ParamPoint> vertices) {
+  // 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, e = vertices.size(); j < e; j++) {
+    newChambers.clear();
+
+    PresburgerRelation newRegion = activeRegions[j];
+    ParamPoint newVertex = vertices[j];
+
+    for (unsigned i = 0, f = chambers.size(); i < f; i++) {
+      auto [currentRegion, currentVertices] = 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 = currentRegion.intersect(newRegion);
+      bool isFullDim = intersection.getNumRangeVars() == 0 ||
+                       llvm::any_of(intersection.getAllDisjuncts(),
+                                    [&](IntegerRelation disjunct) -> bool {
+                                      return disjunct.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 = currentRegion.subtract(newRegion);
+        newChambers.push_back(std::make_pair(subtraction, currentVertices));
+
+        currentVertices.push_back(j);
+        newChambers.push_back(std::make_pair(intersection, currentVertices));
+      }
+    }
+
+    // Finally we compute the chamber where only v_j is active by subtracting
+    // all existing chambers from R_j.
+    for (const std::pair<PresburgerRelation, std::vector<unsigned>> &chamber :
+         newChambers)
+      newRegion = newRegion.subtract(chamber.first);
+    newChambers.push_back(std::make_pair(newRegion, std::vector({j})));
+
+    // We filter `chambers` to remove empty regions.
+    chambers.clear();
+    for (const std::pair<PresburgerRelation, std::vector<unsigned>> &chamber :
+         newChambers) {
+      auto [r, v] = chamber;
+      bool isEmpty = llvm::all_of(
+          r.getAllDisjuncts(),
+          [&](IntegerRelation disjunct) -> bool { return disjunct.isEmpty(); });
+      if (!isEmpty)
+        chambers.push_back(chamber);
+    }
----------------
Abhinav271828 wrote:

oh there's a mistake, this should be outside the loop
but it isn't handled by the full dim thing, a lot of the chambers we add are defined by set difference so they may be empty

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