[CodeGen] Let PassBuilder support machine passes (#76320)

`PassBuilder` would be a better place to parse MIR pipeline. We can
reuse the code to support parsing pass with parameters and targets can
reuse `registerPassBuilderCallbacks` to register the target specific
passes. `PassBuilder` also has ability to check whether a Pass is a
machine pass.

Author: paperchalice

llvm/include/llvm/Passes/PassBuilder.h

@@ -16,6 +16,7 @@
 #define LLVM_PASSES_PASSBUILDER_H
 
 #include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Passes/OptimizationLevel.h"
 #include "llvm/Support/Error.h"
@@ -165,6 +166,14 @@ class PassBuilder {
   /// additional analyses.
   void registerLoopAnalyses(LoopAnalysisManager &LAM);
 
+  /// Registers all available machine function analysis passes.
+  ///
+  /// This is an interface that can be used to populate a \c
+  /// MachineFunctionAnalysisManager with all registered function analyses.
+  /// Callers can still manually register any additional analyses. Callers can
+  /// also pre-register analyses and this will not override those.
+  void registerMachineFunctionAnalyses(MachineFunctionAnalysisManager &MFAM);
+
   /// Construct the core LLVM function canonicalization and simplification
   /// pipeline.
   ///
@@ -352,6 +361,18 @@ class PassBuilder {
   Error parsePassPipeline(LoopPassManager &LPM, StringRef PipelineText);
   /// @}}
 
+  /// Parse a textual MIR pipeline into the provided \c MachineFunctionPass
+  /// manager.
+  /// The format of the textual machine pipeline is a comma separated list of
+  /// machine pass names:
+  ///
+  ///   machine-funciton-pass,machine-module-pass,...
+  ///
+  /// There is no need to specify the pass nesting, and this function
+  /// currently cannot handle the pass nesting.
+  Error parsePassPipeline(MachineFunctionPassManager &MFPM,
+                          StringRef PipelineText);
+
   /// Parse a textual alias analysis pipeline into the provided AA manager.
   ///
   /// The format of the textual AA pipeline is a comma separated list of AA
@@ -520,6 +541,10 @@ class PassBuilder {
       const std::function<void(ModuleAnalysisManager &)> &C) {
     ModuleAnalysisRegistrationCallbacks.push_back(C);
   }
+  void registerAnalysisRegistrationCallback(
+      const std::function<void(MachineFunctionAnalysisManager &)> &C) {
+    MachineFunctionAnalysisRegistrationCallbacks.push_back(C);
+  }
   /// @}}
 
   /// {{@ Register pipeline parsing callbacks with this pass builder instance.
@@ -546,6 +571,11 @@ class PassBuilder {
                                ArrayRef<PipelineElement>)> &C) {
     ModulePipelineParsingCallbacks.push_back(C);
   }
+  void registerPipelineParsingCallback(
+      const std::function<bool(StringRef Name, MachineFunctionPassManager &)>
+          &C) {
+    MachinePipelineParsingCallbacks.push_back(C);
+  }
   /// @}}
 
   /// Register a callback for a top-level pipeline entry.
@@ -616,8 +646,12 @@ class PassBuilder {
   Error parseCGSCCPass(CGSCCPassManager &CGPM, const PipelineElement &E);
   Error parseFunctionPass(FunctionPassManager &FPM, const PipelineElement &E);
   Error parseLoopPass(LoopPassManager &LPM, const PipelineElement &E);
+  Error parseMachinePass(MachineFunctionPassManager &MFPM,
+                         const PipelineElement &E);
   bool parseAAPassName(AAManager &AA, StringRef Name);
 
+  Error parseMachinePassPipeline(MachineFunctionPassManager &MFPM,
+                                 ArrayRef<PipelineElement> Pipeline);
   Error parseLoopPassPipeline(LoopPassManager &LPM,
                               ArrayRef<PipelineElement> Pipeline);
   Error parseFunctionPassPipeline(FunctionPassManager &FPM,
@@ -699,6 +733,11 @@ class PassBuilder {
   // AA callbacks
   SmallVector<std::function<bool(StringRef Name, AAManager &AA)>, 2>
       AAParsingCallbacks;
+  // Machine pass callbackcs
+  SmallVector<std::function<void(MachineFunctionAnalysisManager &)>, 2>
+      MachineFunctionAnalysisRegistrationCallbacks;
+  SmallVector<std::function<bool(StringRef, MachineFunctionPassManager &)>, 2>
+      MachinePipelineParsingCallbacks;
 };
 
 /// This utility template takes care of adding require<> and invalidate<>

llvm/lib/Passes/PassBuilder.cpp

@@ -392,6 +392,12 @@ PassBuilder::PassBuilder(TargetMachine *TM, PipelineTuningOptions PTO,
 #define CGSCC_ANALYSIS(NAME, CREATE_PASS)                                      \
   PIC->addClassToPassName(decltype(CREATE_PASS)::name(), NAME);
 #include "PassRegistry.def"
+
+#define MACHINE_FUNCTION_ANALYSIS(NAME, PASS_NAME, CONSTRUCTOR)                \
+  PIC->addClassToPassName(PASS_NAME::name(), NAME);
+#define MACHINE_FUNCTION_PASS(NAME, PASS_NAME, CONSTRUCTOR)                    \
+  PIC->addClassToPassName(PASS_NAME::name(), NAME);
+#include "llvm/CodeGen/MachinePassRegistry.def"
   }
 }
 
@@ -427,6 +433,17 @@ void PassBuilder::registerFunctionAnalyses(FunctionAnalysisManager &FAM) {
     C(FAM);
 }
 
+void PassBuilder::registerMachineFunctionAnalyses(
+    MachineFunctionAnalysisManager &MFAM) {
+
+#define MACHINE_FUNCTION_ANALYSIS(NAME, PASS_NAME, CONSTRUCTOR)                \
+  MFAM.registerPass([&] { return PASS_NAME(); });
+#include "llvm/CodeGen/MachinePassRegistry.def"
+
+  for (auto &C : MachineFunctionAnalysisRegistrationCallbacks)
+    C(MFAM);
+}
+
 void PassBuilder::registerLoopAnalyses(LoopAnalysisManager &LAM) {
 #define LOOP_ANALYSIS(NAME, CREATE_PASS)                                       \
   LAM.registerPass([&] { return CREATE_PASS; });
@@ -1781,6 +1798,33 @@ Error PassBuilder::parseLoopPass(LoopPassManager &LPM,
                                  inconvertibleErrorCode());
 }
 
+Error PassBuilder::parseMachinePass(MachineFunctionPassManager &MFPM,
+                                    const PipelineElement &E) {
+  StringRef Name = E.Name;
+  if (!E.InnerPipeline.empty())
+    return make_error<StringError>("invalid pipeline",
+                                   inconvertibleErrorCode());
+
+#define MACHINE_MODULE_PASS(NAME, PASS_NAME, CONSTRUCTOR)                      \
+  if (Name == NAME) {                                                          \
+    MFPM.addPass(PASS_NAME());                                                 \
+    return Error::success();                                                   \
+  }
+#define MACHINE_FUNCTION_PASS(NAME, PASS_NAME, CONSTRUCTOR)                    \
+  if (Name == NAME) {                                                          \
+    MFPM.addPass(PASS_NAME());                                                 \
+    return Error::success();                                                   \
+  }
+#include "llvm/CodeGen/MachinePassRegistry.def"
+
+  for (auto &C : MachinePipelineParsingCallbacks)
+    if (C(Name, MFPM))
+      return Error::success();
+  return make_error<StringError>(
+      formatv("unknown machine pass '{0}'", Name).str(),
+      inconvertibleErrorCode());
+}
+
 bool PassBuilder::parseAAPassName(AAManager &AA, StringRef Name) {
 #define MODULE_ALIAS_ANALYSIS(NAME, CREATE_PASS)                               \
   if (Name == NAME) {                                                          \
@@ -1802,6 +1846,15 @@ bool PassBuilder::parseAAPassName(AAManager &AA, StringRef Name) {
   return false;
 }
 
+Error PassBuilder::parseMachinePassPipeline(
+    MachineFunctionPassManager &MFPM, ArrayRef<PipelineElement> Pipeline) {
+  for (const auto &Element : Pipeline) {
+    if (auto Err = parseMachinePass(MFPM, Element))
+      return Err;
+  }
+  return Error::success();
+}
+
 Error PassBuilder::parseLoopPassPipeline(LoopPassManager &LPM,
                                          ArrayRef<PipelineElement> Pipeline) {
   for (const auto &Element : Pipeline) {
@@ -1961,6 +2014,20 @@ Error PassBuilder::parsePassPipeline(LoopPassManager &CGPM,
   return Error::success();
 }
 
+Error PassBuilder::parsePassPipeline(MachineFunctionPassManager &MFPM,
+                                     StringRef PipelineText) {
+  auto Pipeline = parsePipelineText(PipelineText);
+  if (!Pipeline || Pipeline->empty())
+    return make_error<StringError>(
+        formatv("invalid machine pass pipeline '{0}'", PipelineText).str(),
+        inconvertibleErrorCode());
+
+  if (auto Err = parseMachinePassPipeline(MFPM, *Pipeline))
+    return Err;
+
+  return Error::success();
+}
+
 Error PassBuilder::parseAAPipeline(AAManager &AA, StringRef PipelineText) {
   // If the pipeline just consists of the word 'default' just replace the AA
   // manager with our default one.
@@ -2055,6 +2122,21 @@ void PassBuilder::printPassNames(raw_ostream &OS) {
   OS << "Loop analyses:\n";
 #define LOOP_ANALYSIS(NAME, CREATE_PASS) printPassName(NAME, OS);
 #include "PassRegistry.def"
+
+  OS << "Machine module passes (WIP):\n";
+#define MACHINE_MODULE_PASS(NAME, PASS_NAME, CONSTRUCTOR)                      \
+  printPassName(NAME, OS);
+#include "llvm/CodeGen/MachinePassRegistry.def"
+
+  OS << "Machine function passes (WIP):\n";
+#define MACHINE_FUNCTION_PASS(NAME, PASS_NAME, CONSTRUCTOR)                    \
+  printPassName(NAME, OS);
+#include "llvm/CodeGen/MachinePassRegistry.def"
+
+  OS << "Machine function analyses (WIP):\n";
+#define MACHINE_FUNCTION_ANALYSIS(NAME, PASS_NAME, CONSTRUCTOR)                \
+  printPassName(NAME, OS);
+#include "llvm/CodeGen/MachinePassRegistry.def"
 }
 
 void PassBuilder::registerParseTopLevelPipelineCallback(

llvm/unittests/MIR/CMakeLists.txt

@@ -6,13 +6,15 @@ set(LLVM_LINK_COMPONENTS
   FileCheck
   MC
   MIRParser
+  Passes
   Support
   Target
   TargetParser
   )
 
 add_llvm_unittest(MIRTests
   MachineMetadata.cpp
+  PassBuilderCallbacksTest.cpp
   )
 
 target_link_libraries(MIRTests PRIVATE LLVMTestingSupport)