mlir/MathExtras: consolidate with llvm/MathExtras (#95087)

This patch is part of a project to move the Presburger library into
LLVM.

Author: artagnon

llvm/include/llvm/Support/MathExtras.h

@@ -424,11 +424,43 @@ template <uint64_t Align> constexpr inline uint64_t alignTo(uint64_t Value) {
   return (Value + Align - 1) / Align * Align;
 }
 
-/// Returns the integer ceil(Numerator / Denominator).
+/// Returns the integer ceil(Numerator / Denominator). Unsigned integer version.
 inline uint64_t divideCeil(uint64_t Numerator, uint64_t Denominator) {
   return alignTo(Numerator, Denominator) / Denominator;
 }
 
+/// Returns the integer ceil(Numerator / Denominator). Signed integer version.
+inline int64_t divideCeilSigned(int64_t Numerator, int64_t Denominator) {
+  assert(Denominator && "Division by zero");
+  if (!Numerator)
+    return 0;
+  // C's integer division rounds towards 0.
+  int64_t X = (Denominator > 0) ? -1 : 1;
+  bool SameSign = (Numerator > 0) == (Denominator > 0);
+  return SameSign ? ((Numerator + X) / Denominator) + 1
+                  : Numerator / Denominator;
+}
+
+/// Returns the integer floor(Numerator / Denominator). Signed integer version.
+inline int64_t divideFloorSigned(int64_t Numerator, int64_t Denominator) {
+  assert(Denominator && "Division by zero");
+  if (!Numerator)
+    return 0;
+  // C's integer division rounds towards 0.
+  int64_t X = (Denominator > 0) ? -1 : 1;
+  bool SameSign = (Numerator > 0) == (Denominator > 0);
+  return SameSign ? Numerator / Denominator
+                  : -((-Numerator + X) / Denominator) - 1;
+}
+
+/// Returns the remainder of the Euclidean division of LHS by RHS. Result is
+/// always non-negative.
+inline int64_t mod(int64_t Numerator, int64_t Denominator) {
+  assert(Denominator >= 1 && "Mod by non-positive number");
+  int64_t Mod = Numerator % Denominator;
+  return Mod < 0 ? Mod + Denominator : Mod;
+}
+
 /// Returns the integer nearest(Numerator / Denominator).
 inline uint64_t divideNearest(uint64_t Numerator, uint64_t Denominator) {
   return (Numerator + (Denominator / 2)) / Denominator;

llvm/unittests/Support/MathExtrasTest.cpp

@@ -434,8 +434,39 @@ TEST(MathExtras, IsShiftedInt) {
   EXPECT_FALSE((isShiftedInt<6, 10>(int64_t(1) << 15)));
 }
 
-template <typename T>
-class OverflowTest : public ::testing::Test { };
+TEST(MathExtras, DivideCeilSigned) {
+  EXPECT_EQ(divideCeilSigned(14, 3), 5);
+  EXPECT_EQ(divideCeilSigned(15, 3), 5);
+  EXPECT_EQ(divideCeilSigned(14, -3), -4);
+  EXPECT_EQ(divideCeilSigned(-14, -3), 5);
+  EXPECT_EQ(divideCeilSigned(-14, 3), -4);
+  EXPECT_EQ(divideCeilSigned(-15, 3), -5);
+  EXPECT_EQ(divideCeilSigned(0, 3), 0);
+  EXPECT_EQ(divideCeilSigned(0, -3), 0);
+}
+
+TEST(MathExtras, DivideFloorSigned) {
+  EXPECT_EQ(divideFloorSigned(14, 3), 4);
+  EXPECT_EQ(divideFloorSigned(15, 3), 5);
+  EXPECT_EQ(divideFloorSigned(14, -3), -5);
+  EXPECT_EQ(divideFloorSigned(-14, -3), 4);
+  EXPECT_EQ(divideFloorSigned(-14, 3), -5);
+  EXPECT_EQ(divideFloorSigned(-15, 3), -5);
+  EXPECT_EQ(divideFloorSigned(0, 3), 0);
+  EXPECT_EQ(divideFloorSigned(0, -3), 0);
+}
+
+TEST(MathExtras, Mod) {
+  EXPECT_EQ(mod(1, 14), 1);
+  EXPECT_EQ(mod(-1, 14), 13);
+  EXPECT_EQ(mod(14, 3), 2);
+  EXPECT_EQ(mod(15, 3), 0);
+  EXPECT_EQ(mod(-14, 3), 1);
+  EXPECT_EQ(mod(-15, 3), 0);
+  EXPECT_EQ(mod(0, 3), 0);
+}
+
+template <typename T> class OverflowTest : public ::testing::Test {};
 
 using OverflowTestTypes = ::testing::Types<signed char, short, int, long,
                                            long long>;
@@ -560,5 +591,4 @@ TYPED_TEST(OverflowTest, MulResultZero) {
   EXPECT_FALSE(MulOverflow<TypeParam>(0, -5, Result));
   EXPECT_EQ(Result, TypeParam(0));
 }
-
 } // namespace

mlir/include/mlir/Analysis/Presburger/Fraction.h

@@ -15,7 +15,6 @@
 #define MLIR_ANALYSIS_PRESBURGER_FRACTION_H
 
 #include "mlir/Analysis/Presburger/MPInt.h"
-#include "mlir/Support/MathExtras.h"
 
 namespace mlir {
 namespace presburger {

mlir/include/mlir/Analysis/Presburger/MPInt.h

@@ -17,23 +17,17 @@
 #define MLIR_ANALYSIS_PRESBURGER_MPINT_H
 
 #include "mlir/Analysis/Presburger/SlowMPInt.h"
-#include "mlir/Support/MathExtras.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include <numeric>
 
 namespace mlir {
 namespace presburger {
-
-/// Redefine these functions, which operate on 64-bit ints, to also be part of
-/// the mlir::presburger namespace. This is useful because this file defines
-/// identically-named functions that operate on MPInts, which would otherwie
-/// become the only candidates of overload resolution when calling e.g. ceilDiv
-/// from the mlir::presburger namespace. So to access the 64-bit overloads, an
-/// explict call to mlir::ceilDiv would be required. These using declarations
-/// allow overload resolution to transparently call the right function.
-using ::mlir::ceilDiv;
-using ::mlir::floorDiv;
-using ::mlir::mod;
+using ::llvm::ArrayRef;
+using ::llvm::divideCeilSigned;
+using ::llvm::divideFloorSigned;
+using ::llvm::mod;
 
 namespace detail {
 /// If builtin intrinsics for overflow-checked arithmetic are available,
@@ -375,7 +369,7 @@ LLVM_ATTRIBUTE_ALWAYS_INLINE MPInt ceilDiv(const MPInt &lhs, const MPInt &rhs) {
   if (LLVM_LIKELY(lhs.isSmall() && rhs.isSmall())) {
     if (LLVM_UNLIKELY(detail::divWouldOverflow(lhs.getSmall(), rhs.getSmall())))
       return -lhs;
-    return MPInt(ceilDiv(lhs.getSmall(), rhs.getSmall()));
+    return MPInt(divideCeilSigned(lhs.getSmall(), rhs.getSmall()));
   }
   return MPInt(ceilDiv(detail::SlowMPInt(lhs), detail::SlowMPInt(rhs)));
 }
@@ -384,7 +378,7 @@ LLVM_ATTRIBUTE_ALWAYS_INLINE MPInt floorDiv(const MPInt &lhs,
   if (LLVM_LIKELY(lhs.isSmall() && rhs.isSmall())) {
     if (LLVM_UNLIKELY(detail::divWouldOverflow(lhs.getSmall(), rhs.getSmall())))
       return -lhs;
-    return MPInt(floorDiv(lhs.getSmall(), rhs.getSmall()));
+    return MPInt(divideFloorSigned(lhs.getSmall(), rhs.getSmall()));
   }
   return MPInt(floorDiv(detail::SlowMPInt(lhs), detail::SlowMPInt(rhs)));
 }

mlir/include/mlir/Analysis/Presburger/SlowMPInt.h

@@ -18,7 +18,6 @@
 #ifndef MLIR_ANALYSIS_PRESBURGER_SLOWMPINT_H
 #define MLIR_ANALYSIS_PRESBURGER_SLOWMPINT_H
 
-#include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/Support/raw_ostream.h"

mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h

@@ -17,7 +17,7 @@
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
-#include "mlir/Support/MathExtras.h"
+#include "llvm/Support/MathExtras.h"
 
 namespace mlir {
 namespace mesh {
@@ -114,7 +114,7 @@ inline int64_t shardDimension(int64_t dimSize, int64_t shardCount) {
     return ShapedType::kDynamic;
 
   assert(dimSize % shardCount == 0);
-  return ceilDiv(dimSize, shardCount);
+  return llvm::divideCeilSigned(dimSize, shardCount);
 }
 
 // Get the size of an unsharded dimension.

mlir/include/mlir/Support/MathExtras.h

@@ -16,7 +16,6 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/LLVM.h"
-#include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

@@ -36,7 +36,6 @@
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Support/LogicalResult.h"
-#include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SmallVector.h"

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

@@ -12,7 +12,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/Support/MathExtras.h"
+#include "llvm/Support/MathExtras.h"
 
 using namespace mlir;
 
@@ -363,9 +363,9 @@ void UnrankedMemRefDescriptor::computeSizes(
   // Initialize shared constants.
   Value one = createIndexAttrConstant(builder, loc, indexType, 1);
   Value two = createIndexAttrConstant(builder, loc, indexType, 2);
-  Value indexSize =
-      createIndexAttrConstant(builder, loc, indexType,
-                              ceilDiv(typeConverter.getIndexTypeBitwidth(), 8));
+  Value indexSize = createIndexAttrConstant(
+      builder, loc, indexType,
+      llvm::divideCeilSigned(typeConverter.getIndexTypeBitwidth(), 8));
 
   sizes.reserve(sizes.size() + values.size());
   for (auto [desc, addressSpace] : llvm::zip(values, addressSpaces)) {
@@ -378,7 +378,8 @@ void UnrankedMemRefDescriptor::computeSizes(
     // to data layout) into the unranked descriptor.
     Value pointerSize = createIndexAttrConstant(
         builder, loc, indexType,
-        ceilDiv(typeConverter.getPointerBitwidth(addressSpace), 8));
+        llvm::divideCeilSigned(typeConverter.getPointerBitwidth(addressSpace),
+                               8));
     Value doublePointerSize =
         builder.create<LLVM::MulOp>(loc, indexType, two, pointerSize);