[L0] MAX_COMPUTE_UNITS using ze_eu_count_ext_t

For some recovery SKUs, MAX_COMPUTE_COUNT calculation does not provide the correct number of EUs. Now we will use ze_eu_count_t when available. Signed-off-by: Zhang, Winston <[email protected]>
oneapi-src · Jan 24, 2025 · 22c2056 · 22c2056
1 parent 04cb8f1
commit 22c2056
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/source/adapters/level_zero/device.cpp b/source/adapters/level_zero/device.cpp
@@ -867,6 +867,14 @@ ur_result_t urDeviceGetInfo(
     return ReturnValue(int32_t(ZeDeviceNumIndices));
   } break;
   case UR_DEVICE_INFO_GPU_EU_COUNT: {
+    ze_device_properties_t DeviceProp = {ZE_STRUCTURE_TYPE_DEVICE_PROPERTIES};
+    ze_eu_count_ext_t EuCountDesc = {ZE_STRUCTURE_TYPE_EU_COUNT_EXT};
+    DeviceProp.pNext = &EuCountDesc;
+    ZE2UR_CALL(zeDeviceGetProperties, (ZeDevice, &DeviceProp));
+    if (DeviceProp.pNext &&
+        EuCountDesc.stype == ZE_STRUCTURE_TYPE_EU_COUNT_EXT) {
+      return ReturnValue(uint32_t{EuCountDesc.numTotalEUs});
+    }
     uint32_t count = Device->ZeDeviceProperties->numEUsPerSubslice *
                      Device->ZeDeviceProperties->numSubslicesPerSlice *
                      Device->ZeDeviceProperties->numSlices;