[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])