diff --git a/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp b/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp --- a/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp +++ b/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp @@ -388,13 +388,22 @@ if (FirstNonPHI->isEHPad() || Target->isLandingPad()) continue; + // Remember edge probabilities if needed. + SmallVector EdgeProbabilities; + if (ShouldUpdateAnalysis) { + EdgeProbabilities.reserve(Target->getTerminator()->getNumSuccessors()); + for (unsigned I = 0, E = Target->getTerminator()->getNumSuccessors(); + I < E; ++I) + EdgeProbabilities.emplace_back(BPI->getEdgeProbability(Target, I)); + BPI->eraseBlock(Target); + } + BasicBlock *BodyBlock = Target->splitBasicBlock(FirstNonPHI, ".split"); if (ShouldUpdateAnalysis) { // Copy the BFI/BPI from Target to BodyBlock. for (unsigned I = 0, E = BodyBlock->getTerminator()->getNumSuccessors(); I < E; ++I) - BPI->setEdgeProbability(BodyBlock, I, - BPI->getEdgeProbability(Target, I)); + BPI->setEdgeProbability(BodyBlock, I, EdgeProbabilities[I]); BFI->setBlockFreq(BodyBlock, BFI->getBlockFreq(Target).getFrequency()); } // It's possible Target was its own successor through an indirectbr. @@ -423,7 +432,6 @@ BlockFrequency NewBlockFreqForTarget = BFI->getBlockFreq(Target) - BlockFreqForDirectSucc; BFI->setBlockFreq(Target, NewBlockFreqForTarget.getFrequency()); - BPI->eraseBlock(Target); } // Ok, now fix up the PHIs. We know the two blocks only have PHIs, and that diff --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp --- a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp +++ b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp @@ -7,6 +7,9 @@ //===----------------------------------------------------------------------===// #include "llvm/Transforms/Utils/BasicBlockUtils.h" +#include "llvm/Analysis/BlockFrequencyInfo.h" +#include "llvm/Analysis/BranchProbabilityInfo.h" +#include "llvm/Analysis/LoopInfo.h" #include "llvm/Analysis/PostDominators.h" #include "llvm/AsmParser/Parser.h" #include "llvm/IR/BasicBlock.h" @@ -136,3 +139,45 @@ EXPECT_TRUE(DT.verify()); EXPECT_TRUE(PDT.verify()); } + +TEST(BasicBlockUtils, SplitIndirectBrCriticalEdge) { + LLVMContext C; + + std::unique_ptr M = + parseIR(C, "define void @crit_edge(i8* %cond0, i1 %cond1) {\n" + "entry:\n" + " indirectbr i8* %cond0, [label %bb0, label %bb1]\n" + "bb0:\n" + " br label %bb1\n" + "bb1:\n" + " %p = phi i32 [0, %bb0], [0, %entry]\n" + " br i1 %cond1, label %bb2, label %bb3\n" + "bb2:\n" + " ret void\n" + "bb3:\n" + " ret void\n" + "}\n"); + + auto *F = M->getFunction("crit_edge"); + DominatorTree DT(*F); + LoopInfo LI(DT); + BranchProbabilityInfo BPI(*F, LI); + BlockFrequencyInfo BFI(*F, BPI, LI); + + auto Block = [&F](StringRef BBName) -> const BasicBlock & { + for (auto &BB : *F) + if (BB.getName() == BBName) + return BB; + llvm_unreachable("Block not found"); + }; + + bool Split = SplitIndirectBrCriticalEdges(*F, &BPI, &BFI); + + EXPECT_TRUE(Split); + + // Check that successors of the split block get their probability correct. + BasicBlock *SplitBB = Block("bb1").getTerminator()->getSuccessor(0); + EXPECT_EQ(2u, SplitBB->getTerminator()->getNumSuccessors()); + EXPECT_EQ(BranchProbability(1, 2), BPI.getEdgeProbability(SplitBB, 0u)); + EXPECT_EQ(BranchProbability(1, 2), BPI.getEdgeProbability(SplitBB, 1u)); +}