[Clang] Add codegen option to add passbuilder callback functions (#70171)

Author: wsmoses

clang/examples/CMakeLists.txt

@@ -4,6 +4,7 @@ if(NOT CLANG_BUILD_EXAMPLES)
 endif()
 
 if(CLANG_PLUGIN_SUPPORT)
+  add_subdirectory(LLVMPrintFunctionNames)
   add_subdirectory(PrintFunctionNames)
   add_subdirectory(AnnotateFunctions)
   add_subdirectory(Attribute)

clang/examples/LLVMPrintFunctionNames/CMakeLists.txt

@@ -0,0 +1,23 @@
+# If we don't need RTTI or EH, there's no reason to export anything
+# from the plugin.
+if(NOT MSVC) # MSVC mangles symbols differently, and
+    # PrintLLVMFunctionNames.export contains C++ symbols.
+  if(NOT LLVM_REQUIRES_RTTI)
+    if(NOT LLVM_REQUIRES_EH)
+      set(LLVM_EXPORTED_SYMBOL_FILE ${CMAKE_CURRENT_SOURCE_DIR}/LLVMPrintFunctionNames.exports)
+    endif()
+  endif()
+endif()
+
+add_llvm_library(LLVMPrintFunctionNames MODULE LLVMPrintFunctionNames.cpp PLUGIN_TOOL clang)
+
+if(WIN32 OR CYGWIN)
+  set(LLVM_LINK_COMPONENTS
+    Support
+  )
+  clang_target_link_libraries(LLVMPrintFunctionNames PRIVATE
+    clangAST
+    clangBasic
+    clangFrontend
+    )
+endif()

clang/examples/LLVMPrintFunctionNames/LLVMPrintFunctionNames.cpp

@@ -0,0 +1,78 @@
+//===- LLVMPrintFunctionNames.cpp
+//---------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Example clang plugin which simply prints the names of all the functions
+// within the generated LLVM code.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/AST/AST.h"
+#include "clang/AST/ASTConsumer.h"
+#include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/FrontendPluginRegistry.h"
+#include "clang/Sema/Sema.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Passes/OptimizationLevel.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/raw_ostream.h"
+using namespace clang;
+
+namespace {
+
+class PrintPass final : public llvm::AnalysisInfoMixin<PrintPass> {
+  friend struct llvm::AnalysisInfoMixin<PrintPass>;
+
+private:
+  static llvm::AnalysisKey key;
+
+public:
+  using Result = llvm::PreservedAnalyses;
+
+  Result run(llvm::Module &M, llvm::ModuleAnalysisManager &MAM) {
+    for (auto &F : M)
+      llvm::outs() << "[PrintPass] Found function: " << F.getName() << "\n";
+    return llvm::PreservedAnalyses::all();
+  }
+  static bool isRequired() { return true; }
+};
+
+void PrintCallback(llvm::PassBuilder &PB) {
+  PB.registerPipelineStartEPCallback(
+      [](llvm::ModulePassManager &MPM, llvm::OptimizationLevel) {
+        MPM.addPass(PrintPass());
+      });
+}
+
+class LLVMPrintFunctionsConsumer : public ASTConsumer {
+public:
+  LLVMPrintFunctionsConsumer(CompilerInstance &Instance) {
+    Instance.getCodeGenOpts().PassBuilderCallbacks.push_back(PrintCallback);
+  }
+};
+
+class LLVMPrintFunctionNamesAction : public PluginASTAction {
+protected:
+  std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
+                                                 llvm::StringRef) override {
+    return std::make_unique<LLVMPrintFunctionsConsumer>(CI);
+  }
+  bool ParseArgs(const CompilerInstance &,
+                 const std::vector<std::string> &) override {
+    return true;
+  }
+  PluginASTAction::ActionType getActionType() override {
+    return AddBeforeMainAction;
+  }
+};
+
+} // namespace
+
+static FrontendPluginRegistry::Add<LLVMPrintFunctionNamesAction>
+    X("llvm-print-fns", "print function names, llvm level");

clang/examples/LLVMPrintFunctionNames/LLVMPrintFunctionNames.exports

@@ -26,6 +26,9 @@
 #include <string>
 #include <vector>
 
+namespace llvm {
+class PassBuilder;
+}
 namespace clang {
 
 /// Bitfields of CodeGenOptions, split out from CodeGenOptions to ensure
@@ -408,6 +411,9 @@ class CodeGenOptions : public CodeGenOptionsBase {
   /// List of dynamic shared object files to be loaded as pass plugins.
   std::vector<std::string> PassPlugins;
 
+  /// List of pass builder callbacks.
+  std::vector<std::function<void(llvm::PassBuilder &)>> PassBuilderCallbacks;
+
   /// Path to allowlist file specifying which objects
   /// (files, functions) should exclusively be instrumented
   /// by sanitizer coverage pass.

clang/include/clang/Basic/CodeGenOptions.h

@@ -908,6 +908,8 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
           << PluginFN << toString(PassPlugin.takeError());
     }
   }
+  for (auto PassCallback : CodeGenOpts.PassBuilderCallbacks)
+    PassCallback(PB);
 #define HANDLE_EXTENSION(Ext)                                                  \
   get##Ext##PluginInfo().RegisterPassBuilderCallbacks(PB);
 #include "llvm/Support/Extension.def"

clang/lib/CodeGen/BackendUtil.cpp

@@ -100,6 +100,7 @@ if(CLANG_BUILD_EXAMPLES AND CLANG_PLUGIN_SUPPORT)
     CallSuperAttr
     PluginsOrder
     PrintFunctionNames
+    LLVMPrintFunctionNames
     )
 endif ()
 

clang/test/CMakeLists.txt

@@ -0,0 +1,5 @@
+// RUN: %clang_cc1 -load %llvmshlibdir/LLVMPrintFunctionNames%pluginext -S -o /dev/null -emit-llvm %s 2>&1 | FileCheck %s
+// REQUIRES: plugins, examples
+
+// CHECK: [PrintPass] Found function: x 
+int x(int y){ return y; }