[Openmp-commits] [openmp] [Libomptarget] Rework image checking further (PR #76120)
Joseph Huber via Openmp-commits
openmp-commits at lists.llvm.org
Thu Dec 21 06:11:52 PST 2023
https://github.com/jhuber6 updated https://github.com/llvm/llvm-project/pull/76120
>From f92576ca7acbf2a645ba414b72e8ef3bccf1d320 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Wed, 20 Dec 2023 21:31:50 -0600
Subject: [PATCH] [Libomptarget] Rework image checking further
Summary:
In the future, we may have more checks for different kinds of inputs,
e.g. SPIR-V. This patch simply reworks the handling to be more generic
and do the magic detection up-front. The checks inside the routines are
now asserts so we don't spend time checking this stuff over and over
again.
This patch also tweaked the bitcode check. I used a different function
to get the Lazy-IR module now, as it returns the raw expected value
rather than the SM diganostic.
No functionality change intended.
---
.../plugins-nextgen/common/include/JIT.h | 2 +-
.../common/include/PluginInterface.h | 2 +-
.../plugins-nextgen/common/src/JIT.cpp | 24 ++++-----
.../common/src/PluginInterface.cpp | 52 +++++++++++--------
.../plugins-nextgen/common/src/Utils/ELF.cpp | 3 +-
5 files changed, 42 insertions(+), 41 deletions(-)
diff --git a/openmp/libomptarget/plugins-nextgen/common/include/JIT.h b/openmp/libomptarget/plugins-nextgen/common/include/JIT.h
index 3ec4424f856a00..b22197b8920838 100644
--- a/openmp/libomptarget/plugins-nextgen/common/include/JIT.h
+++ b/openmp/libomptarget/plugins-nextgen/common/include/JIT.h
@@ -57,7 +57,7 @@ struct JITEngine {
/// Return true if \p Image is a bitcode image that can be JITed for the given
/// architecture.
- bool checkBitcodeImage(const __tgt_device_image &Image);
+ Expected<bool> checkBitcodeImage(StringRef Buffer) const;
private:
/// Compile the bitcode image \p Image and generate the binary image that can
diff --git a/openmp/libomptarget/plugins-nextgen/common/include/PluginInterface.h b/openmp/libomptarget/plugins-nextgen/common/include/PluginInterface.h
index cf02783d8b3388..b85dc146d86d2f 100644
--- a/openmp/libomptarget/plugins-nextgen/common/include/PluginInterface.h
+++ b/openmp/libomptarget/plugins-nextgen/common/include/PluginInterface.h
@@ -1067,7 +1067,7 @@ struct GenericPluginTy {
/// Top level interface to verify if a given ELF image can be executed on a
/// given target. Returns true if the \p Image is compatible with the plugin.
- Expected<bool> checkELFImage(__tgt_device_image &Image) const;
+ Expected<bool> checkELFImage(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
diff --git a/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp b/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp
index 08080c9d6091bb..7275be4edfca5b 100644
--- a/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/src/JIT.cpp
@@ -330,24 +330,18 @@ JITEngine::process(const __tgt_device_image &Image,
return &Image;
}
-bool JITEngine::checkBitcodeImage(const __tgt_device_image &Image) {
+Expected<bool> JITEngine::checkBitcodeImage(StringRef Buffer) const {
TimeTraceScope TimeScope("Check bitcode image");
- if (!isImageBitcode(Image))
- return false;
-
- StringRef Data(reinterpret_cast<const char *>(Image.ImageStart),
- target::getPtrDiff(Image.ImageEnd, Image.ImageStart));
- auto MB = MemoryBuffer::getMemBuffer(Data, /*BufferName=*/"",
- /*RequiresNullTerminator=*/false);
- if (!MB)
- return false;
+ assert(identify_magic(Buffer) == file_magic::bitcode &&
+ "Input is not bitcode");
LLVMContext Context;
- SMDiagnostic Diagnostic;
- std::unique_ptr<Module> M =
- llvm::getLazyIRModule(std::move(MB), Diagnostic, Context,
- /*ShouldLazyLoadMetadata=*/true);
+ auto ModuleOrErr = getLazyBitcodeModule(MemoryBufferRef(Buffer, ""), Context,
+ /*ShouldLazyLoadMetadata=*/true);
+ if (!ModuleOrErr)
+ return ModuleOrErr.takeError();
+ Module &M = **ModuleOrErr;
- return M && Triple(M->getTargetTriple()).getArch() == TT.getArch();
+ return Triple(M.getTargetTriple()).getArch() == TT.getArch();
}
diff --git a/openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp b/openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp
index 178c60a77ab51f..be9ace571f54fd 100644
--- a/openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp
@@ -1632,16 +1632,13 @@ Error GenericPluginTy::deinitDevice(int32_t DeviceId) {
return Plugin::success();
}
-Expected<bool> GenericPluginTy::checkELFImage(__tgt_device_image &Image) const {
- StringRef Buffer(reinterpret_cast<const char *>(Image.ImageStart),
- target::getPtrDiff(Image.ImageEnd, Image.ImageStart));
-
+Expected<bool> GenericPluginTy::checkELFImage(StringRef Image) const {
// First check if this image is a regular ELF file.
- if (!utils::elf::isELF(Buffer))
+ if (!utils::elf::isELF(Image))
return false;
// Check if this image is an ELF with a matching machine value.
- auto MachineOrErr = utils::elf::checkMachine(Buffer, getMagicElfBits());
+ auto MachineOrErr = utils::elf::checkMachine(Image, getMagicElfBits());
if (!MachineOrErr)
return MachineOrErr.takeError();
@@ -1649,7 +1646,7 @@ Expected<bool> GenericPluginTy::checkELFImage(__tgt_device_image &Image) const {
return false;
// Perform plugin-dependent checks for the specific architecture if needed.
- return isELFCompatible(Buffer);
+ return isELFCompatible(Image);
}
const bool llvm::omp::target::plugin::libomptargetSupportsRPC() {
@@ -1678,27 +1675,38 @@ int32_t __tgt_rtl_init_plugin() {
return OFFLOAD_SUCCESS;
}
-int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *TgtImage) {
- // TODO: We should be able to perform a trivial ELF machine check without
- // initializing the plugin first to save time if the plugin is not needed.
+int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image) {
if (!Plugin::isActive())
return false;
- // Check if this is a valid ELF with a matching machine and processor.
- auto MatchOrErr = Plugin::get().checkELFImage(*TgtImage);
- if (Error Err = MatchOrErr.takeError()) {
+ StringRef Buffer(reinterpret_cast<const char *>(Image->ImageStart),
+ target::getPtrDiff(Image->ImageEnd, Image->ImageStart));
+
+ auto HandleError = [&](Error Err) -> bool {
[[maybe_unused]] std::string ErrStr = toString(std::move(Err));
- DP("Failure to check validity of image %p: %s", TgtImage, ErrStr.c_str());
+ DP("Failure to check validity of image %p: %s", Image, ErrStr.c_str());
+ return false;
+ };
+ switch (identify_magic(Buffer)) {
+ case file_magic::elf:
+ case file_magic::elf_relocatable:
+ case file_magic::elf_executable:
+ case file_magic::elf_shared_object:
+ case file_magic::elf_core: {
+ auto MatchOrErr = Plugin::get().checkELFImage(Buffer);
+ if (Error Err = MatchOrErr.takeError())
+ return HandleError(std::move(Err));
+ return *MatchOrErr;
+ }
+ case file_magic::bitcode: {
+ auto MatchOrErr = Plugin::get().getJIT().checkBitcodeImage(Buffer);
+ if (Error Err = MatchOrErr.takeError())
+ return HandleError(std::move(Err));
+ return *MatchOrErr;
+ }
+ default:
return false;
- } else if (*MatchOrErr) {
- return true;
}
-
- // Check if this is a valid LLVM-IR file with matching triple.
- if (Plugin::get().getJIT().checkBitcodeImage(*TgtImage))
- return true;
-
- return false;
}
int32_t __tgt_rtl_supports_empty_images() {
diff --git a/openmp/libomptarget/plugins-nextgen/common/src/Utils/ELF.cpp b/openmp/libomptarget/plugins-nextgen/common/src/Utils/ELF.cpp
index bdac6c1db5d23a..c84c3bad5def0a 100644
--- a/openmp/libomptarget/plugins-nextgen/common/src/Utils/ELF.cpp
+++ b/openmp/libomptarget/plugins-nextgen/common/src/Utils/ELF.cpp
@@ -37,8 +37,7 @@ bool utils::elf::isELF(StringRef Buffer) {
}
Expected<bool> utils::elf::checkMachine(StringRef Object, uint16_t EMachine) {
- if (!isELF(Object))
- return createError("Input is not an ELF.");
+ assert(isELF(Object) && "Input is not an ELF!");
Expected<ELF64LEObjectFile> ElfOrErr =
ELF64LEObjectFile::create(MemoryBufferRef(Object, /*Identifier=*/""),
More information about the Openmp-commits
mailing list