[MLIR][Presburger] IntegerRelation::truncate: fix bug when truncating equalities

This was truncating inequalities instead of equalities.

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D122811

diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 50ced2a..de95511 100644 (file)
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -142,7 +142,7 @@ void IntegerRelation::truncate(const CountsSnapshot &counts) {
   truncateIdKind(IdKind::Symbol, counts);
   truncateIdKind(IdKind::Local, counts);
   removeInequalityRange(counts.getNumIneqs(), getNumInequalities());
-  removeInequalityRange(counts.getNumEqs(), getNumEqualities());
+  removeEqualityRange(counts.getNumEqs(), getNumEqualities());
 }
 
 unsigned IntegerRelation::insertId(IdKind kind, unsigned pos, unsigned num) {

diff --git a/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp
index e7bbac3..4d55036 100644 (file)
--- a/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp
@@ -1202,3 +1202,13 @@ TEST(IntegerPolyhedronTest, containsPointNoLocal) {
   EXPECT_TRUE(poly3.containsPointNoLocal({0, 0}));
   EXPECT_FALSE(poly3.containsPointNoLocal({1, 0}));
 }
+
+TEST(IntegerPolyhedronTest, truncateEqualityRegressionTest) {
+  // IntegerRelation::truncate was truncating inequalities to the number of
+  // equalities.
+  IntegerRelation set(1);
+  IntegerRelation::CountsSnapshot snapshot = set.getCounts();
+  set.addEquality({1, 0});
+  set.truncate(snapshot);
+  EXPECT_EQ(set.getNumEqualities(), 0u);
+}