diff --git a/llvm/include/llvm/Support/BlockFrequency.h b/llvm/include/llvm/Support/BlockFrequency.h index 997d11e47ab7..dae520b54e48 100644 --- a/llvm/include/llvm/Support/BlockFrequency.h +++ b/llvm/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/llvm/lib/Support/BlockFrequency.cpp b/llvm/lib/Support/BlockFrequency.cpp index d1f8408dfcbb..00cf75bd5cf3 100644 --- a/llvm/lib/Support/BlockFrequency.cpp +++ b/llvm/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/llvm/unittests/Support/BlockFrequencyTest.cpp b/llvm/unittests/Support/BlockFrequencyTest.cpp index ffdea2c1790f..c3184515046a 100644 --- a/llvm/unittests/Support/BlockFrequencyTest.cpp +++ b/llvm/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); +} + }