Add RegAlloc support

Author: paperchalice

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -56,6 +56,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"
@@ -121,7 +122,7 @@ namespace llvm {
     }                                                                          \
     static AnalysisKey Key;                                                    \
   };
-#include "llvm/CodeGen/MachinePassRegistry.def"
+#include "llvm/Passes/MachinePassRegistry.def"
 
 /// This class provides access to building LLVM's passes.
 ///
@@ -142,9 +143,10 @@ namespace llvm {
 
 template <typename DerivedT> class CodeGenPassBuilder {
 public:
-  explicit CodeGenPassBuilder(LLVMTargetMachine &TM, CGPassBuilderOption Opts,
+  explicit CodeGenPassBuilder(LLVMTargetMachine &TM, PassBuilder &PB,
+                              CGPassBuilderOption Opts,
                               PassInstrumentationCallbacks *PIC)
-      : TM(TM), Opt(Opts), PIC(PIC) {
+      : TM(TM), PB(PB), Opt(Opts), PIC(PIC) {
     // Target could set CGPassBuilderOption::MISchedPostRA to true to achieve
     //     substitutePass(&PostRASchedulerID, &PostMachineSchedulerID)
 
@@ -164,28 +166,6 @@ template <typename DerivedT> class CodeGenPassBuilder {
                       raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
                       CodeGenFileType FileType) const;
 
-  /// Parse single non-target-specific MIR pass
-  /// @param Name the pass name
-  /// @return true if failed
-  bool parseMIRPass(MachineFunctionPassManager &MFPM, StringRef Name) const;
-
-  /// Parse MIR pass pipeline. Unlike IR pass pipeline,
-  /// there is only one pass manager for machine function
-  /// so there is no need to specify the pass nesting.
-  /// @param Text a comma separated pass name list
-  Error parseMIRPipeline(MachineFunctionPassManager &MFPM,
-                         StringRef Text) const {
-    for (auto [LHS, RHS] = Text.split(','); LHS != "";
-         std::tie(LHS, RHS) = RHS.split(',')) {
-      if (parseMIRPass(MFPM, LHS) && derived().parseTargetMIRPass(MFPM, LHS)) {
-        return createStringError(
-            std::make_error_code(std::errc::invalid_argument),
-            Twine('\"') + Twine(LHS) + Twine("\" pass could not be found."));
-      }
-    }
-    return Error::success();
-  }
-
   void registerModuleAnalyses(ModuleAnalysisManager &) const;
   void registerFunctionAnalyses(FunctionAnalysisManager &) const;
   void registerMachineFunctionAnalyses(MachineFunctionAnalysisManager &) const;
@@ -279,6 +259,8 @@ template <typename DerivedT> class CodeGenPassBuilder {
           });
     }
 
+    MachineFunctionPassManager &getPM() { return PM; }
+
     MachineFunctionPassManager releasePM() { return std::move(PM); }
 
   private:
@@ -311,14 +293,10 @@ template <typename DerivedT> class CodeGenPassBuilder {
   }
 
   LLVMTargetMachine &TM;
+  PassBuilder &PB;
   CGPassBuilderOption Opt;
   PassInstrumentationCallbacks *PIC;
 
-  /// Target override these hooks to parse target-specific analyses.
-  void registerTargetAnalysis(ModuleAnalysisManager &) const {}
-  void registerTargetAnalysis(FunctionAnalysisManager &) const {}
-  void registerTargetAnalysis(MachineFunctionAnalysisManager &) const {}
-
   template <typename TMC> TMC &getTM() const { return static_cast<TMC &>(TM); }
   CodeGenOptLevel getOptLevel() const { return TM.getOptLevel(); }
 
@@ -556,14 +534,6 @@ template <typename DerivedT> 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;
-
   /// Add core register allocator passes which do the actual register assignment
   /// and rewriting. \returns Error::success() if any passes were added.
   Error addRegAssignAndRewriteFast(AddMachinePass &addPass) const;
@@ -1257,54 +1227,16 @@ void CodeGenPassBuilder<Derived>::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>
-void CodeGenPassBuilder<Derived>::addTargetRegisterAllocator(
-    AddMachinePass &addPass, bool Optimized) const {
-  if (Optimized)
-    addPass(RAGreedyPass());
-  else
-    addPass(RAFastPass());
-}
-
-/// 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>
-void CodeGenPassBuilder<Derived>::addRegAllocPass(AddMachinePass &addPass,
-                                                  bool Optimized) const {
-  if (Opt.RegAlloc == RegAllocType::Default)
-    // With no -regalloc= override, ask the target for a regalloc pass.
-    derived().addTargetRegisterAllocator(addPass, Optimized);
-  else if (Opt.RegAlloc == RegAllocType::Basic)
-    addPass(RABasicPass());
-  else if (Opt.RegAlloc == RegAllocType::Fast)
-    addPass(RAFastPass());
-  else if (Opt.RegAlloc == RegAllocType::Greedy)
-    addPass(RAGreedyPass());
-  else if (Opt.RegAlloc == RegAllocType::PBQP)
-    addPass(RAPBQPPass());
-  else
-    llvm_unreachable("unknonwn register allocator type");
-}
-
 template <typename Derived>
 Error CodeGenPassBuilder<Derived>::addRegAssignAndRewriteFast(
     AddMachinePass &addPass) const {
-  if (Opt.RegAlloc != RegAllocType::Default &&
-      Opt.RegAlloc != RegAllocType::Fast)
+  if (Opt.RegAlloc != "default" && !Opt.RegAlloc.starts_with("fast"))
     return make_error<StringError>(
         "Must use fast (default) register allocator for unoptimized regalloc.",
         inconvertibleErrorCode());
 
-  addPass(RegAllocPass(false));
+  if (Error Err = PB.parseRegAllocPass(addPass.getPM(), Opt.RegAlloc, false))
+    return Err;
 
   // Allow targets to change the register assignments after
   // fast register allocation.
@@ -1315,7 +1247,8 @@ template <typename DerivedT>
 Error CodeGenPassBuilder<DerivedT>::addRegAssignAndRewriteOptimized(
     AddMachinePass &addPass) const {
   // Add the selected register allocation pass.
-  addRegAllocPass(addPass, true);
+  if (Error Err = PB.parseRegAllocPass(addPass.getPM(), Opt.RegAlloc, true))
+    return Err;
   // Allow targets to change the register assignments before rewriting.
   addPreRewrite(addPass);
 
@@ -1367,22 +1300,22 @@ Error CodeGenPassBuilder<Derived>::addOptimizedRegAlloc(
   // PreRA instruction scheduling.
   addPass(MachineSchedulerPass());
 
-  Error E = derived().addRegAssignAndRewriteOptimized(addPass);
-  if (!E) {
-    // Allow targets to expand pseudo instructions depending on the choice of
-    // registers before MachineCopyPropagation.
-    derived().addPostRewrite(addPass);
+  if (Error Err = derived().addRegAssignAndRewriteOptimized(addPass))
+    return Err;
 
-    // Copy propagate to forward register uses and try to eliminate COPYs that
-    // were not coalesced.
-    addPass(MachineCopyPropagationPass());
+  // Allow targets to expand pseudo instructions depending on the choice of
+  // registers before MachineCopyPropagation.
+  derived().addPostRewrite(addPass);
 
-    // Run post-ra machine LICM to hoist reloads / remats.
-    //
-    // FIXME: can this move into MachineLateOptimization?
-    addPass(MachineLICMPass());
-  }
-  return E;
+  // Copy propagate to forward register uses and try to eliminate COPYs that
+  // were not coalesced.
+  addPass(MachineCopyPropagationPass());
+
+  // Run post-ra machine LICM to hoist reloads / remats.
+  //
+  // FIXME: can this move into MachineLateOptimization?
+  addPass(MachineLICMPass());
+  return Error::success();
 }
 
 //===---------------------------------------------------------------------===//

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -129,6 +129,23 @@ MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis,
 // MACHINE_FUNCTION_PASS("free-machine-function", FreeMachineFunctionPass, ())
 #undef MACHINE_FUNCTION_PASS
 
+#ifndef RA_PASS
+#define RA_PASS(NAME, PASS_NAME, CONSTRUCTOR)
+#endif
+// RA_PASS("ra-basic", RABasicPass, ())
+#undef RA_PASS
+
+#ifndef RA_PASS_WITH_PARAMS
+#define RA_PASS_WITH_PARAMS(NAME, PASS_NAME, CREATE_PASS, PARSER, PARAMS)
+#endif
+// RA_PASS_WITH_PARAMS("fast", RAFastPass,
+//   [](Option Opts){ return RAFastPass(Opts); }, parseRAFastPassOptions, "")
+// RA_PASS_WITH_PARAMS("greedy", RAGreedyPass,
+// [](Option Opts) { return RAGreedyPass(Opts); }, parseRAGreedyPassOptions, "")
+// RA_PASS_WITH_PARAMS("pbqp", RAPBQPPass,
+//   [](Option Opts) { return  RAPBQPPass(Opts); }, parseRAPBQPPassOptions, "")
+#undef RA_PASS_WITH_PARAMS
+
 // After a pass is converted to new pass manager, its entry should be moved from
 // dummy table to the normal one. For example, for a machine function pass,
 // DUMMY_MACHINE_FUNCTION_PASS to MACHINE_FUNCTION_PASS.
@@ -225,15 +242,10 @@ DUMMY_MACHINE_FUNCTION_PASS("processimpdefs", ProcessImplicitDefsPass, ())
 DUMMY_MACHINE_FUNCTION_PASS("prologepilog", PrologEpilogInserterPass, ())
 DUMMY_MACHINE_FUNCTION_PASS("prologepilog-code", PrologEpilogCodeInserterPass,
                             ())
-DUMMY_MACHINE_FUNCTION_PASS("ra-basic", RABasicPass, ())
-DUMMY_MACHINE_FUNCTION_PASS("ra-fast", RAFastPass, ())
-DUMMY_MACHINE_FUNCTION_PASS("ra-greedy", RAGreedyPass, ())
-DUMMY_MACHINE_FUNCTION_PASS("ra-pbqp", RAPBQPPass, ())
 DUMMY_MACHINE_FUNCTION_PASS("reg-usage-collector", RegUsageInfoCollectorPass,
                             ())
 DUMMY_MACHINE_FUNCTION_PASS("reg-usage-propagation",
                             RegUsageInfoPropagationPass, ())
-DUMMY_MACHINE_FUNCTION_PASS("regalloc", RegAllocPass, ())
 DUMMY_MACHINE_FUNCTION_PASS("regallocscoringpass", RegAllocScoringPass, ())
 DUMMY_MACHINE_FUNCTION_PASS("regbankselect", RegBankSelectPass, ())
 DUMMY_MACHINE_FUNCTION_PASS("removeredundantdebugvalues",

llvm/include/llvm/Passes/PassBuilder.h

@@ -370,6 +370,11 @@ class PassBuilder {
   Error parsePassPipeline(MachineFunctionPassManager &MFPM,
                           StringRef PipelineText);
 
+  /// Parse textual regalloc pass into the provided \c MachineFunctionPass
+  /// the PassText should contain only one regalloc pass text.
+  Error parseRegAllocPass(MachineFunctionPassManager &MFPM, StringRef PassText,
+                          bool Optimized);
+
   /// Parse a textual alias analysis pipeline into the provided AA manager.
   ///
   /// The format of the textual AA pipeline is a comma separated list of AA
@@ -575,6 +580,12 @@ class PassBuilder {
   }
   /// @}}
 
+  void registerRegAllocParsingCallback(
+      const std::function<bool(StringRef, MachineFunctionPassManager &, bool)>
+          &C) {
+    RegAllocPassParsingCallbacks.push_back(C);
+  }
+
   /// Register a callback for a top-level pipeline entry.
   ///
   /// If the PassManager type is not given at the top level of the pipeline
@@ -735,6 +746,9 @@ class PassBuilder {
       MachineFunctionAnalysisRegistrationCallbacks;
   SmallVector<std::function<bool(StringRef, MachineFunctionPassManager &)>, 2>
       MachinePipelineParsingCallbacks;
+  SmallVector<
+      std::function<bool(StringRef, MachineFunctionPassManager &, bool)>, 2>
+      RegAllocPassParsingCallbacks;
 };
 
 /// This utility template takes care of adding require<> and invalidate<>

llvm/include/llvm/Target/TargetMachine.h

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

llvm/lib/Passes/PassBuilder.cpp

@@ -2081,6 +2081,43 @@ Error PassBuilder::parsePassPipeline(MachineFunctionPassManager &MFPM,
   return Error::success();
 }
 
+Error PassBuilder::parseRegAllocPass(MachineFunctionPassManager &MFPM,
+                                     StringRef PassText, bool Optimized) {
+  // Targets may have their own default pipeline.
+  for (const auto &C : RegAllocPassParsingCallbacks) {
+    if (C(PassText, MFPM, Optimized))
+      return Error::success();
+  }
+
+  if (PassText == "default") {
+    if (Optimized) {
+      // TODO: MFPM.addPass(GreedyRegisterAllocatorPass());
+    } else {
+      // TODO: MFPM.addPass(FastRegisterAllocatorPass());
+    }
+    return Error::success();
+  }
+
+#define RA_PASS(NAME, PASS_NAME, CONSTRUCTOR)                                  \
+  if (PassText == NAME) {                                                      \
+    MFPM.addPass(PASS_NAME CONSTRUCTOR);                                       \
+    return Error::success();                                                   \
+  }
+#define RA_PASS_WITH_PARAMS(NAME, PASS_NAME, CREATE_PASS, PARSER, PARAMS)      \
+  if (checkParametrizedPassName(PassText, NAME)) {                             \
+    auto Params = parsePassParameters(PARSER, PassText, NAME);                 \
+    if (!Params)                                                               \
+      return Params.takeError();                                               \
+    MFPM.addPass(CREATE_PASS(Params.get()));                                   \
+    return Error::success();                                                   \
+  }
+#include "llvm/Passes/MachinePassRegistry.def"
+
+  return make_error<StringError>(
+      formatv("invalid regalloc pass '{0}'", PassText).str(),
+      inconvertibleErrorCode());
+}
+
 Error PassBuilder::parseAAPipeline(AAManager &AA, StringRef PipelineText) {
   // If the pipeline just consists of the word 'default' just replace the AA
   // manager with our default one.

llvm/lib/Target/X86/CMakeLists.txt

@@ -102,6 +102,7 @@ add_llvm_target(X86CodeGen ${sources}
   IRPrinter
   Instrumentation
   MC
+  Passes
   ProfileData
   Scalar
   SelectionDAG

llvm/lib/Target/X86/X86CodeGenPassBuilder.cpp

@@ -21,10 +21,10 @@ namespace {
 
 class X86CodeGenPassBuilder : public CodeGenPassBuilder<X86CodeGenPassBuilder> {
 public:
-  explicit X86CodeGenPassBuilder(LLVMTargetMachine &TM,
+  explicit X86CodeGenPassBuilder(LLVMTargetMachine &TM, PassBuilder &PB,
                                  CGPassBuilderOption Opts,
                                  PassInstrumentationCallbacks *PIC)
-      : CodeGenPassBuilder(TM, Opts, PIC) {}
+      : CodeGenPassBuilder(TM, PB, Opts, PIC) {}
   void addPreISel(AddIRPass &addPass) const;
   void addAsmPrinter(AddMachinePass &, CreateMCStreamer) const;
   Error addInstSelector(AddMachinePass &) const;
@@ -47,10 +47,9 @@ Error X86CodeGenPassBuilder::addInstSelector(AddMachinePass &) const {
 } // namespace
 
 Error X86TargetMachine::buildCodeGenPipeline(
-    ModulePassManager &MPM, MachineFunctionPassManager &MFPM,
-    MachineFunctionAnalysisManager &, raw_pwrite_stream &Out,
-    raw_pwrite_stream *DwoOut, CodeGenFileType FileType,
+    ModulePassManager &MPM, MachineFunctionPassManager &MFPM, PassBuilder &PB,
+    raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut, CodeGenFileType FileType,
     CGPassBuilderOption Opt, PassInstrumentationCallbacks *PIC) {
-  auto CGPB = X86CodeGenPassBuilder(*this, Opt, PIC);
+  auto CGPB = X86CodeGenPassBuilder(*this, PB, Opt, PIC);
   return CGPB.buildPipeline(MPM, MFPM, Out, DwoOut, FileType);
 }

llvm/lib/Target/X86/X86TargetMachine.h

@@ -59,9 +59,9 @@ class X86TargetMachine final : public LLVMTargetMachine {
                             const TargetSubtargetInfo *STI) const override;
 
   Error buildCodeGenPipeline(ModulePassManager &, MachineFunctionPassManager &,
-                             MachineFunctionAnalysisManager &,
-                             raw_pwrite_stream &, raw_pwrite_stream *,
-                             CodeGenFileType, CGPassBuilderOption,
+                             PassBuilder &, raw_pwrite_stream &,
+                             raw_pwrite_stream *, CodeGenFileType,
+                             CGPassBuilderOption,
                              PassInstrumentationCallbacks *) override;
 
   bool isJIT() const { return IsJIT; }

llvm/test/tools/llc/new-pm/pipeline.ll

@@ -1,5 +1,4 @@
 ; RUN: llc -mtriple=x86_64-pc-linux-gnu -enable-new-pm -print-pipeline-passes -filetype=null %s | FileCheck %s
 
 ; CHECK: require<profile-summary>,require<collector-metadata>
-; CHECK: MachineSanitizerBinaryMetadata,FreeMachineFunctionPass
-
+; CHECK: StackFrameLayoutAnalysisPass,FreeMachineFunctionPass

llvm/tools/llc/NewPMDriver.cpp

@@ -191,7 +191,7 @@ int llvm::compileModuleWithNewPM(
     ModulePassManager MPM;
     MachineFunctionPassManager MFPM;
 
-    ExitOnErr(LLVMTM.buildCodeGenPipeline(MPM, MFPM, MFAM, *OS,
+    ExitOnErr(LLVMTM.buildCodeGenPipeline(MPM, MFPM, PB, *OS,
                                           DwoOut ? &DwoOut->os() : nullptr,
                                           FileType, Opt, &PIC));