[BPI] Keep BPI available in loop passes through LoopStandardAnalysisResults

This is analogous to D86156 (which preserves "lossy" BFI in loop
passes). Lossy means that the analysis preserved may not be up to date
with regards to new blocks that are added in loop passes, but BPI will
not contain stale pointers to basic blocks that are deleted by the loop
passes.

This is achieved through BasicBlockCallbackVH in BPI, which calls
eraseBlock that updates the data structures in BPI whenever a basic
block is deleted.

This patch does not have any changes in the upstream pipeline, since
none of the loop passes in the pipeline use BPI currently.
However, since BPI wasn't previously preserved in loop passes, the loop
predication pass was invoking BPI *on the entire
function* every time it ran in an LPM.  This caused massive compile time
in our downstream LPM invocation which contained loop predication.

See updated test with an invocation of a loop-pipeline containing loop
predication and -debug-pass turned ON.

Reviewed-By: asbirlea, modimo
Differential Revision: https://reviews.llvm.org/D110438

Author: annamthomas

llvm/include/llvm/Analysis/LoopAnalysisManager.h

@@ -58,6 +58,7 @@ struct LoopStandardAnalysisResults {
   TargetLibraryInfo &TLI;
   TargetTransformInfo &TTI;
   BlockFrequencyInfo *BFI;
+  BranchProbabilityInfo *BPI;
   MemorySSA *MSSA;
 };
 

llvm/include/llvm/Transforms/Scalar/LoopPassManager.h

@@ -421,10 +421,12 @@ class FunctionToLoopPassAdaptor
   explicit FunctionToLoopPassAdaptor(std::unique_ptr<PassConceptT> Pass,
                                      bool UseMemorySSA = false,
                                      bool UseBlockFrequencyInfo = false,
+                                     bool UseBranchProbabilityInfo = false,
                                      bool LoopNestMode = false)
       : Pass(std::move(Pass)), LoopCanonicalizationFPM(),
         UseMemorySSA(UseMemorySSA),
         UseBlockFrequencyInfo(UseBlockFrequencyInfo),
+        UseBranchProbabilityInfo(UseBranchProbabilityInfo),
         LoopNestMode(LoopNestMode) {
     LoopCanonicalizationFPM.addPass(LoopSimplifyPass());
     LoopCanonicalizationFPM.addPass(LCSSAPass());
@@ -446,6 +448,7 @@ class FunctionToLoopPassAdaptor
 
   bool UseMemorySSA = false;
   bool UseBlockFrequencyInfo = false;
+  bool UseBranchProbabilityInfo = false;
   const bool LoopNestMode;
 };
 
@@ -457,13 +460,14 @@ template <typename LoopPassT>
 inline std::enable_if_t<is_detected<HasRunOnLoopT, LoopPassT>::value,
                         FunctionToLoopPassAdaptor>
 createFunctionToLoopPassAdaptor(LoopPassT &&Pass, bool UseMemorySSA = false,
-                                bool UseBlockFrequencyInfo = false) {
+                                bool UseBlockFrequencyInfo = false,
+                                bool UseBranchProbabilityInfo = false) {
   using PassModelT =
       detail::PassModel<Loop, LoopPassT, PreservedAnalyses, LoopAnalysisManager,
                         LoopStandardAnalysisResults &, LPMUpdater &>;
   return FunctionToLoopPassAdaptor(
       std::make_unique<PassModelT>(std::forward<LoopPassT>(Pass)), UseMemorySSA,
-      UseBlockFrequencyInfo, false);
+      UseBlockFrequencyInfo, UseBranchProbabilityInfo, false);
 }
 
 /// If \p Pass is a loop-nest pass, \p Pass will first be wrapped into a
@@ -472,24 +476,26 @@ template <typename LoopNestPassT>
 inline std::enable_if_t<!is_detected<HasRunOnLoopT, LoopNestPassT>::value,
                         FunctionToLoopPassAdaptor>
 createFunctionToLoopPassAdaptor(LoopNestPassT &&Pass, bool UseMemorySSA = false,
-                                bool UseBlockFrequencyInfo = false) {
+                                bool UseBlockFrequencyInfo = false,
+                                bool UseBranchProbabilityInfo = false) {
   LoopPassManager LPM;
   LPM.addPass(std::forward<LoopNestPassT>(Pass));
   using PassModelT =
       detail::PassModel<Loop, LoopPassManager, PreservedAnalyses,
                         LoopAnalysisManager, LoopStandardAnalysisResults &,
                         LPMUpdater &>;
   return FunctionToLoopPassAdaptor(std::make_unique<PassModelT>(std::move(LPM)),
-                                   UseMemorySSA, UseBlockFrequencyInfo, true);
+                                   UseMemorySSA, UseBlockFrequencyInfo,
+                                   UseBranchProbabilityInfo, true);
 }
 
 /// If \p Pass is an instance of \c LoopPassManager, the returned adaptor will
 /// be in loop-nest mode if the pass manager contains only loop-nest passes.
 template <>
 inline FunctionToLoopPassAdaptor
-createFunctionToLoopPassAdaptor<LoopPassManager>(LoopPassManager &&LPM,
-                                                 bool UseMemorySSA,
-                                                 bool UseBlockFrequencyInfo) {
+createFunctionToLoopPassAdaptor<LoopPassManager>(
+    LoopPassManager &&LPM, bool UseMemorySSA, bool UseBlockFrequencyInfo,
+    bool UseBranchProbabilityInfo) {
   // Check if LPM contains any loop pass and if it does not, returns an adaptor
   // in loop-nest mode.
   using PassModelT =
@@ -499,7 +505,7 @@ createFunctionToLoopPassAdaptor<LoopPassManager>(LoopPassManager &&LPM,
   bool LoopNestMode = (LPM.getNumLoopPasses() == 0);
   return FunctionToLoopPassAdaptor(std::make_unique<PassModelT>(std::move(LPM)),
                                    UseMemorySSA, UseBlockFrequencyInfo,
-                                   LoopNestMode);
+                                   UseBranchProbabilityInfo, LoopNestMode);
 }
 
 /// Pass for printing a loop's contents as textual IR.

llvm/lib/Passes/PassBuilder.cpp

@@ -1302,8 +1302,11 @@ Error PassBuilder::parseFunctionPass(FunctionPassManager &FPM,
       bool UseMemorySSA = (Name == "loop-mssa");
       bool UseBFI = llvm::any_of(
           InnerPipeline, [](auto Pipeline) { return Pipeline.Name == "licm"; });
+      bool UseBPI = llvm::any_of(InnerPipeline, [](auto Pipeline) {
+        return Pipeline.Name == "loop-predication";
+      });
       FPM.addPass(createFunctionToLoopPassAdaptor(std::move(LPM), UseMemorySSA,
-                                                  UseBFI));
+                                                  UseBFI, UseBPI));
       return Error::success();
     }
     if (auto Count = parseRepeatPassName(Name)) {

llvm/lib/Transforms/Scalar/LoopDistribute.cpp

@@ -1057,8 +1057,8 @@ PreservedAnalyses LoopDistributePass::run(Function &F,
   auto &LAM = AM.getResult<LoopAnalysisManagerFunctionProxy>(F).getManager();
   std::function<const LoopAccessInfo &(Loop &)> GetLAA =
       [&](Loop &L) -> const LoopAccessInfo & {
-    LoopStandardAnalysisResults AR = {AA,  AC,  DT,      LI,     SE,
-                                      TLI, TTI, nullptr, nullptr};
+    LoopStandardAnalysisResults AR = {AA,  AC,  DT,      LI,      SE,
+                                      TLI, TTI, nullptr, nullptr, nullptr};
     return LAM.getResult<LoopAccessAnalysis>(L, AR);
   };
 

llvm/lib/Transforms/Scalar/LoopLoadElimination.cpp

@@ -721,8 +721,8 @@ PreservedAnalyses LoopLoadEliminationPass::run(Function &F,
   auto &LAM = AM.getResult<LoopAnalysisManagerFunctionProxy>(F).getManager();
   bool Changed = eliminateLoadsAcrossLoops(
       F, LI, DT, BFI, PSI, &SE, &AC, [&](Loop &L) -> const LoopAccessInfo & {
-        LoopStandardAnalysisResults AR = {AA,  AC,  DT,      LI,  SE,
-                                          TLI, TTI, nullptr, nullptr};
+        LoopStandardAnalysisResults AR = {AA,  AC,  DT,      LI,      SE,
+                                          TLI, TTI, nullptr, nullptr, nullptr};
         return LAM.getResult<LoopAccessAnalysis>(L, AR);
       });
 

llvm/lib/Transforms/Scalar/LoopPassManager.cpp

@@ -10,6 +10,7 @@
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/Analysis/ScalarEvolutionAliasAnalysis.h"
@@ -207,6 +208,10 @@ PreservedAnalyses FunctionToLoopPassAdaptor::run(Function &F,
   BlockFrequencyInfo *BFI = UseBlockFrequencyInfo && F.hasProfileData()
                                 ? (&AM.getResult<BlockFrequencyAnalysis>(F))
                                 : nullptr;
+  BranchProbabilityInfo *BPI =
+      UseBranchProbabilityInfo && F.hasProfileData()
+          ? (&AM.getResult<BranchProbabilityAnalysis>(F))
+          : nullptr;
   LoopStandardAnalysisResults LAR = {AM.getResult<AAManager>(F),
                                      AM.getResult<AssumptionAnalysis>(F),
                                      AM.getResult<DominatorTreeAnalysis>(F),
@@ -215,6 +220,7 @@ PreservedAnalyses FunctionToLoopPassAdaptor::run(Function &F,
                                      AM.getResult<TargetLibraryAnalysis>(F),
                                      AM.getResult<TargetIRAnalysis>(F),
                                      BFI,
+                                     BPI,
                                      MSSA};
 
   // Setup the loop analysis manager from its proxy. It is important that
@@ -335,6 +341,8 @@ PreservedAnalyses FunctionToLoopPassAdaptor::run(Function &F,
   PA.preserve<ScalarEvolutionAnalysis>();
   if (UseBlockFrequencyInfo && F.hasProfileData())
     PA.preserve<BlockFrequencyAnalysis>();
+  if (UseBranchProbabilityInfo && F.hasProfileData())
+    PA.preserve<BranchProbabilityAnalysis>();
   if (UseMemorySSA)
     PA.preserve<MemorySSAAnalysis>();
   return PA;

llvm/lib/Transforms/Scalar/LoopPredication.cpp

@@ -367,14 +367,10 @@ PreservedAnalyses LoopPredicationPass::run(Loop &L, LoopAnalysisManager &AM,
                                            LoopStandardAnalysisResults &AR,
                                            LPMUpdater &U) {
   Function *F = L.getHeader()->getParent();
-  // For the new PM, we also can't use BranchProbabilityInfo as an analysis
-  // pass. Function analyses need to be preserved across loop transformations
-  // but BPI is not preserved, hence a newly built one is needed.
-  BranchProbabilityInfo BPI(*F, AR.LI, &AR.TLI, &AR.DT, nullptr);
   std::unique_ptr<MemorySSAUpdater> MSSAU;
   if (AR.MSSA)
     MSSAU = std::make_unique<MemorySSAUpdater>(AR.MSSA);
-  LoopPredication LP(&AR.AA, &AR.DT, &AR.SE, &AR.LI, &BPI,
+  LoopPredication LP(&AR.AA, &AR.DT, &AR.SE, &AR.LI, AR.BPI,
                      MSSAU ? MSSAU.get() : nullptr);
   if (!LP.runOnLoop(&L))
     return PreservedAnalyses::all();

llvm/lib/Transforms/Utils/LoopVersioning.cpp

@@ -357,8 +357,8 @@ PreservedAnalyses LoopVersioningPass::run(Function &F,
 
   auto &LAM = AM.getResult<LoopAnalysisManagerFunctionProxy>(F).getManager();
   auto GetLAA = [&](Loop &L) -> const LoopAccessInfo & {
-    LoopStandardAnalysisResults AR = {AA,  AC,  DT,      LI,  SE,
-                                      TLI, TTI, nullptr, nullptr};
+    LoopStandardAnalysisResults AR = {AA,  AC,  DT,      LI,      SE,
+                                      TLI, TTI, nullptr, nullptr, nullptr};
     return LAM.getResult<LoopAccessAnalysis>(L, AR);
   };
 

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -10557,8 +10557,8 @@ PreservedAnalyses LoopVectorizePass::run(Function &F,
     auto &LAM = AM.getResult<LoopAnalysisManagerFunctionProxy>(F).getManager();
     std::function<const LoopAccessInfo &(Loop &)> GetLAA =
         [&](Loop &L) -> const LoopAccessInfo & {
-      LoopStandardAnalysisResults AR = {AA,  AC,  DT,      LI,  SE,
-                                        TLI, TTI, nullptr, nullptr};
+      LoopStandardAnalysisResults AR = {AA,  AC,  DT,      LI,      SE,
+                                        TLI, TTI, nullptr, nullptr, nullptr};
       return LAM.getResult<LoopAccessAnalysis>(L, AR);
     };
     auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);

llvm/test/Transforms/LoopPredication/preserve-bpi.ll

@@ -1,24 +1,19 @@
-; RUN: opt -mtriple=x86_64 -passes='require<scalar-evolution>,require<branch-prob>,loop-mssa(loop-predication,licm,simple-loop-unswitch<nontrivial>,loop-simplifycfg)' -debug-pass-manager -debug-only=branch-prob -S < %s 2>&1 | FileCheck %s
+; RUN: opt -mtriple=x86_64 -passes='loop-mssa(loop-predication,licm,simple-loop-unswitch<nontrivial>,loop-simplifycfg)' -debug-pass-manager -debug-only=branch-prob -S < %s 2>&1 | FileCheck %s
 
 ; REQUIRES: asserts
 
+; This test is to solely check that we do not run BPI every single time loop
+; predication is invoked (since BPI is preserved as part of
+; LoopStandardAnalysisResults).
 declare void @llvm.experimental.guard(i1, ...)
 
-; CHECK: Running pass: RequireAnalysisPass<{{.*}}BranchProbabilityAnalysis
-; CHECK-NEXT: Running analysis: BranchProbabilityAnalysis on unsigned_loop_0_to_n_ult_check
-; CHECK-NEXT: Running analysis: PostDominatorTreeAnalysis on unsigned_loop_0_to_n_ult_check
-; CHECK-NEXT: ---- Branch Probability Info : unsigned_loop_0_to_n_ult_check ----
-; CHECK: Running pass: LoopSimplifyPass on unsigned_loop_0_to_n_ult_check
 ; CHECK: Running pass: LoopPredicationPass on Loop at depth 1
-; CHECK-NEXT: ---- Branch Probability Info : unsigned_loop_0_to_n_ult_check ----
-; CHECK: Running pass: LICMPass on Loop at depth 1
+; CHECK-NEXT: Running pass: LICMPass on Loop at depth 1
 ; CHECK-NEXT: Running pass: SimpleLoopUnswitchPass on Loop at depth 1
 ; CHECK-NEXT: Running pass: LoopPredicationPass on Loop at depth 1
-; CHECK-NEXT: ---- Branch Probability Info : unsigned_loop_0_to_n_ult_check ----
-; CHECK: Running pass: LICMPass on Loop at depth 1
+; CHECK-NEXT: Running pass: LICMPass on Loop at depth 1
 ; CHECK-NEXT: Running pass: SimpleLoopUnswitchPass on Loop at depth 1
 ; CHECK-NEXT: Running pass: LoopSimplifyCFGPass on Loop at depth 1
-; CHECK-NEXT: Invalidating analysis: BranchProbabilityAnalysis on unsigned_loop_0_to_n_ult_check
 
 
 define i32 @unsigned_loop_0_to_n_ult_check(i32* %array, i32 %length, i32 %n) {