From f074bbb04a617f366093a860e5b0976d80129ff4 Mon Sep 17 00:00:00 2001
From: Arjun P <arjunpitchanathan@gmail.com>
Date: Fri, 26 Nov 2021 16:29:55 +0530
Subject: [PATCH] [MLIR] Simplex::pivot: also update the redundant rows when
 pivoting

Previously, the pivot function would only update the non-redundant rows when
pivoting. This is incorrect because in some cases, when rolling back past a
`detectRedundant` call, the basis being used could be different from that which
was used at the time of returning from the `detectRedundant` call. Therefore,
it is important to update the redundant rows as well during pivots. This could
also be triggered by pivots that occur when testing successive constraints for
being redundant in `detectRedundant` after some initial constraints are marked redundant.

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D114614
---
 mlir/lib/Analysis/Presburger/Simplex.cpp      |  2 +-
 .../Analysis/Presburger/SimplexTest.cpp       | 23 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Analysis/Presburger/Simplex.cpp b/mlir/lib/Analysis/Presburger/Simplex.cpp
index a56213c09c00..b04060109f2c 100644
--- a/mlir/lib/Analysis/Presburger/Simplex.cpp
+++ b/mlir/lib/Analysis/Presburger/Simplex.cpp
@@ -240,7 +240,7 @@ void Simplex::pivot(unsigned pivotRow, unsigned pivotCol) {
   }
   normalizeRow(pivotRow);
 
-  for (unsigned row = nRedundant; row < nRow; ++row) {
+  for (unsigned row = 0; row < nRow; ++row) {
     if (row == pivotRow)
       continue;
     if (tableau(row, pivotCol) == 0) // Nothing to do.
diff --git a/mlir/unittests/Analysis/Presburger/SimplexTest.cpp b/mlir/unittests/Analysis/Presburger/SimplexTest.cpp
index 2f1ce819c672..d1d725476cc8 100644
--- a/mlir/unittests/Analysis/Presburger/SimplexTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/SimplexTest.cpp
@@ -373,6 +373,29 @@ TEST(SimplexTest, isMarkedRedundantTiledLoopNestConstraints) {
   EXPECT_FALSE(simplex.isMarkedRedundant(5));
 }
 
+TEST(Simplextest, pivotRedundantRegressionTest) {
+  Simplex simplex(2);
+  simplex.addInequality({-1, 0, -1}); // x <= -1.
+  unsigned snapshot = simplex.getSnapshot();
+
+  simplex.addInequality({-1, 0, -2}); // x <= -2.
+  simplex.addInequality({-3, 0, -6});
+
+  // This first marks x <= -1 as redundant. Then it performs some more pivots
+  // to check if the other constraints are redundant. Pivot must update the
+  // non-redundant rows as well, otherwise these pivots result in an incorrect
+  // tableau state. In particular, after the rollback below, some rows that are
+  // NOT marked redundant will have an incorrect state.
+  simplex.detectRedundant();
+
+  // After the rollback, the only remaining constraint is x <= -1.
+  // The maximum value of x should be -1.
+  simplex.rollback(snapshot);
+  Optional<Fraction> maxX =
+      simplex.computeOptimum(Simplex::Direction::Up, {1, 0, 0});
+  EXPECT_TRUE(maxX.hasValue() && *maxX == Fraction(-1, 1));
+}
+
 TEST(SimplexTest, addInequality_already_redundant) {
   Simplex simplex(1);
   simplex.addInequality({1, -1}); // x >= 1.