[Mlir-commits] [mlir] [MLIR][Presburger] Implement computation of generating function for unimodular cones (PR #77235)

[Arjun P]

================
@@ -63,3 +64,82 @@ MPInt mlir::presburger::detail::getIndex(ConeV cone) {
 
   return cone.determinant();
 }
+
+/// Compute the generating function for a unimodular cone.
+/// This consists of a single term of the form
+/// sign * x^num / prod_j (1 - x^den_j)
+///
+/// sign is either +1 or -1.
+/// den_j is defined as the set of generators of the 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(
+    ParamPoint vertex, int sign, ConeH cone) {
+  // `cone` must be unimodular.
+  assert(getIndex(getDual(cone)) == 1 && "input cone is not unimodular!");
+
+  unsigned numVar = cone.getNumVars();
+  unsigned numIneq = cone.getNumInequalities();
+
+  // Thus its ray matrix, U, is the inverse of the
+  // transpose of its inequality matrix, `cone`.
+  // The last column of the inequality matrix is null,
+  // so we remove it to obtain a square matrix.
+  FracMatrix transp = FracMatrix(cone.getInequalities()).transpose();
+  transp.removeRow(numVar);
+
+  FracMatrix generators(numVar, numIneq);
+  transp.determinant(/*inverse=*/&generators); // This is the U-matrix.
+
+  // The powers in the denominator of the generating
+  // function are given by the generators of the cone,
+  // i.e., the rows of the matrix U.
+  std::vector<Point> denominator(numIneq);
+  ArrayRef<Fraction> row;
+  for (auto i : llvm::seq<int>(0, numVar)) {
+    row = generators.getRow(i);
+    denominator[i] = Point(row);
+  }
+
+  // The vertex is v \in Z^{d x (n+1)}
+  // We need to find affine functions of parameters λ_i(p)
+  // such that v = Σ λ_i(p)*u_i,
+  // where u_i are the rows of U (generators)
+  // The λ_i are given by the columns of Λ = v^T U^{-1}, and
+  // we have transp = U^{-1}.
+  // Then the exponent in the numerator will be
+  // Σ -floor(-λ_i(p))*u_i.
+  // Thus we store the (exponent of the) numerator as the affine function -Λ,
+  // since the generators u_i are already stored as the exponent of the
+  // denominator. Note that the outer -1 will have to be accounted for, as it is
+  // not stored. See end for an example.
+
+  unsigned numColumns = vertex.getNumColumns();
+  unsigned numRows = vertex.getNumRows();
+  ParamPoint numerator(numColumns, numRows);
+  SmallVector<Fraction> ithCol(numRows);
+  for (auto i : llvm::seq<int>(0, numColumns)) {
+    for (auto j : llvm::seq<int>(0, numRows))
+      ithCol[j] = vertex(j, i);
+    numerator.setRow(i, transp.preMultiplyWithRow(ithCol));
+    numerator.negateRow(i);
+  }
+
+  return GeneratingFunction(numColumns - 1, SmallVector<int>(1, sign),
+                            std::vector({numerator}),
+                            std::vector({denominator}));
+
+  // Suppose the vertex is given by the matrix [ 2  2   0], with 2 params
+  //                                           [-1 -1/2 1]
+  // and the cone has H-representation [0  -1] => U-matrix [ 2 -1]
+  //                                   [-1 -2]             [-1  0]
+  // Therefore Λ will be given by [ 1    0 ] and the negation of this will be
+  // stored as the numerator.
+  //                              [ 1/2 -1 ]
+  //                              [ -1  -2 ]
+
+  // Algebraically, the numerator exponent is
+  // [ -2 ⌊ -N - M/2 +1 ⌋ + 1 ⌊ 0 +M +2 ⌋ ] -> first  COLUMN of U is [2, -1]
+  // [  1 ⌊ -N - M/2 +1 ⌋ + 0 ⌊ 0 +M +2 ⌋ ] -> second COLUMN of U is [-1, 0]
----------------
Superty wrote:

I cannot understand this representation on the left side if they're supposed to be expressions then don't write it as `0 +M +2`, it looks like a matrix?

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