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

[Arjun P]

================
@@ -233,4 +234,86 @@ TEST(BarvinokTest, computeNumTerms) {
     for (unsigned j = 0; j < 2; j++)
       for (unsigned k = 0; k < 2; k++)
         EXPECT_EQ(count[i][j][k], 1);
-}
\ No newline at end of file
+}
+
+/// We define some simple polyhedra with unimodular tangent cones and verify
+/// that the returned generating functions correspond to those calculated by
+/// hand.
+TEST(BarvinokTest, computeNumTermsPolytope) {
+  // A cube of side 1.
+  IntMatrix ineqs = makeIntMatrix(6, 4,
+                                  {{1, 0, 0, 0},
+                                   {0, 1, 0, 0},
+                                   {0, 0, 1, 0},
+                                   {-1, 0, 0, 1},
+                                   {0, -1, 0, 1},
+                                   {0, 0, -1, 1}});
+  PolyhedronH poly = defineHRep(3);
+  for (unsigned i = 0; i < 6; i++)
+    poly.addInequality(ineqs.getRow(i));
+
+  std::vector<std::pair<PresburgerSet, GeneratingFunction>> count =
+      computePolytopeGeneratingFunction(poly);
+  // There is only one chamber, as it is non-parametric.
+  EXPECT_EQ(count.size(), 1u);
+
+  GeneratingFunction gf = count[0].second;
+  EXPECT_EQ_REPR_GENERATINGFUNCTION(
+      gf,
+      GeneratingFunction(
+          0, {1, 1, 1, 1, 1, 1, 1, 1},
+          {makeFracMatrix(1, 3, {{1, 1, 1}}), makeFracMatrix(1, 3, {{0, 1, 1}}),
+           makeFracMatrix(1, 3, {{0, 1, 1}}), makeFracMatrix(1, 3, {{0, 0, 1}}),
+           makeFracMatrix(1, 3, {{0, 1, 1}}), makeFracMatrix(1, 3, {{0, 0, 1}}),
+           makeFracMatrix(1, 3, {{0, 0, 1}}),
+           makeFracMatrix(1, 3, {{0, 0, 0}})},
+          {{{-1, 0, 0}, {0, -1, 0}, {0, 0, -1}},
+           {{1, 0, 0}, {0, -1, 0}, {0, 0, -1}},
+           {{0, 1, 0}, {-1, 0, 0}, {0, 0, -1}},
+           {{1, 0, 0}, {0, 1, 0}, {0, 0, -1}},
+           {{0, 0, 1}, {-1, 0, 0}, {0, -1, 0}},
+           {{1, 0, 0}, {0, 0, 1}, {0, -1, 0}},
+           {{0, 1, 0}, {0, 0, 1}, {-1, 0, 0}},
+           {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}}}));
+
+  // A right-angled triangle with side p.
+  ineqs = makeIntMatrix(3, 4, {{1, 0, 0, 0}, {0, 1, 0, 0}, {-1, -1, 1, 0}});
+  poly = defineHRep(2, 1);
+  for (unsigned i = 0; i < 3; i++)
+    poly.addInequality(ineqs.getRow(i));
+
+  count = computePolytopeGeneratingFunction(poly);
+  // There is only one chamber: p ≥ 0
+  EXPECT_EQ(count.size(), 2u);
+
+  gf = count[0].second;
+  EXPECT_EQ_REPR_GENERATINGFUNCTION(
+      gf, GeneratingFunction(
+              1, {1, 1, 1},
+              {makeFracMatrix(2, 2, {{0, 1}, {0, 0}}),
+               makeFracMatrix(2, 2, {{0, 1}, {0, 0}}),
+               makeFracMatrix(2, 2, {{0, 0}, {0, 0}})},
+              {{{-1, 1}, {-1, 0}}, {{1, -1}, {0, -1}}, {{1, 0}, {0, 1}}}));
+
+  // Cartesian product of a cube with side M and a right triangle with side N.
+  ineqs = makeIntMatrix(9, 8,
+                        {{1, 0, 0, 0, 0, 1, 0, 0},
+                         {0, 1, 0, 0, 0, 1, 0, 0},
+                         {0, 0, 1, 0, 0, 1, 0, 0},
+                         {-1, 0, 0, 0, 0, 1, 0, 0},
+                         {0, -1, 0, 0, 0, 1, 0, 0},
+                         {0, 0, -1, 0, 0, 1, 0, 0},
+                         {0, 0, 0, 1, 0, 0, 0, 0},
+                         {0, 0, 0, 0, 1, 0, 0, 0},
+                         {0, 0, 0, -1, -1, 0, 1, 0}});
+  poly = defineHRep(5, 2);
+  for (unsigned i = 0; i < 9; i++)
+    poly.addInequality(ineqs.getRow(i));
----------------
Superty wrote:

You should add a todo to do this for your other tests as well.

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