[llvm] [Offload] Only initialize a plugin if it is needed (PR #92765)
Joseph Huber via llvm-commits
llvm-commits at lists.llvm.org
Thu May 23 06:06:29 PDT 2024
https://github.com/jhuber6 updated https://github.com/llvm/llvm-project/pull/92765
>From de2b654317657539d780dda8a6d56e0a5a41f81a Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Mon, 20 May 2024 10:06:35 -0500
Subject: [PATCH] [Offload] Only initialize a plugin if it is needed
Summary:
Initializing the plugins requires initializing the runtime like CUDA or
HSA. This has a considerable overhead on most platforms, so we should
only actually initialize a plugin if it is needed by any image that is
loaded.
---
offload/plugins-nextgen/common/include/JIT.h | 4 ---
.../common/include/PluginInterface.h | 12 ++++++-
offload/plugins-nextgen/common/src/JIT.cpp | 16 ---------
.../common/src/PluginInterface.cpp | 34 +++++++++++++++----
offload/src/PluginManager.cpp | 34 +++++++++++++------
5 files changed, 63 insertions(+), 37 deletions(-)
diff --git a/offload/plugins-nextgen/common/include/JIT.h b/offload/plugins-nextgen/common/include/JIT.h
index b22197b892083..4414926a6178f 100644
--- a/offload/plugins-nextgen/common/include/JIT.h
+++ b/offload/plugins-nextgen/common/include/JIT.h
@@ -55,10 +55,6 @@ struct JITEngine {
process(const __tgt_device_image &Image,
target::plugin::GenericDeviceTy &Device);
- /// Return true if \p Image is a bitcode image that can be JITed for the given
- /// architecture.
- Expected<bool> checkBitcodeImage(StringRef Buffer) const;
-
private:
/// Compile the bitcode image \p Image and generate the binary image that can
/// be loaded to the target device of the triple \p Triple architecture \p
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index 83f6e8d76fec7..eda6a4fd541e9 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -1052,6 +1052,10 @@ struct GenericPluginTy {
/// given target. Returns true if the \p Image is compatible with the plugin.
Expected<bool> checkELFImage(StringRef Image) const;
+ /// Return true if the \p Image can be compiled to run on the platform's
+ /// target architecture.
+ Expected<bool> checkBitcodeImage(StringRef Image) const;
+
/// Indicate if an image is compatible with the plugin devices. Notice that
/// this function may be called before actually initializing the devices. So
/// we could not move this function into GenericDeviceTy.
@@ -1066,8 +1070,11 @@ struct GenericPluginTy {
public:
// TODO: This plugin interface needs to be cleaned up.
+ /// Returns true if the plugin has been initialized.
+ int32_t is_initialized() const;
+
/// Returns non-zero if the provided \p Image can be executed by the runtime.
- int32_t is_valid_binary(__tgt_device_image *Image);
+ int32_t is_valid_binary(__tgt_device_image *Image, bool Initialized = true);
/// Initialize the device inside of the plugin.
int32_t init_device(int32_t DeviceId);
@@ -1187,6 +1194,9 @@ struct GenericPluginTy {
void **KernelPtr);
private:
+ /// Indicates if the platform runtime has been fully initialized.
+ bool Initialized = false;
+
/// Number of devices available for the plugin.
int32_t NumDevices = 0;
diff --git a/offload/plugins-nextgen/common/src/JIT.cpp b/offload/plugins-nextgen/common/src/JIT.cpp
index 9d58e6060646b..9dbba1459839d 100644
--- a/offload/plugins-nextgen/common/src/JIT.cpp
+++ b/offload/plugins-nextgen/common/src/JIT.cpp
@@ -323,19 +323,3 @@ JITEngine::process(const __tgt_device_image &Image,
return &Image;
}
-
-Expected<bool> JITEngine::checkBitcodeImage(StringRef Buffer) const {
- TimeTraceScope TimeScope("Check bitcode image");
-
- assert(identify_magic(Buffer) == file_magic::bitcode &&
- "Input is not bitcode");
-
- LLVMContext Context;
- auto ModuleOrErr = getLazyBitcodeModule(MemoryBufferRef(Buffer, ""), Context,
- /*ShouldLazyLoadMetadata=*/true);
- if (!ModuleOrErr)
- return ModuleOrErr.takeError();
- Module &M = **ModuleOrErr;
-
- return Triple(M.getTargetTriple()).getArch() == TT.getArch();
-}
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index 550ebc9c28b25..913721a15d713 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -24,6 +24,7 @@
#include "omp-tools.h"
#endif
+#include "llvm/Bitcode/BitcodeReader.h"
#include "llvm/Frontend/OpenMP/OMPConstants.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/JSON.h"
@@ -1495,6 +1496,7 @@ Error GenericPluginTy::init() {
if (!NumDevicesOrErr)
return NumDevicesOrErr.takeError();
+ Initialized = true;
NumDevices = *NumDevicesOrErr;
if (NumDevices == 0)
return Plugin::success();
@@ -1578,14 +1580,27 @@ Expected<bool> GenericPluginTy::checkELFImage(StringRef Image) const {
if (!MachineOrErr)
return MachineOrErr.takeError();
- if (!*MachineOrErr)
+ return MachineOrErr;
+}
+
+Expected<bool> GenericPluginTy::checkBitcodeImage(StringRef Image) const {
+ if (identify_magic(Image) != file_magic::bitcode)
return false;
- // Perform plugin-dependent checks for the specific architecture if needed.
- return isELFCompatible(Image);
+ LLVMContext Context;
+ auto ModuleOrErr = getLazyBitcodeModule(MemoryBufferRef(Image, ""), Context,
+ /*ShouldLazyLoadMetadata=*/true);
+ if (!ModuleOrErr)
+ return ModuleOrErr.takeError();
+ Module &M = **ModuleOrErr;
+
+ return Triple(M.getTargetTriple()).getArch() == getTripleArch();
}
-int32_t GenericPluginTy::is_valid_binary(__tgt_device_image *Image) {
+int32_t GenericPluginTy::is_initialized() const { return Initialized; }
+
+int32_t GenericPluginTy::is_valid_binary(__tgt_device_image *Image,
+ bool Initialized) {
StringRef Buffer(reinterpret_cast<const char *>(Image->ImageStart),
target::getPtrDiff(Image->ImageEnd, Image->ImageStart));
@@ -1603,10 +1618,17 @@ int32_t GenericPluginTy::is_valid_binary(__tgt_device_image *Image) {
auto MatchOrErr = checkELFImage(Buffer);
if (Error Err = MatchOrErr.takeError())
return HandleError(std::move(Err));
- return *MatchOrErr;
+ if (!Initialized || !*MatchOrErr)
+ return *MatchOrErr;
+
+ // Perform plugin-dependent checks for the specific architecture if needed.
+ auto CompatibleOrErr = isELFCompatible(Buffer);
+ if (Error Err = CompatibleOrErr.takeError())
+ return HandleError(std::move(Err));
+ return *CompatibleOrErr;
}
case file_magic::bitcode: {
- auto MatchOrErr = getJIT().checkBitcodeImage(Buffer);
+ auto MatchOrErr = checkBitcodeImage(Buffer);
if (Error Err = MatchOrErr.takeError())
return HandleError(std::move(Err));
return *MatchOrErr;
diff --git a/offload/src/PluginManager.cpp b/offload/src/PluginManager.cpp
index 191afa345641a..f72007849e36e 100644
--- a/offload/src/PluginManager.cpp
+++ b/offload/src/PluginManager.cpp
@@ -34,15 +34,8 @@ void PluginManager::init() {
// Attempt to create an instance of each supported plugin.
#define PLUGIN_TARGET(Name) \
do { \
- auto Plugin = std::unique_ptr<GenericPluginTy>(createPlugin_##Name()); \
- if (auto Err = Plugin->init()) { \
- [[maybe_unused]] std::string InfoMsg = toString(std::move(Err)); \
- DP("Failed to init plugin: %s\n", InfoMsg.c_str()); \
- } else { \
- DP("Registered plugin %s with %d visible device(s)\n", \
- Plugin->getName(), Plugin->number_of_devices()); \
- Plugins.emplace_back(std::move(Plugin)); \
- } \
+ Plugins.emplace_back( \
+ std::unique_ptr<GenericPluginTy>(createPlugin_##Name())); \
} while (false);
#include "Shared/Targets.def"
@@ -160,6 +153,27 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
if (Entry.flags == OMP_REGISTER_REQUIRES)
PM->addRequirements(Entry.data);
+ // Initialize all the plugins that have associated images.
+ for (auto &Plugin : Plugins) {
+ if (Plugin->is_initialized())
+ continue;
+
+ // Extract the exectuable image and extra information if availible.
+ for (int32_t i = 0; i < Desc->NumDeviceImages; ++i) {
+ if (!Plugin->is_valid_binary(&Desc->DeviceImages[i],
+ /*Initialized=*/false))
+ continue;
+
+ if (auto Err = Plugin->init()) {
+ [[maybe_unused]] std::string InfoMsg = toString(std::move(Err));
+ DP("Failed to init plugin: %s\n", InfoMsg.c_str());
+ } else {
+ DP("Registered plugin %s with %d visible device(s)\n",
+ Plugin->getName(), Plugin->number_of_devices());
+ }
+ }
+ }
+
// Extract the exectuable image and extra information if availible.
for (int32_t i = 0; i < Desc->NumDeviceImages; ++i)
PM->addDeviceImage(*Desc, Desc->DeviceImages[i]);
@@ -177,7 +191,7 @@ void PluginManager::registerLib(__tgt_bin_desc *Desc) {
if (!R.number_of_devices())
continue;
- if (!R.is_valid_binary(Img)) {
+ if (!R.is_valid_binary(Img, /*Initialized=*/true)) {
DP("Image " DPxMOD " is NOT compatible with RTL %s!\n",
DPxPTR(Img->ImageStart), R.getName());
continue;
More information about the llvm-commits
mailing list