Skip to content

[Libomptarget][NFCI] Move logic out of PluginAdaptorTy #86971

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We always look them up together, use a single map.


// 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