1 files changed, 41 insertions, 30 deletions

diff --git a/llvm/lib/CodeGen/MachineBlockPlacement.cpp b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
index fcddf346cf6..fba33eb93d5 100644
--- a/llvm/lib/CodeGen/MachineBlockPlacement.cpp
+++ b/llvm/lib/CodeGen/MachineBlockPlacement.cpp
@@ -380,11 +380,19 @@ MachineBlockPlacement::selectBestSuccessor(MachineBasicBlock *BB,
   const BranchProbability HotProb(4, 5); // 80%
 
   MachineBasicBlock *BestSucc = nullptr;
-  auto BestProb = BranchProbability::getZero();
-
-  // Adjust edge probabilities by excluding edges pointing to blocks that is
-  // either not in BlockFilter or is already in the current chain. Consider the
-  // following CFG:
+  // FIXME: Due to the performance of the probability and weight routines in
+  // the MBPI analysis, we manually compute probabilities using the edge
+  // weights. This is suboptimal as it means that the somewhat subtle
+  // definition of edge weight semantics is encoded here as well. We should
+  // improve the MBPI interface to efficiently support query patterns such as
+  // this.
+  uint32_t BestWeight = 0;
+  uint32_t WeightScale = 0;
+  uint32_t SumWeight = MBPI->getSumForBlock(BB, WeightScale);
+
+  // Adjust sum of weights by excluding weights on edges pointing to blocks that
+  // is either not in BlockFilter or is already in the current chain. Consider
+  // the following CFG:
   //
   //     --->A
   //     |  / \
@@ -398,7 +406,7 @@ MachineBlockPlacement::selectBestSuccessor(MachineBasicBlock *BB,
   // HotProb). If we exclude E that is not in BlockFilter when calculating the
   // probability of C->D, D will be selected and we will get A C D B as the
   // layout of this loop.
-  auto AdjustedSumProb = BranchProbability::getOne();
+  uint32_t AdjustedSumWeight = SumWeight;
   SmallVector<MachineBasicBlock *, 4> Successors;
   for (MachineBasicBlock *Succ : BB->successors()) {
     bool SkipSucc = false;
@@ -416,20 +424,15 @@ MachineBlockPlacement::selectBestSuccessor(MachineBasicBlock *BB,
       }
     }
     if (SkipSucc)
-      AdjustedSumProb -= MBPI->getEdgeProbability(BB, Succ);
+      AdjustedSumWeight -= MBPI->getEdgeWeight(BB, Succ) / WeightScale;
     else
       Successors.push_back(Succ);
   }
 
   DEBUG(dbgs() << "Attempting merge from: " << getBlockName(BB) << "\n");
   for (MachineBasicBlock *Succ : Successors) {
-    BranchProbability SuccProb;
-    uint32_t SuccProbN = MBPI->getEdgeProbability(BB, Succ).getNumerator();
-    uint32_t SuccProbD = AdjustedSumProb.getNumerator();
-    if (SuccProbN >= SuccProbD)
-      SuccProb = BranchProbability::getOne();
-    else
-      SuccProb = BranchProbability(SuccProbN, SuccProbD);
+    uint32_t SuccWeight = MBPI->getEdgeWeight(BB, Succ);
+    BranchProbability SuccProb(SuccWeight / WeightScale, AdjustedSumWeight);
 
     // If we outline optional branches, look whether Succ is unavoidable, i.e.
     // dominates all terminators of the MachineFunction. If it does, other
@@ -467,7 +470,7 @@ MachineBlockPlacement::selectBestSuccessor(MachineBasicBlock *BB,
 
       // Make sure that a hot successor doesn't have a globally more
       // important predecessor.
-      auto RealSuccProb = MBPI->getEdgeProbability(BB, Succ);
+      BranchProbability RealSuccProb(SuccWeight / WeightScale, SumWeight);
       BlockFrequency CandidateEdgeFreq =
           MBFI->getBlockFreq(BB) * RealSuccProb * HotProb.getCompl();
       bool BadCFGConflict = false;
@@ -493,10 +496,10 @@ MachineBlockPlacement::selectBestSuccessor(MachineBasicBlock *BB,
                  << " (prob)"
                  << (SuccChain.LoopPredecessors != 0 ? " (CFG break)" : "")
                  << "\n");
-    if (BestSucc && BestProb >= SuccProb)
+    if (BestSucc && BestWeight >= SuccWeight)
       continue;
     BestSucc = Succ;
-    BestProb = SuccProb;
+    BestWeight = SuccWeight;
   }
   return BestSucc;
 }
@@ -725,6 +728,11 @@ MachineBlockPlacement::findBestLoopExit(MachineFunction &F, MachineLoop &L,
     MachineBasicBlock *OldExitingBB = ExitingBB;
     BlockFrequency OldBestExitEdgeFreq = BestExitEdgeFreq;
     bool HasLoopingSucc = false;
+    // FIXME: Due to the performance of the probability and weight routines in
+    // the MBPI analysis, we use the internal weights and manually compute the
+    // probabilities to avoid quadratic behavior.
+    uint32_t WeightScale = 0;
+    uint32_t SumWeight = MBPI->getSumForBlock(MBB, WeightScale);
     for (MachineBasicBlock *Succ : MBB->successors()) {
       if (Succ->isEHPad())
         continue;
@@ -738,10 +746,10 @@ MachineBlockPlacement::findBestLoopExit(MachineFunction &F, MachineLoop &L,
         continue;
       }
 
-      auto SuccProb = MBPI->getEdgeProbability(MBB, Succ);
+      uint32_t SuccWeight = MBPI->getEdgeWeight(MBB, Succ);
       if (LoopBlockSet.count(Succ)) {
         DEBUG(dbgs() << "    looping: " << getBlockName(MBB) << " -> "
-                     << getBlockName(Succ) << " (" << SuccProb << ")\n");
+                     << getBlockName(Succ) << " (" << SuccWeight << ")\n");
         HasLoopingSucc = true;
         continue;
       }
@@ -753,6 +761,7 @@ MachineBlockPlacement::findBestLoopExit(MachineFunction &F, MachineLoop &L,
           BlocksExitingToOuterLoop.insert(MBB);
       }
 
+      BranchProbability SuccProb(SuccWeight / WeightScale, SumWeight);
       BlockFrequency ExitEdgeFreq = MBFI->getBlockFreq(MBB) * SuccProb;
       DEBUG(dbgs() << "    exiting: " << getBlockName(MBB) << " -> "
                    << getBlockName(Succ) << " [L:" << SuccLoopDepth << "] (";
@@ -895,17 +904,21 @@ void MachineBlockPlacement::rotateLoopWithProfile(
   // edge from the tail of the loop chain.
   SmallVector<std::pair<MachineBasicBlock *, BlockFrequency>, 4> ExitsWithFreq;
   for (auto BB : LoopChain) {
-    auto LargestExitEdgeProb = BranchProbability::getZero();
+    uint32_t LargestExitEdgeWeight = 0;
     for (auto *Succ : BB->successors()) {
       BlockChain *SuccChain = BlockToChain[Succ];
       if (!LoopBlockSet.count(Succ) &&
           (!SuccChain || Succ == *SuccChain->begin())) {
-        auto SuccProb = MBPI->getEdgeProbability(BB, Succ);
-        LargestExitEdgeProb = std::max(LargestExitEdgeProb, SuccProb);
+        uint32_t SuccWeight = MBPI->getEdgeWeight(BB, Succ);
+        LargestExitEdgeWeight = std::max(LargestExitEdgeWeight, SuccWeight);
       }
     }
-    if (LargestExitEdgeProb > BranchProbability::getZero()) {
-      auto ExitFreq = MBFI->getBlockFreq(BB) * LargestExitEdgeProb;
+    if (LargestExitEdgeWeight > 0) {
+      uint32_t WeightScale = 0;
+      uint32_t SumWeight = MBPI->getSumForBlock(BB, WeightScale);
+      auto ExitFreq =
+          MBFI->getBlockFreq(BB) *
+          BranchProbability(LargestExitEdgeWeight / WeightScale, SumWeight);
       ExitsWithFreq.emplace_back(BB, ExitFreq);
     }
   }
@@ -1277,16 +1290,14 @@ void MachineBlockPlacement::buildCFGChains(MachineFunction &F) {
       }
 
       // If PrevBB has a two-way branch, try to re-order the branches
-      // such that we branch to the successor with higher probability first.
+      // such that we branch to the successor with higher weight first.
       if (TBB && !Cond.empty() && FBB &&
-          MBPI->getEdgeProbability(PrevBB, FBB) >
-              MBPI->getEdgeProbability(PrevBB, TBB) &&
+          MBPI->getEdgeWeight(PrevBB, FBB) > MBPI->getEdgeWeight(PrevBB, TBB) &&
           !TII->ReverseBranchCondition(Cond)) {
         DEBUG(dbgs() << "Reverse order of the two branches: "
                      << getBlockName(PrevBB) << "\n");
-        DEBUG(dbgs() << "    Edge probability: "
-                     << MBPI->getEdgeProbability(PrevBB, FBB) << " vs "
-                     << MBPI->getEdgeProbability(PrevBB, TBB) << "\n");
+        DEBUG(dbgs() << "    Edge weight: " << MBPI->getEdgeWeight(PrevBB, FBB)
+                     << " vs " << MBPI->getEdgeWeight(PrevBB, TBB) << "\n");
         DebugLoc dl; // FIXME: this is nowhere
         TII->RemoveBranch(*PrevBB);
         TII->InsertBranch(*PrevBB, FBB, TBB, Cond, dl);