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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-analysis

Author: Enna1 (Enna1)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/93451.diff


2 Files Affected:

- (modified) llvm/include/llvm/Analysis/BranchProbabilityInfo.h (+3) 
- (modified) llvm/lib/Analysis/BranchProbabilityInfo.cpp (+6-5) 


``````````diff
diff --git a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
index 91e1872e9bd6f..9713629e1bf02 100644
--- a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
+++ b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
@@ -343,6 +343,9 @@ class BranchProbabilityInfo {
   /// Keeps mapping of a loop to estimated weight to enter the loop.
   SmallDenseMap<LoopData, uint32_t> EstimatedLoopWeight;
 
+  /// Keeps mapping of a Loop to its "exit" blocks.
+  SmallDenseMap<LoopData, SmallVector<BasicBlock *, 4>> LoopExitBlocks;
+
   /// Helper to construct LoopBlock for \p BB.
   LoopBlock getLoopBlock(const BasicBlock *BB) const {
     return LoopBlock(BB, *LI, *SccI.get());
diff --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index cd3e3a4991327..4191fb64b88ec 100644
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -828,12 +828,13 @@ 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);
+      if (!LoopExitBlocks.count(LD))
+        getLoopExitBlocks(LoopBB, LoopExitBlocks[LD]);
+      const SmallVectorImpl<BasicBlock *> &Exits = LoopExitBlocks[LD];
       auto LoopWeight = getMaxEstimatedEdgeWeight(
           LoopBB, make_range(Exits.begin(), Exits.end()));
 
@@ -842,7 +843,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);
       }

``````````

</details>


https://github.com/llvm/llvm-project/pull/93451