[MLIR][Presburger] support a heuristic for the "cut case" in coalesce

[webmiche]

Summary

This patch introduces the cut case. If one polytope has only cutting and
redundant inequalities for the other and the facet of the cutting
inequalities are contained within the other polytope, then the polytopes are
be combined into a polytope consisting only of their respective
redundant constraints.

[Groverkss]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

382–384

397–399

I think if i != j, it is already implied that the size of the arrays is at least 2.

400–402

Please add an assert to check that i != j.

417–422

445–446

You should copy number of local ids too.

448–452

You can use range based for loop here.

591–601

Could you write a comment when this case is taken?

[webmiche]

addressed comments

[Groverkss]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

440–441

591–592

[arjunp]

Please make all the free functions static.

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

385

I think this should return a bool, it's not really a success/failure result.

387

Pass the cached simp instead of constructin from scratch? You can snapshot, add the equality, check redundancy, then rollback and return

389–393

!isRedundantInequality would be shorter

389–393

can use std::all_of

439–445

drop braces for single-line bodies

456–460

++k

[bondhugula]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

389–393

Please drop trivial braces. Please use llvm::all_of.

439–444

Rewrite to use llvm::any_of/all_of.

[webmiche]

Addressed comments and rewrote handling of equalities after finding a case that was impossible to handle currently. The new version types both inequalities that make an equality and considers their types as well.

[webmiche]

address a minor comment I missed previously

[webmiche]

In the earlier version, the negated inequality vector was freed before being used. This patch introduces a fix that tracks all negated equalities in a vector throughout coalescePair. In this way, they don't get freed before the new polytope is constructed.

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

391

Why this cast?

433

I feel this name can be more descriptive. What about just calling it coalescePairCutCase?

447

Please capture by reference.

448–449

Why these copies?

502–507

This section is repeated three times in this patch. IMO, you can simply rename typeInequalities to typeInequality, make it take an ArrayRef to an ineq instead of an integer polyhedron, and move the rest of the logic into cutCase directly.

509–512

You can use the static function getComplementIneq

[webmiche]

Address further comments. In particular, made free functions static, refactor typing of inequalities and equalities.

[webmiche]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

391

When not capturing by reference, the captured variables are implicitly const. That's why I needed this cast (and the copies further down). Capturing by reference solved the problem.

509–512

Not quite. For t(x), getComplementIneq returns -t(x) - 1, whereas I require -t(x). I guess I can use getComplementIneq and then change the constant, but that seems wasteful IMO.

[arjunp]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

479

497

500

Can this copy be removed?

509–512

Oops, my bad. You can use the function getNegatedCoeffs instead.

530

545–561

You can use poly here

[webmiche]

addressed remaining comments.

[arjunp]

LGTM. Please update the documentation too for the variable name change:

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

474

[webmiche]

addressed last comment

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

for (const IntegerPolyhedron &poly : integerPolyhedrons) {
     Optional<uint64_t> volume = poly.computeVolume();
     if (!volume)
       return {};
     result += *volume;
   }
   return result;
 }
 
+/// Given an IntegerPolyhedron `p` and one of its inequalities `ineq`, check
+/// that all inequalities of `cut` are redundant for the facet of `p` where
+/// `ineq` holds as an equality. `simp` must be the Simplex constructed from
+/// `p`.
+bool containedFacet(ArrayRef<int64_t> ineq, Simplex &simp, IntegerPolyhedron &p,
+                    ArrayRef<ArrayRef<int64_t>> cut) {
+  unsigned snapshot = simp.getSnapshot();
+  simp.addEquality(ineq);
+  if (!llvm::all_of(cut, [simp](ArrayRef<int64_t> curr) {
+        return ((Simplex)simp).isRedundantInequality(curr);
+      })) {
+    simp.rollback(snapshot);
+    return false;
+  }
+  simp.rollback(snapshot);
+  return true;
+}
+
+/// Adds `poly` to `polyhedrons` and removes the polyhedrons at position `i` and
+/// `j`. Updates `simplices` to reflect the changes. `i` and `j` cannot be
+/// equal.
+void addCoalescedPolyhedron(SmallVectorImpl<IntegerPolyhedron> &polyhedrons,
+                            unsigned i, unsigned j, IntegerPolyhedron &poly,
+                            SmallVectorImpl<Simplex> &simplices) {
+  assert(i != j && "The indices must refer to different polyhedra");
+  IntegerPolyhedron newSet = poly;
+
+  unsigned n = polyhedrons.size();
+  polyhedrons[i] = polyhedrons[n - 1];
+  polyhedrons[j] = polyhedrons[n - 2];
+  polyhedrons.pop_back();
+  polyhedrons[n - 2] = newSet;
+
+  simplices[i] = simplices[n - 1];
+  simplices[j] = simplices[n - 2];
+  simplices.pop_back();
+  simplices[n - 2] = Simplex(newSet);
+}
+
+/// Given two polyhedra `a` and `b` at positions `i` and `j` in `polyhedrons`
+/// and `redundantIneqsA` being the inequalities of `a` that are redundant for
+/// `b` (similarly for `cuttingIneqsA`, `redundantIneqsB`, and `cuttingIneqsB`),
+/// checks whether the facets of all cutting inequalites of `a` are contained in
+/// `b`. If so, a new polyhedron consisting of all redundant inequalites of `a`
+/// and `b` and all equalities of both is created.
+///
+/// An example of this case:
+///    ___________        ___________
+///   /   /  |   /       /          /
+///   \   \  |  /   ==>  \         /
+///    \   \ | /          \       /
+///     \___\|/            \_____/
+///
+///
+LogicalResult cutCase(SmallVectorImpl<IntegerPolyhedron> &polyhedrons,
+                      SmallVectorImpl<Simplex> &simplices, unsigned i,
+                      unsigned j, ArrayRef<ArrayRef<int64_t>> redundantIneqsA,
+                      ArrayRef<ArrayRef<int64_t>> cuttingIneqsA,
+                      ArrayRef<ArrayRef<int64_t>> redundantIneqsB,
+                      ArrayRef<ArrayRef<int64_t>> cuttingIneqsB) {
+  /// All inequalities of `b` need to be redundant. We already know that the
+  /// redundant ones are, so only the cutting ones remain to be checked.
+  Simplex &simp = simplices[i];
+  IntegerPolyhedron &poly = polyhedrons[i];
+  if (!llvm::all_of(cuttingIneqsA,
+                    [simp, poly, cuttingIneqsB](ArrayRef<int64_t> curr) {
+                      Simplex s = (Simplex)simp;
+                      IntegerPolyhedron p = (IntegerPolyhedron)poly;
+                      return containedFacet(curr, s, p, cuttingIneqsB);
+                    }))
+    return failure();
+  IntegerPolyhedron newSet(polyhedrons[i].getNumDimIds(),
+                           polyhedrons[i].getNumSymbolIds(),
+                           polyhedrons[i].getNumLocalIds());
+
+  for (ArrayRef<int64_t> curr : redundantIneqsA)
+    newSet.addInequality(curr);
+
+  for (ArrayRef<int64_t> curr : redundantIneqsB)
+    newSet.addInequality(curr);
+
+  for (unsigned k = 0, e = polyhedrons[i].getNumEqualities(); k < e; ++k)
+    newSet.addEquality(polyhedrons[i].getEquality(k));
+
+  for (unsigned k = 0, e = polyhedrons[j].getNumEqualities(); k < e; ++k)
+    newSet.addEquality(polyhedrons[j].getEquality(k));
+
+  addCoalescedPolyhedron(polyhedrons, i, j, newSet, simplices);
+  return success();
+}
+
 /// Types the inequalities of `p` according to their `IneqType` for `simp` into
 /// `redundantIneqs` and `cuttingIneqs`. Returns success, if no separate
 /// inequalities were encountered. Otherwise, returns failure.
 LogicalResult
 typeInequalities(const IntegerPolyhedron &p, Simplex &simp,
                  SmallVectorImpl<ArrayRef<int64_t>> &redundantIneqs,
                  SmallVectorImpl<ArrayRef<int64_t>> &cuttingIneqs) {
   for (unsigned i = 0, e = p.getNumInequalities(); i < e; ++i) {
     Simplex::IneqType type = simp.findIneqType(p.getInequality(i));
     if (type == Simplex::IneqType::Redundant)
       redundantIneqs.push_back(p.getInequality(i));
     else if (type == Simplex::IneqType::Cut)
       cuttingIneqs.push_back(p.getInequality(i));
     else
       return failure();
   }
   return success();
 }
 
+/// Types the equalities of `p`, i.e. for an equality `coeffs` == 0, assigns
+/// both `coeffs` >= 0 and -`coeffs` >= 0 according to their `IneqType` for
+/// `simp` into `redundantIneqs` and `cuttingIneqs`. Returns success, if no
+/// separate inequalities were encountered. Otherwise, returns failure.
+LogicalResult typeEqualities(const IntegerPolyhedron &p, Simplex &simp,
+                             SmallVectorImpl<ArrayRef<int64_t>> &redundantIneqs,
+                             SmallVectorImpl<ArrayRef<int64_t>> &cuttingIneqs,
+                             SmallVector<SmallVector<int64_t, 2>, 2> &negEqs) {
+  for (unsigned i = 0, e = p.getNumEqualities(); i < e; ++i) {
+    Simplex::IneqType type = simp.findIneqType(p.getEquality(i));
+    if (type == Simplex::IneqType::Redundant)
+      redundantIneqs.push_back(p.getEquality(i));
+    else if (type == Simplex::IneqType::Cut)
+      cuttingIneqs.push_back(p.getEquality(i));
+    else
+      return failure();
+
+    SmallVector<int64_t, 2> neg;
+    for (int64_t curr : p.getEquality(i))
+      neg.push_back(-curr);
+    negEqs.push_back(neg);
+    ArrayRef<int64_t> inv(negEqs.back());
+    type = simp.findIneqType(inv);
+    if (type == Simplex::IneqType::Redundant)
+      redundantIneqs.push_back(inv);
+    else if (type == Simplex::IneqType::Cut)
+      cuttingIneqs.push_back(inv);
+    else
+      return failure();
+  }
+
+  return success();
+}
+
 /// Replaces the element at position `i` with the last element and erases the
 /// last element for both `polyhedrons` and `simplices`.
 void erasePolyhedron(unsigned i,
                      SmallVectorImpl<IntegerPolyhedron> &polyhedrons,
                      SmallVectorImpl<Simplex> &simplices) {
   assert(simplices.size() == polyhedrons.size() &&
          "simplices and polyhedrons must be equally as long");
   polyhedrons[i] = polyhedrons.back();
   polyhedrons.pop_back();
   simplices[i] = simplices.back();
   simplices.pop_back();
 }
 
 /// Attempts to coalesce the two IntegerPolyhedrons at position `i` and `j` in
 /// `polyhedrons` in-place. Returns whether the polyhedrons were successfully
 /// coalesced. The simplices in `simplices` need to be the ones constructed from
 /// `polyhedrons`. At this point, there are no empty polyhedrons in
 /// `polyhedrons` left.
 LogicalResult coalescePair(unsigned i, unsigned j,
                            SmallVectorImpl<IntegerPolyhedron> &polyhedrons,
                            SmallVectorImpl<Simplex> &simplices) {
 
   IntegerPolyhedron &a = polyhedrons[i];
   IntegerPolyhedron &b = polyhedrons[j];
   assert(a.getNumLocalIds() == 0 && b.getNumLocalIds() == 0 &&
          "Locals are not yet supported!");
   Simplex &simpA = simplices[i];
   Simplex &simpB = simplices[j];
 
-  // Check that all equalities are redundant in a (and in b).
-  bool onlyRedundantEqsA = true;
-  for (unsigned k = 0, e = a.getNumEqualities(); k < e; ++k)
-    if (!simpB.isRedundantEquality(a.getEquality(k))) {
-      onlyRedundantEqsA = false;
-      break;
-    }
-
-  bool onlyRedundantEqsB = true;
-  for (unsigned k = 0, e = b.getNumEqualities(); k < e; ++k)
-    if (!simpA.isRedundantEquality(b.getEquality(k))) {
-      onlyRedundantEqsB = false;
-      break;
-    }
-
-  // If there are non-redundant equalities for both, exit early.
-  if (!onlyRedundantEqsB && !onlyRedundantEqsA)
-    return failure();
-
   SmallVector<ArrayRef<int64_t>, 2> redundantIneqsA;
   SmallVector<ArrayRef<int64_t>, 2> cuttingIneqsA;
+  SmallVector<SmallVector<int64_t, 2>, 2> negEqs;
 
   // Organize all inequalities of `a` according to their type for `b` into
   // `redundantIneqsA` and `cuttingIneqsA` (and vice versa for all inequalities
   // of `b` according to their type in `a`). If a separate inequality is
   // encountered during typing, the two IntegerPolyhedrons cannot be coalesced.
   if (typeInequalities(a, simpB, redundantIneqsA, cuttingIneqsA).failed())
     return failure();
 
+  if (typeEqualities(a, simpB, redundantIneqsA, cuttingIneqsA, negEqs).failed())
+    return failure();
+
   SmallVector<ArrayRef<int64_t>, 2> redundantIneqsB;
   SmallVector<ArrayRef<int64_t>, 2> cuttingIneqsB;
 
   if (typeInequalities(b, simpA, redundantIneqsB, cuttingIneqsB).failed())
     return failure();
 
+  if (typeEqualities(b, simpA, redundantIneqsB, cuttingIneqsB, negEqs).failed())
+    return failure();
+
   // If there are no cutting inequalities of `a` and all equalities of `a` are
   // redundant, then all constraints of `a` are redundant making `b` contained
   // within a (and vice versa for `b`).
-  if (cuttingIneqsA.empty() && onlyRedundantEqsA) {
+  if (cuttingIneqsA.empty()) {
     erasePolyhedron(j, polyhedrons, simplices);
     return success();
   }
 
-  if (cuttingIneqsB.empty() && onlyRedundantEqsB) {
+  if (cuttingIneqsB.empty()) {
     erasePolyhedron(i, polyhedrons, simplices);
     return success();
   }
 
+  // Try to apply the cut case if all equalities of `a` are redundant (and `b`
+  // respectively).
+  if (cutCase(polyhedrons, simplices, i, j, redundantIneqsA, cuttingIneqsA,
+              redundantIneqsB, cuttingIneqsB)
+          .succeeded())
+    return success();
+
+  if (cutCase(polyhedrons, simplices, j, i, redundantIneqsB, cuttingIneqsB,
+              redundantIneqsA, cuttingIneqsA)
+          .succeeded())
+    return success();
+
   return failure();
 }
 
 PresburgerSet PresburgerSet::coalesce() const {
   PresburgerSet newSet =
       PresburgerSet::getEmptySet(getNumDimIds(), getNumSymbolIds());
   SmallVector<IntegerPolyhedron, 2> polyhedrons = integerPolyhedrons;
   SmallVector<Simplex, 2> simplices;

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

 }
 
 TEST(SetTest, coalesceCutOneDim) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       1, {
              "(x) : ( x >= 0, -x + 3 >= 0)",
              "(x) : ( x - 2 >= 0, -x + 4 >= 0)",
          });
-  expectCoalesce(2, set);
+  expectCoalesce(1, set);
 }
 
 TEST(SetTest, coalesceSeparateOneDim) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       1, {"(x) : ( x >= 0, -x + 2 >= 0)", "(x) : ( x - 3 >= 0, -x + 4 >= 0)"});
   expectCoalesce(2, set);
 }
 
+TEST(SetTest, coalesceAdjEq) {
+  PresburgerSet set = parsePresburgerSetFromPolyStrings(
+      1, {"(x) : ( x == 0)", "(x) : ( x - 1 == 0)"});
+  expectCoalesce(2, set);
+}
+
 TEST(SetTest, coalesceContainedTwoDim) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       2, {
              "(x,y) : (x >= 0, -x + 3 >= 0, y >= 0, -y + 3 >= 0)",
              "(x,y) : (x >= 0, -x + 3 >= 0, y - 2 >= 0, -y + 3 >= 0)",
          });
   expectCoalesce(1, set);
 }
 
 TEST(SetTest, coalesceCutTwoDim) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       2, {
              "(x,y) : (x >= 0, -x + 3 >= 0, y >= 0, -y + 2 >= 0)",
              "(x,y) : (x >= 0, -x + 3 >= 0, y - 1 >= 0, -y + 3 >= 0)",
          });
+  expectCoalesce(1, set);
+}
+
+TEST(SetTest, coalesceEqStickingOut) {
+  PresburgerSet set = parsePresburgerSetFromPolyStrings(
+      2, {
+             "(x,y) : (x >= 0, -x + 2 >= 0, y >= 0, -y + 2 >= 0)",
+             "(x,y) : (y - 1 == 0, x >= 0, -x + 3 >= 0)",
+         });
   expectCoalesce(2, set);
 }
 
 TEST(SetTest, coalesceSeparateTwoDim) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       2, {
              "(x,y) : (x >= 0, -x + 3 >= 0, y >= 0, -y + 1 >= 0)",
              "(x,y) : (x >= 0, -x + 3 >= 0, y - 2 >= 0, -y + 3 >= 0)",
          });
   expectCoalesce(2, set);
 }
 
 TEST(SetTest, coalesceContainedEq) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       2, {
              "(x,y) : (x >= 0, -x + 3 >= 0, x - y == 0)",
              "(x,y) : (x - 1 >= 0, -x + 2 >= 0, x - y == 0)",
          });
   expectCoalesce(1, set);
 }
 
 TEST(SetTest, coalesceCuttingEq) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       2, {
-             "(x,y) : (x - 1 >= 0, -x + 3 >= 0, x - y == 0)",
+             "(x,y) : (x + 1 >= 0, -x + 1 >= 0, x - y == 0)",
              "(x,y) : (x >= 0, -x + 2 >= 0, x - y == 0)",
          });
-  expectCoalesce(2, set);
+  expectCoalesce(1, set);
 }
 
 TEST(SetTest, coalesceSeparateEq) {
   PresburgerSet set = parsePresburgerSetFromPolyStrings(
       2, {
              "(x,y) : (x - 3 >= 0, -x + 4 >= 0, x - y == 0)",
              "(x,y) : (x >= 0, -x + 1 >= 0, x - y == 0)",
          });