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 83f6e8d76fec75..eda6a4fd541e96 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 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. @@ -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 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 550ebc9c28b250..913721a15d7131 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 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) { StringRef Buffer(reinterpret_cast(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 191afa345641a1..f72007849e36e4 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(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" @@ -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;