[mlir][Support] Avoid multiplication in floorDiv / ceilDiv

[herhut]

Summary

Using comparisons instead avoids potential overflow.

[herhut]

Add tests and fix behaviour for zero case.

mlir/include/mlir/Support/MathExtras.h

 namespace mlir {
 
 /// Returns the result of MLIR's ceildiv operation on constants. The RHS is
 /// expected to be non-zero.
 inline int64_t ceilDiv(int64_t lhs, int64_t rhs) {
   assert(rhs != 0);
   // C/C++'s integer division rounds towards 0.
   int64_t x = (rhs > 0) ? -1 : 1;
-  return (lhs * rhs > 0) ? ((lhs + x) / rhs) + 1 : -(-lhs / rhs);
+  return ((lhs != 0) && (lhs > 0) == (rhs > 0)) ? ((lhs + x) / rhs) + 1
+                                                : -(-lhs / rhs);
 }
 
 /// Returns the result of MLIR's floordiv operation on constants. The RHS is
 /// expected to be non-zero.
 inline int64_t floorDiv(int64_t lhs, int64_t rhs) {
   assert(rhs != 0);
   // C/C++'s integer division rounds towards 0.
   int64_t x = (rhs < 0) ? 1 : -1;
-  return (lhs * rhs < 0) ? -((-lhs + x) / rhs) - 1 : lhs / rhs;
+  return ((lhs != 0) && ((lhs < 0) != (rhs < 0))) ? -((-lhs + x) / rhs) - 1
+                                                  : lhs / rhs;
 }
 
 /// Returns MLIR's mod operation on constants. MLIR's mod operation yields the
 /// remainder of the Euclidean division of 'lhs' by 'rhs', and is therefore not
 /// C's % operator.  The RHS is always expected to be positive, and the result
 /// is always non-negative.
 inline int64_t mod(int64_t lhs, int64_t rhs) {
   assert(rhs >= 1);

mlir/unittests/Support/MathExtrasTest.cpp

 using namespace mlir;
 using ::testing::Eq;
 
 TEST(MathExtrasTest, CeilDivTest) {
   EXPECT_THAT(ceilDiv(14, 3), Eq(5));
   EXPECT_THAT(ceilDiv(14, -3), Eq(-4));
   EXPECT_THAT(ceilDiv(-14, -3), Eq(5));
   EXPECT_THAT(ceilDiv(-14, 3), Eq(-4));
+  EXPECT_THAT(ceilDiv(0, 3), Eq(0));
+  EXPECT_THAT(ceilDiv(0, -3), Eq(0));
 }
 
 TEST(MathExtrasTest, FloorDivTest) {
   EXPECT_THAT(floorDiv(14, 3), Eq(4));
   EXPECT_THAT(floorDiv(14, -3), Eq(-5));
   EXPECT_THAT(floorDiv(-14, -3), Eq(4));
   EXPECT_THAT(floorDiv(-14, 3), Eq(-5));
+  EXPECT_THAT(floorDiv(0, 3), Eq(0));
+  EXPECT_THAT(floorDiv(0, -3), Eq(0));
 }