Enable new passmanager plugin support for LTO.

This should make both static and dynamic NewPM plugins work with LTO.
And as a bonus, it makes static linking of OldPM plugins more reliable
for plugins with both an OldPM and NewPM interface.

I only implemented the command-line flag to specify NewPM plugins in
llvm-lto2, to show it works. Support can be added for other tools later.

Differential Revision: https://reviews.llvm.org/D76866

lld/tools/lld/CMakeLists.txt

@@ -5,7 +5,6 @@ set(LLVM_LINK_COMPONENTS
 add_lld_tool(lld
   lld.cpp
 
-  ENABLE_PLUGINS
   SUPPORT_PLUGINS
   )
 export_executable_symbols_for_plugins(lld)

llvm/examples/Bye/Bye.cpp

@@ -45,7 +45,7 @@ static RegisterPass<LegacyBye> X("goodbye", "Good Bye World Pass",
 
 /* Legacy PM Registration */
 static llvm::RegisterStandardPasses RegisterBye(
-    llvm::PassManagerBuilder::EP_EarlyAsPossible,
+    llvm::PassManagerBuilder::EP_VectorizerStart,
     [](const llvm::PassManagerBuilder &Builder,
        llvm::legacy::PassManagerBase &PM) { PM.add(new LegacyBye()); });
 
@@ -58,6 +58,15 @@ llvm::PassPluginLibraryInfo getByePluginInfo() {
                    llvm::PassBuilder::OptimizationLevel Level) {
                   PM.addPass(Bye());
                 });
+            PB.registerPipelineParsingCallback(
+                [](StringRef Name, llvm::FunctionPassManager &PM,
+                   ArrayRef<llvm::PassBuilder::PipelineElement>) {
+                  if (Name == "goodbye") {
+                    PM.addPass(Bye());
+                    return true;
+                  }
+                  return false;
+                });
           }};
 }
 

llvm/include/llvm/LTO/Config.h

@@ -41,6 +41,7 @@ struct Config {
   std::string CPU;
   TargetOptions Options;
   std::vector<std::string> MAttrs;
+  std::vector<std::string> PassPlugins;
   Optional<Reloc::Model> RelocModel = Reloc::PIC_;
   Optional<CodeModel::Model> CodeModel = None;
   CodeGenOpt::Level CGOptLevel = CodeGenOpt::Default;

llvm/lib/LTO/CMakeLists.txt

@@ -11,6 +11,8 @@ add_llvm_component_library(LLVMLTO
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/LTO
 
+  ENABLE_PLUGINS
+
   DEPENDS
   intrinsics_gen
   llvm_vcsrevision_h

llvm/lib/LTO/LTOBackend.cpp

@@ -28,6 +28,7 @@
 #include "llvm/MC/SubtargetFeature.h"
 #include "llvm/Object/ModuleSymbolTable.h"
 #include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
@@ -127,6 +128,29 @@ Error Config::addSaveTemps(std::string OutputFileName,
   return Error::success();
 }
 
+#define HANDLE_EXTENSION(Ext)                                                  \
+  llvm::PassPluginLibraryInfo get##Ext##PluginInfo();
+#include "llvm/Support/Extension.def"
+
+static void RegisterPassPlugins(ArrayRef<std::string> PassPlugins,
+                                PassBuilder &PB) {
+#define HANDLE_EXTENSION(Ext)                                                  \
+  get##Ext##PluginInfo().RegisterPassBuilderCallbacks(PB);
+#include "llvm/Support/Extension.def"
+
+  // Load requested pass plugins and let them register pass builder callbacks
+  for (auto &PluginFN : PassPlugins) {
+    auto PassPlugin = PassPlugin::Load(PluginFN);
+    if (!PassPlugin) {
+      errs() << "Failed to load passes from '" << PluginFN
+             << "'. Request ignored.\n";
+      continue;
+    }
+
+    PassPlugin->registerPassBuilderCallbacks(PB);
+  }
+}
+
 namespace {
 
 std::unique_ptr<TargetMachine>
@@ -181,6 +205,8 @@ static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
   if (auto Err = PB.parseAAPipeline(AA, "default"))
     report_fatal_error("Error parsing default AA pipeline");
 
+  RegisterPassPlugins(Conf.PassPlugins, PB);
+
   LoopAnalysisManager LAM(Conf.DebugPassManager);
   FunctionAnalysisManager FAM(Conf.DebugPassManager);
   CGSCCAnalysisManager CGAM(Conf.DebugPassManager);
@@ -228,8 +254,8 @@ static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
   // FIXME (davide): verify the output.
 }
 
-static void runNewPMCustomPasses(Module &Mod, TargetMachine *TM,
-                                 std::string PipelineDesc,
+static void runNewPMCustomPasses(const Config &Conf, Module &Mod,
+                                 TargetMachine *TM, std::string PipelineDesc,
                                  std::string AAPipelineDesc,
                                  bool DisableVerify) {
   PassBuilder PB(TM);
@@ -241,6 +267,8 @@ static void runNewPMCustomPasses(Module &Mod, TargetMachine *TM,
       report_fatal_error("unable to parse AA pipeline description '" +
                          AAPipelineDesc + "': " + toString(std::move(Err)));
 
+  RegisterPassPlugins(Conf.PassPlugins, PB);
+
   LoopAnalysisManager LAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
@@ -307,7 +335,7 @@ bool opt(const Config &Conf, TargetMachine *TM, unsigned Task, Module &Mod,
          const ModuleSummaryIndex *ImportSummary) {
   // FIXME: Plumb the combined index into the new pass manager.
   if (!Conf.OptPipeline.empty())
-    runNewPMCustomPasses(Mod, TM, Conf.OptPipeline, Conf.AAPipeline,
+    runNewPMCustomPasses(Conf, Mod, TM, Conf.OptPipeline, Conf.AAPipeline,
                          Conf.DisableVerify);
   else if (Conf.UseNewPM)
     runNewPMPasses(Conf, Mod, TM, Conf.OptLevel, IsThinLTO, ExportSummary,

llvm/test/Feature/load_extension.ll

@@ -1,7 +1,15 @@
+; REQUIRES: x86-registered-target
 ; RUN: opt %s %loadbye -goodbye -wave-goodbye -disable-output 2>&1 | FileCheck %s
+; RUN: opt %s %loadnewpmbye %loadbye -passes="goodbye" -wave-goodbye -disable-output 2>&1 | FileCheck %s
+; RUN: opt -module-summary %s -o %t.o
+; RUN: llvm-lto2 run %t.o %loadbye -wave-goodbye -o %t -r %t.o,somefunk,plx -r %t.o,junk,plx 2>&1 | FileCheck %s
+; RUN: llvm-lto2 run %t.o %loadbye %loadnewpmbye -wave-goodbye -use-new-pm -o %t -r %t.o,somefunk,plx -r %t.o,junk,plx 2>&1 | FileCheck %s
+; RUN: llvm-lto2 run %t.o %loadbye %loadnewpmbye -opt-pipeline="goodbye" -wave-goodbye -use-new-pm -o %t -r %t.o,somefunk,plx -r %t.o,junk,plx 2>&1 | FileCheck %s
 ; REQUIRES: plugins, examples
 ; CHECK: Bye
 
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
 @junk = global i32 0
 
 define i32* @somefunk() {

llvm/test/lit.cfg.py

@@ -203,11 +203,17 @@ def get_asan_rtlib():
 if config.linked_bye_extension:
     config.substitutions.append(('%llvmcheckext', 'CHECK-EXT'))
     config.substitutions.append(('%loadbye', ''))
+    config.substitutions.append(('%loadnewpmbye', ''))
 else:
     config.substitutions.append(('%llvmcheckext', 'CHECK-NOEXT'))
     config.substitutions.append(('%loadbye',
                                  '-load={}/Bye{}'.format(config.llvm_shlib_dir,
-                                                                  config.llvm_shlib_ext)))
+                                                         config.llvm_shlib_ext)))
+    config.substitutions.append(('%loadnewpmbye',
+                                 '-load-pass-plugin={}/Bye{}'
+                                 .format(config.llvm_shlib_dir,
+                                         config.llvm_shlib_ext)))
+
 
 # Static libraries are not built if BUILD_SHARED_LIBS is ON.
 if not config.build_shared_libs and not config.link_llvm_dylib:

llvm/tools/llvm-lto2/CMakeLists.txt

@@ -21,3 +21,4 @@ add_llvm_tool(llvm-lto2
   DEPENDS
   intrinsics_gen
   )
+export_executable_symbols_for_plugins(llvm-lto2)

llvm/tools/llvm-lto2/llvm-lto2.cpp

@@ -20,9 +20,11 @@
 #include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/LTO/Caching.h"
 #include "llvm/LTO/LTO.h"
+#include "llvm/Passes/PassPlugin.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/PluginLoader.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/Threading.h"
 
@@ -142,6 +144,10 @@ static cl::opt<bool>
 static cl::opt<std::string>
     StatsFile("stats-file", cl::desc("Filename to write statistics to"));
 
+static cl::list<std::string>
+    PassPlugins("load-pass-plugin",
+                cl::desc("Load passes from plugin library"));
+
 static void check(Error E, std::string Msg) {
   if (!E)
     return;
@@ -252,6 +258,8 @@ static int run(int argc, char **argv) {
 
   Conf.OptLevel = OptLevel - '0';
   Conf.UseNewPM = UseNewPM;
+  for (auto &PluginFN : PassPlugins)
+    Conf.PassPlugins.push_back(PluginFN);
   switch (CGOptLevel) {
   case '0':
     Conf.CGOptLevel = CodeGenOpt::None;