src, tests: fix engine creation failure for sub-devices + test

src/gpu/ocl/ocl_utils.cpp

@@ -83,10 +83,25 @@ status_t get_ocl_device_index(size_t *index, cl_device_id device) {
     std::vector<cl_device_id> ocl_devices;
     CHECK(get_ocl_devices(&ocl_devices, CL_DEVICE_TYPE_GPU));
 
-    auto it = std::find(ocl_devices.begin(), ocl_devices.end(), device);
-    if (it == ocl_devices.end()) return status::invalid_arguments;
-    *index = it - ocl_devices.begin();
-    return status::success;
+    // Search the top level device unconditionally
+    auto parent_device = device;
+    auto top_level_device = device;
+    while (parent_device) {
+        top_level_device = parent_device;
+        OCL_CHECK(clGetDeviceInfo(top_level_device, CL_DEVICE_PARENT_DEVICE,
+                sizeof(cl_device_id), &parent_device, nullptr));
+    }
+
+    // Find the top level device in the list
+    auto it = std::find(
+            ocl_devices.begin(), ocl_devices.end(), top_level_device);
+    if (it != ocl_devices.end()) {
+        *index = it - ocl_devices.begin();
+        return status::success;
+    } else {
+        *index = SIZE_MAX;
+        return status::invalid_arguments;
+    }
 }
 
 status_t get_ocl_kernel_arg_type(compute::scalar_type_t *type,

src/sycl/sycl_engine.hpp

@@ -100,13 +100,35 @@ inline std::vector<cl::sycl::device> get_sycl_devices(
 inline status_t get_sycl_device_index(
         size_t *index, const cl::sycl::device &dev) {
     auto dev_type = dev.get_info<cl::sycl::info::device::device_type>();
-    auto devices = get_sycl_devices(dev_type, get_sycl_backend(dev));
+    auto backend = get_sycl_backend(dev);
+    auto devices = get_sycl_devices(dev_type, backend);
 
-    auto it = std::find_if(devices.begin(), devices.end(),
-            [&](const cl::sycl::device &d) { return are_equal(d, dev); });
-    if (it == devices.end()) return status::invalid_arguments;
-    *index = it - devices.begin();
-    return status::success;
+    auto is_subdevice = [&backend](const cl::sycl::device &d) {
+        // TODO: remove this work around once Level-Zero is fixed
+        if (backend == backend_t::level0) return false;
+        return d.get_info<cl::sycl::info::device::partition_type_property>()
+                != cl::sycl::info::partition_property::no_partition;
+    };
+
+    // Search the top level device
+    auto parent_device = dev;
+    while (is_subdevice(parent_device)) {
+        parent_device
+                = parent_device
+                          .get_info<cl::sycl::info::device::parent_device>();
+    }
+
+    // Find the top level device in the list
+    auto it = std::find(devices.begin(), devices.end(), parent_device);
+    if (it != devices.end()) {
+        *index = it - devices.begin();
+        return status::success;
+    } else {
+        *index = SIZE_MAX;
+        // TODO: remove this work around once Level-Zero is fixed
+        if (backend == backend_t::level0) return status::success;
+        return status::invalid_arguments;
+    }
 }
 
 class sycl_engine_factory_t : public engine_factory_t {

tests/gtests/sycl/api/test_engine.cpp

@@ -167,6 +167,34 @@ TEST(sycl_engine_test, HostDevice) {
     }
 }
 
+TEST_P(sycl_engine_test, SubDevice) {
+    auto param = GetParam();
+
+    SKIP_IF(param.expected_status != dnnl_success,
+            "Don't test for failed scenarios");
+    SKIP_IF(!gpu_dev.get(), "Non GPU doesn't support sub-devices");
+
+    auto &dev = *gpu_dev.get();
+    auto max_sub_devices
+            = dev.get_info<info::device::partition_max_sub_devices>();
+    SKIP_IF(max_sub_devices < 2, "This GPU doesn't support sub-devices");
+
+    auto sub_dev = dev.create_sub_devices<
+            info::partition_property::partition_by_affinity_domain>(
+            info::partition_affinity_domain::next_partitionable);
+    context sub_ctx(sub_dev);
+
+    catch_expected_failures(
+            [&]() {
+                for (const auto &sub_dev_i : sub_dev) {
+                    engine eng;
+                    ASSERT_NO_THROW(eng
+                            = sycl_interop::make_engine(sub_dev_i, sub_ctx));
+                }
+            },
+            param.expected_status != dnnl_success, param.expected_status);
+}
+
 INSTANTIATE_TEST_SUITE_P(Simple, sycl_engine_test,
         ::testing::Values(sycl_engine_test_params {dev_kind::gpu, ctx_kind::gpu,
                                   dnnl_success},