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

Tcc100 · web-flow · commit 56a8655f4a9c · 2025-07-10T12:43:09.000+02:00
This PR exposes the backend pass config to plugins via a callback.
Plugin authors can register a callback that is being triggered before
the target backend adds their passes to the pipeline. In the callback
they then get access to the `TargetMachine`, the `PassManager`, and the
`TargetPassConfig`. This allows plugins to call
`TargetPassConfig::insertPass`, which is honored in the subsequent
`addPass` of the main backend. We implemented this using the legacy pass
manager since backends still use it as the default.
diff --git a/llvm/docs/WritingAnLLVMPass.rst b/llvm/docs/WritingAnLLVMPass.rst
@@ -442,6 +442,28 @@ 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++
+
+  RegisterTargetPassConfigCallback X{[](auto &TM, auto &PM, auto *TPC) {
+    TPC->insertPass(/* ... */);
+    TPC->substitutePass(/* ... */);
+  }};
+
+Note that passes still have to be registered:
+
+.. code-block:: c++
+
+  __attribute__((constructor)) static void initCodeGenPlugin() {
+    initializeExamplePass(*PassRegistry::getPassRegistry());
+  }
+
 .. _writing-an-llvm-pass-interaction:
 
 Specifying interactions between passes
diff --git a/llvm/include/llvm/Target/RegisterTargetPassConfigCallback.h b/llvm/include/llvm/Target/RegisterTargetPassConfigCallback.h
@@ -0,0 +1,37 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// This file contains the registry for PassConfigCallbacks that enable changes
+/// to the TargetPassConfig during the initialization of TargetMachine.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TARGET_REGISTERTARGETPASSCONFIGCALLBACK_H
+#define LLVM_TARGET_REGISTERTARGETPASSCONFIGCALLBACK_H
+
+#include "TargetMachine.h"
+
+namespace llvm {
+
+using PassConfigCallback =
+    std::function<void(TargetMachine &, PassManagerBase &, TargetPassConfig *)>;
+
+class RegisterTargetPassConfigCallback {
+public:
+  PassConfigCallback Callback;
+
+  explicit RegisterTargetPassConfigCallback(PassConfigCallback &&C);
+  ~RegisterTargetPassConfigCallback();
+};
+
+void invokeGlobalTargetPassConfigCallbacks(TargetMachine &TM,
+                                           PassManagerBase &PM,
+                                           TargetPassConfig *PassConfig);
+} // namespace llvm
+
+#endif // LLVM_TARGET_REGISTERTARGETPASSCONFIGCALLBACK_H
diff --git a/llvm/lib/CodeGen/CodeGenTargetMachineImpl.cpp b/llvm/lib/CodeGen/CodeGenTargetMachineImpl.cpp
@@ -30,6 +30,7 @@
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FormattedStream.h"
+#include "llvm/Target/RegisterTargetPassConfigCallback.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 using namespace llvm;
@@ -122,6 +123,7 @@ addPassesToGenerateCode(CodeGenTargetMachineImpl &TM, PassManagerBase &PM,
   PassConfig->setDisableVerify(DisableVerify);
   PM.add(PassConfig);
   PM.add(&MMIWP);
+  invokeGlobalTargetPassConfigCallbacks(TM, PM, PassConfig);
 
   if (PassConfig->addISelPasses())
     return nullptr;
diff --git a/llvm/lib/Target/CMakeLists.txt b/llvm/lib/Target/CMakeLists.txt
@@ -3,6 +3,7 @@ list(APPEND LLVM_COMMON_DEPENDS intrinsics_gen)
 list(APPEND LLVM_TABLEGEN_FLAGS -I ${LLVM_MAIN_SRC_DIR}/lib/Target)
 
 add_llvm_component_library(LLVMTarget
+  RegisterTargetPassConfigCallback.cpp
   Target.cpp
   TargetLoweringObjectFile.cpp
   TargetMachine.cpp
diff --git a/llvm/lib/Target/RegisterTargetPassConfigCallback.cpp b/llvm/lib/Target/RegisterTargetPassConfigCallback.cpp
@@ -0,0 +1,39 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// This file contains the registry for PassConfigCallbacks that enable changes
+/// to the TargetPassConfig during the initialization of TargetMachine.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Target/RegisterTargetPassConfigCallback.h"
+
+namespace llvm {
+// TargetPassConfig callbacks
+static SmallVector<RegisterTargetPassConfigCallback *, 1>
+    TargetPassConfigCallbacks{};
+
+void invokeGlobalTargetPassConfigCallbacks(TargetMachine &TM,
+                                           PassManagerBase &PM,
+                                           TargetPassConfig *PassConfig) {
+  for (const RegisterTargetPassConfigCallback *Reg : TargetPassConfigCallbacks)
+    Reg->Callback(TM, PM, PassConfig);
+}
+
+RegisterTargetPassConfigCallback::RegisterTargetPassConfigCallback(
+    PassConfigCallback &&C)
+    : Callback(std::move(C)) {
+  TargetPassConfigCallbacks.push_back(this);
+}
+
+RegisterTargetPassConfigCallback::~RegisterTargetPassConfigCallback() {
+  const auto &It = find(TargetPassConfigCallbacks, this);
+  if (It != TargetPassConfigCallbacks.end())
+    TargetPassConfigCallbacks.erase(It);
+}
+} // namespace llvm
diff --git a/llvm/test/Other/codegen-plugin-loading.ll b/llvm/test/Other/codegen-plugin-loading.ll
@@ -0,0 +1,7 @@
+; RUN: llc -load %llvmshlibdir/../unittests/CodeGen/CGPluginTest/CGTestPlugin%pluginext %s -o - | FileCheck %s
+; REQUIRES: native, system-linux, llvm-dylib
+
+; CHECK: CodeGen Test Pass running on main
+define void @main() {
+  ret void
+}
diff --git a/llvm/unittests/CodeGen/CGPluginTest/CMakeLists.txt b/llvm/unittests/CodeGen/CGPluginTest/CMakeLists.txt
@@ -0,0 +1,18 @@
+set(LLVM_LINK_COMPONENTS
+  ${LLVM_TARGETS_TO_BUILD}
+  Support
+  Core
+  Target
+  CodeGen
+  )
+
+add_llvm_unittest(CGPluginTest
+  PluginTest.cpp
+  Plugin/CodeGenTestPass.cpp
+
+  EXPORT_SYMBOLS
+  )
+target_link_libraries(CGPluginTest PRIVATE LLVMTestingSupport)
+
+unset(LLVM_LINK_COMPONENTS)
+add_subdirectory(Plugin)
diff --git a/llvm/unittests/CodeGen/CGPluginTest/Plugin/CMakeLists.txt b/llvm/unittests/CodeGen/CGPluginTest/Plugin/CMakeLists.txt
@@ -0,0 +1,17 @@
+if (NOT WIN32 AND NOT CYGWIN)
+  unset(LLVM_LINK_COMPONENTS)
+  add_llvm_library(CGTestPlugin MODULE BUILDTREE_ONLY
+    CodeGenTestPass.cpp
+    Plugin.cpp
+    )
+
+  # Put PLUGIN next to the unit test executable.
+  set_output_directory(CGTestPlugin
+    BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}/../
+    LIBRARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}/../
+    )
+  set_target_properties(CGTestPlugin PROPERTIES FOLDER "Tests")
+
+  add_dependencies(CGTestPlugin intrinsics_gen vt_gen)
+  add_dependencies(CGPluginTest CGTestPlugin)
+endif ()
diff --git a/llvm/unittests/CodeGen/CGPluginTest/Plugin/CodeGenTestPass.cpp b/llvm/unittests/CodeGen/CGPluginTest/Plugin/CodeGenTestPass.cpp
@@ -0,0 +1,34 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "CodeGenTestPass.h"
+
+#include "llvm/CodeGen/MachineFrameInfo.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "codegen-test"
+#define CODEGEN_TEST_NAME "CodeGen Test Pass"
+
+CodeGenTest::CodeGenTest() : MachineFunctionPass(ID) {}
+
+bool CodeGenTest::runOnMachineFunction(MachineFunction &MF) {
+  outs() << CODEGEN_TEST_NAME << " running on " << MF.getName()
+         << "\n"; // used for the lit test
+  if (RunCallback)
+    RunCallback();
+  return true;
+}
+
+StringRef CodeGenTest::getPassName() const { return CODEGEN_TEST_NAME; }
+
+char CodeGenTest::ID = 0;
+std::function<void()> CodeGenTest::RunCallback;
+
+INITIALIZE_PASS(CodeGenTest, DEBUG_TYPE, CODEGEN_TEST_NAME, false, false)
diff --git a/llvm/unittests/CodeGen/CGPluginTest/Plugin/CodeGenTestPass.h b/llvm/unittests/CodeGen/CGPluginTest/Plugin/CodeGenTestPass.h
@@ -0,0 +1,33 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CODEGEN_TEST_PASS
+#define CODEGEN_TEST_PASS
+
+#include <llvm/CodeGen/MachineFunctionPass.h>
+
+using namespace llvm;
+
+namespace llvm {
+void initializeCodeGenTestPass(PassRegistry &);
+} // namespace llvm
+
+class CodeGenTest : public MachineFunctionPass {
+public:
+  static char ID;
+
+  CodeGenTest();
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  StringRef getPassName() const override;
+
+  static std::function<void()> RunCallback;
+};
+
+#endif // CODEGEN_TEST_PASS
diff --git a/llvm/unittests/CodeGen/CGPluginTest/Plugin/Plugin.cpp b/llvm/unittests/CodeGen/CGPluginTest/Plugin/Plugin.cpp
@@ -0,0 +1,26 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "CodeGenTestPass.h"
+
+#include <llvm/CodeGen/Passes.h>
+#include <llvm/CodeGen/TargetPassConfig.h>
+#include <llvm/Target/RegisterTargetPassConfigCallback.h>
+
+using namespace llvm;
+
+namespace {
+[[maybe_unused]] RegisterTargetPassConfigCallback X{
+    [](auto &TM, auto &PM, auto *TPC) {
+      TPC->insertPass(&GCLoweringID, &CodeGenTest::ID);
+    }};
+} // namespace
+
+__attribute__((constructor)) static void initCodeGenPlugin() {
+  initializeCodeGenTestPass(*PassRegistry::getPassRegistry());
+}
diff --git a/llvm/unittests/CodeGen/CGPluginTest/PluginTest.cpp b/llvm/unittests/CodeGen/CGPluginTest/PluginTest.cpp
@@ -0,0 +1,105 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "Plugin/CodeGenTestPass.h"
+
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
+#include "llvm/Config/config.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/DynamicLibrary.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/RegisterTargetPassConfigCallback.h"
+#include "llvm/Target/TargetMachine.h"
+
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+void anchor() {}
+
+std::string libPath(const std::string &Name) {
+  const auto &Argvs = testing::internal::GetArgvs();
+  const char *Argv0 = Argvs.size() > 0 ? Argvs[0].c_str() : "CGPluginTest";
+  void *Ptr = (void *)(intptr_t)anchor;
+  std::string Path = sys::fs::getMainExecutable(Argv0, Ptr);
+  SmallString<256> Buf{sys::path::parent_path(Path)};
+  sys::path::append(Buf, (Name + LLVM_PLUGIN_EXT).c_str());
+  return std::string(Buf.str());
+}
+} // namespace
+
+namespace llvm {
+class CGPluginTests : public testing::Test {
+protected:
+  static void SetUpTestCase() {
+    InitializeAllTargetMCs();
+    InitializeAllTargetInfos();
+    InitializeAllTargets();
+  }
+};
+} // namespace llvm
+
+TEST_F(CGPluginTests, LoadPlugin) {
+#if !defined(LLVM_ENABLE_PLUGINS)
+  // Skip the test if plugins are disabled.
+  GTEST_SKIP();
+#endif
+
+  auto PluginPath = libPath("CGTestPlugin");
+  ASSERT_NE("", PluginPath);
+
+  std::string Error;
+  auto Library = sys::DynamicLibrary::getLibrary(PluginPath.c_str(), &Error);
+  ASSERT_TRUE(Library.isValid()) << Error;
+  sys::DynamicLibrary::closeLibrary(Library);
+}
+
+TEST_F(CGPluginTests, ExecuteCallback) {
+#if !defined(LLVM_ENABLE_PLUGINS)
+  // Skip the test if plugins are disabled.
+  GTEST_SKIP();
+#endif
+
+  volatile bool CallbackExecuted = false;
+  volatile bool MPassExecuted = false;
+
+  RegisterTargetPassConfigCallback X{[&](auto &TM, auto &PM, auto *TPC) {
+    CallbackExecuted = true;
+    TPC->insertPass(&GCLoweringID, &CodeGenTest::ID);
+  }};
+
+  CodeGenTest::RunCallback = [&] { MPassExecuted = true; };
+
+  TargetOptions Options;
+  std::unique_ptr<MCContext> MCC;
+  for (auto T : TargetRegistry::targets()) {
+    if (!T.hasTargetMachine())
+      continue;
+    Triple TT{T.getName(), "", ""};
+    auto *TM = T.createTargetMachine(TT, "", "", Options, std::nullopt,
+                                     std::nullopt, CodeGenOptLevel::Default);
+    ASSERT_TRUE(TM);
+
+    legacy::PassManager PM;
+    MCC.reset(new MCContext(TT, TM->getMCAsmInfo(), TM->getMCRegisterInfo(),
+                            TM->getMCSubtargetInfo()));
+    auto *PtrMCC = MCC.get();
+    CallbackExecuted = false;
+    MPassExecuted = false;
+    if (TM->addPassesToEmitMC(PM, PtrMCC, outs()) == true)
+      continue;
+    ASSERT_TRUE(CallbackExecuted) << T.getName() << " callback failed";
+    ASSERT_TRUE(MPassExecuted) << T.getName() << " MachinePass failed";
+  }
+}
diff --git a/llvm/unittests/CodeGen/CMakeLists.txt b/llvm/unittests/CodeGen/CMakeLists.txt
@@ -51,5 +51,6 @@ add_llvm_unittest(CodeGenTests
   )
 
 add_subdirectory(GlobalISel)
+add_subdirectory(CGPluginTest)
 
 target_link_libraries(CodeGenTests PRIVATE LLVMTestingSupport)
