From 02e7f285d8b1f7151c45f522a6478eadad6cb0d4 Mon Sep 17 00:00:00 2001
From: aarongreig <aaron.greig@codeplay.com>
Date: Fri, 19 Apr 2024 10:35:04 +0100
Subject: [PATCH] Merge pull request #1495 from nrspruit/kernel_work_group_size

[L0] Add support for reading maxGroupSize from kernel prop extension
---
 source/adapters/level_zero/common.cpp |  5 +++++
 source/adapters/level_zero/kernel.cpp | 18 ++++++++++++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/source/adapters/level_zero/common.cpp b/source/adapters/level_zero/common.cpp
index af79adeb5d..7ec5f9ef6b 100644
--- 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
index c40e4ef0e3..2cdf4b2cd2 100644
--- 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 {