[libsycl] Add device image registration & compatibility check (#187528)

This is part of the SYCL support upstreaming effort. The relevant RFCs
can be found here:


https://discourse.llvm.org/t/rfc-add-full-support-for-the-sycl-programming-model/74080
https://discourse.llvm.org/t/rfc-sycl-runtime-upstreaming/74479

---------

Signed-off-by: Tikhomirova, Kseniya <kseniya.tikhomirova@intel.com>

libsycl/docs/index.rst

@@ -113,3 +113,8 @@ TODO for added SYCL classes
   * add aligned functions (blocked by liboffload support)
   * forward templated funcs to alignment methods (rewrite current impl)
   * handle sub devices once they are implemented (blocked by liboffload support)
+
+
+* general opens:
+
+  * define a way to report errors from object dtors.

libsycl/src/CMakeLists.txt

@@ -17,6 +17,7 @@ function(add_sycl_rt_library LIB_TARGET_NAME LIB_OBJ_NAME LIB_OUTPUT_NAME)
   add_dependencies(${LIB_OBJ_NAME}
     sycl-headers
     LLVMOffload
+    LLVMObject
   )
 
   target_include_directories(${LIB_OBJ_NAME}
@@ -24,6 +25,7 @@ function(add_sycl_rt_library LIB_TARGET_NAME LIB_OBJ_NAME LIB_OUTPUT_NAME)
       ${CMAKE_CURRENT_SOURCE_DIR}
       ${LIBSYCL_BUILD_INCLUDE_DIR}
       $<TARGET_PROPERTY:LLVMOffload,INTERFACE_INCLUDE_DIRECTORIES>
+      ${LLVM_MAIN_INCLUDE_DIR}
   )
 
   set_target_properties(${LIB_TARGET_NAME}
@@ -68,6 +70,7 @@ function(add_sycl_rt_library LIB_TARGET_NAME LIB_OBJ_NAME LIB_OUTPUT_NAME)
       ${CMAKE_DL_LIBS}
       ${CMAKE_THREAD_LIBS_INIT}
       LLVMOffload
+      LLVMObject
   )
 
   set_target_properties(${LIB_TARGET_NAME}
@@ -93,6 +96,7 @@ set(LIBSYCL_SOURCES
     "detail/device_impl.cpp"
     "detail/global_objects.cpp"
     "detail/platform_impl.cpp"
+    "detail/program_manager.cpp"
     "detail/queue_impl.cpp"
     "detail/offload/offload_utils.cpp"
     "detail/offload/offload_topology.cpp"

libsycl/src/detail/device_binary_structures.hpp

@@ -0,0 +1,91 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _LIBSYCL_DEVICE_BINARY_STRUCTURES
+#define _LIBSYCL_DEVICE_BINARY_STRUCTURES
+
+#include <sycl/__impl/detail/config.hpp>
+
+#include <llvm/Frontend/Offloading/Utility.h>
+#include <llvm/Object/OffloadBinary.h>
+
+#include <cstdint>
+
+_LIBSYCL_BEGIN_NAMESPACE_SYCL
+
+namespace detail {
+
+/// Target identification strings.
+///
+/// A device type represented by a particular target
+/// triple requires specific binary images. We need
+/// to map the image type onto the device target triple.
+
+/// Unknown target.
+static constexpr char DeviceBinaryTripleUnknown[] = "unknown-unknown-unknown";
+/// SPIR-V with 64-bit pointers.
+static constexpr char DeviceBinaryTripleSPIRV64[] = "spirv64-unknown-unknown";
+
+/// Device binary descriptor version supported by this library.
+static constexpr uint16_t SupportedDevicyBinaryVersion = 3;
+
+/// This struct is a record of the device binary information.
+///  It must match the __tgt_device_image structure generated by the
+///  compiler when their `Version` fields match.
+struct __sycl_tgt_device_image {
+  uint16_t Version;
+  /// The type of offload model the binary employs. See `OffloadKind`. Only
+  /// OFK_SYCL is supported by libsycl.
+  uint8_t OffloadKind;
+  /// Format of the binary data, see `ImageKind`.
+  uint8_t ImageFormat;
+  /// A null-terminated string representation of the device's target
+  /// architecture. Must hold one of _LIBSYCL_DEVICE_BINARY_TARGET_* values.
+  const char *TripleString;
+  /// A null-terminated string of target- and compiler-specific options
+  /// that are suggested to use to "compile" program at runtime.
+  const char *CompileOptions;
+  /// A null-terminated string of target- and compiler-specific options
+  /// that are suggested to use to "link" program at runtime.
+  const char *LinkOptions;
+  /// Pointer to the target code start.
+  const unsigned char *ImageStart;
+  /// Pointer to the target code end.
+  const unsigned char *ImageEnd;
+  /// The offload entry table
+  llvm::offloading::EntryTy *EntriesBegin;
+  llvm::offloading::EntryTy *EntriesEnd;
+  // TODO: properties are not supported now.
+  /// Array of property sets.
+  void *PropertiesBegin;
+  void *PropertiesEnd;
+};
+
+/// Version of offload binaries descriptor `__sycl_tgt_bin_desc` supported by
+/// libsycl.
+static constexpr uint16_t SupportedOffloadBinaryVersion = 1;
+
+/// This struct is a record of all the device code that may be offloaded.
+/// It must match the `__tgt_bin_desc` structure generated by
+/// the compiler when their `Version` fields match.
+struct __sycl_tgt_bin_desc {
+  /// Version of the structure.
+  uint16_t Version;
+  /// Number of device binaries in this descriptor.
+  uint16_t NumDeviceBinaries;
+  /// Device binaries data.
+  __sycl_tgt_device_image *DeviceImages;
+  /// The offload entry table (not used, for compatibility with OpenMP).
+  llvm::offloading::EntryTy *HostEntriesBegin;
+  llvm::offloading::EntryTy *HostEntriesEnd;
+};
+
+} // namespace detail
+_LIBSYCL_END_NAMESPACE_SYCL
+
+#endif // _LIBSYCL_DEVICE_BINARY_STRUCTURES

libsycl/src/detail/device_image_wrapper.hpp

@@ -0,0 +1,52 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _LIBSYCL_DEVICE_IMAGE_WRAPPER
+#define _LIBSYCL_DEVICE_IMAGE_WRAPPER
+
+#include <sycl/__impl/detail/config.hpp>
+
+#include <detail/device_binary_structures.hpp>
+
+_LIBSYCL_BEGIN_NAMESPACE_SYCL
+namespace detail {
+
+/// A wrapper of __sycl_tgt_device_image structure to help with its fields
+/// parsing, iteration over data and data transformation.
+class DeviceImageWrapper {
+public:
+  DeviceImageWrapper(const __sycl_tgt_device_image &Bin) : MBin(&Bin) {}
+  // Explicitly delete copy constructor/operator= to avoid unintentional copies.
+  DeviceImageWrapper(const DeviceImageWrapper &) = delete;
+  DeviceImageWrapper &operator=(const DeviceImageWrapper &) = delete;
+
+  DeviceImageWrapper(DeviceImageWrapper &&) = default;
+  DeviceImageWrapper &operator=(DeviceImageWrapper &&) = default;
+
+  ~DeviceImageWrapper() = default;
+
+  /// \return a reference to the corresponding raw __sycl_tgt_device_image
+  /// object.
+  const __sycl_tgt_device_image &getRawData() const { return *get(); }
+
+  /// \return the size of the corresponding device image data in bytes.
+  size_t getSize() const {
+    return static_cast<size_t>(MBin->ImageEnd - MBin->ImageStart);
+  }
+
+protected:
+  const __sycl_tgt_device_image *get() const { return MBin; }
+
+  __sycl_tgt_device_image const *MBin{};
+};
+
+} // namespace detail
+
+_LIBSYCL_END_NAMESPACE_SYCL
+
+#endif // _LIBSYCL_DEVICE_IMAGE_WRAPPER

libsycl/src/detail/device_impl.hpp

@@ -115,7 +115,7 @@ public:
       static_assert(false && "Info descriptor is not properly supported");
   }
 
-  ol_device_handle_t getOLHandle() { return MOffloadDevice; }
+  ol_device_handle_t getHandle() { return MOffloadDevice; }
 
 private:
   ol_device_handle_t MOffloadDevice = {};

libsycl/src/detail/global_objects.cpp

@@ -53,6 +53,3 @@ std::vector<PlatformImplUPtr> &getPlatformCache() {
 
 } // namespace detail
 _LIBSYCL_END_NAMESPACE_SYCL
-
-extern "C" void __sycl_register_lib(void *) {}
-extern "C" void __sycl_unregister_lib(void *) {}

libsycl/src/detail/global_objects.hpp

@@ -16,19 +16,6 @@
 #include <mutex>
 #include <vector>
 
-// +++ Entry points referenced by the offload wrapper object {
-
-/// Executed as a part of current module's (.exe, .dll) static initialization.
-/// Registers device executable images with the runtime.
-extern "C" _LIBSYCL_EXPORT void __sycl_register_lib(void *);
-
-/// Executed as a part of current module's (.exe, .dll) static
-/// de-initialization.
-/// Unregisters device executable images with the runtime.
-extern "C" _LIBSYCL_EXPORT void __sycl_unregister_lib(void *);
-
-// +++ }
-
 _LIBSYCL_BEGIN_NAMESPACE_SYCL
 
 namespace detail {

libsycl/src/detail/kernel_id.hpp

@@ -0,0 +1,77 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _LIBSYCL_KERNEL_ID
+#define _LIBSYCL_KERNEL_ID
+
+#include <sycl/__impl/detail/config.hpp>
+#include <sycl/__impl/detail/obj_utils.hpp>
+
+#include <memory>
+#include <string>
+
+_LIBSYCL_BEGIN_NAMESPACE_SYCL
+
+namespace detail {
+/// The class is the implementation counterpart for sycl::kernel_id, which
+/// represents a kernel identificator.
+class KernelIdImpl {
+public:
+  KernelIdImpl(std::string_view Name) : MName(Name) {}
+  KernelIdImpl() {}
+  /// \return a null-terminated string representing the name of the kernel this
+  /// id stands for.
+  const char *get_name() { return MName.data(); }
+
+private:
+  std::string MName;
+};
+} // namespace detail
+
+// TODO: It is not exported now, but is a part of SYCL spec.
+/// Kernel identifier.
+class kernel_id {
+public:
+  kernel_id() = delete;
+
+  kernel_id(const kernel_id &rhs) = default;
+
+  kernel_id(kernel_id &&rhs) = default;
+
+  kernel_id &operator=(const kernel_id &rhs) = default;
+
+  kernel_id &operator=(kernel_id &&rhs) = default;
+
+  friend bool operator==(const kernel_id &lhs, const kernel_id &rhs) {
+    return lhs.impl == rhs.impl;
+  }
+
+  friend bool operator!=(const kernel_id &lhs, const kernel_id &rhs) {
+    return !(lhs == rhs);
+  }
+
+  /// \returns a null-terminated string that contains the kernel name.
+  const char *get_name() const noexcept { return impl->get_name(); }
+
+private:
+  kernel_id(const char *Name);
+
+  kernel_id(const std::shared_ptr<detail::KernelIdImpl> &Impl)
+      : impl(std::move(Impl)) {}
+
+  std::shared_ptr<detail::KernelIdImpl> impl;
+  friend sycl::detail::ImplUtils;
+};
+
+_LIBSYCL_END_NAMESPACE_SYCL
+
+template <>
+struct std::hash<sycl::kernel_id>
+    : public sycl::detail::HashBase<sycl::kernel_id> {};
+
+#endif // _LIBSYCL_KERNEL_ID

libsycl/src/detail/program_manager.cpp

@@ -0,0 +1,154 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include <detail/program_manager.hpp>
+
+#include <sycl/__impl/exception.hpp>
+
+#include <detail/device_impl.hpp>
+#include <detail/offload/offload_utils.hpp>
+
+#include <cstring>
+
+_LIBSYCL_BEGIN_NAMESPACE_SYCL
+namespace detail {
+
+static inline bool checkFatBinVersion(const __sycl_tgt_bin_desc &FatbinDesc) {
+  return FatbinDesc.Version == SupportedOffloadBinaryVersion;
+}
+
+static inline bool
+checkDeviceImageValidity(const __sycl_tgt_device_image &DeviceImage) {
+  return (DeviceImage.Version == SupportedDevicyBinaryVersion) &&
+         (DeviceImage.OffloadKind == llvm::object::OFK_SYCL) &&
+         (DeviceImage.ImageFormat == llvm::object::IMG_SPIRV);
+}
+
+void ProgramManager::addImages(__sycl_tgt_bin_desc *FatbinDesc) {
+  assert(FatbinDesc && "Device images descriptor can't be nullptr");
+
+  if (!checkFatBinVersion(*FatbinDesc))
+    throw sycl::exception(sycl::make_error_code(sycl::errc::runtime),
+                          "Incompatible version of device images descriptor.");
+  if (!FatbinDesc->NumDeviceBinaries)
+    return;
+
+  std::lock_guard<std::mutex> Guard(MImageCollectionMutex);
+  for (int I = 0; I < FatbinDesc->NumDeviceBinaries; ++I) {
+    const auto &RawDeviceImage = FatbinDesc->DeviceImages[I];
+    if (!checkDeviceImageValidity(RawDeviceImage))
+      throw sycl::exception(sycl::make_error_code(sycl::errc::runtime),
+                            "Incompatible device image.");
+
+    const llvm::offloading::EntryTy *EntriesB = RawDeviceImage.EntriesBegin;
+    const llvm::offloading::EntryTy *EntriesE = RawDeviceImage.EntriesEnd;
+    // Ignore "empty" device image.
+    if (EntriesB == EntriesE)
+      continue;
+
+    std::unique_ptr<DeviceImageWrapper> NewImageWrapper =
+        std::make_unique<DeviceImageWrapper>(RawDeviceImage);
+
+    for (auto EntriesIt = EntriesB; EntriesIt != EntriesE; ++EntriesIt) {
+      auto Name = EntriesIt->SymbolName;
+      auto KernelIDIt = MKernelNameToID.find(Name);
+      if (KernelIDIt == MKernelNameToID.end()) {
+        sycl::kernel_id KernelID =
+            detail::createSyclObjFromImpl<sycl::kernel_id>(
+                std::make_shared<detail::KernelIdImpl>(Name));
+        KernelIDIt = MKernelNameToID.insert(
+            MKernelNameToID.end(),
+            std::make_pair(std::string_view(Name), KernelID));
+      }
+
+      MKernelIDToDevImageJIT.insert(
+          std::make_pair(KernelIDIt->second, NewImageWrapper.get()));
+    }
+
+    MDeviceImageWrappers.insert(
+        std::make_pair(&RawDeviceImage, std::move(NewImageWrapper)));
+  }
+}
+
+void ProgramManager::removeImages(__sycl_tgt_bin_desc *FatbinDesc) {
+  assert(FatbinDesc && "Device images descriptor can't be nullptr");
+
+  if (!checkFatBinVersion(*FatbinDesc) || FatbinDesc->NumDeviceBinaries == 0)
+    return;
+
+  std::lock_guard<std::mutex> Guard(MImageCollectionMutex);
+  for (int I = 0; I < FatbinDesc->NumDeviceBinaries; ++I) {
+    const auto &RawDeviceImage = FatbinDesc->DeviceImages[I];
+
+    auto DevImageIt = MDeviceImageWrappers.find(&RawDeviceImage);
+    if (DevImageIt == MDeviceImageWrappers.end())
+      continue;
+
+    const llvm::offloading::EntryTy *EntriesB = RawDeviceImage.EntriesBegin;
+    const llvm::offloading::EntryTy *EntriesE = RawDeviceImage.EntriesEnd;
+    // Ignore "empty" device image
+    if (EntriesB == EntriesE)
+      continue;
+
+    for (auto EntriesIt = EntriesB; EntriesIt != EntriesE; ++EntriesIt) {
+      if (auto KernelIDIt = MKernelNameToID.find(EntriesIt->SymbolName);
+          KernelIDIt != MKernelNameToID.end()) {
+        MKernelIDToDevImageJIT.erase(KernelIDIt->second);
+        MKernelNameToID.erase(KernelIDIt);
+      }
+    }
+
+    MDeviceImageWrappers.erase(DevImageIt);
+  }
+}
+
+static bool isImageTargetCompatible(const DeviceImageWrapper &Image,
+                                    const DeviceImpl &Device) {
+  sycl::backend BE = Device.getBackend();
+  const char *Target = Image.getRawData().TripleString;
+
+  return (strcmp(Target, DeviceBinaryTripleSPIRV64) == 0) &&
+         (BE == sycl::backend::level_zero);
+}
+
+DeviceImageWrapper *ProgramManager::getDeviceImage(std::string_view KernelName,
+                                                   const kernel_id &KernelID,
+                                                   DeviceImpl &Device) {
+  std::lock_guard<std::mutex> Guard(MImageCollectionMutex);
+  auto [Begin, End] = MKernelIDToDevImageJIT.equal_range(KernelID);
+  if (Begin != End) {
+    bool IsValid{};
+    // TODO: with AOT (not implemented yet), we need to analyze and check
+    // olIsValidBinary for AOT binaries first.
+    for (auto It = Begin; It != End; ++It) {
+      if (isImageTargetCompatible(*It->second, Device)) {
+        callAndThrow(olIsValidBinary, Device.getHandle(),
+                     It->second->getRawData().ImageStart, It->second->getSize(),
+                     &IsValid);
+        if (IsValid)
+          return It->second;
+      }
+    }
+  }
+
+  throw exception(make_error_code(errc::runtime),
+                  "No kernel named " + std::string(KernelName) + " was found");
+}
+
+} // namespace detail
+_LIBSYCL_END_NAMESPACE_SYCL
+
+extern "C" _LIBSYCL_EXPORT void
+__sycl_register_lib(sycl::detail::__sycl_tgt_bin_desc *FatbinDesc) {
+  sycl::detail::ProgramManager::getInstance().addImages(FatbinDesc);
+}
+
+extern "C" _LIBSYCL_EXPORT void
+__sycl_unregister_lib(sycl::detail::__sycl_tgt_bin_desc *FatbinDesc) {
+  sycl::detail::ProgramManager::getInstance().removeImages(FatbinDesc);
+}

libsycl/src/detail/program_manager.hpp

@@ -0,0 +1,97 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _LIBSYCL_PROGRAM_MANAGER
+#define _LIBSYCL_PROGRAM_MANAGER
+
+#include <sycl/__impl/detail/config.hpp>
+
+#include <detail/device_binary_structures.hpp>
+#include <detail/device_image_wrapper.hpp>
+#include <detail/kernel_id.hpp>
+
+#include <mutex>
+#include <unordered_map>
+
+// +++ Entry points referenced by the offload wrapper object {
+
+/// Executed as a part of a module's (.exe, .dll) static initialization.
+/// Registers device executable images with the runtime.
+extern "C" _LIBSYCL_EXPORT void
+__sycl_register_lib(sycl::detail::__sycl_tgt_bin_desc *FatbinDesc);
+
+/// Executed as a part of current module's (.exe, .dll) static
+/// de-initialization.
+/// Unregisters device executable images with the runtime.
+extern "C" _LIBSYCL_EXPORT void
+__sycl_unregister_lib(sycl::detail::__sycl_tgt_bin_desc *FatbinDesc);
+
+// +++ }
+
+_LIBSYCL_BEGIN_NAMESPACE_SYCL
+
+namespace detail {
+
+class DeviceImpl;
+
+/// A class to manage programs and kernels.
+class ProgramManager {
+
+public:
+  static ProgramManager &getInstance() {
+    static ProgramManager PM{};
+    return PM;
+  }
+
+  /// Parses raw device images data and prepares internal structures for
+  /// effective kernel/program creation.
+  /// \param FatbinDesc a record of all the device code that may be offloaded,
+  /// generated by compiler and offloading tools.
+  /// \throw sycl::exception with sycl::errc::runtime if a device image
+  /// descriptor has an incompatible version or if a device image has an
+  /// incompatible version, target or kind.
+  void addImages(__sycl_tgt_bin_desc *FatbinDesc);
+
+  /// Removes all entries of the data in FatbinDesc from internal structures.
+  /// \param FatbinDesc a record of all the device code that may be offloaded,
+  /// generated by compiler and offloading tools. Must match the pointer and
+  /// data passed to addImages.
+  void removeImages(__sycl_tgt_bin_desc *FatbinDesc);
+
+private:
+  ProgramManager() = default;
+  ~ProgramManager() = default;
+  ProgramManager(ProgramManager const &) = delete;
+  ProgramManager &operator=(ProgramManager const &) = delete;
+
+  /// Searches for a device image that contains the requested kernel and is
+  /// compatible with the requested device.
+  /// \param KernelName a null-terminated string representing the name of the
+  /// kernel to obtain a device image for.
+  /// \param KernelID a kernel id matching KernelName.
+  /// \param DeviceImpl a device with which device image must be compatible.
+  /// \throw sycl::exception with sycl::errc::runtime if the device image
+  /// validation failed in liboffload or if no compatible image was found.
+  DeviceImageWrapper *getDeviceImage(std::string_view KernelName,
+                                     const kernel_id &KernelID,
+                                     DeviceImpl &Device);
+
+  // Filled by addImages(...).
+  std::unordered_map<std::string_view, kernel_id> MKernelNameToID;
+  std::unordered_map<kernel_id, DeviceImageWrapper *> MKernelIDToDevImageJIT;
+  // Controls lifetime of device image ptr and wrapper.
+  std::unordered_map<const __sycl_tgt_device_image *,
+                     std::unique_ptr<DeviceImageWrapper>>
+      MDeviceImageWrappers;
+  std::mutex MImageCollectionMutex;
+};
+
+} // namespace detail
+_LIBSYCL_END_NAMESPACE_SYCL
+
+#endif // _LIBSYCL_PROGRAM_MANAGER

libsycl/src/detail/queue_impl.cpp

@@ -6,9 +6,10 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include <detail/device_impl.hpp>
 #include <detail/queue_impl.hpp>
 
+#include <detail/device_impl.hpp>
+
 _LIBSYCL_BEGIN_NAMESPACE_SYCL
 
 namespace detail {

libsycl/src/usm_functions.cpp

@@ -103,7 +103,7 @@ void *malloc(std::size_t numBytes, const device &syclDevice,
 
   void *Ptr{};
   auto Result = detail::callNoCheck(
-      olMemAlloc, detail::getSyclObjImpl(syclDevice)->getOLHandle(),
+      olMemAlloc, detail::getSyclObjImpl(syclDevice)->getHandle(),
       detail::getOlAllocType(kind), numBytes, &Ptr);
   return detail::isFailed(Result) ? nullptr : Ptr;
 }