Skip to content

Commit

Permalink
[Libomptarget][NFCI] Move logic out of PluginAdaptorTy (#86971)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
jhuber6 committed Mar 29, 2024
1 parent 76ba3c8 commit ea707ba
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 61 deletions.
49 changes: 15 additions & 34 deletions openmp/libomptarget/include/PluginManager.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand All @@ -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;
Expand Down
59 changes: 32 additions & 27 deletions openmp/libomptarget/src/PluginManager.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down Expand Up @@ -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());
Expand All @@ -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.
Expand All @@ -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, {});
Expand All @@ -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.
}
}
Expand Down Expand Up @@ -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();
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit ea707ba

Please sign in to comment.