[NewPM][PassInstrumentation] Add PreservedAnalyses parameter to AfterPass* callbacks

[yrouban]

Summary

Both AfterPass and AfterPassInvalidated pass instrumentation callbacks get additional parameter of type PreservedAnalyses.
This patch was created by @fedor.sergeev. I have just slightly changed it.

[fedor.sergeev]

LGTM.
Other people had plenty of time to look at this.

You can proceed with integration.

[yrouban]

just rebased.

[ychen]

I would suggest removing 'NFC' from the title.

llvm/lib/IR/PassTimingInfo.cpp

270–272

Remove PA

[yrouban]

rebased + addressed comments (minor)

llvm/include/llvm/Analysis/CGSCCPassManager.h

for (LazyCallGraph::Node *N : Nodes) {
         continue;
 
       PreservedAnalyses PassPA;
       {
         TimeTraceScope TimeScope(Pass.name());
         PassPA = Pass.run(F, FAM);
       }
 
-      PI.runAfterPass<Function>(Pass, F);
+      PI.runAfterPass<Function>(Pass, F, PassPA);
 
       // We know that the function pass couldn't have invalidated any other
       // function's analyses (that's the contract of a function pass), so
       // directly handle the function analysis manager's invalidation here.
       FAM.invalidate(F, PassPA);
 
       // Then intersect the preserved set so that invalidation of module
       // analyses will eventually occur when the module pass completes.
PreservedAnalyses run(LazyCallGraph::SCC &InitialC, CGSCCAnalysisManager &AM,
     for (int Iteration = 0;; ++Iteration) {
 
       if (!PI.runBeforePass<LazyCallGraph::SCC>(Pass, *C))
         continue;
 
       PreservedAnalyses PassPA = Pass.run(*C, AM, CG, UR);
 
       if (UR.InvalidatedSCCs.count(C))
-        PI.runAfterPassInvalidated<LazyCallGraph::SCC>(Pass);
+        PI.runAfterPassInvalidated<LazyCallGraph::SCC>(Pass, PassPA);
       else
-        PI.runAfterPass<LazyCallGraph::SCC>(Pass, *C);
+        PI.runAfterPass<LazyCallGraph::SCC>(Pass, *C, PassPA);
 
       // If the SCC structure has changed, bail immediately and let the outer
       // CGSCC layer handle any iteration to reflect the refined structure.
       if (UR.UpdatedC && UR.UpdatedC != C) {
         PA.intersect(std::move(PassPA));
         break;
       }
 
do {
 
           PreservedAnalyses PassPA;
           {
             TimeTraceScope TimeScope(Pass.name());
             PassPA = Pass.run(*C, CGAM, CG, UR);
           }
 
           if (UR.InvalidatedSCCs.count(C))
-            PI.runAfterPassInvalidated<LazyCallGraph::SCC>(Pass);
+            PI.runAfterPassInvalidated<LazyCallGraph::SCC>(Pass, PassPA);
           else
-            PI.runAfterPass<LazyCallGraph::SCC>(Pass, *C);
+            PI.runAfterPass<LazyCallGraph::SCC>(Pass, *C, PassPA);
 
           // Update the SCC and RefSCC if necessary.
           C = UR.UpdatedC ? UR.UpdatedC : C;
           RC = UR.UpdatedRC ? UR.UpdatedRC : RC;
 
           if (UR.UpdatedC) {
             // If we're updating the SCC, also update the FAM inside the proxy's
             // result.

llvm/include/llvm/IR/PassInstrumentation.h

public:
   // We call all BeforePassFuncs to determine if a pass should run or not.
   // BeforeNonSkippedPassFuncs are called only if the pass should run.
   // TODO: currently AfterPassInvalidated does not accept IRUnit, since passing
   // already invalidated IRUnit is unsafe. There are ways to handle invalidated
   // IRUnits in a safe way, and we might pursue that as soon as there is a
   // useful instrumentation that needs it.
   using BeforePassFunc = bool(StringRef, Any);
   using BeforeNonSkippedPassFunc = void(StringRef, Any);
-  using AfterPassFunc = void(StringRef, Any);
-  using AfterPassInvalidatedFunc = void(StringRef);
+  using AfterPassFunc = void(StringRef, Any, const PreservedAnalyses &);
+  using AfterPassInvalidatedFunc = void(StringRef, const PreservedAnalyses &);
   using BeforeAnalysisFunc = void(StringRef, Any);
   using AfterAnalysisFunc = void(StringRef, Any);
 
 public:
   PassInstrumentationCallbacks() {}
 
   /// Copying PassInstrumentationCallbacks is not intended.
   PassInstrumentationCallbacks(const PassInstrumentationCallbacks &) = delete;
bool runBeforePass(const PassT &Pass, const IRUnitT &IR) const {
 
     return ShouldRun;
   }
 
   /// AfterPass instrumentation point - takes \p Pass instance that has
   /// just been executed and constant reference to \p IR it operates on.
   /// \p IR is guaranteed to be valid at this point.
   template <typename IRUnitT, typename PassT>
-  void runAfterPass(const PassT &Pass, const IRUnitT &IR) const {
+  void runAfterPass(const PassT &Pass, const IRUnitT &IR,
+                    const PreservedAnalyses &PA) const {
     if (Callbacks)
       for (auto &C : Callbacks->AfterPassCallbacks)
-        C(Pass.name(), llvm::Any(&IR));
+        C(Pass.name(), llvm::Any(&IR), PA);
   }
 
   /// AfterPassInvalidated instrumentation point - takes \p Pass instance
   /// that has just been executed. For use when IR has been invalidated
   /// by \p Pass execution.
   template <typename IRUnitT, typename PassT>
-  void runAfterPassInvalidated(const PassT &Pass) const {
+  void runAfterPassInvalidated(const PassT &Pass,
+                               const PreservedAnalyses &PA) const {
     if (Callbacks)
       for (auto &C : Callbacks->AfterPassInvalidatedCallbacks)
-        C(Pass.name());
+        C(Pass.name(), PA);
   }
 
   /// BeforeAnalysis instrumentation point - takes \p Analysis instance
   /// to be executed and constant reference to IR it operates on.
   template <typename IRUnitT, typename PassT>
   void runBeforeAnalysis(const PassT &Analysis, const IRUnitT &IR) const {
     if (Callbacks)
       for (auto &C : Callbacks->BeforeAnalysisCallbacks)

llvm/include/llvm/IR/PassManager.h

for (unsigned Idx = 0, Size = Passes.size(); Idx != Size; ++Idx) {
       PreservedAnalyses PassPA;
       {
         TimeTraceScope TimeScope(P->name(), IR.getName());
         PassPA = P->run(IR, AM, ExtraArgs...);
       }
 
       // Call onto PassInstrumentation's AfterPass callbacks immediately after
       // running the pass.
-      PI.runAfterPass<IRUnitT>(*P, IR);
+      PI.runAfterPass<IRUnitT>(*P, IR, PassPA);
 
       // Update the analysis manager as each pass runs and potentially
       // invalidates analyses.
       AM.invalidate(IR, PassPA);
 
       // Finally, intersect the preserved analyses to compute the aggregate
       // preserved set for this pass manager.
       PA.intersect(std::move(PassPA));
for (Function &F : M) {
         continue;
 
       PreservedAnalyses PassPA;
       {
         TimeTraceScope TimeScope(Pass.name(), F.getName());
         PassPA = Pass.run(F, FAM);
       }
 
-      PI.runAfterPass(Pass, F);
+      PI.runAfterPass(Pass, F, PassPA);
 
       // We know that the function pass couldn't have invalidated any other
       // function's analyses (that's the contract of a function pass), so
       // directly handle the function analysis manager's invalidation here.
       FAM.invalidate(F, PassPA);
 
       // Then intersect the preserved set so that invalidation of module
       // analyses will eventually occur when the module pass completes.
PreservedAnalyses run(IRUnitT &IR, AnalysisManagerT &AM, Ts &&... Args) {
 
     auto PA = PreservedAnalyses::all();
     for (int i = 0; i < Count; ++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);
+      PreservedAnalyses IterPA = P.run(IR, AM, std::forward<Ts>(Args)...);
+      PA.intersect(IterPA);
+      PI.runAfterPass(P, IR, IterPA);
     }
     return PA;
   }
 
 private:
   int Count;
   PassT P;
 };

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

PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM) {
     // manager handles all the invalidation at that layer.
     PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(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);
+      PI.runAfterPass<Function>(LoopCanonicalizationFPM, F, PA);
     }
 
     // Get the loop structure for this function
     LoopInfo &LI = AM.getResult<LoopAnalysis>(F);
 
     // If there are no loops, there is nothing to do here.
     if (LI.empty())
       return PA;
#endif
       PreservedAnalyses PassPA;
       {
         TimeTraceScope TimeScope(Pass.name());
         PassPA = Pass.run(*L, LAM, LAR, Updater);
       }
 
       // Do not pass deleted Loop into the instrumentation.
       if (Updater.skipCurrentLoop())
-        PI.runAfterPassInvalidated<Loop>(Pass);
+        PI.runAfterPassInvalidated<Loop>(Pass, PassPA);
       else
-        PI.runAfterPass<Loop>(Pass, *L);
+        PI.runAfterPass<Loop>(Pass, *L, PassPA);
 
       // FIXME: We should verify the set of analyses relevant to Loop passes
       // are preserved.
 
       // If the loop hasn't been deleted, we need to handle invalidation here.
       if (!Updater.skipCurrentLoop())
         // We know that the loop pass couldn't have invalidated any other
         // loop's analyses (that's the contract of a loop pass), so directly

llvm/lib/Analysis/CGSCCPassManager.cpp

for (auto &Pass : Passes) {
 
     PreservedAnalyses PassPA;
     {
       TimeTraceScope TimeScope(Pass->name());
       PassPA = Pass->run(*C, AM, G, UR);
     }
 
     if (UR.InvalidatedSCCs.count(C))
-      PI.runAfterPassInvalidated<LazyCallGraph::SCC>(*Pass);
+      PI.runAfterPassInvalidated<LazyCallGraph::SCC>(*Pass, PassPA);
     else
-      PI.runAfterPass<LazyCallGraph::SCC>(*Pass, *C);
+      PI.runAfterPass<LazyCallGraph::SCC>(*Pass, *C, PassPA);
 
     // Update the SCC if necessary.
     C = UR.UpdatedC ? UR.UpdatedC : C;
     if (UR.UpdatedC) {
       // If C is updated, also create a proxy and update FAM inside the result.
       auto *ResultFAMCP =
           &AM.getResult<FunctionAnalysisManagerCGSCCProxy>(*C, G);
       ResultFAMCP->updateFAM(FAM);

llvm/lib/IR/PassTimingInfo.cpp

 
 void TimePassesHandler::registerCallbacks(PassInstrumentationCallbacks &PIC) {
   if (!Enabled)
     return;
 
   PIC.registerBeforeNonSkippedPassCallback(
       [this](StringRef P, Any) { this->runBeforePass(P); });
   PIC.registerAfterPassCallback(
-      [this](StringRef P, Any) { this->runAfterPass(P); });
+      [this](StringRef P, Any, const PreservedAnalyses &) {
+        this->runAfterPass(P);
+      });
   PIC.registerAfterPassInvalidatedCallback(
-      [this](StringRef P) { this->runAfterPass(P); });
+      [this](StringRef P, const PreservedAnalyses &) {
+        this->runAfterPass(P);
+      });
   PIC.registerBeforeAnalysisCallback(
       [this](StringRef P, Any) { this->runBeforePass(P); });
   PIC.registerAfterAnalysisCallback(
       [this](StringRef P, Any) { this->runAfterPass(P); });
 }
 
 } // namespace llvm

llvm/lib/Passes/StandardInstrumentations.cpp

void PrintIRInstrumentation::registerCallbacks(
   // for later use in AfterPassInvalidated.
   StoreModuleDesc = llvm::forcePrintModuleIR() && llvm::shouldPrintAfterPass();
   if (llvm::shouldPrintBeforePass() || StoreModuleDesc)
     PIC.registerBeforeNonSkippedPassCallback(
         [this](StringRef P, Any IR) { this->printBeforePass(P, IR); });
 
   if (llvm::shouldPrintAfterPass()) {
     PIC.registerAfterPassCallback(
-        [this](StringRef P, Any IR) { this->printAfterPass(P, IR); });
+        [this](StringRef P, Any IR, const PreservedAnalyses &) {
+          this->printAfterPass(P, IR);
+        });
     PIC.registerAfterPassInvalidatedCallback(
-        [this](StringRef P) { this->printAfterPassInvalidated(P); });
+        [this](StringRef P, const PreservedAnalyses &) {
+          this->printAfterPassInvalidated(P);
+        });
   }
 }
 
 void OptNoneInstrumentation::registerCallbacks(
     PassInstrumentationCallbacks &PIC) {
   PIC.registerBeforePassCallback(
       [this](StringRef P, Any IR) { return this->skip(P, IR); });
 }

llvm/lib/Transforms/Scalar/LoopPassManager.cpp

for (auto &Pass : Passes) {
     PreservedAnalyses PassPA;
     {
       TimeTraceScope TimeScope(Pass->name(), L.getName());
       PassPA = Pass->run(L, AM, AR, U);
     }
 
     // do not pass deleted Loop into the instrumentation
     if (U.skipCurrentLoop())
-      PI.runAfterPassInvalidated<Loop>(*Pass);
+      PI.runAfterPassInvalidated<Loop>(*Pass, PassPA);
     else
-      PI.runAfterPass<Loop>(*Pass, L);
+      PI.runAfterPass<Loop>(*Pass, L, PassPA);
 
     // If the loop was deleted, abort the run and return to the outer walk.
     if (U.skipCurrentLoop()) {
       PA.intersect(std::move(PassPA));
       break;
     }
 
 #ifndef NDEBUG

llvm/unittests/IR/PassBuilderCallbacksTest.cpp

-//===- unittests/IR/PassBuilderCallbacksTest.cpp - PB Callback Tests --===//
+//===- Standard pass instrumentations handling ----------------*- C++ -*--===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
+/// \file
+///
+/// This file defines IR-printing pass instrumentation callbacks as well as
+/// StandardInstrumentations class that manages standard pass instrumentations.
+///
+//===----------------------------------------------------------------------===//
 
-#include "llvm/Testing/Support/Error.h"
-#include <functional>
-#include <gmock/gmock.h>
-#include <gtest/gtest.h>
-#include <llvm/ADT/Any.h>
-#include <llvm/Analysis/CGSCCPassManager.h>
-#include <llvm/Analysis/LoopAnalysisManager.h>
-#include <llvm/AsmParser/Parser.h>
-#include <llvm/IR/LLVMContext.h>
-#include <llvm/IR/PassInstrumentation.h>
-#include <llvm/IR/PassManager.h>
-#include <llvm/Passes/PassBuilder.h>
-#include <llvm/Support/Regex.h>
-#include <llvm/Support/SourceMgr.h>
-#include <llvm/Transforms/Scalar/LoopPassManager.h>
+#include "llvm/ADT/Any.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/Analysis/CallGraphSCCPass.h"
+#include "llvm/Analysis/LazyCallGraph.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRPrintingPasses.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassInstrumentation.h"
+#include "llvm/Passes/StandardInstrumentations.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 
+// TODO: remove once all required passes are marked as such.
+static cl::opt<bool>
+    EnableOptnone("enable-npm-optnone", cl::init(false),
+                  cl::desc("Enable skipping optional passes optnone functions "
+                           "under new pass manager"));
+
 namespace {
-using testing::AnyNumber;
-using testing::AtLeast;
-using testing::DoDefault;
-using testing::Not;
-using testing::Return;
-using testing::Expectation;
-using testing::Invoke;
-using testing::WithArgs;
-using testing::_;
 
-/// A CRTP base for analysis mock handles
-///
-/// This class reconciles mocking with the value semantics implementation of the
-/// AnalysisManager. Analysis mock handles should derive from this class and
-/// call \c setDefault() in their constroctur for wiring up the defaults defined
-/// by this base with their mock run() and invalidate() implementations.
-template <typename DerivedT, typename IRUnitT,
-          typename AnalysisManagerT = AnalysisManager<IRUnitT>,
-          typename... ExtraArgTs>
-class MockAnalysisHandleBase {
-public:
-  class Analysis : public AnalysisInfoMixin<Analysis> {
-    friend AnalysisInfoMixin<Analysis>;
-    friend MockAnalysisHandleBase;
-    static AnalysisKey Key;
-
-    DerivedT *Handle;
-
-    Analysis(DerivedT &Handle) : Handle(&Handle) {
-      static_assert(std::is_base_of<MockAnalysisHandleBase, DerivedT>::value,
-                    "Must pass the derived type to this template!");
-    }
-
-  public:
-    class Result {
-      friend MockAnalysisHandleBase;
-
-      DerivedT *Handle;
-
-      Result(DerivedT &Handle) : Handle(&Handle) {}
-
-    public:
-      // Forward invalidation events to the mock handle.
-      bool invalidate(IRUnitT &IR, const PreservedAnalyses &PA,
-                      typename AnalysisManagerT::Invalidator &Inv) {
-        return Handle->invalidate(IR, PA, Inv);
-      }
-    };
-
-    Result run(IRUnitT &IR, AnalysisManagerT &AM, ExtraArgTs... ExtraArgs) {
-      return Handle->run(IR, AM, ExtraArgs...);
-    }
-  };
-
-  Analysis getAnalysis() { return Analysis(static_cast<DerivedT &>(*this)); }
-  typename Analysis::Result getResult() {
-    return typename Analysis::Result(static_cast<DerivedT &>(*this));
-  }
-  static StringRef getName() { return llvm::getTypeName<DerivedT>(); }
-
-protected:
-  // FIXME: MSVC seems unable to handle a lambda argument to Invoke from within
-  // the template, so we use a boring static function.
-  static bool invalidateCallback(IRUnitT &IR, const PreservedAnalyses &PA,
-                                 typename AnalysisManagerT::Invalidator &Inv) {
-    auto PAC = PA.template getChecker<Analysis>();
-    return !PAC.preserved() &&
-           !PAC.template preservedSet<AllAnalysesOn<IRUnitT>>();
-  }
-
-  /// Derived classes should call this in their constructor to set up default
-  /// mock actions. (We can't do this in our constructor because this has to
-  /// run after the DerivedT is constructed.)
-  void setDefaults() {
-    ON_CALL(static_cast<DerivedT &>(*this),
-            run(_, _, testing::Matcher<ExtraArgTs>(_)...))
-        .WillByDefault(Return(this->getResult()));
-    ON_CALL(static_cast<DerivedT &>(*this), invalidate(_, _, _))
-        .WillByDefault(Invoke(&invalidateCallback));
+/// Extracting Module out of \p IR unit. Also fills a textual description
+/// of \p IR for use in header when printing.
+Optional<std::pair<const Module *, std::string>> unwrapModule(Any IR) {
+  if (any_isa<const Module *>(IR))
+    return std::make_pair(any_cast<const Module *>(IR), std::string());
+
+  if (any_isa<const Function *>(IR)) {
+    const Function *F = any_cast<const Function *>(IR);
+    if (!llvm::isFunctionInPrintList(F->getName()))
+      return None;
+    const Module *M = F->getParent();
+    return std::make_pair(M, formatv(" (function: {0})", F->getName()).str());
   }
-};
 
-/// A CRTP base for pass mock handles
-///
-/// This class reconciles mocking with the value semantics implementation of the
-/// PassManager. Pass mock handles should derive from this class and
-/// call \c setDefault() in their constroctur for wiring up the defaults defined
-/// by this base with their mock run() and invalidate() implementations.
-template <typename DerivedT, typename IRUnitT, typename AnalysisManagerT,
-          typename... ExtraArgTs>
-AnalysisKey MockAnalysisHandleBase<DerivedT, IRUnitT, AnalysisManagerT,
-                                   ExtraArgTs...>::Analysis::Key;
-
-template <typename DerivedT, typename IRUnitT,
-          typename AnalysisManagerT = AnalysisManager<IRUnitT>,
-          typename... ExtraArgTs>
-class MockPassHandleBase {
-public:
-  class Pass : public PassInfoMixin<Pass> {
-    friend MockPassHandleBase;
-
-    DerivedT *Handle;
-
-    Pass(DerivedT &Handle) : Handle(&Handle) {
-      static_assert(std::is_base_of<MockPassHandleBase, DerivedT>::value,
-                    "Must pass the derived type to this template!");
-    }
-
-  public:
-    PreservedAnalyses run(IRUnitT &IR, AnalysisManagerT &AM,
-                          ExtraArgTs... ExtraArgs) {
-      return Handle->run(IR, AM, ExtraArgs...);
-    }
-  };
-
-  static StringRef getName() { return llvm::getTypeName<DerivedT>(); }
-
-  Pass getPass() { return Pass(static_cast<DerivedT &>(*this)); }
-
-protected:
-  /// Derived classes should call this in their constructor to set up default
-  /// mock actions. (We can't do this in our constructor because this has to
-  /// run after the DerivedT is constructed.)
-  void setDefaults() {
-    ON_CALL(static_cast<DerivedT &>(*this),
-            run(_, _, testing::Matcher<ExtraArgTs>(_)...))
-        .WillByDefault(Return(PreservedAnalyses::all()));
-  }
-};
-
-/// Mock handles for passes for the IRUnits Module, CGSCC, Function, Loop.
-/// These handles define the appropriate run() mock interface for the respective
-/// IRUnit type.
-template <typename IRUnitT> struct MockPassHandle;
-template <>
-struct MockPassHandle<Loop>
-    : MockPassHandleBase<MockPassHandle<Loop>, Loop, LoopAnalysisManager,
-                         LoopStandardAnalysisResults &, LPMUpdater &> {
-  MOCK_METHOD4(run,
-               PreservedAnalyses(Loop &, LoopAnalysisManager &,
-                                 LoopStandardAnalysisResults &, LPMUpdater &));
-  static void invalidateLoop(Loop &L, LoopAnalysisManager &,
-                             LoopStandardAnalysisResults &,
-                             LPMUpdater &Updater) {
-    Updater.markLoopAsDeleted(L, L.getName());
-  }
-  MockPassHandle() { setDefaults(); }
-};
-
-template <>
-struct MockPassHandle<Function>
-    : MockPassHandleBase<MockPassHandle<Function>, Function> {
-  MOCK_METHOD2(run, PreservedAnalyses(Function &, FunctionAnalysisManager &));
-
-  MockPassHandle() { setDefaults(); }
-};
-
-template <>
-struct MockPassHandle<LazyCallGraph::SCC>
-    : MockPassHandleBase<MockPassHandle<LazyCallGraph::SCC>, LazyCallGraph::SCC,
-                         CGSCCAnalysisManager, LazyCallGraph &,
-                         CGSCCUpdateResult &> {
-  MOCK_METHOD4(run,
-               PreservedAnalyses(LazyCallGraph::SCC &, CGSCCAnalysisManager &,
-                                 LazyCallGraph &G, CGSCCUpdateResult &UR));
-
-  static void invalidateSCC(LazyCallGraph::SCC &C, CGSCCAnalysisManager &,
-                            LazyCallGraph &, CGSCCUpdateResult &UR) {
-    UR.InvalidatedSCCs.insert(&C);
-  }
-
-  MockPassHandle() { setDefaults(); }
-};
-
-template <>
-struct MockPassHandle<Module>
-    : MockPassHandleBase<MockPassHandle<Module>, Module> {
-  MOCK_METHOD2(run, PreservedAnalyses(Module &, ModuleAnalysisManager &));
-
-  MockPassHandle() { setDefaults(); }
-};
-
-/// Mock handles for analyses for the IRUnits Module, CGSCC, Function, Loop.
-/// These handles define the appropriate run() and invalidate() mock interfaces
-/// for the respective IRUnit type.
-template <typename IRUnitT> struct MockAnalysisHandle;
-template <>
-struct MockAnalysisHandle<Loop>
-    : MockAnalysisHandleBase<MockAnalysisHandle<Loop>, Loop,
-                             LoopAnalysisManager,
-                             LoopStandardAnalysisResults &> {
-
-  MOCK_METHOD3_T(run, typename Analysis::Result(Loop &, LoopAnalysisManager &,
-                                                LoopStandardAnalysisResults &));
-
-  MOCK_METHOD3_T(invalidate, bool(Loop &, const PreservedAnalyses &,
-                                  LoopAnalysisManager::Invalidator &));
-
-  MockAnalysisHandle<Loop>() { this->setDefaults(); }
-};
-
-template <>
-struct MockAnalysisHandle<Function>
-    : MockAnalysisHandleBase<MockAnalysisHandle<Function>, Function> {
-  MOCK_METHOD2(run, Analysis::Result(Function &, FunctionAnalysisManager &));
-
-  MOCK_METHOD3(invalidate, bool(Function &, const PreservedAnalyses &,
-                                FunctionAnalysisManager::Invalidator &));
-
-  MockAnalysisHandle<Function>() { setDefaults(); }
-};
-
-template <>
-struct MockAnalysisHandle<LazyCallGraph::SCC>
-    : MockAnalysisHandleBase<MockAnalysisHandle<LazyCallGraph::SCC>,
-                             LazyCallGraph::SCC, CGSCCAnalysisManager,
-                             LazyCallGraph &> {
-  MOCK_METHOD3(run, Analysis::Result(LazyCallGraph::SCC &,
-                                     CGSCCAnalysisManager &, LazyCallGraph &));
-
-  MOCK_METHOD3(invalidate, bool(LazyCallGraph::SCC &, const PreservedAnalyses &,
-                                CGSCCAnalysisManager::Invalidator &));
-
-  MockAnalysisHandle<LazyCallGraph::SCC>() { setDefaults(); }
-};
-
-template <>
-struct MockAnalysisHandle<Module>
-    : MockAnalysisHandleBase<MockAnalysisHandle<Module>, Module> {
-  MOCK_METHOD2(run, Analysis::Result(Module &, ModuleAnalysisManager &));
-
-  MOCK_METHOD3(invalidate, bool(Module &, const PreservedAnalyses &,
-                                ModuleAnalysisManager::Invalidator &));
-
-  MockAnalysisHandle<Module>() { setDefaults(); }
-};
-
-static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
-  SMDiagnostic Err;
-  return parseAssemblyString(IR, Err, C);
-}
-
-/// Helper for HasName matcher that returns getName both for IRUnit and
-/// for IRUnit pointer wrapper into llvm::Any (wrapped by PassInstrumentation).
-template <typename IRUnitT> std::string getName(const IRUnitT &IR) {
-  return std::string(IR.getName());
-}
-
-template <> std::string getName(const StringRef &name) {
-  return std::string(name);
-}
-
-template <> std::string getName(const llvm::Any &WrappedIR) {
-  if (any_isa<const Module *>(WrappedIR))
-    return any_cast<const Module *>(WrappedIR)->getName().str();
-  if (any_isa<const Function *>(WrappedIR))
-    return any_cast<const Function *>(WrappedIR)->getName().str();
-  if (any_isa<const Loop *>(WrappedIR))
-    return any_cast<const Loop *>(WrappedIR)->getName().str();
-  if (any_isa<const LazyCallGraph::SCC *>(WrappedIR))
-    return any_cast<const LazyCallGraph::SCC *>(WrappedIR)->getName();
-  return "<UNKNOWN>";
+  if (any_isa<const LazyCallGraph::SCC *>(IR)) {
+    const LazyCallGraph::SCC *C = any_cast<const LazyCallGraph::SCC *>(IR);
+    for (const LazyCallGraph::Node &N : *C) {
+      const Function &F = N.getFunction();
+      if (!F.isDeclaration() && isFunctionInPrintList(F.getName())) {
+        const Module *M = F.getParent();
+        return std::make_pair(M, formatv(" (scc: {0})", C->getName()).str());
       }
-/// Define a custom matcher for objects which support a 'getName' method.
-///
-/// LLVM often has IR objects or analysis objects which expose a name
-/// and in tests it is convenient to match these by name for readability.
-/// Usually, this name is either a StringRef or a plain std::string. This
-/// matcher supports any type exposing a getName() method of this form whose
-/// return value is compatible with an std::ostream. For StringRef, this uses
-/// the shift operator defined above.
-///
-/// It should be used as:
-///
-///   HasName("my_function")
-///
-/// No namespace or other qualification is required.
-MATCHER_P(HasName, Name, "") {
-  *result_listener << "has name '" << getName(arg) << "'";
-  return Name == getName(arg);
-}
-
-MATCHER_P(HasNameRegex, Name, "") {
-  *result_listener << "has name '" << getName(arg) << "'";
-  llvm::Regex r(Name);
-  return r.match(getName(arg));
-}
-
-struct MockPassInstrumentationCallbacks {
-  PassInstrumentationCallbacks Callbacks;
-
-  MockPassInstrumentationCallbacks() {
-    ON_CALL(*this, runBeforePass(_, _)).WillByDefault(Return(true));
-  }
-  MOCK_METHOD2(runBeforePass, bool(StringRef PassID, llvm::Any));
-  MOCK_METHOD2(runBeforeNonSkippedPass, void(StringRef PassID, llvm::Any));
-  MOCK_METHOD2(runAfterPass, void(StringRef PassID, llvm::Any));
-  MOCK_METHOD1(runAfterPassInvalidated, void(StringRef PassID));
-  MOCK_METHOD2(runBeforeAnalysis, void(StringRef PassID, llvm::Any));
-  MOCK_METHOD2(runAfterAnalysis, void(StringRef PassID, llvm::Any));
-
-  void registerPassInstrumentation() {
-    Callbacks.registerBeforePassCallback([this](StringRef P, llvm::Any IR) {
-      return this->runBeforePass(P, IR);
-    });
-    Callbacks.registerBeforeNonSkippedPassCallback(
-        [this](StringRef P, llvm::Any IR) {
-          this->runBeforeNonSkippedPass(P, IR);
-        });
-    Callbacks.registerAfterPassCallback(
-        [this](StringRef P, llvm::Any IR) { this->runAfterPass(P, IR); });
-    Callbacks.registerAfterPassInvalidatedCallback(
-        [this](StringRef P) { this->runAfterPassInvalidated(P); });
-    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); });
+    }
+    return None;
   }
 
-  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, runBeforeNonSkippedPass(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());
+  if (any_isa<const Loop *>(IR)) {
+    const Loop *L = any_cast<const Loop *>(IR);
+    const Function *F = L->getHeader()->getParent();
+    if (!isFunctionInPrintList(F->getName()))
+      return None;
+    const Module *M = F->getParent();
+    std::string LoopName;
+    raw_string_ostream ss(LoopName);
+    L->getHeader()->printAsOperand(ss, false);
+    return std::make_pair(M, formatv(" (loop: {0})", ss.str()).str());
   }
-};
 
-template <typename PassManagerT> class PassBuilderCallbacksTest;
+  llvm_unreachable("Unknown IR unit");
+}
 
-/// This test fixture is shared between all the actual tests below and
-/// takes care of setting up appropriate defaults.
-///
-/// The template specialization serves to extract the IRUnit and AM types from
-/// the given PassManagerT.
-template <typename TestIRUnitT, typename... ExtraPassArgTs,
-          typename... ExtraAnalysisArgTs>
-class PassBuilderCallbacksTest<PassManager<
-    TestIRUnitT, AnalysisManager<TestIRUnitT, ExtraAnalysisArgTs...>,
-    ExtraPassArgTs...>> : public testing::Test {
-protected:
-  using IRUnitT = TestIRUnitT;
-  using AnalysisManagerT = AnalysisManager<TestIRUnitT, ExtraAnalysisArgTs...>;
-  using PassManagerT =
-      PassManager<TestIRUnitT, AnalysisManagerT, ExtraPassArgTs...>;
-  using AnalysisT = typename MockAnalysisHandle<IRUnitT>::Analysis;
-
-  LLVMContext Context;
-  std::unique_ptr<Module> M;
-
-  MockPassInstrumentationCallbacks CallbacksHandle;
-
-  PassBuilder PB;
-  ModulePassManager PM;
-  LoopAnalysisManager LAM;
-  FunctionAnalysisManager FAM;
-  CGSCCAnalysisManager CGAM;
-  ModuleAnalysisManager AM;
-
-  MockPassHandle<IRUnitT> PassHandle;
-  MockAnalysisHandle<IRUnitT> AnalysisHandle;
-
-  static PreservedAnalyses getAnalysisResult(IRUnitT &U, AnalysisManagerT &AM,
-                                             ExtraAnalysisArgTs &&... Args) {
-    (void)AM.template getResult<AnalysisT>(
-        U, std::forward<ExtraAnalysisArgTs>(Args)...);
-    return PreservedAnalyses::all();
-  }
-
-  PassBuilderCallbacksTest()
-      : M(parseIR(Context,
-                  "declare void @bar()\n"
-                  "define void @foo(i32 %n) {\n"
-                  "entry:\n"
-                  "  br label %loop\n"
-                  "loop:\n"
-                  "  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop ]\n"
-                  "  %iv.next = add i32 %iv, 1\n"
-                  "  tail call void @bar()\n"
-                  "  %cmp = icmp eq i32 %iv, %n\n"
-                  "  br i1 %cmp, label %exit, label %loop\n"
-                  "exit:\n"
-                  "  ret void\n"
-                  "}\n")),
-        CallbacksHandle(),
-        PB(nullptr, PipelineTuningOptions(), None, &CallbacksHandle.Callbacks),
-        PM(true), LAM(true), FAM(true), CGAM(true), AM(true) {
+void printIR(const Function *F, StringRef Banner,
+             StringRef Extra = StringRef()) {
+  if (!llvm::isFunctionInPrintList(F->getName()))
+    return;
+  dbgs() << Banner << Extra << "\n" << static_cast<const Value &>(*F);
+}
 
-    EXPECT_TRUE(&CallbacksHandle.Callbacks ==
-                PB.getPassInstrumentationCallbacks());
+void printIR(const Module *M, StringRef Banner, StringRef Extra = StringRef()) {
+  if (llvm::isFunctionInPrintList("*") || llvm::forcePrintModuleIR()) {
+    dbgs() << Banner << Extra << "\n";
+    M->print(dbgs(), nullptr, false);
+  } else {
+    for (const auto &F : M->functions()) {
+      printIR(&F, Banner, Extra);
+    }
+  }
+}
 
-    /// Register a callback for analysis registration.
-    ///
-    /// The callback is a function taking a reference to an AnalyisManager
-    /// object. When called, the callee gets to register its own analyses with
-    /// this PassBuilder instance.
-    PB.registerAnalysisRegistrationCallback([this](AnalysisManagerT &AM) {
-      // Register our mock analysis
-      AM.registerPass([this] { return AnalysisHandle.getAnalysis(); });
-    });
+void printIR(const LazyCallGraph::SCC *C, StringRef Banner,
+             StringRef Extra = StringRef()) {
+  bool BannerPrinted = false;
+  for (const LazyCallGraph::Node &N : *C) {
+    const Function &F = N.getFunction();
+    if (!F.isDeclaration() && llvm::isFunctionInPrintList(F.getName())) {
+      if (!BannerPrinted) {
+        dbgs() << Banner << Extra << "\n";
+        BannerPrinted = true;
+      }
+      F.print(dbgs());
+    }
+  }
+}
+void printIR(const Loop *L, StringRef Banner) {
+  const Function *F = L->getHeader()->getParent();
+  if (!llvm::isFunctionInPrintList(F->getName()))
+    return;
+  llvm::printLoop(const_cast<Loop &>(*L), dbgs(), std::string(Banner));
+}
 
-    /// Register a callback for pipeline parsing.
-    ///
-    /// During parsing of a textual pipeline, the PassBuilder will call these
-    /// callbacks for each encountered pass name that it does not know. This
-    /// includes both simple pass names as well as names of sub-pipelines. In
-    /// the latter case, the InnerPipeline is not empty.
-    PB.registerPipelineParsingCallback(
-        [this](StringRef Name, PassManagerT &PM,
-               ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
-          /// Handle parsing of the names of analysis utilities such as
-          /// require<test-analysis> and invalidate<test-analysis> for our
-          /// analysis mock handle
-          if (parseAnalysisUtilityPasses<AnalysisT>("test-analysis", Name, PM))
-            return true;
+/// Generic IR-printing helper that unpacks a pointer to IRUnit wrapped into
+/// llvm::Any and does actual print job.
+void unwrapAndPrint(Any IR, StringRef Banner, bool ForceModule = false) {
+  if (ForceModule) {
+    if (auto UnwrappedModule = unwrapModule(IR))
+      printIR(UnwrappedModule->first, Banner, UnwrappedModule->second);
+    return;
+  }
 
-          /// Parse the name of our pass mock handle
-          if (Name == "test-transform") {
-            PM.addPass(PassHandle.getPass());
-            return true;
-          }
-          return false;
-        });
+  if (any_isa<const Module *>(IR)) {
+    const Module *M = any_cast<const Module *>(IR);
+    assert(M && "module should be valid for printing");
+    printIR(M, Banner);
+    return;
 
-    /// Register builtin analyses and cross-register the analysis proxies
-    PB.registerModuleAnalyses(AM);
-    PB.registerCGSCCAnalyses(CGAM);
-    PB.registerFunctionAnalyses(FAM);
-    PB.registerLoopAnalyses(LAM);
-    PB.crossRegisterProxies(LAM, FAM, CGAM, AM);
-  }
-};
-
-using ModuleCallbacksTest = PassBuilderCallbacksTest<ModulePassManager>;
-using CGSCCCallbacksTest = PassBuilderCallbacksTest<CGSCCPassManager>;
-using FunctionCallbacksTest = PassBuilderCallbacksTest<FunctionPassManager>;
-using LoopCallbacksTest = PassBuilderCallbacksTest<LoopPassManager>;
+  if (any_isa<const Function *>(IR)) {
+    const Function *F = any_cast<const Function *>(IR);
+    assert(F && "function should be valid for printing");
+    printIR(F, Banner);
+    return;
+  }
 
-/// Test parsing of the name of our mock pass for all IRUnits.
-///
-/// The pass should by default run our mock analysis and then preserve it.
-TEST_F(ModuleCallbacksTest, Passes) {
-  EXPECT_CALL(AnalysisHandle, run(HasName("<string>"), _));
-  EXPECT_CALL(PassHandle, run(HasName("<string>"), _))
-      .WillOnce(Invoke(getAnalysisResult));
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-
-  PM.run(*M, AM);
-}
-
-TEST_F(ModuleCallbacksTest, InstrumentedPasses) {
-  EXPECT_CALL(AnalysisHandle, run(HasName("<string>"), _));
-  EXPECT_CALL(PassHandle, run(HasName("<string>"), _))
-      .WillOnce(Invoke(getAnalysisResult));
-
-  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, runBeforePass(HasNameRegex("MockPassHandle"),
-                                             HasName("<string>")))
-      .InSequence(PISequence);
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeNonSkippedPass(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);
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-
-  PM.run(*M, AM);
-}
-
-TEST_F(ModuleCallbacksTest, InstrumentedSkippedPasses) {
-  CallbacksHandle.registerPassInstrumentation();
-  // Non-mock instrumentation run here can safely be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
-
-  // Skip all passes by returning false. Pass managers and adaptor passes are
-  // also passes that observed by the callbacks.
-  EXPECT_CALL(CallbacksHandle, runBeforePass(_, _))
-      .WillRepeatedly(Return(false));
-
-  EXPECT_CALL(AnalysisHandle, run(HasName("<string>"), _)).Times(0);
-  EXPECT_CALL(PassHandle, run(HasName("<string>"), _)).Times(0);
-
-  // As the pass is skipped there is no nonskippedpass/afterPass,
-  // beforeAnalysis/afterAnalysis as well.
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeNonSkippedPass(HasNameRegex("MockPassHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle, runAfterPass(HasNameRegex("MockPassHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), _))
-      .Times(0);
-
-  // Order is important here. `Adaptor` expectations should be checked first
-  // because the its argument contains 'PassManager' (for example:
-  // ModuleToFunctionPassAdaptor{{.*}}PassManager{{.*}}). Check
-  // `runBeforeNonSkippedPass` and `runAfterPass` to show that they are not
-  // skipped.
-  //
-  // Pass managers are not ignored.
-  // 5 = (1) ModulePassManager + (2) FunctionPassMangers + (1) LoopPassManager +
-  //     (1) CGSCCPassManager
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeNonSkippedPass(HasNameRegex("PassManager"), _))
-      .Times(5);
-  EXPECT_CALL(
-      CallbacksHandle,
-      runBeforeNonSkippedPass(HasNameRegex("ModuleToFunctionPassAdaptor"), _))
-      .Times(1);
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeNonSkippedPass(
-                  HasNameRegex("ModuleToPostOrderCGSCCPassAdaptor"), _))
-      .Times(1);
-  EXPECT_CALL(
-      CallbacksHandle,
-      runBeforeNonSkippedPass(HasNameRegex("CGSCCToFunctionPassAdaptor"), _))
-      .Times(1);
-  EXPECT_CALL(
-      CallbacksHandle,
-      runBeforeNonSkippedPass(HasNameRegex("FunctionToLoopPassAdaptor"), _))
-      .Times(1);
-
-  // The `runAfterPass` checks are the same as these of
-  // `runBeforeNonSkippedPass`.
-  EXPECT_CALL(CallbacksHandle, runAfterPass(HasNameRegex("PassManager"), _))
-      .Times(5);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPass(HasNameRegex("ModuleToFunctionPassAdaptor"), _))
-      .Times(1);
-  EXPECT_CALL(
-      CallbacksHandle,
-      runAfterPass(HasNameRegex("ModuleToPostOrderCGSCCPassAdaptor"), _))
-      .Times(1);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPass(HasNameRegex("CGSCCToFunctionPassAdaptor"), _))
-      .Times(1);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPass(HasNameRegex("FunctionToLoopPassAdaptor"), _))
-      .Times(1);
-
-  // Ignore analyses introduced by adaptor passes.
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeAnalysis(Not(HasNameRegex("MockAnalysisHandle")), _))
-      .Times(AnyNumber());
-  EXPECT_CALL(CallbacksHandle,
-              runAfterAnalysis(Not(HasNameRegex("MockAnalysisHandle")), _))
-      .Times(AnyNumber());
-
-  // Register Funtion and Loop version of "test-transform" for testing
-  PB.registerPipelineParsingCallback(
-      [](StringRef Name, FunctionPassManager &FPM,
-         ArrayRef<PassBuilder::PipelineElement>) {
-        if (Name == "test-transform") {
-          FPM.addPass(MockPassHandle<Function>().getPass());
-          return true;
-        }
-        return false;
-      });
-  PB.registerPipelineParsingCallback(
-      [](StringRef Name, LoopPassManager &LPM,
-         ArrayRef<PassBuilder::PipelineElement>) {
-        if (Name == "test-transform") {
-          LPM.addPass(MockPassHandle<Loop>().getPass());
-          return true;
-        }
-        return false;
-      });
+  if (any_isa<const LazyCallGraph::SCC *>(IR)) {
+    const LazyCallGraph::SCC *C = any_cast<const LazyCallGraph::SCC *>(IR);
+    assert(C && "scc should be valid for printing");
+    std::string Extra = std::string(formatv(" (scc: {0})", C->getName()));
+    printIR(C, Banner, Extra);
+    return;
+
+  if (any_isa<const Loop *>(IR)) {
+    const Loop *L = any_cast<const Loop *>(IR);
+    assert(L && "Loop should be valid for printing");
+    printIR(L, Banner);
+    return;
+  }
+  llvm_unreachable("Unknown wrapped IR type");
 
-  StringRef PipelineText = "test-transform,function(test-transform),cgscc("
-                           "function(loop(test-transform)))";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-
-  PM.run(*M, AM);
-}
-
-TEST_F(FunctionCallbacksTest, Passes) {
-  EXPECT_CALL(AnalysisHandle, run(HasName("foo"), _));
-  EXPECT_CALL(PassHandle, run(HasName("foo"), _))
-      .WillOnce(Invoke(getAnalysisResult));
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(FunctionCallbacksTest, InstrumentedPasses) {
-  CallbacksHandle.registerPassInstrumentation();
-  // Non-mock instrumentation not specifically mentioned below can be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("foo");
-
-  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,
-              runBeforePass(HasNameRegex("MockPassHandle"), HasName("foo")))
-      .InSequence(PISequence);
-  EXPECT_CALL(
-      CallbacksHandle,
-      runBeforeNonSkippedPass(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);
-
-  // Our mock pass does not invalidate IR.
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPassInvalidated(HasNameRegex("MockPassHandle")))
-      .Times(0);
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(FunctionCallbacksTest, InstrumentedSkippedPasses) {
-  CallbacksHandle.registerPassInstrumentation();
-  // Non-mock instrumentation run here can safely be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("foo");
-
-  // Skip the pass by returning false.
-  EXPECT_CALL(CallbacksHandle,
-              runBeforePass(HasNameRegex("MockPassHandle"), HasName("foo")))
-      .WillOnce(Return(false));
-
-  EXPECT_CALL(AnalysisHandle, run(HasName("foo"), _)).Times(0);
-  EXPECT_CALL(PassHandle, run(HasName("foo"), _)).Times(0);
-
-  // As the pass is skipped there is no afterPass, beforeAnalysis/afterAnalysis
-  // as well.
-  EXPECT_CALL(CallbacksHandle, runAfterPass(HasNameRegex("MockPassHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPassInvalidated(HasNameRegex("MockPassHandle")))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle, runAfterPass(HasNameRegex("MockPassHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), _))
-      .Times(0);
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(LoopCallbacksTest, Passes) {
-  EXPECT_CALL(AnalysisHandle, run(HasName("loop"), _, _));
-  EXPECT_CALL(PassHandle, run(HasName("loop"), _, _, _))
-      .WillOnce(WithArgs<0, 1, 2>(Invoke(getAnalysisResult)));
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(LoopCallbacksTest, InstrumentedPasses) {
-  CallbacksHandle.registerPassInstrumentation();
-  // Non-mock instrumentation not specifically mentioned below can be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("foo");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("loop");
-
-  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,
-              runBeforePass(HasNameRegex("MockPassHandle"), HasName("loop")))
-      .InSequence(PISequence);
-  EXPECT_CALL(
-      CallbacksHandle,
-      runBeforeNonSkippedPass(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);
-
-  // Our mock pass does not invalidate IR.
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPassInvalidated(HasNameRegex("MockPassHandle")))
-      .Times(0);
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(LoopCallbacksTest, InstrumentedInvalidatingPasses) {
-  CallbacksHandle.registerPassInstrumentation();
-  // Non-mock instrumentation not specifically mentioned below can be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("foo");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("loop");
-
-  EXPECT_CALL(AnalysisHandle, run(HasName("loop"), _, _));
-  EXPECT_CALL(PassHandle, run(HasName("loop"), _, _, _))
-      .WillOnce(DoAll(WithArgs<0, 1, 2, 3>(Invoke(PassHandle.invalidateLoop)),
-                      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,
-              runBeforePass(HasNameRegex("MockPassHandle"), HasName("loop")))
-      .InSequence(PISequence);
-  EXPECT_CALL(
-      CallbacksHandle,
-      runBeforeNonSkippedPass(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,
-              runAfterPassInvalidated(HasNameRegex("MockPassHandle")))
-      .InSequence(PISequence);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPassInvalidated(HasNameRegex("^PassManager")))
-      .InSequence(PISequence);
-
-  // Our mock pass invalidates IR, thus normal runAfterPass is never called.
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPass(HasNameRegex("MockPassHandle"), HasName("loop")))
-      .Times(0);
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(LoopCallbacksTest, InstrumentedSkippedPasses) {
-  CallbacksHandle.registerPassInstrumentation();
-  // Non-mock instrumentation run here can safely be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("foo");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("loop");
-
-  // Skip the pass by returning false.
-  EXPECT_CALL(CallbacksHandle,
-              runBeforePass(HasNameRegex("MockPassHandle"), HasName("loop")))
-      .WillOnce(Return(false));
-
-  EXPECT_CALL(AnalysisHandle, run(HasName("loop"), _, _)).Times(0);
-  EXPECT_CALL(PassHandle, run(HasName("loop"), _, _, _)).Times(0);
-
-  // As the pass is skipped there is no afterPass, beforeAnalysis/afterAnalysis
-  // as well.
-  EXPECT_CALL(CallbacksHandle, runAfterPass(HasNameRegex("MockPassHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPassInvalidated(HasNameRegex("MockPassHandle")))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), _))
-      .Times(0);
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(CGSCCCallbacksTest, Passes) {
-  EXPECT_CALL(AnalysisHandle, run(HasName("(foo)"), _, _));
-  EXPECT_CALL(PassHandle, run(HasName("(foo)"), _, _, _))
-      .WillOnce(WithArgs<0, 1, 2>(Invoke(getAnalysisResult)));
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(CGSCCCallbacksTest, InstrumentedPasses) {
-  CallbacksHandle.registerPassInstrumentation();
-  // Non-mock instrumentation not specifically mentioned below can be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("foo");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("(foo)");
-
-  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,
-              runBeforePass(HasNameRegex("MockPassHandle"), HasName("(foo)")))
-      .InSequence(PISequence);
-  EXPECT_CALL(
-      CallbacksHandle,
-      runBeforeNonSkippedPass(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);
-
-  // Our mock pass does not invalidate IR.
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPassInvalidated(HasNameRegex("MockPassHandle")))
-      .Times(0);
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(CGSCCCallbacksTest, InstrumentedInvalidatingPasses) {
-  CallbacksHandle.registerPassInstrumentation();
-  // Non-mock instrumentation not specifically mentioned below can be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("foo");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("(foo)");
-
-  EXPECT_CALL(AnalysisHandle, run(HasName("(foo)"), _, _));
-  EXPECT_CALL(PassHandle, run(HasName("(foo)"), _, _, _))
-      .WillOnce(DoAll(WithArgs<0, 1, 2, 3>(Invoke(PassHandle.invalidateSCC)),
-                      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,
-              runBeforePass(HasNameRegex("MockPassHandle"), HasName("(foo)")))
-      .InSequence(PISequence);
-  EXPECT_CALL(
-      CallbacksHandle,
-      runBeforeNonSkippedPass(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,
-              runAfterPassInvalidated(HasNameRegex("MockPassHandle")))
-      .InSequence(PISequence);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPassInvalidated(HasNameRegex("^PassManager")))
-      .InSequence(PISequence);
-
-  // Our mock pass does invalidate IR, thus normal runAfterPass is never called.
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPass(HasNameRegex("MockPassHandle"), HasName("(foo)")))
-      .Times(0);
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-}
-
-TEST_F(CGSCCCallbacksTest, InstrumentedSkippedPasses) {
-  CallbacksHandle.registerPassInstrumentation();
-  // Non-mock instrumentation run here can safely be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("foo");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("(foo)");
-
-  // Skip the pass by returning false.
-  EXPECT_CALL(CallbacksHandle,
-              runBeforePass(HasNameRegex("MockPassHandle"), HasName("(foo)")))
-      .WillOnce(Return(false));
-
-  // neither Analysis nor Pass are called.
-  EXPECT_CALL(AnalysisHandle, run(HasName("(foo)"), _, _)).Times(0);
-  EXPECT_CALL(PassHandle, run(HasName("(foo)"), _, _, _)).Times(0);
-
-  // As the pass is skipped there is no afterPass, beforeAnalysis/afterAnalysis
-  // as well.
-  EXPECT_CALL(CallbacksHandle, runAfterPass(HasNameRegex("MockPassHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterPassInvalidated(HasNameRegex("MockPassHandle")))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), _))
-      .Times(0);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), _))
-      .Times(0);
-
-  StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
+} // namespace
+
+PrintIRInstrumentation::~PrintIRInstrumentation() {
+  assert(ModuleDescStack.empty() && "ModuleDescStack is not empty at exit");
 }
 
-/// Test parsing of the names of analysis utilities for our mock analysis
-/// for all IRUnits.
-///
-/// We first require<>, then invalidate<> it, expecting the analysis to be run
-/// once and subsequently invalidated.
-TEST_F(ModuleCallbacksTest, AnalysisUtilities) {
-  EXPECT_CALL(AnalysisHandle, run(HasName("<string>"), _));
-  EXPECT_CALL(AnalysisHandle, invalidate(HasName("<string>"), _, _));
+void PrintIRInstrumentation::pushModuleDesc(StringRef PassID, Any IR) {
+  assert(StoreModuleDesc);
+  const Module *M = nullptr;
+  std::string Extra;
+  if (auto UnwrappedModule = unwrapModule(IR))
+    std::tie(M, Extra) = UnwrappedModule.getValue();
+  ModuleDescStack.emplace_back(M, Extra, PassID);
+}
 
-  StringRef PipelineText = "require<test-analysis>,invalidate<test-analysis>";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
+PrintIRInstrumentation::PrintModuleDesc
+PrintIRInstrumentation::popModuleDesc(StringRef PassID) {
+  assert(!ModuleDescStack.empty() && "empty ModuleDescStack");
+  PrintModuleDesc ModuleDesc = ModuleDescStack.pop_back_val();
+  assert(std::get<2>(ModuleDesc).equals(PassID) && "malformed ModuleDescStack");
+  return ModuleDesc;
 }
 
-TEST_F(CGSCCCallbacksTest, PassUtilities) {
-  EXPECT_CALL(AnalysisHandle, run(HasName("(foo)"), _, _));
-  EXPECT_CALL(AnalysisHandle, invalidate(HasName("(foo)"), _, _));
+void PrintIRInstrumentation::printBeforePass(StringRef PassID, Any IR) {
+  if (PassID.startswith("PassManager<") || PassID.contains("PassAdaptor<"))
+    return;
+
+  // Saving Module for AfterPassInvalidated operations.
+  // Note: here we rely on a fact that we do not change modules while
+  // traversing the pipeline, so the latest captured module is good
+  // for all print operations that has not happen yet.
+  if (StoreModuleDesc && llvm::shouldPrintAfterPass(PassID))
+    pushModuleDesc(PassID, IR);
 
-  StringRef PipelineText = "require<test-analysis>,invalidate<test-analysis>";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
+  if (!llvm::shouldPrintBeforePass(PassID))
+    return;
+
+  SmallString<20> Banner = formatv("*** IR Dump Before {0} ***", PassID);
+  unwrapAndPrint(IR, Banner, llvm::forcePrintModuleIR());
+  return;
 }
 
-TEST_F(FunctionCallbacksTest, AnalysisUtilities) {
-  EXPECT_CALL(AnalysisHandle, run(HasName("foo"), _));
-  EXPECT_CALL(AnalysisHandle, invalidate(HasName("foo"), _, _));
+void PrintIRInstrumentation::printAfterPass(StringRef PassID, Any IR) {
+  if (PassID.startswith("PassManager<") || PassID.contains("PassAdaptor<"))
+    return;
+
+  if (!llvm::shouldPrintAfterPass(PassID))
+    return;
 
-  StringRef PipelineText = "require<test-analysis>,invalidate<test-analysis>";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
+  if (StoreModuleDesc)
+    popModuleDesc(PassID);
+
+  SmallString<20> Banner = formatv("*** IR Dump After {0} ***", PassID);
+  unwrapAndPrint(IR, Banner, llvm::forcePrintModuleIR());
 }
 
-TEST_F(LoopCallbacksTest, PassUtilities) {
-  EXPECT_CALL(AnalysisHandle, run(HasName("loop"), _, _));
-  EXPECT_CALL(AnalysisHandle, invalidate(HasName("loop"), _, _));
+void PrintIRInstrumentation::printAfterPassInvalidated(StringRef PassID) {
+  if (!StoreModuleDesc || !llvm::shouldPrintAfterPass(PassID))
+    return;
+
+  if (PassID.startswith("PassManager<") || PassID.contains("PassAdaptor<"))
+    return;
 
-  StringRef PipelineText = "require<test-analysis>,invalidate<test-analysis>";
+  const Module *M;
+  std::string Extra;
+  StringRef StoredPassID;
+  std::tie(M, Extra, StoredPassID) = popModuleDesc(PassID);
+  // Additional filtering (e.g. -filter-print-func) can lead to module
+  // printing being skipped.
+  if (!M)
+    return;
 
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
+  SmallString<20> Banner =
+      formatv("*** IR Dump After {0} *** invalidated: ", PassID);
+  printIR(M, Banner, Extra);
 }
 
-/// Test parsing of the top-level pipeline.
-///
-/// The ParseTopLevelPipeline callback takes over parsing of the entire pipeline
-/// from PassBuilder if it encounters an unknown pipeline entry at the top level
-/// (i.e., the first entry on the pipeline).
-/// This test parses a pipeline named 'another-pipeline', whose only elements
-/// may be the test-transform pass or the analysis utilities
-TEST_F(ModuleCallbacksTest, ParseTopLevelPipeline) {
-  PB.registerParseTopLevelPipelineCallback([this](
-      ModulePassManager &MPM, ArrayRef<PassBuilder::PipelineElement> Pipeline,
-      bool VerifyEachPass, bool DebugLogging) {
-    auto &FirstName = Pipeline.front().Name;
-    auto &InnerPipeline = Pipeline.front().InnerPipeline;
-    if (FirstName == "another-pipeline") {
-      for (auto &E : InnerPipeline) {
-        if (parseAnalysisUtilityPasses<AnalysisT>("test-analysis", E.Name, PM))
-          continue;
-
-        if (E.Name == "test-transform") {
-          PM.addPass(PassHandle.getPass());
-          continue;
-        }
-        return false;
-      }
-    }
+void PrintIRInstrumentation::registerCallbacks(
+    PassInstrumentationCallbacks &PIC) {
+  // BeforePass callback is not just for printing, it also saves a Module
+  // for later use in AfterPassInvalidated.
+  StoreModuleDesc = llvm::forcePrintModuleIR() && llvm::shouldPrintAfterPass();
+  if (llvm::shouldPrintBeforePass() || StoreModuleDesc)
+    PIC.registerBeforeNonSkippedPassCallback(
+        [this](StringRef P, Any IR) { this->printBeforePass(P, IR); });
+
+  if (llvm::shouldPrintAfterPass()) {
+    PIC.registerAfterPassCallback(
+        [this](StringRef P, Any IR, const PreservedAnalyses &) {
+          this->printAfterPass(P, IR);
+        });
+    PIC.registerAfterPassInvalidatedCallback(
+        [this](StringRef P, const PreservedAnalyses &) {
+          this->printAfterPassInvalidated(P);
+        });
+
+void OptNoneInstrumentation::registerCallbacks(
+    PassInstrumentationCallbacks &PIC) {
+  PIC.registerBeforePassCallback(
+      [this](StringRef P, Any IR) { return this->skip(P, IR); });
+
+bool OptNoneInstrumentation::skip(StringRef PassID, Any IR) {
+  if (!EnableOptnone)
     return true;
-  });
+  const Function *F = nullptr;
+  if (any_isa<const Function *>(IR)) {
+    F = any_cast<const Function *>(IR);
+  } else if (any_isa<const Loop *>(IR)) {
+    F = any_cast<const Loop *>(IR)->getHeader()->getParent();
+  }
+  if (F && F->hasOptNone()) {
+    if (DebugLogging)
+      dbgs() << "Skipping pass: " << PassID << " (optnone)\n";
+    return false;
+  }
+  return true;
+}
 
-  EXPECT_CALL(AnalysisHandle, run(HasName("<string>"), _));
-  EXPECT_CALL(PassHandle, run(HasName("<string>"), _))
-      .WillOnce(Invoke(getAnalysisResult));
-  EXPECT_CALL(AnalysisHandle, invalidate(HasName("<string>"), _, _));
-
-  StringRef PipelineText =
-      "another-pipeline(test-transform,invalidate<test-analysis>)";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Succeeded())
-      << "Pipeline was: " << PipelineText;
-  PM.run(*M, AM);
-
-  /// Test the negative case
-  PipelineText = "another-pipeline(instcombine)";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText, true), Failed())
-      << "Pipeline was: " << PipelineText;
+void StandardInstrumentations::registerCallbacks(
+    PassInstrumentationCallbacks &PIC) {
+  PrintIR.registerCallbacks(PIC);
+  TimePasses.registerCallbacks(PIC);
+  OptNone.registerCallbacks(PIC);
 }
-} // end anonymous namespace

llvm/unittests/IR/TimePassesTest.cpp

TEST(TimePassesTest, CustomOut) {
   std::unique_ptr<TimePassesHandler> TimePasses =
       std::make_unique<TimePassesHandler>(true);
   TimePasses->setOutStream(ReportStream);
   TimePasses->registerCallbacks(PIC);
 
   // Pretending that passes are running to trigger the timers.
   PI.runBeforePass(Pass1, M);
   PI.runBeforePass(Pass2, M);
-  PI.runAfterPass(Pass2, M);
-  PI.runAfterPass(Pass1, M);
+  PI.runAfterPass(Pass2, M, PreservedAnalyses::all());
+  PI.runAfterPass(Pass1, M, PreservedAnalyses::all());
 
   // Generating report.
   TimePasses->print();
 
   // There should be Pass1 and Pass2 in the report
   EXPECT_FALSE(TimePassesStr.empty());
   EXPECT_TRUE(TimePassesStr.str().contains("report"));
   EXPECT_TRUE(TimePassesStr.str().contains("Pass1"));
   EXPECT_TRUE(TimePassesStr.str().contains("Pass2"));
 
   // Clear and generate report again.
   TimePassesStr.clear();
   TimePasses->print();
   // Since we did not run any passes since last print, report should be empty.
   EXPECT_TRUE(TimePassesStr.empty());
 
   // Now trigger just a single pass to populate timers again.
   PI.runBeforePass(Pass2, M);
-  PI.runAfterPass(Pass2, M);
+  PI.runAfterPass(Pass2, M, PreservedAnalyses::all());
 
   // Generate report by deleting the handler.
   TimePasses.reset();
 
   // There should be Pass2 in this report and no Pass1.
   EXPECT_FALSE(TimePassesStr.str().empty());
   EXPECT_TRUE(TimePassesStr.str().contains("report"));
   EXPECT_FALSE(TimePassesStr.str().contains("Pass1"));
   EXPECT_TRUE(TimePassesStr.str().contains("Pass2"));
 }
 
 } // end anonymous namespace