@@ -45,6 +45,57 @@ TEST_P(urContextCreateTest, InvalidEnumeration) {
45
45
urContextCreate (1 , &device, &properties, context.ptr ()));
46
46
}
47
47
48
+ TEST_P (urContextCreateTest, SuccessParentAndSubDevices) {
49
+ if (!uur::hasDevicePartitionSupport (device,
50
+ UR_DEVICE_PARTITION_BY_AFFINITY_DOMAIN)) {
51
+ GTEST_SKIP () << " Device \' " << device
52
+ << " \' does not support partitioning by affinity domain.\n " ;
53
+ }
54
+
55
+ ur_device_affinity_domain_flags_t flag = UR_DEVICE_AFFINITY_DOMAIN_FLAG_NUMA;
56
+ ur_device_affinity_domain_flags_t supported_flags{0 };
57
+ ASSERT_SUCCESS (
58
+ uur::GetDevicePartitionAffinityDomainFlags (device, supported_flags));
59
+ if (!(flag & supported_flags)) {
60
+ GTEST_SKIP () << static_cast <ur_device_affinity_domain_flag_t >(flag)
61
+ << " is not supported by the device: \' " << device << " \' .\n " ;
62
+ }
63
+
64
+ ur_device_partition_property_t prop =
65
+ uur::makePartitionByAffinityDomain (flag);
66
+
67
+ ur_device_partition_properties_t properties{
68
+ UR_STRUCTURE_TYPE_DEVICE_PARTITION_PROPERTIES,
69
+ nullptr ,
70
+ &prop,
71
+ 1 ,
72
+ };
73
+
74
+ // Get the number of devices that will be created
75
+ uint32_t n_devices = 0 ;
76
+ ASSERT_SUCCESS (
77
+ urDevicePartition (device, &properties, 0 , nullptr , &n_devices));
78
+ ASSERT_NE (n_devices, 0 );
79
+
80
+ std::vector<ur_device_handle_t > sub_devices (n_devices);
81
+ ASSERT_SUCCESS (urDevicePartition (device, &properties,
82
+ static_cast <uint32_t >(sub_devices.size ()),
83
+ sub_devices.data (), nullptr ));
84
+
85
+ std::vector<ur_device_handle_t > all_devices;
86
+ all_devices.push_back (device);
87
+ all_devices.insert (all_devices.end (), sub_devices.begin (), sub_devices.end ());
88
+ uur::raii::Context context = nullptr ;
89
+ ASSERT_SUCCESS (urContextCreate (static_cast <uint32_t >(all_devices.size ()),
90
+ all_devices.data (), nullptr , context.ptr ()));
91
+ ASSERT_NE (nullptr , context);
92
+
93
+ for (auto sub_device : sub_devices) {
94
+ ASSERT_NE (sub_device, nullptr );
95
+ ASSERT_SUCCESS (urDeviceRelease (sub_device));
96
+ }
97
+ }
98
+
48
99
using urContextCreateMultiDeviceTest = uur::urAllDevicesTest;
49
100
UUR_INSTANTIATE_PLATFORM_TEST_SUITE (urContextCreateMultiDeviceTest);
50
101
0 commit comments