[Mlir-commits] [llvm] [mlir] MathExtras: add overflow query for signed-div (PR #97901)

[Ramkumar Ramachandra]

https://github.com/artagnon created https://github.com/llvm/llvm-project/pull/97901

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.

>From f4c879cdd76757a9b9cb0ce7f70bbbd5e03e6350 Mon Sep 17 00:00:00 2001
From: Ramkumar Ramachandra <ramkumar.ramachandra at codasip.com>
Date: Sat, 6 Jul 2024 16:06:58 +0100
Subject: [PATCH] MathExtras: add overflow query for signed-div

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.
---
 llvm/include/llvm/ADT/DynamicAPInt.h   | 17 ++++++-----------
 llvm/include/llvm/Support/MathExtras.h | 13 +++++++++----
 mlir/lib/IR/AffineExpr.cpp             | 11 +++--------
 3 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/llvm/include/llvm/ADT/DynamicAPInt.h b/llvm/include/llvm/ADT/DynamicAPInt.h
index f312f776df971..9fe93a60d7653 100644
--- a/llvm/include/llvm/ADT/DynamicAPInt.h
+++ b/llvm/include/llvm/ADT/DynamicAPInt.h
@@ -231,13 +231,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.
 
 /// ---------------------------------------------------------------------------
@@ -338,7 +331,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());
   }
@@ -353,7 +346,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()));
   }
@@ -363,7 +357,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()));
   }
@@ -473,7 +468,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
index f6b1fdb6aba9e..fcc90cb43927f 100644
--- a/llvm/include/llvm/Support/MathExtras.h
+++ b/llvm/include/llvm/Support/MathExtras.h
@@ -413,9 +413,15 @@ 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 trying to negate the minimal signed value.
+constexpr bool divideSignedWouldOverflow(int64_t Numerator,
+                                         int64_t Denominator) {
+  return Numerator == std::numeric_limits<int64_t>::min() && Denominator == -1;
+}
+
 /// Returns the integer ceil(Numerator / Denominator). Signed version.
-/// Guaranteed to never overflow, unless Numerator is INT64_MIN and Denominator
-/// is -1.
+/// Overflows when divideSignedWouldOverflow returns true.
 template <typename U, typename V, typename T = common_sint<U, V>>
 constexpr T divideCeilSigned(U Numerator, V Denominator) {
   assert(Denominator && "Division by zero");
@@ -429,8 +435,7 @@ 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.
+/// Overflows when divideSignedWouldOverflow returns true.
 template <typename U, typename V, typename T = common_sint<U, V>>
 constexpr T divideFloorSigned(U Numerator, V Denominator) {
   assert(Denominator && "Division by zero");
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 798398464da8d..bfb7c4849356e 100644
--- 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());