[Libomptarget][NFCI] Move logic out of PluginAdaptorTy

Summary:
This patch removes most of the special handling from the
`PluginAdaptorTy` in preparation for changing this to be the
`GenericPluginTy`. Doing this requires that the OpenMP specific handling
of stuff like device offsets be contained within the OpenMP plugin
manager. Generally this was uninvasive expect for the change to tracking
the offset and size of the used devices. The eaiest way I could think to
do this was to use some maps, which double as indicators for which
plugins have devices active. This should not affect the logic.

Author: jhuber6

openmp/libomptarget/include/PluginManager.h

@@ -45,24 +45,6 @@ struct PluginAdaptorTy {
   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 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; }
-
-  /// 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;
-
   /// Name of the shared object file representing the plugin.
   std::string Name;
 
@@ -76,16 +58,6 @@ struct PluginAdaptorTy {
 #include "Shared/PluginAPI.inc"
 #undef PLUGIN_API_HANDLE
 
-  llvm::DenseSet<const __tgt_device_image *> UsedImages;
-
-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);
@@ -120,6 +92,11 @@ struct PluginManager {
         std::make_unique<DeviceImageTy>(TgtBinDesc, TgtDeviceImage));
   }
 
+  /// Initialize as many devices as possible for this plugin adaptor. Devices
+  /// that fail to initialize are ignored. Returns the offset the devices were
+  /// registered at.
+  void initDevices(PluginAdaptorTy &RTL);
+
   /// Return the device presented to the user as device \p DeviceNo if it is
   /// initialized and ready. Otherwise return an error explaining the problem.
   llvm::Expected<DeviceTy &> getDevice(uint32_t DeviceNo);
@@ -169,12 +146,7 @@ struct PluginManager {
     return Devices.getExclusiveAccessor();
   }
 
-  int getNumUsedPlugins() const {
-    int NCI = 0;
-    for (auto &P : PluginAdaptors)
-      NCI += P->isUsed();
-    return NCI;
-  }
+  int getNumUsedPlugins() const { return DeviceOffsets.size(); }
 
   // Initialize all plugins.
   void initAllPlugins();
@@ -195,6 +167,15 @@ struct PluginManager {
   // List of all plugin adaptors, in use or not.
   llvm::SmallVector<std::unique_ptr<PluginAdaptorTy>> PluginAdaptors;
 
+  // Mapping of plugin adaptors to offsets in the device table.
+  llvm::DenseMap<const PluginAdaptorTy *, int32_t> DeviceOffsets;
+
+  // Mapping of plugin adaptors to the number of used devices.
+  llvm::DenseMap<const PluginAdaptorTy *, int32_t> DeviceUsed;
+
+  // Set of all device images currently in use.
+  llvm::DenseSet<const __tgt_device_image *> UsedImages;
+
   /// Executable images and information extracted from the input images passed
   /// to the runtime.
   llvm::SmallVector<std::unique_ptr<DeviceImageTy>> DeviceImages;

openmp/libomptarget/src/PluginManager.cpp

@@ -78,7 +78,7 @@ Error PluginAdaptorTy::init() {
   }
 
   // No devices are supported by this RTL?
-  NumberOfPluginDevices = number_of_devices();
+  int32_t NumberOfPluginDevices = number_of_devices();
   if (!NumberOfPluginDevices) {
     return createStringError(inconvertibleErrorCode(),
                              "No devices supported in this RTL\n");
@@ -110,32 +110,33 @@ void PluginManager::init() {
   DP("RTLs loaded!\n");
 }
 
-void PluginAdaptorTy::initDevices(PluginManager &PM) {
-  if (isUsed())
+void PluginManager::initDevices(PluginAdaptorTy &RTL) {
+  // If this RTL has already been initialized.
+  if (PM->DeviceOffsets.contains(&RTL))
     return;
   TIMESCOPE();
 
   // If this RTL is not already in use, initialize it.
-  assert(getNumberOfPluginDevices() > 0 &&
+  assert(RTL.number_of_devices() > 0 &&
          "Tried to initialize useless plugin adaptor");
 
   // Initialize the device information for the RTL we are about to use.
-  auto ExclusiveDevicesAccessor = PM.getExclusiveDevicesAccessor();
+  auto ExclusiveDevicesAccessor = getExclusiveDevicesAccessor();
 
   // Initialize the index of this RTL and save it in the used RTLs.
-  DeviceOffset = ExclusiveDevicesAccessor->size();
+  int32_t DeviceOffset = ExclusiveDevicesAccessor->size();
 
-  // If possible, set the device identifier offset in the plugin.
-  if (set_device_offset)
-    set_device_offset(DeviceOffset);
+  // Set the device identifier offset in the plugin.
+  RTL.set_device_offset(DeviceOffset);
 
-  int32_t NumPD = getNumberOfPluginDevices();
+  int32_t NumberOfUserDevices = 0;
+  int32_t NumPD = RTL.number_of_devices();
   ExclusiveDevicesAccessor->reserve(DeviceOffset + NumPD);
   // Auto zero-copy is a per-device property. We need to ensure
   // that all devices are suggesting to use it.
   bool UseAutoZeroCopy = !(NumPD == 0);
   for (int32_t PDevI = 0, UserDevId = DeviceOffset; PDevI < NumPD; PDevI++) {
-    auto Device = std::make_unique<DeviceTy>(this, UserDevId, PDevI);
+    auto Device = std::make_unique<DeviceTy>(&RTL, UserDevId, PDevI);
     if (auto Err = Device->init()) {
       DP("Skip plugin known device %d: %s\n", PDevI,
          toString(std::move(Err)).c_str());
@@ -153,20 +154,23 @@ void PluginAdaptorTy::initDevices(PluginManager &PM) {
   // If all devices suggest to use it, change requirment flags to trigger
   // zero-copy behavior when mapping memory.
   if (UseAutoZeroCopy)
-    PM.addRequirements(OMPX_REQ_AUTO_ZERO_COPY);
+    addRequirements(OMPX_REQ_AUTO_ZERO_COPY);
 
+  DeviceOffsets[&RTL] = DeviceOffset;
+  DeviceUsed[&RTL] = NumberOfUserDevices;
   DP("Plugin adaptor " DPxMOD " has index %d, exposes %d out of %d devices!\n",
-     DPxPTR(LibraryHandler.get()), DeviceOffset, NumberOfUserDevices,
-     NumberOfPluginDevices);
+     DPxPTR(RTL.LibraryHandler.get()), DeviceOffset, NumberOfUserDevices,
+     RTL.number_of_devices());
 }
 
 void PluginManager::initAllPlugins() {
   for (auto &R : PluginAdaptors)
-    R->initDevices(*this);
+    initDevices(*R);
 }
 
 static void registerImageIntoTranslationTable(TranslationTable &TT,
-                                              PluginAdaptorTy &RTL,
+                                              int32_t DeviceOffset,
+                                              int32_t NumberOfUserDevices,
                                               __tgt_device_image *Image) {
 
   // same size, as when we increase one, we also increase the other.
@@ -175,8 +179,7 @@ static void registerImageIntoTranslationTable(TranslationTable &TT,
 
   // Resize the Targets Table and Images to accommodate the new targets if
   // required
-  unsigned TargetsTableMinimumSize =
-      RTL.DeviceOffset + RTL.getNumberOfUserDevices();
+  unsigned TargetsTableMinimumSize = DeviceOffset + NumberOfUserDevices;
 
   if (TT.TargetsTable.size() < TargetsTableMinimumSize) {
     TT.DeviceTables.resize(TargetsTableMinimumSize, {});
@@ -186,11 +189,11 @@ static void registerImageIntoTranslationTable(TranslationTable &TT,
   }
 
   // Register the image in all devices for this target type.
-  for (int32_t I = 0; I < RTL.getNumberOfUserDevices(); ++I) {
+  for (int32_t I = 0; I < NumberOfUserDevices; ++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;
-      TT.TargetsTable[RTL.DeviceOffset + I] =
+    if (TT.TargetsImages[DeviceOffset + I] != Image) {
+      TT.TargetsImages[DeviceOffset + I] = Image;
+      TT.TargetsTable[DeviceOffset + I] =
           0; // lazy initialization of target table.
     }
   }
@@ -228,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());
 
-      R.initDevices(*this);
+      PM->initDevices(R);
 
       // Initialize (if necessary) translation table for this library.
       PM->TrlTblMtx.lock();
@@ -246,8 +249,10 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
 
       DP("Registering image " DPxMOD " with RTL %s!\n", DPxPTR(Img->ImageStart),
          R.Name.c_str());
-      registerImageIntoTranslationTable(TransTable, R, Img);
-      R.UsedImages.insert(Img);
+
+      registerImageIntoTranslationTable(TransTable, PM->DeviceOffsets[&R],
+                                        PM->DeviceUsed[&R], Img);
+      PM->UsedImages.insert(Img);
 
       PM->TrlTblMtx.unlock();
       FoundRTL = &R;
@@ -283,11 +288,11 @@ void PluginManager::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 : PM->pluginAdaptors()) {
-      if (!R.isUsed())
+      if (!DeviceOffsets.contains(&R))
         continue;
 
       // Ensure that we do not use any unused images associated with this RTL.
-      if (!R.UsedImages.contains(Img))
+      if (!UsedImages.contains(Img))
         continue;
 
       FoundRTL = &R;