[CodeGen][NewPM] Handle --regalloc-npm option

The final register allocation pass is now determined by `PassBuilder`.
However, register allocation passes are now wrapped by pass manager, start/stop options will no longer works for these passes, becuase I'm not sure will LLVM allow plugins to replace register allocator.

Author: paperchalice

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -59,10 +59,12 @@
 #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"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Target/CGPassBuilderOption.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/CFGuard.h"
@@ -116,9 +118,8 @@ namespace llvm {
 template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
 public:
   explicit CodeGenPassBuilder(TargetMachineT &TM,
-                              const CGPassBuilderOption &Opts,
-                              PassInstrumentationCallbacks *PIC)
-      : TM(TM), Opt(Opts), PIC(PIC) {
+                              const CGPassBuilderOption &Opts, PassBuilder &PB)
+      : TM(TM), Opt(Opts), PB(PB), PIC(PB.getPassInstrumentationCallbacks()) {
     // Target could set CGPassBuilderOption::MISchedPostRA to true to achieve
     //     substitutePass(&PostRASchedulerID, &PostMachineSchedulerID)
 
@@ -253,6 +254,7 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
 
   TargetMachineT &TM;
   CGPassBuilderOption Opt;
+  PassBuilder &PB;
   PassInstrumentationCallbacks *PIC;
 
   template <typename TMC> TMC &getTM() const { return static_cast<TMC &>(TM); }
@@ -453,13 +455,9 @@ template <typename DerivedT, typename TargetMachineT> class CodeGenPassBuilder {
   /// Utilities for targets to add passes to the pass manager.
   ///
 
-  /// createTargetRegisterAllocator - Create the register allocator pass for
-  /// this target at the current optimization level.
-  void addTargetRegisterAllocator(AddMachinePass &, bool Optimized) const;
-
   /// addMachinePasses helper to create the target-selected or overriden
   /// regalloc pass.
-  void addRegAllocPass(AddMachinePass &, bool Optimized) const;
+  Error addRegAllocPass(AddMachinePass &, StringRef FilterName = "all") const;
 
   /// Add core register alloator passes which do the actual register assignment
   /// and rewriting. \returns true if any passes were added.
@@ -521,6 +519,9 @@ Error CodeGenPassBuilder<Derived, TargetMachineT>::buildPipeline(
     return StartStopInfo.takeError();
   setStartStopPasses(*StartStopInfo);
 
+  if (auto Err = PB.parseRegAllocOpt(Opt.RegAlloc))
+    return Err;
+
   bool PrintAsm = TargetPassConfig::willCompleteCodeGenPipeline();
   bool PrintMIR = !PrintAsm && FileType != CodeGenFileType::Null;
 
@@ -1025,45 +1026,40 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addMachineSSAOptimization(
 /// Register Allocation Pass Configuration
 //===---------------------------------------------------------------------===//
 
-/// Instantiate the default register allocator pass for this target for either
-/// the optimized or unoptimized allocation path. This will be added to the pass
-/// manager by addFastRegAlloc in the unoptimized case or addOptimizedRegAlloc
-/// in the optimized case.
-///
-/// A target that uses the standard regalloc pass order for fast or optimized
-/// allocation may still override this for per-target regalloc
-/// selection. But -regalloc=... always takes precedence.
-template <typename Derived, typename TargetMachineT>
-void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
-    AddMachinePass &addPass, bool Optimized) const {
-  if (Optimized)
-    addPass(RAGreedyPass());
-  else
-    addPass(RegAllocFastPass());
-}
-
 /// 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.
 template <typename Derived, typename TargetMachineT>
-void CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPass(
-    AddMachinePass &addPass, bool Optimized) const {
-  // TODO: Parse Opt.RegAlloc to add register allocator.
+Error CodeGenPassBuilder<Derived, TargetMachineT>::addRegAllocPass(
+    AddMachinePass &addPass, StringRef FilterName) const {
+  auto &RegAllocMap = PB.getRegAllocMap();
+  if (RegAllocMap.contains("none"))
+    return Error::success();
+
+  if (!RegAllocMap.contains(FilterName)) {
+    return make_error<StringError>(
+        formatv("No register allocator for register class filter '{0}'",
+                FilterName)
+            .str(),
+        inconvertibleErrorCode());
+  }
+
+  addPass(std::move(RegAllocMap[FilterName]));
+  return Error::success();
 }
 
 template <typename Derived, typename TargetMachineT>
 Error CodeGenPassBuilder<Derived, TargetMachineT>::addRegAssignmentFast(
     AddMachinePass &addPass) const {
-  // TODO: Ensure allocator is default or fast.
-  addRegAllocPass(addPass, false);
-  return Error::success();
+  return addRegAllocPass(addPass);
 }
 
 template <typename Derived, typename TargetMachineT>
 Error CodeGenPassBuilder<Derived, TargetMachineT>::addRegAssignmentOptimized(
     AddMachinePass &addPass) const {
   // Add the selected register allocation pass.
-  addRegAllocPass(addPass, true);
+  if (auto Err = addRegAllocPass(addPass))
+    return Err;
 
   // Allow targets to change the register assignments before rewriting.
   derived().addPreRewrite(addPass);

llvm/include/llvm/Passes/PassBuilder.h

@@ -107,6 +107,7 @@ class PassBuilder {
   PipelineTuningOptions PTO;
   std::optional<PGOOptions> PGOOpt;
   PassInstrumentationCallbacks *PIC;
+  StringMap<MachineFunctionPassManager> RegAllocMap;
 
 public:
   /// A struct to capture parsed pass pipeline names.
@@ -582,12 +583,28 @@ class PassBuilder {
 
   /// Register callbacks to parse target specific filter field if regalloc pass
   /// needs it. E.g. AMDGPU requires regalloc passes can handle sgpr and vgpr
-  /// separately.
+  /// separately. Currently "all" and "none" are preserved filter name.
   void registerRegClassFilterParsingCallback(
       const std::function<RegClassFilterFunc(StringRef)> &C) {
     RegClassFilterParsingCallbacks.push_back(C);
   }
 
+  /// Parse command line option `--regalloc-npm`
+  /// Should only be called by CodeGenPassBuilder.
+  Error parseRegAllocOpt(StringRef Text);
+
+  /// Target hook to set default regalloc.
+  void setDefaultRegAllocBuilder(
+      const std::function<void(StringMap<MachineFunctionPassManager> &)> &C) {
+    DefaultRegAllocBuilder = C;
+  }
+
+  /// Used by CodeGenPassBuilder to add correct regalloc pass.
+  /// Should only be called by CodeGenPassBuilder.
+  StringMap<MachineFunctionPassManager> &getRegAllocMap() {
+    return RegAllocMap;
+  }
+
   /// Register a callback for a top-level pipeline entry.
   ///
   /// If the PassManager type is not given at the top level of the pipeline
@@ -807,6 +824,8 @@ class PassBuilder {
   // Callbacks to parse `filter` parameter in register allocation passes
   SmallVector<std::function<RegClassFilterFunc(StringRef)>, 2>
       RegClassFilterParsingCallbacks;
+  std::function<void(StringMap<MachineFunctionPassManager> &)>
+      DefaultRegAllocBuilder;
 };
 
 /// This utility template takes care of adding require<> and invalidate<>

llvm/include/llvm/Target/TargetMachine.h

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

llvm/lib/Passes/PassBuilder.cpp

@@ -409,13 +409,33 @@ class RequireAllMachineFunctionPropertiesPass
 
 } // namespace
 
+static void defaultRegAllocBuilder(TargetMachine *TM,
+                                   StringMap<MachineFunctionPassManager> &M) {
+  if (!TM)
+    return;
+
+  MachineFunctionPassManager MFPM;
+  auto Opts = getCGPassBuilderOption();
+  if (Opts.OptimizeRegAlloc.value_or(TM->getOptLevel() !=
+                                     CodeGenOptLevel::None)) {
+    // TODO: Add greedy register allocator.
+  } else {
+    MFPM.addPass(RegAllocFastPass());
+  }
+  M["all"] = std::move(MFPM);
+}
+
 PassBuilder::PassBuilder(TargetMachine *TM, PipelineTuningOptions PTO,
                          std::optional<PGOOptions> PGOOpt,
                          PassInstrumentationCallbacks *PIC)
     : TM(TM), PTO(PTO), PGOOpt(PGOOpt), PIC(PIC) {
   bool ShouldPopulateClassToPassNames = PIC && shouldPopulateClassToPassNames();
-  if (TM)
+  if (TM) {
+    DefaultRegAllocBuilder = [TM](StringMap<MachineFunctionPassManager> &M) {
+      defaultRegAllocBuilder(TM, M);
+    };
     TM->registerPassBuilderCallbacks(*this, ShouldPopulateClassToPassNames);
+  }
   if (ShouldPopulateClassToPassNames) {
 #define MODULE_PASS(NAME, CREATE_PASS)                                         \
   PIC->addClassToPassName(decltype(CREATE_PASS)::name(), NAME);
@@ -2219,6 +2239,18 @@ Error PassBuilder::parseAAPipeline(AAManager &AA, StringRef PipelineText) {
   return Error::success();
 }
 
+static StringRef getFilterName(StringRef PassName) {
+  PassName = PassName.drop_until([](char C) { return C == '<'; });
+  StringRef Params = PassName.drop_front().drop_back();
+  while (!Params.empty()) {
+    StringRef ParamName;
+    std::tie(ParamName, Params) = Params.split(';');
+    if (ParamName.consume_front("filter="))
+      return ParamName;
+  }
+  return "all";
+}
+
 RegClassFilterFunc PassBuilder::parseRegAllocFilter(StringRef FilterName) {
   if (FilterName == "all")
     return allocateAllRegClasses;
@@ -2228,6 +2260,45 @@ RegClassFilterFunc PassBuilder::parseRegAllocFilter(StringRef FilterName) {
   return nullptr;
 }
 
+Error PassBuilder::parseRegAllocOpt(StringRef Text) {
+  assert(TM && "Need target machine to parse this option!");
+  if (RegAllocMap.empty())
+    DefaultRegAllocBuilder(RegAllocMap);
+
+  MachineFunctionPassManager MFPM;
+  if (Text == "default") {
+    // Add nothing when target inserts "none" into the map.
+    if (RegAllocMap.contains("none"))
+      RegAllocMap["all"] = MachineFunctionPassManager();
+    return Error::success();
+  }
+
+  if (RegAllocMap.contains("none")) {
+    return make_error<StringError>(
+        "Target doesn't support register allocation!",
+        inconvertibleErrorCode());
+  }
+
+  bool IsOptimized = TM->getOptLevel() != CodeGenOptLevel::None;
+  while (!Text.empty()) {
+    StringRef PassName;
+    std::tie(PassName, Text) = Text.split(',');
+    if (!IsOptimized &&
+        !PassBuilder::checkParametrizedPassName(PassName, "regallocfast")) {
+      return make_error<StringError>(
+          "Must use fast (default) register allocator for "
+          "unoptimized regalloc.",
+          inconvertibleErrorCode());
+    }
+    // FIXME: Should only accept reg-alloc passes.
+    if (auto Err = parsePassPipeline(MFPM, PassName))
+      return Err;
+    RegAllocMap[getFilterName(PassName)] = std::move(MFPM);
+    MFPM = MachineFunctionPassManager();
+  }
+  return Error::success();
+}
+
 static void printPassName(StringRef PassName, raw_ostream &OS) {
   OS << "  " << PassName << "\n";
 }

llvm/lib/Target/AMDGPU/AMDGPUCodeGenPassBuilder.cpp

@@ -14,9 +14,8 @@
 using namespace llvm;
 
 AMDGPUCodeGenPassBuilder::AMDGPUCodeGenPassBuilder(
-    AMDGPUTargetMachine &TM, const CGPassBuilderOption &Opts,
-    PassInstrumentationCallbacks *PIC)
-    : CodeGenPassBuilder(TM, Opts, PIC) {
+    AMDGPUTargetMachine &TM, const CGPassBuilderOption &Opts, PassBuilder &PB)
+    : CodeGenPassBuilder(TM, Opts, PB) {
   Opt.RequiresCodeGenSCCOrder = true;
   // Exceptions and StackMaps are not supported, so these passes will never do
   // anything.
@@ -40,3 +39,13 @@ Error AMDGPUCodeGenPassBuilder::addInstSelector(AddMachinePass &addPass) const {
   addPass(AMDGPUISelDAGToDAGPass(TM));
   return Error::success();
 }
+
+Error AMDGPUCodeGenPassBuilder::addRegAssignmentFast(
+    AddMachinePass &addPass) const {
+  if (auto Err = addRegAllocPass(addPass, "sgpr"))
+    return Err;
+  // TODO: Add other passes.
+  if (auto Err = addRegAllocPass(addPass, "vgpr"))
+    return Err;
+  return Error::success();
+}

llvm/lib/Target/AMDGPU/AMDGPUCodeGenPassBuilder.h

@@ -20,12 +20,12 @@ class AMDGPUCodeGenPassBuilder
     : public CodeGenPassBuilder<AMDGPUCodeGenPassBuilder, AMDGPUTargetMachine> {
 public:
   AMDGPUCodeGenPassBuilder(AMDGPUTargetMachine &TM,
-                           const CGPassBuilderOption &Opts,
-                           PassInstrumentationCallbacks *PIC);
+                           const CGPassBuilderOption &Opts, PassBuilder &PB);
 
   void addPreISel(AddIRPass &addPass) const;
   void addAsmPrinter(AddMachinePass &, CreateMCStreamer) const;
   Error addInstSelector(AddMachinePass &) const;
+  Error addRegAssignmentFast(AddMachinePass &addPass) const;
 };
 
 } // namespace llvm

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -46,6 +46,7 @@
 #include "llvm/CodeGen/GlobalISel/RegBankSelect.h"
 #include "llvm/CodeGen/MIRParser/MIParser.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/RegAllocFast.h"
 #include "llvm/CodeGen/RegAllocRegistry.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
@@ -650,11 +651,13 @@ parseAMDGPUAtomicOptimizerStrategy(StringRef Params) {
   return make_error<StringError>("invalid parameter", inconvertibleErrorCode());
 }
 
-Error AMDGPUTargetMachine::buildCodeGenPipeline(
-    ModulePassManager &MPM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
-    CodeGenFileType FileType, const CGPassBuilderOption &Opts,
-    PassInstrumentationCallbacks *PIC) {
-  AMDGPUCodeGenPassBuilder CGPB(*this, Opts, PIC);
+Error AMDGPUTargetMachine::buildCodeGenPipeline(ModulePassManager &MPM,
+                                                raw_pwrite_stream &Out,
+                                                raw_pwrite_stream *DwoOut,
+                                                CodeGenFileType FileType,
+                                                const CGPassBuilderOption &Opts,
+                                                PassBuilder &PB) {
+  AMDGPUCodeGenPassBuilder CGPB(*this, Opts, PB);
   return CGPB.buildPipeline(MPM, Out, DwoOut, FileType);
 }
 
@@ -749,6 +752,30 @@ void AMDGPUTargetMachine::registerPassBuilderCallbacks(
           return onlyAllocateVGPRs;
         return nullptr;
       });
+
+  PB.setDefaultRegAllocBuilder(
+      [TM = this](StringMap<MachineFunctionPassManager> &RegAllocMap) {
+        auto Opts = getCGPassBuilderOption();
+        if (Opts.OptimizeRegAlloc.value_or(TM->getOptLevel() !=
+                                           CodeGenOptLevel::None)) {
+          // TODO: Add greedy register allocator.
+        } else {
+          RegAllocFastPassOptions Opts;
+          Opts.Filter = onlyAllocateSGPRs;
+          Opts.FilterName = "sgpr";
+          Opts.ClearVRegs = false;
+          MachineFunctionPassManager MFPM;
+          MFPM.addPass(RegAllocFastPass(Opts));
+          RegAllocMap["sgpr"] = std::move(MFPM);
+
+          Opts.Filter = onlyAllocateVGPRs;
+          Opts.FilterName = "vgpr";
+          Opts.ClearVRegs = true;
+          MFPM = MachineFunctionPassManager();
+          MFPM.addPass(RegAllocFastPass(Opts));
+          RegAllocMap["vgpr"] = std::move(MFPM);
+        }
+      });
 }
 
 int64_t AMDGPUTargetMachine::getNullPointerValue(unsigned AddrSpace) {

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h

@@ -56,7 +56,7 @@ class AMDGPUTargetMachine : public LLVMTargetMachine {
                              raw_pwrite_stream *DwoOut,
                              CodeGenFileType FileType,
                              const CGPassBuilderOption &Opts,
-                             PassInstrumentationCallbacks *PIC) override;
+                             PassBuilder &PB) override;
 
   void registerPassBuilderCallbacks(PassBuilder &PB,
                                     bool PopulateClassToPassNames) override;