[CodeGen][NPM] Support generic regalloc-npm option

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -103,6 +103,7 @@
 #include "llvm/IRPrinter/IRPrintingPasses.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCTargetOptions.h"
+#include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
@@ -120,6 +121,7 @@
 #include "llvm/Transforms/Utils/EntryExitInstrumenter.h"
 #include "llvm/Transforms/Utils/LowerInvoke.h"
 #include <cassert>
+#include <tuple>
 #include <type_traits>
 #include <utility>
 
@@ -161,8 +163,9 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
 public:
   explicit CodeGenPassBuilder(TargetMachineT &TM,
                               const CGPassBuilderOption &Opts,
-                              PassInstrumentationCallbacks *PIC)
-      : TM(TM), Opt(Opts), PIC(PIC) {
+                              PassInstrumentationCallbacks *PIC,
+                              PassBuilder &PB)
+      : TM(TM), Opt(Opts), PIC(PIC), PB(PB) {
     // Target could set CGPassBuilderOption::MISchedPostRA to true to achieve
     //     substitutePass(&PostRASchedulerID, &PostMachineSchedulerID)
 
@@ -298,6 +301,7 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
   TargetMachineT &TM;
   CGPassBuilderOption Opt;
   PassInstrumentationCallbacks *PIC;
+  PassBuilder &PB;
 
   template <typename TMC> TMC &getTM() const { return static_cast<TMC &>(TM); }
   CodeGenOptLevel getOptLevel() const { return TM.getOptLevel(); }
@@ -505,6 +509,15 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
   /// addMachinePasses helper to create the target-selected or overriden
   /// regalloc pass.
   void addRegAllocPass(AddMachinePass &, bool Optimized) const;
+  /// Read the --regalloc-npm option to add the next pass in line.
+  /// Returns false if no pass is left in the option.
+  bool addRegAllocPassFromOpt(AddMachinePass &,
+                              StringRef MatchPassTo = StringRef{}) const;
+  /// Add the next pass in the cli option or the pass specified if no pass is
+  /// left in the option.
+  template <typename RegAllocPassBuilderT>
+  void addRegAllocPassOrOpt(AddMachinePass &,
+                            RegAllocPassBuilderT PassBuilder) const;
 
   /// Add core register alloator passes which do the actual register assignment
   /// and rewriting. \returns true if any passes were added.
@@ -601,6 +614,11 @@ Error CodeGenPassBuilder<Derived, TargetMachineT>::buildPipeline(
   if (PrintMIR)
     addPass(PrintMIRPass(Out), /*Force=*/true);
 
+  if (!Opt.RegAllocPipeline.empty())
+    return make_error<StringError>(
+        "extra passes in regalloc pipeline: " + Opt.RegAllocPipeline,
+        std::make_error_code(std::errc::invalid_argument));
+
   return verifyStartStop(*StartStopInfo);
 }
 
@@ -1098,6 +1116,48 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
     addPass(RegAllocFastPass());
 }
 
+template <typename Derived, typename TargetMachineT>
+template <typename RegAllocPassBuilderT>
+void CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPassOrOpt(
+    AddMachinePass &addPass, RegAllocPassBuilderT PassBuilder) const {
+  if (!addRegAllocPassFromOpt(addPass))
+    addPass(std::move(PassBuilder()));
+}
+
+template <typename Derived, typename TargetMachineT>
+bool CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPassFromOpt(
+    AddMachinePass &addPass, StringRef MatchPassTo) const {
+  if (!Opt.RegAllocPipeline.empty()) {
+    StringRef PassOpt;
+    std::tie(PassOpt, Opt.RegAllocPipeline) = Opt.RegAllocPipeline.split(',');
+    // Reuse the registered parser to parse the pass name.
+#define RA_PASS_WITH_PARAMS(NAME, CLASS, CREATE_PASS, PARSER, PARAMS)          \
+  if (PB.checkParametrizedPassName(PassOpt, NAME)) {                           \
+    auto Params = PB.parsePassParameters(PARSER, PassOpt, NAME,                \
+                                         const_cast<const PassBuilder &>(PB)); \
+    if (!Params) {                                                             \
+      auto Err = Params.takeError();                                           \
+      ExitOnError()(std::move(Err));                                           \
+    }                                                                          \
+    if (!MatchPassTo.empty()) {                                                \
+      if (MatchPassTo != CLASS)                                                \
+        report_fatal_error("expected " +                                       \
+                               PIC->getPassNameForClassName(MatchPassTo) +     \
+                               " in option --regalloc-npm",                    \
+                           false);                                             \
+    }                                                                          \
+    addPass(CREATE_PASS(Params.get()));                                        \
+    return true;                                                               \
+  }
+#include "llvm/Passes/MachinePassRegistry.def"
+    if (PassOpt != "default") {
+      report_fatal_error("unknown register allocator pass: " + PassOpt, false);
+    }
+  }
+  // If user did not give a specific pass, use the default provided.
+  return false;
+}
+
 /// Find and instantiate the register allocation pass requested by this target
 /// at the current optimization level.  Different register allocators are
 /// defined as separate passes because they may require different analysis.
@@ -1108,22 +1168,13 @@ template <typename Derived, typename TargetMachineT>
 void CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPass(
     AddMachinePass &addPass, bool Optimized) const {
   // Use the specified -regalloc-npm={basic|greedy|fast|pbqp}
-  if (Opt.RegAlloc > RegAllocType::Default) {
-    switch (Opt.RegAlloc) {
-    case RegAllocType::Fast:
-      addPass(RegAllocFastPass());
-      break;
-    case RegAllocType::Greedy:
-      addPass(RAGreedyPass());
-      break;
-    default:
-      report_fatal_error("register allocator not supported yet", false);
-    }
-    return;
+  StringRef RegAllocPassName;
+  if (!Optimized)
+    RegAllocPassName = RegAllocFastPass::name();
+
+  if (!addRegAllocPassFromOpt(addPass, RegAllocPassName)) {
+    derived().addTargetRegisterAllocator(addPass, Optimized);
   }
-  // -regalloc=default or unspecified, so pick based on the optimization level
-  // or ask the target for the regalloc pass.
-  derived().addTargetRegisterAllocator(addPass, Optimized);
 }
 
 template <typename Derived, typename TargetMachineT>

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -222,8 +222,8 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
 MACHINE_FUNCTION_PASS_WITH_PARAMS(
     "branch-folder", "BranchFolderPass",
     [](bool EnableTailMerge) { return BranchFolderPass(EnableTailMerge); },
-    [](StringRef Params) {
-      return parseSinglePassOption(Params, "enable-tail-merge",
+    [](StringRef Params, const PassBuilder &) {
+      return PassBuilder::parseSinglePassOption(Params, "enable-tail-merge",
                                    "BranchFolderPass");
     },
     "enable-tail-merge")
@@ -233,8 +233,8 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
     [](bool ShouldEmitDebugEntryValues) {
       return LiveDebugValuesPass(ShouldEmitDebugEntryValues);
     },
-    [](StringRef Params) {
-      return parseSinglePassOption(Params, "emit-debug-entry-values",
+    [](StringRef Params, const PassBuilder &) {
+      return PassBuilder::parseSinglePassOption(Params, "emit-debug-entry-values",
                                    "LiveDebugValuesPass");
     },
     "emit-debug-entry-values")
@@ -247,27 +247,36 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
     parseMachineSinkingPassOptions, "enable-sink-fold")
 
 MACHINE_FUNCTION_PASS_WITH_PARAMS(
+    "virt-reg-rewriter", "VirtRegRewriterPass",
+    [](bool ClearVirtRegs) { return VirtRegRewriterPass(ClearVirtRegs); },
+    parseVirtRegRewriterPassOptions, "no-clear-vregs;clear-vregs")
+
+#ifndef RA_PASS_WITH_PARAMS
+// Define MachineFunction passes that are register allocators.
+// This is to differentiate them from other MachineFunction passes
+// to be used in the --regalloc-npm option.
+#define RA_PASS_WITH_PARAMS(NAME, CLASS, CREATE_PASS, PARSER, PARAMS)          \
+  MACHINE_FUNCTION_PASS_WITH_PARAMS(NAME, CLASS, CREATE_PASS, PARSER, PARAMS)
+#endif
+
+RA_PASS_WITH_PARAMS(
     "regallocfast", "RegAllocFastPass",
     [](RegAllocFastPass::Options Opts) { return RegAllocFastPass(Opts); },
-    [PB = this](StringRef Params) {
-      return parseRegAllocFastPassOptions(*PB, Params);
+    [](StringRef Params, const PassBuilder &PB) {
+      return parseRegAllocFastPassOptions(PB, Params);
     },
     "filter=reg-filter;no-clear-vregs")
 
 // 'all' is the default filter.
-MACHINE_FUNCTION_PASS_WITH_PARAMS(
+RA_PASS_WITH_PARAMS(
     "greedy", "RAGreedyPass",
     [](RAGreedyPass::Options Opts) { return RAGreedyPass(Opts); },
-    [PB = this](StringRef Params) {
-      return parseRegAllocGreedyFilterFunc(*PB, Params);
-    }, "reg-filter"
-)
-
-MACHINE_FUNCTION_PASS_WITH_PARAMS(
-    "virt-reg-rewriter", "VirtRegRewriterPass",
-    [](bool ClearVirtRegs) { return VirtRegRewriterPass(ClearVirtRegs); },
-    parseVirtRegRewriterPassOptions, "no-clear-vregs;clear-vregs")
+    [](StringRef Params, const PassBuilder &PB) {
+      return parseRegAllocGreedyFilterFunc(PB, Params);
+    },
+    "reg-filter")
 
+#undef RA_PASS_WITH_PARAMS
 #undef MACHINE_FUNCTION_PASS_WITH_PARAMS
 
 // After a pass is converted to new pass manager, its entry should be moved from

llvm/include/llvm/Passes/PassBuilder.h

@@ -18,6 +18,7 @@
 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/RegAllocCommon.h"
+#include "llvm/CodeGen/RegAllocGreedyPass.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Passes/OptimizationLevel.h"
 #include "llvm/Support/Error.h"
@@ -397,7 +398,7 @@ class PassBuilder {
 
   /// Parse RegAllocFilterName to get RegAllocFilterFunc.
   std::optional<RegAllocFilterFunc>
-  parseRegAllocFilter(StringRef RegAllocFilterName);
+  parseRegAllocFilter(StringRef RegAllocFilterName) const;
 
   /// Print pass names.
   void printPassNames(raw_ostream &OS);
@@ -688,11 +689,13 @@ class PassBuilder {
   /// parameter list in a form of a custom parameters type, all wrapped into
   /// Expected<> template class.
   ///
-  template <typename ParametersParseCallableT>
+  template <typename ParametersParseCallableT, typename... ExtraArgs>
   static auto parsePassParameters(ParametersParseCallableT &&Parser,
-                                  StringRef Name, StringRef PassName)
-      -> decltype(Parser(StringRef{})) {
-    using ParametersT = typename decltype(Parser(StringRef{}))::value_type;
+                                  StringRef Name, StringRef PassName,
+                                  ExtraArgs &&...Args)
+      -> decltype(Parser(StringRef{}, std::forward<ExtraArgs>(Args)...)) {
+    using ParametersT = typename decltype(Parser(
+        StringRef{}, std::forward<ExtraArgs>(Args)...))::value_type;
 
     StringRef Params = Name;
     if (!Params.consume_front(PassName)) {
@@ -704,7 +707,8 @@ class PassBuilder {
       llvm_unreachable("invalid format for parametrized pass name");
     }
 
-    Expected<ParametersT> Result = Parser(Params);
+    Expected<ParametersT> Result =
+        Parser(Params, std::forward<ExtraArgs>(Args)...);
     assert((Result || Result.template errorIsA<StringError>()) &&
            "Pass parameter parser can only return StringErrors.");
     return Result;
@@ -980,6 +984,18 @@ class NoOpLoopAnalysis : public AnalysisInfoMixin<NoOpLoopAnalysis> {
 /// Common option used by multiple tools to print pipeline passes
 extern cl::opt<bool> PrintPipelinePasses;
 
+Expected<RAGreedyPass::Options>
+parseRegAllocGreedyFilterFunc(const PassBuilder &PB, StringRef Params);
+
+Expected<RegAllocFastPass::Options>
+parseRegAllocFastPassOptions(const PassBuilder &PB, StringRef Params);
+
+Expected<bool> parseMachineSinkingPassOptions(StringRef Params,
+                                              const PassBuilder &);
+Expected<bool> parseMachineBlockPlacementPassOptions(StringRef Params,
+                                                     const PassBuilder &);
+Expected<bool> parseVirtRegRewriterPassOptions(StringRef Params,
+                                               const PassBuilder &);
 }
 
 #endif

llvm/include/llvm/Passes/TargetPassRegistry.inc

@@ -83,7 +83,7 @@ if (PIC) {
 
 #define ADD_PASS_WITH_PARAMS(NAME, CREATE_PASS, PARSER)                        \
   if (PassBuilder::checkParametrizedPassName(Name, NAME)) {                    \
-    auto Params = PassBuilder::parsePassParameters(PARSER, Name, NAME);        \
+    auto Params = PassBuilder::parsePassParameters(PARSER, Name, NAME, PB);    \
     if (!Params) {                                                             \
       errs() << NAME ": " << toString(Params.takeError()) << '\n';             \
       return false;                                                            \

llvm/include/llvm/Target/CGPassBuilderOption.h

@@ -21,22 +21,6 @@
 namespace llvm {
 
 enum class RunOutliner { TargetDefault, AlwaysOutline, NeverOutline };
-enum class RegAllocType { Unset, Default, Basic, Fast, Greedy, PBQP };
-
-class RegAllocTypeParser : public cl::parser<RegAllocType> {
-public:
-  RegAllocTypeParser(cl::Option &O) : cl::parser<RegAllocType>(O) {}
-  void initialize() {
-    cl::parser<RegAllocType>::initialize();
-    addLiteralOption("default", RegAllocType::Default,
-                     "Default register allocator");
-    addLiteralOption("pbqp", RegAllocType::PBQP, "PBQP register allocator");
-    addLiteralOption("fast", RegAllocType::Fast, "Fast register allocator");
-    addLiteralOption("basic", RegAllocType::Basic, "Basic register allocator");
-    addLiteralOption("greedy", RegAllocType::Greedy,
-                     "Greedy register allocator");
-  }
-};
 
 // Not one-on-one but mostly corresponding to commandline options in
 // TargetPassConfig.cpp.
@@ -70,7 +54,7 @@ struct CGPassBuilderOption {
   bool RequiresCodeGenSCCOrder = false;
 
   RunOutliner EnableMachineOutliner = RunOutliner::TargetDefault;
-  RegAllocType RegAlloc = RegAllocType::Unset;
+  mutable StringRef RegAllocPipeline;
   std::optional<GlobalISelAbortMode> EnableGlobalISelAbort;
   std::string FSProfileFile;
   std::string FSRemappingFile;

llvm/include/llvm/Target/TargetMachine.h

@@ -472,7 +472,8 @@ class TargetMachine {
   virtual Error buildCodeGenPipeline(ModulePassManager &, raw_pwrite_stream &,
                                      raw_pwrite_stream *, CodeGenFileType,
                                      const CGPassBuilderOption &,
-                                     PassInstrumentationCallbacks *) {
+                                     PassInstrumentationCallbacks *,
+                                     PassBuilder &) {
     return make_error<StringError>("buildCodeGenPipeline is not overridden",
                                    inconvertibleErrorCode());
   }