[OpenMP][NFC] Modernize the plugin handling (#74034)

This basically moves code around again, but this time to provide cleaner
interfaces and remove duplication. PluginAdaptorManagerTy is almost all
gone after this.

5 files changed, 170 insertions, 174 deletions

diff --git a/openmp/libomptarget/include/PluginManager.h b/openmp/libomptarget/include/PluginManager.h
index c92884d..720f4f7 100644
--- a/openmp/libomptarget/include/PluginManager.h
+++ b/openmp/libomptarget/include/PluginManager.h
@@ -20,23 +20,34 @@
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/DynamicLibrary.h"
 
 #include <list>
 #include <mutex>
+#include <string>
 
 struct PluginAdaptorTy {
-  int32_t Idx = -1;             // RTL index, index is the number of devices
-                                // of other RTLs that were registered before,
-                                // i.e. the OpenMP index of the first device
-                                // to be registered with this RTL.
-  int32_t NumberOfDevices = -1; // Number of devices this RTL deals with.
+  PluginAdaptorTy(const std::string &Name);
 
-  std::unique_ptr<llvm::sys::DynamicLibrary> LibraryHandler;
+  bool isUsed() const { return DeviceOffset >= 0; }
+
+  /// Return the number of devices available to this plugin.
+  int32_t getNumDevices() const { return NumberOfDevices; }
+
+  /// RTL index, index is the number of devices of other RTLs that were
+  /// registered before, i.e. the OpenMP index of the first device to be
+  /// registered with this RTL.
+  int32_t DeviceOffset = -1;
 
-#ifdef OMPTARGET_DEBUG
-  std::string RTLName;
-#endif
+  /// Number of devices this RTL deals with.
+  int32_t NumberOfDevices = -1;
+
+  /// Name of the shared object file representing the plugin.
+  std::string Name;
+
+  /// Access to the shared object file representing the plugin.
+  std::unique_ptr<llvm::sys::DynamicLibrary> LibraryHandler;
 
 #define PLUGIN_API_HANDLE(NAME, MANDATORY)                                     \
   using NAME##_ty = decltype(__tgt_rtl_##NAME);                                \
@@ -45,9 +56,6 @@ struct PluginAdaptorTy {
 #include "Shared/PluginAPI.inc"
 #undef PLUGIN_API_HANDLE
 
-  // Are there images associated with this RTL.
-  bool IsUsed = false;
-
   llvm::DenseSet<const __tgt_device_image *> UsedImages;
 
   // Mutex for thread-safety when calling RTL interface functions.
@@ -58,13 +66,6 @@ struct PluginAdaptorTy {
 
 /// RTLs identified in the system.
 struct PluginAdaptorManagerTy {
-  // List of the detected runtime libraries.
-  std::list<PluginAdaptorTy> AllRTLs;
-
-  // Array of pointers to the detected runtime libraries that have compatible
-  // binaries.
-  llvm::SmallVector<PluginAdaptorTy *> UsedRTLs;
-
   int64_t RequiresFlags = OMP_REQ_UNDEFINED;
 
   explicit PluginAdaptorManagerTy() = default;
@@ -72,27 +73,19 @@ struct PluginAdaptorManagerTy {
   // Register the clauses of the requires directive.
   void registerRequires(int64_t Flags);
 
-  // Initialize RTL if it has not been initialized
-  void initRTLonce(PluginAdaptorTy &RTL);
-
-  // Initialize all RTLs
-  void initAllRTLs();
-
   // Register a shared library with all (compatible) RTLs.
   void registerLib(__tgt_bin_desc *Desc);
 
   // Unregister a shared library from all RTLs.
   void unregisterLib(__tgt_bin_desc *Desc);
-
-  // not thread-safe, called from global constructor (i.e. once)
-  void loadRTLs();
-
-private:
-  static bool attemptLoadRTL(const std::string &RTLName, PluginAdaptorTy &RTL);
 };
 
 /// Struct for the data required to handle plugins
 struct PluginManager {
+  PluginManager() {}
+
+  void init();
+
   /// RTLs identified on the host
   PluginAdaptorManagerTy RTLs;
 
@@ -141,9 +134,30 @@ struct PluginManager {
     return Devices.size();
   }
 
+  int getNumUsedPlugins() const {
+    int NCI = 0;
+    for (auto &P : PluginAdaptors)
+      NCI += P.isUsed();
+    return NCI;
+  }
+
+  // Initialize \p Plugin if it has not been initialized.
+  void initPlugin(PluginAdaptorTy &Plugin);
+
+  // Initialize all plugins.
+  void initAllPlugins();
+
+  /// Iterator range for all plugin adaptors (in use or not, but always valid).
+  auto pluginAdaptors() {
+    return llvm::make_range(PluginAdaptors.begin(), PluginAdaptors.end());
+  }
+
 private:
   bool RTLsLoaded = false;
   llvm::SmallVector<__tgt_bin_desc *> DelayedBinDesc;
+
+  // List of all plugin adaptors, in use or not.
+  std::list<PluginAdaptorTy> PluginAdaptors;
 };
 
 extern PluginManager *PM;
diff --git a/openmp/libomptarget/src/PluginManager.cpp b/openmp/libomptarget/src/PluginManager.cpp
index 260aecd..82cca5e 100644
--- a/openmp/libomptarget/src/PluginManager.cpp
+++ b/openmp/libomptarget/src/PluginManager.cpp
@@ -12,4 +12,105 @@
 
 #include "PluginManager.h"
 
+using namespace llvm;
+using namespace llvm::sys;
+
 PluginManager *PM;
+
+// List of all plugins that can support offloading.
+static const char *RTLNames[] = {
+    /* PowerPC target       */ "libomptarget.rtl.ppc64",
+    /* x86_64 target        */ "libomptarget.rtl.x86_64",
+    /* CUDA target          */ "libomptarget.rtl.cuda",
+    /* AArch64 target       */ "libomptarget.rtl.aarch64",
+    /* AMDGPU target        */ "libomptarget.rtl.amdgpu",
+};
+
+PluginAdaptorTy::PluginAdaptorTy(const std::string &Name) : Name(Name) {
+  DP("Attempting to load library '%s'...\n", Name.c_str());
+
+  std::string ErrMsg;
+  LibraryHandler = std::make_unique<DynamicLibrary>(
+      DynamicLibrary::getPermanentLibrary(Name.c_str(), &ErrMsg));
+
+  if (!LibraryHandler->isValid()) {
+    // Library does not exist or cannot be found.
+    DP("Unable to load library '%s': %s!\n", Name.c_str(), ErrMsg.c_str());
+    return;
+  }
+
+  DP("Successfully loaded library '%s'!\n", Name.c_str());
+
+#define PLUGIN_API_HANDLE(NAME, MANDATORY)                                     \
+  NAME = reinterpret_cast<decltype(NAME)>(                                     \
+      LibraryHandler->getAddressOfSymbol(GETNAME(__tgt_rtl_##NAME)));          \
+  if (MANDATORY && !NAME) {                                                    \
+    DP("Invalid plugin as necessary interface is not found.\n");               \
+    return;                                                                    \
+  }
+
+#include "Shared/PluginAPI.inc"
+#undef PLUGIN_API_HANDLE
+
+  // Remove plugin on failure to call optional init_plugin
+  int32_t Rc = init_plugin();
+  if (Rc != OFFLOAD_SUCCESS) {
+    DP("Unable to initialize library '%s': %u!\n", Name.c_str(), Rc);
+    return;
+  }
+
+  // No devices are supported by this RTL?
+  NumberOfDevices = number_of_devices();
+  if (!NumberOfDevices) {
+    DP("No devices supported in this RTL\n");
+    return;
+  }
+
+  DP("Registered '%s' with %d devices!\n", Name.c_str(), NumberOfDevices);
+}
+
+void PluginManager::init() {
+  DP("Loading RTLs...\n");
+
+  // Attempt to open all the plugins and, if they exist, check if the interface
+  // is correct and if they are supporting any devices.
+  for (const char *Name : RTLNames) {
+    PluginAdaptors.emplace_back(std::string(Name) + ".so");
+    if (PluginAdaptors.back().getNumDevices() <= 0)
+      PluginAdaptors.pop_back();
+  }
+
+  DP("RTLs loaded!\n");
+}
+
+void PluginManager::initPlugin(PluginAdaptorTy &Plugin) {
+  // If this RTL is not already in use, initialize it.
+  if (Plugin.isUsed() || !Plugin.NumberOfDevices)
+    return;
+
+  // Initialize the device information for the RTL we are about to use.
+  const size_t Start = Devices.size();
+  Devices.reserve(Start + Plugin.NumberOfDevices);
+  for (int32_t DeviceId = 0; DeviceId < Plugin.NumberOfDevices; DeviceId++) {
+    Devices.push_back(std::make_unique<DeviceTy>(&Plugin));
+    // global device ID
+    Devices[Start + DeviceId]->DeviceID = Start + DeviceId;
+    // RTL local device ID
+    Devices[Start + DeviceId]->RTLDeviceID = DeviceId;
+  }
+
+  // Initialize the index of this RTL and save it in the used RTLs.
+  Plugin.DeviceOffset = Start;
+
+  // If possible, set the device identifier offset in the plugin.
+  if (Plugin.set_device_offset)
+    Plugin.set_device_offset(Start);
+
+  DP("RTL " DPxMOD " has index %d!\n", DPxPTR(Plugin.LibraryHandler.get()),
+     Plugin.DeviceOffset);
+}
+
+void PluginManager::initAllPlugins() {
+  for (auto &R : PluginAdaptors)
+    initPlugin(R);
+}
diff --git a/openmp/libomptarget/src/interface.cpp b/openmp/libomptarget/src/interface.cpp
index ee1bd49..3d82b02 100644
--- a/openmp/libomptarget/src/interface.cpp
+++ b/openmp/libomptarget/src/interface.cpp
@@ -54,7 +54,7 @@ EXTERN void __tgt_register_lib(__tgt_bin_desc *Desc) {
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Initialize all available devices without registering any image
-EXTERN void __tgt_init_all_rtls() { PM->RTLs.initAllRTLs(); }
+EXTERN void __tgt_init_all_rtls() { PM->initAllPlugins(); }
 
 ////////////////////////////////////////////////////////////////////////////////
 /// unloads a target shared library
@@ -426,7 +426,7 @@ EXTERN void __tgt_push_mapper_component(void *RtMapperHandle, void *Base,
 EXTERN void __tgt_set_info_flag(uint32_t NewInfoLevel) {
   std::atomic<uint32_t> &InfoLevel = getInfoLevelInternal();
   InfoLevel.store(NewInfoLevel);
-  for (auto &R : PM->RTLs.AllRTLs) {
+  for (auto &R : PM->pluginAdaptors()) {
     if (R.set_info_flag)
       R.set_info_flag(NewInfoLevel);
   }
diff --git a/openmp/libomptarget/src/omptarget.cpp b/openmp/libomptarget/src/omptarget.cpp
index 9d75fd3..fb854a4 100644
--- a/openmp/libomptarget/src/omptarget.cpp
+++ b/openmp/libomptarget/src/omptarget.cpp
@@ -300,7 +300,7 @@ void handleTargetOutcome(bool Success, ident_t *Loc) {
         FAILURE_MESSAGE("Consult https://openmp.llvm.org/design/Runtimes.html "
                         "for debugging options.\n");
 
-      if (PM->RTLs.UsedRTLs.empty()) {
+      if (!PM->getNumUsedPlugins()) {
         llvm::SmallVector<llvm::StringRef> Archs;
         llvm::transform(PM->Images, std::back_inserter(Archs),
                         [](const auto &X) {
diff --git a/openmp/libomptarget/src/rtl.cpp b/openmp/libomptarget/src/rtl.cpp
index 3cc7ac3..52ea764 100644
--- a/openmp/libomptarget/src/rtl.cpp
+++ b/openmp/libomptarget/src/rtl.cpp
@@ -32,15 +32,6 @@ using namespace llvm;
 using namespace llvm::sys;
 using namespace llvm::omp::target;
 
-// List of all plugins that can support offloading.
-static const char *RTLNames[] = {
-    /* PowerPC target       */ "libomptarget.rtl.ppc64",
-    /* x86_64 target        */ "libomptarget.rtl.x86_64",
-    /* CUDA target          */ "libomptarget.rtl.cuda",
-    /* AArch64 target       */ "libomptarget.rtl.aarch64",
-    /* AMDGPU target        */ "libomptarget.rtl.amdgpu",
-};
-
 #ifdef OMPT_SUPPORT
 extern void ompt::connectLibrary();
 #endif
@@ -55,8 +46,9 @@ __attribute__((constructor(101))) void init() {
   ompt::connectLibrary();
 #endif
 
+  PM->init();
+
   Profiler::get();
-  PM->RTLs.loadRTLs();
   PM->registerDelayedLibraries();
 }
 
@@ -65,85 +57,6 @@ __attribute__((destructor(101))) void deinit() {
   delete PM;
 }
 
-void PluginAdaptorManagerTy::loadRTLs() {
-  // Parse environment variable OMP_TARGET_OFFLOAD (if set)
-  PM->TargetOffloadPolicy =
-      (kmp_target_offload_kind_t)__kmpc_get_target_offload();
-  if (PM->TargetOffloadPolicy == tgt_disabled) {
-    return;
-  }
-
-  DP("Loading RTLs...\n");
-
-  // Attempt to open all the plugins and, if they exist, check if the interface
-  // is correct and if they are supporting any devices.
-  for (const char *Name : RTLNames) {
-    AllRTLs.emplace_back();
-
-    PluginAdaptorTy &RTL = AllRTLs.back();
-
-    const std::string BaseRTLName(Name);
-    if (!attemptLoadRTL(BaseRTLName + ".so", RTL))
-      AllRTLs.pop_back();
-  }
-
-  DP("RTLs loaded!\n");
-}
-
-bool PluginAdaptorManagerTy::attemptLoadRTL(const std::string &RTLName, PluginAdaptorTy &RTL) {
-  const char *Name = RTLName.c_str();
-
-  DP("Loading library '%s'...\n", Name);
-
-  std::string ErrMsg;
-  auto DynLibrary = std::make_unique<sys::DynamicLibrary>(
-      sys::DynamicLibrary::getPermanentLibrary(Name, &ErrMsg));
-
-  if (!DynLibrary->isValid()) {
-    // Library does not exist or cannot be found.
-    DP("Unable to load library '%s': %s!\n", Name, ErrMsg.c_str());
-    return false;
-  }
-
-  DP("Successfully loaded library '%s'!\n", Name);
-
-#define PLUGIN_API_HANDLE(NAME, MANDATORY)                                     \
-  *((void **)&RTL.NAME) =                                                      \
-      DynLibrary->getAddressOfSymbol(GETNAME(__tgt_rtl_##NAME));               \
-  if (MANDATORY && !RTL.NAME) {                                                \
-    DP("Invalid plugin as necessary interface is not found.\n");               \
-    return false;                                                              \
-  }
-
-#include "Shared/PluginAPI.inc"
-#undef PLUGIN_API_HANDLE
-
-  // Remove plugin on failure to call optional init_plugin
-  int32_t Rc = RTL.init_plugin();
-  if (Rc != OFFLOAD_SUCCESS) {
-    DP("Unable to initialize library '%s': %u!\n", Name, Rc);
-    return false;
-  }
-
-  // No devices are supported by this RTL?
-  if (!(RTL.NumberOfDevices = RTL.number_of_devices())) {
-    // The RTL is invalid! Will pop the object from the RTLs list.
-    DP("No devices supported in this RTL\n");
-    return false;
-  }
-
-#ifdef OMPTARGET_DEBUG
-  RTL.RTLName = Name;
-#endif
-
-  DP("Registering RTL %s supporting %d devices!\n", Name, RTL.NumberOfDevices);
-
-  RTL.LibraryHandler = std::move(DynLibrary);
-
-  // Successfully loaded
-  return true;
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 // Functionality for registering libs
 
@@ -157,7 +70,7 @@ static void registerImageIntoTranslationTable(TranslationTable &TT,
 
   // Resize the Targets Table and Images to accommodate the new targets if
   // required
-  unsigned TargetsTableMinimumSize = RTL.Idx + RTL.NumberOfDevices;
+  unsigned TargetsTableMinimumSize = RTL.DeviceOffset + RTL.NumberOfDevices;
 
   if (TT.TargetsTable.size() < TargetsTableMinimumSize) {
     TT.TargetsImages.resize(TargetsTableMinimumSize, 0);
@@ -167,9 +80,10 @@ static void registerImageIntoTranslationTable(TranslationTable &TT,
   // Register the image in all devices for this target type.
   for (int32_t I = 0; I < RTL.NumberOfDevices; ++I) {
     // If we are changing the image we are also invalidating the target table.
-    if (TT.TargetsImages[RTL.Idx + I] != Image) {
-      TT.TargetsImages[RTL.Idx + I] = Image;
-      TT.TargetsTable[RTL.Idx + I] = 0; // lazy initialization of target table.
+    if (TT.TargetsImages[RTL.DeviceOffset + I] != Image) {
+      TT.TargetsImages[RTL.DeviceOffset + I] = Image;
+      TT.TargetsTable[RTL.DeviceOffset + I] =
+          0; // lazy initialization of target table.
     }
   }
 }
@@ -182,7 +96,7 @@ static void registerGlobalCtorsDtorsForImage(__tgt_bin_desc *Desc,
                                              PluginAdaptorTy *RTL) {
 
   for (int32_t I = 0; I < RTL->NumberOfDevices; ++I) {
-    DeviceTy &Device = *PM->Devices[RTL->Idx + I];
+    DeviceTy &Device = *PM->Devices[RTL->DeviceOffset + I];
     Device.PendingGlobalsMtx.lock();
     Device.HasPendingGlobals = true;
     for (__tgt_offload_entry *Entry = Img->EntriesBegin;
@@ -291,39 +205,6 @@ void PluginAdaptorManagerTy::registerRequires(int64_t Flags) {
      Flags, RequiresFlags);
 }
 
-void PluginAdaptorManagerTy::initRTLonce(PluginAdaptorTy &R) {
-  // If this RTL is not already in use, initialize it.
-  if (R.IsUsed || !R.NumberOfDevices)
-    return;
-
-  // Initialize the device information for the RTL we are about to use.
-  const size_t Start = PM->Devices.size();
-  PM->Devices.reserve(Start + R.NumberOfDevices);
-  for (int32_t DeviceId = 0; DeviceId < R.NumberOfDevices; DeviceId++) {
-    PM->Devices.push_back(std::make_unique<DeviceTy>(&R));
-    // global device ID
-    PM->Devices[Start + DeviceId]->DeviceID = Start + DeviceId;
-    // RTL local device ID
-    PM->Devices[Start + DeviceId]->RTLDeviceID = DeviceId;
-  }
-
-  // Initialize the index of this RTL and save it in the used RTLs.
-  R.Idx = Start;
-  R.IsUsed = true;
-  UsedRTLs.push_back(&R);
-
-  // If possible, set the device identifier offset
-  if (R.set_device_offset)
-    R.set_device_offset(Start);
-
-  DP("RTL " DPxMOD " has index %d!\n", DPxPTR(R.LibraryHandler.get()), R.Idx);
-}
-
-void PluginAdaptorManagerTy::initAllRTLs() {
-  for (auto &R : AllRTLs)
-    initRTLonce(R);
-}
-
 void PluginAdaptorManagerTy::registerLib(__tgt_bin_desc *Desc) {
   PM->RTLsMtx.lock();
 
@@ -342,23 +223,23 @@ void PluginAdaptorManagerTy::registerLib(__tgt_bin_desc *Desc) {
 
     // Scan the RTLs that have associated images until we find one that supports
     // the current image.
-    for (auto &R : AllRTLs) {
+    for (auto &R : PM->pluginAdaptors()) {
       if (R.is_valid_binary_info) {
         if (!R.is_valid_binary_info(Img, Info)) {
           DP("Image " DPxMOD " is NOT compatible with RTL %s!\n",
-             DPxPTR(Img->ImageStart), R.RTLName.c_str());
+             DPxPTR(Img->ImageStart), R.Name.c_str());
           continue;
         }
       } else if (!R.is_valid_binary(Img)) {
         DP("Image " DPxMOD " is NOT compatible with RTL %s!\n",
-           DPxPTR(Img->ImageStart), R.RTLName.c_str());
+           DPxPTR(Img->ImageStart), R.Name.c_str());
         continue;
       }
 
       DP("Image " DPxMOD " is compatible with RTL %s!\n",
-         DPxPTR(Img->ImageStart), R.RTLName.c_str());
+         DPxPTR(Img->ImageStart), R.Name.c_str());
 
-      initRTLonce(R);
+      PM->initPlugin(R);
 
       // Initialize (if necessary) translation table for this library.
       PM->TrlTblMtx.lock();
@@ -375,7 +256,7 @@ void PluginAdaptorManagerTy::registerLib(__tgt_bin_desc *Desc) {
           (PM->HostEntriesBeginToTransTable)[Desc->HostEntriesBegin];
 
       DP("Registering image " DPxMOD " with RTL %s!\n", DPxPTR(Img->ImageStart),
-         R.RTLName.c_str());
+         R.Name.c_str());
       registerImageIntoTranslationTable(TransTable, R, Img);
       R.UsedImages.insert(Img);
 
@@ -411,20 +292,20 @@ void PluginAdaptorManagerTy::unregisterLib(__tgt_bin_desc *Desc) {
 
     // Scan the RTLs that have associated images until we find one that supports
     // the current image. We only need to scan RTLs that are already being used.
-    for (auto *R : UsedRTLs) {
-
-      assert(R->IsUsed && "Expecting used RTLs.");
+    for (auto &R : PM->pluginAdaptors()) {
+      if (!R.isUsed())
+        continue;
 
       // Ensure that we do not use any unused images associated with this RTL.
-      if (!R->UsedImages.contains(Img))
+      if (!R.UsedImages.contains(Img))
         continue;
 
-      FoundRTL = R;
+      FoundRTL = &R;
 
       // Execute dtors for static objects if the device has been used, i.e.
       // if its PendingCtors list has been emptied.
       for (int32_t I = 0; I < FoundRTL->NumberOfDevices; ++I) {
-        DeviceTy &Device = *PM->Devices[FoundRTL->Idx + I];
+        DeviceTy &Device = *PM->Devices[FoundRTL->DeviceOffset + I];
         Device.PendingGlobalsMtx.lock();
         if (Device.PendingCtorsDtors[Desc].PendingCtors.empty()) {
           AsyncInfoTy AsyncInfo(Device);
@@ -445,7 +326,7 @@ void PluginAdaptorManagerTy::unregisterLib(__tgt_bin_desc *Desc) {
       }
 
       DP("Unregistered image " DPxMOD " from RTL " DPxMOD "!\n",
-         DPxPTR(Img->ImageStart), DPxPTR(R->LibraryHandler.get()));
+         DPxPTR(Img->ImageStart), DPxPTR(R.LibraryHandler.get()));
 
       break;
     }