Merge pull request #1495 from nrspruit/kernel_work_group_size

[L0] Add support for reading maxGroupSize from kernel prop extension
oneapi-src · Apr 19, 2024 · 02e7f28 · 02e7f28
1 parent ce58dfd
commit 02e7f28
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/source/adapters/level_zero/common.cpp b/source/adapters/level_zero/common.cpp
@@ -215,6 +215,11 @@ ze_structure_type_t
 getZeStructureType<ze_relaxed_allocation_limits_exp_desc_t>() {
   return ZE_STRUCTURE_TYPE_RELAXED_ALLOCATION_LIMITS_EXP_DESC;
 }
+template <>
+ze_structure_type_t
+getZeStructureType<ze_kernel_max_group_size_properties_ext_t>() {
+  return ZE_STRUCTURE_TYPE_KERNEL_MAX_GROUP_SIZE_EXT_PROPERTIES;
+}
 template <> ze_structure_type_t getZeStructureType<ze_host_mem_alloc_desc_t>() {
   return ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC;
 }

diff --git a/source/adapters/level_zero/kernel.cpp b/source/adapters/level_zero/kernel.cpp
@@ -574,10 +574,20 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetGroupInfo(
     return ReturnValue(GlobalWorkSize);
   }
   case UR_KERNEL_GROUP_INFO_WORK_GROUP_SIZE: {
-    // As of right now, L0 is missing API to query kernel and device specific
-    // max work group size.
-    return ReturnValue(
-        uint64_t{Device->ZeDeviceComputeProperties->maxTotalGroupSize});
+    ZeStruct<ze_kernel_max_group_size_properties_ext_t> workGroupProperties;
+    workGroupProperties.maxGroupSize = 0;
+
+    ZeStruct<ze_kernel_properties_t> kernelProperties;
+    kernelProperties.pNext = &workGroupProperties;
+
+    auto ZeResult = ZE_CALL_NOCHECK(
+        zeKernelGetProperties,
+        (Kernel->ZeKernelMap[Device->ZeDevice], &kernelProperties));
+    if (ZeResult || workGroupProperties.maxGroupSize == 0) {
+      return ReturnValue(
+          uint64_t{Device->ZeDeviceComputeProperties->maxTotalGroupSize});
+    }
+    return ReturnValue(workGroupProperties.maxGroupSize);
   }
   case UR_KERNEL_GROUP_INFO_COMPILE_WORK_GROUP_SIZE: {
     struct {