diff --git a/llvm/include/llvm/ADT/APInt.h b/llvm/include/llvm/ADT/APInt.h --- a/llvm/include/llvm/ADT/APInt.h +++ b/llvm/include/llvm/ADT/APInt.h @@ -2181,6 +2181,13 @@ /// Return A sign-divided by B, rounded by the given rounding mode. APInt RoundingSDiv(const APInt &A, const APInt &B, APInt::Rounding RM); +/// Return A logically shifted right by B, rounded by the given rounding mode. +APInt RoundingLShr(const APInt &A, const APInt &B, APInt::Rounding RM); + +/// Return A arithmetically shifted right by B, rounded by the given rounding +/// mode. +APInt RoundingAShr(const APInt &A, const APInt &B, APInt::Rounding RM); + /// Let q(n) = An^2 + Bn + C, and BW = bit width of the value range /// (e.g. 32 for i32). /// This function finds the smallest number n, such that diff --git a/llvm/lib/Support/APInt.cpp b/llvm/lib/Support/APInt.cpp --- a/llvm/lib/Support/APInt.cpp +++ b/llvm/lib/Support/APInt.cpp @@ -2797,6 +2797,47 @@ llvm_unreachable("Unknown APInt::Rounding enum"); } +APInt llvm::APIntOps::RoundingLShr(const APInt &A, const APInt &B, + APInt::Rounding RM) { + // Currently udivrem always rounds down. + switch (RM) { + case APInt::Rounding::DOWN: + case APInt::Rounding::TOWARD_ZERO: + return A.lshr(B); + case APInt::Rounding::UP: { + APInt Quo = A.lshr(B); + bool HasRem = A != Quo.shl(B); + if (!HasRem) + return Quo; + return Quo + 1; + } + } + llvm_unreachable("Unknown APInt::Rounding enum"); +} + +APInt llvm::APIntOps::RoundingAShr(const APInt &A, const APInt &B, + APInt::Rounding RM) { + switch (RM) { + case APInt::Rounding::TOWARD_ZERO: + case APInt::Rounding::UP: { + APInt Quo = A.ashr(B); + APInt RoundedA = Quo.shl(B); + bool HasRem = A != RoundedA; + if (!HasRem) + return Quo; + if (RM == APInt::Rounding::UP) + return Quo + 1; + if (A.isNegative()) + return Quo + 1; + return Quo; + } + // Currently ashr rounds down. + case APInt::Rounding::DOWN: + return A.ashr(B); + } + llvm_unreachable("Unknown APInt::Rounding enum"); +} + Optional llvm::APIntOps::SolveQuadraticEquationWrap(APInt A, APInt B, APInt C, unsigned RangeWidth) { diff --git a/llvm/unittests/ADT/APIntTest.cpp b/llvm/unittests/ADT/APIntTest.cpp --- a/llvm/unittests/ADT/APIntTest.cpp +++ b/llvm/unittests/ADT/APIntTest.cpp @@ -2526,16 +2526,16 @@ { APInt Quo = APIntOps::RoundingSDiv(A, B, APInt::Rounding::UP); auto Prod = Quo.sext(16) * B.sext(16); - EXPECT_TRUE(Prod.uge(A)); - if (Prod.ugt(A)) { + EXPECT_TRUE(Prod.sge(A)); + if (Prod.sgt(A)) { EXPECT_TRUE(((Quo - 1).sext(16) * B.sext(16)).ult(A)); } } { APInt Quo = APIntOps::RoundingSDiv(A, B, APInt::Rounding::DOWN); auto Prod = Quo.sext(16) * B.sext(16); - EXPECT_TRUE(Prod.ule(A)); - if (Prod.ult(A)) { + EXPECT_TRUE(Prod.sle(A)); + if (Prod.slt(A)) { EXPECT_TRUE(((Quo + 1).sext(16) * B.sext(16)).ugt(A)); } } @@ -2547,6 +2547,115 @@ } } +TEST(APIntTest, RoundingLShr) { + for (uint64_t Ai = 0; Ai <= 255; Ai++) { + APInt A(8, Ai); + + for (uint64_t Bi = 0; Bi <= 7; Bi++) { + APInt B(8, Bi); + + if (Bi == 0) { + EXPECT_EQ(A, APIntOps::RoundingLShr(A, B, APInt::Rounding::UP)); + EXPECT_EQ(A, APIntOps::RoundingLShr(A, B, APInt::Rounding::DOWN)); + EXPECT_EQ(A, + APIntOps::RoundingLShr(A, B, APInt::Rounding::TOWARD_ZERO)); + } + + { + APInt Quo = APIntOps::RoundingLShr(A, B, APInt::Rounding::UP); + auto Prod = Quo.zext(16) << B.zext(16); + EXPECT_TRUE(Prod.uge(Ai)); + if (Prod.ugt(Ai)) { + EXPECT_TRUE(((Quo - 1).zext(16) << B.zext(16)).ult(Ai)); + } + } + { + APInt Quo = A.lshr(B); + EXPECT_EQ(Quo, + APIntOps::RoundingLShr(A, B, APInt::Rounding::TOWARD_ZERO)); + EXPECT_EQ(Quo, APIntOps::RoundingLShr(A, B, APInt::Rounding::DOWN)); + } + } + } +} + +TEST(APIntTest, RoundingAShr) { + for (int64_t Ai = -128; Ai <= 127; Ai++) { + APInt A(8, Ai); + + for (uint64_t Bi = 0; Bi <= 7; Bi++) { + APInt B(8, Bi); + + if (Bi == 0) { + EXPECT_EQ(A, APIntOps::RoundingAShr(A, B, APInt::Rounding::UP)); + EXPECT_EQ(A, APIntOps::RoundingAShr(A, B, APInt::Rounding::DOWN)); + EXPECT_EQ(A, + APIntOps::RoundingAShr(A, B, APInt::Rounding::TOWARD_ZERO)); + } + { + APInt Quo = APIntOps::RoundingAShr(A, B, APInt::Rounding::UP); + auto Prod = Quo.sext(16) << B.sext(16); + EXPECT_TRUE(Prod.sge(Ai)); + if (Prod.sgt(Ai)) { + EXPECT_TRUE(((Quo - 1).sext(16) << B.sext(16)).slt(Ai)); + } + } + { + APInt Quo = APIntOps::RoundingAShr(A, B, APInt::Rounding::TOWARD_ZERO); + auto Prod = Quo.sext(16) << B.sext(16); + EXPECT_TRUE(Prod.ule(Ai)); + if (Prod.ult(Ai)) { + EXPECT_TRUE(((Quo + 1).sext(16) << B.sext(16)).sgt(Ai)); + } + } + { + APInt Quo = A.ashr(B); + EXPECT_EQ(Quo, APIntOps::RoundingAShr(A, B, APInt::Rounding::DOWN)); + } + } + } +} + +TEST(APIntTest, RoundingLShrVsUDiv) { + for (uint64_t Ai = 0; Ai <= 255; Ai++) { + APInt A(8, Ai); + + for (uint64_t Bi = 0; Bi <= 7; Bi++) { + APInt B(8, Bi); + + APInt Divisor(8, 1 << Bi); + + EXPECT_EQ(APIntOps::RoundingUDiv(A, Divisor, APInt::Rounding::UP), + APIntOps::RoundingLShr(A, B, APInt::Rounding::UP)); + EXPECT_EQ(APIntOps::RoundingUDiv(A, Divisor, APInt::Rounding::DOWN), + APIntOps::RoundingLShr(A, B, APInt::Rounding::DOWN)); + EXPECT_EQ( + APIntOps::RoundingUDiv(A, Divisor, APInt::Rounding::TOWARD_ZERO), + APIntOps::RoundingLShr(A, B, APInt::Rounding::TOWARD_ZERO)); + } + } +} + +TEST(APIntTest, RoundingAShrVsSDiv) { + for (uint64_t Ai = 0; Ai <= 255; Ai++) { + APInt A(8, Ai); + + for (uint64_t Bi = 0; Bi <= 6; Bi++) { + APInt B(8, Bi); + + APInt Divisor(8, 1 << Bi); + + EXPECT_EQ(APIntOps::RoundingSDiv(A, Divisor, APInt::Rounding::UP), + APIntOps::RoundingAShr(A, B, APInt::Rounding::UP)); + EXPECT_EQ(APIntOps::RoundingSDiv(A, Divisor, APInt::Rounding::DOWN), + APIntOps::RoundingAShr(A, B, APInt::Rounding::DOWN)); + EXPECT_EQ( + APIntOps::RoundingSDiv(A, Divisor, APInt::Rounding::TOWARD_ZERO), + APIntOps::RoundingAShr(A, B, APInt::Rounding::TOWARD_ZERO)); + } + } +} + TEST(APIntTest, umul_ov) { const std::pair Overflows[] = { {0x8000000000000000, 2},