From dc6e68abc69fdac9e7475adeaf598af9eafe4e4d Mon Sep 17 00:00:00 2001 From: Ron Lieberman Date: Sun, 26 May 2024 12:16:58 -0500 Subject: [PATCH] [Offload] Only initialize a plugin if it is needed (#92765) relands 21f3a609 [Offload] Only initialize a plugin if it is needed (#92765) includes f284af [Offload][Fix] Fix lazy initialization with multiple images includes fix in patchset 3 to resolve 3 smoke tests: currently fails smoke: multi-image multi-image3 targetid_multi_image Change-Id: Ic1890f6e765fd9a91d00f1b85e36f26382b8c608 --- 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 | 42 ++++++++++++++++--- offload/src/PluginManager.cpp | 34 ++++++++++----- 5 files changed, 71 insertions(+), 37 deletions(-) diff --git a/offload/plugins-nextgen/common/include/JIT.h b/offload/plugins-nextgen/common/include/JIT.h index b22197b8920838..4414926a6178f2 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 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 197f26d8ffa3c8..ecd47827000b9f 100644 --- a/offload/plugins-nextgen/common/include/PluginInterface.h +++ b/offload/plugins-nextgen/common/include/PluginInterface.h @@ -1164,6 +1164,10 @@ struct GenericPluginTy { /// given target. Returns true if the \p Image is compatible with the plugin. Expected checkELFImage(StringRef Image) const; + /// Return true if the \p Image can be compiled to run on the platform's + /// target architecture. + Expected 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. @@ -1186,8 +1190,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); /// Checks if the image is not supported. void check_invalid_image(__tgt_device_image *InvalidImage); @@ -1352,6 +1359,9 @@ struct GenericPluginTy { bool isEagerMaps); 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 9d58e6060646ba..9dbba1459839d1 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 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 e1d65e600611fa..2572a3d8cf9401 100644 --- a/offload/plugins-nextgen/common/src/PluginInterface.cpp +++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp @@ -27,6 +27,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" @@ -1607,6 +1608,7 @@ Error GenericPluginTy::init() { if (!NumDevicesOrErr) return NumDevicesOrErr.takeError(); + Initialized = true; NumDevices = *NumDevicesOrErr; if (NumDevices == 0) return Plugin::success(); @@ -1692,14 +1694,27 @@ Expected GenericPluginTy::checkELFImage(StringRef Image) const { if (!MachineOrErr) return MachineOrErr.takeError(); - if (!*MachineOrErr) + return MachineOrErr; +} + +Expected 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) { auto T = logger::log(__func__, Image); int32_t R = [&]() { StringRef Buffer(reinterpret_cast(Image->ImageStart), @@ -1719,10 +1734,25 @@ 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::offload_binary: { + // 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 4ac80fdd79046e..f8d1fbec50ad18 100644 --- a/offload/src/PluginManager.cpp +++ b/offload/src/PluginManager.cpp @@ -36,15 +36,8 @@ void PluginManager::init() { // Attempt to create an instance of each supported plugin. #define PLUGIN_TARGET(Name) \ do { \ - auto Plugin = std::unique_ptr(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(createPlugin_##Name())); \ } while (false); #include "Shared/Targets.def" @@ -199,6 +192,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) { + // Extract the exectuable image and extra information if availible. + for (int32_t i = 0; i < Desc->NumDeviceImages; ++i) { + if (Plugin->is_initialized()) + continue; + + 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]); @@ -217,7 +231,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;