[Spec] fix urKernelSuggestMaxCooperativeGroupCountExp

Add extra param: ur_device_handle_t

It is necessary to implement this function on L0 for
kernels that are build for multiple devices. Right now,
the implementation only works when the kernel is created
from a native handle.

Ref: #2262

Author: igchor

include/ur_api.h

@@ -9434,12 +9434,14 @@ urEnqueueCooperativeKernelLaunchExp(
 ///     - ::UR_RESULT_ERROR_ADAPTER_SPECIFIC
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
 ///         + `NULL == hKernel`
+///         + `NULL == hDevice`
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
 ///         + `NULL == pGroupCountRet`
 ///     - ::UR_RESULT_ERROR_INVALID_KERNEL
 UR_APIEXPORT ur_result_t UR_APICALL
 urKernelSuggestMaxCooperativeGroupCountExp(
     ur_kernel_handle_t hKernel,     ///< [in] handle of the kernel object
+    ur_device_handle_t hDevice,     ///< [in] handle of the device object
     size_t localWorkSize,           ///< [in] number of local work-items that will form a work-group when the
                                     ///< kernel is launched
     size_t dynamicSharedMemorySize, ///< [in] size of dynamic shared memory, for each work-group, in bytes,
@@ -10687,6 +10689,7 @@ typedef struct ur_kernel_set_specialization_constants_params_t {
 ///     allowing the callback the ability to modify the parameter's value
 typedef struct ur_kernel_suggest_max_cooperative_group_count_exp_params_t {
     ur_kernel_handle_t *phKernel;
+    ur_device_handle_t *phDevice;
     size_t *plocalWorkSize;
     size_t *pdynamicSharedMemorySize;
     uint32_t **ppGroupCountRet;

include/ur_ddi.h

@@ -651,6 +651,7 @@ typedef ur_result_t(UR_APICALL *ur_pfnGetKernelProcAddrTable_t)(
 /// @brief Function-pointer for urKernelSuggestMaxCooperativeGroupCountExp
 typedef ur_result_t(UR_APICALL *ur_pfnKernelSuggestMaxCooperativeGroupCountExp_t)(
     ur_kernel_handle_t,
+    ur_device_handle_t,
     size_t,
     size_t,
     uint32_t *);

include/ur_print.hpp

@@ -12214,6 +12214,12 @@ inline std::ostream &operator<<(std::ostream &os, [[maybe_unused]] const struct
     ur::details::printPtr(os,
                           *(params->phKernel));
 
+    os << ", ";
+    os << ".hDevice = ";
+
+    ur::details::printPtr(os,
+                          *(params->phDevice));
+
     os << ", ";
     os << ".localWorkSize = ";
 

scripts/core/exp-cooperative-kernels.yml

@@ -78,6 +78,9 @@ params:
     - type: $x_kernel_handle_t
       name: hKernel
       desc: "[in] handle of the kernel object"
+    - type: $x_device_handle_t
+      name: hDevice
+      desc: "[in] handle of the device object"
     - type: size_t
       name: localWorkSize
       desc: "[in] number of local work-items that will form a work-group when the kernel is launched"

source/adapters/cuda/kernel.cpp

@@ -190,10 +190,13 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetNativeHandle(
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
-    ur_kernel_handle_t hKernel, size_t localWorkSize,
-    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
+    ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
+    size_t localWorkSize, size_t dynamicSharedMemorySize,
+    uint32_t *pGroupCountRet) {
   UR_ASSERT(hKernel, UR_RESULT_ERROR_INVALID_KERNEL);
 
+  std::ignore = hDevice;
+
   // We need to set the active current device for this kernel explicitly here,
   // because the occupancy querying API does not take device parameter.
   ur_device_handle_t Device = hKernel->getProgram()->getDevice();

source/adapters/hip/kernel.cpp

@@ -169,9 +169,11 @@ urKernelGetNativeHandle(ur_kernel_handle_t, ur_native_handle_t *) {
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
-    ur_kernel_handle_t hKernel, size_t localWorkSize,
-    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
+    ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
+    size_t localWorkSize, size_t dynamicSharedMemorySize,
+    uint32_t *pGroupCountRet) {
   std::ignore = hKernel;
+  std::ignore = hDevice;
   std::ignore = localWorkSize;
   std::ignore = dynamicSharedMemorySize;
   std::ignore = pGroupCountRet;

source/adapters/level_zero/kernel.cpp

@@ -1051,14 +1051,17 @@ ur_result_t urKernelGetNativeHandle(
 }
 
 ur_result_t urKernelSuggestMaxCooperativeGroupCountExp(
-    ur_kernel_handle_t hKernel, size_t localWorkSize,
-    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
+    ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
+    size_t localWorkSize, size_t dynamicSharedMemorySize,
+    uint32_t *pGroupCountRet) {
   (void)localWorkSize;
   (void)dynamicSharedMemorySize;
   std::shared_lock<ur_shared_mutex> Guard(hKernel->Mutex);
   uint32_t TotalGroupCount = 0;
+  ze_kernel_handle_t ZeKernel;
+  UR_CALL(getZeKernel(hDevice->ZeDevice, hKernel, &ZeKernel));
   ZE2UR_CALL(zeKernelSuggestMaxCooperativeGroupCount,
-             (hKernel->ZeKernel, &TotalGroupCount));
+             (ZeKernel, &TotalGroupCount));
   *pGroupCountRet = TotalGroupCount;
   return UR_RESULT_SUCCESS;
 }

source/adapters/level_zero/ur_interface_loader.hpp

@@ -687,8 +687,9 @@ ur_result_t urEnqueueCooperativeKernelLaunchExp(
     const size_t *pLocalWorkSize, uint32_t numEventsInWaitList,
     const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent);
 ur_result_t urKernelSuggestMaxCooperativeGroupCountExp(
-    ur_kernel_handle_t hKernel, size_t localWorkSize,
-    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet);
+    ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
+    size_t localWorkSize, size_t dynamicSharedMemorySize,
+    uint32_t *pGroupCountRet);
 ur_result_t urEnqueueTimestampRecordingExp(
     ur_queue_handle_t hQueue, bool blocking, uint32_t numEventsInWaitList,
     const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent);

source/adapters/level_zero/v2/api.cpp

@@ -568,8 +568,9 @@ ur_result_t urCommandBufferCommandGetInfoExp(
 }
 
 ur_result_t urKernelSuggestMaxCooperativeGroupCountExp(
-    ur_kernel_handle_t hKernel, size_t localWorkSize,
-    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
+    ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
+    size_t localWorkSize, size_t dynamicSharedMemorySize,
+    uint32_t *pGroupCountRet) {
   logger::error("{} function not implemented!", __FUNCTION__);
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }

source/adapters/mock/ur_mockddi.cpp

@@ -10003,6 +10003,7 @@ __urdlllocal ur_result_t UR_APICALL urEnqueueCooperativeKernelLaunchExp(
 /// @brief Intercept function for urKernelSuggestMaxCooperativeGroupCountExp
 __urdlllocal ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
     ur_kernel_handle_t hKernel, ///< [in] handle of the kernel object
+    ur_device_handle_t hDevice, ///< [in] handle of the device object
     size_t
         localWorkSize, ///< [in] number of local work-items that will form a work-group when the
                        ///< kernel is launched
@@ -10014,7 +10015,8 @@ __urdlllocal ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
     ur_result_t result = UR_RESULT_SUCCESS;
 
     ur_kernel_suggest_max_cooperative_group_count_exp_params_t params = {
-        &hKernel, &localWorkSize, &dynamicSharedMemorySize, &pGroupCountRet};
+        &hKernel, &hDevice, &localWorkSize, &dynamicSharedMemorySize,
+        &pGroupCountRet};
 
     auto beforeCallback = reinterpret_cast<ur_mock_callback_t>(
         mock::getCallbacks().get_before_callback(