From e3578923bfffbd3fa63f2221d723d23cb08e1185 Mon Sep 17 00:00:00 2001
From: "Wu, Yingcong" <yingcong.wu@intel.com>
Date: Fri, 9 May 2025 06:03:36 +0200
Subject: [PATCH] respect workdim

Signed-off-by: Wu, Yingcong <yingcong.wu@intel.com>
---
 .../source/adapters/level_zero/kernel.cpp          | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/unified-runtime/source/adapters/level_zero/kernel.cpp b/unified-runtime/source/adapters/level_zero/kernel.cpp
index 8658adc58bcf0..71da5d2e7d2f3 100644
--- a/unified-runtime/source/adapters/level_zero/kernel.cpp
+++ b/unified-runtime/source/adapters/level_zero/kernel.cpp
@@ -279,15 +279,11 @@ ur_result_t urEnqueueCooperativeKernelLaunchExp(
 
   if (LocalWorkSize) {
     // L0
-    UR_ASSERT(LocalWorkSize[0] < (std::numeric_limits<uint32_t>::max)(),
-              UR_RESULT_ERROR_INVALID_VALUE);
-    UR_ASSERT(LocalWorkSize[1] < (std::numeric_limits<uint32_t>::max)(),
-              UR_RESULT_ERROR_INVALID_VALUE);
-    UR_ASSERT(LocalWorkSize[2] < (std::numeric_limits<uint32_t>::max)(),
-              UR_RESULT_ERROR_INVALID_VALUE);
-    WG[0] = static_cast<uint32_t>(LocalWorkSize[0]);
-    WG[1] = static_cast<uint32_t>(LocalWorkSize[1]);
-    WG[2] = static_cast<uint32_t>(LocalWorkSize[2]);
+    for (int I = 0; I < WorkDim; I++) {
+      UR_ASSERT(LocalWorkSize[I] < (std::numeric_limits<uint32_t>::max)(),
+                UR_RESULT_ERROR_INVALID_VALUE);
+      WG[I] = static_cast<uint32_t>(LocalWorkSize[I]);
+    }
   } else {
     // We can't call to zeKernelSuggestGroupSize if 64-bit GlobalWorkSize
     // values do not fit to 32-bit that the API only supports currently.