SplitIndirectBrCriticalEdges: Fix Branch Probability update

[yrouban]

Summary

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).

[hjyamauchi]

Why considering the cases where "target basic block happens to have unset a probability for any of its successors" specifically?

Does the new test TEST(BasicBlockUtils, SplitIndirectBrCriticalEdge) fail before the change in BreakCriticalEdges.cpp (where the target block doesn't seem to have an unset successor)?

[yrouban]

In D78806#2002145, @yamauchi wrote:

Does the new test TEST(BasicBlockUtils, SplitIndirectBrCriticalEdge) fail ..

Yes. It fails with 100% on both edges.

[ebrevnov]

Functionally the change looks ok.
I think we should improve current API of BP (not in this change though) because it's easy to break invariant that sum of all probabilities is 1.
We better disallow cases when only part of edges have probability set and others unset. 1/NumSuccesors will likely be incorrect value in most cases with uneven distribution.

[hjyamauchi]

llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp

146

Can you clang-format this code (as the lint check)?

169

Relying on the block name ".split" (that SplitIndirectBrCriticalEdges internally uses) may be a bit fragile. How about finding the block that's the predecessor of bb2/bb3?

[yrouban]

llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp

146

ok. but this will look different from the other test cases

[yrouban]

addressed the comments about the test

[hjyamauchi]

What happens in the newly added test seems that none of the successor probabilities are ever set (BranchProbabilityInfo::Probs is empty throughout). The thing is that BPI does not internally store the edge probabilities unless they are non-default, unevenly-split ones (1/NumOfSuccessors) and relies on the real-time state of the blocks/CFG for default cases, which is likely different from what the author of SplitIndirectBrCriticalEdges expected. In contrast, BFI stores the block frequencies at the time it was calculated, which stays there after the blocks/CFG are modified. When the block "bb1" gets split, the number of its successors goes from two (50% each) to one (100%). Then, this 100% probability incorrectly gets copied to each of the two successors of the new block.

[hjyamauchi]

LGTM.

llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp

for (BasicBlock *Target : Targets) {
     if (!IBRPred || OtherPreds.empty())
       continue;
 
     // Don't even think about ehpads/landingpads.
     Instruction *FirstNonPHI = Target->getFirstNonPHI();
     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.
     // In this case, the indirectbr now comes from BodyBlock.
     if (IBRPred == Target)
       IBRPred = BodyBlock;
 
     // At this point Target only has PHIs, and BodyBlock has the rest of the
for (BasicBlock *Pred : OtherPreds) {
         BlockFreqForDirectSucc += BFI->getBlockFreq(Src) *
             BPI->getEdgeProbability(Src, DirectSucc);
     }
     if (ShouldUpdateAnalysis) {
       BFI->setBlockFreq(DirectSucc, BlockFreqForDirectSucc.getFrequency());
       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
     // they are clones, so the number of PHIs are the same.
     // (a) Remove the edge coming from IBRPred from the "Direct" PHI
     // (b) Leave that as the only edge in the "Indirect" PHI.
     // (c) Merge the two in the body block.
     BasicBlock::iterator Indirect = Target->begin(),

llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp

 //===- BasicBlockUtils.cpp - Unit tests for BasicBlockUtils ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
TEST(BasicBlockUtils, SplitCriticalEdge) {
   DominatorTree DT(*F);
   PostDominatorTree PDT(*F);
 
   CriticalEdgeSplittingOptions CESO(&DT, nullptr, nullptr, &PDT);
   EXPECT_EQ(1u, SplitAllCriticalEdges(*F, CESO));
   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));
+}