[Driver][Frontend] Introduce load-pass-plugin option

Allow dynamic loading of LLVM passes via `load-pass-plugin`
option passed to the Swift compiler driver.

include/swift/AST/DiagnosticsIRGen.def

@@ -65,5 +65,8 @@ ERROR(temporary_allocation_alignment_not_power_of_2,none,
 ERROR(explosion_size_oveflow,none,
       "explosion size too large", ())
 
+ERROR(unable_to_load_pass_plugin,none,
+      "unable to load plugin '%0': '%1'", (StringRef, StringRef))
+
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/AST/IRGenOptions.h

@@ -511,6 +511,9 @@ class IRGenOptions {
   /// Emit a .casid file next to the object file if CAS Backend is used.
   bool EmitCASIDFile;
 
+  /// Paths to the pass plugins registered via -load-pass-plugin.
+  std::vector<std::string> LLVMPassPlugins;
+
   IRGenOptions()
       : DWARFVersion(2),
         OutputKind(IRGenOutputKind::LLVMAssemblyAfterOptimization),

include/swift/Option/Options.td

@@ -1701,6 +1701,11 @@ def use_interface_for_module: Separate<["-", "--"], "use-interface-for-module">,
   HelpText<"Prefer loading these modules via interface">,
   MetaVarName<"<name>">;
 
+def load_pass_plugin_EQ : Joined<["-"], "load-pass-plugin=">,
+  Flags<[FrontendOption, ArgumentIsPath]>,
+  HelpText<"Load LLVM pass plugin from a dynamic shared object file.">,
+  MetaVarName<"<path>">;
+
 // ONLY SUPPORTED IN NEW DRIVER
 
 // These flags only exist here so that the old driver doesn't fail with unknown

include/swift/Subsystems.h

@@ -252,7 +252,10 @@ namespace swift {
   /// Given an already created LLVM module, construct a pass pipeline and run
   /// the Swift LLVM Pipeline upon it. This will include the emission of LLVM IR
   /// if requested (\out is not null).
-  void performLLVMOptimizations(const IRGenOptions &Opts, llvm::Module *Module,
+  void performLLVMOptimizations(const IRGenOptions &Opts,
+                                DiagnosticEngine &Diags,
+                                llvm::sys::Mutex *DiagMutex,
+                                llvm::Module *Module,
                                 llvm::TargetMachine *TargetMachine,
                                 llvm::raw_pwrite_stream *out);
 

lib/Driver/ToolChains.cpp

@@ -320,6 +320,7 @@ void ToolChain::addCommonFrontendArgs(const OutputInfo &OI,
   inputArgs.AddLastArg(arguments, options::OPT_enable_experimental_cxx_interop);
   inputArgs.AddLastArg(arguments, options::OPT_cxx_interoperability_mode);
   inputArgs.AddLastArg(arguments, options::OPT_enable_builtin_module);
+  inputArgs.AddLastArg(arguments, options::OPT_load_pass_plugin_EQ);
 
   // Pass on any build config options
   inputArgs.AddAllArgs(arguments, options::OPT_D);

lib/DriverTool/swift_llvm_opt_main.cpp

@@ -215,7 +215,10 @@ int swift_llvm_opt_main(ArrayRef<const char *> argv, void *MainAddr) {
     Opts.OutputKind = IRGenOutputKind::LLVMAssemblyAfterOptimization;
 
     // Then perform the optimizations.
-    performLLVMOptimizations(Opts, M.get(), TM.get(), &Out->os());
+    SourceManager SM;
+    DiagnosticEngine Diags(SM);
+    performLLVMOptimizations(Opts, Diags, nullptr, M.get(), TM.get(),
+                             &Out->os());
   } else {
     std::string Pipeline = PassPipeline;
     llvm::TargetLibraryInfoImpl TLII(ModuleTriple);

lib/Frontend/CompilerInvocation.cpp

@@ -2608,6 +2608,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
     }
   }
 
+  for (const Arg *A : Args.filtered(OPT_load_pass_plugin_EQ)) {
+    Opts.LLVMPassPlugins.push_back(A->getValue());
+  }
+
   for (const Arg *A : Args.filtered(OPT_verify_type_layout)) {
     Opts.VerifyTypeLayoutNames.push_back(A->getValue());
   }

lib/IRGen/IRGen.cpp

@@ -68,6 +68,7 @@
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -194,7 +195,23 @@ static void align(llvm::Module *Module) {
     }
 }
 
+template <typename... ArgTypes>
+void diagnoseSync(
+    DiagnosticEngine &Diags, llvm::sys::Mutex *DiagMutex, SourceLoc Loc,
+    Diag<ArgTypes...> ID,
+    typename swift::detail::PassArgument<ArgTypes>::type... Args) {
+  if (DiagMutex)
+    DiagMutex->lock();
+
+  Diags.diagnose(Loc, ID, std::move(Args)...);
+
+  if (DiagMutex)
+    DiagMutex->unlock();
+}
+
 void swift::performLLVMOptimizations(const IRGenOptions &Opts,
+                                     DiagnosticEngine &Diags,
+                                     llvm::sys::Mutex *DiagMutex,
                                      llvm::Module *Module,
                                      llvm::TargetMachine *TargetMachine,
                                      llvm::raw_pwrite_stream *out) {
@@ -244,6 +261,18 @@ void swift::performLLVMOptimizations(const IRGenOptions &Opts,
 
   PassBuilder PB(TargetMachine, PTO, PGOOpt, &PIC);
 
+  // Attempt to load pass plugins and register their callbacks with PB.
+  for (const auto &PluginFile : Opts.LLVMPassPlugins) {
+    Expected<PassPlugin> PassPlugin = PassPlugin::Load(PluginFile);
+    if (PassPlugin) {
+      PassPlugin->registerPassBuilderCallbacks(PB);
+    } else {
+      diagnoseSync(Diags, DiagMutex, SourceLoc(),
+                   diag::unable_to_load_pass_plugin, PluginFile,
+                   toString(PassPlugin.takeError()));
+    }
+  }
+
   // Register the AA manager first so that our version is the one used.
   FAM.registerPass([&] {
     auto AA = PB.buildDefaultAAPipeline();
@@ -512,20 +541,6 @@ static void countStatsPostIRGen(UnifiedStatsReporter &Stats,
   }
 }
 
-template<typename ...ArgTypes>
-void
-diagnoseSync(DiagnosticEngine &Diags, llvm::sys::Mutex *DiagMutex,
-             SourceLoc Loc, Diag<ArgTypes...> ID,
-             typename swift::detail::PassArgument<ArgTypes>::type... Args) {
-  if (DiagMutex)
-    DiagMutex->lock();
-
-  Diags.diagnose(Loc, ID, std::move(Args)...);
-
-  if (DiagMutex)
-    DiagMutex->unlock();
-}
-
 /// Run the LLVM passes. In multi-threaded compilation this will be done for
 /// multiple LLVM modules in parallel.
 bool swift::performLLVM(const IRGenOptions &Opts,
@@ -594,7 +609,7 @@ bool swift::performLLVM(const IRGenOptions &Opts,
     assert(Opts.OutputKind == IRGenOutputKind::Module && "no output specified");
   }
 
-  performLLVMOptimizations(Opts, Module, TargetMachine,
+  performLLVMOptimizations(Opts, Diags, DiagMutex, Module, TargetMachine,
                            OutputFile ? &OutputFile->getOS() : nullptr);
 
   if (Stats) {
@@ -1697,7 +1712,7 @@ GeneratedModule OptimizedIRRequest::evaluate(Evaluator &evaluator,
   if (!irMod)
     return irMod;
 
-  performLLVMOptimizations(desc.Opts, irMod.getModule(),
+  performLLVMOptimizations(desc.Opts, ctx.Diags, nullptr, irMod.getModule(),
                            irMod.getTargetMachine(), desc.out);
   return irMod;
 }

unittests/Driver/CMakeLists.txt

@@ -1,11 +1,27 @@
-add_swift_unittest(SwiftDriverTests
+set(SWIFT_DRIVER_TESTS_SOURCE_FILES
   FineGrainedDependencyGraphTests.cpp
   MockingFineGrainedDependencyGraphs.cpp
   UnitTestSourceFileDepGraphFactory.cpp
 )
 
+set(SWIFT_DRIVER_TESTS_LIBRARIES
+  swiftAST
+  swiftClangImporter
+  swiftSema
+  swiftDriver
+)
+
+# PassPlugin not supported on Windows
+if(NOT SWIFT_HOST_VARIANT STREQUAL "windows")
+  list(APPEND SWIFT_DRIVER_TESTS_SOURCE_FILES
+    PassPluginTest.cpp)
+  list(APPEND SWIFT_DRIVER_TESTS_LIBRARIES
+    LLVMTestingSupport)
+  add_subdirectory(PassPluginInput)
+endif()
+
+add_swift_unittest(SwiftDriverTests
+  ${SWIFT_DRIVER_TESTS_SOURCE_FILES})
+
 target_link_libraries(SwiftDriverTests PRIVATE
-   swiftAST
-   swiftClangImporter
-   swiftSema
-   swiftDriver)
+  ${SWIFT_DRIVER_TESTS_LIBRARIES})

unittests/Driver/PassPluginInput/CMakeLists.txt

@@ -0,0 +1,13 @@
+add_library(InputTestPlugin SHARED
+  TestPlugin.cpp)
+
+# Put PLUGIN next to the unit test executable.
+set_output_directory(InputTestPlugin
+  BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}/../
+  LIBRARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}/../
+)
+
+llvm_map_components_to_libnames(LLVM_LIBRARIES Core)
+target_link_libraries(InputTestPlugin PRIVATE ${LLVM_LIBRARIES})
+
+set_target_properties(InputTestPlugin PROPERTIES FOLDER "Tests")

unittests/Driver/PassPluginInput/TestPlugin.cpp

@@ -0,0 +1,48 @@
+//===---------------------- TestPlugin.cpp ----------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/PassPlugin.h"
+
+#include "../TestPlugin.h"
+
+using namespace llvm;
+
+namespace {
+
+struct TestPluginPass : PassInfoMixin<TestPluginPass> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    return PreservedAnalyses::all();
+  }
+};
+
+} // namespace
+
+PassPluginLibraryInfo getTestPluginInfo() {
+  return {LLVM_PLUGIN_API_VERSION, TEST_PLUGIN_NAME, TEST_PLUGIN_VERSION,
+          [](PassBuilder &PB) {
+            PB.registerPipelineParsingCallback(
+                [](StringRef Name, ModulePassManager &PM,
+                   ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
+                  if (Name == "plugin-pass") {
+                    PM.addPass(TestPluginPass());
+                    return true;
+                  }
+                  return false;
+                });
+          }};
+}
+
+extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
+llvmGetPassPluginInfo() {
+  return getTestPluginInfo();
+}

unittests/Driver/PassPluginTest.cpp

@@ -0,0 +1,49 @@
+//===-------------------- PassPluginTest.cpp --------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Config/config.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/PassPlugin.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+#include "TestPlugin.h"
+
+using namespace llvm;
+
+static std::string LibPath(const std::string &Name) {
+  const auto &Argvs = testing::internal::GetArgvs();
+  const char *Argv0 = Argvs.size() > 0 ? Argvs[0].c_str() : "PassPluginTest";
+  std::string Path = sys::fs::getMainExecutable(Argv0, nullptr);
+  llvm::SmallString<256> Buf{sys::path::parent_path(Path)};
+  sys::path::append(Buf, (Name + LLVM_PLUGIN_EXT).c_str());
+  return std::string(Buf.str());
+}
+
+TEST(PassPluginTest, LoadPlugin) {
+  auto PluginPath = LibPath("libInputTestPlugin");
+  ASSERT_NE("", PluginPath);
+
+  Expected<PassPlugin> Plugin = PassPlugin::Load(PluginPath);
+  ASSERT_TRUE(!!Plugin) << "Plugin path: " << PluginPath;
+
+  ASSERT_EQ(TEST_PLUGIN_NAME, Plugin->getPluginName());
+  ASSERT_EQ(TEST_PLUGIN_VERSION, Plugin->getPluginVersion());
+
+  PassBuilder PB;
+  ModulePassManager PM;
+  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, "plugin-pass"), Failed());
+
+  Plugin->registerPassBuilderCallbacks(PB);
+  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, "plugin-pass"), Succeeded());
+}

unittests/Driver/TestPlugin.h

@@ -0,0 +1,14 @@
+//===----------------------- TestPlugin.h -----------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#define TEST_PLUGIN_NAME "TestPlugin"
+#define TEST_PLUGIN_VERSION "0.1-unit"