[NPM] Resolve llvmGetPassPluginInfo to the plugin being loaded

Dynamically loaded plugins for the new pass manager are initialised by
calling llvmGetPassPluginInfo. This is defined as a weak symbol so that
it is continually redefined by each plugin that is loaded. When loading
a plugin from a shared library, the intention is that
llvmGetPassPluginInfo will be resolved to the definition in the most
recent plugin. However, using a global search for this resolution can
fail in situations where multiple plugins are loaded.

Currently:

* If a plugin does not define llvmGetPassPluginInfo, then it will be
  silently resolved to the previous plugin's definition.

* If loading the same plugin twice with another in between, e.g. plugin
  A/plugin B/plugin A, then the second load of plugin A will resolve to
  llvmGetPassPluginInfo in plugin B.

* The previous case can also occur when a dynamic library defines both
  NPM and legacy plugins; the legacy plugins are loaded first and then
  with `-fplugin=A -fpass-plugin=B -fpass-plugin=A`: A will be loaded as
  a legacy plugin and define llvmGetPassPluginInfo; B will be loaded
  and redefine it; and finally when A is loaded as an NPM plugin it will
  be resolved to the definition from B.

Instead of searching globally, restrict the symbol lookup to the library
that is currently being loaded.

Differential Revision: https://reviews.llvm.org/D104916

lib/Passes/PassPlugin.cpp

@@ -23,8 +23,11 @@ Expected<PassPlugin> PassPlugin::Load(const std::string &Filename) {
                                    inconvertibleErrorCode());
 
   PassPlugin P{Filename, Library};
+
+  // llvmGetPassPluginInfo should be resolved to the definition from the plugin
+  // we are currently loading.
   intptr_t getDetailsFn =
-      (intptr_t)Library.SearchForAddressOfSymbol("llvmGetPassPluginInfo");
+      (intptr_t)Library.getAddressOfSymbol("llvmGetPassPluginInfo");
 
   if (!getDetailsFn)
     // If the symbol isn't found, this is probably a legacy plugin, which is an

unittests/Passes/CMakeLists.txt

@@ -1,5 +1,5 @@
 # Needed by LLVM's CMake checks because this file defines multiple targets.
-set(LLVM_OPTIONAL_SOURCES PluginsTest.cpp TestPlugin.cpp PassBuilderBindingsTest.cpp)
+set(LLVM_OPTIONAL_SOURCES PluginsTest.cpp TestPlugin.cpp DoublerPlugin.cpp PassBuilderBindingsTest.cpp)
 
 # If plugins are disabled, this test will disable itself at runtime. Otherwise,
 # reconfiguring with plugins disabled will leave behind a stale executable.
@@ -20,19 +20,20 @@ if (NOT WIN32)
   target_link_libraries(PluginsTests PRIVATE LLVMTestingSupport)
 
   set(LLVM_LINK_COMPONENTS)
-  add_llvm_library(TestPlugin MODULE BUILDTREE_ONLY
-    TestPlugin.cpp
-    )
+  foreach(PLUGIN TestPlugin DoublerPlugin)
+    add_llvm_library(${PLUGIN} MODULE BUILDTREE_ONLY ${PLUGIN}.cpp)
 
-  # Put plugin next to the unit test executable.
-  set_output_directory(TestPlugin
-    BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}
-    LIBRARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}
-    )
-  set_target_properties(TestPlugin PROPERTIES FOLDER "Tests")
+    # Put PLUGIN next to the unit test executable.
+    set_output_directory(${PLUGIN}
+      BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}
+      LIBRARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}
+      )
+    set_target_properties(${PLUGIN} PROPERTIES FOLDER "Tests")
+
+    add_dependencies(${PLUGIN} intrinsics_gen)
+    add_dependencies(PluginsTests ${PLUGIN})
+  endforeach()
 
-  add_dependencies(TestPlugin intrinsics_gen)
-  add_dependencies(PluginsTests TestPlugin)
 endif()
 
 set(LLVM_LINK_COMPONENTS Support Passes Core Target native AllTargetsInfos)
@@ -40,3 +41,4 @@ add_llvm_unittest(PassesBindingsTests
 	PassBuilderBindingsTest.cpp
 	)
 target_link_libraries(PassesBindingsTests PRIVATE LLVMTestingSupport)
+

unittests/Passes/DoublerPlugin.cpp

@@ -0,0 +1,44 @@
+//===- unittests/Passes/DoublerPlugin.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/PassPlugin.h"
+
+using namespace llvm;
+
+struct DoublerModulePass : public PassInfoMixin<DoublerModulePass> {
+
+  // Double the value of the initializer
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    auto *GV = cast<GlobalVariable>(M.getNamedValue("doubleme"));
+    auto *Init = GV->getInitializer();
+    auto *Init2 = ConstantExpr::getAdd(Init, Init);
+    GV->setInitializer(Init2);
+
+    return PreservedAnalyses::none();
+  }
+
+  static void registerCallbacks(PassBuilder &PB) {
+    PB.registerPipelineParsingCallback(
+        [](StringRef Name, ModulePassManager &PM,
+           ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
+          if (Name == "doubler-pass") {
+            PM.addPass(DoublerModulePass());
+            return true;
+          }
+          return false;
+        });
+  }
+};
+
+extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK
+llvmGetPassPluginInfo() {
+  return {LLVM_PLUGIN_API_VERSION, "DoublerPlugin", "2.2-unit",
+          DoublerModulePass::registerCallbacks};
+}

unittests/Passes/PluginsTest.cpp

@@ -7,7 +7,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/AsmParser/Parser.h"
 #include "llvm/Config/config.h"
+#include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
@@ -58,3 +60,80 @@ TEST(PluginsTests, LoadPlugin) {
   Plugin->registerPassBuilderCallbacks(PB);
   ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, "plugin-pass"), Succeeded());
 }
+
+// Test that llvmGetPassPluginInfo from DoublerPlugin is called twice with
+// -fpass-plugin=DoublerPlugin -fpass-plugin=TestPlugin
+// -fpass-plugin=DoublerPlugin.
+TEST(PluginsTests, LoadMultiplePlugins) {
+#if !defined(LLVM_ENABLE_PLUGINS)
+  // Disable the test if plugins are disabled.
+  return;
+#endif
+
+  auto DoublerPluginPath = LibPath("DoublerPlugin");
+  auto TestPluginPath = LibPath("TestPlugin");
+  ASSERT_NE("", DoublerPluginPath);
+  ASSERT_NE("", TestPluginPath);
+
+  Expected<PassPlugin> DoublerPlugin1 = PassPlugin::Load(DoublerPluginPath);
+  ASSERT_TRUE(!!DoublerPlugin1)
+      << "Plugin path: " << DoublerPlugin1->getFilename();
+
+  Expected<PassPlugin> TestPlugin = PassPlugin::Load(TestPluginPath);
+  ASSERT_TRUE(!!TestPlugin) << "Plugin path: " << TestPlugin->getFilename();
+
+  // If llvmGetPassPluginInfo is resolved as a weak symbol taking into account
+  // all loaded symbols, the second call to PassPlugin::Load will actually
+  // return the llvmGetPassPluginInfo from the most recently loaded plugin, in
+  // this case TestPlugin.
+  Expected<PassPlugin> DoublerPlugin2 = PassPlugin::Load(DoublerPluginPath);
+  ASSERT_TRUE(!!DoublerPlugin2)
+      << "Plugin path: " << DoublerPlugin2->getFilename();
+
+  ASSERT_EQ("DoublerPlugin", DoublerPlugin1->getPluginName());
+  ASSERT_EQ("2.2-unit", DoublerPlugin1->getPluginVersion());
+  ASSERT_EQ(TEST_PLUGIN_NAME, TestPlugin->getPluginName());
+  ASSERT_EQ(TEST_PLUGIN_VERSION, TestPlugin->getPluginVersion());
+  // Check that the plugin name/version is set correctly when loaded a second
+  // time
+  ASSERT_EQ("DoublerPlugin", DoublerPlugin2->getPluginName());
+  ASSERT_EQ("2.2-unit", DoublerPlugin2->getPluginVersion());
+
+  PassBuilder PB;
+  ModulePassManager PM;
+  const char *PipelineText = "module(doubler-pass,plugin-pass,doubler-pass)";
+  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText), Failed());
+  TestPlugin->registerPassBuilderCallbacks(PB);
+  DoublerPlugin1->registerPassBuilderCallbacks(PB);
+  DoublerPlugin2->registerPassBuilderCallbacks(PB);
+  ASSERT_THAT_ERROR(PB.parsePassPipeline(PM, PipelineText), Succeeded());
+
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M =
+      parseAssemblyString(R"IR(@doubleme = constant i32 7)IR", Err, C);
+
+  // Check that the initial value is 7
+  {
+    auto *GV = M->getNamedValue("doubleme");
+    auto *Init = cast<GlobalVariable>(GV)->getInitializer();
+    auto *CI = cast<ConstantInt>(Init);
+    ASSERT_EQ(CI->getSExtValue(), 7);
+  }
+
+  ModuleAnalysisManager MAM;
+  // Register required pass instrumentation analysis.
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(); });
+  PM.run(*M, MAM);
+
+  // Check that the final value is 28 because DoublerPlugin::run was called
+  // twice, indicating that the llvmGetPassPluginInfo and registerCallbacks
+  // were correctly called.
+  {
+    // Check the value was doubled twice
+    auto *GV = M->getNamedValue("doubleme");
+    auto *Init = cast<GlobalVariable>(GV)->getInitializer();
+    auto *CI = cast<ConstantInt>(Init);
+    ASSERT_EQ(CI->getSExtValue(), 28);
+  }
+}

unittests/Passes/TestPlugin.cpp

@@ -1,4 +1,4 @@
-//===- unittests/Passes/Plugins/Plugin.cpp --------------------------------===//
+//===- unittests/Passes/TestPlugin.cpp --------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -17,22 +17,22 @@ struct TestModulePass : public PassInfoMixin<TestModulePass> {
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
     return PreservedAnalyses::all();
   }
-};
 
-void registerCallbacks(PassBuilder &PB) {
-  PB.registerPipelineParsingCallback(
-      [](StringRef Name, ModulePassManager &PM,
-         ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
-        if (Name == "plugin-pass") {
-          PM.addPass(TestModulePass());
-          return true;
-        }
-        return false;
-      });
-}
+  static void registerCallbacks(PassBuilder &PB) {
+    PB.registerPipelineParsingCallback(
+        [](StringRef Name, ModulePassManager &PM,
+           ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
+          if (Name == "plugin-pass") {
+            PM.addPass(TestModulePass());
+            return true;
+          }
+          return false;
+        });
+  }
+};
 
 extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK
 llvmGetPassPluginInfo() {
   return {LLVM_PLUGIN_API_VERSION, TEST_PLUGIN_NAME, TEST_PLUGIN_VERSION,
-          registerCallbacks};
+          TestModulePass::registerCallbacks};
 }