[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 30, 2025 · 3dbf8b2 · 3dbf8b2
1 parent 14f4a3b
commit 3dbf8b2
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/source/adapters/level_zero/device.cpp b/source/adapters/level_zero/device.cpp
@@ -873,6 +873,18 @@ ur_result_t urDeviceGetInfo(
     return ReturnValue(int32_t(ZeDeviceNumIndices));
   } break;
   case UR_DEVICE_INFO_GPU_EU_COUNT: {
+    if (Device->Platform->ZeDriverEuCountExtensionFound) {
+      ze_device_properties_t DeviceProp = {};
+      DeviceProp.stype = ZE_STRUCTURE_TYPE_DEVICE_PROPERTIES;
+      ze_eu_count_ext_t EuCountDesc = {};
+      EuCountDesc.stype = ZE_STRUCTURE_TYPE_EU_COUNT_EXT;
+      DeviceProp.pNext = (void *)&EuCountDesc;
+      ZE2UR_CALL(zeDeviceGetProperties, (ZeDevice, &DeviceProp));
+      if (EuCountDesc.numTotalEUs > 0) {
+        return ReturnValue(uint32_t{EuCountDesc.numTotalEUs});
+      }
+    }
+
     uint32_t count = Device->ZeDeviceProperties->numEUsPerSubslice *
                      Device->ZeDeviceProperties->numSubslicesPerSlice *
                      Device->ZeDeviceProperties->numSlices;

diff --git a/source/adapters/level_zero/platform.cpp b/source/adapters/level_zero/platform.cpp
@@ -270,6 +270,12 @@ ur_result_t ur_platform_handle_t_::initialize() {
         ZeIntelExternalSemaphoreExtensionSupported = true;
       }
     }
+    if (strncmp(extension.name, ZE_EU_COUNT_EXT_NAME,
+                strlen(ZE_EU_COUNT_EXT_NAME) + 1) == 0) {
+      if (extension.version == ZE_EU_COUNT_EXT_VERSION_1_0) {
+        ZeDriverEuCountExtensionFound = true;
+      }
+    }
     zeDriverExtensionMap[extension.name] = extension.version;
   }
 

diff --git a/source/adapters/level_zero/platform.hpp b/source/adapters/level_zero/platform.hpp
@@ -61,6 +61,7 @@ struct ur_platform_handle_t_ : public _ur_platform {
   bool ZeDriverModuleProgramExtensionFound{false};
   bool ZeDriverEventPoolCountingEventsExtensionFound{false};
   bool zeDriverImmediateCommandListAppendFound{false};
+  bool ZeDriverEuCountExtensionFound{false};
 
   // Cache UR devices for reuse
   std::vector<std::unique_ptr<ur_device_handle_t_>> URDevicesCache;