[PATCH] D78806: SplitIndirectBrCriticalEdges: Fix Branch Probability update

[Yevgeny Rouban via Phabricator via llvm-commits]

yrouban updated this revision to Diff 261477.
yrouban added a comment.

addressed the comments about the test


CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D78806/new/

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,45 @@
   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));
+}
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.261477.patch
Type: text/x-patch
Size: 3892 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200502/c968e428/attachment-0001.bin>