[Libomptarget] Remove __tgt_image_info and use the ELF directly (llvm#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.

Reverts following patches for reimplementation:
   53859b6
   b7e137c

Change-Id: I37fc701d76d99464d7bd877168cbf912f622b634

Author: saiislam

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

@@ -56,12 +56,6 @@ void __tgt_rtl_set_up_env(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.
@@ -249,8 +243,6 @@ int32_t __tgt_rtl_initialize_record_replay(int32_t DeviceId, int64_t MemorySize,
 
 bool __tgt_rtl_requested_prepopulate_gpu_page_table();
 
-bool __tgt_rtl_exists_valid_binary_for_RTL(void *, void *);
-
 bool __tgt_rtl_is_system_supporting_managed_memory();
 
 int32_t __tgt_rtl_launch_kernel_sync(int32_t, void *, void **, ptrdiff_t *,

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);
@@ -48,7 +47,6 @@ PLUGIN_API_HANDLE(data_notify_mapped, false);
 PLUGIN_API_HANDLE(data_notify_unmapped, false);
 PLUGIN_API_HANDLE(set_device_offset, false);
 PLUGIN_API_HANDLE(initialize_record_replay, false);
-PLUGIN_API_HANDLE(exists_valid_binary_for_RTL, true);
 PLUGIN_API_HANDLE(has_apu_device, true);
 PLUGIN_API_HANDLE(has_USM_capable_dGPU, true);
 PLUGIN_API_HANDLE(are_allocations_for_maps_on_apus_disabled, true);

openmp/libomptarget/include/omptarget.h

@@ -322,12 +322,6 @@ void __tgt_register_lib(__tgt_bin_desc *Desc);
 
 /// Initialize all RTLs at once
 void __tgt_init_all_rtls();
-/// adds an image information struct, called for each image
-void __tgt_register_image_info(__tgt_image_info *imageInfo);
-
-/// gets pointer to image information for specified image number
-/// Returns nullptr for apps built with old version of compiler
-__tgt_image_info *__tgt_get_image_info(uint32_t image_num);
 
 /// removes a target shared library from the target execution image
 void __tgt_unregister_lib(__tgt_bin_desc *Desc);

openmp/libomptarget/include/rtl.h

@@ -13,13 +13,198 @@
 #ifndef _OMPTARGET_RTL_H
 #define _OMPTARGET_RTL_H
 
-#include "llvm/ADT/SmallVector.h"
-
 #include "omptarget.h"
+#include <llvm/ADT/DenseSet.h>
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/DynamicLibrary.h"
 
 #include <cstdint>
+#include <list>
+#include <map>
+#include <mutex>
+#include <string>
+#include <vector>
 #include <map>
 
+// Forward declarations.
+struct DeviceTy;
+struct __tgt_bin_desc;
+
+struct RTLInfoTy {
+  typedef int32_t(init_plugin_ty)();
+  typedef int32_t(deinit_plugin_ty)();
+  typedef int32_t(is_valid_binary_ty)(void *);
+  typedef int32_t(is_valid_binary_info_ty)(void *, void *);
+  typedef int32_t(is_data_exchangable_ty)(int32_t, int32_t);
+  typedef int32_t(number_of_devices_ty)();
+  typedef bool(has_apu_device_ty)();
+  typedef bool(has_USM_capable_dGPU_ty)();
+  typedef bool(are_allocations_for_maps_on_apus_disabled_ty)();
+  typedef bool(requested_prepopulate_gpu_page_table_ty)();
+  typedef bool(is_no_maps_check_ty)();
+  typedef bool(is_fine_grained_memory_enabled_ty)();
+  typedef bool(is_system_supporting_managed_memory_ty)();
+  typedef int32_t(init_device_ty)(int32_t);
+  typedef int32_t(deinit_device_ty)(int32_t);
+  typedef int32_t(number_of_team_procs_ty)(int32_t);
+  typedef __tgt_target_table *(load_binary_ty)(int32_t, void *);
+  typedef void *(data_alloc_ty)(int32_t, int64_t, void *, int32_t);
+
+  typedef int32_t(data_submit_ty)(int32_t, void *, void *, int64_t);
+  typedef int32_t(data_submit_async_ty)(int32_t, void *, void *, int64_t,
+                                        __tgt_async_info *);
+  typedef int32_t(data_retrieve_ty)(int32_t, void *, void *, int64_t);
+  typedef int32_t(data_retrieve_async_ty)(int32_t, void *, void *, int64_t,
+                                          __tgt_async_info *);
+  typedef int32_t(data_exchange_ty)(int32_t, void *, int32_t, void *, int64_t);
+  typedef int32_t(data_exchange_async_ty)(int32_t, void *, int32_t, void *,
+                                          int64_t, __tgt_async_info *);
+  typedef int32_t(data_delete_ty)(int32_t, void *, int32_t);
+  typedef int32_t(launch_kernel_sync_ty)(int32_t, void *, void **, ptrdiff_t *,
+                                         const KernelArgsTy *);
+  typedef int32_t(launch_kernel_ty)(int32_t, void *, void **, ptrdiff_t *,
+                                    const KernelArgsTy *, __tgt_async_info *);
+  typedef int64_t(init_requires_ty)(int64_t);
+  typedef int32_t(synchronize_ty)(int32_t, __tgt_async_info *);
+  typedef int32_t(query_async_ty)(int32_t, __tgt_async_info *);
+  typedef int32_t (*register_lib_ty)(__tgt_bin_desc *);
+  typedef int32_t(supports_empty_images_ty)();
+  typedef void(print_device_info_ty)(int32_t);
+  typedef void(set_info_flag_ty)(uint32_t);
+  typedef int32_t(create_event_ty)(int32_t, void **);
+  typedef int32_t(record_event_ty)(int32_t, void *, __tgt_async_info *);
+  typedef int32_t(wait_event_ty)(int32_t, void *, __tgt_async_info *);
+  typedef int32_t(sync_event_ty)(int32_t, void *);
+  typedef int32_t(destroy_event_ty)(int32_t, void *);
+  typedef int(set_coarse_grain_mem_region_ty)(int32_t, void *, int64_t);
+  typedef int(prepopulate_page_table_ty)(int32_t, void *, int64_t);
+  typedef int32_t(query_coarse_grain_mem_region_ty)(int32_t, void *, int64_t);
+  typedef int32_t(enable_access_to_all_agents_ty)(void *, int32_t);
+  typedef int32_t(release_async_info_ty)(int32_t, __tgt_async_info *);
+  typedef int32_t(init_async_info_ty)(int32_t, __tgt_async_info **);
+  typedef int64_t(init_device_into_ty)(int64_t, __tgt_device_info *,
+                                       const char **);
+  typedef int32_t(data_lock_ty)(int32_t, void *, int64_t, void **);
+  typedef int32_t(data_unlock_ty)(int32_t, void *);
+  typedef int32_t(data_notify_mapped_ty)(int32_t, void *, int64_t);
+  typedef int32_t(data_notify_unmapped_ty)(int32_t, void *);
+  typedef int32_t(set_device_offset_ty)(int32_t);
+  typedef int32_t(activate_record_replay_ty)(int32_t, uint64_t, void *, bool,
+                                             bool);
+  typedef void(set_up_env_ty)(void);
+
+  int32_t Idx = -1;             // RTL index, index is the number of devices
+                                // of other RTLs that were registered before,
+                                // i.e. the OpenMP index of the first device
+                                // to be registered with this RTL.
+  int32_t NumberOfDevices = -1; // Number of devices this RTL deals with.
+
+  std::unique_ptr<llvm::sys::DynamicLibrary> LibraryHandler;
+
+#ifdef OMPTARGET_DEBUG
+  std::string RTLName;
+#endif
+
+  // Functions implemented in the RTL.
+  init_plugin_ty *init_plugin = nullptr;
+  deinit_plugin_ty *deinit_plugin = nullptr;
+  is_valid_binary_ty *is_valid_binary = nullptr;
+  is_valid_binary_info_ty *is_valid_binary_info = nullptr;
+  is_data_exchangable_ty *is_data_exchangable = nullptr;
+  number_of_devices_ty *number_of_devices = nullptr;
+  has_apu_device_ty *has_apu_device = nullptr;
+  has_USM_capable_dGPU_ty *has_USM_capable_dGPU = nullptr;
+  are_allocations_for_maps_on_apus_disabled_ty
+      *are_allocations_for_maps_on_apus_disabled = nullptr;
+  requested_prepopulate_gpu_page_table_ty
+      *requested_prepopulate_gpu_page_table = nullptr;
+  is_no_maps_check_ty *is_no_maps_check = nullptr;
+  is_fine_grained_memory_enabled_ty *is_fine_grained_memory_enabled = nullptr;
+  is_system_supporting_managed_memory_ty *is_system_supporting_managed_memory =
+      nullptr;
+  init_device_ty *init_device = nullptr;
+  deinit_device_ty *deinit_device = nullptr;
+  number_of_team_procs_ty *number_of_team_procs = nullptr;
+  load_binary_ty *load_binary = nullptr;
+  data_alloc_ty *data_alloc = nullptr;
+  data_submit_ty *data_submit = nullptr;
+  data_submit_async_ty *data_submit_async = nullptr;
+  data_retrieve_ty *data_retrieve = nullptr;
+  data_retrieve_async_ty *data_retrieve_async = nullptr;
+  data_exchange_ty *data_exchange = nullptr;
+  data_exchange_async_ty *data_exchange_async = nullptr;
+  data_delete_ty *data_delete = nullptr;
+  launch_kernel_sync_ty *launch_kernel_sync = nullptr;
+  launch_kernel_ty *launch_kernel = nullptr;
+  init_requires_ty *init_requires = nullptr;
+  synchronize_ty *synchronize = nullptr;
+  query_async_ty *query_async = nullptr;
+  register_lib_ty register_lib = nullptr;
+  register_lib_ty unregister_lib = nullptr;
+  supports_empty_images_ty *supports_empty_images = nullptr;
+  set_info_flag_ty *set_info_flag = nullptr;
+  print_device_info_ty *print_device_info = nullptr;
+  create_event_ty *create_event = nullptr;
+  record_event_ty *record_event = nullptr;
+  wait_event_ty *wait_event = nullptr;
+  sync_event_ty *sync_event = nullptr;
+  destroy_event_ty *destroy_event = nullptr;
+  init_async_info_ty *init_async_info = nullptr;
+  init_device_into_ty *init_device_info = nullptr;
+  release_async_info_ty *release_async_info = nullptr;
+  data_lock_ty *data_lock = nullptr;
+  data_unlock_ty *data_unlock = nullptr;
+  set_coarse_grain_mem_region_ty *set_coarse_grain_mem_region = nullptr;
+  prepopulate_page_table_ty *prepopulate_page_table = nullptr;
+  query_coarse_grain_mem_region_ty *query_coarse_grain_mem_region = nullptr;
+  enable_access_to_all_agents_ty *enable_access_to_all_agents = nullptr;
+  data_notify_mapped_ty *data_notify_mapped = nullptr;
+  data_notify_unmapped_ty *data_notify_unmapped = nullptr;
+  set_device_offset_ty *set_device_offset = nullptr;
+  activate_record_replay_ty *activate_record_replay = nullptr;
+  set_up_env_ty *set_up_env = nullptr;
+
+  // Are there images associated with this RTL.
+  bool IsUsed = false;
+
+  llvm::DenseSet<const __tgt_device_image *> UsedImages;
+
+  // Mutex for thread-safety when calling RTL interface functions.
+  // It is easier to enforce thread-safety at the libomptarget level,
+  // so that developers of new RTLs do not have to worry about it.
+  std::mutex Mtx;
+};
+
+/// RTLs identified in the system.
+struct RTLsTy {
+  // List of the detected runtime libraries.
+  std::list<RTLInfoTy> AllRTLs;
+
+  // Array of pointers to the detected runtime libraries that have compatible
+  // binaries.
+  llvm::SmallVector<RTLInfoTy *> UsedRTLs;
+
+  explicit RTLsTy() = default;
+
+  // Register the clauses of the requires directive.
+  void registerRequires(int64_t Flags);
+
+  // Initialize RTL if it has not been initialized
+  void initRTLonce(RTLInfoTy &RTL);
+
+  // Initialize all RTLs
+  void initAllRTLs();
+
+  // Register a shared library with all (compatible) RTLs.
+  void registerLib(__tgt_bin_desc *Desc);
+
+  // Unregister a shared library from all RTLs.
+  void unregisterLib(__tgt_bin_desc *Desc);
+
+  // not thread-safe, called from global constructor (i.e. once)
+  void loadRTLs();
+};
+
 /// Map between the host entry begin and the translation table. Each
 /// registered library gets one TranslationTable. Use the map from
 /// __tgt_offload_entry so that we may quickly determine whether we

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

@@ -24,6 +24,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"
@@ -4027,17 +4028,15 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     }
   }
 
-  void checkInvalidImage(__tgt_image_info *Info,
-                              __tgt_device_image *TgtImage) override final {
-
-    utils::checkImageCompatibilityWithSystemXnackMode(TgtImage,
-                                                      IsXnackEnabled());
-  }
-
   /// Check whether the image is compatible with an AMDGPU device.
-  Expected<bool>
-  isImageCompatible(__tgt_image_info *Info,
-                    __tgt_device_image *TgtImage) 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;
@@ -4061,10 +4060,17 @@ 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 true;
     }
 
+    // Check if the system's XNACK mode matches the one required by the
+    // image. Print a warning if not.
+    // utils::checkImageCompatibilityWithSystemXnackMode(TgtImage,
+    //                                                   IsXnackEnabled());
+
     return false;
   }