diff --git a/openmp/libomptarget/include/PluginManager.h b/openmp/libomptarget/include/PluginManager.h index 77684285ddf15..eece7525e25e7 100644 --- a/openmp/libomptarget/include/PluginManager.h +++ b/openmp/libomptarget/include/PluginManager.h @@ -45,24 +45,6 @@ struct PluginAdaptorTy { static llvm::Expected> 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 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 DL); @@ -120,6 +92,11 @@ struct PluginManager { std::make_unique(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 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> PluginAdaptors; + // Mapping of plugin adaptors to offsets in the device table. + llvm::DenseMap DeviceOffsets; + + // Mapping of plugin adaptors to the number of used devices. + llvm::DenseMap DeviceUsed; + + // Set of all device images currently in use. + llvm::DenseSet UsedImages; + /// Executable images and information extracted from the input images passed /// to the runtime. llvm::SmallVector> DeviceImages; diff --git a/openmp/libomptarget/src/PluginManager.cpp b/openmp/libomptarget/src/PluginManager.cpp index 928913275332c..792cae3e3dd58 100644 --- a/openmp/libomptarget/src/PluginManager.cpp +++ b/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(this, UserDevId, PDevI); + auto Device = std::make_unique(&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;