[llvm] f779ec7 - [BPI] Cache LoopExitBlocks to improve compile time (#93451)

[via llvm-commits]

Author: Enna1
Date: 2024-06-03T10:21:07+08:00
New Revision: f779ec7c13bdfccd29655d13a325f34c60797a76

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

LOG: [BPI] Cache LoopExitBlocks to improve compile time (#93451)

The `LoopBlock` stored in `LoopWorkList` consist of basic block and its
loop data information. When iterate `LoopWorkList`, if estimated weight
of a loop is not stored in `EstimatedLoopWeight`, `getLoopExitBlocks()`
is called to get all exit blocks of the loop. The estimated weight of a
loop is calculated by iterating over edges leading from basic block to
all exit blocks of the loop. If at least one edge has unknown estimated
weight, the estimated weight of loop is unknown and will not be stored
in `EstimatedLoopWeight`. `LoopWorkList` can contain different blocks in
a same loop, so there is wasted work that calls `getLoopExitBlocks()`
for same loop multiple times.

Since computing the exit blocks of loop is expensive and the loop
structure is not mutated in Branch Probability Analysis, we can cache
the result and improve compile time.

With this change, the overall compile time for a file containing a very
large loop is dropped by around 82%.

Added: 
    

Modified: 
    llvm/lib/Analysis/BranchProbabilityInfo.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index cd3e3a4991327..50dcd5f45233f 100644
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -810,6 +810,7 @@ void BranchProbabilityInfo::computeEestimateBlockWeight(
     const Function &F, DominatorTree *DT, PostDominatorTree *PDT) {
   SmallVector<BasicBlock *, 8> BlockWorkList;
   SmallVector<LoopBlock, 8> LoopWorkList;
+  SmallDenseMap<LoopData, SmallVector<BasicBlock *, 4>> LoopExitBlocks;
 
   // By doing RPO we make sure that all predecessors already have weights
   // calculated before visiting theirs successors.
@@ -828,12 +829,14 @@ void BranchProbabilityInfo::computeEestimateBlockWeight(
   do {
     while (!LoopWorkList.empty()) {
       const LoopBlock LoopBB = LoopWorkList.pop_back_val();
-
-      if (EstimatedLoopWeight.count(LoopBB.getLoopData()))
+      const LoopData LD = LoopBB.getLoopData();
+      if (EstimatedLoopWeight.count(LD))
         continue;
 
-      SmallVector<BasicBlock *, 4> Exits;
-      getLoopExitBlocks(LoopBB, Exits);
+      auto Res = LoopExitBlocks.try_emplace(LD);
+      SmallVectorImpl<BasicBlock *> &Exits = Res.first->second;
+      if (Res.second)
+        getLoopExitBlocks(LoopBB, Exits);
       auto LoopWeight = getMaxEstimatedEdgeWeight(
           LoopBB, make_range(Exits.begin(), Exits.end()));
 
@@ -842,7 +845,7 @@ void BranchProbabilityInfo::computeEestimateBlockWeight(
         if (LoopWeight <= static_cast<uint32_t>(BlockExecWeight::UNREACHABLE))
           LoopWeight = static_cast<uint32_t>(BlockExecWeight::LOWEST_NON_ZERO);
 
-        EstimatedLoopWeight.insert({LoopBB.getLoopData(), *LoopWeight});
+        EstimatedLoopWeight.insert({LD, *LoopWeight});
         // Add all blocks entering the loop into working list.
         getLoopEnterBlocks(LoopBB, BlockWorkList);
       }