diff --git a/include/llvm/Support/BlockFrequency.h b/include/llvm/Support/BlockFrequency.h index 997d11e47ab..dae520b54e4 100644 --- a/include/llvm/Support/BlockFrequency.h +++ b/include/llvm/Support/BlockFrequency.h @@ -55,6 +55,9 @@ public: BlockFrequency &operator+=(const BlockFrequency &Freq); const BlockFrequency operator+(const BlockFrequency &Freq) const; + /// \brief Shift block frequency to the right by count digits saturating to 1. + BlockFrequency &operator>>=(const unsigned count); + /// \brief Scale the given BlockFrequency by N/D. Return the remainder from /// the division by D. Upon overflow, the routine will saturate. uint32_t scale(const BranchProbability &Prob); diff --git a/lib/Support/BlockFrequency.cpp b/lib/Support/BlockFrequency.cpp index d1f8408dfcb..00cf75bd5cf 100644 --- a/lib/Support/BlockFrequency.cpp +++ b/lib/Support/BlockFrequency.cpp @@ -145,6 +145,18 @@ BlockFrequency::operator+(const BlockFrequency &Prob) const { return Freq; } +BlockFrequency &BlockFrequency::operator>>=(const unsigned count) { + // Frequency can never be 0 by design. + assert(Frequency != 0); + + // Shift right by count. + Frequency >>= count; + + // Saturate to 1 if we are 0. + Frequency |= Frequency == 0; + return *this; +} + uint32_t BlockFrequency::scale(const BranchProbability &Prob) { return scale(Prob.getNumerator(), Prob.getDenominator()); } diff --git a/unittests/Support/BlockFrequencyTest.cpp b/unittests/Support/BlockFrequencyTest.cpp index ffdea2c1790..c3184515046 100644 --- a/unittests/Support/BlockFrequencyTest.cpp +++ b/unittests/Support/BlockFrequencyTest.cpp @@ -237,4 +237,12 @@ TEST(BlockFrequencyTest, ProbabilityCompare) { EXPECT_FALSE(BigZero >= BigOne); } +TEST(BlockFrequencyTest, SaturatingRightShift) { + BlockFrequency Freq(0x10080ULL); + Freq >>= 2; + EXPECT_EQ(Freq.getFrequency(), 0x4020ULL); + Freq >>= 20; + EXPECT_EQ(Freq.getFrequency(), 0x1ULL); +} + }