[LLVM] Add 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.

Front-end adjustments will follow as separate patches.

[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

• 186984 tests passed
• 4920 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 still think it can be a little surprising, that returning true from a "PreCodeGenCallback" skips the entire backend. But I guess the relevant invariant is: if exit code indicates success, then one input created one output. And that holds, even if the callback didn't write anything. The output file will be empty, but it is updated on disk. So, conceptually this is fine.

The default init of the struct members is not great and it requires the ignore -Wreturn-type-c-linkage, but it prevents an API break for downstream users and that is more important.

It would be nice to see a symmetric "PostCodeGenCallback", but I understand that we have no use-case and we don't introduce things just for symmetry. We can still change that if we ever get NPM in codegen. Pass-plugin support in llc is a nice preparation for it btw.

If this lands, can we make it consistent and roll it out to all plugin users in-tree? (Before the end of the release cycle?) I can imagine there will be a few more edge-case to fix.

[weliveindetail]

Assertion inside fails for a EmitLL action, happens with -flto. There is no test for clang. It would be very easy with https://discourse.llvm.org/t/rfc-a-reference-pass-plugin-in-llvm/89073/12 ;-)

[aengelke]

Thanks for catching this. To keep things simple and manageable, I'm going to split the Clang part into a separate PR with tests.

[aengelke]

I put up #171872 for the Clang part, which now includes tests and in turn is based on #171868 to make plugins reasonably testable.

If there're no further comments or objections on the LLVM part, I'm going to merge this early next week.

[weliveindetail]

LGTM, thanks! It would be great to have documentation for this in the next release. Pass plugin docs are here, but maybe there is a better place since this is isn't related to NPM:
https://llvm.org/docs/WritingAnLLVMNewPMPass.html#registering-passes-as-plugins

[nikic]

Looks reasonable to me.

[mikaelholmen]

I think this RUN line is opening a file codegen-plugin.o in the local directory (which may be write protected) for the output. Can we change it to stdout (or dev/null if we're nott interested in it)?
Maybe just do "< %s" as in the other RUN lines?

[aengelke]

Good point, fixed with f24c0ec. Thanks.

[mikaelholmen]

Good point, fixed with f24c0ec. Thanks.

Nice, thank you!