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

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

Change-Id: I2c9db48af4b047277bbde393124716f19e9002d8

Author: jdoerfert

openmp/libomptarget/include/PluginManager.h

@@ -20,27 +20,38 @@
 
 #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>
 
 // Forward declarations.
 struct DeviceTy;
 struct __tgt_bin_desc;
 
 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);                                \
@@ -49,9 +60,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.
@@ -62,41 +70,26 @@ 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;
 
   // 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;
 
@@ -145,9 +138,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;

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[] = {
+    /* AMDGPU target        */ "libomptarget.rtl.amdgpu",
+    /* CUDA target          */ "libomptarget.rtl.cuda",
+    /* x86_64 target        */ "libomptarget.rtl.x86_64",
+    /* AArch64 target       */ "libomptarget.rtl.aarch64",
+    /* PowerPC target       */ "libomptarget.rtl.ppc64",
+};
+
+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);
+}

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
@@ -448,7 +448,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);
   }

openmp/libomptarget/src/omptarget.cpp

@@ -312,7 +312,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) {