[SYCL] Add runtime support for split device binaries (#759)

This patch enables usage of multiple split device binaries per OS
module (executable or shared object file). The required binary is
chosen based on the entry table (filled by clang-offload-wrapper)
that lists all kernels contained within.

Signed-off-by: Sergey Semenov <[email protected]>

Author: sergey-semenov

sycl/include/CL/sycl/detail/common.hpp

@@ -246,6 +246,15 @@ size_t getLinearIndex(const T<Dims> &Index, const U<Dims> &Range) {
   return LinearIndex;
 }
 
+// Kernel set ID, used to group kernels (represented by OSModule & kernel name
+// pairs) into disjoint sets based on the kernel distribution among device
+// images.
+using KernelSetId = size_t;
+// Kernel set ID for kernels contained within the SPIRV file specified via
+// environment.
+constexpr KernelSetId SpvFileKSId = 0;
+constexpr KernelSetId LastKSId = SpvFileKSId;
+
 } // namespace detail
 } // namespace sycl
 } // namespace cl

sycl/include/CL/sycl/detail/context_impl.hpp

@@ -105,7 +105,7 @@ class context_impl {
   /// Gets cached programs.
   ///
   /// @return a map of cached programs.
-  std::map<OSModuleHandle, RT::PiProgram> &getCachedPrograms() {
+  std::map<KernelSetId, RT::PiProgram> &getCachedPrograms() {
     return MCachedPrograms;
   }
 
@@ -128,7 +128,7 @@ class context_impl {
   platform MPlatform;
   bool MPluginInterop;
   bool MHostContext;
-  std::map<OSModuleHandle, RT::PiProgram> MCachedPrograms;
+  std::map<KernelSetId, RT::PiProgram> MCachedPrograms;
   std::map<RT::PiProgram, std::map<string_class, RT::PiKernel>> MCachedKernels;
   std::shared_ptr<usm::USMDispatcher> MUSMDispatch;
 };

sycl/include/CL/sycl/detail/pi.h

@@ -229,8 +229,16 @@ typedef _pi_sampler_addressing_mode pi_sampler_addressing_mode;
 typedef _pi_sampler_filter_mode     pi_sampler_filter_mode;
 typedef _pi_sampler_info            pi_sampler_info;
 
-// Opaque data type for compatibility with OpenMP.
-typedef void * _pi_offload_entry;
+// Entry type, matches OpenMP for compatibility
+struct _pi_offload_entry_struct {
+  void *addr;
+  char *name;
+  size_t size;
+  int32_t flags;
+  int32_t reserved;
+};
+
+typedef _pi_offload_entry_struct * _pi_offload_entry;
 
 /// Types of device binary.
 typedef uint8_t pi_device_binary_type;
@@ -302,7 +310,7 @@ struct pi_device_binary_struct {
   const unsigned char *BinaryStart;
   /// Pointer to the target code end
   const unsigned char *BinaryEnd;
-  /// the offload entry table (not used, for compatibility with OpenMP)
+  /// the offload entry table
   _pi_offload_entry EntriesBegin;
   _pi_offload_entry EntriesEnd;
 };

sycl/include/CL/sycl/detail/program_impl.hpp

@@ -62,8 +62,7 @@ class program_impl {
     if (!is_host()) {
       DevicesSorted = sort_devices_by_cl_device_id(Devices);
     }
-    check_device_feature_support<
-        info::device::is_linker_available>(Devices);
+    check_device_feature_support<info::device::is_linker_available>(Devices);
     for (const auto &Prg : ProgramList) {
       Prg->throw_if_state_is_not(program_state::compiled);
       if (Prg->Context != Context) {
@@ -183,10 +182,9 @@ class program_impl {
   template <typename KernelT>
   void compile_with_kernel_type(string_class CompileOptions = "") {
     throw_if_state_is_not(program_state::none);
-    // TODO Check for existence of kernel
     if (!is_host()) {
       OSModuleHandle M = OSUtil::getOSModuleHandle(AddressInThisModule);
-      create_cl_program_with_il(M);
+      create_pi_program_with_kernel_name(M, KernelInfo<KernelT>::getName());
       compile(CompileOptions);
     }
     State = program_state::compiled;
@@ -206,17 +204,16 @@ class program_impl {
   template <typename KernelT>
   void build_with_kernel_type(string_class BuildOptions = "") {
     throw_if_state_is_not(program_state::none);
-    // TODO Check for existence of kernel
     if (!is_host()) {
       OSModuleHandle M = OSUtil::getOSModuleHandle(AddressInThisModule);
       // If there are no build options, program can be safely cached
       if (is_cacheable_with_options(BuildOptions)) {
         IsProgramAndKernelCachingAllowed = true;
-        Program =
-            ProgramManager::getInstance().getBuiltOpenCLProgram(M, Context);
+        Program = ProgramManager::getInstance().getBuiltPIProgram(
+            M, Context, KernelInfo<KernelT>::getName());
         PI_CALL(piProgramRetain)(Program);
       } else {
-        create_cl_program_with_il(M);
+        create_pi_program_with_kernel_name(M, KernelInfo<KernelT>::getName());
         build(BuildOptions);
       }
     }
@@ -237,8 +234,7 @@ class program_impl {
   void link(string_class LinkOptions = "") {
     throw_if_state_is_not(program_state::compiled);
     if (!is_host()) {
-      check_device_feature_support<
-          info::device::is_linker_available>(Devices);
+      check_device_feature_support<info::device::is_linker_available>(Devices);
       vector_class<RT::PiDevice> Devices(get_pi_devices());
       PI_CALL_THROW(piProgramLink, compile_program_error)(
           detail::getSyclObjImpl(Context)->getHandleRef(), Devices.size(),
@@ -339,8 +335,7 @@ class program_impl {
 
 private:
   template <info::device param>
-  void check_device_feature_support(
-      const vector_class<device> &devices) {
+  void check_device_feature_support(const vector_class<device> &devices) {
     for (const auto &device : devices) {
       if (!device.get_info<param>()) {
         throw feature_not_supported(
@@ -349,9 +344,12 @@ class program_impl {
     }
   }
 
-  void create_cl_program_with_il(OSModuleHandle M) {
+  void create_pi_program_with_kernel_name(OSModuleHandle M,
+                                          const string_class &KernelName) {
     assert(!Program && "This program already has an encapsulated PI program");
-    Program = ProgramManager::getInstance().createOpenCLProgram(M, Context);
+    ProgramManager &PM = ProgramManager::getInstance();
+    DeviceImage &Img = PM.getDeviceImage(M, KernelName, Context);
+    Program = PM.createPIProgram(Img, Context);
   }
 
   void create_cl_program_with_source(const string_class &Source) {
@@ -364,8 +362,7 @@ class program_impl {
   }
 
   void compile(const string_class &Options) {
-    check_device_feature_support<
-        info::device::is_compiler_available>(Devices);
+    check_device_feature_support<info::device::is_compiler_available>(Devices);
     vector_class<RT::PiDevice> Devices(get_pi_devices());
     RT::PiResult Err = PI_CALL_NOCHECK(piProgramCompile)(
         Program, Devices.size(), Devices.data(), Options.c_str(), 0, nullptr,
@@ -380,8 +377,7 @@ class program_impl {
   }
 
   void build(const string_class &Options) {
-    check_device_feature_support<
-        info::device::is_compiler_available>(Devices);
+    check_device_feature_support<info::device::is_compiler_available>(Devices);
     vector_class<RT::PiDevice> Devices(get_pi_devices());
     RT::PiResult Err =
         PI_CALL_NOCHECK(piProgramBuild)(Program, Devices.size(), Devices.data(),
@@ -402,6 +398,12 @@ class program_impl {
     return PiDevices;
   }
 
+  bool is_cacheable() const { return IsProgramAndKernelCachingAllowed; }
+
+  static bool is_cacheable_with_options(const string_class &Options) {
+    return Options.empty();
+  }
+
   bool has_cl_kernel(const string_class &KernelName) const {
     size_t Size;
     PI_CALL(piProgramGetInfo)(Program, CL_PROGRAM_KERNEL_NAMES, 0, nullptr,
@@ -420,15 +422,6 @@ class program_impl {
     return false;
   }
 
-  bool is_cacheable() const {
-    return IsProgramAndKernelCachingAllowed;
-  }
-
-  static bool
-  is_cacheable_with_options(const string_class &Options) {
-    return Options.empty();
-  }
-
   RT::PiKernel get_pi_kernel(const string_class &KernelName) const {
     RT::PiKernel Kernel;
 

sycl/include/CL/sycl/detail/program_manager/program_manager.hpp

@@ -48,11 +48,11 @@ class ProgramManager {
   // Returns the single instance of the program manager for the entire process.
   // Can only be called after staticInit is done.
   static ProgramManager &getInstance();
-  RT::PiProgram createOpenCLProgram(OSModuleHandle M, const context &Context,
-                                 DeviceImage **I = nullptr) {
-    return loadProgram(M, Context, I);
-  }
-  RT::PiProgram getBuiltOpenCLProgram(OSModuleHandle M, const context &Context);
+  DeviceImage &getDeviceImage(OSModuleHandle M, const string_class &KernelName,
+                              const context &Context);
+  RT::PiProgram createPIProgram(const DeviceImage &Img, const context &Context);
+  RT::PiProgram getBuiltPIProgram(OSModuleHandle M, const context &Context,
+                                  const string_class &KernelName);
   RT::PiKernel getOrCreateKernel(OSModuleHandle M, const context &Context,
                                   const string_class &KernelName);
   RT::PiProgram getClProgramFromClKernel(RT::PiKernel Kernel);
@@ -63,26 +63,65 @@ class ProgramManager {
   static string_class getProgramBuildLog(const RT::PiProgram &Program);
 
 private:
-  RT::PiProgram loadProgram(OSModuleHandle M, const context &Context,
-                            DeviceImage **I = nullptr);
-  void build(RT::PiProgram Program, const string_class &Options = "",
-             std::vector<RT::PiDevice> Devices = std::vector<RT::PiDevice>());
-
-  ProgramManager() = default;
+  ProgramManager();
   ~ProgramManager() = default;
   ProgramManager(ProgramManager const &) = delete;
   ProgramManager &operator=(ProgramManager const &) = delete;
 
+  DeviceImage &getDeviceImage(OSModuleHandle M, KernelSetId KSId,
+                              const context &Context);
+  void build(RT::PiProgram Program, const string_class &Options,
+             std::vector<RT::PiDevice> Devices);
+  /// Provides a new kernel set id for grouping kernel names together
+  KernelSetId getNextKernelSetId() const;
+  /// Returns the kernel set associated with the kernel, handles some special
+  /// cases (when reading images from file or using images with no entry info)
+  KernelSetId getKernelSetId(OSModuleHandle M,
+                             const string_class &KernelName) const;
+  /// Returns the format of the binary image
+  RT::PiDeviceBinaryType getFormat(const DeviceImage &Img) const;
+  /// Dumps image to current directory
+  void dumpImage(const DeviceImage &Img, KernelSetId KSId) const;
+
+  /// The three maps below are used during kernel resolution. Any kernel is
+  /// identified by its name and the OS module it's coming from, allowing
+  /// kernels with identical names in different OS modules. The following
+  /// assumption is made: for any two device images in a SYCL application their
+  /// kernel sets are either identical or disjoint.
+  /// Based on this assumption, m_KernelSets is used to group kernels together
+  /// into sets by assigning a set ID to them during device image registration.
+  /// This ID is then mapped to a vector of device images containing kernels
+  /// from the set (m_DeviceImages).
+  /// An exception is made for device images with no entry information: a
+  /// special kernel set ID is used for them which is assigned to just the OS
+  /// module. These kernel set ids are stored in m_OSModuleKernelSets and device
+  /// images associated with them are assumed to contain all kernels coming from
+  /// that OS module.
+
   /// Keeps all available device executable images added via \ref addImages.
-  /// Organizes the images as a map from a module handle (.exe .dll) to the
-  /// vector of images coming from the module.
+  /// Organizes the images as a map from a kernel set id to the vector of images
+  /// containing kernels from that set.
+  /// Access must be guarded by the \ref Sync::getGlobalLock()
+  std::map<KernelSetId, std::unique_ptr<std::vector<DeviceImage *>>> m_DeviceImages;
+
+  using StrToKSIdMap = std::map<string_class, KernelSetId>;
+  /// Maps names of kernels from a specific OS module (.exe .dll) to their set
+  /// id (the sets are disjoint).
+  /// Access must be guarded by the \ref Sync::getGlobalLock()
+  std::map<OSModuleHandle, StrToKSIdMap> m_KernelSets;
+
+  /// Keeps kernel sets for OS modules containing images without entry info.
+  /// Such images are assumed to contain all kernel associated with the module.
   /// Access must be guarded by the \ref Sync::getGlobalLock()
-  std::map<OSModuleHandle, std::unique_ptr<std::vector<DeviceImage *>>>
-      m_DeviceImages;
+  std::map<OSModuleHandle, KernelSetId> m_OSModuleKernelSets;
+
   /// Keeps device images not bound to a particular module. Program manager
   /// allocated memory for these images, so they are auto-freed in destructor.
   /// No image can out-live the Program manager.
   std::vector<std::unique_ptr<DeviceImage, ImageDeleter>> m_OrphanDeviceImages;
+
+  /// True iff a SPIRV file has been specified with an environment variable
+  bool m_UseSpvFile = false;
 };
 } // namespace detail
 } // namespace sycl

sycl/plugins/opencl/pi_opencl.cpp

@@ -68,10 +68,10 @@ pi_result OCL(piDevicesGet)(pi_platform platform, pi_device_type device_type,
   return cast<pi_result>(result);
 }
 
-pi_result OCL(piextDeviceSelectBinary)(
-    pi_device device, // TODO: does this need to be context?
-    pi_device_binary *images, pi_uint32 num_images,
-    pi_device_binary *selected_image) {
+pi_result OCL(piextDeviceSelectBinary)(pi_device device,
+                                       pi_device_binary *images,
+                                       pi_uint32 num_images,
+                                       pi_device_binary *selected_image) {
 
   // TODO: this is a bare-bones implementation for choosing a device image
   // that would be compatible with the targeted device. An AOT-compiled

sycl/source/detail/context_impl.cpp

@@ -86,9 +86,8 @@ context_impl::~context_impl() {
   // Release all programs and kernels created with this context
   for (auto ProgIt : MCachedPrograms) {
     RT::PiProgram ToBeDeleted = ProgIt.second;
-    for (auto KernIt : MCachedKernels[ToBeDeleted]) {
+    for (auto KernIt : MCachedKernels[ToBeDeleted])
       PI_CALL(piKernelRelease)(KernIt.second);
-    }
     PI_CALL(piProgramRelease)(ToBeDeleted);
   }
 }

sycl/source/detail/os_util.cpp

@@ -56,6 +56,8 @@ struct ModuleInfo {
   const char *Name;     // out
 };
 
+constexpr OSModuleHandle OSUtil::ExeModuleHandle;
+
 static int callback(struct dl_phdr_info *Info, size_t Size, void *Data) {
   auto Base = reinterpret_cast<unsigned char *>(Info->dlpi_addr);
   auto MI = reinterpret_cast<ModuleInfo *>(Data);