[clang] [llvm] [CodeGen] Expose the extensibility of PassConfig to plugins (PR #139059)

[via cfe-commits]

https://github.com/Tcc100 updated https://github.com/llvm/llvm-project/pull/139059

>From 0b9f452fff11853207e0eab6108e47e8c7469295 Mon Sep 17 00:00:00 2001
From: T <T>
Date: Thu, 8 May 2025 11:56:00 +0200
Subject: [PATCH] [CodeGen] Expose the extensibility of PassConfig to plugins

---
 clang/test/CodeGen/passconfighook.cpp         | 56 +++++++++++++++++++
 llvm/docs/WritingAnLLVMPass.rst               | 17 ++++++
 llvm/include/llvm/Target/TargetMachine.h      | 14 +++++
 llvm/lib/CodeGen/CodeGenTargetMachineImpl.cpp |  3 +
 llvm/lib/Target/TargetMachine.cpp             |  3 +
 5 files changed, 93 insertions(+)
 create mode 100644 clang/test/CodeGen/passconfighook.cpp

diff --git a/clang/test/CodeGen/passconfighook.cpp b/clang/test/CodeGen/passconfighook.cpp
new file mode 100644
index 0000000000000..094df123f3de1
--- /dev/null
+++ b/clang/test/CodeGen/passconfighook.cpp
@@ -0,0 +1,56 @@
+// RUN: %clangxx -shared -fPIC -I??/install/include -L%llvmshlibdir %s -o %t.so
+// RUN: %clangxx -O3 -DMAIN -Xclang -load -Xclang %t.so %s -o %t-main | FileCheck %s
+
+#ifndef MAIN
+
+#include <llvm/Target/TargetMachine.h>
+#include <llvm/CodeGen/TargetPassConfig.h>
+#include <llvm/CodeGen/MachineFunctionPass.h>
+#include <llvm/CodeGen/Passes.h>
+
+#define DEBUG_TYPE "codegen-test"
+#define CODEGEN_TEST_NAME "CodeGen Test Pass"
+
+using namespace llvm;
+
+namespace llvm {
+    void initializeCodeGenTestPass(PassRegistry &);
+} // namespace llvm
+
+class CodeGenTest : public MachineFunctionPass {
+public:
+    static char ID;
+
+    CodeGenTest(): MachineFunctionPass(ID) {
+    }
+
+    bool runOnMachineFunction(MachineFunction &MF) override {
+        outs() << "[CodeGen] CodeGenTest::runOnMachineFunction" << "\n";
+        return true;
+    }
+
+    StringRef getPassName() const override {
+        return CODEGEN_TEST_NAME;
+    }
+};
+
+char CodeGenTest::ID = 0;
+INITIALIZE_PASS(CodeGenTest, DEBUG_TYPE, CODEGEN_TEST_NAME, false, false)
+
+__attribute__((constructor)) static void initCodeGenPlugin() {
+    initializeCodeGenTestPass(*PassRegistry::getPassRegistry());
+
+    TargetMachine::registerTargetPassConfigCallback([](auto &TM, auto &PM, auto *TPC) {
+        outs() << "registerTargetPassConfigCallback\n";
+        TPC->insertPass(&GCLoweringID, &CodeGenTest::ID);
+    });
+}
+
+#else
+
+// CHECK: CodeGenTest::runOnMachineFunction
+int main(int argc, char **argv) {
+    return 0;
+}
+
+#endif
diff --git a/llvm/docs/WritingAnLLVMPass.rst b/llvm/docs/WritingAnLLVMPass.rst
index 484227bac38b5..770f5f6acd115 100644
--- a/llvm/docs/WritingAnLLVMPass.rst
+++ b/llvm/docs/WritingAnLLVMPass.rst
@@ -442,6 +442,23 @@ in certain circumstances (such as calling the ``Pass::dump()`` from a
 debugger), so it should only be used to enhance debug output, it should not be
 depended on.
 
+Scheduling a MachineFunctionPass
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Backends create a ``TargetPassConfig`` and use ``addPass`` to schedule
+``MachineFunctionPass``\ es. External plugins can register a callback to modify
+and insert additional passes:
+
+.. code-block:: c++
+
+  TargetMachine::registerTargetPassConfigCallback(
+    [](TargetMachine &TM, PassManager &PM, TargetPassConfig *TPC) {
+      TPC->insertPass(/* ... */);
+      TPC->substitutePass(/* ... */);
+    }
+  );
+
+
 .. _writing-an-llvm-pass-interaction:
 
 Specifying interactions between passes
diff --git a/llvm/include/llvm/Target/TargetMachine.h b/llvm/include/llvm/Target/TargetMachine.h
index 906926729ed74..bcc1ce29b8282 100644
--- a/llvm/include/llvm/Target/TargetMachine.h
+++ b/llvm/include/llvm/Target/TargetMachine.h
@@ -20,10 +20,12 @@
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/PGOOptions.h"
 #include "llvm/Target/CGPassBuilderOption.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/TargetParser/Triple.h"
+#include <functional>
 #include <optional>
 #include <string>
 #include <utility>
@@ -72,6 +74,10 @@ namespace yaml {
 struct MachineFunctionInfo;
 } // namespace yaml
 
+class TargetMachine;
+using PassConfigCallback =
+    std::function<void(TargetMachine &, PassManagerBase &, TargetPassConfig *)>;
+
 //===----------------------------------------------------------------------===//
 ///
 /// Primary interface to the complete machine description for the target
@@ -119,6 +125,9 @@ class TargetMachine {
   std::optional<PGOOptions> PGOOption;
 
 public:
+  static ManagedStatic<SmallVector<PassConfigCallback, 1>>
+      TargetPassConfigCallbacks;
+
   mutable TargetOptions Options;
 
   TargetMachine(const TargetMachine &) = delete;
@@ -518,6 +527,11 @@ class TargetMachine {
 
   // MachineRegisterInfo callback function
   virtual void registerMachineRegisterInfoCallback(MachineFunction &MF) const {}
+
+  // TargetPassConfig callback function
+  static void registerTargetPassConfigCallback(const PassConfigCallback &C) {
+    TargetPassConfigCallbacks->push_back(C);
+  }
 };
 
 } // end namespace llvm
diff --git a/llvm/lib/CodeGen/CodeGenTargetMachineImpl.cpp b/llvm/lib/CodeGen/CodeGenTargetMachineImpl.cpp
index 4a3503a2da7db..336f1db776036 100644
--- a/llvm/lib/CodeGen/CodeGenTargetMachineImpl.cpp
+++ b/llvm/lib/CodeGen/CodeGenTargetMachineImpl.cpp
@@ -119,6 +119,9 @@ addPassesToGenerateCode(CodeGenTargetMachineImpl &TM, PassManagerBase &PM,
   PM.add(PassConfig);
   PM.add(&MMIWP);
 
+  for (auto& C : *TargetMachine::TargetPassConfigCallbacks)
+    C(TM, PM, PassConfig);
+
   if (PassConfig->addISelPasses())
     return nullptr;
   PassConfig->addMachinePasses();
diff --git a/llvm/lib/Target/TargetMachine.cpp b/llvm/lib/Target/TargetMachine.cpp
index 69b6e26e602f6..c43e2ba00f733 100644
--- a/llvm/lib/Target/TargetMachine.cpp
+++ b/llvm/lib/Target/TargetMachine.cpp
@@ -332,3 +332,6 @@ std::pair<int, int> TargetMachine::parseBinutilsVersion(StringRef Version) {
     Version.consumeInteger(10, Ret.second);
   return Ret;
 }
+
+// TargetPassConfig callbacks
+ManagedStatic<SmallVector<PassConfigCallback, 1>> TargetMachine::TargetPassConfigCallbacks{};