[RFC][LLVM][Clang] Add LLVM plugin hook for back-ends

[aengelke]

Add a mechanism to permit plugins running code between optimizations and
the back-end passes. Implement this through the LLVM plug-in mechanism
to make permit plugins to be written independently of the front-end.

The primary motivation for this point is TPDE-LLVM, which substitutes
the LLVM back-end (optionally falling back to it for unsupported IR). We
have been distributing a Clang patch; but requiring a custom-build
toolchain is impracticable for many users.

[vgvassilev]

That looks like a reasonable direction to me. Any chance for a test?

[github-actions]

This is another comment for testing the issue write workflow that was placed in a separate file

[github-actions]

This is a comment for testing the issue write workflow

[aengelke]

Added plugin support to llc and test the plugin behavior through that.

[github-actions]

🐧 Linux x64 Test Results

• 193594 tests passed
• 6242 tests skipped

✅ The build succeeded and all tests passed.

[efriedma-quic]

This looks like what I was expecting; the architecture seems fine.

[llvmbot]

@llvm/pr-subscribers-clang-codegen

Author: Alexis Engelke (aengelke)

Changes

Add a mechanism to permit plugins running code between optimizations and
the back-end passes. Implement this through the LLVM plug-in mechanism
to make permit plugins to be written independently of the front-end.

The primary motivation for this point is TPDE-LLVM, which substitutes
the LLVM back-end (optionally falling back to it for unsupported IR). We
have been distributing a Clang patch; but requiring a custom-build
toolchain is impracticable for many users.

---

Full diff: https://github.com/llvm/llvm-project/pull/170846.diff

6 Files Affected:

• (modified) clang/lib/CodeGen/BackendUtil.cpp (+23-10)
• (modified) llvm/examples/Bye/Bye.cpp (+35-15)
• (modified) llvm/include/llvm/Passes/PassPlugin.h (+34-5)
• (modified) llvm/lib/Passes/PassPlugin.cpp (-5)
• (added) llvm/test/Feature/codegen-plugin.ll (+18)
• (modified) llvm/tools/llc/llc.cpp (+28-5)

diff --git a/clang/lib/CodeGen/BackendUtil.cpp b/clang/lib/CodeGen/BackendUtil.cpp
index af3480d5755f1..451c6c990bf40 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -148,6 +148,7 @@ class EmitAssemblyHelper {
   const LangOptions &LangOpts;
   llvm::Module *TheModule;
   IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS;
+  llvm::SmallVector<llvm::PassPlugin> Plugins;
 
   std::unique_ptr<raw_pwrite_stream> OS;
 
@@ -1017,16 +1018,9 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
     }
 #endif
   }
-  // Attempt to load pass plugins and register their callbacks with PB.
-  for (auto &PluginFN : CodeGenOpts.PassPlugins) {
-    auto PassPlugin = PassPlugin::Load(PluginFN);
-    if (PassPlugin) {
-      PassPlugin->registerPassBuilderCallbacks(PB);
-    } else {
-      Diags.Report(diag::err_fe_unable_to_load_plugin)
-          << PluginFN << toString(PassPlugin.takeError());
-    }
-  }
+  // Register plugin callbacks with PB.
+  for (auto &Plugin : Plugins)
+    Plugin.registerPassBuilderCallbacks(PB);
   for (const auto &PassCallback : CodeGenOpts.PassBuilderCallbacks)
     PassCallback(PB);
 #define HANDLE_EXTENSION(Ext)                                                  \
@@ -1265,6 +1259,14 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
 void EmitAssemblyHelper::RunCodegenPipeline(
     BackendAction Action, std::unique_ptr<raw_pwrite_stream> &OS,
     std::unique_ptr<llvm::ToolOutputFile> &DwoOS) {
+  // Invoke pre-codegen callback from plugin, which might want to take over the
+  // entire code generation itself.
+  for (auto &Plugin : Plugins) {
+    CodeGenFileType CGFT = getCodeGenFileType(Action);
+    if (Plugin.invokePreCodeGenCallback(*TheModule, *TM, CGFT, *OS))
+      return;
+  }
+
   // We still use the legacy PM to run the codegen pipeline since the new PM
   // does not work with the codegen pipeline.
   // FIXME: make the new PM work with the codegen pipeline.
@@ -1328,6 +1330,17 @@ void EmitAssemblyHelper::emitAssembly(BackendAction Action,
   // Before executing passes, print the final values of the LLVM options.
   cl::PrintOptionValues();
 
+  // Attempt to load pass plugins.
+  for (auto &PluginFN : CodeGenOpts.PassPlugins) {
+    auto PassPlugin = PassPlugin::Load(PluginFN);
+    if (PassPlugin) {
+      Plugins.push_back(std::move(*PassPlugin));
+    } else {
+      Diags.Report(diag::err_fe_unable_to_load_plugin)
+          << PluginFN << toString(PassPlugin.takeError());
+    }
+  }
+
   std::unique_ptr<llvm::ToolOutputFile> ThinLinkOS, DwoOS;
   RunOptimizationPipeline(Action, OS, ThinLinkOS, BC);
   RunCodegenPipeline(Action, OS, DwoOS);
diff --git a/llvm/examples/Bye/Bye.cpp b/llvm/examples/Bye/Bye.cpp
index d88bf9e490e9c..4d612e2350a01 100644
--- a/llvm/examples/Bye/Bye.cpp
+++ b/llvm/examples/Bye/Bye.cpp
@@ -11,6 +11,9 @@ using namespace llvm;
 static cl::opt<bool> Wave("wave-goodbye", cl::init(false),
                           cl::desc("wave good bye"));
 
+static cl::opt<bool> LastWords("last-words", cl::init(false),
+                               cl::desc("say last words (suppress codegen)"));
+
 namespace {
 
 bool runBye(Function &F) {
@@ -35,6 +38,37 @@ struct Bye : PassInfoMixin<Bye> {
   }
 };
 
+void registerPassBuilderCallbacks(PassBuilder &PB) {
+  PB.registerVectorizerStartEPCallback(
+      [](llvm::FunctionPassManager &PM, 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;
+      });
+}
+
+bool preCodeGenCallback(Module &M, TargetMachine &, CodeGenFileType CGFT,
+                        raw_pwrite_stream &OS) {
+  if (LastWords) {
+    if (CGFT != CodeGenFileType::AssemblyFile) {
+      // Test error emission.
+      M.getContext().emitError("last words unsupported for binary output");
+      return false;
+    }
+    OS << "CodeGen Bye\n";
+    return true; // Suppress remaining compilation pipeline.
+  }
+  // Do nothing.
+  return false;
+}
+
 } // namespace
 
 char LegacyBye::ID = 0;
@@ -46,21 +80,7 @@ static RegisterPass<LegacyBye> X("goodbye", "Good Bye World Pass",
 /* New PM Registration */
 llvm::PassPluginLibraryInfo getByePluginInfo() {
   return {LLVM_PLUGIN_API_VERSION, "Bye", LLVM_VERSION_STRING,
-          [](PassBuilder &PB) {
-            PB.registerVectorizerStartEPCallback(
-                [](llvm::FunctionPassManager &PM, 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;
-                });
-          }};
+          registerPassBuilderCallbacks, preCodeGenCallback};
 }
 
 #ifndef LLVM_BYE_LINK_INTO_TOOLS
diff --git a/llvm/include/llvm/Passes/PassPlugin.h b/llvm/include/llvm/Passes/PassPlugin.h
index 947504bc207a7..c1840b0fabfdb 100644
--- a/llvm/include/llvm/Passes/PassPlugin.h
+++ b/llvm/include/llvm/Passes/PassPlugin.h
@@ -14,6 +14,7 @@
 #define LLVM_PASSES_PASSPLUGIN_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/CodeGen.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/Error.h"
@@ -21,7 +22,9 @@
 #include <string>
 
 namespace llvm {
+class Module;
 class PassBuilder;
+class TargetMachine;
 
 /// \macro LLVM_PLUGIN_API_VERSION
 /// Identifies the API version understood by this plugin.
@@ -30,14 +33,15 @@ class PassBuilder;
 /// against that of the plugin. A mismatch is an error. The supported version
 /// will be incremented for ABI-breaking changes to the \c PassPluginLibraryInfo
 /// struct, i.e. when callbacks are added, removed, or reordered.
-#define LLVM_PLUGIN_API_VERSION 1
+#define LLVM_PLUGIN_API_VERSION 2
 
 extern "C" {
 /// Information about the plugin required to load its passes
 ///
 /// This struct defines the core interface for pass plugins and is supposed to
-/// be filled out by plugin implementors. LLVM-side users of a plugin are
-/// expected to use the \c PassPlugin class below to interface with it.
+/// be filled out by plugin implementors. Unused function pointers can be set to
+/// nullptr. LLVM-side users of a plugin are expected to use the \c PassPlugin
+/// class below to interface with it.
 struct PassPluginLibraryInfo {
   /// The API version understood by this plugin, usually \c
   /// LLVM_PLUGIN_API_VERSION
@@ -49,7 +53,14 @@ struct PassPluginLibraryInfo {
 
   /// The callback for registering plugin passes with a \c PassBuilder
   /// instance
-  void (*RegisterPassBuilderCallbacks)(PassBuilder &);
+  void (*RegisterPassBuilderCallbacks)(PassBuilder &) = nullptr;
+
+  /// Callback called before running the back-end passes on the module. The
+  /// callback can generate code itself by writing the expected output to OS and
+  /// returning true to prevent the default pipeline and further plugin
+  /// callbacks from running.
+  bool (*PreCodeGenCallback)(Module &, TargetMachine &, CodeGenFileType,
+                             raw_pwrite_stream &OS) = nullptr;
 };
 }
 
@@ -80,7 +91,17 @@ class PassPlugin {
 
   /// Invoke the PassBuilder callback registration
   void registerPassBuilderCallbacks(PassBuilder &PB) const {
-    Info.RegisterPassBuilderCallbacks(PB);
+    if (Info.RegisterPassBuilderCallbacks)
+      Info.RegisterPassBuilderCallbacks(PB);
+  }
+
+  /// Invoke the pre-codegen callback.
+  bool invokePreCodeGenCallback(Module &M, TargetMachine &TM,
+                                CodeGenFileType CGFT,
+                                raw_pwrite_stream &OS) const {
+    if (Info.PreCodeGenCallback)
+      return Info.PreCodeGenCallback(M, TM, CGFT, OS);
+    return false;
   }
 
 private:
@@ -93,6 +114,11 @@ class PassPlugin {
 };
 }
 
+// The function returns a struct with default initializers.
+#ifdef __clang__
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wreturn-type-c-linkage"
+#endif
 /// The public entry point for a pass plugin.
 ///
 /// When a plugin is loaded by the driver, it will call this entry point to
@@ -109,5 +135,8 @@ class PassPlugin {
 /// ```
 extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK
 llvmGetPassPluginInfo();
+#ifdef __clang__
+#pragma clang diagnostic pop
+#endif
 
 #endif /* LLVM_PASSES_PASSPLUGIN_H */
diff --git a/llvm/lib/Passes/PassPlugin.cpp b/llvm/lib/Passes/PassPlugin.cpp
index 6182cbbb1fddd..201f5eef080c3 100644
--- a/llvm/lib/Passes/PassPlugin.cpp
+++ b/llvm/lib/Passes/PassPlugin.cpp
@@ -45,10 +45,5 @@ Expected<PassPlugin> PassPlugin::Load(const std::string &Filename) {
             Twine(LLVM_PLUGIN_API_VERSION) + ".",
         inconvertibleErrorCode());
 
-  if (!P.Info.RegisterPassBuilderCallbacks)
-    return make_error<StringError>(Twine("Empty entry callback in plugin '") +
-                                       Filename + "'.'",
-                                   inconvertibleErrorCode());
-
   return P;
 }
diff --git a/llvm/test/Feature/codegen-plugin.ll b/llvm/test/Feature/codegen-plugin.ll
new file mode 100644
index 0000000000000..2c6a4d5ac9bb4
--- /dev/null
+++ b/llvm/test/Feature/codegen-plugin.ll
@@ -0,0 +1,18 @@
+; REQUIRES: x86-registered-target
+; RUN: llc < %s %loadnewpmbye | FileCheck %s --check-prefix=CHECK-ASM
+; RUN: llc < %s %loadnewpmbye -last-words | FileCheck %s --check-prefix=CHECK-ACTIVE
+; RUN: not llc %s %loadnewpmbye -last-words -filetype=obj 2>&1 | FileCheck %s --check-prefix=CHECK-ERR
+; REQUIRES: plugins, examples
+; UNSUPPORTED: target={{.*windows.*}}
+; CHECK-ASM: somefunk:
+; CHECK-ACTIVE: CodeGen Bye
+; CHECK-ERR: error: last words unsupported for binary output
+
+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 ptr @somefunk() {
+  ret ptr @junk
+}
+
diff --git a/llvm/tools/llc/llc.cpp b/llvm/tools/llc/llc.cpp
index 613780ecbfb40..9f5bec2eeae62 100644
--- a/llvm/tools/llc/llc.cpp
+++ b/llvm/tools/llc/llc.cpp
@@ -40,6 +40,7 @@
 #include "llvm/MC/MCTargetOptionsCommandFlags.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Pass.h"
+#include "llvm/Passes/PassPlugin.h"
 #include "llvm/Remarks/HotnessThresholdParser.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -213,6 +214,9 @@ static cl::opt<std::string> RemarksFormat(
     cl::desc("The format used for serializing remarks (default: YAML)"),
     cl::value_desc("format"), cl::init("yaml"));
 
+static cl::list<std::string> PassPlugins("load-pass-plugin",
+                                         cl::desc("Load plugin library"));
+
 static cl::opt<bool> EnableNewPassManager(
     "enable-new-pm", cl::desc("Enable the new pass manager"), cl::init(false));
 
@@ -286,8 +290,8 @@ static void setPGOOptions(TargetMachine &TM) {
     TM.setPGOOption(PGOOpt);
 }
 
-static int compileModule(char **argv, LLVMContext &Context,
-                         std::string &OutputFilename);
+static int compileModule(char **argv, SmallVectorImpl<PassPlugin> &,
+                         LLVMContext &Context, std::string &OutputFilename);
 
 [[noreturn]] static void reportError(Twine Msg, StringRef Filename = "") {
   SmallString<256> Prefix;
@@ -396,6 +400,14 @@ int main(int argc, char **argv) {
   // Initialize debugging passes.
   initializeScavengerTestPass(*Registry);
 
+  SmallVector<PassPlugin, 1> PluginList;
+  PassPlugins.setCallback([&](const std::string &PluginPath) {
+    auto Plugin = PassPlugin::Load(PluginPath);
+    if (!Plugin)
+      reportFatalUsageError(Plugin.takeError());
+    PluginList.emplace_back(Plugin.get());
+  });
+
   // Register the Target and CPU printer for --version.
   cl::AddExtraVersionPrinter(sys::printDefaultTargetAndDetectedCPU);
   // Register the target printer for --version.
@@ -447,7 +459,7 @@ int main(int argc, char **argv) {
   // Compile the module TimeCompilations times to give better compile time
   // metrics.
   for (unsigned I = TimeCompilations; I; --I)
-    if (int RetVal = compileModule(argv, Context, OutputFilename))
+    if (int RetVal = compileModule(argv, PluginList, Context, OutputFilename))
       return RetVal;
 
   if (RemarksFile)
@@ -485,8 +497,8 @@ static bool addPass(PassManagerBase &PM, const char *argv0, StringRef PassName,
   return false;
 }
 
-static int compileModule(char **argv, LLVMContext &Context,
-                         std::string &OutputFilename) {
+static int compileModule(char **argv, SmallVectorImpl<PassPlugin> &PluginList,
+                         LLVMContext &Context, std::string &OutputFilename) {
   // Load the module to be compiled...
   SMDiagnostic Err;
   std::unique_ptr<Module> M;
@@ -707,6 +719,17 @@ static int compileModule(char **argv, LLVMContext &Context,
   // flags.
   codegen::setFunctionAttributes(CPUStr, FeaturesStr, *M);
 
+  for (auto &Plugin : PluginList) {
+    CodeGenFileType CGFT = codegen::getFileType();
+    if (Plugin.invokePreCodeGenCallback(*M, *Target, CGFT, Out->os())) {
+      // TODO: Deduplicate code with below and the NewPMDriver.
+      if (Context.getDiagHandlerPtr()->HasErrors)
+        exit(1);
+      Out->keep();
+      return 0;
+    }
+  }
+
   if (mc::getExplicitRelaxAll() &&
       codegen::getFileType() != CodeGenFileType::ObjectFile)
     WithColor::warning(errs(), argv[0])

[vgvassilev]

LGTM! Thank you, @aengelke!

[weliveindetail]

Please note that this is the first-ever change in the pass-plugin API. This isn't a breaking change, but we might need them if we need to fix things down the line. This change will also affect lld as well as out-of-tree frontends like Rust and Swift.

[weliveindetail]

I like the idea to have pre/post opt/codegen callbacks. If we ever get the new pass manager in the Codegen pipeline, that would help plugins to handle the pipeline change. Could we generalize this into something like registerPipelinePhaseCallbacks()?

[aengelke]

I would consider this to be out-of-scope for now. Back-end pipelines are, currently, target-specific and any support for plugins modifying these pipelines is likely to need a lot of further discussion (esp. as many back-end passes not really self-contained). There was some discussion on CodeGenPassBuilder to add callbacks for targets earlier this year, but it seems work on that has stalled. Additionally, the proposed plugin API already permits building and running a custom pass manager, so right now there would be no benefit of speculatively building a more generic API.

If there's a desire to add external hooks for CodeGenPassBuilder, I think this should be a separate function similar to the RegisterPassBuilderCallbacks that exists so far.

[weliveindetail]

What about the other places in-tree that load plugins: Flang, clang-linker-wrapper, libLTO, opt and lld? Do they need adjustment?

[aengelke]

Need no, they continue to work just fine.

• LTO: yes, might be good to have at some point in the future.
• Flang: likewise.
• clang-linker-wrapper: if there's interest, this can be added. While it currently seems to load plugins, it doesn't seem to call any methods on them.
• MLIR ModuleToObject::translateToISA also only seems to be used for GPUs and MLIR doesn't do anything related to plugins right now.

opt doesn't call and lld doesn't directly call back-ends.