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

[via llvm-commits]

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

>From f6b0a73402493eeecb75145c2e37eeff244b44f3 Mon Sep 17 00:00:00 2001
From: "xumingjie.enna1" <xumingjie.enna1 at bytedance.com>
Date: Mon, 27 May 2024 11:42:13 +0800
Subject: [PATCH 1/2] [BPI] Cache LoopExitBlocks to improve compile time

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%.
---
 llvm/include/llvm/Analysis/BranchProbabilityInfo.h |  3 +++
 llvm/lib/Analysis/BranchProbabilityInfo.cpp        | 11 ++++++-----
 2 files changed, 9 insertions(+), 5 deletions(-)

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);
       }

>From 139efcd4217976aa6ad10f6ed213c6774a158a4c Mon Sep 17 00:00:00 2001
From: "xumingjie.enna1" <xumingjie.enna1 at bytedance.com>
Date: Mon, 27 May 2024 19:24:36 +0800
Subject: [PATCH 2/2] address nikic's comments

---
 llvm/include/llvm/Analysis/BranchProbabilityInfo.h | 3 ---
 llvm/lib/Analysis/BranchProbabilityInfo.cpp        | 8 +++++---
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
index 9713629e1bf02..91e1872e9bd6f 100644
--- a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
+++ b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h
@@ -343,9 +343,6 @@ 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 4191fb64b88ec..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.
@@ -832,9 +833,10 @@ void BranchProbabilityInfo::computeEestimateBlockWeight(
       if (EstimatedLoopWeight.count(LD))
         continue;
 
-      if (!LoopExitBlocks.count(LD))
-        getLoopExitBlocks(LoopBB, LoopExitBlocks[LD]);
-      const SmallVectorImpl<BasicBlock *> &Exits = LoopExitBlocks[LD];
+      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()));