[OpenMP] Reorganize the initialization of PluginAdaptorTy (#74397)

This introduces checked errors into the creation and initialization of
`PluginAdaptorTy`. We also allow the adaptor to "hide" devices from the
user if the initialization failed. The new organization avoids the
"initOnce" stuff but we still do not eagerly initialize the plugin
devices (I think we should merge `PluginAdaptorTy::initDevices` into
`PluginAdaptorTy::init`)

Author: jdoerfert

openmp/libomptarget/include/PluginManager.h

@@ -34,13 +34,29 @@
 #include <mutex>
 #include <string>
 
+struct PluginManager;
+
+/// Plugin adaptors should be created via `PluginAdaptorTy::create` which will
+/// invoke the constructor and call `PluginAdaptorTy::init`. Eventual errors are
+/// reported back to the caller, otherwise a valid and initialized adaptor is
+/// returned.
 struct PluginAdaptorTy {
-  PluginAdaptorTy(const std::string &Name);
+  /// Try to create a plugin adaptor from a filename.
+  static llvm::Expected<std::unique_ptr<PluginAdaptorTy>>
+  create(const std::string &Name);
+
+  /// Initialize as many devices as possible for this plugin adaptor. Devices
+  /// that fail to initialize are ignored.
+  void initDevices(PluginManager &PM);
 
   bool isUsed() const { return DeviceOffset >= 0; }
 
-  /// Return the number of devices available to this plugin.
-  int32_t getNumDevices() const { return NumberOfDevices; }
+  /// Return the number of devices visible to the underlying plugin.
+  int32_t getNumberOfPluginDevices() const { return NumberOfPluginDevices; }
+
+  /// Return the number of devices successfully initialized and visible to the
+  /// user.
+  int32_t getNumberOfUserDevices() const { return NumberOfUserDevices; }
 
   /// Add all offload entries described by \p DI to the devices managed by this
   /// plugin.
@@ -51,9 +67,6 @@ struct PluginAdaptorTy {
   /// registered with this RTL.
   int32_t DeviceOffset = -1;
 
-  /// Number of devices this RTL deals with.
-  int32_t NumberOfDevices = -1;
-
   /// Name of the shared object file representing the plugin.
   std::string Name;
 
@@ -73,6 +86,22 @@ struct PluginAdaptorTy {
   // It is easier to enforce thread-safety at the libomptarget level,
   // so that developers of new RTLs do not have to worry about it.
   std::mutex Mtx;
+
+private:
+  /// Number of devices the underling plugins sees.
+  int32_t NumberOfPluginDevices = -1;
+
+  /// Number of devices exposed to the user. This can be less than the number of
+  /// devices for the plugin if some failed to initialize.
+  int32_t NumberOfUserDevices = 0;
+
+  /// Create a plugin adaptor for filename \p Name with a dynamic library \p DL.
+  PluginAdaptorTy(const std::string &Name,
+                  std::unique_ptr<llvm::sys::DynamicLibrary> DL);
+
+  /// Initialize the plugin adaptor, this can fail in which case the adaptor is
+  /// useless.
+  llvm::Error init();
 };
 
 /// Struct for the data required to handle plugins
@@ -150,20 +179,15 @@ struct PluginManager {
   int getNumUsedPlugins() const {
     int NCI = 0;
     for (auto &P : PluginAdaptors)
-      NCI += P.isUsed();
+      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());
-  }
+  auto pluginAdaptors() { return llvm::make_pointee_range(PluginAdaptors); }
 
   /// Return the user provided requirements.
   int64_t getRequirements() const { return Requirements.getRequirements(); }
@@ -176,7 +200,7 @@ struct PluginManager {
   llvm::SmallVector<__tgt_bin_desc *> DelayedBinDesc;
 
   // List of all plugin adaptors, in use or not.
-  std::list<PluginAdaptorTy> PluginAdaptors;
+  llvm::SmallVector<std::unique_ptr<PluginAdaptorTy>> PluginAdaptors;
 
   /// Executable images and information extracted from the input images passed
   /// to the runtime.

openmp/libomptarget/include/device.h

@@ -51,8 +51,6 @@ struct DeviceTy {
   PluginAdaptorTy *RTL;
   int32_t RTLDeviceID;
 
-  bool IsInit;
-  std::once_flag InitFlag;
   bool HasMappedGlobalData = false;
 
   /// Host data to device map type with a wrapper key indirection that allows
@@ -72,13 +70,16 @@ struct DeviceTy {
 
   std::mutex PendingGlobalsMtx;
 
-  DeviceTy(PluginAdaptorTy *RTL);
+  DeviceTy(PluginAdaptorTy *RTL, int32_t DeviceID, int32_t RTLDeviceID);
   // DeviceTy is not copyable
   DeviceTy(const DeviceTy &D) = delete;
   DeviceTy &operator=(const DeviceTy &D) = delete;
 
   ~DeviceTy();
 
+  /// Try to initialize the device and return any failure.
+  llvm::Error init();
+
   // Return true if data can be copied to DstDevice directly
   bool isDataExchangable(const DeviceTy &DstDevice);
 
@@ -145,8 +146,6 @@ struct DeviceTy {
   int associatePtr(void *HstPtrBegin, void *TgtPtrBegin, int64_t Size);
   int disassociatePtr(void *HstPtrBegin);
 
-  // calls to RTL
-  int32_t initOnce();
   __tgt_target_table *loadBinary(__tgt_device_image *Img);
 
   // device memory allocation/deallocation routines
@@ -237,9 +236,6 @@ struct DeviceTy {
   void dumpOffloadEntries();
 
 private:
-  // Call to RTL
-  void init(); // To be called only via DeviceTy::initOnce()
-
   /// Deinitialize the device (and plugin).
   void deinit();
 

openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp

@@ -1712,7 +1712,7 @@ int32_t __tgt_rtl_number_of_devices() { return Plugin::get().getNumDevices(); }
 
 int64_t __tgt_rtl_init_requires(int64_t RequiresFlags) {
   Plugin::get().setRequiresFlag(RequiresFlags);
-  return RequiresFlags;
+  return OFFLOAD_SUCCESS;
 }
 
 int32_t __tgt_rtl_is_data_exchangable(int32_t SrcDeviceId,

openmp/libomptarget/src/PluginManager.cpp

@@ -15,6 +15,7 @@
 
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
+#include <memory>
 
 using namespace llvm;
 using namespace llvm::sys;
@@ -30,27 +31,43 @@ static const char *RTLNames[] = {
     /* AMDGPU target        */ "libomptarget.rtl.amdgpu",
 };
 
-PluginAdaptorTy::PluginAdaptorTy(const std::string &Name) : Name(Name) {
+Expected<std::unique_ptr<PluginAdaptorTy>>
+PluginAdaptorTy::create(const std::string &Name) {
   DP("Attempting to load library '%s'...\n", Name.c_str());
 
   std::string ErrMsg;
-  LibraryHandler = std::make_unique<DynamicLibrary>(
+  auto 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;
+    return createStringError(inconvertibleErrorCode(),
+                             "Unable to load library '%s': %s!\n", Name.c_str(),
+                             ErrMsg.c_str());
   }
 
   DP("Successfully loaded library '%s'!\n", Name.c_str());
+  auto PluginAdaptor = std::unique_ptr<PluginAdaptorTy>(
+      new PluginAdaptorTy(Name, std::move(LibraryHandler)));
+  if (auto Err = PluginAdaptor->init())
+    return Err;
+  return PluginAdaptor;
+}
+
+PluginAdaptorTy::PluginAdaptorTy(const std::string &Name,
+                                 std::unique_ptr<llvm::sys::DynamicLibrary> DL)
+    : Name(Name), LibraryHandler(std::move(DL)) {}
+
+Error PluginAdaptorTy::init() {
 
 #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;                                                                    \
+    return createStringError(inconvertibleErrorCode(),                         \
+                             "Invalid plugin as necessary interface function " \
+                             "(%s) was not found.\n",                          \
+                             NAME);                                            \
   }
 
 #include "Shared/PluginAPI.inc"
@@ -59,22 +76,25 @@ PluginAdaptorTy::PluginAdaptorTy(const std::string &Name) : Name(Name) {
   // 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;
+    return createStringError(inconvertibleErrorCode(),
+                             "Unable to initialize library '%s': %u!\n",
+                             Name.c_str(), Rc);
   }
 
   // No devices are supported by this RTL?
-  NumberOfDevices = number_of_devices();
-  if (!NumberOfDevices) {
-    DP("No devices supported in this RTL\n");
-    return;
+  NumberOfPluginDevices = number_of_devices();
+  if (!NumberOfPluginDevices) {
+    return createStringError(inconvertibleErrorCode(),
+                             "No devices supported in this RTL\n");
   }
 
-  DP("Registered '%s' with %d devices!\n", Name.c_str(), NumberOfDevices);
+  DP("Registered '%s' with %d plugin visible devices!\n", Name.c_str(),
+     NumberOfPluginDevices);
+  return Error::success();
 }
 
 void PluginAdaptorTy::addOffloadEntries(DeviceImageTy &DI) {
-  for (int32_t I = 0; I < NumberOfDevices; ++I) {
+  for (int32_t I = 0, E = getNumberOfUserDevices(); I < E; ++I) {
     auto DeviceOrErr = PM->getDevice(DeviceOffset + I);
     if (!DeviceOrErr)
       FATAL_MESSAGE(DeviceOffset + I, "%s",
@@ -92,45 +112,61 @@ void PluginManager::init() {
   // 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();
+    auto PluginAdaptorOrErr =
+        PluginAdaptorTy::create(std::string(Name) + ".so");
+    if (!PluginAdaptorOrErr) {
+      [[maybe_unused]] std::string InfoMsg =
+          toString(PluginAdaptorOrErr.takeError());
+      DP("%s", InfoMsg.c_str());
+    } else {
+      PluginAdaptors.push_back(std::move(*PluginAdaptorOrErr));
+    }
   }
 
   DP("RTLs loaded!\n");
 }
 
-void PluginManager::initPlugin(PluginAdaptorTy &Plugin) {
-  // If this RTL is not already in use, initialize it.
-  if (Plugin.isUsed() || !Plugin.NumberOfDevices)
+void PluginAdaptorTy::initDevices(PluginManager &PM) {
+  if (isUsed())
     return;
 
+  // If this RTL is not already in use, initialize it.
+  assert(getNumberOfPluginDevices() > 0 &&
+         "Tried to initialize useless plugin adaptor");
+
   // Initialize the device information for the RTL we are about to use.
-  auto ExclusiveDevicesAccessor = getExclusiveDevicesAccessor();
-  const size_t Start = ExclusiveDevicesAccessor->size();
-  ExclusiveDevicesAccessor->reserve(Start + Plugin.NumberOfDevices);
-  for (int32_t DeviceId = 0; DeviceId < Plugin.NumberOfDevices; DeviceId++) {
-    ExclusiveDevicesAccessor->push_back(std::make_unique<DeviceTy>(&Plugin));
-    // global device ID
-    (*ExclusiveDevicesAccessor)[Start + DeviceId]->DeviceID = Start + DeviceId;
-    // RTL local device ID
-    (*ExclusiveDevicesAccessor)[Start + DeviceId]->RTLDeviceID = DeviceId;
-  }
+  auto ExclusiveDevicesAccessor = PM.getExclusiveDevicesAccessor();
 
   // Initialize the index of this RTL and save it in the used RTLs.
-  Plugin.DeviceOffset = Start;
+  DeviceOffset = ExclusiveDevicesAccessor->size();
 
   // If possible, set the device identifier offset in the plugin.
-  if (Plugin.set_device_offset)
-    Plugin.set_device_offset(Start);
+  if (set_device_offset)
+    set_device_offset(DeviceOffset);
+
+  int32_t NumPD = getNumberOfPluginDevices();
+  ExclusiveDevicesAccessor->reserve(DeviceOffset + NumPD);
+  for (int32_t PDevI = 0, UserDevId = DeviceOffset; PDevI < NumPD; PDevI++) {
+    auto Device = std::make_unique<DeviceTy>(this, UserDevId, PDevI);
+    if (auto Err = Device->init()) {
+      DP("Skip plugin known device %d: %s\n", PDevI,
+         toString(std::move(Err)).c_str());
+      continue;
+    }
+
+    ExclusiveDevicesAccessor->push_back(std::move(Device));
+    ++NumberOfUserDevices;
+    ++UserDevId;
+  }
 
-  DP("RTL " DPxMOD " has index %d!\n", DPxPTR(Plugin.LibraryHandler.get()),
-     Plugin.DeviceOffset);
+  DP("Plugin adaptor " DPxMOD " has index %d, exposes %d out of %d devices!\n",
+     DPxPTR(LibraryHandler.get()), DeviceOffset, NumberOfUserDevices,
+     NumberOfPluginDevices);
 }
 
 void PluginManager::initAllPlugins() {
   for (auto &R : PluginAdaptors)
-    initPlugin(R);
+    R->initDevices(*this);
 }
 
 static void registerImageIntoTranslationTable(TranslationTable &TT,
@@ -143,15 +179,16 @@ static void registerImageIntoTranslationTable(TranslationTable &TT,
 
   // Resize the Targets Table and Images to accommodate the new targets if
   // required
-  unsigned TargetsTableMinimumSize = RTL.DeviceOffset + RTL.NumberOfDevices;
+  unsigned TargetsTableMinimumSize =
+      RTL.DeviceOffset + RTL.getNumberOfUserDevices();
 
   if (TT.TargetsTable.size() < TargetsTableMinimumSize) {
     TT.TargetsImages.resize(TargetsTableMinimumSize, 0);
     TT.TargetsTable.resize(TargetsTableMinimumSize, 0);
   }
 
   // Register the image in all devices for this target type.
-  for (int32_t I = 0; I < RTL.NumberOfDevices; ++I) {
+  for (int32_t I = 0; I < RTL.getNumberOfUserDevices(); ++I) {
     // If we are changing the image we are also invalidating the target table.
     if (TT.TargetsImages[RTL.DeviceOffset + I] != Image) {
       TT.TargetsImages[RTL.DeviceOffset + I] = Image;
@@ -194,7 +231,7 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
       DP("Image " DPxMOD " is compatible with RTL %s!\n",
          DPxPTR(Img->ImageStart), R.Name.c_str());
 
-      PM->initPlugin(R);
+      R.initDevices(*this);
 
       // Initialize (if necessary) translation table for this library.
       PM->TrlTblMtx.lock();
@@ -263,7 +300,7 @@ void PluginManager::unregisterLib(__tgt_bin_desc *Desc) {
 
       // 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) {
+      for (int32_t I = 0; I < FoundRTL->getNumberOfUserDevices(); ++I) {
         auto DeviceOrErr = PM->getDevice(FoundRTL->DeviceOffset + I);
         if (!DeviceOrErr)
           FATAL_MESSAGE(FoundRTL->DeviceOffset + I, "%s",
@@ -337,17 +374,5 @@ Expected<DeviceTy &> PluginManager::getDevice(uint32_t DeviceNo) {
         "Device number '%i' out of range, only %i devices available", DeviceNo,
         ExclusiveDevicesAccessor->size());
 
-  DeviceTy &Device = *(*ExclusiveDevicesAccessor)[DeviceNo];
-
-  DP("Is the device %d (local ID %d) initialized? %d\n", DeviceNo,
-     Device.RTLDeviceID, Device.IsInit);
-
-  // Init the device if not done before
-  if (!Device.IsInit && Device.initOnce() != OFFLOAD_SUCCESS) {
-    return createStringError(inconvertibleErrorCode(),
-                             "Failed to init device %d\n", DeviceNo);
-  }
-
-  DP("Device %d is ready to use.\n", DeviceNo);
-  return Device;
+  return *(*ExclusiveDevicesAccessor)[DeviceNo];
 }