1212#include " ur_level_zero.hpp"
1313#include < algorithm>
1414#include < climits>
15+ #include < optional>
1516
1617UR_APIEXPORT ur_result_t UR_APICALL urDeviceGet (
1718 ur_platform_handle_t Platform, // /< [in] handle of the platform instance
@@ -353,8 +354,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(
353354 UR_DEVICE_AFFINITY_DOMAIN_FLAG_NEXT_PARTITIONABLE));
354355 case UR_DEVICE_INFO_PARTITION_TYPE: {
355356 // For root-device there is no partitioning to report.
356- if (pSize && !Device->isSubDevice ()) {
357- *pSize = 0 ;
357+ if (Device->SubDeviceCreationProperty == std::nullopt ||
358+ !Device->isSubDevice ()) {
359+ if (pSize)
360+ *pSize = 0 ;
358361 return UR_RESULT_SUCCESS;
359362 }
360363
@@ -365,7 +368,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(
365368 return ReturnValue (cslice);
366369 }
367370
368- return ReturnValue (Device->SubDeviceCreationProperty );
371+ return ReturnValue (* Device->SubDeviceCreationProperty );
369372 }
370373 // Everything under here is not supported yet
371374 case UR_EXT_DEVICE_INFO_OPENCL_C_VERSION:
@@ -1218,16 +1221,14 @@ UR_APIEXPORT ur_result_t UR_APICALL urDevicePartition(
12181221 UR_ASSERT (NumDevices == EffectiveNumDevices, UR_RESULT_ERROR_INVALID_VALUE);
12191222
12201223 for (uint32_t I = 0 ; I < NumDevices; I++) {
1221- Device->SubDevices [I]->SubDeviceCreationProperty =
1222- Properties->pProperties [0 ];
1223- if (Properties->pProperties [0 ].type ==
1224- UR_DEVICE_PARTITION_BY_AFFINITY_DOMAIN) {
1224+ auto prop = Properties->pProperties [0 ];
1225+ if (prop.type == UR_DEVICE_PARTITION_BY_AFFINITY_DOMAIN) {
12251226 // In case the value is NEXT_PARTITIONABLE, we need to change it to the
12261227 // chosen domain. This will always be NUMA since that's the only domain
12271228 // supported by level zero.
1228- Device->SubDevices [I]->SubDeviceCreationProperty .value .affinity_domain =
1229- UR_DEVICE_AFFINITY_DOMAIN_FLAG_NUMA;
1229+ prop.value .affinity_domain = UR_DEVICE_AFFINITY_DOMAIN_FLAG_NUMA;
12301230 }
1231+ Device->SubDevices [I]->SubDeviceCreationProperty = prop;
12311232
12321233 OutDevices[I] = Device->SubDevices [I];
12331234 // reusing the same pi_device needs to increment the reference count
0 commit comments