[Libomptarget] Remove __tgt_image_info and use the ELF directly (#75720)

Summary:
This patch reorganizes a lot of the code used to check for compatibility
with the current environment. The main bulk of this patch involves
moving from using a separate `__tgt_image_info` struct (which just
contains a string for the architecture) to instead simply checking this
information from the ELF directly. Checking information in the ELF is
very inexpensive as creating an ELF file is simply writing a base
pointer.

The main desire to do this was to reorganize everything into the ELF
image. We can then do the majority of these checks without first
initializing the plugin. A future patch will move the first ELF checks
to happen without initializing the plugin so we no longer need to
initialize and plugins that don't have needed images.

This patch also adds a lot more sanity checks for whether or not the ELF
is actually compatible. Such as if the images have a valid ABI, 64-bit
width, executable, etc.

Author: jhuber6

openmp/libomptarget/include/DeviceImage.h

@@ -30,20 +30,13 @@ class DeviceImageTy {
 
   __tgt_bin_desc *BinaryDesc;
   __tgt_device_image Image;
-  __tgt_image_info ImageInfo;
 
 public:
   DeviceImageTy(__tgt_bin_desc &BinaryDesc, __tgt_device_image &Image);
 
   __tgt_device_image &getExecutableImage() { return Image; }
-  __tgt_image_info &getImageInfo() { return ImageInfo; }
   __tgt_bin_desc &getBinaryDesc() { return *BinaryDesc; }
 
-  llvm::StringRef
-  getArch(llvm::StringRef DefaultArch = llvm::StringRef()) const {
-    return ImageInfo.Arch ? ImageInfo.Arch : DefaultArch;
-  }
-
   auto entries() { return llvm::make_pointee_range(OffloadEntries); }
 };
 

openmp/libomptarget/include/Shared/APITypes.h

@@ -46,11 +46,6 @@ struct __tgt_device_info {
   void *Device = nullptr;
 };
 
-/// This struct contains information about a given image.
-struct __tgt_image_info {
-  const char *Arch;
-};
-
 /// This struct is a record of all the host code that may be offloaded to a
 /// target.
 struct __tgt_bin_desc {

openmp/libomptarget/include/Shared/PluginAPI.h

@@ -35,12 +35,6 @@ int32_t __tgt_rtl_number_of_devices(void);
 // having to load the library, which can be expensive.
 int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image);
 
-// This provides the same functionality as __tgt_rtl_is_valid_binary except we
-// also use additional information to determine if the image is valid. This
-// allows us to determine if an image has a compatible architecture.
-int32_t __tgt_rtl_is_valid_binary_info(__tgt_device_image *Image,
-                                       __tgt_image_info *Info);
-
 // Return an integer other than zero if the data can be exchaned from SrcDevId
 // to DstDevId. If it is data exchangable, the device plugin should provide
 // function to move data from source device to destination device directly.

openmp/libomptarget/include/Shared/PluginAPI.inc

@@ -15,7 +15,6 @@
 
 PLUGIN_API_HANDLE(init_plugin, true);
 PLUGIN_API_HANDLE(is_valid_binary, true);
-PLUGIN_API_HANDLE(is_valid_binary_info, false);
 PLUGIN_API_HANDLE(is_data_exchangable, false);
 PLUGIN_API_HANDLE(number_of_devices, true);
 PLUGIN_API_HANDLE(init_device, true);

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

@@ -23,6 +23,7 @@
 #include "Shared/Debug.h"
 #include "Shared/Environment.h"
 #include "Shared/Utils.h"
+#include "Utils/ELF.h"
 
 #include "GlobalHandler.h"
 #include "OpenMP/OMPT/Callback.h"
@@ -3015,7 +3016,15 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
   uint16_t getMagicElfBits() const override { return ELF::EM_AMDGPU; }
 
   /// Check whether the image is compatible with an AMDGPU device.
-  Expected<bool> isImageCompatible(__tgt_image_info *Info) const override {
+  Expected<bool> isELFCompatible(StringRef Image) const override {
+    // Get the associated architecture and flags from the ELF.
+    auto ElfOrErr =
+        ELF64LEObjectFile::create(MemoryBufferRef(Image, /*Identifier=*/""),
+                                  /*InitContent=*/false);
+    if (!ElfOrErr)
+      return ElfOrErr.takeError();
+    std::optional<StringRef> Processor = ElfOrErr->tryGetCPUName();
+
     for (hsa_agent_t Agent : KernelAgents) {
       std::string Target;
       auto Err = utils::iterateAgentISAs(Agent, [&](hsa_isa_t ISA) {
@@ -3038,7 +3047,9 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
       if (Err)
         return std::move(Err);
 
-      if (!utils::isImageCompatibleWithEnv(Info, Target))
+      if (!utils::isImageCompatibleWithEnv(Processor ? *Processor : "",
+                                           ElfOrErr->getPlatformFlags(),
+                                           Target))
         return false;
     }
     return true;

openmp/libomptarget/plugins-nextgen/amdgpu/utils/UtilitiesRTL.h

@@ -13,6 +13,7 @@
 #include <cstdint>
 
 #include "Shared/Debug.h"
+#include "Utils/ELF.h"
 
 #include "omptarget.h"
 
@@ -58,92 +59,58 @@ uint32_t getImplicitArgsSize(uint16_t Version) {
              : sizeof(AMDGPUImplicitArgsTy);
 }
 
-/// Parse a TargetID to get processor arch and feature map.
-/// Returns processor subarch.
-/// Returns TargetID features in \p FeatureMap argument.
-/// If the \p TargetID contains feature+, FeatureMap it to true.
-/// If the \p TargetID contains feature-, FeatureMap it to false.
-/// If the \p TargetID does not contain a feature (default), do not map it.
-StringRef parseTargetID(StringRef TargetID, StringMap<bool> &FeatureMap) {
-  if (TargetID.empty())
-    return llvm::StringRef();
-
-  auto ArchFeature = TargetID.split(":");
-  auto Arch = ArchFeature.first;
-  auto Features = ArchFeature.second;
-  if (Features.empty())
-    return Arch;
-
-  if (Features.contains("sramecc+")) {
-    FeatureMap.insert(std::pair<StringRef, bool>("sramecc", true));
-  } else if (Features.contains("sramecc-")) {
-    FeatureMap.insert(std::pair<StringRef, bool>("sramecc", false));
-  }
-  if (Features.contains("xnack+")) {
-    FeatureMap.insert(std::pair<StringRef, bool>("xnack", true));
-  } else if (Features.contains("xnack-")) {
-    FeatureMap.insert(std::pair<StringRef, bool>("xnack", false));
-  }
-
-  return Arch;
-}
-
-/// Check if an image is compatible with current system's environment.
-bool isImageCompatibleWithEnv(const __tgt_image_info *Info,
-                              StringRef EnvTargetID) {
-  llvm::StringRef ImageTargetID(Info->Arch);
-
-  // Compatible in case of exact match.
-  if (ImageTargetID == EnvTargetID) {
-    DP("Compatible: Exact match \t[Image: %s]\t:\t[Env: %s]\n",
-       ImageTargetID.data(), EnvTargetID.data());
-    return true;
-  }
-
-  // Incompatible if Archs mismatch.
-  StringMap<bool> ImgMap, EnvMap;
-  StringRef ImgArch = utils::parseTargetID(ImageTargetID, ImgMap);
-  StringRef EnvArch = utils::parseTargetID(EnvTargetID, EnvMap);
-
-  // Both EnvArch and ImgArch can't be empty here.
-  if (EnvArch.empty() || ImgArch.empty() || !ImgArch.contains(EnvArch)) {
-    DP("Incompatible: Processor mismatch \t[Image: %s]\t:\t[Env: %s]\n",
-       ImageTargetID.data(), EnvTargetID.data());
+/// Check if an image is compatible with current system's environment. The
+/// system environment is given as a 'target-id' which has the form:
+///
+/// <target-id> := <processor> ( ":" <target-feature> ( "+" | "-" ) )*
+///
+/// If a feature is not specific as '+' or '-' it is assumed to be in an 'any'
+/// and is compatible with either '+' or '-'. The HSA runtime returns this
+/// information using the target-id, while we use the ELF header to determine
+/// these features.
+inline bool isImageCompatibleWithEnv(StringRef ImageArch, uint32_t ImageFlags,
+                                     StringRef EnvTargetID) {
+  StringRef EnvArch = EnvTargetID.split(":").first;
+
+  // Trivial check if the base processors match.
+  if (EnvArch != ImageArch)
     return false;
-  }
 
-  // Incompatible if image has more features than the environment,
-  // irrespective of type or sign of features.
-  if (ImgMap.size() > EnvMap.size()) {
-    DP("Incompatible: Image has more features than the Environment \t[Image: "
-       "%s]\t:\t[Env: %s]\n",
-       ImageTargetID.data(), EnvTargetID.data());
-    return false;
+  // Check if the image is requesting xnack on or off.
+  switch (ImageFlags & EF_AMDGPU_FEATURE_XNACK_V4) {
+  case EF_AMDGPU_FEATURE_XNACK_OFF_V4:
+    // The image is 'xnack-' so the environment must be 'xnack-'.
+    if (!EnvTargetID.contains("xnack-"))
+      return false;
+    break;
+  case EF_AMDGPU_FEATURE_XNACK_ON_V4:
+    // The image is 'xnack+' so the environment must be 'xnack+'.
+    if (!EnvTargetID.contains("xnack+"))
+      return false;
+    break;
+  case EF_AMDGPU_FEATURE_XNACK_UNSUPPORTED_V4:
+  case EF_AMDGPU_FEATURE_XNACK_ANY_V4:
+  default:
+    break;
   }
 
-  // Compatible if each target feature specified by the environment is
-  // compatible with target feature of the image. The target feature is
-  // compatible if the iamge does not specify it (meaning Any), or if it
-  // specifies it with the same value (meaning On or Off).
-  for (const auto &ImgFeature : ImgMap) {
-    auto EnvFeature = EnvMap.find(ImgFeature.first());
-    if (EnvFeature == EnvMap.end() ||
-        (EnvFeature->first() == ImgFeature.first() &&
-         EnvFeature->second != ImgFeature.second)) {
-      DP("Incompatible: Value of Image's non-ANY feature is not matching with "
-         "the Environment's non-ANY feature \t[Image: %s]\t:\t[Env: %s]\n",
-         ImageTargetID.data(), EnvTargetID.data());
+  // Check if the image is requesting sramecc on or off.
+  switch (ImageFlags & EF_AMDGPU_FEATURE_SRAMECC_V4) {
+  case EF_AMDGPU_FEATURE_SRAMECC_OFF_V4:
+    // The image is 'sramecc-' so the environment must be 'sramecc-'.
+    if (!EnvTargetID.contains("sramecc-"))
       return false;
-    }
+    break;
+  case EF_AMDGPU_FEATURE_SRAMECC_ON_V4:
+    // The image is 'sramecc+' so the environment must be 'sramecc+'.
+    if (!EnvTargetID.contains("sramecc+"))
+      return false;
+    break;
+  case EF_AMDGPU_FEATURE_SRAMECC_UNSUPPORTED_V4:
+  case EF_AMDGPU_FEATURE_SRAMECC_ANY_V4:
+    break;
   }
 
-  // Image is compatible if all features of Environment are:
-  //   - either, present in the Image's features map with the same sign,
-  //   - or, the feature is missing from Image's features map i.e. it is
-  //   set to ANY
-  DP("Compatible: Target IDs are compatible \t[Image: %s]\t:\t[Env: %s]\n",
-     ImageTargetID.data(), EnvTargetID.data());
-
   return true;
 }
 

openmp/libomptarget/plugins-nextgen/common/include/PluginInterface.h

@@ -1065,10 +1065,14 @@ struct GenericPluginTy {
     return isValidDeviceId(SrcDeviceId) && isValidDeviceId(DstDeviceId);
   }
 
+  /// 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;
+
   /// 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.
-  virtual Expected<bool> isImageCompatible(__tgt_image_info *Info) const = 0;
+  virtual Expected<bool> isELFCompatible(StringRef Image) const = 0;
 
   /// Indicate whether the plugin supports empty images.
   virtual bool supportsEmptyImages() const { return false; }

openmp/libomptarget/plugins-nextgen/common/src/Utils/ELF.h renamed to openmp/libomptarget/plugins-nextgen/common/include/Utils/ELF.h

@@ -21,9 +21,11 @@
 namespace utils {
 namespace elf {
 
-/// Return non-zero, if the given \p image is an ELF object, which
-/// e_machine matches \p target_id; return zero otherwise.
-int32_t checkMachine(__tgt_device_image *Image, uint16_t TargetId);
+/// Returns true or false if the \p Buffer is an ELF file.
+bool isELF(llvm::StringRef Buffer);
+
+/// Checks if the given \p Object is a valid ELF matching the e_machine value.
+llvm::Expected<bool> checkMachine(llvm::StringRef Object, uint16_t EMachine);
 
 /// Returns a pointer to the given \p Symbol inside of an ELF object.
 llvm::Expected<const void *> getSymbolAddress(

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

@@ -1632,6 +1632,26 @@ 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));
+
+  // First check if this image is a regular ELF file.
+  if (!utils::elf::isELF(Buffer))
+    return false;
+
+  // Check if this image is an ELF with a matching machine value.
+  auto MachineOrErr = utils::elf::checkMachine(Buffer, getMagicElfBits());
+  if (!MachineOrErr)
+    return MachineOrErr.takeError();
+
+  if (!*MachineOrErr)
+    return false;
+
+  // Perform plugin-dependent checks for the specific architecture if needed.
+  return isELFCompatible(Buffer);
+}
+
 const bool llvm::omp::target::plugin::libomptargetSupportsRPC() {
 #ifdef LIBOMPTARGET_RPC_SUPPORT
   return true;
@@ -1659,44 +1679,26 @@ int32_t __tgt_rtl_init_plugin() {
 }
 
 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.
   if (!Plugin::isActive())
     return false;
 
-  if (utils::elf::checkMachine(TgtImage, Plugin::get().getMagicElfBits()))
-    return true;
-
-  return Plugin::get().getJIT().checkBitcodeImage(*TgtImage);
-}
-
-int32_t __tgt_rtl_is_valid_binary_info(__tgt_device_image *TgtImage,
-                                       __tgt_image_info *Info) {
-  if (!Plugin::isActive())
-    return false;
-
-  if (!__tgt_rtl_is_valid_binary(TgtImage))
+  // Check if this is a valid ELF with a matching machine and processor.
+  auto MatchOrErr = Plugin::get().checkELFImage(*TgtImage);
+  if (Error Err = MatchOrErr.takeError()) {
+    [[maybe_unused]] std::string ErrStr = toString(std::move(Err));
+    DP("Failure to check validity of image %p: %s", TgtImage, ErrStr.c_str());
     return false;
-
-  // A subarchitecture was not specified. Assume it is compatible.
-  if (!Info->Arch)
+  } else if (*MatchOrErr) {
     return true;
-
-  // Check the compatibility with all the available devices. Notice the
-  // devices may not be initialized yet.
-  auto CompatibleOrErr = Plugin::get().isImageCompatible(Info);
-  if (!CompatibleOrErr) {
-    // This error should not abort the execution, so we just inform the user
-    // through the debug system.
-    std::string ErrString = toString(CompatibleOrErr.takeError());
-    DP("Failure to check whether image %p is valid: %s\n", TgtImage,
-       ErrString.data());
-    return false;
   }
 
-  bool Compatible = *CompatibleOrErr;
-  DP("Image is %scompatible with current environment: %s\n",
-     (Compatible) ? "" : "not", Info->Arch);
+  // Check if this is a valid LLVM-IR file with matching triple.
+  if (Plugin::get().getJIT().checkBitcodeImage(*TgtImage))
+    return true;
 
-  return Compatible;
+  return false;
 }
 
 int32_t __tgt_rtl_supports_empty_images() {