[MLIR][Presburger] Find particular solutions to parametric equation systems

[AdUhTkJm]

We are solving equation system $Ax = Bp + C$, where $x$ is an integer vector we solve for and $p$ stands for parameters. This is a component of the projecting step of Barvinok's algorithm, which projects lower-dimensional polyhedra into full-dimensional ones.

We do this by finding the smith normal form $D$ of $A$, namely $D = UAV$ for two unimodular matrices U and V. Then we can rewrite the system to
$U^{-1} D V^{-1} x = Bp + C$
which is
$D V^{-1} x = UBp + UC$

Since $D$ is diagonal by definition, it is easy to solve for $y = V^{-1} x$. To obtain $x$, we only need to compute $Vy$.

The tricky part is that $x$ has to be integer-entried. This means for every row $i$, $D_i$ must divide the $i$'th element of the vector $UBp + UC$. This imposes modular constraints on $p$.

It is also possible that $D_i$ is zero, in which case the $i$'th element of $UBp + UC$ must also be zero. This gives an equality constraint on $p$.

These constraints are also computed and returned by the function.

[llvmorg-github-actions]

@llvm/pr-subscribers-mlir-presburger

@llvm/pr-subscribers-mlir

Author: Yue Huang (AdUhTkJm)

Changes

We are solving equation system $Ax = Bp + C$, where $x$ is an integer vector we solve for and $p$ stands for parameters. This is a component of the projecting step of Barvinok's algorithm, which projects lower-dimensional polyhedra into full-dimensional ones.

We do this by finding the smith normal form $D$ of $A$, namely $D = UAV$ for two unimodular matrices U and V. Then we can rewrite the system to
$U^{-1} D V^{-1} x = Bp + C$
which is
$D V^{-1} x = UBp + UC$

Since $D$ is diagonal by definition, it is easy to solve for $y = V^{-1} x$. To obtain $x$, we only need to compute $Vy$.

The tricky part is that $x$ has to be integer-entried. This means for every row $i$, $D_i$ must divide the $i$'th element of the vector $UBp + UC$. This imposes modular constraints on $p$.

It is also possible that $D_i$ is zero, in which case the $i$'th element of $UBp + UC$ must also be zero. This gives an equality constraint on $p$.

These constraints are also computed and returned by the function.

---

Full diff: https://github.com/llvm/llvm-project/pull/196758.diff

3 Files Affected:

• (modified) mlir/include/mlir/Analysis/Presburger/Barvinok.h (+16)
• (modified) mlir/lib/Analysis/Presburger/Barvinok.cpp (+93)
• (modified) mlir/unittests/Analysis/Presburger/BarvinokTest.cpp (+39)

diff --git a/mlir/include/mlir/Analysis/Presburger/Barvinok.h b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
index b09617f0fe88e..952d2e953f7e6 100644
--- a/mlir/include/mlir/Analysis/Presburger/Barvinok.h
+++ b/mlir/include/mlir/Analysis/Presburger/Barvinok.h
@@ -48,6 +48,22 @@ using PolyhedronV = IntMatrix;
 using ConeH = PolyhedronH;
 using ConeV = PolyhedronV;
 
+/// An integer particular solution of equation system Ax = Bp + C.
+struct ParticularSolution {
+  // The constraints that p must satisfy for the system to have an integer
+  // solution.
+  IntegerRelation constraint;
+  // Row i is an affine expression for the i'th variable in vector x, in p.
+  FracMatrix solution;
+};
+
+/// Find particular solution of equation system Ax = Bp + C.
+/// The first argument `eqs` corresponds to matrix A,
+/// and the second argument `constants` corresponds to B and C,
+/// where C is the final column.
+ParticularSolution findParticularSolution(const IntMatrix &eqs,
+                                          const IntMatrix &constants);
+
 inline PolyhedronH defineHRep(int numVars, int numSymbols = 0) {
   // We don't distinguish between domain and range variables, so
   // we set the number of domain variables as 0 and the number of
diff --git a/mlir/lib/Analysis/Presburger/Barvinok.cpp b/mlir/lib/Analysis/Presburger/Barvinok.cpp
index e27658a37ad12..9c2f65ce7ab48 100644
--- a/mlir/lib/Analysis/Presburger/Barvinok.cpp
+++ b/mlir/lib/Analysis/Presburger/Barvinok.cpp
@@ -15,6 +15,99 @@ using namespace mlir;
 using namespace presburger;
 using namespace mlir::presburger::detail;
 
+ParticularSolution
+mlir::presburger::detail::findParticularSolution(const IntMatrix &eqs,
+                                                 const IntMatrix &constants) {
+  auto [u, d, v] = eqs.computeSmithNormalForm();
+
+  // We are solving Ax = Bp + C. Now we've obtained UAV = D, where D is the
+  // Smith Normal Form. Hence Ax = (U^{-1} D V^{-1}) x = Bp + C, which means
+  // D(V^{-1}x) = U(Bp + C). We denote y = V^{-1}x, and we solve for y first.
+  //
+  // As D is diagonal, i.e. D = diag(s_1, ..., s_n, 0, ..., 0),
+  // it is obvious that the equation is reduced to the following equalities:
+  //   -  s_1 y_1 = (UBp + UC)_1
+  //   -  s_2 y_2 = (UBp + UC)_2
+  //   -  ...
+  //
+  // So y_n = (UBp + UC)_n / s_n for every s_n != 0. For those s_n = 0, we can
+  // simply set y_n = 0, as we are only finding *a* solution rather than all
+  // solutions. Note that it is not guaranteed that the values will always be
+  // divisible. If it is not, then the parameters are constrained into a
+  // sublattice. These constraints on p are collected in an IntegerRelation and
+  // are also returned.
+
+  unsigned numRows = eqs.getNumRows();
+  unsigned numCols = eqs.getNumColumns();
+  assert(numRows == constants.getNumRows());
+  unsigned numRhsCols = constants.getNumColumns();
+
+  // Bp + C is stored as [B | C]. Hence to calculate UBp + UC, simply do a
+  // matmul.
+  IntMatrix rhs = u.postMultiply(constants);
+
+  // The result, in a form [y_B | y_C] that represents y_B p + y_C.
+  FracMatrix y(numCols, numRhsCols);
+  SmallVector<unsigned> modIndex, eqIndex;
+  for (unsigned i = 0; i < numRows; i++) {
+    MutableArrayRef row = rhs.getRow(i);
+    // It is possible that there are more rows than columns. But by Smith Normal
+    // Form, the extra rows are all zero.
+    auto s = i >= numCols ? DynamicAPInt() : d(i, i);
+
+    // (UBp + UC)_i is essentially (UB)[i, :] \cdot p + (UC)_i.
+    if (s == 0 && std::any_of(row.begin(), row.end(),
+                              [](const DynamicAPInt &x) { return x != 0; }))
+      // This is an equality constraint, e.g. 0x + 3N - 6 == 0 implies N == 2.
+      eqIndex.push_back(i);
+
+    if (s != 0 &&
+        std::any_of(row.begin(), row.end(),
+                    [&](const DynamicAPInt &x) { return x % s != 0; }))
+      // This is a modular constraint, e.g. 2x + 3N - 5 == 0 implies (3N - 5) %
+      // 2 == 0.
+      modIndex.push_back(i);
+
+    // A constraint s_i y_i = (UB)[i, :] p + (UC)_i is found and must be
+    // enforced.
+    if (s != 0) {
+      for (unsigned j = 0; j < numRhsCols; j++)
+        y(i, j) = Fraction(rhs(i, j), s);
+    }
+  }
+
+  // For modular constraints, we use div-locals to capture them.
+  unsigned numParams = numRhsCols - 1;
+  unsigned numLocals = modIndex.size();
+  unsigned dims = numParams + numLocals;
+  auto space = PresburgerSpace::getSetSpace(numParams, 0, numLocals);
+  IntegerRelation constraints(space);
+
+  // First, we'll create the constraints of `p`.
+  for (auto i : eqIndex) {
+    // These means the RHS must be equal to zero, so we can directly add them
+    // as constraints, modulo the locals.
+    auto row = rhs.getRow(i);
+    SmallVector<DynamicAPInt> eq(dims + 1);
+    std::copy(row.begin(), row.end() - 1, eq.begin());
+    eq.back() = row.back();
+    constraints.addEquality(eq);
+  }
+
+  for (auto [j, i] : llvm::enumerate(modIndex)) {
+    // These will use the j'th local variable to express a modulus constraint.
+    auto row = rhs.getRow(i);
+    SmallVector<DynamicAPInt> eq(dims + 1);
+    std::copy(row.begin(), row.end() - 1, eq.begin());
+    // This is the divisor.
+    eq[numParams + j] = -d(i, i);
+    eq.back() = row.back();
+    constraints.addEquality(eq);
+  }
+
+  return {constraints, v.asFracMatrix().postMultiply(y)};
+}
+
 /// Assuming that the input cone is pointed at the origin,
 /// converts it to its dual in V-representation.
 /// Essentially we just remove the all-zeroes constant column.
diff --git a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
index 4ca999878df2c..12b2af7d31b21 100644
--- a/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/BarvinokTest.cpp
@@ -310,3 +310,42 @@ TEST(BarvinokTest, solveParametricEquations) {
   EXPECT_EQ(solution.at(0, 0), Fraction(1, 2));
   EXPECT_EQ(solution.at(1, 0), 1);
 }
+
+TEST(BarvinokTest, findParticularSolution) {
+  // The system:
+  //   2x + y = p + 4
+  //    x + y = p + 3
+  // Solution should be x = 1, y = p + 2,
+  // and no constraints on p should be involved.
+  {
+    IntMatrix eq = makeIntMatrix(2, 2, {{2, 1}, {1, 1}});
+    IntMatrix constant = makeIntMatrix(2, 2, {{1, 4}, {1, 3}});
+    auto [constraint, solution] = findParticularSolution(eq, constant);
+
+    auto universe = IntegerRelation::getUniverse(constraint.getSpace());
+    EXPECT_TRUE(universe.isEqual(constraint));
+
+    // x = 0p + 1
+    EXPECT_EQ(solution(0, 0), 0);
+    EXPECT_EQ(solution(0, 1), 1);
+    // y = 1p + 2
+    EXPECT_EQ(solution(1, 0), 1);
+    EXPECT_EQ(solution(1, 1), 2);
+  }
+
+  // The system:
+  //   2x = p
+  // Solution should be x = p/2, and p must be divisible by 2.
+  {
+    IntMatrix eq = makeIntMatrix(1, 1, {{2}});
+    IntMatrix constant = makeIntMatrix(1, 2, {{1, 0}});
+    auto [constraint, solution] = findParticularSolution(eq, constant);
+
+    IntegerRelation expected =
+        parseIntegerPolyhedron("(p): (p - 2 * (p floordiv 2) == 0)");
+    EXPECT_TRUE(expected.isEqual(constraint));
+
+    EXPECT_EQ(solution(0, 0), Fraction(1, 2));
+    EXPECT_EQ(solution(0, 1), 0);
+  }
+}