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

Akshat Oke via llvm-commits llvm-commits at lists.llvm.org
Mon Apr 21 23:18:36 PDT 2025


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

>From a0c222362f0982ed9c4c2a9873b62f59fcd271f9 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 1/3] [CodeGen][NPM] Support generic regalloc-npm option

---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h |  84 +++++++++++---
 .../llvm/Passes/MachinePassRegistry.def       |  16 +--
 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   |  23 +++-
 llvm/tools/llc/NewPMDriver.cpp                |   8 +-
 18 files changed, 280 insertions(+), 121 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index d092049022623..808962f523966 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -99,6 +99,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"
@@ -116,6 +117,7 @@
 #include "llvm/Transforms/Utils/EntryExitInstrumenter.h"
 #include "llvm/Transforms/Utils/LowerInvoke.h"
 #include <cassert>
+#include <tuple>
 #include <type_traits>
 #include <utility>
 
@@ -157,8 +159,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)
 
@@ -294,6 +297,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(); }
@@ -501,6 +505,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.
@@ -597,6 +608,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);
 }
 
@@ -1094,6 +1110,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.
@@ -1104,22 +1163,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 fff5083caad1a..9073a60fe12e3 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -215,8 +215,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")
@@ -226,8 +226,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")
@@ -242,8 +242,8 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
 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")
 
@@ -251,8 +251,8 @@ 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 MACHINE_FUNCTION_PASS_WITH_PARAMS
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 5a16c9cafcd7a..ed6f907efffc0 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 a5c624b96304b..0a010f7247950 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1390,39 +1390,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;
@@ -1471,26 +1438,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-");
@@ -1503,8 +1461,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) {
@@ -2223,7 +2225,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()));                                   \
@@ -2489,7 +2492,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 90e3489ced923..730f3f181ff1b 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -773,7 +773,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=");
@@ -787,8 +787,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;
@@ -1076,8 +1076,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);
 }
 
@@ -1963,8 +1963,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
@@ -2208,6 +2208,44 @@ void AMDGPUCodeGenPassBuilder::addPreEmitPass(AddMachinePass &addPass) const {
 
   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..901712efc8e54 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-reg"}));
+    // 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..eeed9f35b7d92 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,23 @@
 # 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: regallocfast
+# 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 --check-prefix=CHECK-FAST
+
+# 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: greedy<tile-reg>
+# CHECK: greedy<all>
+
+# CHECK-FAST: 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.

>From adabf4e3f71992feb62f9b6bc7f0e10fa39b6a42 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Tue, 22 Apr 2025 05:48:30 +0000
Subject: [PATCH 2/3] AS, allow only RA passes

---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h  | 11 +++++------
 .../llvm/Passes/MachinePassRegistry.def        | 18 ++++++++++++++----
 llvm/include/llvm/Target/CGPassBuilderOption.h | 16 ----------------
 .../llc/new-pm/x86_64-regalloc-pipeline.mir    | 11 ++++++++---
 llvm/tools/llc/NewPMDriver.cpp                 | 13 +++++++++----
 5 files changed, 36 insertions(+), 33 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index 808962f523966..1d18262e62dee 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -610,7 +610,7 @@ Error CodeGenPassBuilder<Derived, TargetMachineT>::buildPipeline(
 
   if (!Opt.RegAllocPipeline.empty())
     return make_error<StringError>(
-        "Extra passes in regalloc pipeline: " + Opt.RegAllocPipeline,
+        "extra passes in regalloc pipeline: " + Opt.RegAllocPipeline,
         std::make_error_code(std::errc::invalid_argument));
 
   return verifyStartStop(*StartStopInfo);
@@ -1125,8 +1125,7 @@ bool CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPassFromOpt(
     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)                              \
+#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)); \
@@ -1136,9 +1135,9 @@ bool CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPassFromOpt(
     }                                                                          \
     if (!MatchPassTo.empty()) {                                                \
       if (MatchPassTo != CLASS)                                                \
-        report_fatal_error("Expected " +                                       \
+        report_fatal_error("expected " +                                       \
                                PIC->getPassNameForClassName(MatchPassTo) +     \
-                               " in option -regalloc-npm",                     \
+                               " in option --regalloc-npm",                    \
                            false);                                             \
     }                                                                          \
     addPass(CREATE_PASS(Params.get()));                                        \
@@ -1146,7 +1145,7 @@ bool CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPassFromOpt(
   }
 #include "llvm/Passes/MachinePassRegistry.def"
     if (PassOpt != "default") {
-      report_fatal_error("Unknown register allocator pass: " + PassOpt, false);
+      report_fatal_error("unknown register allocator pass: " + PassOpt, false);
     }
   }
   // If user did not give a specific pass, use the default provided.
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index 9073a60fe12e3..4e9062dd96f12 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -239,7 +239,15 @@ MACHINE_FUNCTION_PASS_WITH_PARAMS(
     },
     parseMachineSinkingPassOptions, "enable-sink-fold")
 
-MACHINE_FUNCTION_PASS_WITH_PARAMS(
+#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); },
     [](StringRef Params, const PassBuilder &PB) {
@@ -248,13 +256,15 @@ MACHINE_FUNCTION_PASS_WITH_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); },
     [](StringRef Params, const PassBuilder &PB) {
       return parseRegAllocGreedyFilterFunc(PB, Params);
-    }, "reg-filter"
-)
+    },
+    "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/Target/CGPassBuilderOption.h b/llvm/include/llvm/Target/CGPassBuilderOption.h
index 96829fbd5f445..8d60decacb018 100644
--- a/llvm/include/llvm/Target/CGPassBuilderOption.h
+++ b/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.
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 eeed9f35b7d92..356224d1187d3 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
@@ -11,7 +11,10 @@
 # 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
+# RUN: not llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -O3 -regalloc-npm='machine-cp' -print-pipeline-passes %s -o - 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID
+# RUN: not llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -passes='machine-sink' -regalloc-npm='greedy' -print-pipeline-passes %s -o - 2>&1 | FileCheck %s --check-prefix=CHECK-BOTH-INVALID
+
+# RUN: not llc -mtriple=x86_64-unknown-linux-gnu -enable-new-pm -O0 -regalloc-npm='greedy' -print-pipeline-passes %s -o - 2>&1 | FileCheck %s --check-prefix=CHECK-ONLY-FAST
 
 # CHECK: greedy<tile-reg>
 # CHECK: greedy<all>
@@ -19,5 +22,7 @@
 # CHECK-FAST: 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
+# CHECK-ERROR: extra passes in regalloc pipeline: basic
+# CHECK-INVALID: unknown register allocator pass: machine-cp
+# CHECK-BOTH-INVALID: --passes and --regalloc-npm cannot be used together
+# CHECK-ONLY-FAST: expected regallocfast in option --regalloc-npm
diff --git a/llvm/tools/llc/NewPMDriver.cpp b/llvm/tools/llc/NewPMDriver.cpp
index 602c651b862bc..4ae3de13d471a 100644
--- a/llvm/tools/llc/NewPMDriver.cpp
+++ b/llvm/tools/llc/NewPMDriver.cpp
@@ -49,9 +49,9 @@
 using namespace llvm;
 
 static cl::opt<std::string>
-    RegAlloc("regalloc-npm",
-             cl::desc("Register allocator to use for new pass manager"),
-             cl::Hidden);
+    RegAllocPasses("regalloc-npm",
+                   cl::desc("Register allocator to use for new pass manager"),
+                   cl::Hidden);
 
 static cl::opt<bool>
     DebugPM("debug-pass-manager", cl::Hidden,
@@ -98,6 +98,11 @@ int llvm::compileModuleWithNewPM(
         << TargetPassConfig::getLimitedCodeGenPipelineReason() << ".\n";
     return 1;
   }
+  if (!PassPipeline.empty() && !RegAllocPasses.empty()) {
+    WithColor::error(errs(), Arg0)
+        << "--passes and --regalloc-npm cannot be used together.\n";
+    return 1;
+  }
 
   raw_pwrite_stream *OS = &Out->os();
 
@@ -105,7 +110,7 @@ int llvm::compileModuleWithNewPM(
   CGPassBuilderOption Opt = getCGPassBuilderOption();
   Opt.DisableVerify = VK != VerifierKind::InputOutput;
   Opt.DebugPM = DebugPM;
-  Opt.RegAllocPipeline = RegAlloc;
+  Opt.RegAllocPipeline = RegAllocPasses;
 
   MachineModuleInfo MMI(Target.get());
 

>From 72d38673dbd85c2c2ffe1d70fe8d6bbf11d13a73 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Tue, 22 Apr 2025 06:18:19 +0000
Subject: [PATCH 3/3] Change to builder callback

---
 llvm/include/llvm/Passes/CodeGenPassBuilder.h  | 14 ++++++++------
 llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp |  9 ++++++---
 llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp  |  5 +++--
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index 1d18262e62dee..3af156ab7d2cb 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -506,12 +506,14 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
   /// 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 return false if there is no pass
+  /// Add the next pass in the cli option or the pass specified if no pass is
   /// left in the option.
-  template <typename RegAllocPassT>
-  void addRegAllocPassOrOpt(AddMachinePass &, RegAllocPassT Pass) const;
+  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.
@@ -1111,11 +1113,11 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
 }
 
 template <typename Derived, typename TargetMachineT>
-template <typename RegAllocPassT>
+template <typename RegAllocPassBuilderT>
 void CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPassOrOpt(
-    AddMachinePass &addPass, RegAllocPassT Pass) const {
+    AddMachinePass &addPass, RegAllocPassBuilderT PassBuilder) const {
   if (!addRegAllocPassFromOpt(addPass))
-    addPass(std::move(Pass));
+    addPass(std::move(PassBuilder()));
 }
 
 template <typename Derived, typename TargetMachineT>
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index 730f3f181ff1b..07afe5943992e 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -2212,7 +2212,8 @@ Error AMDGPUCodeGenPassBuilder::addRegAssignmentOptimized(
     AddMachinePass &addPass) const {
   addPass(GCNPreRALongBranchRegPass());
 
-  addRegAllocPassOrOpt(addPass, RAGreedyPass({onlyAllocateSGPRs, "sgpr"}));
+  addRegAllocPassOrOpt(
+      addPass, []() { return 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
@@ -2232,13 +2233,15 @@ Error AMDGPUCodeGenPassBuilder::addRegAssignmentOptimized(
   addPass(SIPreAllocateWWMRegsPass());
 
   // For allocating other wwm register operands.
-  addRegAllocPassOrOpt(addPass, RAGreedyPass({onlyAllocateWWMRegs, "wwm"}));
+  addRegAllocPassOrOpt(
+      addPass, []() { return RAGreedyPass({onlyAllocateWWMRegs, "wwm"}); });
   addPass(SILowerWWMCopiesPass());
   addPass(VirtRegRewriterPass(false));
   addPass(AMDGPUReserveWWMRegsPass());
 
   // For allocating per-thread VGPRs.
-  addRegAllocPassOrOpt(addPass, RAGreedyPass({onlyAllocateVGPRs, "vgpr"}));
+  addRegAllocPassOrOpt(
+      addPass, []() { return RAGreedyPass({onlyAllocateVGPRs, "vgpr"}); });
 
   // TODO: addPreRewrite();
   addPass(VirtRegRewriterPass(false));
diff --git a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
index 901712efc8e54..14f7dee75b6df 100644
--- a/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
+++ b/llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp
@@ -39,8 +39,9 @@ class X86CodeGenPassBuilder
 Error X86CodeGenPassBuilder::addRegAssignmentOptimized(
     AddMachinePass &addPass) const {
   if (EnableTileRAPass) {
-    addRegAllocPassOrOpt(addPass,
-                         RAGreedyPass({onlyAllocateTileRegisters, "tile-reg"}));
+    addRegAllocPassOrOpt(addPass, []() {
+      return RAGreedyPass({onlyAllocateTileRegisters, "tile-reg"});
+    });
     // TODO: addPass(X86TileConfigPass());
   }
   return Base::addRegAssignmentOptimized(addPass);



More information about the llvm-commits mailing list