[BPI] Keep BPI available in loop passes through LoopStandardAnalysisResults

[anna]

Summary

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.

[modimo]

LGTM pending the added test case demonstrating that BPI is correctly preserved across loop passes. Also from D86156 we found there was a significant compiler speed up using LazyBlockFrequencyInfo for the passes that needed it. This may also help for LoopPredication's usage of BPI.

[asbirlea]

Changes LGTM. Please add the testcase mentioned in the description before pushing.

[anna]

Thank you for the reviews. I will precommit the test case and update this diff. The point to note is that since Loop predication didn't preserve BPI, we were recalculating it every time for the NPM. We are primarily interested in the NPM, so I'll be adding the testcase for this along with -debug-only=branch-prob, which shows that before this change, we calculate BPI every time we call Loop predication.

[anna]

updated precommited testcase and BPI invocation removed from LoopPredication.

[anna]

The removed code in LoopPredication (which removes BPI calculation per Loop predication invocation) shows the testcase difference. I was hoping to keep the loop predication change separate to save some downstream merge headaches, but couldn't show a testcase for the standalone original change in this review.

[modimo]

In D110438#3027462, @anna wrote:

The removed code in LoopPredication (which removes BPI calculation per Loop predication invocation) shows the testcase difference. I was hoping to keep the loop predication change separate to save some downstream merge headaches, but couldn't show a testcase for the standalone original change in this review.

I think it makes sense as part of this change given LoopStandardAnalysisResults now stores BPI, so calculating it separately in LoopPredication is a waste.

Make sure to update the description before pushing, thanks!

[anna]

We need to explicitly check for loop-predication in PassBuilder, so that UseBranchProbabilityInfo is set true.
(Without this change profitability.ll test will fail since LPM invocation by default has UseBranchProbabilityInfo=false).

This is the only change in the diff and analogous to the UseBFI = any_pass = LICM check just above that.

[anna]

llvm/lib/Passes/PassBuilder.cpp

1305

Just for ease of review compared to previous LGTM'ed version: this is the (only) update done vs previous diff.

[asbirlea]

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

1

Does this still need the requires before loop-mssa?

[anna]

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

1

We do not. Thanks. I'll update and land the patch.

llvm/include/llvm/Analysis/LoopAnalysisManager.h

struct LoopStandardAnalysisResults {
   AAResults &AA;
   AssumptionCache ∾
   DominatorTree &DT;
   LoopInfo &LI;
   ScalarEvolution &SE;
   TargetLibraryInfo &TLI;
   TargetTransformInfo &TTI;
   BlockFrequencyInfo *BFI;
+  BranchProbabilityInfo *BPI;
   MemorySSA *MSSA;
 };
 
 /// Extern template declaration for the analysis set for this IR unit.
 extern template class AllAnalysesOn<Loop>;
 
 extern template class AnalysisManager<Loop, LoopStandardAnalysisResults &>;
 /// The loop analysis manager.

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

 public:
   using PassConceptT =
       detail::PassConcept<Loop, LoopAnalysisManager,
                           LoopStandardAnalysisResults &, LPMUpdater &>;
 
   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());
   }
 
   /// Runs the loop passes across every loop in the function.
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
   void printPipeline(raw_ostream &OS,
                      function_ref<StringRef(StringRef)> MapClassName2PassName);
 
   static bool isRequired() { return true; }
 
   bool isLoopNestMode() const { return LoopNestMode; }
 
 private:
   std::unique_ptr<PassConceptT> Pass;
 
   FunctionPassManager LoopCanonicalizationFPM;
 
   bool UseMemorySSA = false;
   bool UseBlockFrequencyInfo = false;
+  bool UseBranchProbabilityInfo = false;
   const bool LoopNestMode;
 };
 
 /// A function to deduce a loop pass type and wrap it in the templated
 /// adaptor.
 ///
 /// If \p Pass is a loop pass, the returned adaptor will be in loop mode.
 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
 /// \c LoopPassManager and the returned adaptor will be in loop-nest mode.
 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 =
       detail::PassModel<Loop, LoopPassManager, PreservedAnalyses,
                         LoopAnalysisManager, LoopStandardAnalysisResults &,
                         LPMUpdater &>;
   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.
 class PrintLoopPass : public PassInfoMixin<PrintLoopPass> {
   raw_ostream &OS;
   std::string Banner;
 
 public:

llvm/lib/Passes/PassBuilder.cpp

if (!InnerPipeline.empty()) {
     if (Name == "loop" || Name == "loop-mssa") {
       LoopPassManager LPM;
       if (auto Err = parseLoopPassPipeline(LPM, InnerPipeline))
         return Err;
       // Add the nested pass manager with the appropriate adaptor.
       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)) {
       FunctionPassManager NestedFPM;
       if (auto Err = parseFunctionPassPipeline(NestedFPM, InnerPipeline))
         return Err;
       FPM.addPass(createRepeatedPass(*Count, std::move(NestedFPM)));
       return Error::success();

llvm/lib/Transforms/Scalar/LoopDistribute.cpp

PreservedAnalyses LoopDistributePass::run(Function &F,
   auto &AA = AM.getResult<AAManager>(F);
   auto &AC = AM.getResult<AssumptionAnalysis>(F);
   auto &TTI = AM.getResult<TargetIRAnalysis>(F);
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(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);
   };
 
   bool Changed = runImpl(F, &LI, &DT, &SE, &ORE, GetLAA);
   if (!Changed)
     return PreservedAnalyses::all();
   PreservedAnalyses PA;
   PA.preserve<LoopAnalysis>();

llvm/lib/Transforms/Scalar/LoopLoadElimination.cpp

PreservedAnalyses LoopLoadEliminationPass::run(Function &F,
   auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
   auto *PSI = MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
   auto *BFI = (PSI && PSI->hasProfileSummary()) ?
       &AM.getResult<BlockFrequencyAnalysis>(F) : nullptr;
 
   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);
       });
 
   if (!Changed)
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
   return PA;
 }

llvm/lib/Transforms/Scalar/LoopPassManager.cpp

 //===- LoopPassManager.cpp - Loop pass management -------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/LoopPassManager.h"
 #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"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/TimeProfiler.h"
 
 using namespace llvm;
if (LI.empty())
     return PA;
 
   // Get the analysis results needed by loop passes.
   MemorySSA *MSSA =
       UseMemorySSA ? (&AM.getResult<MemorySSAAnalysis>(F).getMSSA()) : nullptr;
   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),
                                      AM.getResult<LoopAnalysis>(F),
                                      AM.getResult<ScalarEvolutionAnalysis>(F),
                                      AM.getResult<TargetLibraryAnalysis>(F),
                                      AM.getResult<TargetIRAnalysis>(F),
                                      BFI,
+                                     BPI,
                                      MSSA};
 
   // Setup the loop analysis manager from its proxy. It is important that
   // this is only done when there are loops to process and we have built the
   // LoopStandardAnalysisResults object. The loop analyses cached in this
   // manager have access to those analysis results and so it must invalidate
   // itself when they go away.
   auto &LAMFP = AM.getResult<LoopAnalysisManagerFunctionProxy>(F);
#endif
   PA.preserveSet<AllAnalysesOn<Loop>>();
   PA.preserve<LoopAnalysisManagerFunctionProxy>();
   // We also preserve the set of standard analyses.
   PA.preserve<DominatorTreeAnalysis>();
   PA.preserve<LoopAnalysis>();
   PA.preserve<ScalarEvolutionAnalysis>();
   if (UseBlockFrequencyInfo && F.hasProfileData())
     PA.preserve<BlockFrequencyAnalysis>();
+  if (UseBranchProbabilityInfo && F.hasProfileData())
+    PA.preserve<BranchProbabilityAnalysis>();
   if (UseMemorySSA)
     PA.preserve<MemorySSAAnalysis>();
   return PA;
 }
 
 PrintLoopPass::PrintLoopPass() : OS(dbgs()) {}
 PrintLoopPass::PrintLoopPass(raw_ostream &OS, const std::string &Banner)
     : OS(OS), Banner(Banner) {}
 
 PreservedAnalyses PrintLoopPass::run(Loop &L, LoopAnalysisManager &,
                                      LoopStandardAnalysisResults &,
                                      LPMUpdater &) {
   printLoop(L, OS, Banner);
   return PreservedAnalyses::all();
 }

llvm/lib/Transforms/Scalar/LoopPredication.cpp

 Pass *llvm::createLoopPredicationPass() {
   return new LoopPredicationLegacyPass();
 }
 
 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();
 
   auto PA = getLoopPassPreservedAnalyses();
   if (AR.MSSA)
     PA.preserve<MemorySSAAnalysis>();
   return PA;

llvm/lib/Transforms/Utils/LoopVersioning.cpp

PreservedAnalyses LoopVersioningPass::run(Function &F,
   auto &TTI = AM.getResult<TargetIRAnalysis>(F);
   auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto &AA = AM.getResult<AAManager>(F);
   auto &AC = AM.getResult<AssumptionAnalysis>(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);
   };
 
   if (runImpl(&LI, GetLAA, &DT, &SE))
     return PreservedAnalyses::none();
   return PreservedAnalyses::all();
 }
 } // namespace llvm

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

PreservedAnalyses LoopVectorizePass::run(Function &F,
     auto &AA = AM.getResult<AAManager>(F);
     auto &AC = AM.getResult<AssumptionAnalysis>(F);
     auto &DB = AM.getResult<DemandedBitsAnalysis>(F);
     auto &ORE = AM.getResult<OptimizationRemarkEmitterAnalysis>(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);
     ProfileSummaryInfo *PSI =
         MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
     LoopVectorizeResult Result =
         runImpl(F, SE, LI, TTI, DT, BFI, &TLI, DB, AA, AC, GetLAA, ORE, PSI);
     if (!Result.MadeAnyChange)

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

-; 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) {
 entry:
   %tmp5 = icmp eq i32 %n, 0
   br i1 %tmp5, label %exit, label %loop.preheader
 
 loop.preheader:                                   ; preds = %entry