[Libomptarget] Remove __tgt_image_info and use the ELF directly

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

@@ -1062,10 +1062,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 → 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

@@ -1624,6 +1624,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;
@@ -1651,44 +1671,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() {