[MLIR][Presburger] SlowMPInt::gcd: fix crash when sizes differ

Reviewed By: Groverkss

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

diff --git a/mlir/lib/Analysis/Presburger/SlowMPInt.cpp b/mlir/lib/Analysis/Presburger/SlowMPInt.cpp
index 622aa46..bd1f0e2 100644 (file)
--- a/mlir/lib/Analysis/Presburger/SlowMPInt.cpp
+++ b/mlir/lib/Analysis/Presburger/SlowMPInt.cpp
@@ -221,7 +221,9 @@ SlowMPInt detail::mod(const SlowMPInt &lhs, const SlowMPInt &rhs) {
 
 SlowMPInt detail::gcd(const SlowMPInt &a, const SlowMPInt &b) {
   assert(a >= 0 && b >= 0 && "operands must be non-negative!");
-  return SlowMPInt(llvm::APIntOps::GreatestCommonDivisor(a.val, b.val));
+  unsigned width = getMaxWidth(a.val, b.val);
+  return SlowMPInt(llvm::APIntOps::GreatestCommonDivisor(a.val.sext(width),
+                                                         b.val.sext(width)));
 }
 
 /// Returns the least common multiple of 'a' and 'b'.

diff --git a/mlir/unittests/Analysis/Presburger/MPIntTest.cpp b/mlir/unittests/Analysis/Presburger/MPIntTest.cpp
index aa954a3..3c145d3 100644 (file)
--- a/mlir/unittests/Analysis/Presburger/MPIntTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/MPIntTest.cpp
@@ -190,6 +190,8 @@ TYPED_TEST(IntTest, floorCeilModAbsLcmGcd) {
     EXPECT_EQ(abs(y), y);
     EXPECT_EQ(abs(-y), y);
 
+    EXPECT_EQ(gcd(3 * y, three), three);
+    EXPECT_EQ(lcm(y, three), 3 * y);
     EXPECT_EQ(gcd(2 * y, 3 * y), y);
     EXPECT_EQ(lcm(2 * y, 3 * y), 6 * y);
     EXPECT_EQ(gcd(15 * y, 6 * y), 3 * y);