[NewPM] Add extension points to LTO pipeline in PassBuilder

[EliaGeretto]

Summary

This PR adds two extension points to the default LTO pipeline in PassBuilder, one at the beginning and one at the end. These two extension points already existed in the old pass manager, the aim is to replicate the same functionality in the new one.

[aeubanks]

seems reasonable to me
any tests for O1/O2?

[EliaGeretto]

I added tests for -O1 and -O2.

[aeubanks]

I think new-pm-lto-defaults.ll is better for the O1/O2 tests since we actually see where they are in the pipeline, see new-pm-thinlto-defaults.ll for inspiration
otherwise looks good

[EliaGeretto]

Given that new-pm-lto-defaults.ll was also not a great fit, I decided to move all the tests for the LTO extension points to a new new-pm-lto-ep-callbacks.ll file, which is very similar to new-pm-O0-ep-callbacks.ll but has a more meaningful name.

[aeubanks]

sorry for the nitpicking, but I was thinking of something like https://github.com/aeubanks/llvm-project/commit/e93062bac54abf72b49e58415d682adde1deec7c
seeing exactly where in the pipeline the extension point passes run is kinda nice and actually caused me to move it around a tiny bit (probably doesn't matter, but it's consistent with other default pipelines)

llvm/tools/opt/NewPMDriver.cpp

244

ModulePassManager, ditto below

[EliaGeretto]

No problem, I will add your changes and reupload the patch then. Also, for my benefit, why moving the Early extension point after the Annotation2MetadataPass? I was not sure about the proper position, so I am interested in knowing the logic, if you do not mind.

[EliaGeretto]

I have incorporated the changes proposed by @aeubanks in my patch. I have also moved the O0 tests back to new-pm-O0-ep-callbacks.ll.

[aeubanks]

lgtm

[EliaGeretto]

Also, looking at the correction you made in NewPMDriver.cpp, I have noticed that the same issue is present for the OptimizerLast extension point as well. It accepts module passes, but it is tested with a function pass.

[aeubanks]

do you need me to push this for you?

In D120491#3346048, @EliaGeretto wrote:

Also, looking at the correction you made in NewPMDriver.cpp, I have noticed that the same issue is present for the OptimizerLast extension point as well. It accepts module passes, but it is tested with a function pass.

ah I can fix that separately

[EliaGeretto]

Yes, please, I don't have commit access. You can use Elia Geretto <[e-mail]> for the Author field. Also, I already have a patch for the OptimizerLast stuff, I am going to submit that one as well separately.

llvm/include/llvm/Passes/PassBuilder.h

public:
   ///
   /// This extension point allows adding optimizations at the very end of the
   /// function optimization pipeline.
   void registerOptimizerLastEPCallback(
       const std::function<void(ModulePassManager &, OptimizationLevel)> &C) {
     OptimizerLastEPCallbacks.push_back(C);
   }
 
+  /// Register a callback for a default optimizer pipeline extension point
+  ///
+  /// This extension point allows adding optimizations at the start of the full
+  /// LTO pipeline.
+  void registerFullLinkTimeOptimizationEarlyEPCallback(
+      const std::function<void(ModulePassManager &, OptimizationLevel)> &C) {
+    FullLinkTimeOptimizationEarlyEPCallbacks.push_back(C);
+  }
+
+  /// Register a callback for a default optimizer pipeline extension point
+  ///
+  /// This extension point allows adding optimizations at the end of the full
+  /// LTO pipeline.
+  void registerFullLinkTimeOptimizationLastEPCallback(
+      const std::function<void(ModulePassManager &, OptimizationLevel)> &C) {
+    FullLinkTimeOptimizationLastEPCallbacks.push_back(C);
+  }
+
   /// Register a callback for parsing an AliasAnalysis Name to populate
   /// the given AAManager \p AA
   void registerParseAACallback(
       const std::function<bool(StringRef Name, AAManager &AA)> &C) {
     AAParsingCallbacks.push_back(C);
   }
 
   /// {{@ Register callbacks for analysis registration with this PassBuilder
private:
   SmallVector<std::function<void(FunctionPassManager &, OptimizationLevel)>, 2>
       ScalarOptimizerLateEPCallbacks;
   SmallVector<std::function<void(CGSCCPassManager &, OptimizationLevel)>, 2>
       CGSCCOptimizerLateEPCallbacks;
   SmallVector<std::function<void(FunctionPassManager &, OptimizationLevel)>, 2>
       VectorizerStartEPCallbacks;
   SmallVector<std::function<void(ModulePassManager &, OptimizationLevel)>, 2>
       OptimizerLastEPCallbacks;
+  SmallVector<std::function<void(ModulePassManager &, OptimizationLevel)>, 2>
+      FullLinkTimeOptimizationEarlyEPCallbacks;
+  SmallVector<std::function<void(ModulePassManager &, OptimizationLevel)>, 2>
+      FullLinkTimeOptimizationLastEPCallbacks;
   // Module callbacks
   SmallVector<std::function<void(ModulePassManager &, OptimizationLevel)>, 2>
       PipelineStartEPCallbacks;
   SmallVector<std::function<void(ModulePassManager &, OptimizationLevel)>, 2>
       PipelineEarlySimplificationEPCallbacks;
 
   SmallVector<std::function<void(ModuleAnalysisManager &)>, 2>
       ModuleAnalysisRegistrationCallbacks;

llvm/lib/Passes/PassBuilderPipelines.cpp

return buildPerModuleDefaultPipeline(Level,
                                        /* LTOPreLink */ true);
 }
 
 ModulePassManager
 PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
                                      ModuleSummaryIndex *ExportSummary) {
   ModulePassManager MPM;
 
+  for (auto &C : FullLinkTimeOptimizationEarlyEPCallbacks)
+    C(MPM, Level);
+
   // Convert @llvm.global.annotations to !annotation metadata.
   MPM.addPass(Annotation2MetadataPass());
 
   // Create a function that performs CFI checks for cross-DSO calls with targets
   // in the current module.
   MPM.addPass(CrossDSOCFIPass());
 
   if (Level == OptimizationLevel::O0) {
     // The WPD and LowerTypeTest passes need to run at -O0 to lower type
     // metadata and intrinsics.
     MPM.addPass(WholeProgramDevirtPass(ExportSummary, nullptr));
     MPM.addPass(LowerTypeTestsPass(ExportSummary, nullptr));
     // Run a second time to clean up any type tests left behind by WPD for use
     // in ICP.
     MPM.addPass(LowerTypeTestsPass(nullptr, nullptr, true));
 
+    for (auto &C : FullLinkTimeOptimizationLastEPCallbacks)
+      C(MPM, Level);
+
     // Emit annotation remarks.
     addAnnotationRemarksPass(MPM);
 
     return MPM;
   }
 
   if (PGOOpt && PGOOpt->Action == PGOOptions::SampleUse) {
     // Load sample profile before running the LTO optimization pipeline.
if (Level == OptimizationLevel::O1) {
     // The LowerTypeTestsPass needs to run to lower type metadata and the
     // type.test intrinsics. The pass does nothing if CFI is disabled.
     MPM.addPass(LowerTypeTestsPass(ExportSummary, nullptr));
     // Run a second time to clean up any type tests left behind by WPD for use
     // in ICP (which is performed earlier than this in the regular LTO
     // pipeline).
     MPM.addPass(LowerTypeTestsPass(nullptr, nullptr, true));
 
+    for (auto &C : FullLinkTimeOptimizationLastEPCallbacks)
+      C(MPM, Level);
+
     // Emit annotation remarks.
     addAnnotationRemarksPass(MPM);
 
     return MPM;
   }
 
   // Optimize globals to try and fold them into constants.
   MPM.addPass(GlobalOptPass());
PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
   MPM.addPass(EliminateAvailableExternallyPass());
 
   // Now that we have optimized the program, discard unreachable functions.
   MPM.addPass(GlobalDCEPass());
 
   if (PTO.MergeFunctions)
     MPM.addPass(MergeFunctionsPass());
 
+  for (auto &C : FullLinkTimeOptimizationLastEPCallbacks)
+    C(MPM, Level);
+
   // Emit annotation remarks.
   addAnnotationRemarksPass(MPM);
 
   return MPM;
 }
 
 ModulePassManager PassBuilder::buildO0DefaultPipeline(OptimizationLevel Level,
                                                       bool LTOPreLink) {

llvm/test/Other/new-pm-ep-callbacks.ll

 ; RUN: opt -disable-output -debug-pass-manager -passes-ep-late-loop-optimizations=no-op-loop -passes='default<O0>' 2>&1 < %s | FileCheck %s
 ; RUN: opt -disable-output -debug-pass-manager -passes-ep-loop-optimizer-end=no-op-loop -passes='default<O0>' 2>&1 < %s | FileCheck %s
 ; RUN: opt -disable-output -debug-pass-manager -passes-ep-scalar-optimizer-late=no-op-function -passes='default<O0>' 2>&1 < %s | FileCheck %s
 ; RUN: opt -disable-output -debug-pass-manager -passes-ep-cgscc-optimizer-late=no-op-cgscc -passes='default<O0>' 2>&1 < %s | FileCheck %s
 ; RUN: opt -disable-output -debug-pass-manager -passes-ep-vectorizer-start=no-op-function -passes='default<O0>' 2>&1 < %s | FileCheck %s
 ; RUN: opt -disable-output -debug-pass-manager -passes-ep-pipeline-start=no-op-module -passes='default<O0>' 2>&1 < %s | FileCheck %s
 ; RUN: opt -disable-output -debug-pass-manager -passes-ep-pipeline-early-simplification=no-op-module -passes='default<O0>' 2>&1 < %s | FileCheck %s
 ; RUN: opt -disable-output -debug-pass-manager -passes-ep-optimizer-last=no-op-function -passes='default<O0>' 2>&1 < %s | FileCheck %s
+; RUN: opt -disable-output -debug-pass-manager -passes-ep-full-link-time-optimization-early=no-op-function -passes='lto<O0>' 2>&1 < %s | FileCheck %s
+; RUN: opt -disable-output -debug-pass-manager -passes-ep-full-link-time-optimization-last=no-op-function -passes='lto<O0>' 2>&1 < %s | FileCheck %s
+; RUN: opt -disable-output -debug-pass-manager -passes-ep-full-link-time-optimization-early=no-op-function -passes='lto<O1>' 2>&1 < %s | FileCheck %s
+; RUN: opt -disable-output -debug-pass-manager -passes-ep-full-link-time-optimization-last=no-op-function -passes='lto<O1>' 2>&1 < %s | FileCheck %s
+; RUN: opt -disable-output -debug-pass-manager -passes-ep-full-link-time-optimization-early=no-op-function -passes='lto<O2>' 2>&1 < %s | FileCheck %s
+; RUN: opt -disable-output -debug-pass-manager -passes-ep-full-link-time-optimization-last=no-op-function -passes='lto<O2>' 2>&1 < %s | FileCheck %s
 
 ; CHECK: Running pass: NoOp
 
 declare void @bar() local_unnamed_addr
 
 define void @foo(i32 %n) local_unnamed_addr {
 entry:
   br label %loop

llvm/tools/opt/NewPMDriver.cpp

static cl::opt<std::string> PipelineEarlySimplificationEPPipeline(
     cl::desc("A textual description of the module pass pipeline inserted at "
              "the EarlySimplification extension point into default pipelines"),
     cl::Hidden);
 static cl::opt<std::string> OptimizerLastEPPipeline(
     "passes-ep-optimizer-last",
     cl::desc("A textual description of the module pass pipeline inserted at "
              "the OptimizerLast extension point into default pipelines"),
     cl::Hidden);
+static cl::opt<std::string> FullLinkTimeOptimizationEarlyEPPipeline(
+    "passes-ep-full-link-time-optimization-early",
+    cl::desc("A textual description of the module pass pipeline inserted at "
+             "the FullLinkTimeOptimizationEarly extension point into default "
+             "pipelines"),
+    cl::Hidden);
+static cl::opt<std::string> FullLinkTimeOptimizationLastEPPipeline(
+    "passes-ep-full-link-time-optimization-last",
+    cl::desc("A textual description of the module pass pipeline inserted at "
+             "the FullLinkTimeOptimizationLast extension point into default "
+             "pipelines"),
+    cl::Hidden);
 
 // Individual pipeline tuning options.
 extern cl::opt<bool> DisableLoopUnrolling;
 
 namespace llvm {
 extern cl::opt<PGOKind> PGOKindFlag;
 extern cl::opt<std::string> ProfileFile;
 extern cl::opt<CSPGOKind> CSPGOKindFlag;
PB.registerPipelineEarlySimplificationEPCallback(
           Err(PB.parsePassPipeline(PM, PipelineEarlySimplificationEPPipeline));
         });
   if (tryParsePipelineText<FunctionPassManager>(PB, OptimizerLastEPPipeline))
     PB.registerOptimizerLastEPCallback(
         [&PB](ModulePassManager &PM, OptimizationLevel) {
           ExitOnError Err("Unable to parse OptimizerLastEP pipeline: ");
           Err(PB.parsePassPipeline(PM, OptimizerLastEPPipeline));
         });
+  if (tryParsePipelineText<FunctionPassManager>(
+          PB, FullLinkTimeOptimizationEarlyEPPipeline))
+    PB.registerFullLinkTimeOptimizationEarlyEPCallback(
+        [&PB](ModulePassManager &PM, OptimizationLevel) {
+          ExitOnError Err(
+              "Unable to parse FullLinkTimeOptimizationEarlyEP pipeline: ");
+          Err(PB.parsePassPipeline(PM,
+                                   FullLinkTimeOptimizationEarlyEPPipeline));
+        });
+  if (tryParsePipelineText<FunctionPassManager>(
+          PB, FullLinkTimeOptimizationLastEPPipeline))
+    PB.registerFullLinkTimeOptimizationLastEPCallback(
+        [&PB](ModulePassManager &PM, OptimizationLevel) {
+          ExitOnError Err(
+              "Unable to parse FullLinkTimeOptimizationLastEP pipeline: ");
+          Err(PB.parsePassPipeline(PM, FullLinkTimeOptimizationLastEPPipeline));
+        });
 }
 
 #define HANDLE_EXTENSION(Ext)                                                  \
   llvm::PassPluginLibraryInfo get##Ext##PluginInfo();
 #include "llvm/Support/Extension.def"
 
 bool llvm::runPassPipeline(StringRef Arg0, Module &M, TargetMachine *TM,
                            TargetLibraryInfoImpl *TLII, ToolOutputFile *Out,