3 files changed, 100 insertions, 8 deletions

diff --git a/llvm/lib/CodeGen/MachineBlockPlacement.cpp b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
index 186daa1cb8e..7379a3099b1 100644
--- a/llvm/lib/CodeGen/MachineBlockPlacement.cpp
+++ b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
@@ -125,6 +125,7 @@ BranchFoldPlacement("branch-fold-placement",
               cl::init(true), cl::Hidden);
 
 extern cl::opt<unsigned> StaticLikelyProb;
+extern cl::opt<unsigned> ProfileLikelyProb;
 
 namespace {
 class BlockChain;
@@ -520,13 +521,20 @@ bool MachineBlockPlacement::shouldPredBlockBeOutlined(
     return false;
 }
 
-// FIXME (PGO handling)
-// For now this method just returns a fixed threshold. It needs to be enhanced
-// such that BB and Succ is passed in so that CFG shapes are examined such that
-// the threshold is computed with more precise cost model when PGO is on.
-static BranchProbability getLayoutSuccessorProbThreshold() {
-  BranchProbability HotProb(StaticLikelyProb, 100);
-  return HotProb;
+// When profile is not present, return the StaticLikelyProb.
+// When profile is available, we need to handle the triangle-shape CFG.
+static BranchProbability getLayoutSuccessorProbThreshold(
+      MachineBasicBlock *BB) {
+  if (!BB->getParent()->getFunction()->getEntryCount())
+    return BranchProbability(StaticLikelyProb, 100);
+  if (BB->succ_size() == 2) {
+    const MachineBasicBlock *Succ1 = *BB->succ_begin();
+    const MachineBasicBlock *Succ2 = *(BB->succ_begin() + 1);
+    if (Succ1->isSuccessor(Succ2) || Succ2->isSuccessor(Succ1))
+      return BranchProbability(
+          200 - 2 * ProfileLikelyProb, 200 - ProfileLikelyProb);
+  }
+  return BranchProbability(ProfileLikelyProb, 100);
 }
 
 /// Checks to see if the layout candidate block \p Succ has a better layout
@@ -593,7 +601,7 @@ bool MachineBlockPlacement::hasBetterLayoutPredecessor(
   // edge: Prob(Succ->BB) needs to >= HotProb in order to be selected (without
   // profile data).
 
-  BranchProbability HotProb = getLayoutSuccessorProbThreshold();
+  BranchProbability HotProb = getLayoutSuccessorProbThreshold(BB);
 
   // Forward checking. For case 2, SuccProb will be 1.
   if (SuccProb < HotProb) {
diff --git a/llvm/lib/CodeGen/MachineBranchProbabilityInfo.cpp b/llvm/lib/CodeGen/MachineBranchProbabilityInfo.cpp
index 3554e4efa5f..9e14cb621dd 100644
--- a/llvm/lib/CodeGen/MachineBranchProbabilityInfo.cpp
+++ b/llvm/lib/CodeGen/MachineBranchProbabilityInfo.cpp
@@ -29,6 +29,12 @@ cl::opt<unsigned> StaticLikelyProb(
     cl::desc("branch probability threshold to be considered very likely"),
     cl::init(80), cl::Hidden);
 
+cl::opt<unsigned> ProfileLikelyProb(
+    "profile-likely-prob",
+    cl::desc("branch probability threshold to be considered very likely "
+             "when profile is available"),
+    cl::init(51), cl::Hidden);
+
 char MachineBranchProbabilityInfo::ID = 0;
 
 void MachineBranchProbabilityInfo::anchor() {}
diff --git a/llvm/test/CodeGen/X86/block-placement.ll b/llvm/test/CodeGen/X86/block-placement.ll
index b83180ad509..aeea6c70417 100644
--- a/llvm/test/CodeGen/X86/block-placement.ll
+++ b/llvm/test/CodeGen/X86/block-placement.ll
@@ -1207,4 +1207,82 @@ exit:
   ret void
 }
 
+define void @test_hot_branch_profile(i32* %a) !prof !6 {
+; Test that a hot branch that has a probability a little larger than 60% will
+; break CFG constrains when doing block placement when profile is available.
+; CHECK-LABEL: test_hot_branch_profile:
+; CHECK: %entry
+; CHECK: %then
+; CHECK: %exit
+; CHECK: %else
+
+entry:
+  %gep1 = getelementptr i32, i32* %a, i32 1
+  %val1 = load i32, i32* %gep1
+  %cond1 = icmp ugt i32 %val1, 1
+  br i1 %cond1, label %then, label %else, !prof !5
+
+then:
+  call void @hot_function()
+  br label %exit
+
+else:
+  call void @cold_function()
+  br label %exit
+
+exit:
+  call void @hot_function()
+  ret void
+}
+
+define void @test_hot_branch_triangle_profile(i32* %a) !prof !6 {
+; Test that a hot branch that has a probability a little larger than 80% will
+; break triangle shaped CFG constrains when doing block placement if profile
+; is present.
+; CHECK-LABEL: test_hot_branch_triangle_profile:
+; CHECK: %entry
+; CHECK: %exit
+; CHECK: %then
+
+entry:
+  %gep1 = getelementptr i32, i32* %a, i32 1
+  %val1 = load i32, i32* %gep1
+  %cond1 = icmp ugt i32 %val1, 1
+  br i1 %cond1, label %exit, label %then, !prof !5
+
+then:
+  call void @hot_function()
+  br label %exit
+
+exit:
+  call void @hot_function()
+  ret void
+}
+
+define void @test_hot_branch_triangle_profile_topology(i32* %a) !prof !6 {
+; Test that a hot branch that has a probability between 50% and 66% will not
+; break triangle shaped CFG constrains when doing block placement if profile
+; is present.
+; CHECK-LABEL: test_hot_branch_triangle_profile_topology:
+; CHECK: %entry
+; CHECK: %then
+; CHECK: %exit
+
+entry:
+  %gep1 = getelementptr i32, i32* %a, i32 1
+  %val1 = load i32, i32* %gep1
+  %cond1 = icmp ugt i32 %val1, 1
+  br i1 %cond1, label %exit, label %then, !prof !7
+
+then:
+  call void @hot_function()
+  br label %exit
+
+exit:
+  call void @hot_function()
+  ret void
+}
+
 !5 = !{!"branch_weights", i32 84, i32 16}
+!6 = !{!"function_entry_count", i32 10}
+!7 = !{!"branch_weights", i32 60, i32 40}