[PATCH] D78806: SplitIndirectBrCriticalEdges: Fix Branch Probability update

[Yevgeny Rouban via Phabricator via llvm-commits]

yrouban created this revision.
yrouban added reviewers: craig.topper, arsenm, mkazantsev, asbirlea, ebrevnov, yamauchi.
yrouban added a project: LLVM.
Herald added subscribers: hiraditya, wdng.

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


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D78806

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


Index: llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
===================================================================
--- llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
+++ 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,48 @@
   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.
+  EXPECT_EQ(Block(".split").getTerminator()->getNumSuccessors(), 2u);
+  EXPECT_EQ(BPI.getEdgeProbability(&Block(".split"), 0u),
+            BranchProbability(1, 2));
+  EXPECT_EQ(BPI.getEdgeProbability(&Block(".split"), 1u),
+            BranchProbability(1, 2));
+}
Index: llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
===================================================================
--- llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
+++ llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp
@@ -388,13 +388,22 @@
     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 @@
       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


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D78806.259860.patch
Type: text/x-patch
Size: 3725 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200424/d90dd127/attachment.bin>