[MLIR][Presburger] SlowMPInt: fix bug in ceilDiv, floorDiv where widths weren't harmonized

[arjunp]

Summary

This also adds tests for abs, ceilDiv, floorDiv, mod, gcd and lcm.

[Groverkss]

mlir/unittests/Analysis/Presburger/MPIntTest.cpp

167–172

Could you also add tests when both arguments are over 64bit int?

[arjunp]

Address kunwar's comment

[Groverkss]

LGTM.

mlir/lib/Analysis/Presburger/SlowMPInt.cpp

 SlowMPInt SlowMPInt::operator/(const SlowMPInt &o) const {
   return SlowMPInt(
       runOpWithExpandOnOverflow(val, o.val, std::mem_fn(&APInt::sdiv_ov)));
 }
 SlowMPInt detail::abs(const SlowMPInt &x) { return x >= 0 ? x : -x; }
 SlowMPInt detail::ceilDiv(const SlowMPInt &lhs, const SlowMPInt &rhs) {
   if (rhs == -1)
     return -lhs;
-  return SlowMPInt(
-      llvm::APIntOps::RoundingSDiv(lhs.val, rhs.val, APInt::Rounding::UP));
+  unsigned width = getMaxWidth(lhs.val, rhs.val);
+  return SlowMPInt(llvm::APIntOps::RoundingSDiv(
+      lhs.val.sext(width), rhs.val.sext(width), APInt::Rounding::UP));
 }
 SlowMPInt detail::floorDiv(const SlowMPInt &lhs, const SlowMPInt &rhs) {
   if (rhs == -1)
     return -lhs;
-  return SlowMPInt(
-      llvm::APIntOps::RoundingSDiv(lhs.val, rhs.val, APInt::Rounding::DOWN));
+  unsigned width = getMaxWidth(lhs.val, rhs.val);
+  return SlowMPInt(llvm::APIntOps::RoundingSDiv(
+      lhs.val.sext(width), rhs.val.sext(width), APInt::Rounding::DOWN));
 }
 // The RHS is always expected to be positive, and the result
 /// is always non-negative.
 SlowMPInt detail::mod(const SlowMPInt &lhs, const SlowMPInt &rhs) {
   assert(rhs >= 1 && "mod is only supported for positive divisors!");
   return lhs % rhs < 0 ? lhs % rhs + rhs : lhs % rhs;
 }
 

mlir/unittests/Analysis/Presburger/MPIntTest.cpp

TYPED_TEST(IntTest, overflows) {
   EXPECT_EQ(-(2 * (-y)), 2 * y); // -(-2^63) overflow.
   x *= x;
   EXPECT_EQ(x, (y * y) / 16);
   y += y;
   y += y;
   y += y;
   y /= 8;
   EXPECT_EQ(y, 1ll << 62);
-  int64_t min = std::numeric_limits<int64_t>::min();
-  int64_t max = std::numeric_limits<int64_t>::max();
-  TypeParam z(min);
+
+  TypeParam min(std::numeric_limits<int64_t>::min());
+  TypeParam one(1);
+  EXPECT_EQ(floorDiv(min, -one), -min);
+  EXPECT_EQ(ceilDiv(min, -one), -min);
+  EXPECT_EQ(abs(min), -min);
+
+  TypeParam z = min;
   z /= -1;
-  EXPECT_EQ(z, -TypeParam(min));
+  EXPECT_EQ(z, -min);
   TypeParam w(min);
   --w;
   EXPECT_EQ(w, TypeParam(min) - 1);
   TypeParam u(min);
   u -= 1;
   EXPECT_EQ(u, w);
-  TypeParam v(max);
+
+  TypeParam max(std::numeric_limits<int64_t>::max());
+  TypeParam v = max;
   ++v;
-  EXPECT_EQ(v, TypeParam(max) + 1);
-  TypeParam t(max);
+  EXPECT_EQ(v, max + 1);
+  TypeParam t = max;
   t += 1;
   EXPECT_EQ(t, v);
 }
+
+TYPED_TEST(IntTest, floorCeilModAbsLcmGcd) {
+  TypeParam x(1ll << 50), one(1), two(2), three(3);
+
+  // Run on small values and large values.
+  for (const TypeParam &y : {x, x * x}) {
+    EXPECT_EQ(floorDiv(3 * y, three), y);
+    EXPECT_EQ(ceilDiv(3 * y, three), y);
+    EXPECT_EQ(floorDiv(3 * y - 1, three), y - 1);
+    EXPECT_EQ(ceilDiv(3 * y - 1, three), y);
+    EXPECT_EQ(floorDiv(3 * y - 2, three), y - 1);
+    EXPECT_EQ(ceilDiv(3 * y - 2, three), y);
+
+    EXPECT_EQ(mod(3 * y, three), 0);
+    EXPECT_EQ(mod(3 * y + 1, three), one);
+    EXPECT_EQ(mod(3 * y + 2, three), two);
+
+    EXPECT_EQ(floorDiv(3 * y, y), 3);
+    EXPECT_EQ(ceilDiv(3 * y, y), 3);
+    EXPECT_EQ(floorDiv(3 * y - 1, y), 2);
+    EXPECT_EQ(ceilDiv(3 * y - 1, y), 3);
+    EXPECT_EQ(floorDiv(3 * y - 2, y), 2);
+    EXPECT_EQ(ceilDiv(3 * y - 2, y), 3);
+
+    EXPECT_EQ(mod(3 * y, y), 0);
+    EXPECT_EQ(mod(3 * y + 1, y), 1);
+    EXPECT_EQ(mod(3 * y + 2, y), 2);
+
+    EXPECT_EQ(abs(y), y);
+    EXPECT_EQ(abs(-y), 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);
+    EXPECT_EQ(lcm(15 * y, 6 * y), 30 * y);
+  }
+}