[NewPM] Add callback style codegen pass builder

Author: paperchalice

llvm/include/llvm/CodeGen/TargetPassConfig.h

@@ -190,8 +190,7 @@ class TargetPassConfig : public ImmutablePass {
 
   /// Returns pass name in `-stop-before` or `-stop-after`
   /// NOTE: New pass manager migration only
-  static Expected<StartStopInfo>
-  getStartStopInfo(PassInstrumentationCallbacks &PIC);
+  static Expected<StartStopInfo> getStartStopInfo();
 
   void setDisableVerify(bool Disable) { setOpt(DisableVerify, Disable); }
 

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -523,7 +523,7 @@ template <typename Derived, typename TargetMachineT>
 Error CodeGenPassBuilder<Derived, TargetMachineT>::buildPipeline(
     ModulePassManager &MPM, raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
     CodeGenFileType FileType) const {
-  auto StartStopInfo = TargetPassConfig::getStartStopInfo(*PIC);
+  auto StartStopInfo = TargetPassConfig::getStartStopInfo();
   if (!StartStopInfo)
     return StartStopInfo.takeError();
   setStartStopPasses(*StartStopInfo);

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -191,6 +191,7 @@ DUMMY_MACHINE_MODULE_PASS("mir-strip-debug", StripDebugMachineModulePass)
 #ifndef DUMMY_MACHINE_FUNCTION_PASS
 #define DUMMY_MACHINE_FUNCTION_PASS(NAME, PASS_NAME)
 #endif
+DUMMY_MACHINE_FUNCTION_PASS("bb-path-cloning", BasicBlockPathCloningPass)
 DUMMY_MACHINE_FUNCTION_PASS("bbsections-prepare", BasicBlockSectionsPass)
 DUMMY_MACHINE_FUNCTION_PASS("bbsections-profile-reader", BasicBlockSectionsProfileReaderPass)
 DUMMY_MACHINE_FUNCTION_PASS("block-placement", MachineBlockPlacementPass)
@@ -210,6 +211,8 @@ DUMMY_MACHINE_FUNCTION_PASS("fixup-statepoint-caller-saved", FixupStatepointCall
 DUMMY_MACHINE_FUNCTION_PASS("fs-profile-loader", MIRProfileLoaderNewPass)
 DUMMY_MACHINE_FUNCTION_PASS("funclet-layout", FuncletLayoutPass)
 DUMMY_MACHINE_FUNCTION_PASS("gc-empty-basic-blocks", GCEmptyBasicBlocksPass)
+DUMMY_MACHINE_FUNCTION_PASS("gc-machine-code-insersion",
+                            GCMachineCodeInsertionPass)
 DUMMY_MACHINE_FUNCTION_PASS("implicit-null-checks", ImplicitNullChecksPass)
 DUMMY_MACHINE_FUNCTION_PASS("init-undef-pass", InitUndefPass)
 DUMMY_MACHINE_FUNCTION_PASS("instruction-select", InstructionSelectPass)
@@ -261,6 +264,8 @@ DUMMY_MACHINE_FUNCTION_PASS("stack-slot-coloring", StackSlotColoringPass)
 DUMMY_MACHINE_FUNCTION_PASS("stackmap-liveness", StackMapLivenessPass)
 DUMMY_MACHINE_FUNCTION_PASS("tailduplication", TailDuplicatePass)
 DUMMY_MACHINE_FUNCTION_PASS("unpack-mi-bundles", UnpackMachineBundlesPass)
+DUMMY_MACHINE_FUNCTION_PASS("unreachable-mbb-elimination",
+                            UnreachableMachineBlockElimPass)
 DUMMY_MACHINE_FUNCTION_PASS("virtregrewriter", VirtRegRewriterPass)
 DUMMY_MACHINE_FUNCTION_PASS("xray-instrumentation", XRayInstrumentationPass)
 #undef DUMMY_MACHINE_FUNCTION_PASS

llvm/include/llvm/Passes/PassBuilder.h

@@ -15,14 +15,17 @@
 #ifndef LLVM_PASSES_PASSBUILDER_H
 #define LLVM_PASSES_PASSBUILDER_H
 
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/RegAllocCommon.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Passes/OptimizationLevel.h"
+#include "llvm/Support/CodeGen.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/PGOOptions.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Target/CGPassBuilderOption.h"
 #include "llvm/Transforms/IPO/Inliner.h"
 #include "llvm/Transforms/IPO/ModuleInliner.h"
 #include "llvm/Transforms/Instrumentation.h"
@@ -35,6 +38,7 @@ class StringRef;
 class AAManager;
 class TargetMachine;
 class ModuleSummaryIndex;
+class MCContext;
 template <typename T> class IntrusiveRefCntPtr;
 namespace vfs {
 class FileSystem;
@@ -107,6 +111,7 @@ class PassBuilder {
   TargetMachine *TM;
   PipelineTuningOptions PTO;
   std::optional<PGOOptions> PGOOpt;
+  CGPassBuilderOption CGPBO;
   PassInstrumentationCallbacks *PIC;
 
 public:
@@ -307,6 +312,24 @@ class PassBuilder {
   /// TargetMachine::registerDefaultAliasAnalyses().
   AAManager buildDefaultAAPipeline();
 
+  /// Build CodeGen pass pipeline.
+  ///
+  /// {{@
+  Expected<ModulePassManager>
+  buildDefaultCodeGenPipeline(raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
+                              CodeGenFileType FileType, MCContext &Ctx);
+  Error buildDefaultCodeGenPipeline(ModulePassManager &MPM,
+                                    raw_pwrite_stream &Out,
+                                    raw_pwrite_stream *DwoOut,
+                                    CodeGenFileType FileType, MCContext &Ctx);
+  Error addRegAllocPass(MachineFunctionPassManager &MFPM,
+                        StringRef Filter = "all");
+  // TODO: Add method to build MC emission pipeline.
+  template <typename... PassTs> void disablePass() {
+    (DisabledPasses.insert(PassTs::name()), ...);
+  }
+  /// @}}
+
   /// Parse a textual pass pipeline description into a \c
   /// ModulePassManager.
   ///
@@ -519,6 +542,133 @@ class PassBuilder {
     FullLinkTimeOptimizationLastEPCallbacks.push_back(C);
   }
 
+  /// Register target specific callbacks to extend codegen pipeline.
+  /// {{@
+
+  /// If target want its own pipeline, use this callback.
+  void setCustomCodeGenPipelineBuilderCallback(
+      const std::function<Error(ModulePassManager &, raw_pwrite_stream &,
+                                raw_pwrite_stream *, CodeGenFileType,
+                                MCContext &)>
+          C) {
+    CustomCodeGenPipelineBuilderCallback = C;
+  }
+
+  void registerCodeGenIREarlyEPCallback(
+      const std::function<void(ModulePassManager &)> C) {
+    CodeGenIREarlyEPCallbacks.push_back(C);
+  }
+
+  void registerCodeGenIRLateEPCallback(
+      const std::function<void(ModulePassManager &)> C) {
+    CodeGenIRLateEPCallbacks.push_back(C);
+  }
+
+  void registerCodeGenIRFinalEPCallback(
+      const std::function<void(FunctionPassManager &)> C) {
+    CodeGenIRFinalEPCallbacks.push_back(C);
+  }
+
+  void registerGCLoweringEPCallback(
+      const std::function<void(FunctionPassManager &)> C) {
+    GCLoweringEPCallbacks.push_back(C);
+  }
+
+  void registerISelPrepareEPCallback(
+      const std::function<void(ModulePassManager &)> &C) {
+    ISelPrepareEPCallbacks.push_back(C);
+  }
+
+  void registerMachineSSAOptimizationEarlyEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    MachineSSAOptimizationEarlyEPCallbacks.push_back(C);
+  }
+
+  void registerILPOptsEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    ILPOptsEPCallbacks.push_back(C);
+  }
+
+  void registerMachineSSAOptimizationLastEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    MachineSSAOptimizationLastEPCallbacks.push_back(C);
+  }
+
+  void registerPreRegAllocEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    PreRegAllocEPCallbacks.push_back(C);
+  }
+
+  void registerPostRegAllocEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    PostRegAllocEPCallbacks.push_back(C);
+  }
+
+  void registerPreRegBankSelectEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    PreRegBankSelectEPCallbacks.push_back(C);
+  }
+
+  void registerPreGlobalInstructionSelectEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    PreGlobalInstructionSelectEPCallbacks.push_back(C);
+  }
+
+  void registerPostGlobalInstructionSelectEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    PostGlobalInstructionSelectEPCallbacks.push_back(C);
+  }
+
+  void registerMachineLateOptimizationEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    MachineLateOptimizationEPCallbacks.push_back(C);
+  }
+
+  void registerPreSched2EPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    PreSched2EPCallbacks.push_back(C);
+  }
+
+  void registerPostRewriteEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    PostRewriteEPCallbacks.push_back(C);
+  }
+
+  void registerPreEmitEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    PreEmitEPCallbacks.push_back(C);
+  }
+
+  void registerPostBBSectionsEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    PostBBSectionsEPCallbacks.push_back(C);
+  }
+
+  void registerMIEmitEPCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    MIEmitEPCallbacks.push_back(C);
+  }
+
+  void setAddInstSelectorCallback(
+      const std::function<void(MachineFunctionPassManager &)> &C) {
+    AddInstSelectorCallback = C;
+  }
+
+  void setCustomFastRegAllocPipelineBuilder(
+      const std::function<Error(MachineFunctionPassManager &)> &C) {
+    AddFastRegAllocCallback = C;
+  }
+
+  void setCustomOptimizedRegAllocPipelineBuilder(
+      const std::function<Error(MachineFunctionPassManager &)> &C) {
+    AddOptimizedRegAllocCallback = C;
+  }
+  ///@}}
+
+  // New pass manager migration methods, don't use them
+  // outside llvm!
+  CGPassBuilderOption &getCGPBO() { return CGPBO; }
+
   /// Register a callback for parsing an AliasAnalysis Name to populate
   /// the given AAManager \p AA
   void registerParseAACallback(
@@ -640,6 +790,30 @@ class PassBuilder {
   void invokePipelineEarlySimplificationEPCallbacks(ModulePassManager &MPM,
                                                     OptimizationLevel Level);
 
+  void invokeCodeGenIREarlyEPCallbacks(ModulePassManager &MPM);
+  void invokeCodeGenIRLateEPCallbacks(ModulePassManager &MPM);
+  void invokeCodeGenIRFinalEPCallbacks(FunctionPassManager &FPM);
+  void invokeGCLoweringEPCallbacks(FunctionPassManager &FPM);
+  void invokeISelPrepareEPCallbacks(ModulePassManager &MPM);
+  void invokeMachineSSAOptimizationEarlyEPCallbacks(
+      MachineFunctionPassManager &MFPM);
+  void
+  invokeMachineSSAOptimizationLastEPCallbacks(MachineFunctionPassManager &MFPM);
+  void invokePreRegAllocEPCallbacks(MachineFunctionPassManager &MFPM);
+  void invokePostRegAllocEPCallbacks(MachineFunctionPassManager &MFPM);
+  void invokePreRegBankSelectEPCallbacks(MachineFunctionPassManager &MFPM);
+  void
+  invokePreGlobalInstructionSelectEPCallbacks(MachineFunctionPassManager &MFPM);
+  void invokePostGlobalInstructionSelectEPCallbacks(
+      MachineFunctionPassManager &MFPM);
+  void invokeILPOptsEPCallbacks(MachineFunctionPassManager &MFPM);
+  void
+  invokeMachineLateOptimizationEPCallbacks(MachineFunctionPassManager &MFPM);
+  void invokePreEmitEPCallbacks(MachineFunctionPassManager &MFPM);
+  void invokePostBBSectionsEPCallbacks(MachineFunctionPassManager &MFPM);
+  void invokeMIEmitEPCallbacks(MachineFunctionPassManager &MFPM);
+  void invokePreSched2EPCallbacks(MachineFunctionPassManager &MFPM);
+
   static bool checkParametrizedPassName(StringRef Name, StringRef PassName) {
     if (!Name.consume_front(PassName))
       return false;
@@ -704,6 +878,21 @@ class PassBuilder {
   void addVectorPasses(OptimizationLevel Level, FunctionPassManager &FPM,
                        bool IsFullLTO);
 
+  Error addExceptionHandlingPasses(FunctionPassManager &FPM);
+
+  Error addInstructionSelectorPasses(MachineFunctionPassManager &MFPM);
+
+  void addMachineSSAOptimizationPasses(MachineFunctionPassManager &MFPM);
+
+  Error addMachinePasses(ModulePassManager &MPM, FunctionPassManager &FPM,
+                         MachineFunctionPassManager &MFPM);
+
+  Error addRegisterAllocatorPasses(MachineFunctionPassManager &MFPM);
+
+  Error parseRegAllocOption(StringRef Text);
+
+  bool isOptimizedRegAlloc() const;
+
   static std::optional<std::vector<PipelineElement>>
   parsePipelineText(StringRef Text);
 
@@ -766,6 +955,58 @@ class PassBuilder {
   SmallVector<std::function<void(ModulePassManager &, OptimizationLevel)>, 2>
       PipelineEarlySimplificationEPCallbacks;
 
+  // CodeGen extension point callbacks
+  std::function<Error(ModulePassManager &, raw_pwrite_stream &,
+                      raw_pwrite_stream *, CodeGenFileType, MCContext &)>
+      CustomCodeGenPipelineBuilderCallback;
+
+  SmallVector<std::function<void(ModulePassManager &)>, 2>
+      CodeGenIREarlyEPCallbacks;
+  SmallVector<std::function<void(ModulePassManager &)>, 2>
+      CodeGenIRLateEPCallbacks;
+  SmallVector<std::function<void(FunctionPassManager &)>, 2>
+      CodeGenIRFinalEPCallbacks;
+  SmallVector<std::function<void(FunctionPassManager &)>, 2>
+      GCLoweringEPCallbacks;
+  SmallVector<std::function<void(ModulePassManager &)>, 2>
+      ISelPrepareEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      MachineSSAOptimizationEarlyEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      MachineSSAOptimizationLastEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      PreRegAllocEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      PostRegAllocEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      PreRegBankSelectEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      PreGlobalInstructionSelectEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      PostGlobalInstructionSelectEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      ILPOptsEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      MachineLateOptimizationEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      PreSched2EPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      PostRewriteEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      PreEmitEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      PostBBSectionsEPCallbacks;
+  SmallVector<std::function<void(MachineFunctionPassManager &)>, 2>
+      MIEmitEPCallbacks;
+
+  std::function<void(MachineFunctionPassManager &)> AddInstSelectorCallback;
+  std::function<Error(MachineFunctionPassManager &)> AddFastRegAllocCallback;
+  std::function<Error(MachineFunctionPassManager &)>
+      AddOptimizedRegAllocCallback;
+  StringSet<> DisabledPasses;
+  StringMap<MachineFunctionPassManager> RegAllocPasses;
+  // TODO: Add methods in LLVMTargetMachine so we can get rid of it.
+
   SmallVector<std::function<void(ModuleAnalysisManager &)>, 2>
       ModuleAnalysisRegistrationCallbacks;
   SmallVector<std::function<bool(StringRef, ModulePassManager &,

llvm/include/llvm/Target/CGPassBuilderOption.h

@@ -41,6 +41,9 @@ struct CGPassBuilderOption {
   bool PrintLSR = false;
   bool DisableMergeICmps = false;
   bool DisablePartialLibcallInlining = false;
+  bool DisableReplaceWithVecLib = false;
+  bool DisableLayoutFSProfileLoader = false;
+  bool DisablePrologEpilogInserterPass = false;
   bool DisableConstantHoisting = false;
   bool DisableSelectOptimize = true;
   bool DisableAtExitBasedGlobalDtorLowering = false;

llvm/lib/CodeGen/MachinePassManager.cpp

@@ -105,6 +105,18 @@ bool MachineFunctionAnalysisManagerFunctionProxy::Result::invalidate(
   return false;
 }
 
+void FunctionToMachineFunctionPassAdaptor::eraseIf(
+    function_ref<bool(StringRef)> Pred) {
+  StringRef PassName = Pass->name();
+  if (PassName.contains("PassManager") || PassName.ends_with("PassAdaptor")) {
+    Pass->eraseIf(Pred);
+    if (Pass->isEmpty())
+      Pass.reset();
+  } else if (Pred(PassName)) {
+    Pass.reset();
+  }
+}
+
 PreservedAnalyses
 FunctionToMachineFunctionPassAdaptor::run(Function &F,
                                           FunctionAnalysisManager &FAM) {