[New PM][PassInstrumentation] enhancing PassInstrumentation with PassManager tracking

[fedor.sergeev]

Summary

In order to implement -time-passes functionality we need to be able to pause
PassManager's own timer when starting execution (and timer) for nested passes.

For that we need to be able to track which PassManager is currently being executed.

Adding startPassManager/endPassManager points to PassInstrumentation and call them
in all the pass managers/pass adaptors. Also adding startIteration point for
repeating adaptors (e.g RepeatPass) since that was pretty much the only
remaining piece of nontrivial pass execution tracking we were lacking
(albeit not needed for -time-passes).

Introducing PassExecutionLocation which is a wrapper for PassExecutionCounter
to provide a stable interface to be able to extend it later with extra information
(e.g. pass execution stack) if needed.
All the callbacks now take PassExecutionLocation.

Majority of the changes are rather mechanical ones.

[fedor.sergeev]

rebased and somewhat simplified by removing PassManager tracking inside the PassInstrumentation object
itself. It is left for discretion of the callbacks implementing particular instrumentations.
Subsequent -time-passes implementation will use it this way.

[fedor.sergeev]

rebased

[fedor.sergeev]

adding unittest changes, making unittest essentially the same as shown in D51973

[fedor.sergeev]

rebased

[fedor.sergeev]

adding const to PassManager parameter of callbacks.

[fedor.sergeev]

rebased

[fedor.sergeev]

rebased

[fedor.sergeev]

rebased

[philip.pfaffe]

I don't fully understand the rationale behind this. How is this different compared to a before/after instrumentation matching the PassManagers?

[fedor.sergeev]

In D51632#1239228, @philip.pfaffe wrote:

I don't fully understand the rationale behind this. How is this different compared to a before/after instrumentation matching the PassManagers?

To be honest, the rationale is nearly dead.
It was useful for -time-passes when we were passing instance pointers and were doing nontrivial things with PM timers.
Right now we just ignore pass managers there.
So the only reason for these interface existence is that I would like to be able to move DebugLogging functionality to PassInstrumentation
and start/finish points are the only places in code where we have DebugLogging that is not yet covered by PassInstrumentation.

That hints these points might be useful for nontrivial instrumentations in future.
I can defer this patch until I prepare DebugLogging replacement patch.

[fedor.sergeev]

changing status to reflect that there is no active work on this patch atm.

include/llvm/Analysis/CGSCCPassManager.h

public:
   ModuleToPostOrderCGSCCPassAdaptor &
   operator=(ModuleToPostOrderCGSCCPassAdaptor RHS) {
     swap(*this, RHS);
     return *this;
   }
 
   /// Runs the CGSCC pass across every SCC in the module.
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
+    // Request PassInstrumentation from analysis manager, will use it to run
+    // instrumenting callbacks for the passes later.
+    PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(M);
+    PI.startPassManager(*this, M);
+
     // Setup the CGSCC analysis manager from its proxy.
     CGSCCAnalysisManager &CGAM =
         AM.getResult<CGSCCAnalysisManagerModuleProxy>(M).getManager();
 
     // Get the call graph for this module.
     LazyCallGraph &CG = AM.getResult<LazyCallGraphAnalysis>(M);
 
     // We keep worklists to allow us to push more work onto the pass manager as
     // the passes are run.
     SmallPriorityWorklist<LazyCallGraph::RefSCC *, 1> RCWorklist;
     SmallPriorityWorklist<LazyCallGraph::SCC *, 1> CWorklist;
 
     // Keep sets for invalidated SCCs and RefSCCs that should be skipped when
     // iterating off the worklists.
     SmallPtrSet<LazyCallGraph::RefSCC *, 4> InvalidRefSCCSet;
     SmallPtrSet<LazyCallGraph::SCC *, 4> InvalidSCCSet;
 
     SmallDenseSet<std::pair<LazyCallGraph::Node *, LazyCallGraph::SCC *>, 4>
         InlinedInternalEdges;
 
     CGSCCUpdateResult UR = {RCWorklist,          CWorklist, InvalidRefSCCSet,
                             InvalidSCCSet,       nullptr,   nullptr,
                             InlinedInternalEdges};
 
-    // Request PassInstrumentation from analysis manager, will use it to run
-    // instrumenting callbacks for the passes later.
-    PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(M);
-
     PreservedAnalyses PA = PreservedAnalyses::all();
     CG.buildRefSCCs();
     for (auto RCI = CG.postorder_ref_scc_begin(),
               RCE = CG.postorder_ref_scc_end();
          RCI != RCE;) {
       assert(RCWorklist.empty() &&
              "Should always start with an empty RefSCC worklist");
       // The postorder_ref_sccs range we are walking is lazily constructed, so
PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
     }
 
     // By definition we preserve the call garph, all SCC analyses, and the
     // analysis proxies by handling them above and in any nested pass managers.
     PA.preserveSet<AllAnalysesOn<LazyCallGraph::SCC>>();
     PA.preserve<LazyCallGraphAnalysis>();
     PA.preserve<CGSCCAnalysisManagerModuleProxy>();
     PA.preserve<FunctionAnalysisManagerModuleProxy>();
+
+    PI.finishPassManager(*this, M);
     return PA;
   }
 
 private:
   CGSCCPassT Pass;
 };
 
 /// A function to deduce a function pass type and wrap it in the
public:
   CGSCCToFunctionPassAdaptor &operator=(CGSCCToFunctionPassAdaptor RHS) {
     swap(*this, RHS);
     return *this;
   }
 
   /// Runs the function pass across every function in the module.
   PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
                         LazyCallGraph &CG, CGSCCUpdateResult &UR) {
+    // Request PassInstrumentation from analysis manager, will use it to run
+    // instrumenting callbacks for the passes later.
+    PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(C, CG);
+    PI.startPassManager(*this, C);
+
     // Setup the function analysis manager from its proxy.
     FunctionAnalysisManager &FAM =
         AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
 
     SmallVector<LazyCallGraph::Node *, 4> Nodes;
     for (LazyCallGraph::Node &N : C)
       Nodes.push_back(&N);
 
for (LazyCallGraph::Node *N : Nodes) {
       // Skip nodes from other SCCs. These may have been split out during
       // processing. We'll eventually visit those SCCs and pick up the nodes
       // there.
       if (CG.lookupSCC(*N) != CurrentC)
         continue;
 
       Function &F = N->getFunction();
 
-      PassInstrumentation PI = FAM.getResult<PassInstrumentationAnalysis>(F);
       if (!PI.runBeforePass<Function>(Pass, F))
         continue;
 
       PreservedAnalyses PassPA = Pass.run(F, FAM);
 
       PI.runAfterPass<Function>(Pass, F);
 
       // We know that the function pass couldn't have invalidated any other
PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
     // proxy, but that's OK because we've taken care to invalidate analyses in
     // the function analysis manager incrementally above.
     PA.preserveSet<AllAnalysesOn<Function>>();
     PA.preserve<FunctionAnalysisManagerCGSCCProxy>();
 
     // We've also ensured that we updated the call graph along the way.
     PA.preserve<LazyCallGraphAnalysis>();
 
+    PI.finishPassManager(*this, C);
     return PA;
   }
 
 private:
   FunctionPassT Pass;
 };
 
 /// A function to deduce a function pass type and wrap it in the
public:
   explicit DevirtSCCRepeatedPass(PassT Pass, int MaxIterations)
       : Pass(std::move(Pass)), MaxIterations(MaxIterations) {}
 
   /// Runs the wrapped pass up to \c MaxIterations on the SCC, iterating
   /// whenever an indirect call is refined.
   PreservedAnalyses run(LazyCallGraph::SCC &InitialC, CGSCCAnalysisManager &AM,
                         LazyCallGraph &CG, CGSCCUpdateResult &UR) {
     PreservedAnalyses PA = PreservedAnalyses::all();
+    // Request PassInstrumentation from analysis manager, will use it to run
+    // instrumenting callbacks for the passes later.
     PassInstrumentation PI =
         AM.getResult<PassInstrumentationAnalysis>(InitialC, CG);
+    PI.startPassManager(*this, InitialC);
 
     // The SCC may be refined while we are running passes over it, so set up
     // a pointer that we can update.
     LazyCallGraph::SCC *C = &InitialC;
 
     // Collect value handles for all of the indirect call sites.
     SmallVector<WeakTrackingVH, 8> CallHandles;
 
auto ScanSCC = [](LazyCallGraph::SCC &C,
 
       return CallCounts;
     };
 
     // Populate the initial call handles and get the initial call counts.
     auto CallCounts = ScanSCC(*C, CallHandles);
 
     for (int Iteration = 0;; ++Iteration) {
+      PI.startIteration(*this, Iteration);
 
       if (!PI.runBeforePass<LazyCallGraph::SCC>(Pass, *C))
         continue;
 
       PreservedAnalyses PassPA = Pass.run(*C, AM, CG, UR);
 
       PI.runAfterPass<LazyCallGraph::SCC>(Pass, *C);
 
for (int Iteration = 0;; ++Iteration) {
       CallCounts = std::move(NewCallCounts);
 
       // Update the analysis manager with each run and intersect the total set
       // of preserved analyses so we're ready to iterate.
       AM.invalidate(*C, PassPA);
       PA.intersect(std::move(PassPA));
     }
 
+    PI.finishPassManager(*this, InitialC);
+
     // Note that we don't add any preserved entries here unlike a more normal
     // "pass manager" because we only handle invalidation *between* iterations,
     // not after the last iteration.
     return PA;
   }
 
 private:
   PassT Pass;

include/llvm/IR/PassInstrumentation.h

 class PassInstrumentationCallbacks {
 public:
   // Before/After callbacks accept IRUnits, so they need to take them
   // as pointers, wrapped with llvm::Any
   using BeforePassFunc = bool(StringRef, Any);
   using AfterPassFunc = void(StringRef, Any);
   using BeforeAnalysisFunc = void(StringRef, Any);
   using AfterAnalysisFunc = void(StringRef, Any);
+  using PassManagerFunc = void(StringRef, Any);
+  using PMIterationFunc = void(StringRef, int);
 
 public:
   PassInstrumentationCallbacks() {}
 
   /// Copying PassInstrumentationCallbacks is not intended.
   PassInstrumentationCallbacks(const PassInstrumentationCallbacks &) = delete;
   void operator=(const PassInstrumentationCallbacks &) = delete;
 
void registerBeforeAnalysisCallback(CallableT C) {
     BeforeAnalysisCallbacks.emplace_back(std::move(C));
   }
 
   template <typename CallableT>
   void registerAfterAnalysisCallback(CallableT C) {
     AfterAnalysisCallbacks.emplace_back(std::move(C));
   }
 
+  template <typename CallableT>
+  void registerStartPassManagerCallback(CallableT C) {
+    StartPMCallbacks.emplace_back(std::move(C));
+  }
+
+  template <typename CallableT>
+  void registerFinishPassManagerCallback(CallableT C) {
+    FinishPMCallbacks.emplace_back(std::move(C));
+  }
+
+  void
+  registerStartIterationCallback(llvm::unique_function<PMIterationFunc> &&C) {
+    StartIterationCallbacks.emplace_back(std::move(C));
+  }
+
 private:
   friend class PassInstrumentation;
 
   SmallVector<llvm::unique_function<BeforePassFunc>, 4> BeforePassCallbacks;
   SmallVector<llvm::unique_function<AfterPassFunc>, 4> AfterPassCallbacks;
   SmallVector<llvm::unique_function<BeforeAnalysisFunc>, 4>
       BeforeAnalysisCallbacks;
   SmallVector<llvm::unique_function<AfterAnalysisFunc>, 4>
       AfterAnalysisCallbacks;
+  SmallVector<llvm::unique_function<PassManagerFunc>, 2> StartPMCallbacks;
+  SmallVector<llvm::unique_function<PassManagerFunc>, 2> FinishPMCallbacks;
+  SmallVector<llvm::unique_function<PMIterationFunc>, 2>
+      StartIterationCallbacks;
 };
 
 /// This class provides instrumentation entry points for the Pass Manager,
 /// doing calls to callbacks registered in PassInstrumentationCallbacks.
 class PassInstrumentation {
   PassInstrumentationCallbacks *Callbacks;
 
 public:
public:
   /// that has just been executed and constant reference to IR it operated on.
   template <typename IRUnitT, typename PassT>
   void runAfterAnalysis(const PassT &Analysis, const IRUnitT &IR) const {
     if (Callbacks)
       for (auto &C : Callbacks->AfterAnalysisCallbacks)
         C(Analysis.name(), llvm::Any(&IR));
   }
 
+  /// startPassManager instrumentation point - takes \p PassManager instance
+  /// that is about to start its operations and constant reference to IR
+  /// it operates on.
+  template <typename IRUnitT, typename PassT>
+  void startPassManager(const PassT &PassManager, const IRUnitT &IR) const {
+    if (Callbacks)
+      for (auto &C : Callbacks->StartPMCallbacks)
+        C(PassManager.name(), llvm::Any(&IR));
+  }
+
+  /// finishPassManager instrumentation point - takes \p PassManager instance
+  /// that is about to finish its operations, as well as constant reference
+  /// to IR it operates on.
+  template <typename IRUnitT, typename PassT>
+  void finishPassManager(const PassT &PassManager, const IRUnitT &IR) const {
+    if (Callbacks)
+      for (auto &C : Callbacks->FinishPMCallbacks)
+        C(PassManager.name(), llvm::Any(&IR));
+  }
+
+  /// startIteration instrumentation point - takes \p PassManager instance
+  /// that starts a new iteration through passes and constant reference
+  /// to IR it operates on.
+  template <typename PassT>
+  void startIteration(const PassT &PassManager, int Iteration) const {
+    if (Callbacks)
+      for (auto &C : Callbacks->StartIterationCallbacks)
+        C(PassManager.name(), Iteration);
+  }
+
   /// Handle invalidation from the pass manager when PassInstrumentation
   /// is used as the result of PassInstrumentationAnalysis.
   ///
   /// On attempt to invalidate just return false. There is nothing to become
   /// invalid here.
   template <typename IRUnitT, typename... ExtraArgsT>
   bool invalidate(IRUnitT &, const class llvm::PreservedAnalyses &,
                   ExtraArgsT...) {
     return false;
   }
 };
 
 } // namespace llvm
 
 #endif

include/llvm/IR/PassManager.h

PreservedAnalyses run(IRUnitT &IR, AnalysisManagerT &AM,
 
     // Request PassInstrumentation from analysis manager, will use it to run
     // instrumenting callbacks for the passes later.
     // Here we use std::tuple wrapper over getResult which helps to extract
     // AnalysisManager's arguments out of the whole ExtraArgs set.
     PassInstrumentation PI =
         detail::getAnalysisResult<PassInstrumentationAnalysis>(
             AM, IR, std::tuple<ExtraArgTs...>(ExtraArgs...));
-
+    PI.startPassManager(*this, IR);
     if (DebugLogging)
       dbgs() << "Starting " << getTypeName<IRUnitT>() << " pass manager run.\n";
 
     for (unsigned Idx = 0, Size = Passes.size(); Idx != Size; ++Idx) {
       auto *P = Passes[Idx].get();
       if (DebugLogging)
         dbgs() << "Running pass: " << P->name() << " on " << IR.getName()
                << "\n";
PreservedAnalyses run(IRUnitT &IR, AnalysisManagerT &AM,
     }
 
     // Invalidation was handled after each pass in the above loop for the
     // current unit of IR. Therefore, the remaining analysis results in the
     // AnalysisManager are preserved. We mark this with a set so that we don't
     // need to inspect each one individually.
     PA.preserveSet<AllAnalysesOn<IRUnitT>>();
 
+    PI.finishPassManager(*this, IR);
     if (DebugLogging)
       dbgs() << "Finished " << getTypeName<IRUnitT>() << " pass manager run.\n";
 
     return PA;
   }
 
   template <typename PassT> void addPass(PassT Pass) {
     using PassModelT =

 class ModuleToFunctionPassAdaptor
     : public PassInfoMixin<ModuleToFunctionPassAdaptor<FunctionPassT>> {
 public:
   explicit ModuleToFunctionPassAdaptor(FunctionPassT Pass)
       : Pass(std::move(Pass)) {}
 
   /// Runs the function pass across every function in the module.
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
-    FunctionAnalysisManager &FAM =
-        AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
-
     // Request PassInstrumentation from analysis manager, will use it to run
     // instrumenting callbacks for the passes later.
     PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(M);
+    PI.startPassManager(*this, M);
+
+    FunctionAnalysisManager &FAM =
+        AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
 
     PreservedAnalyses PA = PreservedAnalyses::all();
     for (Function &F : M) {
       if (F.isDeclaration())
         continue;
 
       // Check the PassInstrumentation's BeforePass callbacks before running the
       // pass, skip its execution completely if asked to (callback returns
PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
 
     // The FunctionAnalysisManagerModuleProxy is preserved because (we assume)
     // the function passes we ran didn't add or remove any functions.
     //
     // We also preserve all analyses on Functions, because we did all the
     // invalidation we needed to do above.
     PA.preserveSet<AllAnalysesOn<Function>>();
     PA.preserve<FunctionAnalysisManagerModuleProxy>();
+
+    PI.finishPassManager(*this, M);
     return PA;
   }
 
 private:
   FunctionPassT Pass;
 };
 
 /// A function to deduce a function pass type and wrap it in the
PreservedAnalyses run(IRUnitT &IR, AnalysisManagerT &AM, Ts &&... Args) {
     // Request PassInstrumentation from analysis manager, will use it to run
     // instrumenting callbacks for the passes later.
     // Here we use std::tuple wrapper over getResult which helps to extract
     // AnalysisManager's arguments out of the whole Args set.
     PassInstrumentation PI =
         detail::getAnalysisResult<PassInstrumentationAnalysis>(
             AM, IR, std::tuple<Ts...>(Args...));
 
+    PI.startPassManager(*this, IR);
     auto PA = PreservedAnalyses::all();
     for (int i = 0; i < Count; ++i) {
+      // Tell PassInstrumentation that we started a new iteration.
+      PI.startIteration(*this, i);
+
       // Check the PassInstrumentation's BeforePass callbacks before running the
       // pass, skip its execution completely if asked to (callback returns
       // false).
       if (!PI.runBeforePass<IRUnitT>(P, IR))
         continue;
       PA.intersect(P.run(IR, AM, std::forward<Ts>(Args)...));
       PI.runAfterPass(P, IR);
     }
+    PI.finishPassManager(*this, IR);
     return PA;
   }
 
 private:
   int Count;
   PassT P;
 };
 
 template <typename PassT>
 RepeatedPass<PassT> createRepeatedPass(int Count, PassT P) {
   return RepeatedPass<PassT>(Count, std::move(P));
 }
 
 } // end namespace llvm
 
 #endif // LLVM_IR_PASSMANAGER_H

include/llvm/Transforms/Scalar/LoopPassManager.h

public:
 
   /// Runs the loop passes across every loop in the function.
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM) {
     // Before we even compute any loop analyses, first run a miniature function
     // pass pipeline to put loops into their canonical form. Note that we can
     // directly build up function analyses after this as the function pass
     // manager handles all the invalidation at that layer.
     PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(F);
+    PI.startPassManager(*this, F);
 
     PreservedAnalyses PA = PreservedAnalyses::all();
     // Check the PassInstrumentation's BeforePass callbacks before running the
     // canonicalization pipeline.
     if (PI.runBeforePass<Function>(LoopCanonicalizationFPM, F)) {
       PA = LoopCanonicalizationFPM.run(F, AM);
       PI.runAfterPass<Function>(LoopCanonicalizationFPM, F);
     }
#endif
     if (EnableMSSALoopDependency)
       PA.preserve<MemorySSAAnalysis>();
     // FIXME: What we really want to do here is preserve an AA category, but
     // that concept doesn't exist yet.
     PA.preserve<AAManager>();
     PA.preserve<BasicAA>();
     PA.preserve<GlobalsAA>();
     PA.preserve<SCEVAA>();
+
+    PI.finishPassManager(*this, F);
     return PA;
   }
 
 private:
   LoopPassT Pass;
 
   FunctionPassManager LoopCanonicalizationFPM;
 };

lib/Analysis/CGSCCPassManager.cpp

 PassManager<LazyCallGraph::SCC, CGSCCAnalysisManager, LazyCallGraph &,
             CGSCCUpdateResult &>::run(LazyCallGraph::SCC &InitialC,
                                       CGSCCAnalysisManager &AM,
                                       LazyCallGraph &G, CGSCCUpdateResult &UR) {
   // Request PassInstrumentation from analysis manager, will use it to run
   // instrumenting callbacks for the passes later.
   PassInstrumentation PI =
       AM.getResult<PassInstrumentationAnalysis>(InitialC, G);
+  PI.startPassManager(*this, InitialC);
 
   PreservedAnalyses PA = PreservedAnalyses::all();
 
   if (DebugLogging)
     dbgs() << "Starting CGSCC pass manager run.\n";
 
   // The SCC may be refined while we are running passes over it, so set up
   // a pointer that we can update.
PassManager<LazyCallGraph::SCC, CGSCCAnalysisManager, LazyCallGraph &,
   }
 
   // Invalidation was handled after each pass in the above loop for the current
   // SCC. Therefore, the remaining analysis results in the AnalysisManager are
   // preserved. We mark this with a set so that we don't need to inspect each
   // one individually.
   PA.preserveSet<AllAnalysesOn<LazyCallGraph::SCC>>();
 
+  PI.finishPassManager(*this, InitialC);
   if (DebugLogging)
     dbgs() << "Finished CGSCC pass manager run.\n";
 
   return PA;
 }
 
 bool CGSCCAnalysisManagerModuleProxy::Result::invalidate(
     Module &M, const PreservedAnalyses &PA,

lib/Transforms/Scalar/LoopPassManager.cpp

PassManager<Loop, LoopAnalysisManager, LoopStandardAnalysisResults &,
   PreservedAnalyses PA = PreservedAnalyses::all();
 
   if (DebugLogging)
     dbgs() << "Starting Loop pass manager run.\n";
 
   // Request PassInstrumentation from analysis manager, will use it to run
   // instrumenting callbacks for the passes later.
   PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(L, AR);
+  PI.startPassManager(*this, L);
   for (auto &Pass : Passes) {
     if (DebugLogging)
       dbgs() << "Running pass: " << Pass->name() << " on " << L;
 
     // Check the PassInstrumentation's BeforePass callbacks before running the
     // pass, skip its execution completely if asked to (callback returns false).
     if (!PI.runBeforePass<Loop>(*Pass, L))
       continue;
#endif
   // analysis results shouldn't be impacted by runs over this loop. Therefore,
   // the remaining analysis results in the AnalysisManager are preserved. We
   // mark this with a set so that we don't need to inspect each one
   // individually.
   // FIXME: This isn't correct! This loop and all nested loops' analyses should
   // be preserved, but unrolling should invalidate the parent loop's analyses.
   PA.preserveSet<AllAnalysesOn<Loop>>();
 
+  PI.finishPassManager(*this, L);
+
   if (DebugLogging)
     dbgs() << "Finished Loop pass manager run.\n";
 
   return PA;
 }
 }
 
 PrintLoopPass::PrintLoopPass() : OS(dbgs()) {}

unittests/IR/PassBuilderCallbacksTest.cpp

 }
 
 struct MockPassInstrumentationCallbacks {
   PassInstrumentationCallbacks Callbacks;
 
   MockPassInstrumentationCallbacks() {
     ON_CALL(*this, runBeforePass(_, _)).WillByDefault(Return(true));
   }
+
   MOCK_METHOD2(runBeforePass, bool(StringRef PassID, llvm::Any));
   MOCK_METHOD2(runAfterPass, void(StringRef PassID, llvm::Any));
   MOCK_METHOD2(runBeforeAnalysis, void(StringRef PassID, llvm::Any));
   MOCK_METHOD2(runAfterAnalysis, void(StringRef PassID, llvm::Any));
+  MOCK_METHOD2(startPassManager, void(StringRef PM, llvm::Any));
+  MOCK_METHOD2(finishPassManager, void(StringRef PM, llvm::Any));
 
   void registerPassInstrumentation() {
     Callbacks.registerBeforePassCallback([this](StringRef P, llvm::Any IR) {
       return this->runBeforePass(P, IR);
     });
     Callbacks.registerAfterPassCallback(
         [this](StringRef P, llvm::Any IR) { this->runAfterPass(P, IR); });
     Callbacks.registerBeforeAnalysisCallback([this](StringRef P, llvm::Any IR) {
       return this->runBeforeAnalysis(P, IR);
     });
     Callbacks.registerAfterAnalysisCallback(
         [this](StringRef P, llvm::Any IR) { this->runAfterAnalysis(P, IR); });
+    Callbacks.registerStartPassManagerCallback(
+        [this](StringRef P, llvm::Any IR) { this->startPassManager(P, IR); });
+    Callbacks.registerFinishPassManagerCallback(
+        [this](StringRef P, llvm::Any IR) { this->finishPassManager(P, IR); });
   }
 
   void ignoreNonMockPassInstrumentation(StringRef IRName) {
     // Generic EXPECT_CALLs are needed to match instrumentation on unimportant
     // parts of a pipeline that we do not care about (e.g. various passes added
     // by default by PassBuilder - Verifier pass etc).
     // Make sure to avoid ignoring Mock passes/analysis, we definitely want
     // to check these explicitly.
     EXPECT_CALL(*this,
                 runBeforePass(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
     EXPECT_CALL(*this, runAfterPass(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
     EXPECT_CALL(*this,
                 runBeforeAnalysis(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
     EXPECT_CALL(*this,
                 runAfterAnalysis(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
+    EXPECT_CALL(*this,
+                startPassManager(Not(HasNameRegex("Mock")), HasName(IRName)))
+        .Times(AnyNumber());
+    EXPECT_CALL(*this,
+                finishPassManager(Not(HasNameRegex("Mock")), HasName(IRName)))
+        .Times(AnyNumber());
   }
 };
 
 template <typename PassManagerT> class PassBuilderCallbacksTest;
 
 /// This test fixture is shared between all the actual tests below and
 /// takes care of setting up appropriate defaults.
 ///
TEST_F(ModuleCallbacksTest, InstrumentedPasses) {
 
   CallbacksHandle.registerPassInstrumentation();
   // Non-mock instrumentation not specifically mentioned below can be ignored.
   CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
 
   // PassInstrumentation calls should happen in-sequence, in the same order
   // as passes/analyses are scheduled.
   ::testing::Sequence PISequence;
+  EXPECT_CALL(CallbacksHandle, startPassManager(_, HasName("<string>")))
+      .InSequence(PISequence);
   EXPECT_CALL(CallbacksHandle, runBeforePass(HasNameRegex("MockPassHandle"),
                                              HasName("<string>")))
       .InSequence(PISequence);
   EXPECT_CALL(CallbacksHandle,
               runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"),
                                 HasName("<string>")))
       .InSequence(PISequence);
   EXPECT_CALL(
       CallbacksHandle,
       runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), HasName("<string>")))
       .InSequence(PISequence);
   EXPECT_CALL(CallbacksHandle,
               runAfterPass(HasNameRegex("MockPassHandle"), HasName("<string>")))
       .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("<string>")))
+      .InSequence(PISequence);
 
   StringRef PipelineText = "test-transform";
   ASSERT_TRUE(PB.parsePassPipeline(PM, PipelineText, true))
       << "Pipeline was: " << PipelineText;
 
   PM.run(*M, AM);
 }
 
TEST_F(FunctionCallbacksTest, InstrumentedPasses) {
 
   EXPECT_CALL(AnalysisHandle, run(HasName("foo"), _));
   EXPECT_CALL(PassHandle, run(HasName("foo"), _))
       .WillOnce(Invoke(getAnalysisResult));
 
   // PassInstrumentation calls should happen in-sequence, in the same order
   // as passes/analyses are scheduled.
   ::testing::Sequence PISequence;
+  EXPECT_CALL(CallbacksHandle, startPassManager(HasNameRegex("PassManager"),
+                                                HasName("<string>")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, startPassManager(HasNameRegex("PassAdaptor"),
+                                                HasName("<string>")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle,
+              startPassManager(HasNameRegex("PassManager"), HasName("foo")))
+      .InSequence(PISequence);
   EXPECT_CALL(CallbacksHandle,
               runBeforePass(HasNameRegex("MockPassHandle"), HasName("foo")))
       .InSequence(PISequence);
   EXPECT_CALL(
       CallbacksHandle,
       runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), HasName("foo")))
       .InSequence(PISequence);
   EXPECT_CALL(
       CallbacksHandle,
       runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), HasName("foo")))
       .InSequence(PISequence);
   EXPECT_CALL(CallbacksHandle,
               runAfterPass(HasNameRegex("MockPassHandle"), HasName("foo")))
       .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("foo")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("<string>")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("<string>")))
+      .InSequence(PISequence);
 
   StringRef PipelineText = "test-transform";
   ASSERT_TRUE(PB.parsePassPipeline(PM, PipelineText, true))
       << "Pipeline was: " << PipelineText;
   PM.run(*M, AM);
 }
 
 TEST_F(FunctionCallbacksTest, InstrumentedSkippedPasses) {
TEST_F(LoopCallbacksTest, InstrumentedPasses) {
   EXPECT_CALL(AnalysisHandle, run(HasName("loop"), _, _));
   EXPECT_CALL(PassHandle, run(HasName("loop"), _, _, _))
       .WillOnce(WithArgs<0, 1, 2>(Invoke(getAnalysisResult)));
 
   // PassInstrumentation calls should happen in-sequence, in the same order
   // as passes/analyses are scheduled.
   ::testing::Sequence PISequence;
   EXPECT_CALL(CallbacksHandle,
+              startPassManager(HasNameRegex("PassManager<.*Module.*>"),
+                               HasName("<string>")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle,
+              startPassManager(HasNameRegex("ModuleToFunctionPassAdaptor"),
+                               HasName("<string>")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle,
+              startPassManager(HasNameRegex("PassManager<.*Function.*>"),
+                               HasName("foo")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle,
+              startPassManager(HasNameRegex("FunctionToLoopPassAdaptor"),
+                               HasName("foo")))
+      .InSequence(PISequence);
+  EXPECT_CALL(
+      CallbacksHandle,
+      startPassManager(HasNameRegex("PassManager<.*Loop.*>"), HasName("loop")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle,
               runBeforePass(HasNameRegex("MockPassHandle"), HasName("loop")))
       .InSequence(PISequence);
   EXPECT_CALL(
       CallbacksHandle,
       runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), HasName("loop")))
       .InSequence(PISequence);
   EXPECT_CALL(
       CallbacksHandle,
       runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), HasName("loop")))
       .InSequence(PISequence);
   EXPECT_CALL(CallbacksHandle,
               runAfterPass(HasNameRegex("MockPassHandle"), HasName("loop")))
       .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("loop")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("foo")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("foo")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("<string>")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("<string>")))
+      .InSequence(PISequence);
 
   StringRef PipelineText = "test-transform";
   ASSERT_TRUE(PB.parsePassPipeline(PM, PipelineText, true))
       << "Pipeline was: " << PipelineText;
   PM.run(*M, AM);
 }
 
 TEST_F(LoopCallbacksTest, InstrumentedSkippedPasses) {
TEST_F(CGSCCCallbacksTest, InstrumentedPasses) {
 
   EXPECT_CALL(AnalysisHandle, run(HasName("(foo)"), _, _));
   EXPECT_CALL(PassHandle, run(HasName("(foo)"), _, _, _))
       .WillOnce(WithArgs<0, 1, 2>(Invoke(getAnalysisResult)));
 
   // PassInstrumentation calls should happen in-sequence, in the same order
   // as passes/analyses are scheduled.
   ::testing::Sequence PISequence;
+  EXPECT_CALL(CallbacksHandle,
+              startPassManager(HasNameRegex("PassManager<.*Module.*>"),
+                               HasName("<string>")))
+      .InSequence(PISequence);
+  EXPECT_CALL(
+      CallbacksHandle,
+      startPassManager(HasNameRegex("ModuleToPostOrderCGSCCPassAdaptor"),
+                       HasName("<string>")))
+      .InSequence(PISequence);
+  EXPECT_CALL(
+      CallbacksHandle,
+      startPassManager(HasNameRegex("PassManager<.*LazyCallGraph::SCC.*>"),
+                       HasName("(foo)")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle,
+              runBeforePass(HasNameRegex("MockPassHandle"), HasName("(foo)")))
+      .InSequence(PISequence);
   EXPECT_CALL(
       CallbacksHandle,
       runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), HasName("(foo)")))
       .InSequence(PISequence);
   EXPECT_CALL(
       CallbacksHandle,
       runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), HasName("(foo)")))
       .InSequence(PISequence);
   EXPECT_CALL(CallbacksHandle,
               runAfterPass(HasNameRegex("MockPassHandle"), HasName("(foo)")))
       .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("(foo)")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("<string>")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, finishPassManager(_, HasName("<string>")))
+      .InSequence(PISequence);
 
   StringRef PipelineText = "test-transform";
   ASSERT_TRUE(PB.parsePassPipeline(PM, PipelineText, true))
       << "Pipeline was: " << PipelineText;
   PM.run(*M, AM);
 }
 
 TEST_F(CGSCCCallbacksTest, InstrumentedSkippedPasses) {