[UR] Fix OpenCL handles generated from sub-devices (#17931)

Sub-devices don't appear in the list of devices, and so the prior OpenCL
implementation would assume that they are invalid. With this patch,
creating a sub-device will add it to a "cache" stored in the platform.
Multiple `DeviceCreateWithNativeHandle` calls with the same native
handle will result in the same UR handle. A test has been added to
verify this.

This unfortunately required shuffling headers around to avoid circular
imports.

This fixes an issue tracked internally by Intel as URT-903

Author: RossBrunton

unified-runtime/source/adapters/opencl/adapter.hpp

@@ -7,6 +7,7 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+#include "device.hpp"
 #include "logger/ur_logger.hpp"
 #include "platform.hpp"
 

unified-runtime/source/adapters/opencl/device.cpp

@@ -1598,6 +1598,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceCreateWithNativeHandle(
     ur_native_handle_t hNativeDevice, ur_adapter_handle_t,
     const ur_device_native_properties_t *pProperties,
     ur_device_handle_t *phDevice) {
+
+  auto SetDeviceProps = [&]() {
+    (*phDevice)->IsNativeHandleOwned =
+        pProperties ? pProperties->isNativeHandleOwned : false;
+  };
+
   cl_device_id NativeHandle = reinterpret_cast<cl_device_id>(hNativeDevice);
 
   uint32_t NumPlatforms = 0;
@@ -1617,12 +1623,43 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceCreateWithNativeHandle(
     for (auto &Device : Devices) {
       if (Device->CLDevice == NativeHandle) {
         *phDevice = Device;
-        (*phDevice)->IsNativeHandleOwned =
-            pProperties ? pProperties->isNativeHandleOwned : false;
+        SetDeviceProps();
         return UR_RESULT_SUCCESS;
       }
     }
   }
+
+  // Handle sub-devices by storing/querying a map stored in the platform
+  cl_device_id Parent = nullptr;
+  CL_RETURN_ON_FAILURE(clGetDeviceInfo(NativeHandle, CL_DEVICE_PARENT_DEVICE,
+                                       sizeof(Parent), &Parent, nullptr));
+  if (Parent != nullptr) {
+    ur_device_handle_t ParentUrHandle;
+    // This will either create a new device handle, or return an existing one
+    UR_RETURN_ON_FAILURE(urDeviceCreateWithNativeHandle(
+        reinterpret_cast<ur_native_handle_t>(Parent), nullptr, nullptr,
+        &ParentUrHandle));
+
+    ur_platform_handle_t PlatformHandle = ParentUrHandle->Platform;
+    assert(PlatformHandle);
+
+    {
+      std::lock_guard lock{PlatformHandle->SubDevicesLock};
+
+      if (PlatformHandle->SubDevices.count(NativeHandle)) {
+        *phDevice = PlatformHandle->SubDevices[NativeHandle];
+      } else {
+        *phDevice = std::make_unique<ur_device_handle_t_>(
+                        NativeHandle, PlatformHandle, ParentUrHandle)
+                        .release();
+        PlatformHandle->SubDevices[NativeHandle] = *phDevice;
+      }
+    }
+
+    SetDeviceProps();
+    return UR_RESULT_SUCCESS;
+  }
+
   return UR_RESULT_ERROR_INVALID_DEVICE;
 }
 

unified-runtime/source/adapters/opencl/device.hpp

@@ -10,6 +10,7 @@
 #pragma once
 
 #include "common.hpp"
+#include "device.hpp"
 #include "platform.hpp"
 
 struct ur_device_handle_t_ {
@@ -35,6 +36,12 @@ struct ur_device_handle_t_ {
   }
 
   ~ur_device_handle_t_() {
+    if (ParentDevice) {
+      // This does not need protected by a lock; this destructor can only run
+      // exactly once. However, to prevent issues with the OpenCL handle being
+      // reused, CLDevice must still be alive here.
+      Platform->SubDevices.erase(CLDevice);
+    }
     if (ParentDevice && IsNativeHandleOwned) {
       clReleaseDevice(CLDevice);
     }

unified-runtime/source/adapters/opencl/platform.cpp

@@ -10,6 +10,7 @@
 
 #include "platform.hpp"
 #include "adapter.hpp"
+#include "device.hpp"
 
 static cl_int mapURPlatformInfoToCL(ur_platform_info_t URPropName) {
 
@@ -191,3 +192,43 @@ urPlatformGetBackendOption(ur_platform_handle_t, const char *pFrontendOption,
   }
   return UR_RESULT_ERROR_INVALID_VALUE;
 }
+
+ur_result_t ur_platform_handle_t_::InitDevices() {
+  if (Devices.empty()) {
+    cl_uint DeviceNum = 0;
+    cl_int Res =
+        clGetDeviceIDs(CLPlatform, CL_DEVICE_TYPE_ALL, 0, nullptr, &DeviceNum);
+
+    // Absorb the CL_DEVICE_NOT_FOUND and just return 0 in num_devices
+    if (Res == CL_DEVICE_NOT_FOUND) {
+      return UR_RESULT_SUCCESS;
+    }
+
+    CL_RETURN_ON_FAILURE(Res);
+
+    std::vector<cl_device_id> CLDevices(DeviceNum);
+    Res = clGetDeviceIDs(CLPlatform, CL_DEVICE_TYPE_ALL, DeviceNum,
+                         CLDevices.data(), nullptr);
+
+    // Absorb the CL_DEVICE_NOT_FOUND and just return 0 in num_devices
+    if (Res == CL_DEVICE_NOT_FOUND) {
+      return UR_RESULT_SUCCESS;
+    }
+
+    CL_RETURN_ON_FAILURE(Res);
+
+    try {
+      Devices.resize(DeviceNum);
+      for (size_t i = 0; i < DeviceNum; i++) {
+        Devices[i] =
+            std::make_unique<ur_device_handle_t_>(CLDevices[i], this, nullptr);
+      }
+    } catch (std::bad_alloc &) {
+      return UR_RESULT_ERROR_OUT_OF_RESOURCES;
+    } catch (...) {
+      return UR_RESULT_ERROR_UNKNOWN;
+    }
+  }
+
+  return UR_RESULT_SUCCESS;
+}

unified-runtime/source/adapters/opencl/platform.hpp

@@ -10,14 +10,17 @@
 #pragma once
 
 #include "common.hpp"
-#include "device.hpp"
 
 #include <vector>
 
+struct ur_device_handle_t_;
+
 struct ur_platform_handle_t_ {
   using native_type = cl_platform_id;
   native_type CLPlatform = nullptr;
   std::vector<std::unique_ptr<ur_device_handle_t_>> Devices;
+  std::map<cl_device_id, ur_device_handle_t> SubDevices;
+  std::mutex SubDevicesLock;
 
   ur_platform_handle_t_(native_type Plat) : CLPlatform(Plat) {}
 
@@ -42,45 +45,7 @@ struct ur_platform_handle_t_ {
     return UR_RESULT_SUCCESS;
   }
 
-  ur_result_t InitDevices() {
-    if (Devices.empty()) {
-      cl_uint DeviceNum = 0;
-      cl_int Res = clGetDeviceIDs(CLPlatform, CL_DEVICE_TYPE_ALL, 0, nullptr,
-                                  &DeviceNum);
-
-      // Absorb the CL_DEVICE_NOT_FOUND and just return 0 in num_devices
-      if (Res == CL_DEVICE_NOT_FOUND) {
-        return UR_RESULT_SUCCESS;
-      }
-
-      CL_RETURN_ON_FAILURE(Res);
-
-      std::vector<cl_device_id> CLDevices(DeviceNum);
-      Res = clGetDeviceIDs(CLPlatform, CL_DEVICE_TYPE_ALL, DeviceNum,
-                           CLDevices.data(), nullptr);
-
-      // Absorb the CL_DEVICE_NOT_FOUND and just return 0 in num_devices
-      if (Res == CL_DEVICE_NOT_FOUND) {
-        return UR_RESULT_SUCCESS;
-      }
-
-      CL_RETURN_ON_FAILURE(Res);
-
-      try {
-        Devices.resize(DeviceNum);
-        for (size_t i = 0; i < DeviceNum; i++) {
-          Devices[i] = std::make_unique<ur_device_handle_t_>(CLDevices[i], this,
-                                                             nullptr);
-        }
-      } catch (std::bad_alloc &) {
-        return UR_RESULT_ERROR_OUT_OF_RESOURCES;
-      } catch (...) {
-        return UR_RESULT_ERROR_UNKNOWN;
-      }
-    }
-
-    return UR_RESULT_SUCCESS;
-  }
+  ur_result_t InitDevices();
 
   ur_result_t getPlatformVersion(oclv::OpenCLVersion &Version) {
     size_t PlatVerSize = 0;

unified-runtime/test/conformance/device/urDeviceCreateWithNativeHandle.cpp

@@ -65,3 +65,36 @@ TEST_P(urDeviceCreateWithNativeHandleTest, InvalidNullPointerDevice) {
       UR_RESULT_ERROR_INVALID_NULL_POINTER,
       urDeviceCreateWithNativeHandle(native_handle, adapter, nullptr, nullptr));
 }
+
+TEST_P(urDeviceCreateWithNativeHandleTest, SubDeviceHandleEquality) {
+  if (!uur::hasDevicePartitionSupport(device, UR_DEVICE_PARTITION_EQUALLY)) {
+    GTEST_SKIP() << "Device: \'" << device
+                 << "\' does not support partitioning equally.";
+  }
+
+  ur_device_partition_property_t property = uur::makePartitionEquallyDesc(2);
+  ur_device_partition_properties_t properties{
+      UR_STRUCTURE_TYPE_DEVICE_PARTITION_PROPERTIES,
+      nullptr,
+      &property,
+      1,
+  };
+
+  std::vector<ur_device_handle_t> sub_devices(2);
+  ASSERT_SUCCESS(
+      urDevicePartition(device, &properties, 2, sub_devices.data(), nullptr));
+
+  ur_native_handle_t native;
+  UUR_ASSERT_SUCCESS_OR_UNSUPPORTED(
+      urDeviceGetNativeHandle(sub_devices[0], &native));
+
+  ur_device_handle_t handle_a;
+  ASSERT_SUCCESS(
+      urDeviceCreateWithNativeHandle(native, adapter, nullptr, &handle_a));
+  ur_device_handle_t handle_b;
+  ASSERT_SUCCESS(
+      urDeviceCreateWithNativeHandle(native, adapter, nullptr, &handle_b));
+
+  ASSERT_EQ(handle_a, handle_b);
+  ASSERT_NE(handle_a, nullptr);
+}