[llvm] [NewPM/CodeGen] Rewrite pass manager nesting (PR #81068)

Arthur Eubanks via llvm-commits llvm-commits at lists.llvm.org
Wed Feb 14 22:59:09 PST 2024


https://github.com/aeubanks updated https://github.com/llvm/llvm-project/pull/81068

>From 3a5f3cbefc8a55e4af7fc57c86745a875bd65e53 Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks at google.com>
Date: Fri, 2 Feb 2024 04:06:28 +0000
Subject: [PATCH 1/2] [NewPM/CodeGen] Rewrite pass manager nesting

Currently the new PM infra for codegen puts everything into a MachineFunctionPassManager. The MachineFunctionPassManager owns both Module passes and MachineFunction passes, and batches adjacent MachineFunction passes like a typical PassManager.

The current MachineFunctionAnalysisManager also directly references a module and function analysis manager to get results.

The initial argument was that the codegen pipeline is relatively "flat", meaning it's mostly machine function passes with a couple of module passes here and there. However, there are a couple of issues with this as compared to a more structured nesting more like the optimization pipeline. For example, it doesn't allow running function passes then machine function passes on a function and its machine function all at once. It also currently requires the caller to split out the IR passes into one pass manager and the MIR passes into another pass manager.

This patch rewrites the new pass manager infra for the codegen pipeline to be more similar to the nesting in the optimization pipeline. Basically, a Function contains a MachineFunction. So we can have Module -> Function -> MachineFunction adaptors. It also rewrites the analysis managers to have inner/outer proxies like the ones in the optimization pipeline. The new pass managers/adaptors/analysis managers can be seen in use in PassManagerTest.cpp.

This allows us to consolidate to just having to add to one ModulePassManager when using the codegen pipeline.

I haven't added the Function -> MachineFunction adaptor in this patch, but it should be added when we merge AddIRPass/AddMachinePass so that we can run IR and MIR passes on a function before proceeding to the next function.

The MachineFunctionProperties infra for MIR verification is still WIP.
---
 .../include/llvm/CodeGen/MachinePassManager.h | 405 ++++++++++--------
 llvm/include/llvm/Passes/CodeGenPassBuilder.h |  82 ++--
 llvm/include/llvm/Passes/PassBuilder.h        |  15 +-
 llvm/include/llvm/Target/TargetMachine.h      |  10 +-
 llvm/lib/CodeGen/MachinePassManager.cpp       | 183 ++++----
 llvm/lib/Passes/PassBuilder.cpp               |  48 ++-
 llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp |   9 +-
 llvm/lib/Target/X86/X86TargetMachine.h        |   7 +-
 llvm/test/tools/llc/new-pm/pipeline.mir       |   3 +-
 llvm/test/tools/llc/new-pm/start-stop.ll      |   3 +-
 llvm/tools/llc/NewPMDriver.cpp                |  87 ++--
 llvm/unittests/CodeGen/PassManagerTest.cpp    | 213 +++------
 .../MIR/PassBuilderCallbacksTest.cpp          | 216 ++++++----
 13 files changed, 666 insertions(+), 615 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/MachinePassManager.h b/llvm/include/llvm/CodeGen/MachinePassManager.h
index a0ad7d7a95a289..7713c55661ccc1 100644
--- a/llvm/include/llvm/CodeGen/MachinePassManager.h
+++ b/llvm/include/llvm/CodeGen/MachinePassManager.h
@@ -25,17 +25,18 @@
 
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/PassManagerInternal.h"
 #include "llvm/Support/Error.h"
 
-#include <map>
-
 namespace llvm {
 class Module;
 class Function;
 class MachineFunction;
 
 extern template class AnalysisManager<MachineFunction>;
+using MachineFunctionAnalysisManager = AnalysisManager<MachineFunction>;
 
 /// A CRTP mix-in that provides informational APIs needed for machine passes.
 ///
@@ -46,217 +47,247 @@ struct MachinePassInfoMixin : public PassInfoMixin<DerivedT> {
   // TODO: Add MachineFunctionProperties support.
 };
 
-/// An AnalysisManager<MachineFunction> that also exposes IR analysis results.
-class MachineFunctionAnalysisManager : public AnalysisManager<MachineFunction> {
-public:
-  using Base = AnalysisManager<MachineFunction>;
+namespace detail {
+struct MachinePassConcept
+    : PassConcept<MachineFunction, MachineFunctionAnalysisManager> {
+  virtual MachineFunctionProperties getRequiredProperties() const = 0;
+  virtual MachineFunctionProperties getSetProperties() const = 0;
+  virtual MachineFunctionProperties getClearedProperties() const = 0;
+};
 
-  MachineFunctionAnalysisManager() : FAM(nullptr), MAM(nullptr) {}
-  MachineFunctionAnalysisManager(FunctionAnalysisManager &FAM,
-                                 ModuleAnalysisManager &MAM)
-      : FAM(&FAM), MAM(&MAM) {}
-  MachineFunctionAnalysisManager(MachineFunctionAnalysisManager &&) = default;
-  MachineFunctionAnalysisManager &
-  operator=(MachineFunctionAnalysisManager &&) = default;
+template <typename PassT> struct MachinePassModel : MachinePassConcept {
+  explicit MachinePassModel(PassT Pass) : Pass(std::move(Pass)) {}
+  // We have to explicitly define all the special member functions because MSVC
+  // refuses to generate them.
+  MachinePassModel(const MachinePassModel &Arg) : Pass(Arg.Pass) {}
+  MachinePassModel(MachinePassModel &&Arg) : Pass(std::move(Arg.Pass)) {}
 
-  /// Get the result of an analysis pass for a Function.
-  ///
-  /// Runs the analysis if a cached result is not available.
-  template <typename PassT> typename PassT::Result &getResult(Function &F) {
-    return FAM->getResult<PassT>(F);
+  friend void swap(MachinePassModel &LHS, MachinePassModel &RHS) {
+    using std::swap;
+    swap(LHS.Pass, RHS.Pass);
   }
 
-  /// Get the cached result of an analysis pass for a Function.
-  ///
-  /// This method never runs the analysis.
-  ///
-  /// \returns null if there is no cached result.
-  template <typename PassT>
-  typename PassT::Result *getCachedResult(Function &F) {
-    return FAM->getCachedResult<PassT>(F);
+  MachinePassModel &operator=(MachinePassModel RHS) {
+    swap(*this, RHS);
+    return *this;
   }
 
-  /// Get the result of an analysis pass for a Module.
-  ///
-  /// Runs the analysis if a cached result is not available.
-  template <typename PassT> typename PassT::Result &getResult(Module &M) {
-    return MAM->getResult<PassT>(M);
+  PreservedAnalyses run(MachineFunction &IR,
+                        MachineFunctionAnalysisManager &AM) override {
+    return Pass.run(IR, AM);
   }
 
-  /// Get the cached result of an analysis pass for a Module.
-  ///
-  /// This method never runs the analysis.
-  ///
-  /// \returns null if there is no cached result.
-  template <typename PassT> typename PassT::Result *getCachedResult(Module &M) {
-    return MAM->getCachedResult<PassT>(M);
+  void printPipeline(
+      raw_ostream &OS,
+      function_ref<StringRef(StringRef)> MapClassName2PassName) override {
+    Pass.printPipeline(OS, MapClassName2PassName);
   }
 
-  /// Get the result of an analysis pass for a MachineFunction.
-  ///
-  /// Runs the analysis if a cached result is not available.
-  using Base::getResult;
+  StringRef name() const override { return PassT::name(); }
 
-  /// Get the cached result of an analysis pass for a MachineFunction.
-  ///
-  /// This method never runs the analysis.
-  ///
-  /// returns null if there is no cached result.
-  using Base::getCachedResult;
-
-  // FIXME: Add LoopAnalysisManager or CGSCCAnalysisManager if needed.
-  FunctionAnalysisManager *FAM;
-  ModuleAnalysisManager *MAM;
-};
+  template <typename T>
+  using has_required_t = decltype(std::declval<T &>().isRequired());
+  template <typename T>
+  static std::enable_if_t<is_detected<has_required_t, T>::value, bool>
+  passIsRequiredImpl() {
+    return T::isRequired();
+  }
+  template <typename T>
+  static std::enable_if_t<!is_detected<has_required_t, T>::value, bool>
+  passIsRequiredImpl() {
+    return false;
+  }
+  bool isRequired() const override { return passIsRequiredImpl<PassT>(); }
+
+  template <typename T>
+  using has_get_required_properties_t =
+      decltype(std::declval<T &>().getRequiredProperties());
+  template <typename T>
+  static std::enable_if_t<is_detected<has_get_required_properties_t, T>::value,
+                          MachineFunctionProperties>
+  getRequiredPropertiesImpl() {
+    return PassT::getRequiredProperties();
+  }
+  template <typename T>
+  static std::enable_if_t<!is_detected<has_get_required_properties_t, T>::value,
+                          MachineFunctionProperties>
+  getRequiredPropertiesImpl() {
+    return MachineFunctionProperties();
+  }
+  MachineFunctionProperties getRequiredProperties() const override {
+    return getRequiredPropertiesImpl<PassT>();
+  }
 
-extern template class PassManager<MachineFunction>;
+  template <typename T>
+  using has_get_set_properties_t =
+      decltype(std::declval<T &>().getSetProperties());
+  template <typename T>
+  static std::enable_if_t<is_detected<has_get_set_properties_t, T>::value,
+                          MachineFunctionProperties>
+  getSetPropertiesImpl() {
+    return PassT::getSetProperties();
+  }
+  template <typename T>
+  static std::enable_if_t<!is_detected<has_get_set_properties_t, T>::value,
+                          MachineFunctionProperties>
+  getSetPropertiesImpl() {
+    return MachineFunctionProperties();
+  }
+  MachineFunctionProperties getSetProperties() const override {
+    return getSetPropertiesImpl<PassT>();
+  }
 
-/// MachineFunctionPassManager adds/removes below features to/from the base
-/// PassManager template instantiation.
-///
-/// - Support passes that implement doInitialization/doFinalization. This is for
-///   machine function passes to work on module level constructs. One such pass
-///   is AsmPrinter.
-///
-/// - Support machine module pass which runs over the module (for example,
-///   MachineOutliner). A machine module pass needs to define the method:
-///
-///   ```Error run(Module &, MachineFunctionAnalysisManager &)```
-///
-///   FIXME: machine module passes still need to define the usual machine
-///          function pass interface, namely,
-///          `PreservedAnalyses run(MachineFunction &,
-///                                 MachineFunctionAnalysisManager &)`
-///          But this interface wouldn't be executed. It is just a placeholder
-///          to satisfy the pass manager type-erased inteface. This
-///          special-casing of machine module pass is due to its limited use
-///          cases and the unnecessary complexity it may bring to the machine
-///          pass manager.
-///
-/// - The base class `run` method is replaced by an alternative `run` method.
-///   See details below.
-///
-/// - Support codegening in the SCC order. Users include interprocedural
-///   register allocation (IPRA).
-class MachineFunctionPassManager
-    : public PassManager<MachineFunction, MachineFunctionAnalysisManager> {
-  using Base = PassManager<MachineFunction, MachineFunctionAnalysisManager>;
+  template <typename T>
+  using has_get_cleared_properties_t =
+      decltype(std::declval<T &>().getClearedProperties());
+  template <typename T>
+  static std::enable_if_t<is_detected<has_get_cleared_properties_t, T>::value,
+                          MachineFunctionProperties>
+  getClearedPropertiesImpl() {
+    return PassT::getClearedProperties();
+  }
+  template <typename T>
+  static std::enable_if_t<!is_detected<has_get_cleared_properties_t, T>::value,
+                          MachineFunctionProperties>
+  getClearedPropertiesImpl() {
+    return MachineFunctionProperties();
+  }
+  MachineFunctionProperties getClearedProperties() const override {
+    return getClearedPropertiesImpl<PassT>();
+  }
 
+  PassT Pass;
+};
+} // namespace detail
+
+using MachineFunctionAnalysisManagerModuleProxy =
+    InnerAnalysisManagerProxy<MachineFunctionAnalysisManager, Module>;
+
+template <>
+bool MachineFunctionAnalysisManagerModuleProxy::Result::invalidate(
+    Module &M, const PreservedAnalyses &PA,
+    ModuleAnalysisManager::Invalidator &Inv);
+extern template class InnerAnalysisManagerProxy<MachineFunctionAnalysisManager,
+                                                Module>;
+
+extern template class OuterAnalysisManagerProxy<ModuleAnalysisManager,
+                                                MachineFunction>;
+/// Provide the \c ModuleAnalysisManager to \c Function proxy.
+using ModuleAnalysisManagerMachineFunctionProxy =
+    OuterAnalysisManagerProxy<ModuleAnalysisManager, MachineFunction>;
+
+class FunctionAnalysisManagerMachineFunctionProxy
+    : public AnalysisInfoMixin<FunctionAnalysisManagerMachineFunctionProxy> {
 public:
-  MachineFunctionPassManager(bool RequireCodeGenSCCOrder = false,
-                             bool VerifyMachineFunction = false)
-      : RequireCodeGenSCCOrder(RequireCodeGenSCCOrder),
-        VerifyMachineFunction(VerifyMachineFunction) {}
-  MachineFunctionPassManager(MachineFunctionPassManager &&) = default;
-  MachineFunctionPassManager &
-  operator=(MachineFunctionPassManager &&) = default;
-
-  /// Run machine passes for a Module.
+  class Result {
+  public:
+    explicit Result(FunctionAnalysisManager &FAM) : FAM(&FAM) {}
+
+    Result(Result &&Arg) : FAM(std::move(Arg.FAM)) {
+      // We have to null out the analysis manager in the moved-from state
+      // because we are taking ownership of the responsibilty to clear the
+      // analysis state.
+      Arg.FAM = nullptr;
+    }
+
+    ~Result() {
+      // FAM is cleared in a moved from state where there is nothing to do.
+      if (!FAM)
+        return;
+
+      // Clear out the analysis manager if we're being destroyed -- it means we
+      // didn't even see an invalidate call when we got invalidated.
+      FAM->clear();
+    }
+
+    Result &operator=(Result &&RHS) {
+      FAM = RHS.FAM;
+      // We have to null out the analysis manager in the moved-from state
+      // because we are taking ownership of the responsibilty to clear the
+      // analysis state.
+      RHS.FAM = nullptr;
+      return *this;
+    }
+
+    /// Accessor for the analysis manager.
+    FunctionAnalysisManager &getManager() { return *FAM; }
+
+    /// Handler for invalidation of the outer IR unit, \c IRUnitT.
+    ///
+    /// If the proxy analysis itself is not preserved, we assume that the set of
+    /// inner IR objects contained in IRUnit may have changed.  In this case,
+    /// we have to call \c clear() on the inner analysis manager, as it may now
+    /// have stale pointers to its inner IR objects.
+    ///
+    /// Regardless of whether the proxy analysis is marked as preserved, all of
+    /// the analyses in the inner analysis manager are potentially invalidated
+    /// based on the set of preserved analyses.
+    bool invalidate(MachineFunction &IR, const PreservedAnalyses &PA,
+                    MachineFunctionAnalysisManager::Invalidator &Inv);
+
+  private:
+    FunctionAnalysisManager *FAM;
+  };
+
+  explicit FunctionAnalysisManagerMachineFunctionProxy(
+      FunctionAnalysisManager &FAM)
+      : FAM(&FAM) {}
+
+  /// Run the analysis pass and create our proxy result object.
   ///
-  /// The intended use is to start the codegen pipeline for a Module. The base
-  /// class's `run` method is deliberately hidden by this due to the observation
-  /// that we don't yet have the use cases of compositing two instances of
-  /// machine pass managers, or compositing machine pass managers with other
-  /// types of pass managers.
-  Error run(Module &M, MachineFunctionAnalysisManager &MFAM);
-
-  template <typename PassT> void addPass(PassT &&Pass) {
-    Base::addPass(std::forward<PassT>(Pass));
-    PassConceptT *P = Passes.back().get();
-    addDoInitialization<PassT>(P);
-    addDoFinalization<PassT>(P);
-
-    // Add machine module pass.
-    addRunOnModule<PassT>(P);
+  /// This doesn't do any interesting work; it is primarily used to insert our
+  /// proxy result object into the outer analysis cache so that we can proxy
+  /// invalidation to the inner analysis manager.
+  Result run(MachineFunction &, MachineFunctionAnalysisManager &) {
+    return Result(*FAM);
   }
 
-private:
-  template <typename PassT>
-  using has_init_t = decltype(std::declval<PassT &>().doInitialization(
-      std::declval<Module &>(),
-      std::declval<MachineFunctionAnalysisManager &>()));
-
-  template <typename PassT>
-  std::enable_if_t<!is_detected<has_init_t, PassT>::value>
-  addDoInitialization(PassConceptT *Pass) {}
-
-  template <typename PassT>
-  std::enable_if_t<is_detected<has_init_t, PassT>::value>
-  addDoInitialization(PassConceptT *Pass) {
-    using PassModelT = detail::PassModel<MachineFunction, PassT,
-                                         MachineFunctionAnalysisManager>;
-    auto *P = static_cast<PassModelT *>(Pass);
-    InitializationFuncs.emplace_back(
-        [=](Module &M, MachineFunctionAnalysisManager &MFAM) {
-          return P->Pass.doInitialization(M, MFAM);
-        });
-  }
+  static AnalysisKey Key;
 
-  template <typename PassT>
-  using has_fini_t = decltype(std::declval<PassT &>().doFinalization(
-      std::declval<Module &>(),
-      std::declval<MachineFunctionAnalysisManager &>()));
-
-  template <typename PassT>
-  std::enable_if_t<!is_detected<has_fini_t, PassT>::value>
-  addDoFinalization(PassConceptT *Pass) {}
-
-  template <typename PassT>
-  std::enable_if_t<is_detected<has_fini_t, PassT>::value>
-  addDoFinalization(PassConceptT *Pass) {
-    using PassModelT = detail::PassModel<MachineFunction, PassT,
-                                         MachineFunctionAnalysisManager>;
-    auto *P = static_cast<PassModelT *>(Pass);
-    FinalizationFuncs.emplace_back(
-        [=](Module &M, MachineFunctionAnalysisManager &MFAM) {
-          return P->Pass.doFinalization(M, MFAM);
-        });
-  }
+private:
+  FunctionAnalysisManager *FAM;
+};
 
-  template <typename PassT>
-  using is_machine_module_pass_t = decltype(std::declval<PassT &>().run(
-      std::declval<Module &>(),
-      std::declval<MachineFunctionAnalysisManager &>()));
-
-  template <typename PassT>
-  using is_machine_function_pass_t = decltype(std::declval<PassT &>().run(
-      std::declval<MachineFunction &>(),
-      std::declval<MachineFunctionAnalysisManager &>()));
-
-  template <typename PassT>
-  std::enable_if_t<!is_detected<is_machine_module_pass_t, PassT>::value>
-  addRunOnModule(PassConceptT *Pass) {}
-
-  template <typename PassT>
-  std::enable_if_t<is_detected<is_machine_module_pass_t, PassT>::value>
-  addRunOnModule(PassConceptT *Pass) {
-    static_assert(is_detected<is_machine_function_pass_t, PassT>::value,
-                  "machine module pass needs to define machine function pass "
-                  "api. sorry.");
-
-    using PassModelT = detail::PassModel<MachineFunction, PassT,
-                                         MachineFunctionAnalysisManager>;
-    auto *P = static_cast<PassModelT *>(Pass);
-    MachineModulePasses.emplace(
-        Passes.size() - 1,
-        [=](Module &M, MachineFunctionAnalysisManager &MFAM) {
-          return P->Pass.run(M, MFAM);
-        });
-  }
+class ModuleToMachineFunctionPassAdaptor
+    : public PassInfoMixin<ModuleToMachineFunctionPassAdaptor> {
+  using MachinePassConcept = detail::MachinePassConcept;
 
-  using FuncTy = Error(Module &, MachineFunctionAnalysisManager &);
-  SmallVector<llvm::unique_function<FuncTy>, 4> InitializationFuncs;
-  SmallVector<llvm::unique_function<FuncTy>, 4> FinalizationFuncs;
+public:
+  explicit ModuleToMachineFunctionPassAdaptor(
+      std::unique_ptr<MachinePassConcept> Pass)
+      : Pass(std::move(Pass)) {}
 
-  using PassIndex = decltype(Passes)::size_type;
-  std::map<PassIndex, llvm::unique_function<FuncTy>> MachineModulePasses;
+  /// Runs the function pass across every function in the module.
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+  void printPipeline(raw_ostream &OS,
+                     function_ref<StringRef(StringRef)> MapClassName2PassName);
 
-  // Run codegen in the SCC order.
-  bool RequireCodeGenSCCOrder;
+  static bool isRequired() { return true; }
 
-  bool VerifyMachineFunction;
+private:
+  std::unique_ptr<MachinePassConcept> Pass;
 };
 
+template <typename MachineFunctionPassT>
+ModuleToMachineFunctionPassAdaptor
+createModuleToMachineFunctionPassAdaptor(MachineFunctionPassT &&Pass) {
+  using PassModelT = detail::MachinePassModel<MachineFunctionPassT>;
+  // Do not use make_unique, it causes too many template instantiations,
+  // causing terrible compile times.
+  return ModuleToMachineFunctionPassAdaptor(
+      std::unique_ptr<detail::MachinePassConcept>(
+          new PassModelT(std::forward<MachineFunctionPassT>(Pass))));
+}
+
+template <>
+PreservedAnalyses
+PassManager<MachineFunction>::run(MachineFunction &,
+                                  AnalysisManager<MachineFunction> &);
+extern template class PassManager<MachineFunction>;
+
+/// Convenience typedef for a pass manager over functions.
+using MachineFunctionPassManager = PassManager<MachineFunction>;
+
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_MACHINEPASSMANAGER_H
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index fa6dbd4a49730e..7bb6c7ced0f907 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -89,12 +89,8 @@ namespace llvm {
 #define DUMMY_MACHINE_MODULE_PASS(NAME, PASS_NAME)                             \
   struct PASS_NAME : public MachinePassInfoMixin<PASS_NAME> {                  \
     template <typename... Ts> PASS_NAME(Ts &&...) {}                           \
-    Error run(Module &, MachineFunctionAnalysisManager &) {                    \
-      return Error::success();                                                 \
-    }                                                                          \
-    PreservedAnalyses run(MachineFunction &,                                   \
-                          MachineFunctionAnalysisManager &) {                  \
-      llvm_unreachable("this api is to make new PM api happy");                \
+    PreservedAnalyses run(Module &, ModuleAnalysisManager &) {                 \
+      return PreservedAnalyses::all();                                         \
     }                                                                          \
   };
 #define DUMMY_MACHINE_FUNCTION_PASS(NAME, PASS_NAME)                           \
@@ -133,17 +129,19 @@ template <typename DerivedT> class CodeGenPassBuilder {
       Opt.OptimizeRegAlloc = getOptLevel() != CodeGenOptLevel::None;
   }
 
-  Error buildPipeline(ModulePassManager &MPM, MachineFunctionPassManager &MFPM,
-                      raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
+  Error buildPipeline(ModulePassManager &MPM, raw_pwrite_stream &Out,
+                      raw_pwrite_stream *DwoOut,
                       CodeGenFileType FileType) const;
 
   void registerModuleAnalyses(ModuleAnalysisManager &) const;
   void registerFunctionAnalyses(FunctionAnalysisManager &) const;
   void registerMachineFunctionAnalyses(MachineFunctionAnalysisManager &) const;
 
-  void registerAnalyses(MachineFunctionAnalysisManager &MFAM) const {
-    registerModuleAnalyses(*MFAM.MAM);
-    registerFunctionAnalyses(*MFAM.FAM);
+  void registerAnalyses(ModuleAnalysisManager &MAM,
+                        FunctionAnalysisManager &FAM,
+                        MachineFunctionAnalysisManager &MFAM) const {
+    registerModuleAnalyses(MAM);
+    registerFunctionAnalyses(FAM);
     registerMachineFunctionAnalyses(MFAM);
   }
 
@@ -160,7 +158,15 @@ template <typename DerivedT> class CodeGenPassBuilder {
   using is_function_pass_t = decltype(std::declval<PassT &>().run(
       std::declval<Function &>(), std::declval<FunctionAnalysisManager &>()));
 
+  template <typename PassT>
+  using is_machine_function_pass_t = decltype(std::declval<PassT &>().run(
+      std::declval<MachineFunction &>(),
+      std::declval<MachineFunctionAnalysisManager &>()));
+
   // Function object to maintain state while adding codegen IR passes.
+  // TODO: add a Function -> MachineFunction adaptor and merge
+  // AddIRPass/AddMachinePass so we can have a function pipeline that runs both
+  // function passes and machine function passes.
   class AddIRPass {
   public:
     AddIRPass(ModulePassManager &MPM, const DerivedT &PB) : MPM(MPM), PB(PB) {}
@@ -207,31 +213,46 @@ template <typename DerivedT> class CodeGenPassBuilder {
   // Function object to maintain state while adding codegen machine passes.
   class AddMachinePass {
   public:
-    AddMachinePass(MachineFunctionPassManager &PM, const DerivedT &PB)
-        : PM(PM), PB(PB) {}
+    AddMachinePass(ModulePassManager &MPM, const DerivedT &PB)
+        : MPM(MPM), PB(PB) {}
+    ~AddMachinePass() {
+      if (!MFPM.isEmpty())
+        MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
+    }
 
-    template <typename PassT> void operator()(PassT &&Pass) {
-      if (!PB.runBeforeAdding(PassT::name()))
+    template <typename PassT>
+    void operator()(PassT &&Pass, StringRef Name = PassT::name()) {
+      static_assert((is_detected<is_machine_function_pass_t, PassT>::value ||
+                     is_detected<is_module_pass_t, PassT>::value) &&
+                    "Only module pass and function pass are supported.");
+
+      if (!PB.runBeforeAdding(Name))
         return;
 
-      PM.addPass(std::forward<PassT>(Pass));
+      // Add Function Pass
+      if constexpr (is_detected<is_machine_function_pass_t, PassT>::value) {
+        MFPM.addPass(std::forward<PassT>(Pass));
 
-      for (auto &C : PB.AfterCallbacks)
-        C(PassT::name());
-    }
+        for (auto &C : PB.AfterCallbacks)
+          C(Name);
+      } else {
+        // Add Module Pass
+        if (!MFPM.isEmpty()) {
+          MPM.addPass(
+              createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
+          MFPM = MachineFunctionPassManager();
+        }
 
-    template <typename PassT> void insertPass(StringRef PassName, PassT Pass) {
-      PB.AfterCallbacks.emplace_back(
-          [this, PassName, Pass = std::move(Pass)](StringRef Name) {
-            if (PassName == Name)
-              this->PM.addPass(std::move(Pass));
-          });
-    }
+        MPM.addPass(std::forward<PassT>(Pass));
 
-    MachineFunctionPassManager releasePM() { return std::move(PM); }
+        for (auto &C : PB.AfterCallbacks)
+          C(Name);
+      }
+    }
 
   private:
-    MachineFunctionPassManager &PM;
+    ModulePassManager &MPM;
+    MachineFunctionPassManager MFPM;
     const DerivedT &PB;
   };
 
@@ -486,8 +507,7 @@ template <typename DerivedT> class CodeGenPassBuilder {
 
 template <typename Derived>
 Error CodeGenPassBuilder<Derived>::buildPipeline(
-    ModulePassManager &MPM, MachineFunctionPassManager &MFPM,
-    raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
+    ModulePassManager &MPM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
     CodeGenFileType FileType) const {
   auto StartStopInfo = TargetPassConfig::getStartStopInfo(*PIC);
   if (!StartStopInfo)
@@ -499,7 +519,7 @@ Error CodeGenPassBuilder<Derived>::buildPipeline(
   addIRPass(RequireAnalysisPass<CollectorMetadataAnalysis, Module>());
   addISelPasses(addIRPass);
 
-  AddMachinePass addPass(MFPM, derived());
+  AddMachinePass addPass(MPM, derived());
   if (auto Err = addCoreISelPasses(addPass))
     return std::move(Err);
 
diff --git a/llvm/include/llvm/Passes/PassBuilder.h b/llvm/include/llvm/Passes/PassBuilder.h
index 10c5b7c00bae33..6822cfdb4957b6 100644
--- a/llvm/include/llvm/Passes/PassBuilder.h
+++ b/llvm/include/llvm/Passes/PassBuilder.h
@@ -133,7 +133,8 @@ class PassBuilder {
   void crossRegisterProxies(LoopAnalysisManager &LAM,
                             FunctionAnalysisManager &FAM,
                             CGSCCAnalysisManager &CGAM,
-                            ModuleAnalysisManager &MAM);
+                            ModuleAnalysisManager &MAM,
+                            MachineFunctionAnalysisManager *MFAM = nullptr);
 
   /// Registers all available module analysis passes.
   ///
@@ -569,9 +570,9 @@ class PassBuilder {
     ModulePipelineParsingCallbacks.push_back(C);
   }
   void registerPipelineParsingCallback(
-      const std::function<bool(StringRef Name, MachineFunctionPassManager &)>
-          &C) {
-    MachinePipelineParsingCallbacks.push_back(C);
+      const std::function<bool(StringRef Name, MachineFunctionPassManager &,
+                               ArrayRef<PipelineElement>)> &C) {
+    MachineFunctionPipelineParsingCallbacks.push_back(C);
   }
   /// @}}
 
@@ -733,8 +734,10 @@ class PassBuilder {
   // Machine pass callbackcs
   SmallVector<std::function<void(MachineFunctionAnalysisManager &)>, 2>
       MachineFunctionAnalysisRegistrationCallbacks;
-  SmallVector<std::function<bool(StringRef, MachineFunctionPassManager &)>, 2>
-      MachinePipelineParsingCallbacks;
+  SmallVector<std::function<bool(StringRef, MachineFunctionPassManager &,
+                                 ArrayRef<PipelineElement>)>,
+              2>
+      MachineFunctionPipelineParsingCallbacks;
 };
 
 /// This utility template takes care of adding require<> and invalidate<>
diff --git a/llvm/include/llvm/Target/TargetMachine.h b/llvm/include/llvm/Target/TargetMachine.h
index 7462f61d32b566..d7ce088cad49f6 100644
--- a/llvm/include/llvm/Target/TargetMachine.h
+++ b/llvm/include/llvm/Target/TargetMachine.h
@@ -34,8 +34,6 @@ using ModulePassManager = PassManager<Module>;
 
 class Function;
 class GlobalValue;
-class MachineFunctionPassManager;
-class MachineFunctionAnalysisManager;
 class MachineModuleInfoWrapperPass;
 class Mangler;
 class MCAsmInfo;
@@ -455,11 +453,9 @@ class LLVMTargetMachine : public TargetMachine {
                       bool DisableVerify = true,
                       MachineModuleInfoWrapperPass *MMIWP = nullptr) override;
 
-  virtual Error buildCodeGenPipeline(ModulePassManager &,
-                                     MachineFunctionPassManager &,
-                                     MachineFunctionAnalysisManager &,
-                                     raw_pwrite_stream &, raw_pwrite_stream *,
-                                     CodeGenFileType, CGPassBuilderOption,
+  virtual Error buildCodeGenPipeline(ModulePassManager &, raw_pwrite_stream &,
+                                     raw_pwrite_stream *, CodeGenFileType,
+                                     CGPassBuilderOption,
                                      PassInstrumentationCallbacks *) {
     return make_error<StringError>("buildCodeGenPipeline is not overridden",
                                    inconvertibleErrorCode());
diff --git a/llvm/lib/CodeGen/MachinePassManager.cpp b/llvm/lib/CodeGen/MachinePassManager.cpp
index d42bbe239830f0..9a750b5bed4339 100644
--- a/llvm/lib/CodeGen/MachinePassManager.cpp
+++ b/llvm/lib/CodeGen/MachinePassManager.cpp
@@ -11,7 +11,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/MachinePassManager.h"
-#include "llvm/CodeGen/FreeMachineFunction.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/IR/PassManagerImpl.h"
@@ -19,99 +18,121 @@
 using namespace llvm;
 
 namespace llvm {
-template class AllAnalysesOn<MachineFunction>;
+
+AnalysisKey FunctionAnalysisManagerMachineFunctionProxy::Key;
+
 template class AnalysisManager<MachineFunction>;
 template class PassManager<MachineFunction>;
+template class InnerAnalysisManagerProxy<MachineFunctionAnalysisManager,
+                                         Module>;
+template class OuterAnalysisManagerProxy<ModuleAnalysisManager,
+                                         MachineFunction>;
+
+bool FunctionAnalysisManagerMachineFunctionProxy::Result::invalidate(
+    MachineFunction &IR, const PreservedAnalyses &PA,
+    MachineFunctionAnalysisManager::Invalidator &Inv) {
+  // MachineFunction passes should not invalidate Function analyses.
+  // TODO: verify that PA doesn't invalidate Function analyses.
+  return false;
+}
 
-Error MachineFunctionPassManager::run(Module &M,
-                                      MachineFunctionAnalysisManager &MFAM) {
-  // MachineModuleAnalysis is a module analysis pass that is never invalidated
-  // because we don't run any module pass in codegen pipeline. This is very
-  // important because the codegen state is stored in MMI which is the analysis
-  // result of MachineModuleAnalysis. MMI should not be recomputed.
-  auto &MMI = MFAM.getResult<MachineModuleAnalysis>(M).getMMI();
-
-  (void)RequireCodeGenSCCOrder;
-  assert(!RequireCodeGenSCCOrder && "not implemented");
-
-  // M is unused here
-  PassInstrumentation PI = MFAM.getResult<PassInstrumentationAnalysis>(M);
-
-  // Add a PIC to verify machine functions.
-  if (VerifyMachineFunction) {
-    // No need to pop this callback later since MIR pipeline is flat which means
-    // current pipeline is the top-level pipeline. Callbacks are not used after
-    // current pipeline.
-    PI.pushBeforeNonSkippedPassCallback([](StringRef PassID, Any IR) {
-      assert(llvm::any_cast<const MachineFunction *>(&IR));
-      const MachineFunction *MF = llvm::any_cast<const MachineFunction *>(IR);
-      assert(MF && "Machine function should be valid for printing");
-      std::string Banner = std::string("After ") + std::string(PassID);
-      verifyMachineFunction(Banner, *MF);
-    });
+template <>
+bool MachineFunctionAnalysisManagerModuleProxy::Result::invalidate(
+    Module &M, const PreservedAnalyses &PA,
+    ModuleAnalysisManager::Invalidator &Inv) {
+  // If literally everything is preserved, we're done.
+  if (PA.areAllPreserved())
+    return false; // This is still a valid proxy.
+
+  // If this proxy isn't marked as preserved, then even if the result remains
+  // valid, the key itself may no longer be valid, so we clear everything.
+  //
+  // Note that in order to preserve this proxy, a module pass must ensure that
+  // the MFAM has been completely updated to handle the deletion of functions.
+  // Specifically, any MFAM-cached results for those functions need to have been
+  // forcibly cleared. When preserved, this proxy will only invalidate results
+  // cached on functions *still in the module* at the end of the module pass.
+  auto PAC = PA.getChecker<MachineFunctionAnalysisManagerModuleProxy>();
+  if (!PAC.preserved() && !PAC.preservedSet<AllAnalysesOn<Module>>()) {
+    InnerAM->clear();
+    return true;
   }
 
-  for (auto &F : InitializationFuncs) {
-    if (auto Err = F(M, MFAM))
-      return Err;
+  // FIXME: be more precise, see
+  // FunctionAnalysisManagerModuleProxy::Result::invalidate.
+  if (!PA.allAnalysesInSetPreserved<AllAnalysesOn<MachineFunction>>()) {
+    InnerAM->clear();
+    return true;
   }
 
-  unsigned Idx = 0;
-  size_t Size = Passes.size();
-  do {
-    // Run machine module passes
-    for (; MachineModulePasses.count(Idx) && Idx != Size; ++Idx) {
-      if (!PI.runBeforePass<Module>(*Passes[Idx], M))
-        continue;
-      if (auto Err = MachineModulePasses.at(Idx)(M, MFAM))
-        return Err;
-      PI.runAfterPass(*Passes[Idx], M, PreservedAnalyses::all());
-    }
-
-    // Finish running all passes.
-    if (Idx == Size)
-      break;
-
-    // Run machine function passes
-
-    // Get index range of machine function passes.
-    unsigned Begin = Idx;
-    for (; !MachineModulePasses.count(Idx) && Idx != Size; ++Idx)
-      ;
-
-    for (Function &F : M) {
-      // Do not codegen any 'available_externally' functions at all, they have
-      // definitions outside the translation unit.
-      if (F.hasAvailableExternallyLinkage())
-        continue;
-
-      MachineFunction &MF = MMI.getOrCreateMachineFunction(F);
-
-      for (unsigned I = Begin, E = Idx; I != E; ++I) {
-        auto *P = Passes[I].get();
+  // Return false to indicate that this result is still a valid proxy.
+  return false;
+}
 
-        if (!PI.runBeforePass<MachineFunction>(*P, MF))
-          continue;
+PreservedAnalyses
+ModuleToMachineFunctionPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
+  auto &MMI = AM.getResult<MachineModuleAnalysis>(M).getMMI();
+  MachineFunctionAnalysisManager &MFAM =
+      AM.getResult<MachineFunctionAnalysisManagerModuleProxy>(M).getManager();
+  PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(M);
+  PreservedAnalyses PA = PreservedAnalyses::all();
+  for (Function &F : M) {
+    // Do not codegen any 'available_externally' functions at all, they have
+    // definitions outside the translation unit.
+    if (F.hasAvailableExternallyLinkage())
+      continue;
+
+    MachineFunction &MF = MMI.getOrCreateMachineFunction(F);
+
+    if (!PI.runBeforePass<MachineFunction>(*Pass, MF))
+      continue;
+    PreservedAnalyses PassPA = Pass->run(MF, MFAM);
+    if (MMI.getMachineFunction(F)) {
+      MFAM.invalidate(MF, PassPA);
+      PI.runAfterPass(*Pass, MF, PassPA);
+    } else {
+      MFAM.clear(MF, F.getName());
+      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PassPA);
+    }
+    PA.intersect(std::move(PassPA));
+  }
 
-        // TODO: EmitSizeRemarks
-        PreservedAnalyses PassPA = P->run(MF, MFAM);
+  return PA;
+}
 
-        // MF is dangling after FreeMachineFunctionPass
-        if (P->name() != FreeMachineFunctionPass::name()) {
-          MFAM.invalidate(MF, PassPA);
+void ModuleToMachineFunctionPassAdaptor::printPipeline(
+    raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
+  OS << "machine-function(";
+  Pass->printPipeline(OS, MapClassName2PassName);
+  OS << ')';
+}
 
-          PI.runAfterPass(*P, MF, PassPA);
-        }
-      }
+template <>
+PreservedAnalyses
+PassManager<MachineFunction>::run(MachineFunction &MF,
+                                  AnalysisManager<MachineFunction> &MFAM) {
+  PassInstrumentation PI = MFAM.getResult<PassInstrumentationAnalysis>(MF);
+  Function &F = MF.getFunction();
+  MachineModuleInfo &MMI =
+      MFAM.getResult<ModuleAnalysisManagerMachineFunctionProxy>(MF)
+          .getCachedResult<MachineModuleAnalysis>(*F.getParent())
+          ->getMMI();
+  PreservedAnalyses PA = PreservedAnalyses::all();
+  for (auto &Pass : Passes) {
+    if (!PI.runBeforePass<MachineFunction>(*Pass, MF))
+      continue;
+
+    PreservedAnalyses PassPA = Pass->run(MF, MFAM);
+    if (MMI.getMachineFunction(F)) {
+      MFAM.invalidate(MF, PassPA);
+      PI.runAfterPass(*Pass, MF, PassPA);
+    } else {
+      MFAM.clear(MF, F.getName());
+      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PassPA);
     }
-  } while (true);
-
-  for (auto &F : FinalizationFuncs) {
-    if (auto Err = F(M, MFAM))
-      return Err;
+    PA.intersect(std::move(PassPA));
   }
-
-  return Error::success();
+  return PA;
 }
 
 } // namespace llvm
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 7c306c4a21daf7..2bcd9523a0531a 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -91,6 +91,7 @@
 #include "llvm/CodeGen/JMCInstrumenter.h"
 #include "llvm/CodeGen/LowerEmuTLS.h"
 #include "llvm/CodeGen/MIRPrinter.h"
+#include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/SafeStack.h"
 #include "llvm/CodeGen/SelectOptimize.h"
 #include "llvm/CodeGen/ShadowStackGCLowering.h"
@@ -1258,6 +1259,28 @@ static bool isFunctionPassName(StringRef Name, CallbacksT &Callbacks) {
   return callbacksAcceptPassName<FunctionPassManager>(Name, Callbacks);
 }
 
+template <typename CallbacksT>
+static bool isMachineFunctionPassName(StringRef Name, CallbacksT &Callbacks) {
+  // Explicitly handle pass manager names.
+  if (Name == "machine-function")
+    return true;
+
+  // Explicitly handle custom-parsed pass names.
+  if (parseRepeatPassName(Name))
+    return true;
+
+#define MACHINE_FUNCTION_PASS(NAME, CREATE_PASS)                               \
+  if (Name == NAME)                                                            \
+    return true;
+#define MACHINE_FUNCTION_ANALYSIS(NAME, CREATE_PASS)                           \
+  if (Name == "require<" NAME ">" || Name == "invalidate<" NAME ">")           \
+    return true;
+
+#include "llvm/Passes/MachinePassRegistry.def"
+
+  return callbacksAcceptPassName<MachineFunctionPassManager>(Name, Callbacks);
+}
+
 template <typename CallbacksT>
 static bool isLoopNestPassName(StringRef Name, CallbacksT &Callbacks,
                                bool &UseMemorySSA) {
@@ -1393,6 +1416,13 @@ Error PassBuilder::parseModulePass(ModulePassManager &MPM,
       MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM)));
       return Error::success();
     }
+    if (Name == "machine-function") {
+      MachineFunctionPassManager MFPM;
+      if (auto Err = parseMachinePassPipeline(MFPM, InnerPipeline))
+        return Err;
+      MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
+      return Error::success();
+    }
     if (auto Params = parseFunctionPipelineName(Name)) {
       if (Params->second)
         return make_error<StringError>(
@@ -1873,8 +1903,8 @@ Error PassBuilder::parseMachinePass(MachineFunctionPassManager &MFPM,
   }
 #include "llvm/Passes/MachinePassRegistry.def"
 
-  for (auto &C : MachinePipelineParsingCallbacks)
-    if (C(Name, MFPM))
+  for (auto &C : MachineFunctionPipelineParsingCallbacks)
+    if (C(Name, MFPM, E.InnerPipeline))
       return Error::success();
   return make_error<StringError>(
       formatv("unknown machine pass '{0}'", Name).str(),
@@ -1941,7 +1971,8 @@ Error PassBuilder::parseCGSCCPassPipeline(CGSCCPassManager &CGPM,
 void PassBuilder::crossRegisterProxies(LoopAnalysisManager &LAM,
                                        FunctionAnalysisManager &FAM,
                                        CGSCCAnalysisManager &CGAM,
-                                       ModuleAnalysisManager &MAM) {
+                                       ModuleAnalysisManager &MAM,
+                                       MachineFunctionAnalysisManager *MFAM) {
   MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
   MAM.registerPass([&] { return CGSCCAnalysisManagerModuleProxy(CGAM); });
   CGAM.registerPass([&] { return ModuleAnalysisManagerCGSCCProxy(MAM); });
@@ -1949,6 +1980,14 @@ void PassBuilder::crossRegisterProxies(LoopAnalysisManager &LAM,
   FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
   FAM.registerPass([&] { return LoopAnalysisManagerFunctionProxy(LAM); });
   LAM.registerPass([&] { return FunctionAnalysisManagerLoopProxy(FAM); });
+  if (MFAM) {
+    MAM.registerPass(
+        [&] { return MachineFunctionAnalysisManagerModuleProxy(*MFAM); });
+    MFAM->registerPass(
+        [&] { return ModuleAnalysisManagerMachineFunctionProxy(MAM); });
+    MFAM->registerPass(
+        [&] { return FunctionAnalysisManagerMachineFunctionProxy(FAM); });
+  }
 }
 
 Error PassBuilder::parseModulePassPipeline(ModulePassManager &MPM,
@@ -1990,6 +2029,9 @@ Error PassBuilder::parsePassPipeline(ModulePassManager &MPM,
                               UseMemorySSA)) {
       Pipeline = {{"function", {{UseMemorySSA ? "loop-mssa" : "loop",
                                  std::move(*Pipeline)}}}};
+    } else if (isMachineFunctionPassName(
+                   FirstName, MachineFunctionPipelineParsingCallbacks)) {
+      Pipeline = {{"machine-function", std::move(*Pipeline)}};
     } else {
       for (auto &C : TopLevelPipelineParsingCallbacks)
         if (C(MPM, *Pipeline))
diff --git a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
index 4a11dd2e31acde..a620ba911ec616 100644
--- a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
+++ b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
@@ -47,10 +47,9 @@ Error X86CodeGenPassBuilder::addInstSelector(AddMachinePass &) const {
 } // namespace
 
 Error X86TargetMachine::buildCodeGenPipeline(
-    ModulePassManager &MPM, MachineFunctionPassManager &MFPM,
-    MachineFunctionAnalysisManager &, raw_pwrite_stream &Out,
-    raw_pwrite_stream *DwoOut, CodeGenFileType FileType,
-    CGPassBuilderOption Opt, PassInstrumentationCallbacks *PIC) {
+    ModulePassManager &MPM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
+    CodeGenFileType FileType, CGPassBuilderOption Opt,
+    PassInstrumentationCallbacks *PIC) {
   auto CGPB = X86CodeGenPassBuilder(*this, Opt, PIC);
-  return CGPB.buildPipeline(MPM, MFPM, Out, DwoOut, FileType);
+  return CGPB.buildPipeline(MPM, Out, DwoOut, FileType);
 }
diff --git a/llvm/lib/Target/X86/X86TargetMachine.h b/llvm/lib/Target/X86/X86TargetMachine.h
index f31c971df9584d..0fd3e47aaefe7e 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.h
+++ b/llvm/lib/Target/X86/X86TargetMachine.h
@@ -58,10 +58,9 @@ class X86TargetMachine final : public LLVMTargetMachine {
   createMachineFunctionInfo(BumpPtrAllocator &Allocator, const Function &F,
                             const TargetSubtargetInfo *STI) const override;
 
-  Error buildCodeGenPipeline(ModulePassManager &, MachineFunctionPassManager &,
-                             MachineFunctionAnalysisManager &,
-                             raw_pwrite_stream &, raw_pwrite_stream *,
-                             CodeGenFileType, CGPassBuilderOption,
+  Error buildCodeGenPipeline(ModulePassManager &, raw_pwrite_stream &,
+                             raw_pwrite_stream *, CodeGenFileType,
+                             CGPassBuilderOption,
                              PassInstrumentationCallbacks *) override;
 
   bool isJIT() const { return IsJIT; }
diff --git a/llvm/test/tools/llc/new-pm/pipeline.mir b/llvm/test/tools/llc/new-pm/pipeline.mir
index c7dda4b6d13560..fcc7d4f8f02e38 100644
--- a/llvm/test/tools/llc/new-pm/pipeline.mir
+++ b/llvm/test/tools/llc/new-pm/pipeline.mir
@@ -1,7 +1,6 @@
 # RUN: llc -mtriple=x86_64-pc-linux-gnu -x mir -passes=no-op-machine-function --print-pipeline-passes -filetype=null < %s | FileCheck %s --match-full-lines
 
-# CHECK: IR pipeline: PrintMIRPreparePass
-# CHECK: MIR pipeline: no-op-machine-function,print,FreeMachineFunctionPass
+# CHECK: machine-function(no-op-machine-function),PrintMIRPreparePass,machine-function(print,FreeMachineFunctionPass)
 
 ---
 name: f
diff --git a/llvm/test/tools/llc/new-pm/start-stop.ll b/llvm/test/tools/llc/new-pm/start-stop.ll
index c25e45d1f7ab9f..4a42bb452c9a1b 100644
--- a/llvm/test/tools/llc/new-pm/start-stop.ll
+++ b/llvm/test/tools/llc/new-pm/start-stop.ll
@@ -1,4 +1,3 @@
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -enable-new-pm -print-pipeline-passes -start-before=mergeicmps -stop-after=gc-lowering -filetype=null %s | FileCheck --match-full-lines %s
 
-; CHECK: IR pipeline: function(mergeicmps,expand-memcmp,gc-lowering)
-
+; CHECK: function(mergeicmps,expand-memcmp,gc-lowering),machine-function(print)
diff --git a/llvm/tools/llc/NewPMDriver.cpp b/llvm/tools/llc/NewPMDriver.cpp
index c3288ef9d0808b..196e50071b6e54 100644
--- a/llvm/tools/llc/NewPMDriver.cpp
+++ b/llvm/tools/llc/NewPMDriver.cpp
@@ -89,30 +89,6 @@ bool LLCDiagnosticHandler::handleDiagnostics(const DiagnosticInfo &DI) {
 
 static llvm::ExitOnError ExitOnErr;
 
-static void RunPasses(bool BOS, ToolOutputFile *Out, Module *M,
-                      LLVMContext &Context, SmallString<0> &Buffer,
-                      ModulePassManager *MPM, ModuleAnalysisManager *MAM,
-                      MachineFunctionPassManager &MFPM,
-                      MachineFunctionAnalysisManager &MFAM) {
-  assert(M && "invalid input module!");
-
-  // Before executing passes, print the final values of the LLVM options.
-  cl::PrintOptionValues();
-
-  if (MPM) {
-    assert(MAM && "expect a ModuleAnalysisManager!");
-    MPM->run(*M, *MAM);
-  }
-
-  ExitOnErr(MFPM.run(*M, MFAM));
-
-  if (Context.getDiagHandlerPtr()->HasErrors)
-    exit(1);
-
-  if (BOS)
-    Out->os() << Buffer;
-}
-
 int llvm::compileModuleWithNewPM(
     StringRef Arg0, std::unique_ptr<Module> M, std::unique_ptr<MIRParser> MIR,
     std::unique_ptr<TargetMachine> Target, std::unique_ptr<ToolOutputFile> Out,
@@ -131,16 +107,6 @@ int llvm::compileModuleWithNewPM(
 
   raw_pwrite_stream *OS = &Out->os();
 
-  // Manually do the buffering rather than using buffer_ostream,
-  // so we can memcmp the contents in CompileTwice mode in future.
-  SmallString<0> Buffer;
-  std::unique_ptr<raw_svector_ostream> BOS;
-  if ((codegen::getFileType() != CodeGenFileType::AssemblyFile &&
-       !Out->os().supportsSeeking())) {
-    BOS = std::make_unique<raw_svector_ostream>(Buffer);
-    OS = BOS.get();
-  }
-
   // Fetch options from TargetPassConfig
   CGPassBuilderOption Opt = getCGPassBuilderOption();
   Opt.DisableVerify = NoVerify;
@@ -158,20 +124,19 @@ int llvm::compileModuleWithNewPM(
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
   ModuleAnalysisManager MAM;
+  MachineFunctionAnalysisManager MFAM;
   PassBuilder PB(Target.get(), PipelineTuningOptions(), std::nullopt, &PIC);
   PB.registerModuleAnalyses(MAM);
   PB.registerCGSCCAnalyses(CGAM);
   PB.registerFunctionAnalyses(FAM);
   PB.registerLoopAnalyses(LAM);
-  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+  PB.registerMachineFunctionAnalyses(MFAM);
+  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM, &MFAM);
 
   FAM.registerPass([&] { return TargetLibraryAnalysis(TLII); });
   MAM.registerPass([&] { return MachineModuleAnalysis(MMI); });
 
-  MachineFunctionAnalysisManager MFAM(FAM, MAM);
-
   ModulePassManager MPM;
-  MachineFunctionPassManager MFPM;
 
   if (!PassPipeline.empty()) {
     // Construct a custom pass pipeline that starts after instruction
@@ -182,49 +147,53 @@ int llvm::compileModuleWithNewPM(
       return 1;
     }
 
-    ExitOnErr(PB.parsePassPipeline(MFPM, PassPipeline));
+    ExitOnErr(PB.parsePassPipeline(MPM, PassPipeline));
     MPM.addPass(PrintMIRPreparePass(*OS));
+    MachineFunctionPassManager MFPM;
     MFPM.addPass(PrintMIRPass(*OS));
     MFPM.addPass(FreeMachineFunctionPass());
+    MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
 
-    auto &MMI = MFAM.getResult<MachineModuleAnalysis>(*M).getMMI();
     if (MIR->parseMachineFunctions(*M, MMI))
       return 1;
   } else {
-    ExitOnErr(LLVMTM.buildCodeGenPipeline(MPM, MFPM, MFAM, *OS,
-                                          DwoOut ? &DwoOut->os() : nullptr,
-                                          FileType, Opt, &PIC));
+    ExitOnErr(LLVMTM.buildCodeGenPipeline(
+        MPM, *OS, DwoOut ? &DwoOut->os() : nullptr, FileType, Opt, &PIC));
 
     auto StartStopInfo = TargetPassConfig::getStartStopInfo(PIC);
     assert(StartStopInfo && "Expect StartStopInfo!");
 
     if (auto StopPassName = StartStopInfo->StopPass; !StopPassName.empty()) {
-      MFPM.addPass(PrintMIRPass(*OS));
-      MFPM.addPass(FreeMachineFunctionPass());
+      MPM.addPass(createModuleToMachineFunctionPassAdaptor(PrintMIRPass(*OS)));
     }
   }
 
   if (PrintPipelinePasses) {
-    std::string IRPipeline;
-    raw_string_ostream IRSOS(IRPipeline);
-    MPM.printPipeline(IRSOS, [&PIC](StringRef ClassName) {
-      auto PassName = PIC.getPassNameForClassName(ClassName);
-      return PassName.empty() ? ClassName : PassName;
-    });
-    outs() << "IR pipeline: " << IRPipeline << '\n';
-
-    std::string MIRPipeline;
-    raw_string_ostream MIRSOS(MIRPipeline);
-    MFPM.printPipeline(MIRSOS, [&PIC](StringRef ClassName) {
+    std::string PipelineStr;
+    raw_string_ostream OS(PipelineStr);
+    MPM.printPipeline(OS, [&PIC](StringRef ClassName) {
       auto PassName = PIC.getPassNameForClassName(ClassName);
       return PassName.empty() ? ClassName : PassName;
     });
-    outs() << "MIR pipeline: " << MIRPipeline << '\n';
+    outs() << PipelineStr << '\n';
     return 0;
   }
 
-  RunPasses(BOS.get(), Out.get(), M.get(), Context, Buffer, &MPM, &MAM, MFPM,
-            MFAM);
+  // Before executing passes, print the final values of the LLVM options.
+  cl::PrintOptionValues();
+
+  MPM.run(*M, MAM);
+
+  if (Context.getDiagHandlerPtr()->HasErrors)
+    exit(1);
+
+  // Before executing passes, print the final values of the LLVM options.
+  cl::PrintOptionValues();
+
+  MPM.run(*M, MAM);
+
+  if (Context.getDiagHandlerPtr()->HasErrors)
+    exit(1);
 
   // Declare success.
   Out->keep();
diff --git a/llvm/unittests/CodeGen/PassManagerTest.cpp b/llvm/unittests/CodeGen/PassManagerTest.cpp
index 28003c2f4b3f1d..4283eb01a9c8f2 100644
--- a/llvm/unittests/CodeGen/PassManagerTest.cpp
+++ b/llvm/unittests/CodeGen/PassManagerTest.cpp
@@ -5,13 +5,18 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+// Test that the various MachineFunction pass managers, adaptors, analyses, and
+// analysis managers work.
+//===----------------------------------------------------------------------===//
 
+#include "llvm/IR/PassManager.h"
 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachinePassManager.h"
+#include "llvm/IR/Analysis.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/MC/TargetRegistry.h"
@@ -34,14 +39,9 @@ class TestFunctionAnalysis : public AnalysisInfoMixin<TestFunctionAnalysis> {
     int InstructionCount;
   };
 
-  /// Run the analysis pass over the function and return a result.
+  /// The number of instructions in the Function.
   Result run(Function &F, FunctionAnalysisManager &AM) {
-    int Count = 0;
-    for (Function::iterator BBI = F.begin(), BBE = F.end(); BBI != BBE; ++BBI)
-      for (BasicBlock::iterator II = BBI->begin(), IE = BBI->end(); II != IE;
-           ++II)
-        ++Count;
-    return Result(Count);
+    return Result(F.getInstructionCount());
   }
 
 private:
@@ -59,13 +59,12 @@ class TestMachineFunctionAnalysis
     int InstructionCount;
   };
 
-  /// Run the analysis pass over the machine function and return a result.
-  Result run(MachineFunction &MF, MachineFunctionAnalysisManager::Base &AM) {
-    auto &MFAM = static_cast<MachineFunctionAnalysisManager &>(AM);
-    // Query function analysis result.
+  Result run(MachineFunction &MF, MachineFunctionAnalysisManager &AM) {
+    FunctionAnalysisManager &FAM =
+        AM.getResult<FunctionAnalysisManagerMachineFunctionProxy>(MF)
+            .getManager();
     TestFunctionAnalysis::Result &FAR =
-        MFAM.getResult<TestFunctionAnalysis>(MF.getFunction());
-    // + 5
+        FAM.getResult<TestFunctionAnalysis>(MF.getFunction());
     return FAR.InstructionCount;
   }
 
@@ -76,90 +75,54 @@ class TestMachineFunctionAnalysis
 
 AnalysisKey TestMachineFunctionAnalysis::Key;
 
-const std::string DoInitErrMsg = "doInitialization failed";
-const std::string DoFinalErrMsg = "doFinalization failed";
-
 struct TestMachineFunctionPass : public PassInfoMixin<TestMachineFunctionPass> {
-  TestMachineFunctionPass(int &Count, std::vector<int> &BeforeInitialization,
-                          std::vector<int> &BeforeFinalization,
-                          std::vector<int> &MachineFunctionPassCount)
-      : Count(Count), BeforeInitialization(BeforeInitialization),
-        BeforeFinalization(BeforeFinalization),
-        MachineFunctionPassCount(MachineFunctionPassCount) {}
-
-  Error doInitialization(Module &M, MachineFunctionAnalysisManager &MFAM) {
-    // Force doInitialization fail by starting with big `Count`.
-    if (Count > 10000)
-      return make_error<StringError>(DoInitErrMsg, inconvertibleErrorCode());
-
-    // + 1
-    ++Count;
-    BeforeInitialization.push_back(Count);
-    return Error::success();
-  }
-  Error doFinalization(Module &M, MachineFunctionAnalysisManager &MFAM) {
-    // Force doFinalization fail by starting with big `Count`.
-    if (Count > 1000)
-      return make_error<StringError>(DoFinalErrMsg, inconvertibleErrorCode());
-
-    // + 1
-    ++Count;
-    BeforeFinalization.push_back(Count);
-    return Error::success();
-  }
+  TestMachineFunctionPass(int &Count, std::vector<int> &Counts)
+      : Count(Count), Counts(Counts) {}
 
   PreservedAnalyses run(MachineFunction &MF,
                         MachineFunctionAnalysisManager &MFAM) {
-    // Query function analysis result.
+    FunctionAnalysisManager &FAM =
+        MFAM.getResult<FunctionAnalysisManagerMachineFunctionProxy>(MF)
+            .getManager();
     TestFunctionAnalysis::Result &FAR =
-        MFAM.getResult<TestFunctionAnalysis>(MF.getFunction());
-    // 3 + 1 + 1 = 5
+        FAM.getResult<TestFunctionAnalysis>(MF.getFunction());
     Count += FAR.InstructionCount;
 
-    // Query module analysis result.
-    MachineModuleInfo &MMI =
-        MFAM.getResult<MachineModuleAnalysis>(*MF.getFunction().getParent())
-            .getMMI();
-    // 1 + 1 + 1 = 3
-    Count += (MMI.getModule() == MF.getFunction().getParent());
-
-    // Query machine function analysis result.
     TestMachineFunctionAnalysis::Result &MFAR =
         MFAM.getResult<TestMachineFunctionAnalysis>(MF);
-    // 3 + 1 + 1 = 5
     Count += MFAR.InstructionCount;
 
-    MachineFunctionPassCount.push_back(Count);
+    Counts.push_back(Count);
 
     return PreservedAnalyses::none();
   }
 
   int &Count;
-  std::vector<int> &BeforeInitialization;
-  std::vector<int> &BeforeFinalization;
-  std::vector<int> &MachineFunctionPassCount;
+  std::vector<int> &Counts;
 };
 
 struct TestMachineModulePass : public PassInfoMixin<TestMachineModulePass> {
-  TestMachineModulePass(int &Count, std::vector<int> &MachineModulePassCount)
-      : Count(Count), MachineModulePassCount(MachineModulePassCount) {}
-
-  Error run(Module &M, MachineFunctionAnalysisManager &MFAM) {
-    MachineModuleInfo &MMI = MFAM.getResult<MachineModuleAnalysis>(M).getMMI();
-    // + 1
-    Count += (MMI.getModule() == &M);
-    MachineModulePassCount.push_back(Count);
-    return Error::success();
-  }
-
-  PreservedAnalyses run(MachineFunction &MF,
-                        MachineFunctionAnalysisManager &AM) {
-    llvm_unreachable(
-        "This should never be reached because this is machine module pass");
+  TestMachineModulePass(int &Count, std::vector<int> &Counts)
+      : Count(Count), Counts(Counts) {}
+
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    MachineModuleInfo &MMI = MAM.getResult<MachineModuleAnalysis>(M).getMMI();
+    FunctionAnalysisManager &FAM =
+        MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+    MachineFunctionAnalysisManager &MFAM =
+        MAM.getResult<MachineFunctionAnalysisManagerModuleProxy>(M)
+            .getManager();
+    for (Function &F : M) {
+      MachineFunction &MF = MMI.getOrCreateMachineFunction(F);
+      Count += FAM.getResult<TestFunctionAnalysis>(F).InstructionCount;
+      Count += MFAM.getResult<TestMachineFunctionAnalysis>(MF).InstructionCount;
+    }
+    Counts.push_back(Count);
+    return PreservedAnalyses::all();
   }
 
   int &Count;
-  std::vector<int> &MachineModulePassCount;
+  std::vector<int> &Counts;
 };
 
 std::unique_ptr<Module> parseIR(LLVMContext &Context, const char *IR) {
@@ -211,102 +174,40 @@ TEST_F(PassManagerTest, Basic) {
   M->setDataLayout(TM->createDataLayout());
 
   MachineModuleInfo MMI(LLVMTM);
+
   LoopAnalysisManager LAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
   ModuleAnalysisManager MAM;
+  MachineFunctionAnalysisManager MFAM;
   PassBuilder PB(TM.get());
   PB.registerModuleAnalyses(MAM);
+  PB.registerCGSCCAnalyses(CGAM);
   PB.registerFunctionAnalyses(FAM);
-  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+  PB.registerLoopAnalyses(LAM);
+  PB.registerMachineFunctionAnalyses(MFAM);
+  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM, &MFAM);
 
   FAM.registerPass([&] { return TestFunctionAnalysis(); });
-  FAM.registerPass([&] { return PassInstrumentationAnalysis(); });
   MAM.registerPass([&] { return MachineModuleAnalysis(MMI); });
-  MAM.registerPass([&] { return PassInstrumentationAnalysis(); });
-
-  MachineFunctionAnalysisManager MFAM;
-  {
-    // Test move assignment.
-    MachineFunctionAnalysisManager NestedMFAM(FAM, MAM);
-    NestedMFAM.registerPass([&] { return PassInstrumentationAnalysis(); });
-    NestedMFAM.registerPass([&] { return TestMachineFunctionAnalysis(); });
-    MFAM = std::move(NestedMFAM);
-  }
+  MFAM.registerPass([&] { return TestMachineFunctionAnalysis(); });
 
   int Count = 0;
-  std::vector<int> BeforeInitialization[2];
-  std::vector<int> BeforeFinalization[2];
-  std::vector<int> TestMachineFunctionCount[2];
-  std::vector<int> TestMachineModuleCount[2];
+  std::vector<int> Counts;
 
+  ModulePassManager MPM;
   MachineFunctionPassManager MFPM;
-  {
-    // Test move assignment.
-    MachineFunctionPassManager NestedMFPM;
-    NestedMFPM.addPass(TestMachineModulePass(Count, TestMachineModuleCount[0]));
-    NestedMFPM.addPass(TestMachineFunctionPass(Count, BeforeInitialization[0],
-                                               BeforeFinalization[0],
-                                               TestMachineFunctionCount[0]));
-    NestedMFPM.addPass(TestMachineModulePass(Count, TestMachineModuleCount[1]));
-    NestedMFPM.addPass(TestMachineFunctionPass(Count, BeforeInitialization[1],
-                                               BeforeFinalization[1],
-                                               TestMachineFunctionCount[1]));
-    MFPM = std::move(NestedMFPM);
-  }
+  MPM.addPass(TestMachineModulePass(Count, Counts));
+  MPM.addPass(createModuleToMachineFunctionPassAdaptor(
+      TestMachineFunctionPass(Count, Counts)));
+  MPM.addPass(TestMachineModulePass(Count, Counts));
+  MFPM.addPass(TestMachineFunctionPass(Count, Counts));
+  MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
+
+  MPM.run(*M, MAM);
 
-  ASSERT_FALSE(errorToBool(MFPM.run(*M, MFAM)));
-
-  // Check first machine module pass
-  EXPECT_EQ(1u, TestMachineModuleCount[0].size());
-  EXPECT_EQ(3, TestMachineModuleCount[0][0]);
-
-  // Check first machine function pass
-  EXPECT_EQ(1u, BeforeInitialization[0].size());
-  EXPECT_EQ(1, BeforeInitialization[0][0]);
-  EXPECT_EQ(3u, TestMachineFunctionCount[0].size());
-  EXPECT_EQ(10, TestMachineFunctionCount[0][0]);
-  EXPECT_EQ(13, TestMachineFunctionCount[0][1]);
-  EXPECT_EQ(16, TestMachineFunctionCount[0][2]);
-  EXPECT_EQ(1u, BeforeFinalization[0].size());
-  EXPECT_EQ(31, BeforeFinalization[0][0]);
-
-  // Check second machine module pass
-  EXPECT_EQ(1u, TestMachineModuleCount[1].size());
-  EXPECT_EQ(17, TestMachineModuleCount[1][0]);
-
-  // Check second machine function pass
-  EXPECT_EQ(1u, BeforeInitialization[1].size());
-  EXPECT_EQ(2, BeforeInitialization[1][0]);
-  EXPECT_EQ(3u, TestMachineFunctionCount[1].size());
-  EXPECT_EQ(24, TestMachineFunctionCount[1][0]);
-  EXPECT_EQ(27, TestMachineFunctionCount[1][1]);
-  EXPECT_EQ(30, TestMachineFunctionCount[1][2]);
-  EXPECT_EQ(1u, BeforeFinalization[1].size());
-  EXPECT_EQ(32, BeforeFinalization[1][0]);
-
-  EXPECT_EQ(32, Count);
-
-  // doInitialization returns error
-  Count = 10000;
-  MFPM.addPass(TestMachineFunctionPass(Count, BeforeInitialization[1],
-                                       BeforeFinalization[1],
-                                       TestMachineFunctionCount[1]));
-  std::string Message;
-  llvm::handleAllErrors(MFPM.run(*M, MFAM), [&](llvm::StringError &Error) {
-    Message = Error.getMessage();
-  });
-  EXPECT_EQ(Message, DoInitErrMsg);
-
-  // doFinalization returns error
-  Count = 1000;
-  MFPM.addPass(TestMachineFunctionPass(Count, BeforeInitialization[1],
-                                       BeforeFinalization[1],
-                                       TestMachineFunctionCount[1]));
-  llvm::handleAllErrors(MFPM.run(*M, MFAM), [&](llvm::StringError &Error) {
-    Message = Error.getMessage();
-  });
-  EXPECT_EQ(Message, DoFinalErrMsg);
+  EXPECT_EQ((std::vector<int>{10, 16, 18, 20, 30, 36, 38, 40}), Counts);
+  EXPECT_EQ(40, Count);
 }
 
 } // namespace
diff --git a/llvm/unittests/MIR/PassBuilderCallbacksTest.cpp b/llvm/unittests/MIR/PassBuilderCallbacksTest.cpp
index 8ecde223cb5105..4b7d7846b0336c 100644
--- a/llvm/unittests/MIR/PassBuilderCallbacksTest.cpp
+++ b/llvm/unittests/MIR/PassBuilderCallbacksTest.cpp
@@ -6,6 +6,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/Analysis/LoopAnalysisManager.h"
+#include "llvm/CodeGen/FreeMachineFunction.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Testing/Support/Error.h"
@@ -96,8 +99,6 @@ MATCHER_P(HasNameRegex, Name, "") {
 }
 
 struct MockPassInstrumentationCallbacks {
-  PassInstrumentationCallbacks Callbacks;
-
   MockPassInstrumentationCallbacks() {
     ON_CALL(*this, runBeforePass(_, _)).WillByDefault(Return(true));
   }
@@ -111,7 +112,7 @@ struct MockPassInstrumentationCallbacks {
   MOCK_METHOD2(runBeforeAnalysis, void(StringRef PassID, llvm::Any));
   MOCK_METHOD2(runAfterAnalysis, void(StringRef PassID, llvm::Any));
 
-  void registerPassInstrumentation() {
+  void registerPassInstrumentation(PassInstrumentationCallbacks &Callbacks) {
     Callbacks.registerShouldRunOptionalPassCallback(
         [this](StringRef P, llvm::Any IR) {
           return this->runBeforePass(P, IR);
@@ -147,7 +148,8 @@ struct MockPassInstrumentationCallbacks {
     // to check these explicitly.
     EXPECT_CALL(*this,
                 runBeforePass(Not(HasNameRegex("Mock")), HasName(IRName)))
-        .Times(AnyNumber());
+        .Times(AnyNumber())
+        .WillRepeatedly(Return(false));
     EXPECT_CALL(
         *this, runBeforeSkippedPass(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
@@ -157,15 +159,9 @@ struct MockPassInstrumentationCallbacks {
     EXPECT_CALL(*this,
                 runAfterPass(Not(HasNameRegex("Mock")), HasName(IRName), _))
         .Times(AnyNumber());
-    EXPECT_CALL(*this, runBeforeAnalysis(HasNameRegex("MachineModuleAnalysis"),
-                                         HasName(IRName)))
-        .Times(AnyNumber());
     EXPECT_CALL(*this,
                 runBeforeAnalysis(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
-    EXPECT_CALL(*this, runAfterAnalysis(HasNameRegex("MachineModuleAnalysis"),
-                                        HasName(IRName)))
-        .Times(AnyNumber());
     EXPECT_CALL(*this,
                 runAfterAnalysis(Not(HasNameRegex("Mock")), HasName(IRName)))
         .Times(AnyNumber());
@@ -202,7 +198,7 @@ template <typename DerivedT> class MockAnalysisHandleBase {
       }
     };
 
-    Result run(MachineFunction &IR, MachineFunctionAnalysisManager::Base &AM) {
+    Result run(MachineFunction &IR, MachineFunctionAnalysisManager &AM) {
       return Handle->run(IR, AM);
     }
   };
@@ -249,7 +245,7 @@ template <typename DerivedT> class MockPassHandleBase {
 
   public:
     PreservedAnalyses run(MachineFunction &IR,
-                          MachineFunctionAnalysisManager::Base &AM) {
+                          MachineFunctionAnalysisManager &AM) {
       return Handle->run(IR, AM);
     }
   };
@@ -270,7 +266,7 @@ template <typename DerivedT> class MockPassHandleBase {
 
 struct MockAnalysisHandle : public MockAnalysisHandleBase<MockAnalysisHandle> {
   MOCK_METHOD2(run, Analysis::Result(MachineFunction &,
-                                     MachineFunctionAnalysisManager::Base &));
+                                     MachineFunctionAnalysisManager &));
 
   MOCK_METHOD3(invalidate, bool(MachineFunction &, const PreservedAnalyses &,
                                 MachineFunctionAnalysisManager::Invalidator &));
@@ -284,7 +280,7 @@ AnalysisKey MockAnalysisHandleBase<DerivedT>::Analysis::Key;
 class MockPassHandle : public MockPassHandleBase<MockPassHandle> {
 public:
   MOCK_METHOD2(run, PreservedAnalyses(MachineFunction &,
-                                      MachineFunctionAnalysisManager::Base &));
+                                      MachineFunctionAnalysisManager &));
 
   MockPassHandle() { setDefaults(); }
 };
@@ -297,50 +293,51 @@ class MachineFunctionCallbacksTest : public testing::Test {
     InitializeAllTargetMCs();
   }
 
+  LLVMContext Context;
+
   std::unique_ptr<LLVMTargetMachine> TM;
   std::unique_ptr<MachineModuleInfo> MMI;
 
-  LLVMContext Context;
   std::unique_ptr<Module> M;
-  std::unique_ptr<MIRParser> MIR;
-
-  MockPassInstrumentationCallbacks CallbacksHandle;
 
-  PassBuilder PB;
-  ModulePassManager PM;
-  MachineFunctionPassManager MFPM;
-  FunctionAnalysisManager FAM;
-  ModuleAnalysisManager AM;
+  PassInstrumentationCallbacks PIC;
+  std::unique_ptr<PassBuilder> PB;
+  ModulePassManager MPM;
   MachineFunctionAnalysisManager MFAM;
+  LoopAnalysisManager LAM;
+  FunctionAnalysisManager FAM;
+  CGSCCAnalysisManager CGAM;
+  ModuleAnalysisManager MAM;
 
+  MockPassInstrumentationCallbacks CallbacksHandle;
   MockPassHandle PassHandle;
   MockAnalysisHandle AnalysisHandle;
 
-  std::unique_ptr<Module> parseMIR(const TargetMachine &TM, StringRef MIRCode,
-                                   MachineModuleInfo &MMI) {
+  static std::unique_ptr<Module> parseMIR(StringRef MIRCode,
+                                          LLVMContext &Context,
+                                          TargetMachine &TM,
+                                          MachineModuleInfo &MMI) {
     SMDiagnostic Diagnostic;
     std::unique_ptr<MemoryBuffer> MBuffer = MemoryBuffer::getMemBuffer(MIRCode);
-    MIR = createMIRParser(std::move(MBuffer), Context);
-    if (!MIR)
-      return nullptr;
+    std::unique_ptr<MIRParser> MIR =
+        createMIRParser(std::move(MBuffer), Context);
+    assert(MIR);
 
     std::unique_ptr<Module> Mod = MIR->parseIRModule();
-    if (!Mod)
-      return nullptr;
+    assert(Mod);
 
+    // Module identifier is used in tests below.
+    Mod->setModuleIdentifier("module");
     Mod->setDataLayout(TM.createDataLayout());
 
-    if (MIR->parseMachineFunctions(*Mod, MMI)) {
-      M.reset();
-      return nullptr;
-    }
+    bool Ret = MIR->parseMachineFunctions(*Mod, MMI);
+    assert(!Ret);
+
     return Mod;
   }
 
   static PreservedAnalyses
-  getAnalysisResult(MachineFunction &U,
-                    MachineFunctionAnalysisManager::Base &AM) {
-    auto &MFAM = static_cast<MachineFunctionAnalysisManager &>(AM);
+  getAnalysisResult(MachineFunction &U, MachineFunctionAnalysisManager &MFAM) {
     MFAM.getResult<MockAnalysisHandle::Analysis>(U);
     return PreservedAnalyses::all();
   }
@@ -356,25 +353,18 @@ class MachineFunctionCallbacksTest : public testing::Test {
             TripleName, "", "", TargetOptions(), std::nullopt)));
     if (!TM)
       GTEST_SKIP();
-    MMI = std::make_unique<MachineModuleInfo>(TM.get());
-    M = parseMIR(*TM, MIRString, *MMI);
-    AM.registerPass([&] { return MachineModuleAnalysis(*MMI); });
-  }
 
-  MachineFunctionCallbacksTest()
-      : CallbacksHandle(), PB(nullptr, PipelineTuningOptions(), std::nullopt,
-                              &CallbacksHandle.Callbacks),
-        PM(), FAM(), AM(), MFAM(FAM, AM) {
-
-    EXPECT_TRUE(&CallbacksHandle.Callbacks ==
-                PB.getPassInstrumentationCallbacks());
+    MMI = std::make_unique<MachineModuleInfo>(TM.get());
+    M = parseMIR(MIRString, Context, *TM, *MMI);
+    PB = std::make_unique<PassBuilder>(TM.get(), PipelineTuningOptions(),
+                                       std::nullopt, &PIC);
 
     /// 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(
+    PB->registerAnalysisRegistrationCallback(
         [this](MachineFunctionAnalysisManager &AM) {
           // Register our mock analysis
           AM.registerPass([this] { return AnalysisHandle.getAnalysis(); });
@@ -386,24 +376,29 @@ class MachineFunctionCallbacksTest : public testing::Test {
     /// 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, MachineFunctionPassManager &PM) {
+    PB->registerPipelineParsingCallback(
+        [this](StringRef Name, MachineFunctionPassManager &PM,
+               ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
           if (parseAnalysisUtilityPasses<MockAnalysisHandle::Analysis>(
                   "test-analysis", Name, PM))
             return true;
 
           /// Parse the name of our pass mock handle
           if (Name == "test-transform") {
-            MFPM.addPass(PassHandle.getPass());
+            PM.addPass(PassHandle.getPass());
             return true;
           }
           return false;
         });
 
     /// Register builtin analyses and cross-register the analysis proxies
-    PB.registerModuleAnalyses(AM);
-    PB.registerFunctionAnalyses(FAM);
-    PB.registerMachineFunctionAnalyses(MFAM);
+    PB->registerModuleAnalyses(MAM);
+    PB->registerCGSCCAnalyses(CGAM);
+    PB->registerFunctionAnalyses(FAM);
+    PB->registerLoopAnalyses(LAM);
+    PB->registerMachineFunctionAnalyses(MFAM);
+    PB->crossRegisterProxies(LAM, FAM, CGAM, MAM, &MFAM);
+    MAM.registerPass([&] { return MachineModuleAnalysis(*MMI); });
   }
 };
 
@@ -412,53 +407,58 @@ TEST_F(MachineFunctionCallbacksTest, Passes) {
   EXPECT_CALL(PassHandle, run(HasName("test"), _)).WillOnce(&getAnalysisResult);
 
   StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(MFPM, PipelineText), Succeeded())
+  ASSERT_THAT_ERROR(PB->parsePassPipeline(MPM, PipelineText), Succeeded())
       << "Pipeline was: " << PipelineText;
-  ASSERT_THAT_ERROR(MFPM.run(*M, MFAM), Succeeded());
+  MPM.run(*M, MAM);
 }
 
 TEST_F(MachineFunctionCallbacksTest, InstrumentedPasses) {
-  CallbacksHandle.registerPassInstrumentation();
+  CallbacksHandle.registerPassInstrumentation(PIC);
   // Non-mock instrumentation not specifically mentioned below can be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
   CallbacksHandle.ignoreNonMockPassInstrumentation("test");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("");
+  CallbacksHandle.ignoreNonMockPassInstrumentation("module");
 
   // 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("test")))
-      .InSequence(PISequence);
+      .InSequence(PISequence)
+      .WillOnce(Return(true));
   EXPECT_CALL(
       CallbacksHandle,
       runBeforeNonSkippedPass(HasNameRegex("MockPassHandle"), HasName("test")))
       .InSequence(PISequence);
-  EXPECT_CALL(CallbacksHandle,
-              runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), _))
+  EXPECT_CALL(
+      CallbacksHandle,
+      runBeforeAnalysis(HasNameRegex("MockAnalysisHandle"), HasName("test")))
       .InSequence(PISequence);
-  EXPECT_CALL(CallbacksHandle,
-              runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), _))
+  EXPECT_CALL(
+      CallbacksHandle,
+      runAfterAnalysis(HasNameRegex("MockAnalysisHandle"), HasName("test")))
       .InSequence(PISequence);
   EXPECT_CALL(CallbacksHandle,
               runAfterPass(HasNameRegex("MockPassHandle"), HasName("test"), _))
       .InSequence(PISequence);
+  EXPECT_CALL(
+      CallbacksHandle,
+      runBeforeSkippedPass(HasNameRegex("MockPassHandle"), HasName("test")))
+      .Times(0);
 
   EXPECT_CALL(AnalysisHandle, run(HasName("test"), _));
   EXPECT_CALL(PassHandle, run(HasName("test"), _)).WillOnce(&getAnalysisResult);
 
   StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(MFPM, PipelineText), Succeeded())
+  ASSERT_THAT_ERROR(PB->parsePassPipeline(MPM, PipelineText), Succeeded())
       << "Pipeline was: " << PipelineText;
-  ASSERT_THAT_ERROR(MFPM.run(*M, MFAM), Succeeded());
+  MPM.run(*M, MAM);
 }
 
 TEST_F(MachineFunctionCallbacksTest, InstrumentedSkippedPasses) {
-  CallbacksHandle.registerPassInstrumentation();
+  CallbacksHandle.registerPassInstrumentation(PIC);
   // Non-mock instrumentation run here can safely be ignored.
-  CallbacksHandle.ignoreNonMockPassInstrumentation("<string>");
   CallbacksHandle.ignoreNonMockPassInstrumentation("test");
-  CallbacksHandle.ignoreNonMockPassInstrumentation("");
+  CallbacksHandle.ignoreNonMockPassInstrumentation("module");
 
   // Skip the pass by returning false.
   EXPECT_CALL(CallbacksHandle,
@@ -495,9 +495,81 @@ TEST_F(MachineFunctionCallbacksTest, InstrumentedSkippedPasses) {
       .Times(0);
 
   StringRef PipelineText = "test-transform";
-  ASSERT_THAT_ERROR(PB.parsePassPipeline(MFPM, PipelineText), Succeeded())
+  ASSERT_THAT_ERROR(PB->parsePassPipeline(MPM, PipelineText), Succeeded())
       << "Pipeline was: " << PipelineText;
-  ASSERT_THAT_ERROR(MFPM.run(*M, MFAM), Succeeded());
+  MPM.run(*M, MAM);
+}
+
+// Check that the Module -> MachineFunction adaptor properly calls
+// runAfterPassInvalidated.
+TEST_F(MachineFunctionCallbacksTest, InstrumentedFreeMFPass) {
+  CallbacksHandle.registerPassInstrumentation(PIC);
+  // Non-mock instrumentation run here can safely be ignored.
+  CallbacksHandle.ignoreNonMockPassInstrumentation("test");
+  CallbacksHandle.ignoreNonMockPassInstrumentation("module");
+
+  ::testing::Sequence PISequence;
+  EXPECT_CALL(
+      CallbacksHandle,
+      runBeforePass(HasNameRegex("FreeMachineFunctionPass"), HasName("test")))
+      .InSequence(PISequence)
+      .WillOnce(Return(true));
+  EXPECT_CALL(CallbacksHandle,
+              runBeforeNonSkippedPass(HasNameRegex("FreeMachineFunctionPass"),
+                                      HasName("test")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, runAfterPassInvalidated(
+                                   HasNameRegex("FreeMachineFunctionPass"), _))
+      .InSequence(PISequence);
+
+  // runAfterPass should not be called since the MachineFunction is no longer
+  // valid after FreeMachineFunctionPass.
+  EXPECT_CALL(CallbacksHandle,
+              runAfterPass(HasNameRegex("FreeMachineFunctionPass"), _, _))
+      .Times(0);
+
+  MPM.addPass(
+      createModuleToMachineFunctionPassAdaptor(FreeMachineFunctionPass()));
+  MPM.run(*M, MAM);
+}
+
+// Check that the Module -> MachineFunction adaptor and MachineFunction pass
+// manager properly call runAfterPassInvalidated.
+TEST_F(MachineFunctionCallbacksTest, InstrumentedFreeMFPass2) {
+  CallbacksHandle.registerPassInstrumentation(PIC);
+  // Non-mock instrumentation run here can safely be ignored.
+  CallbacksHandle.ignoreNonMockPassInstrumentation("test");
+  CallbacksHandle.ignoreNonMockPassInstrumentation("module");
+
+  ::testing::Sequence PISequence;
+  EXPECT_CALL(
+      CallbacksHandle,
+      runBeforePass(HasNameRegex("FreeMachineFunctionPass"), HasName("test")))
+      .InSequence(PISequence)
+      .WillOnce(Return(true));
+  EXPECT_CALL(CallbacksHandle,
+              runBeforeNonSkippedPass(HasNameRegex("FreeMachineFunctionPass"),
+                                      HasName("test")))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle, runAfterPassInvalidated(
+                                   HasNameRegex("FreeMachineFunctionPass"), _))
+      .InSequence(PISequence);
+  EXPECT_CALL(CallbacksHandle,
+              runAfterPassInvalidated(HasNameRegex("PassManager"), _))
+      .InSequence(PISequence);
+
+  // runAfterPass should not be called since the MachineFunction is no longer
+  // valid after FreeMachineFunctionPass.
+  EXPECT_CALL(CallbacksHandle,
+              runAfterPass(HasNameRegex("FreeMachineFunctionPass"), _, _))
+      .Times(0);
+  EXPECT_CALL(CallbacksHandle, runAfterPass(HasNameRegex("PassManager"), _, _))
+      .Times(0);
+
+  MachineFunctionPassManager MFPM;
+  MFPM.addPass(FreeMachineFunctionPass());
+  MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
+  MPM.run(*M, MAM);
 }
 
 } // end anonymous namespace

>From db9ddb18da484b7520461a6f8813bd0be9ef90c9 Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks at google.com>
Date: Thu, 15 Feb 2024 06:57:22 +0000
Subject: [PATCH 2/2] Move MIRPrinter to CodeGenPassBuilder::buildPipeline()

Follow legacy PM asm/MIR printer behavior in llc.
---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h | 30 ++++++++++++-------
 llvm/test/tools/llc/new-pm/start-stop.ll      |  6 ++--
 llvm/tools/llc/NewPMDriver.cpp                |  7 -----
 3 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index df7bb1f67f18b3..592a6163b21100 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -37,6 +37,7 @@
 #include "llvm/CodeGen/InterleavedLoadCombine.h"
 #include "llvm/CodeGen/JMCInstrumenter.h"
 #include "llvm/CodeGen/LowerEmuTLS.h"
+#include "llvm/CodeGen/MIRPrinter.h"
 #include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
 #include "llvm/CodeGen/ReplaceWithVeclib.h"
@@ -208,12 +209,13 @@ template <typename DerivedT> class CodeGenPassBuilder {
     }
 
     template <typename PassT>
-    void operator()(PassT &&Pass, StringRef Name = PassT::name()) {
+    void operator()(PassT &&Pass, bool Force = false,
+                    StringRef Name = PassT::name()) {
       static_assert((is_detected<is_machine_function_pass_t, PassT>::value ||
                      is_detected<is_module_pass_t, PassT>::value) &&
                     "Only module pass and function pass are supported.");
 
-      if (!PB.runBeforeAdding(Name))
+      if (!Force && !PB.runBeforeAdding(Name))
         return;
 
       // Add Function Pass
@@ -492,11 +494,13 @@ Error CodeGenPassBuilder<Derived>::buildPipeline(
   if (!StartStopInfo)
     return StartStopInfo.takeError();
   setStartStopPasses(*StartStopInfo);
-  AddIRPass addIRPass(MPM, derived());
-  // `ProfileSummaryInfo` is always valid.
-  addIRPass(RequireAnalysisPass<ProfileSummaryAnalysis, Module>());
-  addIRPass(RequireAnalysisPass<CollectorMetadataAnalysis, Module>());
-  addISelPasses(addIRPass);
+
+  {
+    AddIRPass addIRPass(MPM, derived());
+    addIRPass(RequireAnalysisPass<ProfileSummaryAnalysis, Module>());
+    addIRPass(RequireAnalysisPass<CollectorMetadataAnalysis, Module>());
+    addISelPasses(addIRPass);
+  }
 
   AddMachinePass addPass(MPM, derived());
   if (auto Err = addCoreISelPasses(addPass))
@@ -505,10 +509,14 @@ Error CodeGenPassBuilder<Derived>::buildPipeline(
   if (auto Err = derived().addMachinePasses(addPass))
     return std::move(Err);
 
-  derived().addAsmPrinter(
-      addPass, [this, &Out, DwoOut, FileType](MCContext &Ctx) {
-        return this->TM.createMCStreamer(Out, DwoOut, FileType, Ctx);
-      });
+  if (TargetPassConfig::willCompleteCodeGenPipeline()) {
+    derived().addAsmPrinter(
+        addPass, [this, &Out, DwoOut, FileType](MCContext &Ctx) {
+          return this->TM.createMCStreamer(Out, DwoOut, FileType, Ctx);
+        });
+  } else if (FileType != CodeGenFileType::Null) {
+    addPass(PrintMIRPass(Out), /*Force=*/true);
+  }
 
   addPass(FreeMachineFunctionPass());
   return verifyStartStop(*StartStopInfo);
diff --git a/llvm/test/tools/llc/new-pm/start-stop.ll b/llvm/test/tools/llc/new-pm/start-stop.ll
index 4a42bb452c9a1b..83ad40103835e4 100644
--- a/llvm/test/tools/llc/new-pm/start-stop.ll
+++ b/llvm/test/tools/llc/new-pm/start-stop.ll
@@ -1,3 +1,5 @@
-; RUN: llc -mtriple=x86_64-pc-linux-gnu -enable-new-pm -print-pipeline-passes -start-before=mergeicmps -stop-after=gc-lowering -filetype=null %s | FileCheck --match-full-lines %s
+; RUN: llc -mtriple=x86_64-pc-linux-gnu -enable-new-pm -print-pipeline-passes -start-before=mergeicmps -stop-after=gc-lowering -filetype=null %s | FileCheck --match-full-lines %s --check-prefix=NULL
+; RUN: llc -mtriple=x86_64-pc-linux-gnu -enable-new-pm -print-pipeline-passes -start-before=mergeicmps -stop-after=gc-lowering -o /dev/null %s | FileCheck --match-full-lines %s --check-prefix=OBJ
 
-; CHECK: function(mergeicmps,expand-memcmp,gc-lowering),machine-function(print)
+; NULL: function(mergeicmps,expand-memcmp,gc-lowering)
+; OBJ: function(mergeicmps,expand-memcmp,gc-lowering),machine-function(print)
diff --git a/llvm/tools/llc/NewPMDriver.cpp b/llvm/tools/llc/NewPMDriver.cpp
index 196e50071b6e54..baeb566cf9819f 100644
--- a/llvm/tools/llc/NewPMDriver.cpp
+++ b/llvm/tools/llc/NewPMDriver.cpp
@@ -159,13 +159,6 @@ int llvm::compileModuleWithNewPM(
   } else {
     ExitOnErr(LLVMTM.buildCodeGenPipeline(
         MPM, *OS, DwoOut ? &DwoOut->os() : nullptr, FileType, Opt, &PIC));
-
-    auto StartStopInfo = TargetPassConfig::getStartStopInfo(PIC);
-    assert(StartStopInfo && "Expect StartStopInfo!");
-
-    if (auto StopPassName = StartStopInfo->StopPass; !StopPassName.empty()) {
-      MPM.addPass(createModuleToMachineFunctionPassAdaptor(PrintMIRPass(*OS)));
-    }
   }
 
   if (PrintPipelinePasses) {



More information about the llvm-commits mailing list