[llvm] b921543 - SplitIndirectBrCriticalEdges: Fix Branch Probability update

[Yevgeny Rouban via llvm-commits]

Author: Yevgeny Rouban
Date: 2020-05-07T15:31:44+07:00
New Revision: b921543c494bfb321114602902aacb9392e840da

URL: https://github.com/llvm/llvm-project/commit/b921543c494bfb321114602902aacb9392e840da
DIFF: https://github.com/llvm/llvm-project/commit/b921543c494bfb321114602902aacb9392e840da.diff

LOG: 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

Added: 
    

Modified: 
    llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
    llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp b/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
index 008cea333e6b..f34d098cadc7 100644
--- a/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
+++ b/llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
@@ -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

diff  --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
index bf8228e99d81..899858e96792 100644
--- 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 @@ 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));
+}