Index: llvm/include/llvm/Analysis/BranchProbabilityInfo.h =================================================================== --- llvm/include/llvm/Analysis/BranchProbabilityInfo.h +++ llvm/include/llvm/Analysis/BranchProbabilityInfo.h @@ -189,6 +189,9 @@ /// 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(); Index: llvm/include/llvm/IR/Instructions.h =================================================================== --- llvm/include/llvm/IR/Instructions.h +++ llvm/include/llvm/IR/Instructions.h @@ -3275,7 +3275,7 @@ /// Swaps the successors of the branch instruction. This also swaps any /// branch weight metadata associated with the instruction so that it /// continues to map correctly to each operand. - void swapSuccessors(); + void swapSuccessors(BranchProbabilityInfo *BPI = nullptr); iterator_range successors() { return make_range( Index: llvm/lib/Analysis/BranchProbabilityInfo.cpp =================================================================== --- llvm/lib/Analysis/BranchProbabilityInfo.cpp +++ llvm/lib/Analysis/BranchProbabilityInfo.cpp @@ -1175,6 +1175,14 @@ } } +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, Index: llvm/lib/IR/Instructions.cpp =================================================================== --- llvm/lib/IR/Instructions.cpp +++ llvm/lib/IR/Instructions.cpp @@ -13,6 +13,7 @@ #include "llvm/IR/Instructions.h" #include "LLVMContextImpl.h" +#include "llvm/Analysis/BranchProbabilityInfo.h" #include "llvm/ADT/SmallBitVector.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/Twine.h" @@ -1440,7 +1441,7 @@ SubclassOptionalData = BI.SubclassOptionalData; } -void BranchInst::swapSuccessors() { +void BranchInst::swapSuccessors(BranchProbabilityInfo *BPI) { assert(isConditional() && "Cannot swap successors of an unconditional branch"); Op<-1>().swap(Op<-2>()); @@ -1448,6 +1449,8 @@ // Update profile metadata if present and it matches our structural // expectations. swapProfMetadata(); + if (BPI && getParent()) + BPI->swapSuccEdgesProbabilities(getParent()); } //===----------------------------------------------------------------------===// Index: llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp =================================================================== --- llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp +++ llvm/unittests/Analysis/BranchProbabilityInfoTest.cpp @@ -115,10 +115,10 @@ auto ProbEdge1 = BPI->getEdgeProbability(LoopHeaderBB, 1U); EXPECT_LT(ProbEdge0, ProbEdge1); - Branch->swapSuccessors(); + Branch->swapSuccessors(BPI); // 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