From 8a67c6ee7c15e270587c6aae7eb7029cd83f038e Mon Sep 17 00:00:00 2001
From: Arjun P <arjunpitchanathan@gmail.com>
Date: Fri, 18 Mar 2022 12:38:54 +0000
Subject: [PATCH] [MLIR][Presburger] simplify removeConstraintsInvolvingRange

---
 mlir/lib/Analysis/Presburger/IntegerRelation.cpp | 47 ++++++++++--------------
 1 file changed, 19 insertions(+), 28 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index c4e377c..1917d9a 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -102,6 +102,23 @@ IntegerRelation::findIntegerLexMin() const {
   return maybeLexMin;
 }
 
+static bool rangeIsZero(ArrayRef<int64_t> range) {
+  return llvm::all_of(range, [](int64_t x) { return x == 0; });
+}
+
+void removeConstraintsInvolvingIdRange(IntegerRelation &poly, unsigned begin,
+                                       unsigned count) {
+  // We loop until i > 0 and index into i - 1 to avoid sign issues.
+  //
+  // We iterate backwards so that whether we remove constraint i - 1 or not, the
+  // next constraint to be tested is always i - 2.
+  for (unsigned i = poly.getNumEqualities(); i > 0; i--)
+    if (!rangeIsZero(poly.getEquality(i - 1).slice(begin, count)))
+      poly.removeEquality(i - 1);
+  for (unsigned i = poly.getNumInequalities(); i > 0; i--)
+    if (!rangeIsZero(poly.getInequality(i - 1).slice(begin, count)))
+      poly.removeInequality(i - 1);
+}
 unsigned IntegerRelation::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
@@ -561,33 +578,6 @@ Matrix IntegerRelation::getBoundedDirections() const {
   return dirs;
 }
 
-bool eqInvolvesSuffixDims(const IntegerRelation &rel, unsigned eqIndex,
-                          unsigned numDims) {
-  for (unsigned e = rel.getNumIds(), j = e - numDims; j < e; ++j)
-    if (rel.atEq(eqIndex, j) != 0)
-      return true;
-  return false;
-}
-bool ineqInvolvesSuffixDims(const IntegerRelation &rel, unsigned ineqIndex,
-                            unsigned numDims) {
-  for (unsigned e = rel.getNumIds(), j = e - numDims; j < e; ++j)
-    if (rel.atIneq(ineqIndex, j) != 0)
-      return true;
-  return false;
-}
-
-void removeConstraintsInvolvingSuffixDims(IntegerRelation &rel,
-                                          unsigned unboundedDims) {
-  // We iterate backwards so that whether we remove constraint i - 1 or not, the
-  // next constraint to be tested is always i - 2.
-  for (unsigned i = rel.getNumEqualities(); i > 0; i--)
-    if (eqInvolvesSuffixDims(rel, i - 1, unboundedDims))
-      rel.removeEquality(i - 1);
-  for (unsigned i = rel.getNumInequalities(); i > 0; i--)
-    if (ineqInvolvesSuffixDims(rel, i - 1, unboundedDims))
-      rel.removeInequality(i - 1);
-}
-
 bool IntegerRelation::isIntegerEmpty() const {
   return !findIntegerSample().hasValue();
 }
@@ -671,7 +661,8 @@ Optional<SmallVector<int64_t, 8>> IntegerRelation::findIntegerSample() const {
   IntegerRelation boundedSet(transformedSet);
   unsigned numBoundedDims = result.first;
   unsigned numUnboundedDims = getNumIds() - numBoundedDims;
-  removeConstraintsInvolvingSuffixDims(boundedSet, numUnboundedDims);
+  removeConstraintsInvolvingIdRange(boundedSet, numBoundedDims,
+                                    numUnboundedDims);
   boundedSet.removeIdRange(numBoundedDims, boundedSet.getNumIds());
 
   // 3) Try to obtain a sample from the bounded set.
-- 
2.7.4