[Libomptarget] Rework image checking further (#76120)

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.

Author: jhuber6

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

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

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();
 }

openmp/libomptarget/plugins-nextgen/common/src/PluginInterface.cpp

@@ -1632,24 +1632,21 @@ 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();
 
   if (!*MachineOrErr)
     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() {

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=*/""),