Do not call urDeviceRetain() in case of subdevices

urDeviceRetain(MDevice) should not be called
in the constructor of device_impl in case of subdevices,
because their reference counter does not drop to zero
and subdevices are not destroyed nor freed now,
what causes memory leaks.

It fixes URT-961.

Signed-off-by: Lukasz Dorau <[email protected]>

Author: ldorau

sycl/source/detail/device_impl.cpp

@@ -22,14 +22,18 @@ namespace detail {
 /// Constructs a SYCL device instance using the provided
 /// UR device instance.
 device_impl::device_impl(ur_device_handle_t Device, platform_impl &Platform,
-                         device_impl::private_tag)
+                         device_impl::private_tag, bool IsSubDevice)
     : MDevice(Device), MPlatform(Platform),
       // No need to set MRootDevice when MAlwaysRootDevice is true
       MRootDevice(Platform.MAlwaysRootDevice
                       ? nullptr
                       : get_info_impl<UR_DEVICE_INFO_PARENT_DEVICE>()),
       // TODO catch an exception and put it to list of asynchronous exceptions:
       MCache{*this} {
+  if (IsSubDevice) {
+    return;
+  }
+
   // Interoperability Constructor already calls DeviceRetain in
   // urDeviceCreateWithNativeHandle.
   getAdapter().call<UrApiKind::urDeviceRetain>(MDevice);
@@ -164,7 +168,7 @@ std::vector<device> device_impl::create_sub_devices(
   std::for_each(SubDevices.begin(), SubDevices.end(),
                 [&res, this](const ur_device_handle_t &a_ur_device) {
                   device sycl_device = detail::createSyclObjFromImpl<device>(
-                      MPlatform.getOrMakeDeviceImpl(a_ur_device));
+                      MPlatform.getOrMakeSubDeviceImpl(a_ur_device));
                   res.push_back(sycl_device);
                 });
   return res;

sycl/source/detail/device_impl.hpp

@@ -429,10 +429,11 @@ class device_impl : public std::enable_shared_from_this<device_impl> {
   /// Constructs a SYCL device instance using the provided
   /// UR device instance.
   //
-  // Must be called through `platform_impl::getOrMakeDeviceImpl` only.
+  // Must be called through `platform_impl::getOrMakeDeviceImpl` or
+  // `platform_impl::getOrMakeSubDeviceImpl` only.
   // `private_tag` ensures that is true.
   explicit device_impl(ur_device_handle_t Device, platform_impl &Platform,
-                       private_tag);
+                       private_tag, bool IsSubDevice);
 
   ~device_impl();
 

sycl/source/detail/platform_impl.cpp

@@ -306,7 +306,20 @@ device_impl &platform_impl::getOrMakeDeviceImpl(ur_device_handle_t UrDevice) {
 
   // Otherwise make the impl
   MDevices.emplace_back(std::make_shared<device_impl>(
-      UrDevice, *this, device_impl::private_tag{}));
+      UrDevice, *this, device_impl::private_tag{}, false /*IsSubDevice*/));
+
+  return *MDevices.back();
+}
+
+device_impl &platform_impl::getOrMakeSubDeviceImpl(ur_device_handle_t UrSubDevice) {
+  const std::lock_guard<std::mutex> Guard(MDeviceMapMutex);
+  // If we've already seen this device, return the impl
+  if (device_impl *Result = getDeviceImplHelper(UrSubDevice))
+    return *Result;
+
+  // Otherwise make the impl
+  MDevices.emplace_back(std::make_shared<device_impl>(
+      UrSubDevice, *this, device_impl::private_tag{}, true /*IsSubDevice*/));
 
   return *MDevices.back();
 }

sycl/source/detail/platform_impl.hpp

@@ -172,6 +172,17 @@ class platform_impl : public std::enable_shared_from_this<platform_impl> {
   /// \return a device_impl* corresponding to the device
   device_impl &getOrMakeDeviceImpl(ur_device_handle_t UrDevice);
 
+  /// Queries the device_impl cache to either return a shared_ptr
+  /// for the device_impl corresponding to the UrSubDevice or add
+  /// a new entry to the cache
+  ///
+  /// \param UrSubDevice is the UrSubDevice whose impl is requested
+  ///
+  /// \param PlatormImpl is the Platform for that Device
+  ///
+  /// \return a device_impl* corresponding to the subdevice
+  device_impl &getOrMakeSubDeviceImpl(ur_device_handle_t UrSubDevice);
+
   /// Queries the cache to see if the specified UR platform has been seen
   /// before.  If so, return the cached platform_impl, otherwise create a new
   /// one and cache it.