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/IR/Instructions.cpp =================================================================== --- llvm/lib/IR/Instructions.cpp +++ llvm/lib/IR/Instructions.cpp @@ -16,6 +16,7 @@ #include "llvm/ADT/SmallBitVector.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/Twine.h" +#include "llvm/Analysis/BranchProbabilityInfo.h" #include "llvm/IR/Attributes.h" #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Constant.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 @@ -83,5 +83,43 @@ EXPECT_TRUE(BPI.isEdgeHot(EntryBB, ExitBB)); } +TEST_F(BranchProbabilityInfoTest, SwapProbabilities) { + StringRef Assembly = R"( + define void @f() { + entry: + br label %loop + + loop: + %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop ] + %iv.next = add i32 %iv, 1 + %cond = icmp slt i32 %iv.next, 10 + br i1 %cond, label %exit, label %loop + + exit: + ret void + } + )"; + LLVMContext Context; + SMDiagnostic Error; + auto M = parseAssemblyString(Assembly, Error, Context); + ASSERT_TRUE(M) << "Bad assembly?"; + + Function *F = M->getFunction("f"); + auto *LoopHeaderBB = F->front().getSingleSuccessor(); + ASSERT_TRUE(LoopHeaderBB != nullptr); + BranchInst *Branch = dyn_cast(LoopHeaderBB->getTerminator()); + ASSERT_TRUE(Branch != nullptr); + // Save the probabilities before successors swapping + BranchProbabilityInfo *BPI = &buildBPI(*F); + auto ProbEdge0 = BPI->getEdgeProbability(LoopHeaderBB, 0U); + auto ProbEdge1 = BPI->getEdgeProbability(LoopHeaderBB, 1U); + EXPECT_LT(ProbEdge0, ProbEdge1); + + Branch->swapSuccessors(); + // 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)); +} + } // end anonymous namespace } // end namespace llvm