[SYCL] Fix broken uniqueness in make_device

sycl/include/CL/sycl/platform.hpp

@@ -26,13 +26,18 @@ class device_selector;
 class device;
 namespace detail {
 class platform_impl;
-}
+device make_device(pi_native_handle NativeHandle, backend Backend);
+} // namespace detail
 
 /// Encapsulates a SYCL platform on which kernels may be executed.
 ///
 /// \ingroup sycl_api
 class __SYCL_EXPORT platform {
 public:
+
+friend device detail::make_device(pi_native_handle NativeHandle,
+                                  backend Backend);
+
   /// Constructs a SYCL platform as a host platform.
   platform();
 

sycl/include/sycl/ext/oneapi/experimental/backend/cuda.hpp

@@ -75,6 +75,17 @@ interop_handle::get_native_context<backend::ext_oneapi_cuda>() const {
 template <>
 inline device make_device<backend::ext_oneapi_cuda>(
     const backend_input_t<backend::ext_oneapi_cuda, device> &BackendObject) {
+      auto plts = platform::get_platforms();
+  for (const auto &plt : plts) {
+    if (plt.get_backend() == backend::ext_oneapi_cuda) {
+      auto devs = plt.get_devices(info::device_type::all);
+      for (auto &dev : devs) {
+        if (BackendObject == get_native<backend::ext_oneapi_cuda>(dev)) {
+          return dev;
+        }
+      }
+    }
+  }
   pi_native_handle NativeHandle = static_cast<pi_native_handle>(BackendObject);
   return ext::oneapi::cuda::make_device(NativeHandle);
 }

sycl/source/backend.cpp

@@ -57,12 +57,21 @@ platform make_platform(pi_native_handle NativeHandle, backend Backend) {
 
 __SYCL_EXPORT device make_device(pi_native_handle NativeHandle,
                                  backend Backend) {
-  const auto &Plugin = getPlugin(Backend);
 
-  pi::PiDevice PiDevice = nullptr;
-  Plugin.call<PiApiKind::piextDeviceCreateWithNativeHandle>(NativeHandle,
-                                                            nullptr, &PiDevice);
-  // Construct the SYCL device from PI device.
+  auto plts = platform::get_platforms();
+  detail::pi::PiDevice PiDevice = nullptr;
+  const auto &Plugin = detail::getPlugin(Backend);
+  for (const auto &plt : plts) {
+    if (plt.get_backend() == Backend) {
+      Plugin.call<detail::PiApiKind::piextDeviceCreateWithNativeHandle>(
+          NativeHandle, nullptr, &PiDevice);
+      auto devImpl = plt.impl->getDeviceImpl(PiDevice, plt.impl);
+      if (devImpl != nullptr) {
+        return detail::createSyclObjFromImpl<device>(devImpl);
+      }
+    }
+  }
+
   return detail::createSyclObjFromImpl<device>(
       std::make_shared<device_impl>(PiDevice, Plugin));
 }

sycl/source/detail/platform_impl.cpp

@@ -212,6 +212,20 @@ static void filterDeviceFilter(std::vector<RT::PiDevice> &PiDevices,
   Plugin.setLastDeviceId(Platform, DeviceNum);
 }
 
+std::shared_ptr<device_impl> platform_impl::getDeviceImpl(
+    RT::PiDevice PiDevice, const std::shared_ptr<platform_impl> &PlatformImpl) {
+  const std::lock_guard<std::mutex> Guard(MDeviceMapMutex);
+
+  // If we've already seen this device, return the impl
+  for (const std::weak_ptr<device_impl> &DeviceWP : MDeviceCache) {
+    if (std::shared_ptr<device_impl> Device = DeviceWP.lock()) {
+      if (Device->getHandleRef() == PiDevice)
+        return Device;
+    }
+  }
+  return nullptr;
+    }
+
 std::shared_ptr<device_impl> platform_impl::getOrMakeDeviceImpl(
     RT::PiDevice PiDevice, const std::shared_ptr<platform_impl> &PlatformImpl) {
   const std::lock_guard<std::mutex> Guard(MDeviceMapMutex);

sycl/source/detail/platform_impl.hpp

@@ -152,6 +152,10 @@ class platform_impl {
   getOrMakeDeviceImpl(RT::PiDevice PiDevice,
                       const std::shared_ptr<platform_impl> &PlatformImpl);
 
+  std::shared_ptr<device_impl>
+  getDeviceImpl(RT::PiDevice PiDevice,
+                const std::shared_ptr<platform_impl> &PlatformImpl);
+
   /// Static functions that help maintain platform uniquess and
   /// equality of comparison