[llvm] dc86900 - [BPI] Add method to swap outgoing edges probabilities

[Max Kazantsev via llvm-commits]

Author: Max Kazantsev
Date: 2023-04-18T20:05:24+07:00
New Revision: dc86900ff31e35e61ec9c5adca0488bf33d11833

URL: https://github.com/llvm/llvm-project/commit/dc86900ff31e35e61ec9c5adca0488bf33d11833
DIFF: https://github.com/llvm/llvm-project/commit/dc86900ff31e35e61ec9c5adca0488bf33d11833.diff

LOG: [BPI] Add method to swap outgoing edges probabilities

The motivation is need to update branch probability info after
swapping successors of branch instruction.

Differential Revision: https://reviews.llvm.org/D148237
Reviewed By: nikic

Added: 
    

Modified: 
    llvm/include/llvm/Analysis/BranchProbabilityInfo.h
    llvm/lib/Analysis/BranchProbabilityInfo.cpp
    llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
index 14d3080b5053..fb02997371bf 100644
--- a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
+++ b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
@@ -189,6 +189,9 @@ class BranchProbabilityInfo {
   /// unset for source.
   void copyEdgeProbabilities(BasicBlock *Src, BasicBlock *Dst);
 
+  /// Swap outgoing edges probabilities for \p Src with branch terminator
+  void swapSuccEdgesProbabilities(const BasicBlock *Src);
+
   static BranchProbability getBranchProbStackProtector(bool IsLikely) {
     static const BranchProbability LikelyProb((1u << 20) - 1, 1u << 20);
     return IsLikely ? LikelyProb : LikelyProb.getCompl();

diff  --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index d01c652333ff..b45deccd913d 100644
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -1175,6 +1175,14 @@ void BranchProbabilityInfo::copyEdgeProbabilities(BasicBlock *Src,
   }
 }
 
+void BranchProbabilityInfo::swapSuccEdgesProbabilities(const BasicBlock *Src) {
+  assert(Src->getTerminator()->getNumSuccessors() == 2);
+  if (!Probs.contains(std::make_pair(Src, 0)))
+    return; // No probability is set for edges from Src
+  assert(Probs.contains(std::make_pair(Src, 1)));
+  std::swap(Probs[std::make_pair(Src, 0)], Probs[std::make_pair(Src, 1)]);
+}
+
 raw_ostream &
 BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
                                             const BasicBlock *Src,

diff  --git a/llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp b/llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp
index 3d8877c7a40e..898dfef83a6e 100644
--- a/llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp
+++ b/llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp
@@ -116,9 +116,10 @@ TEST_F(BranchProbabilityInfoTest, SwapProbabilities) {
   EXPECT_LT(ProbEdge0, ProbEdge1);
 
   Branch->swapSuccessors();
+  BPI->swapSuccEdgesProbabilities(LoopHeaderBB);
   // TODO: Check the probabilities are swapped as well as the edges
-  EXPECT_EQ(ProbEdge0, BPI->getEdgeProbability(LoopHeaderBB, 0U));
-  EXPECT_EQ(ProbEdge1, BPI->getEdgeProbability(LoopHeaderBB, 1U));
+  EXPECT_EQ(ProbEdge0, BPI->getEdgeProbability(LoopHeaderBB, 1U));
+  EXPECT_EQ(ProbEdge1, BPI->getEdgeProbability(LoopHeaderBB, 0U));
 }
 
 } // end anonymous namespace