SplitIndirectBrCriticalEdges: Fix Branch Probability update

Splitting critical edges for indirect branches
the SplitIndirectBrCriticalEdges() function may break branch
probabilities if target basic block happens to have unset
a probability for any of its successors. That is because in
such cases the getEdgeProbability(Target) function returns
probability 1/NumOfSuccessors and it is called after Target
was split (thus Target has a single successor). As the result
the correspondent successor of the split block gets
probability 100% but 1/NumOfSuccessors is expected (or better
be left unset).

Reviewers: yamauchi
Differential Revision: https://reviews.llvm.org/D78806
This commit is contained in:
Yevgeny Rouban 2020-05-07 12:10:48 +07:00
parent 20d67ffeae
commit b921543c49
2 changed files with 56 additions and 3 deletions

View File

@ -388,13 +388,22 @@ bool llvm::SplitIndirectBrCriticalEdges(Function &F,
if (FirstNonPHI->isEHPad() || Target->isLandingPad())
continue;
// Remember edge probabilities if needed.
SmallVector<BranchProbability, 4> 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 @@ bool llvm::SplitIndirectBrCriticalEdges(Function &F,
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

View File

@ -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 @@ TEST(BasicBlockUtils, SplitCriticalEdge) {
EXPECT_TRUE(DT.verify());
EXPECT_TRUE(PDT.verify());
}
TEST(BasicBlockUtils, SplitIndirectBrCriticalEdge) {
LLVMContext C;
std::unique_ptr<Module> 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));
}