bring in [Libomptarget][NFCI] Move logic out of PluginAdaptorTy (llvm#86971)

Change-Id: If41e44d435e53663cbc6baa27322345b9102b540

Author: ronlieb

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

@@ -80,7 +80,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");
@@ -112,26 +112,27 @@ 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.
@@ -140,7 +141,7 @@ void PluginAdaptorTy::initDevices(PluginManager &PM) {
   // They are surfaced per-device because the related properties
   // are computed as such in the plugins.
   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());
@@ -158,11 +159,9 @@ void PluginAdaptorTy::initDevices(PluginManager &PM) {
   // We do not mix APUs with discrete GPUs. Eager maps is set by a host
   // environment variable.
   bool IsAPU = false;
-  bool SupportsUnifiedMemory = false;
   if (ExclusiveDevicesAccessor->size() > 0) {
     auto &Device = *(*ExclusiveDevicesAccessor)[0];
     IsAPU = Device.checkIfAPU();
-    SupportsUnifiedMemory = Device.supportsUnifiedMemory();
   }
   bool EagerMapsRequested = BoolEnvar("OMPX_EAGER_ZERO_COPY_MAPS", false).get();
 
@@ -171,43 +170,46 @@ 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);
 
   // Eager Zero-Copy Maps makes a "copy" execution turn into
   // an automatic zero-copy. It also applies to unified_shared_memory.
   // It is only available on APUs.
   if (IsAPU && EagerMapsRequested) {
-    PM.addRequirements(OMPX_REQ_EAGER_ZERO_COPY_MAPS);
-    if (!(PM.getRequirements() & OMP_REQ_UNIFIED_SHARED_MEMORY))
-      PM.addRequirements(OMPX_REQ_AUTO_ZERO_COPY);
+    addRequirements(OMPX_REQ_EAGER_ZERO_COPY_MAPS);
+    if (!(getRequirements() & OMP_REQ_UNIFIED_SHARED_MEMORY))
+      addRequirements(OMPX_REQ_AUTO_ZERO_COPY);
   }
 
   // sanity checks for zero-copy depend on specific devices: request it here
   if ((ExclusiveDevicesAccessor->size() > 0) &&
-      ((PM.getRequirements() & OMP_REQ_UNIFIED_SHARED_MEMORY) ||
-       (PM.getRequirements() & OMPX_REQ_AUTO_ZERO_COPY))) {
+      ((getRequirements() & OMP_REQ_UNIFIED_SHARED_MEMORY) ||
+       (getRequirements() & OMPX_REQ_AUTO_ZERO_COPY))) {
     // APUs are assumed to be a homogeneous set of GPUs: ask
     // the first device in the system to run a sanity check.
     auto &Device = *(*ExclusiveDevicesAccessor)[0];
     // just skip checks if no devices are found in the system
     Device.zeroCopySanityChecksAndDiag(
-        (PM.getRequirements() & OMP_REQ_UNIFIED_SHARED_MEMORY),
-        (PM.getRequirements() & OMPX_REQ_AUTO_ZERO_COPY),
-        (PM.getRequirements() & OMPX_REQ_EAGER_ZERO_COPY_MAPS));
+        (getRequirements() & OMP_REQ_UNIFIED_SHARED_MEMORY),
+        (getRequirements() & OMPX_REQ_AUTO_ZERO_COPY),
+        (getRequirements() & OMPX_REQ_EAGER_ZERO_COPY_MAPS));
   }
 
+  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.
@@ -216,8 +218,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, {});
@@ -227,11 +228,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.
     }
   }
@@ -270,7 +271,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();
@@ -288,8 +289,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;
@@ -344,11 +347,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;