From 5d02a5a754a8b87b0dfc27014ad6b7db465956ab Mon Sep 17 00:00:00 2001
From: "Kenneth Benzie (Benie)" <k.benzie@codeplay.com>
Date: Thu, 16 May 2024 11:11:37 +0100
Subject: [PATCH] Merge pull request #1611 from
 nrspruit/fix_interop_wg_max_size

[L0] ensure a valid kernel handle for the device when reading max wg
---
 source/adapters/level_zero/kernel.cpp | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/source/adapters/level_zero/kernel.cpp b/source/adapters/level_zero/kernel.cpp
index 9a0f1b1059..44d3654448 100644
--- a/source/adapters/level_zero/kernel.cpp
+++ b/source/adapters/level_zero/kernel.cpp
@@ -582,15 +582,26 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetGroupInfo(
 
     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) {
+    // Set the Kernel to use as the ZeKernel initally for native handle support.
+    // This makes the assumption that this device is the same device where this
+    // kernel was created.
+    auto ZeKernelDevice = Kernel->ZeKernel;
+    auto It = Kernel->ZeKernelMap.find(Device->ZeDevice);
+    if (It != Kernel->ZeKernelMap.end()) {
+      ZeKernelDevice = Kernel->ZeKernelMap[Device->ZeDevice];
+    }
+    if (ZeKernelDevice) {
+      auto ZeResult = ZE_CALL_NOCHECK(zeKernelGetProperties,
+                                      (ZeKernelDevice, &kernelProperties));
+      if (ZeResult || workGroupProperties.maxGroupSize == 0) {
+        return ReturnValue(
+            uint64_t{Device->ZeDeviceComputeProperties->maxTotalGroupSize});
+      }
+      return ReturnValue(workGroupProperties.maxGroupSize);
+    } else {
       return ReturnValue(
           uint64_t{Device->ZeDeviceComputeProperties->maxTotalGroupSize});
     }
-    return ReturnValue(workGroupProperties.maxGroupSize);
   }
   case UR_KERNEL_GROUP_INFO_COMPILE_WORK_GROUP_SIZE: {
     struct {