[Mlir-commits] [mlir] 06eb057 - [MLIR][Presburger] fix vector update in coalesce

[llvmlistbot at llvm.org]

Author: Michel Weber
Date: 2022-03-10T15:42:11+05:30
New Revision: 06eb0577380c4ba027eae5f48553b2063b3feaff

URL: https://github.com/llvm/llvm-project/commit/06eb0577380c4ba027eae5f48553b2063b3feaff
DIFF: https://github.com/llvm/llvm-project/commit/06eb0577380c4ba027eae5f48553b2063b3feaff.diff

LOG: [MLIR][Presburger] fix vector update in coalesce

When `addCoalescedPolyhedron` was called with `j == n - 1`,
the `polyhedrons`-vector was not properly updated (the
`IntegerPolyhedron` at position `n - 2` was "lost"). This patch adds
special handling to that case and a regression testcase.

Reviewed By: Groverkss

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

Added: 
    

Modified: 
    mlir/lib/Analysis/Presburger/PresburgerSet.cpp
    mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Analysis/Presburger/PresburgerSet.cpp b/mlir/lib/Analysis/Presburger/PresburgerSet.cpp
index c53b44f47fd21..07b2575d6f5a5 100644
--- a/mlir/lib/Analysis/Presburger/PresburgerSet.cpp
+++ b/mlir/lib/Analysis/Presburger/PresburgerSet.cpp
@@ -409,15 +409,33 @@ addCoalescedPolyhedron(SmallVectorImpl<IntegerPolyhedron> &polyhedrons,
   assert(i != j && "The indices must refer to 
diff erent 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);
+  if (j == n - 1) {
+    // This case needs special handling since position `n` - 1 is removed from
+    // the vector, hence the `IntegerPolyhedron` at position `n` - 2 is lost
+    // otherwise.
+    polyhedrons[i] = polyhedrons[n - 2];
+    polyhedrons.pop_back();
+    polyhedrons[n - 2] = poly;
+
+    simplices[i] = simplices[n - 2];
+    simplices.pop_back();
+    simplices[n - 2] = Simplex(poly);
+
+  } else {
+    // Other possible edge cases are correct since for `j` or `i` == `n` - 2,
+    // the `IntegerPolyhedron` at position `n` - 2 should be lost. The case
+    // `i` == `n` - 1 makes the first following statement a noop. Hence, in this
+    // case the same thing is done as above, but with `j` rather than `i`.
+    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`

diff  --git a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp b/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp
index 7574bf3ec2a72..ee45692e17d1d 100644
--- a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp
@@ -645,6 +645,17 @@ TEST(SetTest, coalesceDoubleIncrement) {
   expectCoalesce(3, set);
 }
 
+TEST(SetTest, coalesceLastCoalesced) {
+  PresburgerSet set = parsePresburgerSetFromPolyStrings(
+      1, {
+             "(x) : (x == 0)",
+             "(x) : (x - 1 >= 0, -x + 3 >= 0)",
+             "(x) : (x + 2 == 0)",
+             "(x) : (x - 2 >= 0, -x + 4 >= 0)",
+         });
+  expectCoalesce(3, set);
+}
+
 TEST(SetTest, coalesceDiv) {
   PresburgerSet set =
       parsePresburgerSetFromPolyStrings(1, {