[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

159

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

159

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

159

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

71 ↗

(On Diff #436009)

you may want to make it private

72 ↗

(On Diff #436009)

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

75 ↗

(On Diff #436009)

you may want to make it private

openmp/libomptarget/src/rtl.cpp

13–15

I'd do the following.

287–288

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

71 ↗

(On Diff #436009)

It doesn't modify any content so

72 ↗

(On Diff #436009)

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

75 ↗

(On Diff #436009)

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

478

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

478

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

557

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/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 and information extracted from the input images passed
+  /// to the runtime.
+  std::list<std::pair<__tgt_device_image, __tgt_image_info>> 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

 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);
 
+// 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.
 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
 // functions)

openmp/libomptarget/include/rtl.h

 #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_info_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,
struct RTLInfoTy {
   void *LibraryHandler = nullptr;
 
 #ifdef OMPTARGET_DEBUG
   std::string RTLName;
 #endif
 
   // Functions implemented in the RTL.
   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;
   init_device_ty *init_device = nullptr;
   deinit_device_ty *deinit_device = 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;

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 "llvm/Object/OffloadBinary.h"
+
 #include <cassert>
 #include <cstdlib>
 #include <cstring>
 #include <dlfcn.h>
 #include <mutex>
 #include <string>
 
 // List of all plugins that can support offloading.

 #ifdef OMPTARGET_DEBUG
     R.RTLName = Name;
 #endif
 
     DP("Registering RTL %s supporting %d devices!\n", R.RTLName.c_str(),
        R.NumberOfDevices);
 
     // Optional functions
+    *((void **)&R.is_valid_binary_info) =
+        dlsym(DynlibHandle, "__tgt_rtl_is_valid_binary_info");
     *((void **)&R.deinit_device) =
         dlsym(DynlibHandle, "__tgt_rtl_deinit_device");
     *((void **)&R.init_requires) =
         dlsym(DynlibHandle, "__tgt_rtl_init_requires");
     *((void **)&R.data_submit_async) =
         dlsym(DynlibHandle, "__tgt_rtl_data_submit_async");
     *((void **)&R.data_retrieve_async) =
         dlsym(DynlibHandle, "__tgt_rtl_data_retrieve_async");
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) {
+  llvm::StringRef ImageStr(static_cast<char *>(Image->ImageStart),
+                           static_cast<char *>(Image->ImageEnd) -
+                               static_cast<char *>(Image->ImageStart));
+  auto BinaryOrErr =
+      llvm::object::OffloadBinary::create(llvm::MemoryBufferRef(ImageStr, ""));
+  if (!BinaryOrErr) {
+    llvm::consumeError(BinaryOrErr.takeError());
+    return *Image;
+  }
+
+  void *Begin = const_cast<void *>(
+      static_cast<const void *>((*BinaryOrErr)->getImage().bytes_begin()));
+  void *End = const_cast<void *>(
+      static_cast<const void *>((*BinaryOrErr)->getImage().bytes_end()));
+
+  return {Begin, End, Image->EntriesBegin, Image->EntriesEnd};
+}
+
+static __tgt_image_info getImageInfo(__tgt_device_image *Image) {
+  llvm::StringRef ImageStr(static_cast<char *>(Image->ImageStart),
+                           static_cast<char *>(Image->ImageEnd) -
+                               static_cast<char *>(Image->ImageStart));
+  auto BinaryOrErr =
+      llvm::object::OffloadBinary::create(llvm::MemoryBufferRef(ImageStr, ""));
+  if (!BinaryOrErr) {
+    llvm::consumeError(BinaryOrErr.takeError());
+    return __tgt_image_info{};
+  }
+
+  return __tgt_image_info{(*BinaryOrErr)->getArch().data()};
+}
+
 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();
+
+  // Extract the exectuable image and extra information if availible.
+  for (int32_t i = 0; i < Desc->NumDeviceImages; ++i)
+    PM->Images.emplace_back(getExecutableImage(&Desc->DeviceImages[i]),
+                            getImageInfo(&Desc->DeviceImages[i]));
+
   // 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];
+  for (auto &ImageAndInfo : PM->Images) {
+    // Obtain the image and information that was previously extracted.
+    __tgt_device_image *Img = &ImageAndInfo.first;
+    __tgt_image_info *Info = &ImageAndInfo.second;
 
     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_info) {
+        if (!R.is_valid_binary_info(Img, Info)) {
+          DP("Image " DPxMOD " is NOT compatible with RTL %s!\n",
+             DPxPTR(Img->ImageStart), R.RTLName.c_str());
+          continue;
+        }
+      } else if (!R.is_valid_binary(Img)) {
         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::registerLib(__tgt_bin_desc *Desc) {
   DP("Done registering entries!\n");
 }
 
 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];
+  for (auto &ImageAndInfo : PM->Images) {
+    // Obtain the image and information that was previously extracted.
+    __tgt_device_image *Img = &ImageAndInfo.first;
+    __tgt_image_info *Info = &ImageAndInfo.second;
 
     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)) {
-        DP("Image " DPxMOD " is NOT compatible with RTL " DPxMOD "!\n",
-           DPxPTR(Img->ImageStart), DPxPTR(R->LibraryHandler));
+      if (R->is_valid_binary_info) {
+        if (!R->is_valid_binary_info(Img, Info)) {
+          DP("Image " DPxMOD " is NOT compatible with RTL %s!\n",
+             DPxPTR(Img->ImageStart), R->RTLName.c_str());
+          continue;
+        }
+      } else if (!R->is_valid_binary(Img)) {
+        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 " 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");
 }