[SYCL] Fix memory leak for sub-devices (#20370)

Current `device_impl` implementation assumes shared ownership and manages the handle via `urRetain`/`urRelease` in constructor/destructor. Whenever we get temporary device handles via `urDeviceGet`/`urDevicePartition` (which implicitly return device handle with `refcount == 1`) and wrap those handles in `device_impl`, we have to release temporary handles at the end of scope to avoid memory leak. It's done here for devices: https://github.com/intel/llvm/blob/6b29cf120c267118dc8ad737e54e8b844bfb4e06/sycl/source/detail/platform_impl.cpp#L560-L563 Same needs to be done when we create sub-devices. But currently we were missing release calls for temporary handles after creating `device_impl` objects wrapping those handles. This PR fixes that issue by releasing temporary handles at the end of the scope.

Author: againull

sycl/source/detail/device_impl.cpp

@@ -167,6 +167,12 @@ std::vector<device> device_impl::create_sub_devices(
  MPlatform.getOrMakeDeviceImpl(a_ur_device));
  res.push_back(sycl_device);
  });
+ // urDevicePartition returns devices with their reference counts
+ // incremented. Each device_impl wrapper increments the reference count and
+ // decrements it on destruction (shared ownership). So, we have to decrement
+ // the reference count once here to release temporary handles.
+ for (ur_device_handle_t &SubDevice : SubDevices)
+ Adapter.call<UrApiKind::urDeviceRelease>(SubDevice);
  return res;
 }
 

sycl/unittests/context_device/DeviceRefCounter.cpp

@@ -29,6 +29,39 @@ static ur_result_t redefinedDeviceReleaseAfter(void *) {
  return UR_RESULT_SUCCESS;
 }
 
+ur_result_t redefinedDevicePartitionAfter(void *pParams) {
+ auto params = *static_cast<ur_device_partition_params_t *>(pParams);
+ if (*params.pphSubDevices) {
+ for (size_t I = 0; I < *params.pNumDevices; ++I) {
+ *params.pphSubDevices[I] = reinterpret_cast<ur_device_handle_t>(1000 + I);
+ }
+ }
+ if (*params.ppNumDevicesRet)
+ **params.ppNumDevicesRet = *params.pNumDevices;
+
+ DevRefCounter += *params.pNumDevices;
+ return UR_RESULT_SUCCESS;
+}
+
+static constexpr size_t NumSubDevices = 2;
+
+ur_result_t redefinedDeviceGetInfoAfter(void *pParams) {
+ auto params = *static_cast<ur_device_get_info_params_t *>(pParams);
+ if (*params.ppropName == UR_DEVICE_INFO_SUPPORTED_PARTITIONS) {
+ if (*params.ppPropValue) {
+ auto *Result =
+ reinterpret_cast<ur_device_partition_t *>(*params.ppPropValue);
+ *Result = UR_DEVICE_PARTITION_EQUALLY;
+ }
+ if (*params.ppPropSizeRet)
+ **params.ppPropSizeRet = sizeof(ur_device_partition_t);
+ } else if (*params.ppropName == UR_DEVICE_INFO_MAX_COMPUTE_UNITS) {
+ auto *Result = reinterpret_cast<uint32_t *>(*params.ppPropValue);
+ *Result = NumSubDevices;
+ }
+ return UR_RESULT_SUCCESS;
+}
+
 TEST(DevRefCounter, DevRefCounter) {
  {
  sycl::unittest::UrMock<> Mock;
@@ -52,3 +85,32 @@ TEST(DevRefCounter, DevRefCounter) {
  }
  EXPECT_EQ(DevRefCounter, 0);
 }
+
+TEST(SubDevRefCounter, SubDevRefCounter) {
+ {
+ DevRefCounter = 0;
+ sycl::unittest::UrMock<> Mock;
+ mock::getCallbacks().set_after_callback("urDeviceGet",
+ &redefinedDevicesGetAfter);
+ mock::getCallbacks().set_after_callback("urDeviceRetain",
+ &redefinedDeviceRetainAfter);
+ mock::getCallbacks().set_after_callback("urDeviceRelease",
+ &redefinedDeviceReleaseAfter);
+ mock::getCallbacks().set_before_callback("urDevicePartition",
+ &redefinedDevicePartitionAfter);
+ mock::getCallbacks().set_after_callback("urDeviceGetInfo",
+ &redefinedDeviceGetInfoAfter);
+ sycl::platform Plt = sycl::platform();
+
+ auto Devs = Plt.get_devices();
+ if (!Devs.empty()) {
+ auto Subdevs = Devs[0]
+ .create_sub_devices<
+ sycl::info::partition_property::partition_equally>(
+ NumSubDevices);
+ }
+ EXPECT_NE(DevRefCounter, 0);
+ sycl::detail::GlobalHandler::instance().getPlatformCache().clear();
+ }
+ EXPECT_EQ(DevRefCounter, 0);
+}