[NewPM] Add FunctionToMachineFunctionPassAdaptor

Author: paperchalice

llvm/include/llvm/CodeGen/MachinePassManager.h

@@ -129,16 +129,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,15 +201,46 @@ 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(
       std::unique_ptr<ModuleToMachineFunctionPassAdaptor::PassConceptT>(
           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) {

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;
 }

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"

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);
@@ -203,11 +203,15 @@ TEST_F(PassManagerTest, Basic) {
   MPM.addPass(TestMachineModulePass(Count, Counts));
   MFPM.addPass(TestMachineFunctionPass(Count, Counts));
   MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
+  MPM.addPass(createModuleToFunctionPassAdaptor(
+      createFunctionToMachineFunctionPassAdaptor(
+          TestMachineFunctionPass(Count, Counts))));
 
   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