[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.

578–588

Could you write a comment when this case is taken?

[webmiche]

addressed comments

[Groverkss]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

440–441

578–579

[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?

511–516

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.

518–521

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.

518–521

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

388

406

409

Can this copy be removed?

439

454–470

You can use poly here

518–521

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

[webmiche]

addressed remaining comments.

[arjunp]

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

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

383

[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;
 }
 
-/// Types the inequalities of `p` according to their `IneqType` for `simp` into
+/// Given an IntegerPolyhedron `p` and one of its inequalities `ineq`, check
+/// that all inequalities of `cuttingIneqs` are redundant for the facet of `p`
+/// where `ineq` holds as an equality. `simp` must be the Simplex constructed
+/// from `p`.
+static bool isFacetContained(ArrayRef<int64_t> ineq, Simplex &simp,
+                             IntegerPolyhedron &p,
+                             ArrayRef<ArrayRef<int64_t>> cuttingIneqs) {
+  unsigned snapshot = simp.getSnapshot();
+  simp.addEquality(ineq);
+  if (llvm::any_of(cuttingIneqs, [&simp](ArrayRef<int64_t> curr) {
+        return !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.
+static void
+addCoalescedPolyhedron(SmallVectorImpl<IntegerPolyhedron> &polyhedrons,
+                       unsigned i, unsigned j, const IntegerPolyhedron &poly,
+                       SmallVectorImpl<Simplex> &simplices) {
+  assert(i != j && "The indices must refer to different polyhedra");
+
+  unsigned n = polyhedrons.size();
+  polyhedrons[i] = polyhedrons[n - 1];
+  polyhedrons[j] = polyhedrons[n - 2];
+  polyhedrons.pop_back();
+  polyhedrons[n - 2] = poly;
+
+  simplices[i] = simplices[n - 1];
+  simplices[j] = simplices[n - 2];
+  simplices.pop_back();
+  simplices[n - 2] = Simplex(poly);
+}
+
+/// 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:
+///    ___________        ___________
+///   /   /  |   /       /          /
+///   \   \  |  /   ==>  \         /
+///    \   \ | /          \       /
+///     \___\|/            \_____/
+///
+///
+static LogicalResult
+coalescePairCutCase(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::any_of(cuttingIneqsA,
+                   [&simp, &poly, &cuttingIneqsB](ArrayRef<int64_t> curr) {
+                     return !isFacetContained(curr, simp, poly, cuttingIneqsB);
+                   }))
+    return failure();
+  IntegerPolyhedron newSet(poly.getNumDimIds(), poly.getNumSymbolIds(),
+                           poly.getNumLocalIds());
+
+  for (ArrayRef<int64_t> curr : redundantIneqsA)
+    newSet.addInequality(curr);
+
+  for (ArrayRef<int64_t> curr : redundantIneqsB)
+    newSet.addInequality(curr);
+
+  addCoalescedPolyhedron(polyhedrons, i, j, newSet, simplices);
+  return success();
+}
+
+/// Types the inequality `ineq` according to its `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();
-  }
+static LogicalResult
+typeInequality(ArrayRef<int64_t> ineq, Simplex &simp,
+  Simplex::IneqType type = simp.findIneqType(ineq);
+    redundantIneqs.push_back(ineq);
+    cuttingIneqs.push_back(ineq);
+  return success();
+
+/// Types the equality `eq`, i.e. for `eq` == 0, types both `eq` >= 0 and -`eq`
+/// >= 0 according to their `IneqType` for `simp` into `redundantIneqs` and
+/// `cuttingIneqs`. Returns success, if no separate inequalities were
+/// encountered. Otherwise, returns failure.
+static LogicalResult
+typeEquality(ArrayRef<int64_t> eq, Simplex &simp,
+             SmallVectorImpl<ArrayRef<int64_t>> &redundantIneqs,
+             SmallVectorImpl<ArrayRef<int64_t>> &cuttingIneqs,
+             SmallVectorImpl<SmallVector<int64_t, 2>> &negEqs) {
+  if (typeInequality(eq, simp, redundantIneqs, cuttingIneqs).failed())
+    return failure();
+  negEqs.push_back(getNegatedCoeffs(eq));
+  ArrayRef<int64_t> inv(negEqs.back());
+  if (typeInequality(inv, simp, redundantIneqs, cuttingIneqs).failed())
+    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,
+static 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) {
+static LogicalResult
+coalescePair(unsigned i, unsigned j,
 
   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 and equalities 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.
+  for (int k = 0, e = a.getNumInequalities(); k < e; ++k)
+    if (typeInequality(a.getInequality(k), simpB, redundantIneqsA,
+                       cuttingIneqsA)
+            .failed())
+      return failure();
 
-  // 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())
+  for (int k = 0, e = a.getNumEqualities(); k < e; ++k)
+    if (typeEquality(a.getEquality(k), simpB, redundantIneqsA, cuttingIneqsA,
+                     negEqs)
+            .failed())
       return failure();
 
   SmallVector<ArrayRef<int64_t>, 2> redundantIneqsB;
   SmallVector<ArrayRef<int64_t>, 2> cuttingIneqsB;
+  for (int k = 0, e = b.getNumInequalities(); k < e; ++k)
+    if (typeInequality(b.getInequality(k), simpA, redundantIneqsB,
+                       cuttingIneqsB)
+            .failed())
+      return failure();
 
-  if (typeInequalities(b, simpA, redundantIneqsB, cuttingIneqsB).failed())
+  for (int k = 0, e = b.getNumEqualities(); k < e; ++k)
+    if (typeEquality(b.getEquality(k), 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 there are no cutting inequalities of `a`, `b` is contained
+  // within `a` (and vice versa for `b`).
+  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 (coalescePairCutCase(polyhedrons, simplices, i, j, redundantIneqsA,
+                          cuttingIneqsA, redundantIneqsB, cuttingIneqsB)
+          .succeeded())
+    return success();
+
+  if (coalescePairCutCase(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)",
          });