[llvm] [NewPM] Add `FunctionToMachineFunctionPassAdaptor` (PR #88711)

via llvm-commits llvm-commits at lists.llvm.org
Mon Apr 15 23:56:34 PDT 2024


https://github.com/paperchalice updated https://github.com/llvm/llvm-project/pull/88711

>From b6058e09c5555c6c60f4b2568995992d28805f70 Mon Sep 17 00:00:00 2001
From: PaperChalice <liujunchang97 at outlook.com>
Date: Fri, 12 Apr 2024 20:26:51 +0800
Subject: [PATCH] [NewPM][CodeGen] Add `MachineFunctionAnalysis`

---
 .../llvm/CodeGen/FreeMachineFunction.h        |  1 +
 .../FunctionToMachineFunctionAnalysis.h       | 50 +++++++++++
 .../llvm/CodeGen/MIRParser/MIRParser.h        | 12 +++
 llvm/include/llvm/CodeGen/MachineModuleInfo.h |  4 +
 .../include/llvm/CodeGen/MachinePassManager.h | 49 +++++++---
 llvm/lib/CodeGen/CMakeLists.txt               |  1 +
 llvm/lib/CodeGen/FreeMachineFunction.cpp      |  8 +-
 .../FunctionToMachineFunctionAnalysis.cpp     | 49 ++++++++++
 llvm/lib/CodeGen/MIRParser/MIRParser.cpp      | 49 +++++++---
 llvm/lib/CodeGen/MachinePassManager.cpp       | 90 +++++++++++++++----
 llvm/lib/Passes/PassBuilder.cpp               |  1 +
 llvm/lib/Passes/PassRegistry.def              |  3 +
 llvm/lib/Passes/StandardInstrumentations.cpp  |  8 --
 llvm/tools/llc/NewPMDriver.cpp                |  4 +-
 llvm/unittests/CodeGen/PassManagerTest.cpp    | 11 ++-
 .../MIR/PassBuilderCallbacksTest.cpp          | 11 +--
 16 files changed, 287 insertions(+), 64 deletions(-)
 create mode 100644 llvm/include/llvm/CodeGen/FunctionToMachineFunctionAnalysis.h
 create mode 100644 llvm/lib/CodeGen/FunctionToMachineFunctionAnalysis.cpp

diff --git a/llvm/include/llvm/CodeGen/FreeMachineFunction.h b/llvm/include/llvm/CodeGen/FreeMachineFunction.h
index 5f21c6720350bc..9471b7a00e4472 100644
--- a/llvm/include/llvm/CodeGen/FreeMachineFunction.h
+++ b/llvm/include/llvm/CodeGen/FreeMachineFunction.h
@@ -13,6 +13,7 @@
 
 namespace llvm {
 
+// TODO: Convert it to function pass.
 class FreeMachineFunctionPass : public PassInfoMixin<FreeMachineFunctionPass> {
 public:
   PreservedAnalyses run(MachineFunction &MF,
diff --git a/llvm/include/llvm/CodeGen/FunctionToMachineFunctionAnalysis.h b/llvm/include/llvm/CodeGen/FunctionToMachineFunctionAnalysis.h
new file mode 100644
index 00000000000000..a4cfc305ed9ffe
--- /dev/null
+++ b/llvm/include/llvm/CodeGen/FunctionToMachineFunctionAnalysis.h
@@ -0,0 +1,50 @@
+//===- llvm/CodeGen/FunctionToMachineFunctionAnalysis.h ---------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the FunctionToMachineFunctionAnalysis class.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_FUNCTIONTOMACHINEFUNCTIONANALYSIS
+#define LLVM_CODEGEN_FUNCTIONTOMACHINEFUNCTIONANALYSIS
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class MachineFunction;
+class LLVMTargetMachine;
+
+/// This analysis create MachineFunction for given Function.
+/// To release the MachineFunction, users should invalidate it explicitly.
+class FunctionToMachineFunctionAnalysis
+    : public AnalysisInfoMixin<FunctionToMachineFunctionAnalysis> {
+  friend AnalysisInfoMixin<FunctionToMachineFunctionAnalysis>;
+
+  static AnalysisKey Key;
+
+  const LLVMTargetMachine *TM;
+
+public:
+  class Result {
+    std::unique_ptr<MachineFunction> MF;
+
+  public:
+    Result(std::unique_ptr<MachineFunction> MF) : MF(std::move(MF)) {}
+    MachineFunction &getMF() { return *MF; };
+    bool invalidate(Function &, const PreservedAnalyses &PA,
+                    FunctionAnalysisManager::Invalidator &);
+  };
+
+  FunctionToMachineFunctionAnalysis(const LLVMTargetMachine *TM) : TM(TM){};
+  Result run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_FUNCTIONTOMACHINEFUNCTIONANALYSIS
diff --git a/llvm/include/llvm/CodeGen/MIRParser/MIRParser.h b/llvm/include/llvm/CodeGen/MIRParser/MIRParser.h
index e1606e7c0ea72d..ae0938a48a7110 100644
--- a/llvm/include/llvm/CodeGen/MIRParser/MIRParser.h
+++ b/llvm/include/llvm/CodeGen/MIRParser/MIRParser.h
@@ -34,6 +34,9 @@ class MachineModuleInfo;
 class SMDiagnostic;
 class StringRef;
 
+template <typename IRUnitT, typename... ExtraArgTs> class AnalysisManager;
+using ModuleAnalysisManager = AnalysisManager<Module>;
+
 typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
@@ -60,6 +63,15 @@ class MIRParser {
   ///
   /// \returns true if an error occurred.
   bool parseMachineFunctions(Module &M, MachineModuleInfo &MMI);
+
+  /// Parses MachineFunctions in the MIR file and add them as the result
+  /// of MachineFunctionAnalysis in ModulePassManager \p MAM.
+  /// User should register at least MachineFunctionAnalysis,
+  /// MachineModuleAnalysis, FunctionAnalysisManagerModuleProxy and
+  /// PassInstrumentationAnalysis in \p MAM before parsing MIR.
+  ///
+  /// \returns true if an error occurred.
+  bool parseMachineFunctions(Module &M, ModuleAnalysisManager &MAM);
 };
 
 /// This function is the main interface to the MIR serialization format parser.
diff --git a/llvm/include/llvm/CodeGen/MachineModuleInfo.h b/llvm/include/llvm/CodeGen/MachineModuleInfo.h
index 5f66d1ada0d082..92ea3c902ce95e 100644
--- a/llvm/include/llvm/CodeGen/MachineModuleInfo.h
+++ b/llvm/include/llvm/CodeGen/MachineModuleInfo.h
@@ -147,10 +147,14 @@ class MachineModuleInfo {
 
   /// Returns the MachineFunction constructed for the IR function \p F.
   /// Creates a new MachineFunction if none exists yet.
+  /// NOTE: New pass manager clients shall not use this method to get
+  /// the `MachineFunction`, use `MachineFunctionAnalysis` instead.
   MachineFunction &getOrCreateMachineFunction(Function &F);
 
   /// \brief Returns the MachineFunction associated to IR function \p F if there
   /// is one, otherwise nullptr.
+  /// NOTE: New pass manager clients shall not use this method to get
+  /// the `MachineFunction`, use `MachineFunctionAnalysis` instead.
   MachineFunction *getMachineFunction(const Function &F) const;
 
   /// Delete the MachineFunction \p MF and reset the link in the IR Function to
diff --git a/llvm/include/llvm/CodeGen/MachinePassManager.h b/llvm/include/llvm/CodeGen/MachinePassManager.h
index 4f0b6ba2b1e738..ab8b4a5cfbc870 100644
--- a/llvm/include/llvm/CodeGen/MachinePassManager.h
+++ b/llvm/include/llvm/CodeGen/MachinePassManager.h
@@ -76,6 +76,10 @@ struct MachinePassModel
 #endif
 
     auto PA = this->Pass.run(IR, AM);
+    // Machine function passes are not allowed to modify the LLVM
+    // representation, therefore we should preserve all IR analyses.
+    PA.template preserveSet<AllAnalysesOn<Module>>();
+    PA.template preserveSet<AllAnalysesOn<Function>>();
 
     if constexpr (is_detected<has_get_set_properties_t, PassT>::value)
       IR.getProperties().set(PassT::getSetProperties());
@@ -129,16 +133,6 @@ class FunctionAnalysisManagerMachineFunctionProxy
       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
@@ -211,8 +205,7 @@ class ModuleToMachineFunctionPassAdaptor
 template <typename MachineFunctionPassT>
 ModuleToMachineFunctionPassAdaptor
 createModuleToMachineFunctionPassAdaptor(MachineFunctionPassT &&Pass) {
-  using PassModelT = detail::PassModel<MachineFunction, MachineFunctionPassT,
-                                       MachineFunctionAnalysisManager>;
+  using PassModelT = detail::MachinePassModel<MachineFunctionPassT>;
   // Do not use make_unique, it causes too many template instantiations,
   // causing terrible compile times.
   return ModuleToMachineFunctionPassAdaptor(
@@ -220,6 +213,38 @@ createModuleToMachineFunctionPassAdaptor(MachineFunctionPassT &&Pass) {
           new PassModelT(std::forward<MachineFunctionPassT>(Pass))));
 }
 
+class FunctionToMachineFunctionPassAdaptor
+    : public PassInfoMixin<FunctionToMachineFunctionPassAdaptor> {
+public:
+  using PassConceptT =
+      detail::PassConcept<MachineFunction, MachineFunctionAnalysisManager>;
+
+  explicit FunctionToMachineFunctionPassAdaptor(
+      std::unique_ptr<PassConceptT> Pass)
+      : Pass(std::move(Pass)) {}
+
+  /// Runs the machine function pass across every function.
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+  void printPipeline(raw_ostream &OS,
+                     function_ref<StringRef(StringRef)> MapClassName2PassName);
+
+  static bool isRequired() { return true; }
+
+private:
+  std::unique_ptr<PassConceptT> Pass;
+};
+
+template <typename MachineFunctionPassT>
+FunctionToMachineFunctionPassAdaptor
+createFunctionToMachineFunctionPassAdaptor(MachineFunctionPassT &&Pass) {
+  using PassModelT = detail::MachinePassModel<MachineFunctionPassT>;
+  // Do not use make_unique, it causes too many template instantiations,
+  // causing terrible compile times.
+  return FunctionToMachineFunctionPassAdaptor(
+      std::unique_ptr<FunctionToMachineFunctionPassAdaptor::PassConceptT>(
+          new PassModelT(std::forward<MachineFunctionPassT>(Pass))));
+}
+
 template <>
 template <typename PassT>
 void PassManager<MachineFunction>::addPass(PassT &&Pass) {
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index 2c24de60edd43e..6c77ab0b2cef63 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -67,6 +67,7 @@ add_llvm_component_library(LLVMCodeGen
   FixupStatepointCallerSaved.cpp
   FreeMachineFunction.cpp
   FuncletLayout.cpp
+  FunctionToMachineFunctionAnalysis.cpp
   GCMetadata.cpp
   GCMetadataPrinter.cpp
   GCRootLowering.cpp
diff --git a/llvm/lib/CodeGen/FreeMachineFunction.cpp b/llvm/lib/CodeGen/FreeMachineFunction.cpp
index f38f3e63a3f37a..183eceb95c9c1f 100644
--- a/llvm/lib/CodeGen/FreeMachineFunction.cpp
+++ b/llvm/lib/CodeGen/FreeMachineFunction.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/FreeMachineFunction.h"
+#include "llvm/CodeGen/FunctionToMachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 
@@ -15,8 +16,7 @@ using namespace llvm;
 PreservedAnalyses
 FreeMachineFunctionPass::run(MachineFunction &MF,
                              MachineFunctionAnalysisManager &MFAM) {
-  auto &MMI = MF.getMMI();
-  MFAM.invalidate(MF, PreservedAnalyses::none());
-  MMI.deleteMachineFunctionFor(MF.getFunction()); // MF is dangling now.
-  return PreservedAnalyses::none();
+  PreservedAnalyses PA = PreservedAnalyses::none();
+  PA.abandon<FunctionToMachineFunctionAnalysis>();
+  return PA;
 }
diff --git a/llvm/lib/CodeGen/FunctionToMachineFunctionAnalysis.cpp b/llvm/lib/CodeGen/FunctionToMachineFunctionAnalysis.cpp
new file mode 100644
index 00000000000000..4133f1dc4e2625
--- /dev/null
+++ b/llvm/lib/CodeGen/FunctionToMachineFunctionAnalysis.cpp
@@ -0,0 +1,49 @@
+//===- FunctionToMachineFunctionAnalysis.cpp ------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the definitions of the FunctionToMachineFunctionAnalysis
+// members.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/FunctionToMachineFunctionAnalysis.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/Target/TargetMachine.h"
+#include <atomic>
+
+using namespace llvm;
+
+AnalysisKey FunctionToMachineFunctionAnalysis::Key;
+
+bool FunctionToMachineFunctionAnalysis::Result::invalidate(
+    Function &, const PreservedAnalyses &PA,
+    FunctionAnalysisManager::Invalidator &) {
+  // Unless it is invalidated explicitly, it should remain preserved.
+  auto PAC = PA.getChecker<FunctionToMachineFunctionAnalysis>();
+  return !PAC.preservedWhenStateless();
+}
+
+FunctionToMachineFunctionAnalysis::Result
+FunctionToMachineFunctionAnalysis::run(Function &F,
+                                       FunctionAnalysisManager &FAM) {
+  // Next unique number available for a MachineFunction in current module.
+  static std::atomic_uint NextFnNum = 0;
+
+  const TargetSubtargetInfo &STI = *TM->getSubtargetImpl(F);
+  auto &MMI = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F)
+                  .getCachedResult<MachineModuleAnalysis>(*F.getParent())
+                  ->getMMI();
+  auto MF = std::make_unique<MachineFunction>(F, *TM, STI, NextFnNum++, MMI);
+  MF->initTargetMachineFunctionInfo(STI);
+
+  // MRI callback for target specific initializations.
+  TM->registerMachineRegisterInfoCallback(*MF);
+
+  return Result(std::move(MF));
+}
diff --git a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp b/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
index 4d9a8dc5602ba6..2f1de9c3da43d2 100644
--- a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
@@ -16,6 +16,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/AsmParser/SlotMapping.h"
+#include "llvm/CodeGen/FunctionToMachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MIRParser/MIParser.h"
 #include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
@@ -97,13 +98,15 @@ class MIRParserImpl {
   /// Create an empty function with the given name.
   Function *createDummyFunction(StringRef Name, Module &M);
 
-  bool parseMachineFunctions(Module &M, MachineModuleInfo &MMI);
+  bool parseMachineFunctions(Module &M, MachineModuleInfo &MMI,
+                             ModuleAnalysisManager *FAM = nullptr);
 
   /// Parse the machine function in the current YAML document.
   ///
   ///
   /// Return true if an error occurred.
-  bool parseMachineFunction(Module &M, MachineModuleInfo &MMI);
+  bool parseMachineFunction(Module &M, MachineModuleInfo &MMI,
+                            ModuleAnalysisManager *FAM);
 
   /// Initialize the machine function to the state that's described in the MIR
   /// file.
@@ -275,13 +278,14 @@ MIRParserImpl::parseIRModule(DataLayoutCallbackTy DataLayoutCallback) {
   return M;
 }
 
-bool MIRParserImpl::parseMachineFunctions(Module &M, MachineModuleInfo &MMI) {
+bool MIRParserImpl::parseMachineFunctions(Module &M, MachineModuleInfo &MMI,
+                                          ModuleAnalysisManager *MAM) {
   if (NoMIRDocuments)
     return false;
 
   // Parse the machine functions.
   do {
-    if (parseMachineFunction(M, MMI))
+    if (parseMachineFunction(M, MMI, MAM))
       return true;
     In.nextDocument();
   } while (In.setCurrentDocument());
@@ -303,7 +307,8 @@ Function *MIRParserImpl::createDummyFunction(StringRef Name, Module &M) {
   return F;
 }
 
-bool MIRParserImpl::parseMachineFunction(Module &M, MachineModuleInfo &MMI) {
+bool MIRParserImpl::parseMachineFunction(Module &M, MachineModuleInfo &MMI,
+                                         ModuleAnalysisManager *MAM) {
   // Parse the yaml.
   yaml::MachineFunction YamlMF;
   yaml::EmptyContext Ctx;
@@ -327,14 +332,29 @@ bool MIRParserImpl::parseMachineFunction(Module &M, MachineModuleInfo &MMI) {
                    "' isn't defined in the provided LLVM IR");
     }
   }
-  if (MMI.getMachineFunction(*F) != nullptr)
-    return error(Twine("redefinition of machine function '") + FunctionName +
-                 "'");
 
-  // Create the MachineFunction.
-  MachineFunction &MF = MMI.getOrCreateMachineFunction(*F);
-  if (initializeMachineFunction(YamlMF, MF))
-    return true;
+  if (!MAM) {
+    if (MMI.getMachineFunction(*F) != nullptr)
+      return error(Twine("redefinition of machine function '") + FunctionName +
+                   "'");
+
+    // Create the MachineFunction.
+    MachineFunction &MF = MMI.getOrCreateMachineFunction(*F);
+    if (initializeMachineFunction(YamlMF, MF))
+      return true;
+  } else {
+    auto &FAM =
+        MAM->getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+    if (FAM.getCachedResult<FunctionToMachineFunctionAnalysis>(*F))
+      return error(Twine("redefinition of machine function '") + FunctionName +
+                   "'");
+
+    // Create the MachineFunction.
+    MachineFunction &MF =
+        FAM.getResult<FunctionToMachineFunctionAnalysis>(*F).getMF();
+    if (initializeMachineFunction(YamlMF, MF))
+      return true;
+  }
 
   return false;
 }
@@ -1101,6 +1121,11 @@ bool MIRParser::parseMachineFunctions(Module &M, MachineModuleInfo &MMI) {
   return Impl->parseMachineFunctions(M, MMI);
 }
 
+bool MIRParser::parseMachineFunctions(Module &M, ModuleAnalysisManager &MAM) {
+  auto &MMI = MAM.getResult<MachineModuleAnalysis>(M).getMMI();
+  return Impl->parseMachineFunctions(M, MMI, &MAM);
+}
+
 std::unique_ptr<MIRParser> llvm::createMIRParserFromFile(
     StringRef Filename, SMDiagnostic &Error, LLVMContext &Context,
     std::function<void(Function &)> ProcessIRFunction) {
diff --git a/llvm/lib/CodeGen/MachinePassManager.cpp b/llvm/lib/CodeGen/MachinePassManager.cpp
index 2763193b2c306b..3885ebddc263bd 100644
--- a/llvm/lib/CodeGen/MachinePassManager.cpp
+++ b/llvm/lib/CodeGen/MachinePassManager.cpp
@@ -11,6 +11,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/MachinePassManager.h"
+#include "llvm/CodeGen/FreeMachineFunction.h"
+#include "llvm/CodeGen/FunctionToMachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/IR/PassManagerImpl.h"
@@ -69,9 +71,34 @@ bool MachineFunctionAnalysisManagerModuleProxy::Result::invalidate(
   return false;
 }
 
+template <>
+bool MachineFunctionAnalysisManagerFunctionProxy::Result::invalidate(
+    Function &F, const PreservedAnalyses &PA,
+    FunctionAnalysisManager::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.
+  //
+  // Once function changed by a non-trivial pass, we need to do instruction
+  // selection again.
+  auto PAC = PA.getChecker<MachineFunctionAnalysisManagerFunctionProxy>();
+  if (!PAC.preserved() && !PAC.preservedSet<AllAnalysesOn<Function>>()) {
+    InnerAM->clear();
+    return true;
+  }
+
+  // Return false to indicate that this result is still a valid proxy.
+  return false;
+}
+
 PreservedAnalyses
 ModuleToMachineFunctionPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
-  auto &MMI = AM.getResult<MachineModuleAnalysis>(M).getMMI();
+  // Ensure we have a MachineModuleInfo
+  AM.getResult<MachineModuleAnalysis>(M).getMMI();
+  auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   MachineFunctionAnalysisManager &MFAM =
       AM.getResult<MachineFunctionAnalysisManagerModuleProxy>(M).getManager();
   PassInstrumentation PI = AM.getResult<PassInstrumentationAnalysis>(M);
@@ -82,19 +109,51 @@ ModuleToMachineFunctionPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
     if (F.isDeclaration() || F.hasAvailableExternallyLinkage())
       continue;
 
-    MachineFunction &MF = MMI.getOrCreateMachineFunction(F);
+    MachineFunction &MF =
+        FAM.getResult<FunctionToMachineFunctionAnalysis>(F).getMF();
 
     if (!PI.runBeforePass<MachineFunction>(*Pass, MF))
       continue;
     PreservedAnalyses PassPA = Pass->run(MF, MFAM);
-    if (MMI.getMachineFunction(F)) {
-      MFAM.invalidate(MF, PassPA);
+    MFAM.invalidate(MF, PassPA);
+    if (Pass->name() != FreeMachineFunctionPass::name()) {
       PI.runAfterPass(*Pass, MF, PassPA);
+      PA.intersect(std::move(PassPA));
     } else {
-      MFAM.clear(MF, F.getName());
-      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PassPA);
+      PA.intersect(std::move(PassPA));
+      FAM.invalidate(F, PA);
+      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PA);
     }
+  }
+
+  return PA;
+}
+
+PreservedAnalyses
+FunctionToMachineFunctionPassAdaptor::run(Function &F,
+                                          FunctionAnalysisManager &FAM) {
+  // Do not codegen any 'available_externally' functions at all, they have
+  // definitions outside the translation unit.
+  if (F.isDeclaration() || F.hasAvailableExternallyLinkage())
+    return PreservedAnalyses::all();
+
+  auto &MF = FAM.getResult<FunctionToMachineFunctionAnalysis>(F).getMF();
+  auto &MFAM = FAM.getResult<MachineFunctionAnalysisManagerFunctionProxy>(F)
+                   .getManager();
+  auto PI = FAM.getResult<PassInstrumentationAnalysis>(F);
+  PreservedAnalyses PA = PreservedAnalyses::all();
+
+  if (!PI.runBeforePass<MachineFunction>(*Pass, MF))
+    return PreservedAnalyses::all();
+  PreservedAnalyses PassPA = Pass->run(MF, MFAM);
+  MFAM.invalidate(MF, PassPA);
+  if (Pass->name() != FreeMachineFunctionPass::name()) {
+    PI.runAfterPass(*Pass, MF, PassPA);
+    PA.intersect(std::move(PassPA));
+  } else {
     PA.intersect(std::move(PassPA));
+    FAM.invalidate(F, PA);
+    PI.runAfterPassInvalidated<MachineFunction>(*Pass, PA);
   }
 
   return PA;
@@ -112,25 +171,24 @@ 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();
+  FunctionAnalysisManager &FAM =
+      MFAM.getResult<FunctionAnalysisManagerMachineFunctionProxy>(MF)
+          .getManager();
   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);
+    MFAM.invalidate(MF, PassPA);
+    if (Pass->name() != FreeMachineFunctionPass::name()) {
       PI.runAfterPass(*Pass, MF, PassPA);
+      PA.intersect(std::move(PassPA));
     } else {
-      MFAM.clear(MF, F.getName());
-      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PassPA);
+      PA.intersect(std::move(PassPA));
+      FAM.invalidate(MF.getFunction(), PA);
+      PI.runAfterPassInvalidated<MachineFunction>(*Pass, PA);
     }
-    PA.intersect(std::move(PassPA));
   }
   return PA;
 }
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 8d408ca2363a98..0ffe5e8d60edfb 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -83,6 +83,7 @@
 #include "llvm/CodeGen/ExpandLargeFpConvert.h"
 #include "llvm/CodeGen/ExpandMemCmp.h"
 #include "llvm/CodeGen/FreeMachineFunction.h"
+#include "llvm/CodeGen/FunctionToMachineFunctionAnalysis.h"
 #include "llvm/CodeGen/GCMetadata.h"
 #include "llvm/CodeGen/GlobalMerge.h"
 #include "llvm/CodeGen/HardwareLoops.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index d15f58d7adfae4..98ef4b13f1e586 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -254,6 +254,9 @@ FUNCTION_ANALYSIS("demanded-bits", DemandedBitsAnalysis())
 FUNCTION_ANALYSIS("domfrontier", DominanceFrontierAnalysis())
 FUNCTION_ANALYSIS("domtree", DominatorTreeAnalysis())
 FUNCTION_ANALYSIS("func-properties", FunctionPropertiesAnalysis())
+FUNCTION_ANALYSIS("function-to-machine-function",
+                  FunctionToMachineFunctionAnalysis(
+                      static_cast<const LLVMTargetMachine *>(TM)))
 FUNCTION_ANALYSIS("gc-function", GCFunctionAnalysis())
 FUNCTION_ANALYSIS("inliner-size-estimator", InlineSizeEstimatorAnalysis())
 FUNCTION_ANALYSIS("lazy-value-info", LazyValueAnalysis())
diff --git a/llvm/lib/Passes/StandardInstrumentations.cpp b/llvm/lib/Passes/StandardInstrumentations.cpp
index c18b462258623d..63490c83e85f05 100644
--- a/llvm/lib/Passes/StandardInstrumentations.cpp
+++ b/llvm/lib/Passes/StandardInstrumentations.cpp
@@ -297,14 +297,6 @@ void unwrapAndPrint(raw_ostream &OS, Any IR) {
     auto *M = unwrapModule(IR);
     assert(M && "should have unwrapped module");
     printIR(OS, M);
-
-    if (const auto *MF = unwrapIR<MachineFunction>(IR)) {
-      auto &MMI = MF->getMMI();
-      for (const auto &F : *M) {
-        if (auto *MF = MMI.getMachineFunction(F))
-          MF->print(OS);
-      }
-    }
     return;
   }
 
diff --git a/llvm/tools/llc/NewPMDriver.cpp b/llvm/tools/llc/NewPMDriver.cpp
index 6ae1b8db5e115c..e36e77bbd1ad70 100644
--- a/llvm/tools/llc/NewPMDriver.cpp
+++ b/llvm/tools/llc/NewPMDriver.cpp
@@ -121,10 +121,10 @@ int llvm::compileModuleWithNewPM(
   registerCodeGenCallback(PIC, LLVMTM);
 
   LoopAnalysisManager LAM;
+  MachineFunctionAnalysisManager MFAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
   ModuleAnalysisManager MAM;
-  MachineFunctionAnalysisManager MFAM;
   PassBuilder PB(Target.get(), PipelineTuningOptions(), std::nullopt, &PIC);
   PB.registerModuleAnalyses(MAM);
   PB.registerCGSCCAnalyses(CGAM);
@@ -155,7 +155,7 @@ int llvm::compileModuleWithNewPM(
     MFPM.addPass(FreeMachineFunctionPass());
     MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
 
-    if (MIR->parseMachineFunctions(*M, MMI))
+    if (MIR->parseMachineFunctions(*M, MAM))
       return 1;
   } else {
     ExitOnErr(LLVMTM.buildCodeGenPipeline(
diff --git a/llvm/unittests/CodeGen/PassManagerTest.cpp b/llvm/unittests/CodeGen/PassManagerTest.cpp
index 4283eb01a9c8f2..7b0a7f22c8ef33 100644
--- a/llvm/unittests/CodeGen/PassManagerTest.cpp
+++ b/llvm/unittests/CodeGen/PassManagerTest.cpp
@@ -176,10 +176,10 @@ TEST_F(PassManagerTest, Basic) {
   MachineModuleInfo MMI(LLVMTM);
 
   LoopAnalysisManager LAM;
+  MachineFunctionAnalysisManager MFAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
   ModuleAnalysisManager MAM;
-  MachineFunctionAnalysisManager MFAM;
   PassBuilder PB(TM.get());
   PB.registerModuleAnalyses(MAM);
   PB.registerCGSCCAnalyses(CGAM);
@@ -196,6 +196,7 @@ TEST_F(PassManagerTest, Basic) {
   std::vector<int> Counts;
 
   ModulePassManager MPM;
+  FunctionPassManager FPM;
   MachineFunctionPassManager MFPM;
   MPM.addPass(TestMachineModulePass(Count, Counts));
   MPM.addPass(createModuleToMachineFunctionPassAdaptor(
@@ -203,11 +204,15 @@ TEST_F(PassManagerTest, Basic) {
   MPM.addPass(TestMachineModulePass(Count, Counts));
   MFPM.addPass(TestMachineFunctionPass(Count, Counts));
   MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
+  FPM.addPass(createFunctionToMachineFunctionPassAdaptor(
+      TestMachineFunctionPass(Count, Counts)));
+  MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
 
   MPM.run(*M, MAM);
 
-  EXPECT_EQ((std::vector<int>{10, 16, 18, 20, 30, 36, 38, 40}), Counts);
-  EXPECT_EQ(40, Count);
+  EXPECT_EQ((std::vector<int>{10, 16, 18, 20, 30, 36, 38, 40, 46, 48, 50}),
+            Counts);
+  EXPECT_EQ(50, Count);
 }
 
 } // namespace
diff --git a/llvm/unittests/MIR/PassBuilderCallbacksTest.cpp b/llvm/unittests/MIR/PassBuilderCallbacksTest.cpp
index 6fd4e54a929f40..082ff6c3f67ed0 100644
--- a/llvm/unittests/MIR/PassBuilderCallbacksTest.cpp
+++ b/llvm/unittests/MIR/PassBuilderCallbacksTest.cpp
@@ -316,7 +316,7 @@ class MachineFunctionCallbacksTest : public testing::Test {
   static std::unique_ptr<Module> parseMIR(StringRef MIRCode,
                                           LLVMContext &Context,
                                           TargetMachine &TM,
-                                          MachineModuleInfo &MMI) {
+                                          ModuleAnalysisManager &MAM) {
     SMDiagnostic Diagnostic;
     std::unique_ptr<MemoryBuffer> MBuffer = MemoryBuffer::getMemBuffer(MIRCode);
     std::unique_ptr<MIRParser> MIR =
@@ -330,7 +330,7 @@ class MachineFunctionCallbacksTest : public testing::Test {
     Mod->setModuleIdentifier("module");
     Mod->setDataLayout(TM.createDataLayout());
 
-    [[maybe_unused]] bool Ret = MIR->parseMachineFunctions(*Mod, MMI);
+    [[maybe_unused]] bool Ret = MIR->parseMachineFunctions(*Mod, MAM);
     assert(!Ret);
 
     return Mod;
@@ -355,7 +355,6 @@ class MachineFunctionCallbacksTest : public testing::Test {
       GTEST_SKIP();
 
     MMI = std::make_unique<MachineModuleInfo>(TM.get());
-    M = parseMIR(MIRString, Context, *TM, *MMI);
     PB = std::make_unique<PassBuilder>(TM.get(), PipelineTuningOptions(),
                                        std::nullopt, &PIC);
 
@@ -399,6 +398,7 @@ class MachineFunctionCallbacksTest : public testing::Test {
     PB->registerMachineFunctionAnalyses(MFAM);
     PB->crossRegisterProxies(LAM, FAM, CGAM, MAM, &MFAM);
     MAM.registerPass([&] { return MachineModuleAnalysis(*MMI); });
+    M = parseMIR(MIRString, Context, *TM, MAM);
   }
 };
 
@@ -554,9 +554,6 @@ TEST_F(MachineFunctionCallbacksTest, InstrumentedFreeMFPass2) {
   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.
@@ -564,7 +561,7 @@ TEST_F(MachineFunctionCallbacksTest, InstrumentedFreeMFPass2) {
               runAfterPass(HasNameRegex("FreeMachineFunctionPass"), _, _))
       .Times(0);
   EXPECT_CALL(CallbacksHandle, runAfterPass(HasNameRegex("PassManager"), _, _))
-      .Times(0);
+      .Times(1);
 
   MachineFunctionPassManager MFPM;
   MFPM.addPass(FreeMachineFunctionPass());



More information about the llvm-commits mailing list