[Libomptarget] Add support for offloading binaries in libomptarget

[jhuber6]

Summary

The previous path changed the linker wrapper to embed the offloading
binary format inside the target image instead. This will allow us to
more generically bundle metadata with these images, such as requires
clauses or the target architecture it was compiled for.

I wasn't sure how to handle this best, so I introduced a new type that
replaces the old __tgt_device_image struct that we can expand inside
the runtime library. I made the new __tgt_device_binary struct pretty
much the same for now. In the future we could change this struct to
pretty much be the OffloadBinary class in the future.

[jhuber6]

Fixing strcmp.

[jhuber6]

Fix problem with using pointers to thing in a vector, need to reserve the space beforehand so we don't realloc.

[jhuber6]

Typo.

[jhuber6]

Instead of changing the __tgt_device_image, do you think it would be better to just to make a new type that only contains the extracted metadata? We could just pass a pointer to the into struct to the is_valid_binary, might be much cleaner. Then we would make a new __tgt_device_image that points to the actual ELF and use that as we do currently.

[tianshilei1992]

openmp/libomptarget/include/omptarget.h

143

Is it better to define it as follows?

struct __tgt_device_binary {
  __tgt_device_image Image;
  __tgt_image_info Info;
}

[jhuber6]

openmp/libomptarget/include/omptarget.h

143

I did it that way to somewhat reduce noise when changing to the new format. But now that I'm thinking about it, where else will we need this metadata? It might be fine to just parse it out and create new images so we don't change them, e.g.

__tgt_image_info info = getInfoFromImage(image); // Gets the strings out from the offload binary type in the __tgt_device_image`.
__tgt_device_image image = getELFFromImage(image); //Gets the ELF out and makes a new image with begin / end set correctly (Need to make a new one, the pointers are in RO memory).
is_binary_valid(image, info); // Plugin uses info struct if necessary, every other interface is the same.

WDYT?

[tianshilei1992]

openmp/libomptarget/include/omptarget.h

143

That looks neater.
In addition, we also want to have a clear definition for front end as well. I'm not saying the front end is gonna include the header. At least we have a clear definition with proper documentation, instead of nested definition in a function.

[jhuber6]

Chaning to use an additional struct instead. Also implemented the class to parse these.

Let me know if there are runtime considerations here, with the unique pointer allocations and string conversions for the map. The usage of __tgt_is_valid_binary is supposed to be mostly lightweight without invoking the runtime, so let me know if this is too expensive.

[tianshilei1992]

Is it necessary to create __tgt_image_info in libomptarget? If the front end emits a binary compatible with OffloadBinary, all the information a plugin needs is in it right? Can we just include the header and read those information in plugin instead of in libomptarget such that we don't need to change the interface?

openmp/libomptarget/include/OffloadBinary.h

72

you may want to make it private

73

Does it make more sense to return nullptr instead of "" here?

76

you may want to make it private

openmp/libomptarget/src/rtl.cpp

13–15

I'd do the following.

288–289

plus clang-format of course.

[jhuber6]

In D127432#3576385, @tianshilei1992 wrote:

Is it necessary to create __tgt_image_info in libomptarget? If the front end emits a binary compatible with OffloadBinary, all the information a plugin needs is in it right? Can we just include the header and read those information in plugin instead of in libomptarget such that we don't need to change the interface?

I was originally thinking about doing that, but it's much noisier since we'd need to change a lot of code that actually loads the binaries. It would free us from reallocating these structs however. Otherwise we'd need to change each usage of the Image, this just extracts the image at the start and everything else stays the same. I don't know which is the superior approach. I find this is a "cleaner" way to augment the existing functionality while maintaining backwards compatibility, but maybe there's a need for everything else to have access to the metadata in the plugin.

openmp/libomptarget/include/OffloadBinary.h

72

It doesn't modify any content so

73

It's supposed to be a string, so an empty string is more appropriate than a null pointer I feel.

76

It should be private, I'll see if I can do that considering the static create functions calls it.

[jhuber6]

Making suggested changes.

[jdoerfert]

openmp/libomptarget/include/omptargetplugin.h

33

I am unsure if we should/need to create a new api here.

[jhuber6]

Changing the is_valid_binary function to be a new interface. This allows the plguins to optionally define the interface that uses the additional information struct, and we will use it if availible.

[jhuber6]

Changing the vector to a list so we always have stable pointers. Reserving a vector doesn't work because of shared libraries, also storing the extract information in the list as well.

[jhuber6]

ping.

[jhuber6]

A previous revision proposes making libomptarget link with the LLVM libraries. This allows us to use LLVMObject directly without duplicating code. Update this revision with that in mind.

[JonChesterfield]

Looks ok to me. Does propagating const into __tgt_device_image (and possibly further) work to remove the const cast? Iirc the libelf interfaces take void* and libllvm take const, so it's possible we'll be able to gradually move const all the way through the stack.

Strictly I suppose const fields in the type are an abi change, we could static_assert that making them const doesn't change offsetof if feeling paranoid.

[jhuber6]

In D127432#3656370, @JonChesterfield wrote:

Looks ok to me. Does propagating const into __tgt_device_image (and possibly further) work to remove the const cast? Iirc the libelf interfaces take void* and libllvm take const, so it's possible we'll be able to gradually move const all the way through the stack.

Strictly I suppose const fields in the type are an abi change, we could static_assert that making them const doesn't change offsetof if feeling paranoid.

Right now that would require const_cast somewhere else. I believe there are calls to libelf in the x86_64 version of the plugin which requires a char * input. We will probably be able to transition away from those if we are allowed to use LLVM libraries. The data that these images point to are global constants, so writing to them will result in whatever happens when you write to RO memory in an ELF.

[JonChesterfield]

openmp/libomptarget/src/rtl.cpp

458

This turns out to be a bad thing. unregisterLib can be called by the host application exit, at which point the plugin may have been unloaded. Is_valid_binary happens to work ok anyway, though I think that's sketchy, but is_valid_binary_info will call into cuda / hsa after they've been torn down and blow up.

[jhuber6]

openmp/libomptarget/src/rtl.cpp

458

Is there a reason to even call this? We already determined the valid binaries before. We should only need to traverse a list of the valid ones here withotu checking it again.

[JonChesterfield]

openmp/libomptarget/src/rtl.cpp

529

I'm doubtful that calling into the plugin to run destructors is any safer than is_valid_binary_info. I think code is working because libomptarget happens to be torn down before the plugins but nothing particularly ensures that will be the case.

This is however a good place to call a destructor function that we add to the plugins.

So we can add tgt_rtl_init_plugin() and tgt_rtl_dtor_plugin() functions, replace the global state constructors in the plugins with those, call init shortly before init_device and call the dtor from here. That removes this undefined behaviour and fixes the segfault Joseph just noticed.

openmp/libomptarget/include/OffloadBinary.h

+//===---------- OffloadBinary.h - Offload Binary Parser -- C++ ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains a generic class meant to parse data out of an input offloading
+// binary. This binary format embeds an executable or linkable image inside
+// along with some other metadata formatted as a string map.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cstring>
+#include <memory>
+#include <unordered_map>
+
+class OffloadBinary {
+  struct Header {
+    uint8_t Magic[4] = {0x10, 0xFF, 0x10, 0xAD}; // 0x10FF10AD magic bytes.
+    uint32_t Version = 1;                        // Version identifier.
+    uint64_t Size;        // Size in bytes of this entire binary.
+    uint64_t EntryOffset; // Offset of the metadata entry in bytes.
+    uint64_t EntrySize;   // Size of the metadata entry in bytes.
+  };
+
+  struct Entry {
+    uint16_t TheImageKind;   // The kind of the image stored.
+    uint16_t TheOffloadKind; // The producer of this image.
+    uint32_t Flags;          // Additional flags associated with the image.
+    uint64_t StringOffset;   // Offset in bytes to the string map.
+    uint64_t NumStrings;     // Number of entries in the string map.
+    uint64_t ImageOffset;    // Offset in bytes of the actual binary image.
+    uint64_t ImageSize;      // Size in bytes of the binary image.
+  };
+
+  struct StringEntry {
+    uint64_t KeyOffset;   // Offset of the key in the string table.
+    uint64_t ValueOffset; // Offset of the value in the string table.
+  };
+
+public:
+  static bool isValid(char *Buffer) {
+    return !strncmp(Buffer, "\x10\xFF\x10\xAD", sizeof(Header::Magic)) &&
+           reinterpret_cast<Header *>(Buffer)->Version == 1;
+  }
+
+  static std::unique_ptr<OffloadBinary> create(char *Buffer) {
+    if (!isValid(Buffer))
+      return nullptr;
+
+    const Header *TheHeader = reinterpret_cast<const Header *>(Buffer);
+    const Entry *TheEntry =
+        reinterpret_cast<const Entry *>(&Buffer[TheHeader->EntryOffset]);
+
+    return std::make_unique<OffloadBinary>(
+        OffloadBinary(Buffer, TheHeader, TheEntry));
+  }
+
+  uint16_t getImageKind() const { return TheEntry->TheImageKind; }
+  uint16_t getOffloadKind() const { return TheEntry->TheOffloadKind; }
+  uint32_t getVersion() const { return TheHeader->Version; }
+  uint32_t getFlags() const { return TheEntry->Flags; }
+  uint64_t getSize() const { return TheHeader->Size; }
+
+  const char *getTriple() const { return getString("triple"); }
+  const char *getArch() const { return getString("arch"); }
+  char *getImageStart() const { return &Buffer[TheEntry->ImageOffset]; }
+  char *getImageEnd() const { return getImageStart() + getImageSize(); }
+  uint64_t getImageSize() const { return TheEntry->ImageSize; }
+
+private:
+  OffloadBinary(char *Buffer, const Header *TheHeader, const Entry *TheEntry)
+      : Buffer(Buffer), TheHeader(TheHeader), TheEntry(TheEntry) {
+    const StringEntry *StringMapBegin =
+        reinterpret_cast<const StringEntry *>(&Buffer[TheEntry->StringOffset]);
+    for (uint64_t I = 0, E = TheEntry->NumStrings; I != E; ++I) {
+      const char *Key = &Buffer[StringMapBegin[I].KeyOffset];
+      StringData[Key] = &Buffer[StringMapBegin[I].ValueOffset];
+    }
+  }
+
+  const char *getString(const char *Key) const {
+    return (StringData.count(Key)) ? StringData.at(Key) : "";
+  }
+
+  /// Map from keys to offsets in the binary.
+  std::unordered_map<std::string, const char *> StringData;
+  /// Raw pointer to the MemoryBufferRef for convenience.
+  char *Buffer;
+  /// Location of the header within the binary.
+  const Header *TheHeader;
+  /// Location of the metadata entries within the binary.
+  const Entry *TheEntry;
+};

openmp/libomptarget/include/device.h

 /// Struct for the data required to handle plugins
 struct PluginManager {
   PluginManager(bool UseEventsForAtomicTransfers)
       : UseEventsForAtomicTransfers(UseEventsForAtomicTransfers) {}
 
   /// RTLs identified on the host
   RTLsTy RTLs;
 
+  /// Executable images extracted from the input images passed to the runtime.
+  std::vector<__tgt_device_image> Images;
+
   /// Devices associated with RTLs
   std::vector<std::unique_ptr<DeviceTy>> Devices;
   std::mutex RTLsMtx; ///< For RTLs and Devices
 
   /// Translation table retreived from the binary
   HostEntriesBeginToTransTableTy HostEntriesBeginToTransTable;
   std::mutex TrlTblMtx; ///< For Translation Table
   /// Host offload entries in order of image registration

openmp/libomptarget/include/omptarget.h

 /// This struct is a record of the device image information
 struct __tgt_device_image {
   void *ImageStart;                  // Pointer to the target code start
   void *ImageEnd;                    // Pointer to the target code end
   __tgt_offload_entry *EntriesBegin; // Begin of table with all target entries
   __tgt_offload_entry *EntriesEnd;   // End of table (non inclusive)
 };
 
+/// 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 {
   int32_t NumDeviceImages;           // Number of device types supported
   __tgt_device_image *DeviceImages;  // Array of device images (1 per dev. type)
   __tgt_offload_entry *HostEntriesBegin; // Begin of table with all host entries
   __tgt_offload_entry *HostEntriesEnd;   // End of table (non inclusive)
 };
 
 /// This struct contains the offload entries identified by the target runtime
 struct __tgt_target_table {
   __tgt_offload_entry *EntriesBegin; // Begin of the table with all the entries
   __tgt_offload_entry
       *EntriesEnd; // End of the table with all the entries (non inclusive)
 };
 
 // clang-format on

openmp/libomptarget/include/omptargetplugin.h

 // target RTL.
 int32_t __tgt_rtl_number_of_devices(void);
 
 // Return an integer different from zero if the provided device image can be
 // supported by the runtime. The functionality is similar to comparing the
 // result of __tgt__rtl__load__binary to NULL. However, this is meant to be a
 // lightweight query to determine if the RTL is suitable for an image without
 // having to load the library, which can be expensive.
-int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image);
+int32_t __tgt_rtl_is_valid_binary(__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.
 int32_t __tgt_rtl_is_data_exchangable(int32_t SrcDevId, int32_t DstDevId);
 
 // Return an integer other than zero if the plugin can handle images which do
 // not contain target regions and global variables (but can contain other

openmp/libomptarget/include/rtl.h

 #include <string>
 #include <vector>
 
 // Forward declarations.
 struct DeviceTy;
 struct __tgt_bin_desc;
 
 struct RTLInfoTy {
-  typedef int32_t(is_valid_binary_ty)(void *);
+  typedef int32_t(is_valid_binary_ty)(void *, void *);
   typedef int32_t(is_data_exchangable_ty)(int32_t, int32_t);
   typedef int32_t(number_of_devices_ty)();
   typedef int32_t(init_device_ty)(int32_t);
   typedef int32_t(deinit_device_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,

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

 namespace core {
 hsa_status_t allow_access_to_all_gpu_agents(void *ptr) {
   return hsa_amd_agents_allow_access(DeviceInfo.HSAAgents.size(),
                                      &DeviceInfo.HSAAgents[0], NULL, ptr);
 }
 } // namespace core
 
 extern "C" {
-int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *image) {
+int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *image,
+                                  __tgt_image_info *info) {
   return elf_machine_id_is_amdgcn(image);
 }
 
 int __tgt_rtl_number_of_devices() {
   // If the construction failed, no methods are safe to call
   if (DeviceInfo.ConstructionSucceeded) {
     return DeviceInfo.NumberOfDevices;
   } else {

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 DeviceRTLTy DeviceRTL;
 } // namespace
 
 // Exposed library API function
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *image) {
+int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *image,
+                                  __tgt_image_info *info) {
   return elf_check_machine(image, /* EM_CUDA */ 190);
 }
 
 int32_t __tgt_rtl_number_of_devices() { return DeviceRTL.getNumOfDevices(); }
 
 int64_t __tgt_rtl_init_requires(int64_t RequiresFlags) {
   DP("Init requires flags to %" PRId64 "\n", RequiresFlags);
   DeviceRTL.setRequiresFlag(RequiresFlags);

openmp/libomptarget/plugins/generic-elf-64bit/src/rtl.cpp

 };
 
 static RTLDeviceInfoTy DeviceInfo(NUMBER_OF_DEVICES);
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *image) {
+int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *image,
+                                  __tgt_image_info *info) {
 // If we don't have a valid ELF ID we can just fail.
 #if TARGET_ELF_ID < 1
   return 0;
 #else
   return elf_check_machine(image, TARGET_ELF_ID);
 #endif
 }
 

openmp/libomptarget/plugins/remote/src/rtl.cpp

 int32_t __tgt_rtl_register_lib(__tgt_bin_desc *Desc) {
   return Manager->registerLib(Desc);
 }
 
 int32_t __tgt_rtl_unregister_lib(__tgt_bin_desc *Desc) {
   return Manager->unregisterLib(Desc);
 }
 
-int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image) {
+int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image,
+                                  __tgt_image_info *Info) {
   return Manager->isValidBinary(Image);
 }
 
 int32_t __tgt_rtl_number_of_devices() { return Manager->getNumberOfDevices(); }
 
 int32_t __tgt_rtl_init_device(int32_t DeviceId) {
   return Manager->initDevice(DeviceId);
 }

openmp/libomptarget/plugins/ve/src/rtl.cpp

 // target RTL.
 int32_t __tgt_rtl_number_of_devices(void) { return DeviceInfo.NodeIds.size(); }
 
 // Return an integer different from zero if the provided device image can be
 // supported by the runtime. The functionality is similar to comparing the
 // result of __tgt__rtl__load__binary to NULL. However, this is meant to be a
 // lightweight query to determine if the RTL is suitable for an image without
 // having to load the library, which can be expensive.
-int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image) {
+int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image,
+                                  __tgt_image_info *Info) {
 #if TARGET_ELF_ID < 1
   return 0;
 #else
   return elf_check_machine(Image, TARGET_ELF_ID);
 #endif
 }
 
 // Initialize the specified device. In case of success return 0; otherwise

openmp/libomptarget/src/rtl.cpp

 //===----------- rtl.cpp - Target independent OpenMP target RTL -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Functionality for handling RTL plugins.
 //
 //===----------------------------------------------------------------------===//
 
 #include "rtl.h"
 #include "device.h"
 #include "private.h"
 
+#include "OffloadBinary.h"
+
 #include <cassert>
 #include <cstdlib>
 #include <cstring>
 #include <dlfcn.h>
 #include <mutex>
 #include <string>
 
 // List of all plugins that can support offloading.
for (__tgt_offload_entry *entry = img->EntriesBegin;
       if (entry->flags & OMP_DECLARE_TARGET_LINK) {
         DP("The \"link\" attribute is not yet supported!\n");
       }
     }
     Device.PendingGlobalsMtx.unlock();
   }
 }
 
+static __tgt_device_image getExecutableImage(__tgt_device_image *Image) {
+  auto Binary = OffloadBinary::create(static_cast<char *>(Image->ImageStart));
+  if (!Binary)
+    return *Image;
+
+  return {Binary->getImageStart(), Binary->getImageEnd(), Image->EntriesBegin,
+          Image->EntriesEnd};
+}
+
+static __tgt_image_info getImageInfo(__tgt_device_image *Image) {
+  auto Binary = OffloadBinary::create(static_cast<char *>(Image->ImageStart));
+  if (!Binary)
+    return __tgt_image_info{};
+
+  return __tgt_image_info{Binary->getArch()};
+}
+
 void RTLsTy::RegisterRequires(int64_t flags) {
   // TODO: add more elaborate check.
   // Minimal check: only set requires flags if previous value
   // is undefined. This ensures that only the first call to this
   // function will set the requires flags. All subsequent calls
   // will be checked for compatibility.
   assert(flags != OMP_REQ_UNDEFINED &&
          "illegal undefined flag for requires directive!");

 
 void RTLsTy::initAllRTLs() {
   for (auto &R : AllRTLs)
     initRTLonce(R);
 }
 
 void RTLsTy::RegisterLib(__tgt_bin_desc *desc) {
   PM->RTLsMtx.lock();
+
+  // Reserve enough space for the new images so we have stable pointers.
+  PM->Images.reserve(desc->NumDeviceImages);
+
   // Register the images with the RTLs that understand them, if any.
   for (int32_t i = 0; i < desc->NumDeviceImages; ++i) {
     // Obtain the image.
-    __tgt_device_image *img = &desc->DeviceImages[i];
+    PM->Images.emplace_back(getExecutableImage(&desc->DeviceImages[i]));
+    __tgt_device_image *img = &PM->Images.back();
+    __tgt_image_info info = getImageInfo(&desc->DeviceImages[i]);
 
     RTLInfoTy *FoundRTL = nullptr;
 
     // Scan the RTLs that have associated images until we find one that supports
     // the current image.
     for (auto &R : AllRTLs) {
-      if (!R.is_valid_binary(img)) {
+      if (!R.is_valid_binary(img, &info)) {
         DP("Image " DPxMOD " is NOT compatible with RTL %s!\n",
            DPxPTR(img->ImageStart), R.RTLName.c_str());
         continue;
       }
 
       DP("Image " DPxMOD " is compatible with RTL %s!\n",
          DPxPTR(img->ImageStart), R.RTLName.c_str());
 

 
 void RTLsTy::UnregisterLib(__tgt_bin_desc *desc) {
   DP("Unloading target library!\n");
 
   PM->RTLsMtx.lock();
   // Find which RTL understands each image, if any.
   for (int32_t i = 0; i < desc->NumDeviceImages; ++i) {
     // Obtain the image.
-    __tgt_device_image *img = &desc->DeviceImages[i];
+    __tgt_device_image *img = &PM->Images[i];
+    __tgt_image_info info = getImageInfo(&desc->DeviceImages[i]);
 
     RTLInfoTy *FoundRTL = NULL;
 
     // Scan the RTLs that have associated images until we find one that supports
     // the current image. We only need to scan RTLs that are already being used.
     for (auto *R : UsedRTLs) {
 
       assert(R->isUsed && "Expecting used RTLs.");
 
-      if (!R->is_valid_binary(img)) {
+      if (!R->is_valid_binary(img, &info)) {
         DP("Image " DPxMOD " is NOT compatible with RTL " DPxMOD "!\n",
            DPxPTR(img->ImageStart), DPxPTR(R->LibraryHandler));
         continue;
       }
 
       DP("Image " DPxMOD " is compatible with RTL " DPxMOD "!\n",
          DPxPTR(img->ImageStart), DPxPTR(R->LibraryHandler));
 
       FoundRTL = R;
 
if (TransTable != PM->HostEntriesBeginToTransTable.end()) {
     DP("Translation table for descriptor " DPxMOD " cannot be found, probably "
        "it has been already removed.\n",
        DPxPTR(desc->HostEntriesBegin));
   }
 
   PM->TblMapMtx.unlock();
 
   // TODO: Remove RTL and the devices it manages if it's not used anymore?
   // TODO: Write some RTL->unload_image(...) function?
 
   DP("Done unregistering library!\n");
 }