[Instrumentation] Support MachineFunctionProperties

paperchalice · paperchalice · commit 00a57c183b55 · 2024-03-02T20:49:19.000+08:00
diff --git a/llvm/include/llvm/IR/PassInstrumentation.h b/llvm/include/llvm/IR/PassInstrumentation.h
@@ -58,6 +58,13 @@
 
 namespace llvm {
 
+namespace detail {
+
+struct MachinePassConcept;
+
+}
+
+class MachineFunction;
 class PreservedAnalyses;
 class StringRef;
 
@@ -78,7 +85,12 @@ class PassInstrumentationCallbacks {
   using BeforePassFunc = bool(StringRef, Any);
   using BeforeSkippedPassFunc = void(StringRef, Any);
   using BeforeNonSkippedPassFunc = void(StringRef, Any);
+  using BeforeNonSkippedMachineFunctionPassFunc =
+      void(const detail::MachinePassConcept &, MachineFunction &);
   using AfterPassFunc = void(StringRef, Any, const PreservedAnalyses &);
+  using AfterMachineFunctionPassFunc = void(const detail::MachinePassConcept &,
+                                            MachineFunction &,
+                                            const PreservedAnalyses &);
   using AfterPassInvalidatedFunc = void(StringRef, const PreservedAnalyses &);
   using BeforeAnalysisFunc = void(StringRef, Any);
   using AfterAnalysisFunc = void(StringRef, Any);
@@ -107,6 +119,11 @@ class PassInstrumentationCallbacks {
     BeforeNonSkippedPassCallbacks.emplace_back(std::move(C));
   }
 
+  template <typename CallableT>
+  void registerBeforeNonSkippedMachineFunctionPassCallback(CallableT C) {
+    BeforeNonSkippedMachineFunctionPassCallbacks.emplace_back(std::move(C));
+  }
+
   template <typename CallableT>
   void registerAfterPassCallback(CallableT C, bool ToFront = false) {
     if (ToFront)
@@ -115,6 +132,11 @@ class PassInstrumentationCallbacks {
       AfterPassCallbacks.emplace_back(std::move(C));
   }
 
+  template <typename CallableT>
+  void registerAfterMachineFunctionPassCallback(CallableT C) {
+    AfterMachineFunctionPassCallbacks.emplace_back(std::move(C));
+  }
+
   template <typename CallableT>
   void registerAfterPassInvalidatedCallback(CallableT C, bool ToFront = false) {
     if (ToFront)
@@ -166,8 +188,12 @@ class PassInstrumentationCallbacks {
   /// These are run on passes that are about to be run.
   SmallVector<llvm::unique_function<BeforeNonSkippedPassFunc>, 4>
       BeforeNonSkippedPassCallbacks;
+  SmallVector<llvm::unique_function<BeforeNonSkippedMachineFunctionPassFunc>, 4>
+      BeforeNonSkippedMachineFunctionPassCallbacks;
   /// These are run on passes that have just run.
   SmallVector<llvm::unique_function<AfterPassFunc>, 4> AfterPassCallbacks;
+  SmallVector<llvm::unique_function<AfterMachineFunctionPassFunc>, 4>
+      AfterMachineFunctionPassCallbacks;
   /// These are run passes that have just run on invalidated IR.
   SmallVector<llvm::unique_function<AfterPassInvalidatedFunc>, 4>
       AfterPassInvalidatedCallbacks;
@@ -247,6 +273,9 @@ class PassInstrumentation {
     return ShouldRun;
   }
 
+  bool runBeforeMachineFunctionPass(const detail::MachinePassConcept &Pass,
+                                    MachineFunction &MF) const;
+
   /// AfterPass instrumentation point - takes \p Pass instance that has
   /// just been executed and constant reference to \p IR it operates on.
   /// \p IR is guaranteed to be valid at this point.
@@ -258,6 +287,10 @@ class PassInstrumentation {
         C(Pass.name(), llvm::Any(&IR), PA);
   }
 
+  void runAfterMachineFunctionPass(const detail::MachinePassConcept &Pass,
+                                   MachineFunction &MF,
+                                   const PreservedAnalyses &PA) const;
+
   /// AfterPassInvalidated instrumentation point - takes \p Pass instance
   /// that has just been executed. For use when IR has been invalidated
   /// by \p Pass execution.
diff --git a/llvm/include/llvm/IR/PassManager.h b/llvm/include/llvm/IR/PassManager.h
@@ -64,6 +64,8 @@ extern llvm::cl::opt<bool> UseNewDbgInfoFormat;
 
 namespace llvm {
 
+class MachineFunction;
+
 // RemoveDIs: Provide facilities for converting debug-info from one form to
 // another, which are no-ops for everything but modules.
 template <class IRUnitT> inline bool shouldConvertDbgInfo(IRUnitT &IR) {
@@ -271,8 +273,10 @@ class PassManager : public PassInfoMixin<
   LLVM_ATTRIBUTE_MINSIZE
       std::enable_if_t<!std::is_same<PassT, PassManager>::value>
       addPass(PassT &&Pass) {
-    using PassModelT =
-        detail::PassModel<IRUnitT, PassT, AnalysisManagerT, ExtraArgTs...>;
+    using PassModelT = std::conditional_t<
+        std::is_same_v<IRUnitT, MachineFunction>,
+        detail::MachinePassModel<PassT>,
+        detail::PassModel<IRUnitT, PassT, AnalysisManagerT, ExtraArgTs...>>;
     // Do not use make_unique or emplace_back, they cause too many template
     // instantiations, causing terrible compile times.
     Passes.push_back(std::unique_ptr<PassConceptT>(
@@ -298,8 +302,9 @@ class PassManager : public PassInfoMixin<
   static bool isRequired() { return true; }
 
 protected:
-  using PassConceptT =
-      detail::PassConcept<IRUnitT, AnalysisManagerT, ExtraArgTs...>;
+  using PassConceptT = std::conditional_t<
+      std::is_same_v<IRUnitT, MachineFunction>, detail::MachinePassConcept,
+      detail::PassConcept<IRUnitT, AnalysisManagerT, ExtraArgTs...>>;
 
   std::vector<std::unique_ptr<PassConceptT>> Passes;
 };
diff --git a/llvm/include/llvm/IR/PassManagerInternal.h b/llvm/include/llvm/IR/PassManagerInternal.h
@@ -328,6 +328,9 @@ struct AnalysisPassModel
   PassT Pass;
 };
 
+struct MachinePassConcept;
+template <typename PassT> struct MachinePassModel;
+
 } // end namespace detail
 
 } // end namespace llvm
diff --git a/llvm/include/llvm/Passes/StandardInstrumentations.h b/llvm/include/llvm/Passes/StandardInstrumentations.h
@@ -563,6 +563,11 @@ class PrintCrashIRInstrumentation {
   static void SignalHandler(void *);
 };
 
+class MachineFunctionPropertiesInstrumentation {
+public:
+  void registerCallbacks(PassInstrumentationCallbacks &PIC, bool Verify);
+};
+
 /// This class provides an interface to register all the standard pass
 /// instrumentations and manages their state (if any).
 class StandardInstrumentations {
@@ -580,6 +585,7 @@ class StandardInstrumentations {
   PrintCrashIRInstrumentation PrintCrashIR;
   IRChangedTester ChangeTester;
   VerifyInstrumentation Verify;
+  MachineFunctionPropertiesInstrumentation MFProp;
 
   bool VerifyEach;
 
diff --git a/llvm/lib/CodeGen/MachinePassManager.cpp b/llvm/lib/CodeGen/MachinePassManager.cpp
@@ -119,13 +119,13 @@ PassManager<MachineFunction>::run(MachineFunction &MF,
           ->getMMI();
   PreservedAnalyses PA = PreservedAnalyses::all();
   for (auto &Pass : Passes) {
-    if (!PI.runBeforePass<MachineFunction>(*Pass, MF))
+    if (!PI.runBeforeMachineFunctionPass(*Pass, MF))
       continue;
 
     PreservedAnalyses PassPA = Pass->run(MF, MFAM);
     if (MMI.getMachineFunction(F)) {
       MFAM.invalidate(MF, PassPA);
-      PI.runAfterPass(*Pass, MF, PassPA);
+      PI.runAfterMachineFunctionPass(*Pass, MF, PassPA);
     } else {
       MFAM.clear(MF, F.getName());
       PI.runAfterPassInvalidated<MachineFunction>(*Pass, PassPA);
@@ -135,4 +135,25 @@ PassManager<MachineFunction>::run(MachineFunction &MF,
   return PA;
 }
 
+bool PassInstrumentation::runBeforeMachineFunctionPass(
+    const detail::MachinePassConcept &Pass, MachineFunction &MF) const {
+  bool ShouldRun = runBeforePass(Pass, MF);
+  if (!ShouldRun)
+    return false;
+
+  if (Callbacks)
+    for (auto &C : Callbacks->BeforeNonSkippedMachineFunctionPassCallbacks)
+      C(Pass, MF);
+  return true;
+}
+
+void PassInstrumentation::runAfterMachineFunctionPass(
+    const detail::MachinePassConcept &Pass, MachineFunction &MF,
+    const PreservedAnalyses &PA) const {
+  runAfterPass(Pass, MF, PA);
+  if (Callbacks)
+    for (auto &C : Callbacks->AfterMachineFunctionPassCallbacks)
+      C(Pass, MF, PA);
+}
+
 } // namespace llvm
diff --git a/llvm/lib/Passes/StandardInstrumentations.cpp b/llvm/lib/Passes/StandardInstrumentations.cpp
@@ -20,6 +20,7 @@
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Module.h"
@@ -2415,6 +2416,35 @@ void PrintCrashIRInstrumentation::registerCallbacks(
       });
 }
 
+void MachineFunctionPropertiesInstrumentation::registerCallbacks(
+    PassInstrumentationCallbacks &PIC, bool Verify) {
+  if (Verify) {
+    PIC.registerBeforeNonSkippedMachineFunctionPassCallback(
+        [](const detail::MachinePassConcept &Pass, MachineFunction &MF) {
+          auto &MFProps = MF.getProperties();
+          auto RequiredProperties = Pass.getRequiredProperties();
+          if (!MFProps.verifyRequiredProperties(RequiredProperties)) {
+            errs() << "MachineFunctionProperties required by " << Pass.name()
+                   << " pass are not met by function " << MF.getName() << ".\n"
+                   << "Required properties: ";
+            RequiredProperties.print(errs());
+            errs() << "\nCurrent properties: ";
+            MFProps.print(errs());
+            errs() << "\n";
+            report_fatal_error("MachineFunctionProperties check failed");
+          }
+        });
+  }
+
+  PIC.registerAfterMachineFunctionPassCallback(
+      [](const detail::MachinePassConcept &Pass, MachineFunction &MF,
+         const PreservedAnalyses &) {
+        auto &MFProp = MF.getProperties();
+        MFProp.reset(Pass.getClearedProperties());
+        MFProp.set(Pass.getSetProperties());
+      });
+}
+
 void StandardInstrumentations::registerCallbacks(
     PassInstrumentationCallbacks &PIC, ModuleAnalysisManager *MAM) {
   PrintIR.registerCallbacks(PIC);
@@ -2440,6 +2470,7 @@ void StandardInstrumentations::registerCallbacks(
   // AfterCallbacks by its `registerCallbacks`. This is necessary
   // to ensure that other callbacks are not included in the timings.
   TimeProfilingPasses.registerCallbacks(PIC);
+  MFProp.registerCallbacks(PIC, VerifyEach);
 }
 
 template class ChangeReporter<std::string>;
