[llvm] [CodeGen][NPM] Support generic regalloc-npm option (PR #135149)

Akshat Oke via llvm-commits llvm-commits at lists.llvm.org
Thu Apr 10 03:11:16 PDT 2025


https://github.com/optimisan updated https://github.com/llvm/llvm-project/pull/135149

>From ce260c8bd4b06750aa1b35fb12967fce342476be Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 10 Apr 2025 08:58:41 +0000
Subject: [PATCH] [CodeGen][NPM] Support generic regalloc-npm option

---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h |  84 +++++++++++---
 .../llvm/Passes/MachinePassRegistry.def       |  22 ++--
 llvm/include/llvm/Passes/PassBuilder.h        |  26 ++++-
 .../llvm/Passes/TargetPassRegistry.inc        |   2 +-
 .../include/llvm/Target/CGPassBuilderOption.h |   2 +-
 llvm/include/llvm/Target/TargetMachine.h      |   3 +-
 llvm/lib/Passes/PassBuilder.cpp               | 105 +++++++++---------
 .../lib/Target/AMDGPU/AMDGPUTargetMachine.cpp |  52 +++++++--
 llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h  |   6 +-
 llvm/lib/Target/AMDGPU/R600TargetMachine.cpp  |   8 +-
 llvm/lib/Target/AMDGPU/R600TargetMachine.h    |   5 +-
 llvm/lib/Target/BPF/BPFTargetMachine.cpp      |   3 +-
 llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp |  29 ++++-
 llvm/lib/Target/X86/X86TargetMachine.cpp      |  15 ++-
 llvm/lib/Target/X86/X86TargetMachine.h        |   9 +-
 .../test/tools/llc/new-pm/regalloc-amdgpu.mir |   5 +
 .../llc/new-pm/x86_64-regalloc-pipeline.mir   |  18 ++-
 llvm/tools/llc/NewPMDriver.cpp                |   8 +-
 18 files changed, 282 insertions(+), 120 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index 25ca982916ff8..8af5d9e51eddc 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -96,6 +96,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"
@@ -113,6 +114,7 @@
 #include "llvm/Transforms/Utils/EntryExitInstrumenter.h"
 #include "llvm/Transforms/Utils/LowerInvoke.h"
 #include <cassert>
+#include <tuple>
 #include <type_traits>
 #include <utility>
 
@@ -154,8 +156,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)
 
@@ -291,6 +294,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(); }
@@ -498,6 +502,13 @@ 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.
+  bool addRegAllocPassFromOpt(AddMachinePass &,
+                              StringRef MatchPassTo = StringRef{}) const;
+  /// Add the next pass in the cli option, or return false if there is no pass
+  /// left in the option.
+  template <typename RegAllocPassT>
+  void addRegAllocPassOrOpt(AddMachinePass &, RegAllocPassT Pass) const;
 
   /// Add core register alloator passes which do the actual register assignment
   /// and rewriting. \returns true if any passes were added.
@@ -594,6 +605,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);
 }
 
@@ -1088,6 +1104,49 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
     addPass(RegAllocFastPass());
 }
 
+template <typename Derived, typename TargetMachineT>
+template <typename RegAllocPassT>
+void CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPassOrOpt(
+    AddMachinePass &addPass, RegAllocPassT Pass) const {
+  if (!addRegAllocPassFromOpt(addPass))
+    addPass(std::move(Pass));
+}
+
+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 MACHINE_FUNCTION_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.
@@ -1098,22 +1157,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>
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index 3e9e788662900..3a2d2c49bbcb8 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -209,8 +209,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")
@@ -220,8 +220,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")
@@ -233,11 +233,16 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
     },
     parseMachineSinkingPassOptions, "enable-sink-fold")
 
+// #ifndef RA_PASS_WITH_PARAMS
+// #define RA_PASS_WITH_PARAMS(NAME, CLASS, CREATE_PASS, PARSER,    \
+//                                           PARAMS)
+// #endif
+
 MACHINE_FUNCTION_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")
 
@@ -245,10 +250,11 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
 MACHINE_FUNCTION_PASS_WITH_PARAMS(
     "greedy", "RAGreedyPass",
     [](RAGreedyPass::Options Opts) { return RAGreedyPass(Opts); },
-    [PB = this](StringRef Params) {
-      return parseRegAllocGreedyFilterFunc(*PB, Params);
+    [](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
diff --git a/llvm/include/llvm/Passes/PassBuilder.h b/llvm/include/llvm/Passes/PassBuilder.h
index 51ccaa53447d7..ac494aef9aab8 100644
--- a/llvm/include/llvm/Passes/PassBuilder.h
+++ b/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,16 @@ 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 &);
 }
 
 #endif
diff --git a/llvm/include/llvm/Passes/TargetPassRegistry.inc b/llvm/include/llvm/Passes/TargetPassRegistry.inc
index 521913cb25a4a..da9fa7242ab80 100644
--- a/llvm/include/llvm/Passes/TargetPassRegistry.inc
+++ b/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;                                                            \
diff --git a/llvm/include/llvm/Target/CGPassBuilderOption.h b/llvm/include/llvm/Target/CGPassBuilderOption.h
index 51f25c1360b87..96829fbd5f445 100644
--- a/llvm/include/llvm/Target/CGPassBuilderOption.h
+++ b/llvm/include/llvm/Target/CGPassBuilderOption.h
@@ -70,7 +70,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;
diff --git a/llvm/include/llvm/Target/TargetMachine.h b/llvm/include/llvm/Target/TargetMachine.h
index 27eeb415ed644..420078ffa4f46 100644
--- a/llvm/include/llvm/Target/TargetMachine.h
+++ b/llvm/include/llvm/Target/TargetMachine.h
@@ -469,7 +469,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());
   }
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 5cda1517e127d..8aa9fee37f87d 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1385,39 +1385,6 @@ Expected<SmallVector<std::string, 0>> parseInternalizeGVs(StringRef Params) {
   return Expected<SmallVector<std::string, 0>>(std::move(PreservedGVs));
 }
 
-Expected<RegAllocFastPass::Options>
-parseRegAllocFastPassOptions(PassBuilder &PB, StringRef Params) {
-  RegAllocFastPass::Options Opts;
-  while (!Params.empty()) {
-    StringRef ParamName;
-    std::tie(ParamName, Params) = Params.split(';');
-
-    if (ParamName.consume_front("filter=")) {
-      std::optional<RegAllocFilterFunc> Filter =
-          PB.parseRegAllocFilter(ParamName);
-      if (!Filter) {
-        return make_error<StringError>(
-            formatv("invalid regallocfast register filter '{0}' ", ParamName)
-                .str(),
-            inconvertibleErrorCode());
-      }
-      Opts.Filter = *Filter;
-      Opts.FilterName = ParamName;
-      continue;
-    }
-
-    if (ParamName == "no-clear-vregs") {
-      Opts.ClearVRegs = false;
-      continue;
-    }
-
-    return make_error<StringError>(
-        formatv("invalid regallocfast pass parameter '{0}' ", ParamName).str(),
-        inconvertibleErrorCode());
-  }
-  return Opts;
-}
-
 Expected<BoundsCheckingPass::Options>
 parseBoundsCheckingOptions(StringRef Params) {
   BoundsCheckingPass::Options Options;
@@ -1466,26 +1433,17 @@ parseBoundsCheckingOptions(StringRef Params) {
   return Options;
 }
 
-Expected<RAGreedyPass::Options>
-parseRegAllocGreedyFilterFunc(PassBuilder &PB, StringRef Params) {
-  if (Params.empty() || Params == "all")
-    return RAGreedyPass::Options();
-
-  std::optional<RegAllocFilterFunc> Filter = PB.parseRegAllocFilter(Params);
-  if (Filter)
-    return RAGreedyPass::Options{*Filter, Params};
-
-  return make_error<StringError>(
-      formatv("invalid regallocgreedy register filter '{0}' ", Params).str(),
-      inconvertibleErrorCode());
-}
+} // namespace
 
-Expected<bool> parseMachineSinkingPassOptions(StringRef Params) {
+Expected<bool> llvm::parseMachineSinkingPassOptions(StringRef Params,
+                                                    const PassBuilder &) {
   return PassBuilder::parseSinglePassOption(Params, "enable-sink-fold",
                                             "MachineSinkingPass");
 }
 
-Expected<bool> parseMachineBlockPlacementPassOptions(StringRef Params) {
+Expected<bool>
+llvm::parseMachineBlockPlacementPassOptions(StringRef Params,
+                                            const PassBuilder &) {
   bool AllowTailMerge = true;
   if (!Params.empty()) {
     AllowTailMerge = !Params.consume_front("no-");
@@ -1498,8 +1456,52 @@ Expected<bool> parseMachineBlockPlacementPassOptions(StringRef Params) {
   return AllowTailMerge;
 }
 
-} // namespace
+Expected<RegAllocFastPass::Options>
+llvm::parseRegAllocFastPassOptions(const PassBuilder &PB, StringRef Params) {
+  RegAllocFastPass::Options Opts;
+  while (!Params.empty()) {
+    StringRef ParamName;
+    std::tie(ParamName, Params) = Params.split(';');
 
+    if (ParamName.consume_front("filter=")) {
+      std::optional<RegAllocFilterFunc> Filter =
+          PB.parseRegAllocFilter(ParamName);
+      if (!Filter) {
+        return make_error<StringError>(
+            formatv("invalid regallocfast register filter '{0}' ", ParamName)
+                .str(),
+            inconvertibleErrorCode());
+      }
+      Opts.Filter = *Filter;
+      Opts.FilterName = ParamName;
+      continue;
+    }
+
+    if (ParamName == "no-clear-vregs") {
+      Opts.ClearVRegs = false;
+      continue;
+    }
+
+    return make_error<StringError>(
+        formatv("invalid regallocfast pass parameter '{0}' ", ParamName).str(),
+        inconvertibleErrorCode());
+  }
+  return Opts;
+}
+
+Expected<RAGreedyPass::Options>
+llvm::parseRegAllocGreedyFilterFunc(const PassBuilder &PB, StringRef Params) {
+  if (Params.empty() || Params == "all")
+    return RAGreedyPass::Options();
+
+  std::optional<RegAllocFilterFunc> Filter = PB.parseRegAllocFilter(Params);
+  if (Filter)
+    return RAGreedyPass::Options{*Filter, Params};
+
+  return make_error<StringError>(
+      formatv("invalid regallocgreedy register filter '{0}' ", Params).str(),
+      inconvertibleErrorCode());
+}
 /// Tests whether a pass name starts with a valid prefix for a default pipeline
 /// alias.
 static bool startsWithDefaultPipelineAliasPrefix(StringRef Name) {
@@ -2218,7 +2220,8 @@ Error PassBuilder::parseMachinePass(MachineFunctionPassManager &MFPM,
 #define MACHINE_FUNCTION_PASS_WITH_PARAMS(NAME, CLASS, CREATE_PASS, PARSER,    \
                                           PARAMS)                              \
   if (checkParametrizedPassName(Name, NAME)) {                                 \
-    auto Params = parsePassParameters(PARSER, Name, NAME);                     \
+    auto Params = parsePassParameters(PARSER, Name, NAME,                      \
+                                      const_cast<const PassBuilder &>(*this)); \
     if (!Params)                                                               \
       return Params.takeError();                                               \
     MFPM.addPass(CREATE_PASS(Params.get()));                                   \
@@ -2484,7 +2487,7 @@ Error PassBuilder::parseAAPipeline(AAManager &AA, StringRef PipelineText) {
 }
 
 std::optional<RegAllocFilterFunc>
-PassBuilder::parseRegAllocFilter(StringRef FilterName) {
+PassBuilder::parseRegAllocFilter(StringRef FilterName) const {
   if (FilterName == "all")
     return nullptr;
   for (auto &C : RegClassFilterParsingCallbacks)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index c2bcd53644371..b904751a95c96 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -767,7 +767,7 @@ void AMDGPUTargetMachine::registerDefaultAliasAnalyses(AAManager &AAM) {
 }
 
 static Expected<ScanOptions>
-parseAMDGPUAtomicOptimizerStrategy(StringRef Params) {
+parseAMDGPUAtomicOptimizerStrategy(StringRef Params, const PassBuilder &) {
   if (Params.empty())
     return ScanOptions::Iterative;
   Params.consume_front("strategy=");
@@ -781,8 +781,8 @@ parseAMDGPUAtomicOptimizerStrategy(StringRef Params) {
   return make_error<StringError>("invalid parameter", inconvertibleErrorCode());
 }
 
-Expected<AMDGPUAttributorOptions>
-parseAMDGPUAttributorPassOptions(StringRef Params) {
+static Expected<AMDGPUAttributorOptions>
+parseAMDGPUAttributorPassOptions(StringRef Params, const PassBuilder &PB) {
   AMDGPUAttributorOptions Result;
   while (!Params.empty()) {
     StringRef ParamName;
@@ -1064,8 +1064,8 @@ GCNTargetMachine::getTargetTransformInfo(const Function &F) const {
 Error GCNTargetMachine::buildCodeGenPipeline(
     ModulePassManager &MPM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
     CodeGenFileType FileType, const CGPassBuilderOption &Opts,
-    PassInstrumentationCallbacks *PIC) {
-  AMDGPUCodeGenPassBuilder CGPB(*this, Opts, PIC);
+    PassInstrumentationCallbacks *PIC, PassBuilder &PB) {
+  AMDGPUCodeGenPassBuilder CGPB(*this, Opts, PIC, PB);
   return CGPB.buildPipeline(MPM, Out, DwoOut, FileType);
 }
 
@@ -1948,8 +1948,8 @@ bool GCNTargetMachine::parseMachineFunctionInfo(
 
 AMDGPUCodeGenPassBuilder::AMDGPUCodeGenPassBuilder(
     GCNTargetMachine &TM, const CGPassBuilderOption &Opts,
-    PassInstrumentationCallbacks *PIC)
-    : CodeGenPassBuilder(TM, Opts, PIC) {
+    PassInstrumentationCallbacks *PIC, PassBuilder &PB)
+    : CodeGenPassBuilder(TM, Opts, PIC, PB) {
   Opt.MISchedPostRA = true;
   Opt.RequiresCodeGenSCCOrder = true;
   // Exceptions and StackMaps are not supported, so these passes will never do
@@ -2193,6 +2193,44 @@ void AMDGPUCodeGenPassBuilder::addPreEmitPass(AddMachinePass &addPass) const {
 
   // TODO: addPass(BranchRelaxationPass());
 }
+Error AMDGPUCodeGenPassBuilder::addRegAssignmentOptimized(
+    AddMachinePass &addPass) const {
+  addPass(GCNPreRALongBranchRegPass());
+
+  addRegAllocPassOrOpt(addPass, RAGreedyPass({onlyAllocateSGPRs, "sgpr"}));
+
+  // Commit allocated register changes. This is mostly necessary because too
+  // many things rely on the use lists of the physical registers, such as the
+  // verifier. This is only necessary with allocators which use LiveIntervals,
+  // since FastRegAlloc does the replacements itself.
+  // TODO: addPass(VirtRegRewriterPass(false));
+
+  // At this point, the sgpr-regalloc has been done and it is good to have the
+  // stack slot coloring to try to optimize the SGPR spill stack indices before
+  // attempting the custom SGPR spill lowering.
+  addPass(StackSlotColoringPass());
+
+  // Equivalent of PEI for SGPRs.
+  addPass(SILowerSGPRSpillsPass());
+
+  // To Allocate wwm registers used in whole quad mode operations (for shaders).
+  addPass(SIPreAllocateWWMRegsPass());
+
+  // For allocating other wwm register operands.
+  addRegAllocPassOrOpt(addPass, RAGreedyPass({onlyAllocateWWMRegs, "wwm"}));
+  addPass(SILowerWWMCopiesPass());
+  addPass(VirtRegRewriterPass(false));
+  addPass(AMDGPUReserveWWMRegsPass());
+
+  // For allocating per-thread VGPRs.
+  addRegAllocPassOrOpt(addPass, RAGreedyPass({onlyAllocateVGPRs, "vgpr"}));
+
+  // TODO: addPreRewrite();
+  addPass(VirtRegRewriterPass(false));
+
+  // TODO: addPass(AMDGPUMarkLastScratchLoadPass());
+  return Error::success();
+}
 
 bool AMDGPUCodeGenPassBuilder::isPassEnabled(const cl::opt<bool> &Opt,
                                              CodeGenOptLevel Level) const {
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h
index 3df4115324ac2..5a0d6b8d04beb 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h
@@ -102,7 +102,8 @@ class GCNTargetMachine final : public AMDGPUTargetMachine {
                              raw_pwrite_stream *DwoOut,
                              CodeGenFileType FileType,
                              const CGPassBuilderOption &Opts,
-                             PassInstrumentationCallbacks *PIC) override;
+                             PassInstrumentationCallbacks *PIC,
+                             PassBuilder &) override;
 
   void registerMachineRegisterInfoCallback(MachineFunction &MF) const override;
 
@@ -169,7 +170,7 @@ class AMDGPUCodeGenPassBuilder
 public:
   AMDGPUCodeGenPassBuilder(GCNTargetMachine &TM,
                            const CGPassBuilderOption &Opts,
-                           PassInstrumentationCallbacks *PIC);
+                           PassInstrumentationCallbacks *PIC, PassBuilder &PB);
 
   void addIRPasses(AddIRPass &) const;
   void addCodeGenPrepare(AddIRPass &) const;
@@ -181,6 +182,7 @@ class AMDGPUCodeGenPassBuilder
   void addMachineSSAOptimization(AddMachinePass &) const;
   void addPostRegAlloc(AddMachinePass &) const;
   void addPreEmitPass(AddMachinePass &) const;
+  Error addRegAssignmentOptimized(AddMachinePass &) const;
 
   /// Check if a pass is enabled given \p Opt option. The option always
   /// overrides defaults if explicitly used. Otherwise its default will be used
diff --git a/llvm/lib/Target/AMDGPU/R600TargetMachine.cpp b/llvm/lib/Target/AMDGPU/R600TargetMachine.cpp
index 10552a1f0b1bc..47fb08a3bb0f9 100644
--- a/llvm/lib/Target/AMDGPU/R600TargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/R600TargetMachine.cpp
@@ -150,8 +150,8 @@ TargetPassConfig *R600TargetMachine::createPassConfig(PassManagerBase &PM) {
 Error R600TargetMachine::buildCodeGenPipeline(
     ModulePassManager &MPM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
     CodeGenFileType FileType, const CGPassBuilderOption &Opts,
-    PassInstrumentationCallbacks *PIC) {
-  R600CodeGenPassBuilder CGPB(*this, Opts, PIC);
+    PassInstrumentationCallbacks *PIC, PassBuilder &PB) {
+  R600CodeGenPassBuilder CGPB(*this, Opts, PIC, PB);
   return CGPB.buildPipeline(MPM, Out, DwoOut, FileType);
 }
 
@@ -168,8 +168,8 @@ MachineFunctionInfo *R600TargetMachine::createMachineFunctionInfo(
 
 R600CodeGenPassBuilder::R600CodeGenPassBuilder(
     R600TargetMachine &TM, const CGPassBuilderOption &Opts,
-    PassInstrumentationCallbacks *PIC)
-    : CodeGenPassBuilder(TM, Opts, PIC) {
+    PassInstrumentationCallbacks *PIC, PassBuilder &PB)
+    : CodeGenPassBuilder(TM, Opts, PIC, PB) {
   Opt.RequiresCodeGenSCCOrder = true;
 }
 
diff --git a/llvm/lib/Target/AMDGPU/R600TargetMachine.h b/llvm/lib/Target/AMDGPU/R600TargetMachine.h
index eb4cb91cb704d..421b48a200f6e 100644
--- a/llvm/lib/Target/AMDGPU/R600TargetMachine.h
+++ b/llvm/lib/Target/AMDGPU/R600TargetMachine.h
@@ -42,7 +42,8 @@ class R600TargetMachine final : public AMDGPUTargetMachine {
                              raw_pwrite_stream *DwoOut,
                              CodeGenFileType FileType,
                              const CGPassBuilderOption &Opt,
-                             PassInstrumentationCallbacks *PIC) override;
+                             PassInstrumentationCallbacks *PIC,
+                             PassBuilder &) override;
 
   const TargetSubtargetInfo *getSubtargetImpl(const Function &) const override;
 
@@ -65,7 +66,7 @@ class R600CodeGenPassBuilder
     : public CodeGenPassBuilder<R600CodeGenPassBuilder, R600TargetMachine> {
 public:
   R600CodeGenPassBuilder(R600TargetMachine &TM, const CGPassBuilderOption &Opts,
-                         PassInstrumentationCallbacks *PIC);
+                         PassInstrumentationCallbacks *PIC, PassBuilder &PB);
 
   void addPreISel(AddIRPass &addPass) const;
   void addAsmPrinter(AddMachinePass &, CreateMCStreamer) const;
diff --git a/llvm/lib/Target/BPF/BPFTargetMachine.cpp b/llvm/lib/Target/BPF/BPFTargetMachine.cpp
index 4c4e6e27b9a5e..23a698642c8b4 100644
--- a/llvm/lib/Target/BPF/BPFTargetMachine.cpp
+++ b/llvm/lib/Target/BPF/BPFTargetMachine.cpp
@@ -110,7 +110,8 @@ TargetPassConfig *BPFTargetMachine::createPassConfig(PassManagerBase &PM) {
   return new BPFPassConfig(*this, PM);
 }
 
-static Expected<bool> parseBPFPreserveStaticOffsetOptions(StringRef Params) {
+static Expected<bool> parseBPFPreserveStaticOffsetOptions(StringRef Params,
+                                                          const PassBuilder &) {
   return PassBuilder::parseSinglePassOption(Params, "allow-partial",
                                             "BPFPreserveStaticOffsetPass");
 }
diff --git a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
index d979517e12af6..069bcdd5b01c6 100644
--- a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
+++ b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
@@ -26,13 +26,26 @@ class X86CodeGenPassBuilder
 public:
   explicit X86CodeGenPassBuilder(X86TargetMachine &TM,
                                  const CGPassBuilderOption &Opts,
-                                 PassInstrumentationCallbacks *PIC)
-      : CodeGenPassBuilder(TM, Opts, PIC) {}
+                                 PassInstrumentationCallbacks *PIC,
+                                 PassBuilder &PB)
+      : CodeGenPassBuilder(TM, Opts, PIC, PB) {}
+  using Base = CodeGenPassBuilder<X86CodeGenPassBuilder, X86TargetMachine>;
   void addPreISel(AddIRPass &addPass) const;
   void addAsmPrinter(AddMachinePass &, CreateMCStreamer) const;
   Error addInstSelector(AddMachinePass &) const;
+  Error addRegAssignmentOptimized(AddMachinePass &addPass) const;
 };
 
+Error X86CodeGenPassBuilder::addRegAssignmentOptimized(
+    AddMachinePass &addPass) const {
+  if (EnableTileRAPass) {
+    addRegAllocPassOrOpt(
+        addPass, RAGreedyPass({onlyAllocateTileRegisters, "tile-regs"}));
+    // TODO: addPass(X86TileConfigPass());
+  }
+  return Base::addRegAssignmentOptimized(addPass);
+}
+
 void X86CodeGenPassBuilder::addPreISel(AddIRPass &addPass) const {
   // TODO: Add passes pre instruction selection.
 }
@@ -53,12 +66,20 @@ Error X86CodeGenPassBuilder::addInstSelector(AddMachinePass &addPass) const {
 void X86TargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
 #define GET_PASS_REGISTRY "X86PassRegistry.def"
 #include "llvm/Passes/TargetPassRegistry.inc"
+
+  PB.registerRegClassFilterParsingCallback(
+      [](StringRef FilterName) -> RegAllocFilterFunc {
+        if (FilterName == "tile-reg") {
+          return onlyAllocateTileRegisters;
+        }
+        return nullptr;
+      });
 }
 
 Error X86TargetMachine::buildCodeGenPipeline(
     ModulePassManager &MPM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
     CodeGenFileType FileType, const CGPassBuilderOption &Opt,
-    PassInstrumentationCallbacks *PIC) {
-  auto CGPB = X86CodeGenPassBuilder(*this, Opt, PIC);
+    PassInstrumentationCallbacks *PIC, PassBuilder &PB) {
+  auto CGPB = X86CodeGenPassBuilder(*this, Opt, PIC, PB);
   return CGPB.buildPipeline(MPM, Out, DwoOut, FileType);
 }
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index 4cecbbf27aa30..63abf46ba793f 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -57,11 +57,10 @@ using namespace llvm;
 static cl::opt<bool> EnableMachineCombinerPass("x86-machine-combiner",
                                cl::desc("Enable the machine combiner pass"),
                                cl::init(true), cl::Hidden);
-
-static cl::opt<bool>
-    EnableTileRAPass("x86-tile-ra",
-                     cl::desc("Enable the tile register allocation pass"),
-                     cl::init(true), cl::Hidden);
+cl::opt<bool>
+    llvm::EnableTileRAPass("x86-tile-ra",
+                           cl::desc("Enable the tile register allocation pass"),
+                           cl::init(true), cl::Hidden);
 
 extern "C" LLVM_C_ABI void LLVMInitializeX86Target() {
   // Register the target.
@@ -677,9 +676,9 @@ std::unique_ptr<CSEConfigBase> X86PassConfig::getCSEConfig() const {
   return getStandardCSEConfigForOpt(TM->getOptLevel());
 }
 
-static bool onlyAllocateTileRegisters(const TargetRegisterInfo &TRI,
-                                      const MachineRegisterInfo &MRI,
-                                      const Register Reg) {
+bool llvm::onlyAllocateTileRegisters(const TargetRegisterInfo &TRI,
+                                     const MachineRegisterInfo &MRI,
+                                     const Register Reg) {
   const TargetRegisterClass *RC = MRI.getRegClass(Reg);
   return static_cast<const X86RegisterInfo &>(TRI).isTileRegisterClass(RC);
 }
diff --git a/llvm/lib/Target/X86/X86TargetMachine.h b/llvm/lib/Target/X86/X86TargetMachine.h
index ced0a9c71fdd8..71653c85a3630 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.h
+++ b/llvm/lib/Target/X86/X86TargetMachine.h
@@ -22,9 +22,15 @@
 
 namespace llvm {
 
+extern cl::opt<bool> EnableTileRAPass;
+
 class StringRef;
 class TargetTransformInfo;
 
+bool onlyAllocateTileRegisters(const TargetRegisterInfo &TRI,
+                               const MachineRegisterInfo &MRI,
+                               const Register Reg);
+
 class X86TargetMachine final : public CodeGenTargetMachineImpl {
   std::unique_ptr<TargetLoweringObjectFile> TLOF;
   mutable StringMap<std::unique_ptr<X86Subtarget>> SubtargetMap;
@@ -74,7 +80,8 @@ class X86TargetMachine final : public CodeGenTargetMachineImpl {
   Error buildCodeGenPipeline(ModulePassManager &, raw_pwrite_stream &,
                              raw_pwrite_stream *, CodeGenFileType,
                              const CGPassBuilderOption &,
-                             PassInstrumentationCallbacks *) override;
+                             PassInstrumentationCallbacks *,
+                             PassBuilder &) override;
 
   bool isJIT() const { return IsJIT; }
 
diff --git a/llvm/test/tools/llc/new-pm/regalloc-amdgpu.mir b/llvm/test/tools/llc/new-pm/regalloc-amdgpu.mir
index 07f2d350ffd9c..7c8b9a9e04dc1 100644
--- a/llvm/test/tools/llc/new-pm/regalloc-amdgpu.mir
+++ b/llvm/test/tools/llc/new-pm/regalloc-amdgpu.mir
@@ -2,11 +2,16 @@
 # RUN: llc -mtriple=amdgcn --passes='regallocfast<filter=sgpr>,regallocfast<filter=wwm>,regallocfast<filter=vgpr>' --print-pipeline-passes --filetype=null %s | FileCheck %s --check-prefix=PASS
 # RUN: not llc -mtriple=amdgcn --passes='regallocfast<filter=bad-filter>' --print-pipeline-passes --filetype=null %s 2>&1 | FileCheck %s --check-prefix=BAD-FILTER
 
+# RUN: llc -mtriple=amdgcn -enable-new-pm --regalloc-npm="regallocfast<filter=sgpr;no-clear-vregs>,greedy<wwm>" --print-pipeline-passes --filetype=null %s | FileCheck %s --check-prefix=RA-NPM
+
 # PASS: regallocfast<filter=sgpr>
 # PASS: regallocfast<filter=wwm>
 # PASS: regallocfast<filter=vgpr>
 # BAD-FILTER: invalid regallocfast register filter 'bad-filter'
 
+# RA-NPM: regallocfast<filter=sgpr;no-clear-vregs>
+# RA-NPM: greedy<wwm>
+# RA-NPM: greedy<vgpr>
 ---
 name: f
 ...
diff --git a/llvm/test/tools/llc/new-pm/x86_64-regalloc-pipeline.mir b/llvm/test/tools/llc/new-pm/x86_64-regalloc-pipeline.mir
index f1a30bc7964ca..a7237aa82e966 100644
--- a/llvm/test/tools/llc/new-pm/x86_64-regalloc-pipeline.mir
+++ b/llvm/test/tools/llc/new-pm/x86_64-regalloc-pipeline.mir
@@ -1,6 +1,20 @@
 # REQUIRES x86_64-registered-target
-# RUN: llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -O3 -regalloc-npm=fast -print-pipeline-passes %s -o - 2>&1 | FileCheck %s
-# RUN: llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -O3 -regalloc-npm=greedy -print-pipeline-passes %s -o - 2>&1 | FileCheck %s --check-prefix=CHECK-GREEDY
+# Test default pipeline
+# RUN: llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -O3 -print-pipeline-passes %s -o - 2>&1 | FileCheck %s --check-prefix=CHECK-GREEDY
+
+# Test the -regalloc-npm option
+# RUN: llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -O3 -regalloc-npm=regallocfast -print-pipeline-passes %s -o - 2>&1 | FileCheck %s
+
+# RUN: llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -O3 -regalloc-npm='greedy<tile-reg>' -print-pipeline-passes %s -o - 2>&1 | FileCheck %s --check-prefix=CHECK-GREEDY-FILTER
+
+# This test attempts to bypass the tile register filter pass, inserts the default greedy<all> pass and then 
+# attempts to insert an extraneous pass.
+# RUN: not llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -O3 -regalloc-npm='default,default,basic' 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR
+
+# RUN: not llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -O3 -regalloc-npm='invalidpass' -print-pipeline-passes %s -o - 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID
 
 # CHECK: regallocfast
 # CHECK-GREEDY: greedy<all>
+# CHECK-GREEDY-FILTER: greedy<tile-reg>
+# CHECK-ERROR: Extra passes in regalloc pipeline: basic
+# CHECK-INVALID: Unknown register allocator pass: invalidpass
diff --git a/llvm/tools/llc/NewPMDriver.cpp b/llvm/tools/llc/NewPMDriver.cpp
index fa82689ecf9ae..602c651b862bc 100644
--- a/llvm/tools/llc/NewPMDriver.cpp
+++ b/llvm/tools/llc/NewPMDriver.cpp
@@ -48,10 +48,10 @@
 
 using namespace llvm;
 
-static cl::opt<RegAllocType, false, RegAllocTypeParser>
+static cl::opt<std::string>
     RegAlloc("regalloc-npm",
              cl::desc("Register allocator to use for new pass manager"),
-             cl::Hidden, cl::init(RegAllocType::Unset));
+             cl::Hidden);
 
 static cl::opt<bool>
     DebugPM("debug-pass-manager", cl::Hidden,
@@ -105,7 +105,7 @@ int llvm::compileModuleWithNewPM(
   CGPassBuilderOption Opt = getCGPassBuilderOption();
   Opt.DisableVerify = VK != VerifierKind::InputOutput;
   Opt.DebugPM = DebugPM;
-  Opt.RegAlloc = RegAlloc;
+  Opt.RegAllocPipeline = RegAlloc;
 
   MachineModuleInfo MMI(Target.get());
 
@@ -155,7 +155,7 @@ int llvm::compileModuleWithNewPM(
 
   } else {
     ExitOnErr(Target->buildCodeGenPipeline(
-        MPM, *OS, DwoOut ? &DwoOut->os() : nullptr, FileType, Opt, &PIC));
+        MPM, *OS, DwoOut ? &DwoOut->os() : nullptr, FileType, Opt, &PIC, PB));
   }
 
   // If user only wants to print the pipeline, print it before parsing the MIR.



More information about the llvm-commits mailing list