Skip to content

Commit 2dde029

Browse files
[MLIR][Presburger] Implement computation of generating function for unimodular cones (#77235)
We implement a function that computes the generating function corresponding to a unimodular cone. The generating function for a polytope is obtained by summing these generating functions over all tangent cones.
1 parent a1dc813 commit 2dde029

File tree

6 files changed

+139
-1
lines changed

6 files changed

+139
-1
lines changed

mlir/include/mlir/Analysis/Presburger/Barvinok.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
#ifndef MLIR_ANALYSIS_PRESBURGER_BARVINOK_H
2525
#define MLIR_ANALYSIS_PRESBURGER_BARVINOK_H
2626

27+
#include "mlir/Analysis/Presburger/GeneratingFunction.h"
2728
#include "mlir/Analysis/Presburger/IntegerRelation.h"
2829
#include "mlir/Analysis/Presburger/Matrix.h"
2930
#include <optional>
@@ -77,6 +78,11 @@ ConeV getDual(ConeH cone);
7778
/// The returned cone is pointed at the origin.
7879
ConeH getDual(ConeV cone);
7980

81+
/// Compute the generating function for a unimodular cone.
82+
/// The input cone must be unimodular; it assert-fails otherwise.
83+
GeneratingFunction unimodularConeGeneratingFunction(ParamPoint vertex, int sign,
84+
ConeH cone);
85+
8086
} // namespace detail
8187
} // namespace presburger
8288
} // namespace mlir

mlir/include/mlir/Analysis/Presburger/IntegerRelation.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,8 @@ class IntegerRelation {
221221
return getInt64Vec(inequalities.getRow(idx));
222222
}
223223

224+
inline IntMatrix getInequalities() const { return inequalities; }
225+
224226
/// Get the number of vars of the specified kind.
225227
unsigned getNumVarKind(VarKind kind) const {
226228
return space.getNumVarKind(kind);

mlir/include/mlir/Analysis/Presburger/Matrix.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,9 @@ class Matrix {
181181
/// `elems` must be equal to the number of columns.
182182
unsigned appendExtraRow(ArrayRef<T> elems);
183183

184+
// Transpose the matrix without modifying it.
185+
Matrix<T> transpose() const;
186+
184187
/// Print the matrix.
185188
void print(raw_ostream &os) const;
186189
void dump() const;

mlir/lib/Analysis/Presburger/Barvinok.cpp

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
//===----------------------------------------------------------------------===//
88

99
#include "mlir/Analysis/Presburger/Barvinok.h"
10+
#include "llvm/ADT/Sequence.h"
1011

1112
using namespace mlir;
1213
using namespace presburger;
@@ -24,7 +25,7 @@ ConeV mlir::presburger::detail::getDual(ConeH cone) {
2425
// is represented as a row [a1, ..., an, b]
2526
// and that b = 0.
2627

27-
for (unsigned i = 0; i < numIneq; ++i) {
28+
for (auto i : llvm::seq<int>(0, numIneq)) {
2829
assert(cone.atIneq(i, numVar) == 0 &&
2930
"H-representation of cone is not centred at the origin!");
3031
for (unsigned j = 0; j < numVar; ++j) {
@@ -63,3 +64,83 @@ MPInt mlir::presburger::detail::getIndex(ConeV cone) {
6364

6465
return cone.determinant();
6566
}
67+
68+
/// Compute the generating function for a unimodular cone.
69+
/// This consists of a single term of the form
70+
/// sign * x^num / prod_j (1 - x^den_j)
71+
///
72+
/// sign is either +1 or -1.
73+
/// den_j is defined as the set of generators of the cone.
74+
/// num is computed by expressing the vertex as a weighted
75+
/// sum of the generators, and then taking the floor of the
76+
/// coefficients.
77+
GeneratingFunction mlir::presburger::detail::unimodularConeGeneratingFunction(
78+
ParamPoint vertex, int sign, ConeH cone) {
79+
// Consider a cone with H-representation [0 -1].
80+
// [-1 -2]
81+
// Let the vertex be given by the matrix [ 2 2 0], with 2 params.
82+
// [-1 -1/2 1]
83+
84+
// `cone` must be unimodular.
85+
assert(getIndex(getDual(cone)) == 1 && "input cone is not unimodular!");
86+
87+
unsigned numVar = cone.getNumVars();
88+
unsigned numIneq = cone.getNumInequalities();
89+
90+
// Thus its ray matrix, U, is the inverse of the
91+
// transpose of its inequality matrix, `cone`.
92+
// The last column of the inequality matrix is null,
93+
// so we remove it to obtain a square matrix.
94+
FracMatrix transp = FracMatrix(cone.getInequalities()).transpose();
95+
transp.removeRow(numVar);
96+
97+
FracMatrix generators(numVar, numIneq);
98+
transp.determinant(/*inverse=*/&generators); // This is the U-matrix.
99+
// Thus the generators are given by U = [2 -1].
100+
// [-1 0]
101+
102+
// The powers in the denominator of the generating
103+
// function are given by the generators of the cone,
104+
// i.e., the rows of the matrix U.
105+
std::vector<Point> denominator(numIneq);
106+
ArrayRef<Fraction> row;
107+
for (auto i : llvm::seq<int>(0, numVar)) {
108+
row = generators.getRow(i);
109+
denominator[i] = Point(row);
110+
}
111+
112+
// The vertex is v \in Z^{d x (n+1)}
113+
// We need to find affine functions of parameters λ_i(p)
114+
// such that v = Σ λ_i(p)*u_i,
115+
// where u_i are the rows of U (generators)
116+
// The λ_i are given by the columns of Λ = v^T U^{-1}, and
117+
// we have transp = U^{-1}.
118+
// Then the exponent in the numerator will be
119+
// Σ -floor(-λ_i(p))*u_i.
120+
// Thus we store the (exponent of the) numerator as the affine function -Λ,
121+
// since the generators u_i are already stored as the exponent of the
122+
// denominator. Note that the outer -1 will have to be accounted for, as it is
123+
// not stored. See end for an example.
124+
125+
unsigned numColumns = vertex.getNumColumns();
126+
unsigned numRows = vertex.getNumRows();
127+
ParamPoint numerator(numColumns, numRows);
128+
SmallVector<Fraction> ithCol(numRows);
129+
for (auto i : llvm::seq<int>(0, numColumns)) {
130+
for (auto j : llvm::seq<int>(0, numRows))
131+
ithCol[j] = vertex(j, i);
132+
numerator.setRow(i, transp.preMultiplyWithRow(ithCol));
133+
numerator.negateRow(i);
134+
}
135+
// Therefore Λ will be given by [ 1 0 ] and the negation of this will be
136+
// [ 1/2 -1 ]
137+
// [ -1 -2 ]
138+
// stored as the numerator.
139+
// Algebraically, the numerator exponent is
140+
// [ -2 ⌊ - N - M/2 + 1 ⌋ + 1 ⌊ 0 + M + 2 ⌋ ] -> first COLUMN of U is [2, -1]
141+
// [ 1 ⌊ - N - M/2 + 1 ⌋ + 0 ⌊ 0 + M + 2 ⌋ ] -> second COLUMN of U is [-1, 0]
142+
143+
return GeneratingFunction(numColumns - 1, SmallVector<int>(1, sign),
144+
std::vector({numerator}),
145+
std::vector({denominator}));
146+
}

mlir/lib/Analysis/Presburger/Matrix.cpp

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,16 @@ unsigned Matrix<T>::appendExtraRow(ArrayRef<T> elems) {
6262
return row;
6363
}
6464

65+
template <typename T>
66+
Matrix<T> Matrix<T>::transpose() const {
67+
Matrix<T> transp(nColumns, nRows);
68+
for (unsigned row = 0; row < nRows; ++row)
69+
for (unsigned col = 0; col < nColumns; ++col)
70+
transp(col, row) = at(row, col);
71+
72+
return transp;
73+
}
74+
6575
template <typename T>
6676
void Matrix<T>::resizeHorizontally(unsigned newNColumns) {
6777
if (newNColumns < nColumns)

mlir/unittests/Analysis/Presburger/BarvinokTest.cpp

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,3 +46,39 @@ TEST(BarvinokTest, getIndex) {
4646
4, 4, {{4, 2, 5, 1}, {4, 1, 3, 6}, {8, 2, 5, 6}, {5, 2, 5, 7}});
4747
EXPECT_EQ(getIndex(cone), cone.determinant());
4848
}
49+
50+
// The following cones and vertices are randomly generated
51+
// (s.t. the cones are unimodular) and the generating functions
52+
// are computed. We check that the results contain the correct
53+
// matrices.
54+
TEST(BarvinokTest, unimodularConeGeneratingFunction) {
55+
ConeH cone = defineHRep(2);
56+
cone.addInequality({0, -1, 0});
57+
cone.addInequality({-1, -2, 0});
58+
59+
ParamPoint vertex =
60+
makeFracMatrix(2, 3, {{2, 2, 0}, {-1, -Fraction(1, 2), 1}});
61+
62+
GeneratingFunction gf = unimodularConeGeneratingFunction(vertex, 1, cone);
63+
64+
EXPECT_EQ_REPR_GENERATINGFUNCTION(
65+
gf, GeneratingFunction(
66+
2, {1},
67+
{makeFracMatrix(3, 2, {{-1, 0}, {-Fraction(1, 2), 1}, {1, 2}})},
68+
{{{2, -1}, {-1, 0}}}));
69+
70+
cone = defineHRep(3);
71+
cone.addInequality({7, 1, 6, 0});
72+
cone.addInequality({9, 1, 7, 0});
73+
cone.addInequality({8, -1, 1, 0});
74+
75+
vertex = makeFracMatrix(3, 2, {{5, 2}, {6, 2}, {7, 1}});
76+
77+
gf = unimodularConeGeneratingFunction(vertex, 1, cone);
78+
79+
EXPECT_EQ_REPR_GENERATINGFUNCTION(
80+
gf,
81+
GeneratingFunction(
82+
1, {1}, {makeFracMatrix(2, 3, {{-83, -100, -41}, {-22, -27, -15}})},
83+
{{{8, 47, -17}, {-7, -41, 15}, {1, 5, -2}}}));
84+
}

0 commit comments

Comments
 (0)