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

[via llvm-commits]

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

Add `FunctionToMachineFunctionPassAdaptor`, currently codegen pipeline has the pattern `module->function->machine-function` if `MachineOutlinerPass` is not enabled. This adaptor allows `FreeMachineFunction` to be a function pass, then we don't need to handle function analyses invalidation in `MachinePassManager` etc.
Stacked on #88610.
Issue: #84397.

>From 2cc7d8c7a78cdfca7ea93291f67c87d7861dc989 Mon Sep 17 00:00:00 2001
From: PaperChalice <liujunchang97 at outlook.com>
Date: Sat, 13 Apr 2024 20:30:45 +0800
Subject: [PATCH] [NewPM] Add `FunctionToMachineFunctionPassAdaptor`

---
 .../include/llvm/CodeGen/MachinePassManager.h | 43 +++++++++++++
 llvm/lib/CodeGen/MachinePassManager.cpp       | 62 +++++++++++++++++++
 llvm/lib/Passes/PassBuilder.cpp               |  9 +++
 3 files changed, 114 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/MachinePassManager.h b/llvm/include/llvm/CodeGen/MachinePassManager.h
index 4f0b6ba2b1e738..e54b536b54b65d 100644
--- a/llvm/include/llvm/CodeGen/MachinePassManager.h
+++ b/llvm/include/llvm/CodeGen/MachinePassManager.h
@@ -109,6 +109,16 @@ bool MachineFunctionAnalysisManagerModuleProxy::Result::invalidate(
 extern template class InnerAnalysisManagerProxy<MachineFunctionAnalysisManager,
                                                 Module>;
 
+using MachineFunctionAnalysisManagerFunctionProxy =
+    InnerAnalysisManagerProxy<MachineFunctionAnalysisManager, Function>;
+
+template <>
+bool MachineFunctionAnalysisManagerFunctionProxy::Result::invalidate(
+    Function &F, const PreservedAnalyses &PA,
+    FunctionAnalysisManager::Invalidator &Inv);
+extern template class InnerAnalysisManagerProxy<MachineFunctionAnalysisManager,
+                                                Function>;
+
 extern template class OuterAnalysisManagerProxy<ModuleAnalysisManager,
                                                 MachineFunction>;
 /// Provide the \c ModuleAnalysisManager to \c Function proxy.
@@ -220,6 +230,39 @@ 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::PassModel<MachineFunction, MachineFunctionPassT,
+                                       MachineFunctionAnalysisManager>;
+  // 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/MachinePassManager.cpp b/llvm/lib/CodeGen/MachinePassManager.cpp
index 2763193b2c306b..4edd84d256644c 100644
--- a/llvm/lib/CodeGen/MachinePassManager.cpp
+++ b/llvm/lib/CodeGen/MachinePassManager.cpp
@@ -25,6 +25,8 @@ template class AnalysisManager<MachineFunction>;
 template class PassManager<MachineFunction>;
 template class InnerAnalysisManagerProxy<MachineFunctionAnalysisManager,
                                          Module>;
+template class InnerAnalysisManagerProxy<MachineFunctionAnalysisManager,
+                                         Function>;
 template class OuterAnalysisManagerProxy<ModuleAnalysisManager,
                                          MachineFunction>;
 
@@ -69,6 +71,29 @@ 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<FunctionAnalysisManagerModuleProxy>();
+  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();
@@ -100,6 +125,36 @@ ModuleToMachineFunctionPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
   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<MachineFunctionAnalysis>(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;
+}
+
 void ModuleToMachineFunctionPassAdaptor::printPipeline(
     raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
   OS << "machine-function(";
@@ -107,6 +162,13 @@ void ModuleToMachineFunctionPassAdaptor::printPipeline(
   OS << ')';
 }
 
+void FunctionToMachineFunctionPassAdaptor::printPipeline(
+    raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
+  OS << "machine-function(";
+  Pass->printPipeline(OS, MapClassName2PassName);
+  OS << ')';
+}
+
 template <>
 PreservedAnalyses
 PassManager<MachineFunction>::run(MachineFunction &MF,
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 8d408ca2363a98..4526aedb1879d3 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1725,6 +1725,13 @@ Error PassBuilder::parseFunctionPass(FunctionPassManager &FPM,
                                                   UseBFI, UseBPI));
       return Error::success();
     }
+    if (Name == "machine-function") {
+      MachineFunctionPassManager MFPM;
+      if (auto Err = parseMachinePassPipeline(MFPM, InnerPipeline))
+        return Err;
+      FPM.addPass(createFunctionToMachineFunctionPassAdaptor(std::move(MFPM)));
+      return Error::success();
+    }
     if (auto Count = parseRepeatPassName(Name)) {
       FunctionPassManager NestedFPM;
       if (auto Err = parseFunctionPassPipeline(NestedFPM, InnerPipeline))
@@ -1975,6 +1982,8 @@ void PassBuilder::crossRegisterProxies(LoopAnalysisManager &LAM,
   if (MFAM) {
     MAM.registerPass(
         [&] { return MachineFunctionAnalysisManagerModuleProxy(*MFAM); });
+    FAM.registerPass(
+        [&] { return MachineFunctionAnalysisManagerFunctionProxy(*MFAM); });
     MFAM->registerPass(
         [&] { return ModuleAnalysisManagerMachineFunctionProxy(MAM); });
     MFAM->registerPass(