MathExtras: add overflow query for signed-div (llvm#97901)

artagnon · aaryanshukla · commit efcdbd034298 · 2024-07-14T17:59:26.000Z
5221634 (Do not trigger UB during AffineExpr parsing) noticed that divideCeilSigned and divideFloorSigned would overflow when Numerator = INT_MIN, and Denominator = -1. This observation has already been made by DynamicAPInt, and it has code to check this. To avoid checks in multiple callers, centralize this query in MathExtras, and change divideCeilSigned/divideFloorSigned to assert on overflow.
diff --git a/llvm/include/llvm/ADT/DynamicAPInt.h b/llvm/include/llvm/ADT/DynamicAPInt.h
@@ -241,13 +241,6 @@ LLVM_ATTRIBUTE_ALWAYS_INLINE DynamicAPInt dynamicAPIntFromInt64(int64_t X) {
 LLVM_ATTRIBUTE_ALWAYS_INLINE DynamicAPInt mod(const DynamicAPInt &LHS,
                                               const DynamicAPInt &RHS);
 
-namespace detail {
-// Division overflows only when trying to negate the minimal signed value.
-LLVM_ATTRIBUTE_ALWAYS_INLINE bool divWouldOverflow(int64_t X, int64_t Y) {
-  return X == std::numeric_limits<int64_t>::min() && Y == -1;
-}
-} // namespace detail
-
 /// We define the operations here in the header to facilitate inlining.
 
 /// ---------------------------------------------------------------------------
@@ -348,7 +341,7 @@ LLVM_ATTRIBUTE_ALWAYS_INLINE DynamicAPInt
 DynamicAPInt::operator/(const DynamicAPInt &O) const {
   if (LLVM_LIKELY(isSmall() && O.isSmall())) {
     // Division overflows only occur when negating the minimal possible value.
-    if (LLVM_UNLIKELY(detail::divWouldOverflow(getSmall(), O.getSmall())))
+    if (LLVM_UNLIKELY(divideSignedWouldOverflow(getSmall(), O.getSmall())))
       return -*this;
     return DynamicAPInt(getSmall() / O.getSmall());
   }
@@ -363,7 +356,8 @@ LLVM_ATTRIBUTE_ALWAYS_INLINE DynamicAPInt abs(const DynamicAPInt &X) {
 LLVM_ATTRIBUTE_ALWAYS_INLINE DynamicAPInt ceilDiv(const DynamicAPInt &LHS,
                                                   const DynamicAPInt &RHS) {
   if (LLVM_LIKELY(LHS.isSmall() && RHS.isSmall())) {
-    if (LLVM_UNLIKELY(detail::divWouldOverflow(LHS.getSmall(), RHS.getSmall())))
+    if (LLVM_UNLIKELY(
+            divideSignedWouldOverflow(LHS.getSmall(), RHS.getSmall())))
       return -LHS;
     return DynamicAPInt(divideCeilSigned(LHS.getSmall(), RHS.getSmall()));
   }
@@ -373,7 +367,8 @@ LLVM_ATTRIBUTE_ALWAYS_INLINE DynamicAPInt ceilDiv(const DynamicAPInt &LHS,
 LLVM_ATTRIBUTE_ALWAYS_INLINE DynamicAPInt floorDiv(const DynamicAPInt &LHS,
                                                    const DynamicAPInt &RHS) {
   if (LLVM_LIKELY(LHS.isSmall() && RHS.isSmall())) {
-    if (LLVM_UNLIKELY(detail::divWouldOverflow(LHS.getSmall(), RHS.getSmall())))
+    if (LLVM_UNLIKELY(
+            divideSignedWouldOverflow(LHS.getSmall(), RHS.getSmall())))
       return -LHS;
     return DynamicAPInt(divideFloorSigned(LHS.getSmall(), RHS.getSmall()));
   }
@@ -483,7 +478,7 @@ LLVM_ATTRIBUTE_ALWAYS_INLINE DynamicAPInt &
 DynamicAPInt::operator/=(const DynamicAPInt &O) {
   if (LLVM_LIKELY(isSmall() && O.isSmall())) {
     // Division overflows only occur when negating the minimal possible value.
-    if (LLVM_UNLIKELY(detail::divWouldOverflow(getSmall(), O.getSmall())))
+    if (LLVM_UNLIKELY(divideSignedWouldOverflow(getSmall(), O.getSmall())))
       return *this = -*this;
     getSmall() /= O.getSmall();
     return *this;
diff --git a/llvm/include/llvm/Support/MathExtras.h b/llvm/include/llvm/Support/MathExtras.h
@@ -413,12 +413,20 @@ constexpr uint64_t divideCeil(uint64_t Numerator, uint64_t Denominator) {
   return (Numerator - Bias) / Denominator + Bias;
 }
 
+// Check whether divideCeilSigned or divideFloorSigned would overflow. This
+// happens only when Numerator = INT_MIN and Denominator = -1.
+template <typename U, typename V>
+constexpr bool divideSignedWouldOverflow(U Numerator, V Denominator) {
+  return Numerator == std::numeric_limits<U>::min() && Denominator == -1;
+}
+
 /// Returns the integer ceil(Numerator / Denominator). Signed version.
-/// Guaranteed to never overflow, unless Numerator is INT64_MIN and Denominator
-/// is -1.
+/// Overflow is explicitly forbidden with an assert.
 template <typename U, typename V, typename T = common_sint<U, V>>
 constexpr T divideCeilSigned(U Numerator, V Denominator) {
   assert(Denominator && "Division by zero");
+  assert(!divideSignedWouldOverflow(Numerator, Denominator) &&
+         "Divide would overflow");
   if (!Numerator)
     return 0;
   // C's integer division rounds towards 0.
@@ -429,11 +437,12 @@ constexpr T divideCeilSigned(U Numerator, V Denominator) {
 }
 
 /// Returns the integer floor(Numerator / Denominator). Signed version.
-/// Guaranteed to never overflow, unless Numerator is INT64_MIN and Denominator
-/// is -1.
+/// Overflow is explicitly forbidden with an assert.
 template <typename U, typename V, typename T = common_sint<U, V>>
 constexpr T divideFloorSigned(U Numerator, V Denominator) {
   assert(Denominator && "Division by zero");
+  assert(!divideSignedWouldOverflow(Numerator, Denominator) &&
+         "Divide would overflow");
   if (!Numerator)
     return 0;
   // C's integer division rounds towards 0.
diff --git a/llvm/unittests/Support/MathExtrasTest.cpp b/llvm/unittests/Support/MathExtrasTest.cpp
@@ -523,6 +523,12 @@ TEST(MathExtras, DivideCeil) {
             std::numeric_limits<int64_t>::min() / 2 + 1);
   EXPECT_EQ(divideCeilSigned(std::numeric_limits<int64_t>::min(), 1),
             std::numeric_limits<int64_t>::min());
+
+  // Overflow.
+  EXPECT_TRUE(
+      divideSignedWouldOverflow(std::numeric_limits<int8_t>::min(), -1));
+  EXPECT_TRUE(
+      divideSignedWouldOverflow(std::numeric_limits<int64_t>::min(), -1));
 }
 
 TEST(MathExtras, DivideFloorSigned) {
@@ -544,6 +550,8 @@ TEST(MathExtras, DivideFloorSigned) {
             std::numeric_limits<int64_t>::min() / 2);
   EXPECT_EQ(divideFloorSigned(std::numeric_limits<int64_t>::min(), 1),
             std::numeric_limits<int64_t>::min());
+
+  // Same overflow condition, divideSignedWouldOverflow, applies.
 }
 
 TEST(MathExtras, Mod) {
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
@@ -26,6 +26,7 @@ using namespace mlir::detail;
 
 using llvm::divideCeilSigned;
 using llvm::divideFloorSigned;
+using llvm::divideSignedWouldOverflow;
 using llvm::mod;
 
 MLIRContext *AffineExpr::getContext() const { return expr->context; }
@@ -859,11 +860,8 @@ static AffineExpr simplifyFloorDiv(AffineExpr lhs, AffineExpr rhs) {
     return nullptr;
 
   if (lhsConst) {
-    // divideFloorSigned can only overflow in this case:
-    if (lhsConst.getValue() == std::numeric_limits<int64_t>::min() &&
-        rhsConst.getValue() == -1) {
+    if (divideSignedWouldOverflow(lhsConst.getValue(), rhsConst.getValue()))
       return nullptr;
-    }
     return getAffineConstantExpr(
         divideFloorSigned(lhsConst.getValue(), rhsConst.getValue()),
         lhs.getContext());
@@ -921,11 +919,8 @@ static AffineExpr simplifyCeilDiv(AffineExpr lhs, AffineExpr rhs) {
     return nullptr;
 
   if (lhsConst) {
-    // divideCeilSigned can only overflow in this case:
-    if (lhsConst.getValue() == std::numeric_limits<int64_t>::min() &&
-        rhsConst.getValue() == -1) {
+    if (divideSignedWouldOverflow(lhsConst.getValue(), rhsConst.getValue()))
       return nullptr;
-    }
     return getAffineConstantExpr(
         divideCeilSigned(lhsConst.getValue(), rhsConst.getValue()),
         lhs.getContext());
