From 004a655e94cad706e4d682f6cbacd8089b46350a Mon Sep 17 00:00:00 2001
From: Artur Gainullin <artur.gainullin@intel.com>
Date: Mon, 30 Jun 2025 12:29:31 -0700
Subject: [PATCH] [UR] Fix creation of context with parent device and its
 sub-devices

This use case seems to be allowed. For example, according to:
https://intel.github.io/llvm/MultiTileCardWithLevelZero.html#context
"Both root-devices and sub-devices can be within single context, but they all should be of the same SYCL platform."
---
 .../source/adapters/level_zero/device.hpp     |  9 +++-
 .../conformance/context/urContextCreate.cpp   | 51 +++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/unified-runtime/source/adapters/level_zero/device.hpp b/unified-runtime/source/adapters/level_zero/device.hpp
index a8326c0cf668b..aa72c4ec0aa26 100644
--- a/unified-runtime/source/adapters/level_zero/device.hpp
+++ b/unified-runtime/source/adapters/level_zero/device.hpp
@@ -16,6 +16,7 @@
 #include <stdarg.h>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "adapters/level_zero/platform.hpp"
@@ -247,12 +248,16 @@ struct ur_device_handle_t_ : ur_object {
 inline std::vector<ur_device_handle_t>
 CollectDevicesAndSubDevices(const std::vector<ur_device_handle_t> &Devices) {
   std::vector<ur_device_handle_t> DevicesAndSubDevices;
+  std::unordered_set<ur_device_handle_t> Seen;
   std::function<void(const std::vector<ur_device_handle_t> &)>
       CollectDevicesAndSubDevicesRec =
           [&](const std::vector<ur_device_handle_t> &Devices) {
             for (auto &Device : Devices) {
-              DevicesAndSubDevices.push_back(Device);
-              CollectDevicesAndSubDevicesRec(Device->SubDevices);
+              // Only add device if has not been seen before.
+              if (Seen.insert(Device).second) {
+                DevicesAndSubDevices.push_back(Device);
+                CollectDevicesAndSubDevicesRec(Device->SubDevices);
+              }
             }
           };
   CollectDevicesAndSubDevicesRec(Devices);
diff --git a/unified-runtime/test/conformance/context/urContextCreate.cpp b/unified-runtime/test/conformance/context/urContextCreate.cpp
index ca77a43df69b2..173c11174121a 100644
--- a/unified-runtime/test/conformance/context/urContextCreate.cpp
+++ b/unified-runtime/test/conformance/context/urContextCreate.cpp
@@ -45,6 +45,57 @@ TEST_P(urContextCreateTest, InvalidEnumeration) {
                    urContextCreate(1, &device, &properties, context.ptr()));
 }
 
+TEST_P(urContextCreateTest, SuccessParentAndSubDevices) {
+  if (!uur::hasDevicePartitionSupport(device,
+                                      UR_DEVICE_PARTITION_BY_AFFINITY_DOMAIN)) {
+    GTEST_SKIP() << "Device \'" << device
+                 << "\' does not support partitioning by affinity domain.\n";
+  }
+
+  ur_device_affinity_domain_flags_t flag = UR_DEVICE_AFFINITY_DOMAIN_FLAG_NUMA;
+  ur_device_affinity_domain_flags_t supported_flags{0};
+  ASSERT_SUCCESS(
+      uur::GetDevicePartitionAffinityDomainFlags(device, supported_flags));
+  if (!(flag & supported_flags)) {
+    GTEST_SKIP() << static_cast<ur_device_affinity_domain_flag_t>(flag)
+                 << " is not supported by the device: \'" << device << "\'.\n";
+  }
+
+  ur_device_partition_property_t prop =
+      uur::makePartitionByAffinityDomain(flag);
+
+  ur_device_partition_properties_t properties{
+      UR_STRUCTURE_TYPE_DEVICE_PARTITION_PROPERTIES,
+      nullptr,
+      &prop,
+      1,
+  };
+
+  // Get the number of devices that will be created
+  uint32_t n_devices = 0;
+  ASSERT_SUCCESS(
+      urDevicePartition(device, &properties, 0, nullptr, &n_devices));
+  ASSERT_NE(n_devices, 0);
+
+  std::vector<ur_device_handle_t> sub_devices(n_devices);
+  ASSERT_SUCCESS(urDevicePartition(device, &properties,
+                                   static_cast<uint32_t>(sub_devices.size()),
+                                   sub_devices.data(), nullptr));
+
+  std::vector<ur_device_handle_t> all_devices;
+  all_devices.push_back(device);
+  all_devices.insert(all_devices.end(), sub_devices.begin(), sub_devices.end());
+  uur::raii::Context context = nullptr;
+  ASSERT_SUCCESS(urContextCreate(static_cast<uint32_t>(all_devices.size()),
+                                 all_devices.data(), nullptr, context.ptr()));
+  ASSERT_NE(nullptr, context);
+
+  for (auto sub_device : sub_devices) {
+    ASSERT_NE(sub_device, nullptr);
+    ASSERT_SUCCESS(urDeviceRelease(sub_device));
+  }
+}
+
 using urContextCreateMultiDeviceTest = uur::urAllDevicesTest;
 UUR_INSTANTIATE_PLATFORM_TEST_SUITE(urContextCreateMultiDeviceTest);