[Mlir-commits] [mlir] [MLIR][Presburger] Find particular solutions to parametric equation systems (PR #196758)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat May 9 14:36:30 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Yue Huang (AdUhTkJm)
<details>
<summary>Changes</summary>
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
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);
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/196758
More information about the Mlir-commits
mailing list