[SYCL] Add support of multiple devices within a context (#2343)

This patch adds support of multiple devices within a context. Programs
can be created from images or from SPIR-V binaries. Only kernels,
created using invoking kernels functions (parallel_for, single_task,
...) are supported. Kernels, created in OpenCL interoperability mode
(using sycl::program and sycl::kernel functions) are not supported.

Author: dm-vodopyanov

sycl/source/detail/context_impl.hpp

@@ -134,7 +134,8 @@ class context_impl {
   /// more details.
   ///
   /// \returns a map with device library programs.
-  std::map<DeviceLibExt, RT::PiProgram> &getCachedLibPrograms() {
+  std::map<std::pair<DeviceLibExt, RT::PiDevice>, RT::PiProgram> &
+  getCachedLibPrograms() {
     return MCachedLibPrograms;
   }
 
@@ -155,7 +156,8 @@ class context_impl {
   PlatformImplPtr MPlatform;
   bool MHostContext;
   bool MUseCUDAPrimaryContext;
-  std::map<DeviceLibExt, RT::PiProgram> MCachedLibPrograms;
+  std::map<std::pair<DeviceLibExt, RT::PiDevice>, RT::PiProgram>
+      MCachedLibPrograms;
   mutable KernelProgramCache MKernelProgramCache;
 };
 

sycl/source/detail/kernel_program_cache.hpp

@@ -67,15 +67,17 @@ class KernelProgramCache {
   using PiProgramT = std::remove_pointer<RT::PiProgram>::type;
   using PiProgramPtrT = std::atomic<PiProgramT *>;
   using ProgramWithBuildStateT = BuildResult<PiProgramT>;
-  using ProgramCacheKeyT = std::pair<SerializedObj, KernelSetId>;
+  using ProgramCacheKeyT =
+      std::pair<std::pair<SerializedObj, KernelSetId>, RT::PiDevice>;
   using ProgramCacheT = std::map<ProgramCacheKeyT, ProgramWithBuildStateT>;
   using ContextPtr = context_impl *;
 
   using PiKernelT = std::remove_pointer<RT::PiKernel>::type;
 
   using PiKernelPtrT = std::atomic<PiKernelT *>;
   using KernelWithBuildStateT = BuildResult<PiKernelT>;
-  using KernelByNameT = std::map<string_class, KernelWithBuildStateT>;
+  using KernelByNameT =
+      std::map<std::pair<string_class, RT::PiDevice>, KernelWithBuildStateT>;
   using KernelCacheT = std::map<RT::PiProgram, KernelByNameT>;
 
   ~KernelProgramCache();

sycl/source/detail/program_impl.cpp

@@ -28,7 +28,14 @@ program_impl::program_impl(ContextImplPtr Context)
 
 program_impl::program_impl(ContextImplPtr Context,
                            vector_class<device> DeviceList)
-    : MContext(Context), MDevices(DeviceList) {}
+    : MContext(Context), MDevices(DeviceList) {
+  if (Context->getDevices().size() > 1) {
+    throw feature_not_supported(
+        "multiple devices within a context are not supported with "
+        "sycl::program and sycl::kernel",
+        PI_INVALID_OPERATION);
+  }
+}
 
 program_impl::program_impl(
     vector_class<shared_ptr_class<program_impl>> ProgramList,
@@ -51,6 +58,12 @@ program_impl::program_impl(
   }
 
   MContext = ProgramList[0]->MContext;
+  if (MContext->getDevices().size() > 1) {
+    throw feature_not_supported(
+        "multiple devices within a context are not supported with "
+        "sycl::program and sycl::kernel",
+        PI_INVALID_OPERATION);
+  }
   MDevices = ProgramList[0]->MDevices;
   vector_class<device> DevicesSorted;
   if (!is_host()) {
@@ -105,6 +118,13 @@ program_impl::program_impl(ContextImplPtr Context,
                            RT::PiProgram Program)
     : MProgram(Program), MContext(Context), MLinkable(true) {
 
+  if (Context->getDevices().size() > 1) {
+    throw feature_not_supported(
+        "multiple devices within a context are not supported with "
+        "sycl::program and sycl::kernel",
+        PI_INVALID_OPERATION);
+  }
+
   const detail::plugin &Plugin = getPlugin();
   if (MProgram == nullptr) {
     assert(InteropProgram &&
@@ -233,7 +253,7 @@ void program_impl::build_with_kernel_name(string_class KernelName,
     if (is_cacheable_with_options(BuildOptions)) {
       MProgramAndKernelCachingAllowed = true;
       MProgram = ProgramManager::getInstance().getBuiltPIProgram(
-          Module, get_context(), KernelName, this,
+          Module, get_context(), get_devices()[0], KernelName, this,
           /*JITCompilationIsRequired=*/(!BuildOptions.empty()));
       const detail::plugin &Plugin = getPlugin();
       Plugin.call<PiApiKind::piProgramRetain>(MProgram);
@@ -356,7 +376,7 @@ void program_impl::build(const string_class &Options) {
   check_device_feature_support<info::device::is_compiler_available>(MDevices);
   vector_class<RT::PiDevice> Devices(get_pi_devices());
   const detail::plugin &Plugin = getPlugin();
-  ProgramManager::getInstance().flushSpecConstants(*this);
+  ProgramManager::getInstance().flushSpecConstants(*this, get_pi_devices()[0]);
   RT::PiResult Err = Plugin.call_nocheck<PiApiKind::piProgramBuild>(
       MProgram, Devices.size(), Devices.data(), Options.c_str(), nullptr,
       nullptr);
@@ -404,7 +424,8 @@ RT::PiKernel program_impl::get_pi_kernel(const string_class &KernelName) const {
   if (is_cacheable()) {
     std::tie(Kernel, std::ignore) =
         ProgramManager::getInstance().getOrCreateKernel(
-            MProgramModuleHandle, get_context(), KernelName, this);
+            MProgramModuleHandle, get_context(), get_devices()[0], KernelName,
+            this);
     getPlugin().call<PiApiKind::piKernelRetain>(Kernel);
   } else {
     const detail::plugin &Plugin = getPlugin();
@@ -453,9 +474,10 @@ void program_impl::create_pi_program_with_kernel_name(
     bool JITCompilationIsRequired) {
   assert(!MProgram && "This program already has an encapsulated PI program");
   ProgramManager &PM = ProgramManager::getInstance();
+  const device FirstDevice = get_devices()[0];
   RTDeviceBinaryImage &Img = PM.getDeviceImage(
-      Module, KernelName, get_context(), JITCompilationIsRequired);
-  MProgram = PM.createPIProgram(Img, get_context());
+      Module, KernelName, get_context(), FirstDevice, JITCompilationIsRequired);
+  MProgram = PM.createPIProgram(Img, get_context(), FirstDevice);
 }
 
 template <>