From 22c2056403d38a588c989a97b3687a12249e01f9 Mon Sep 17 00:00:00 2001
From: "Zhang, Winston" <winston.zhang@intel.com>
Date: Fri, 24 Jan 2025 15:23:18 -0800
Subject: [PATCH] [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 <winston.zhang@intel.com>
---
 source/adapters/level_zero/device.cpp | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/source/adapters/level_zero/device.cpp b/source/adapters/level_zero/device.cpp
index 6b21507d8e..6844ec16bf 100644
--- 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;