[MLIR] Fix negative gcd in `normalizeDivisionByGCD` function.

When the coefficients of dividend are negative, the gcd may be negative
which will change the sign of dividend and overflow denominator.

Reviewed By: Groverkss

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

2 files changed, 29 insertions, 2 deletions

diff --git a/mlir/lib/Analysis/Presburger/Utils.cpp b/mlir/lib/Analysis/Presburger/Utils.cpp
index 765d4fb6a8df..7b8fe49b23c8 100644
--- a/mlir/lib/Analysis/Presburger/Utils.cpp
+++ b/mlir/lib/Analysis/Presburger/Utils.cpp
@@ -25,7 +25,10 @@ static void normalizeDivisionByGCD(SmallVectorImpl<int64_t> &dividend,
                                    unsigned &divisor) {
   if (divisor == 0 || dividend.empty())
     return;
-  int64_t gcd = llvm::greatestCommonDivisor(dividend.front(), int64_t(divisor));
+  // We take the absolute value of dividend's coefficients to make sure that
+  // `gcd` is positive.
+  int64_t gcd =
+      llvm::greatestCommonDivisor(std::abs(dividend.front()), int64_t(divisor));
 
   // The reason for ignoring the constant term is as follows.
   // For a division:
@@ -35,7 +38,7 @@ static void normalizeDivisionByGCD(SmallVectorImpl<int64_t> &dividend,
   // Since `{a/m}/d` in the dividend satisfies 0 <= {a/m}/d < 1/d, it will not
   // influence the result of the floor division and thus, can be ignored.
   for (size_t i = 1, m = dividend.size() - 1; i < m; i++) {
-    gcd = llvm::greatestCommonDivisor(dividend[i], gcd);
+    gcd = llvm::greatestCommonDivisor(std::abs(dividend[i]), gcd);
     if (gcd == 1)
       return;
   }
diff --git a/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp
index 2e4c13577043..5d1cfb4c6e78 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp
@@ -969,4 +969,28 @@ TEST(IntegerPolyhedronTest, mergeDivisionsConstants) {
   }
 }
 
+TEST(IntegerPolyhedronTest, negativeDividends) {
+  // (x) : (exists y = [-x + 1 / 2], z = [-x - 2 / 3]: y + z >= x).
+  IntegerPolyhedron poly1(1);
+  poly1.addLocalFloorDiv({-1, 1}, 2); // y = [x + 1 / 2].
+  // Normalization test with negative dividends
+  poly1.addLocalFloorDiv({-3, 0, -6}, 9); // z = [3x + 6 / 9] -> [x + 2 / 3].
+  poly1.addInequality({-1, 1, 1, 0});     // y + z >= x.
+
+  // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z <= x).
+  IntegerPolyhedron poly2(1);
+  // Normalization test.
+  poly2.addLocalFloorDiv({-2, 2}, 4);     // y = [-2x + 2 / 4] -> [-x + 1 / 2].
+  poly2.addLocalFloorDiv({-1, 0, -2}, 3); // z = [-x - 2 / 3].
+  poly2.addInequality({1, -1, -1, 0});    // y + z <= x.
+
+  poly1.mergeLocalIds(poly2);
+
+  // Merging triggers normalization.
+  std::vector<SmallVector<int64_t, 8>> divisions = {{-1, 0, 0, 1},
+                                                    {-1, 0, 0, -2}};
+  SmallVector<unsigned, 8> denoms = {2, 3};
+  checkDivisionRepresentation(poly1, divisions, denoms);
+}
+
 } // namespace mlir