mlir/MathExtras: consolidate with llvm/MathExtras

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/lib/Analysis/FlatLinearValueConstraints.cpp

@@ -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/Conversion/FuncToLLVM/FuncToLLVM.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/LLVMCommon/MemRefBuilder.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);
 

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

@@ -25,8 +25,8 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/Pass/Pass.h"
-#include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
+#include "llvm/Support/MathExtras.h"
 #include <optional>
 
 namespace mlir {
@@ -971,8 +971,8 @@ struct MemorySpaceCastOpLowering
                                resultUnderlyingDesc, resultElemPtrType);
 
       int64_t bytesToSkip =
-          2 *
-          ceilDiv(getTypeConverter()->getPointerBitwidth(resultAddrSpace), 8);
+          2 * llvm::divideCeilSigned(
+                  getTypeConverter()->getPointerBitwidth(resultAddrSpace), 8);
       Value bytesToSkipConst = rewriter.create<LLVM::ConstantOp>(
           loc, getIndexType(), rewriter.getIndexAttr(bytesToSkip));
       Value copySize = rewriter.create<LLVM::SubOp>(

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

@@ -21,7 +21,6 @@
 #include "mlir/IR/AffineExprVisitor.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/Dialect/Affine/Analysis/LoopAnalysis.cpp

@@ -18,7 +18,7 @@
 #include "mlir/Dialect/Affine/Analysis/NestedMatcher.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
-#include "mlir/Support/MathExtras.h"
+#include "llvm/Support/MathExtras.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallPtrSet.h"
@@ -47,7 +47,8 @@ void mlir::affine::getTripCountMapAndOperands(
     loopSpan = ub - lb;
     if (loopSpan < 0)
       loopSpan = 0;
-    *tripCountMap = AffineMap::getConstantMap(ceilDiv(loopSpan, step), context);
+    *tripCountMap = AffineMap::getConstantMap(
+        llvm::divideCeilSigned(loopSpan, step), context);
     tripCountOperands->clear();
     return;
   }

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -18,19 +18,23 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ShapedOpInterfaces.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
-#include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
 #include <numeric>
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::affine;
 
+using llvm::divideCeilSigned;
+using llvm::divideFloorSigned;
+using llvm::mod;
+
 #define DEBUG_TYPE "affine-ops"
 
 #include "mlir/Dialect/Affine/IR/AffineOpsDialect.cpp.inc"
@@ -824,19 +828,19 @@ static void simplifyExprAndOperands(AffineExpr &expr, unsigned numDims,
     // lhs floordiv c is a single value lhs is bounded in a range `c` that has
     // the same quotient.
     if (binExpr.getKind() == AffineExprKind::FloorDiv &&
-        floorDiv(lhsLbConstVal, rhsConstVal) ==
-            floorDiv(lhsUbConstVal, rhsConstVal)) {
-      expr =
-          getAffineConstantExpr(floorDiv(lhsLbConstVal, rhsConstVal), context);
+        divideFloorSigned(lhsLbConstVal, rhsConstVal) ==
+            divideFloorSigned(lhsUbConstVal, rhsConstVal)) {
+      expr = getAffineConstantExpr(
+          divideFloorSigned(lhsLbConstVal, rhsConstVal), context);
       return;
     }
     // lhs ceildiv c is a single value if the entire range has the same ceil
     // quotient.
     if (binExpr.getKind() == AffineExprKind::CeilDiv &&
-        ceilDiv(lhsLbConstVal, rhsConstVal) ==
-            ceilDiv(lhsUbConstVal, rhsConstVal)) {
-      expr =
-          getAffineConstantExpr(ceilDiv(lhsLbConstVal, rhsConstVal), context);
+        divideCeilSigned(lhsLbConstVal, rhsConstVal) ==
+            divideCeilSigned(lhsUbConstVal, rhsConstVal)) {
+      expr = getAffineConstantExpr(divideCeilSigned(lhsLbConstVal, rhsConstVal),
+                                   context);
       return;
     }
     // lhs mod c is lhs if the entire range has quotient 0 w.r.t the rhs.

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

@@ -23,7 +23,6 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/IntegerSet.h"
-#include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallPtrSet.h"

mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp

@@ -23,7 +23,6 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
-#include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/STLExtras.h"

mlir/lib/Dialect/Func/IR/FuncOps.cpp

@@ -19,7 +19,6 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/FunctionImplementation.h"
-#include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/MapVector.h"