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

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

Author: antoniofrighetto

include/swift/AST/DiagnosticsIRGen.def

@@ -80,5 +80,8 @@ ERROR(attr_objc_implementation_resilient_property_deployment_target, none,
       "target to %0 %2 or store this value in an object or 'any' type",
       (StringRef, const llvm::VersionTuple, const llvm::VersionTuple))
 
+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

@@ -519,6 +519,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()
       : OutputKind(IRGenOutputKind::LLVMAssemblyAfterOptimization),
         Verify(true), OptMode(OptimizationMode::NotSet),

include/swift/Option/Options.td

@@ -1761,6 +1761,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

@@ -266,7 +266,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

@@ -343,6 +343,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

@@ -214,7 +214,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

@@ -2886,6 +2886,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"
@@ -196,7 +197,21 @@ 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) {
+  std::optional<llvm::sys::ScopedLock> Lock;
+  if (DiagMutex)
+    Lock.emplace(*DiagMutex);
+
+  Diags.diagnose(Loc, ID, std::move(Args)...);
+}
+
 void swift::performLLVMOptimizations(const IRGenOptions &Opts,
+                                     DiagnosticEngine &Diags,
+                                     llvm::sys::Mutex *DiagMutex,
                                      llvm::Module *Module,
                                      llvm::TargetMachine *TargetMachine,
                                      llvm::raw_pwrite_stream *out) {
@@ -238,6 +253,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();
@@ -520,20 +547,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,
@@ -602,7 +615,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) {
@@ -1709,7 +1722,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;
 }

test/Frontend/Inputs/TestPlugin.cpp

@@ -0,0 +1,50 @@
+//===---------------------- TestPlugin.cpp --------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+// Used by -load-pass-plugin
+
+#include "llvm/Pass.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/PassPlugin.h"
+
+using namespace llvm;
+
+namespace {
+
+void runTestPlugin(Function &F) {
+  errs() << "TestPlugin: ";
+  errs().write_escaped(F.getName()) << '\n';
+}
+
+struct TestPluginPass : PassInfoMixin<TestPluginPass> {
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &) {
+    runTestPlugin(F);
+    return PreservedAnalyses::all();
+  }
+};
+
+} // namespace
+
+PassPluginLibraryInfo getTestPluginInfo() {
+  return {LLVM_PLUGIN_API_VERSION, "TestPlugin", LLVM_VERSION_STRING,
+          [](PassBuilder &PB) {
+            PB.registerVectorizerStartEPCallback(
+                [](llvm::FunctionPassManager &PM, OptimizationLevel Level) {
+                  PM.addPass(TestPluginPass());
+            });
+          }};
+}
+
+extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
+llvmGetPassPluginInfo() {
+  return getTestPluginInfo();
+}

test/Frontend/load-pass-plugin.swift

@@ -0,0 +1,16 @@
+// REQUIRES: OS=macosx
+
+// RUN: %target-swift-frontend -load-pass-plugin=nonexistent.dylib %s -emit-ir -o /dev/null 2>&1 | %FileCheck -check-prefix=CHECK-UNABLE-LOAD %s
+// CHECK-UNABLE-LOAD: error: unable to load plugin 'nonexistent.dylib': 'Could not load library{{.*}}'
+
+// RUN: %empty-directory(%t)
+// RUN: %target-clangxx %S/Inputs/TestPlugin.cpp -std=c++17 -stdlib=libc++ \
+// RUN:     -isysroot %sdk -I %llvm_src_root/include -I %llvm_obj_root/include -L %llvm_obj_root/lib -lLLVMSupport \
+// RUN:     -Wl,-undefined -Wl,suppress -Wl,-flat_namespace \
+// RUN:     -dynamiclib -o %t/libTestPlugin.dylib
+
+// RUN: %target-swift-frontend -load-pass-plugin=%t/libTestPlugin.dylib %s -emit-ir -o /dev/null 2>&1 | %swift-demangle | %FileCheck %s
+// CHECK: TestPlugin: main
+// CHECK: TestPlugin: null.empty() -> ()
+
+func empty() {}