Merge pull request oneapi-src#1246 from 0x12CC/cooperative_kernel_fun…

…ctions

[UR] Add default implementation for cooperative kernel functions

include/ur_api.h

@@ -8692,8 +8692,12 @@ urEnqueueCooperativeKernelLaunchExp(
 ///     - ::UR_RESULT_ERROR_INVALID_KERNEL
 UR_APIEXPORT ur_result_t UR_APICALL
 urKernelSuggestMaxCooperativeGroupCountExp(
-    ur_kernel_handle_t hKernel, ///< [in] handle of the kernel object
-    uint32_t *pGroupCountRet    ///< [out] pointer to maximum number of groups
+    ur_kernel_handle_t hKernel,     ///< [in] handle of the kernel 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,
+                                    ///< that will be used when the kernel is launched
+    uint32_t *pGroupCountRet        ///< [out] pointer to maximum number of groups
 );
 
 #if !defined(__GNUC__)
@@ -9641,6 +9645,8 @@ 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;
+    size_t *plocalWorkSize;
+    size_t *pdynamicSharedMemorySize;
     uint32_t **ppGroupCountRet;
 } ur_kernel_suggest_max_cooperative_group_count_exp_params_t;
 

include/ur_ddi.h

@@ -627,6 +627,8 @@ 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,
+    size_t,
+    size_t,
     uint32_t *);
 
 ///////////////////////////////////////////////////////////////////////////////

include/ur_print.hpp

@@ -11399,6 +11399,16 @@ inline std::ostream &operator<<(std::ostream &os, [[maybe_unused]] const struct
     ur::details::printPtr(os,
                           *(params->phKernel));
 
+    os << ", ";
+    os << ".localWorkSize = ";
+
+    os << *(params->plocalWorkSize);
+
+    os << ", ";
+    os << ".dynamicSharedMemorySize = ";
+
+    os << *(params->pdynamicSharedMemorySize);
+
     os << ", ";
     os << ".pGroupCountRet = ";
 

scripts/core/exp-cooperative-kernels.yml

@@ -78,6 +78,12 @@ params:
     - type: $x_kernel_handle_t
       name: hKernel
       desc: "[in] handle of the kernel object"
+    - type: size_t
+      name: localWorkSize
+      desc: "[in] number of local work-items that will form a work-group when the kernel is launched"
+    - type: size_t
+      name: dynamicSharedMemorySize
+      desc: "[in] size of dynamic shared memory, for each work-group, in bytes, that will be used when the kernel is launched"
     - type: "uint32_t*"
       name: "pGroupCountRet"
       desc: "[out] pointer to maximum number of groups"

source/adapters/cuda/enqueue.cpp

@@ -494,6 +494,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
   return Result;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urEnqueueCooperativeKernelLaunchExp(
+    ur_queue_handle_t hQueue, ur_kernel_handle_t hKernel, uint32_t workDim,
+    const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
+    const size_t *pLocalWorkSize, uint32_t numEventsInWaitList,
+    const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent) {
+  return urEnqueueKernelLaunch(hQueue, hKernel, workDim, pGlobalWorkOffset,
+                               pGlobalWorkSize, pLocalWorkSize,
+                               numEventsInWaitList, phEventWaitList, phEvent);
+}
+
 /// Set parameters for general 3D memory copy.
 /// If the source and/or destination is on the device, SrcPtr and/or DstPtr
 /// must be a pointer to a CUdeviceptr

source/adapters/cuda/kernel.cpp

@@ -169,6 +169,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetNativeHandle(
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
+    ur_kernel_handle_t hKernel, size_t localWorkSize,
+    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
+  (void)hKernel;
+  (void)localWorkSize;
+  (void)dynamicSharedMemorySize;
+  *pGroupCountRet = 1;
+  return UR_RESULT_SUCCESS;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgValue(
     ur_kernel_handle_t hKernel, uint32_t argIndex, size_t argSize,
     const ur_kernel_arg_value_properties_t *pProperties,

source/adapters/cuda/ur_interface_loader.cpp

@@ -404,7 +404,8 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetEnqueueExpProcAddrTable(
     return result;
   }
 
-  pDdiTable->pfnCooperativeKernelLaunchExp = nullptr;
+  pDdiTable->pfnCooperativeKernelLaunchExp =
+      urEnqueueCooperativeKernelLaunchExp;
 
   return UR_RESULT_SUCCESS;
 }
@@ -416,7 +417,8 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetKernelExpProcAddrTable(
     return result;
   }
 
-  pDdiTable->pfnSuggestMaxCooperativeGroupCountExp = nullptr;
+  pDdiTable->pfnSuggestMaxCooperativeGroupCountExp =
+      urKernelSuggestMaxCooperativeGroupCountExp;
 
   return UR_RESULT_SUCCESS;
 }

source/adapters/hip/enqueue.cpp

@@ -465,6 +465,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
   return Result;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urEnqueueCooperativeKernelLaunchExp(
+    ur_queue_handle_t hQueue, ur_kernel_handle_t hKernel, uint32_t workDim,
+    const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
+    const size_t *pLocalWorkSize, uint32_t numEventsInWaitList,
+    const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent) {
+  return urEnqueueKernelLaunch(hQueue, hKernel, workDim, pGlobalWorkOffset,
+                               pGlobalWorkSize, pLocalWorkSize,
+                               numEventsInWaitList, phEventWaitList, phEvent);
+}
+
 /// Enqueues a wait on the given queue for all events.
 /// See \ref enqueueEventWait
 ///

source/adapters/hip/kernel.cpp

@@ -158,6 +158,16 @@ urKernelGetNativeHandle(ur_kernel_handle_t, ur_native_handle_t *) {
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
+    ur_kernel_handle_t hKernel, size_t localWorkSize,
+    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
+  (void)hKernel;
+  (void)localWorkSize;
+  (void)dynamicSharedMemorySize;
+  *pGroupCountRet = 1;
+  return UR_RESULT_SUCCESS;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgValue(
     ur_kernel_handle_t hKernel, uint32_t argIndex, size_t argSize,
     const ur_kernel_arg_value_properties_t *, const void *pArgValue) {

source/adapters/hip/ur_interface_loader.cpp

@@ -374,7 +374,8 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetEnqueueExpProcAddrTable(
     return result;
   }
 
-  pDdiTable->pfnCooperativeKernelLaunchExp = nullptr;
+  pDdiTable->pfnCooperativeKernelLaunchExp =
+      urEnqueueCooperativeKernelLaunchExp;
 
   return UR_RESULT_SUCCESS;
 }
@@ -386,7 +387,8 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetKernelExpProcAddrTable(
     return result;
   }
 
-  pDdiTable->pfnSuggestMaxCooperativeGroupCountExp = nullptr;
+  pDdiTable->pfnSuggestMaxCooperativeGroupCountExp =
+      urKernelSuggestMaxCooperativeGroupCountExp;
 
   return UR_RESULT_SUCCESS;
 }

source/adapters/level_zero/kernel.cpp

@@ -264,6 +264,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
   return UR_RESULT_SUCCESS;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urEnqueueCooperativeKernelLaunchExp(
+    ur_queue_handle_t hQueue, ur_kernel_handle_t hKernel, uint32_t workDim,
+    const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
+    const size_t *pLocalWorkSize, uint32_t numEventsInWaitList,
+    const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent) {
+  return urEnqueueKernelLaunch(hQueue, hKernel, workDim, pGlobalWorkOffset,
+                               pGlobalWorkSize, pLocalWorkSize,
+                               numEventsInWaitList, phEventWaitList, phEvent);
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urEnqueueDeviceGlobalVariableWrite(
     ur_queue_handle_t Queue,     ///< [in] handle of the queue to submit to.
     ur_program_handle_t Program, ///< [in] handle of the program containing the
@@ -787,6 +797,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetNativeHandle(
   return UR_RESULT_SUCCESS;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
+    ur_kernel_handle_t hKernel, size_t localWorkSize,
+    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
+  (void)hKernel;
+  (void)localWorkSize;
+  (void)dynamicSharedMemorySize;
+  *pGroupCountRet = 1;
+  return UR_RESULT_SUCCESS;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urKernelCreateWithNativeHandle(
     ur_native_handle_t NativeKernel, ///< [in] the native handle of the kernel.
     ur_context_handle_t Context,     ///< [in] handle of the context object

source/adapters/level_zero/ur_interface_loader.cpp

@@ -451,7 +451,8 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetEnqueueExpProcAddrTable(
     return result;
   }
 
-  pDdiTable->pfnCooperativeKernelLaunchExp = nullptr;
+  pDdiTable->pfnCooperativeKernelLaunchExp =
+      urEnqueueCooperativeKernelLaunchExp;
 
   return UR_RESULT_SUCCESS;
 }
@@ -463,7 +464,8 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetKernelExpProcAddrTable(
     return result;
   }
 
-  pDdiTable->pfnSuggestMaxCooperativeGroupCountExp = nullptr;
+  pDdiTable->pfnSuggestMaxCooperativeGroupCountExp =
+      urKernelSuggestMaxCooperativeGroupCountExp;
 
   return UR_RESULT_SUCCESS;
 }

source/adapters/null/ur_nullddi.cpp

@@ -5443,15 +5443,22 @@ __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
-    uint32_t *pGroupCountRet    ///< [out] pointer to maximum number of groups
+    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,
+    ///< that will be used when the kernel is launched
+    uint32_t *pGroupCountRet ///< [out] pointer to maximum number of groups
     ) try {
     ur_result_t result = UR_RESULT_SUCCESS;
 
     // if the driver has created a custom function, then call it instead of using the generic path
     auto pfnSuggestMaxCooperativeGroupCountExp =
         d_context.urDdiTable.KernelExp.pfnSuggestMaxCooperativeGroupCountExp;
     if (nullptr != pfnSuggestMaxCooperativeGroupCountExp) {
-        result = pfnSuggestMaxCooperativeGroupCountExp(hKernel, pGroupCountRet);
+        result = pfnSuggestMaxCooperativeGroupCountExp(
+            hKernel, localWorkSize, dynamicSharedMemorySize, pGroupCountRet);
     } else {
         // generic implementation
     }

source/adapters/opencl/enqueue.cpp

@@ -41,6 +41,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
   return UR_RESULT_SUCCESS;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urEnqueueCooperativeKernelLaunchExp(
+    ur_queue_handle_t hQueue, ur_kernel_handle_t hKernel, uint32_t workDim,
+    const size_t *pGlobalWorkOffset, const size_t *pGlobalWorkSize,
+    const size_t *pLocalWorkSize, uint32_t numEventsInWaitList,
+    const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent) {
+  return urEnqueueKernelLaunch(hQueue, hKernel, workDim, pGlobalWorkOffset,
+                               pGlobalWorkSize, pLocalWorkSize,
+                               numEventsInWaitList, phEventWaitList, phEvent);
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urEnqueueEventsWait(
     ur_queue_handle_t hQueue, uint32_t numEventsInWaitList,
     const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent) {

source/adapters/opencl/kernel.cpp

@@ -10,6 +10,7 @@
 #include "common.hpp"
 
 #include <algorithm>
+#include <cstddef>
 #include <memory>
 
 UR_APIEXPORT ur_result_t UR_APICALL
@@ -376,6 +377,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetNativeHandle(
   return UR_RESULT_SUCCESS;
 }
 
+UR_APIEXPORT ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
+    ur_kernel_handle_t hKernel, size_t localWorkSize,
+    size_t dynamicSharedMemorySize, uint32_t *pGroupCountRet) {
+  (void)hKernel;
+  (void)localWorkSize;
+  (void)dynamicSharedMemorySize;
+  *pGroupCountRet = 1;
+  return UR_RESULT_SUCCESS;
+}
+
 UR_APIEXPORT ur_result_t UR_APICALL urKernelCreateWithNativeHandle(
     ur_native_handle_t hNativeKernel, ur_context_handle_t, ur_program_handle_t,
     const ur_kernel_native_properties_t *pProperties,

source/adapters/opencl/ur_interface_loader.cpp

@@ -395,7 +395,8 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetEnqueueExpProcAddrTable(
     return result;
   }
 
-  pDdiTable->pfnCooperativeKernelLaunchExp = nullptr;
+  pDdiTable->pfnCooperativeKernelLaunchExp =
+      urEnqueueCooperativeKernelLaunchExp;
 
   return UR_RESULT_SUCCESS;
 }
@@ -407,7 +408,8 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetKernelExpProcAddrTable(
     return result;
   }
 
-  pDdiTable->pfnSuggestMaxCooperativeGroupCountExp = nullptr;
+  pDdiTable->pfnSuggestMaxCooperativeGroupCountExp =
+      urKernelSuggestMaxCooperativeGroupCountExp;
 
   return UR_RESULT_SUCCESS;
 }

source/loader/layers/tracing/ur_trcddi.cpp

@@ -6037,7 +6037,13 @@ __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
-    uint32_t *pGroupCountRet    ///< [out] pointer to maximum number of groups
+    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,
+    ///< that will be used when the kernel is launched
+    uint32_t *pGroupCountRet ///< [out] pointer to maximum number of groups
 ) {
     auto pfnSuggestMaxCooperativeGroupCountExp =
         context.urDdiTable.KernelExp.pfnSuggestMaxCooperativeGroupCountExp;
@@ -6047,13 +6053,13 @@ __urdlllocal ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
     }
 
     ur_kernel_suggest_max_cooperative_group_count_exp_params_t params = {
-        &hKernel, &pGroupCountRet};
+        &hKernel, &localWorkSize, &dynamicSharedMemorySize, &pGroupCountRet};
     uint64_t instance = context.notify_begin(
         UR_FUNCTION_KERNEL_SUGGEST_MAX_COOPERATIVE_GROUP_COUNT_EXP,
         "urKernelSuggestMaxCooperativeGroupCountExp", &params);
 
-    ur_result_t result =
-        pfnSuggestMaxCooperativeGroupCountExp(hKernel, pGroupCountRet);
+    ur_result_t result = pfnSuggestMaxCooperativeGroupCountExp(
+        hKernel, localWorkSize, dynamicSharedMemorySize, pGroupCountRet);
 
     context.notify_end(
         UR_FUNCTION_KERNEL_SUGGEST_MAX_COOPERATIVE_GROUP_COUNT_EXP,

source/loader/layers/validation/ur_valddi.cpp

@@ -8827,7 +8827,13 @@ __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
-    uint32_t *pGroupCountRet    ///< [out] pointer to maximum number of groups
+    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,
+    ///< that will be used when the kernel is launched
+    uint32_t *pGroupCountRet ///< [out] pointer to maximum number of groups
 ) {
     auto pfnSuggestMaxCooperativeGroupCountExp =
         context.urDdiTable.KernelExp.pfnSuggestMaxCooperativeGroupCountExp;
@@ -8851,8 +8857,8 @@ __urdlllocal ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
         refCountContext.logInvalidReference(hKernel);
     }
 
-    ur_result_t result =
-        pfnSuggestMaxCooperativeGroupCountExp(hKernel, pGroupCountRet);
+    ur_result_t result = pfnSuggestMaxCooperativeGroupCountExp(
+        hKernel, localWorkSize, dynamicSharedMemorySize, pGroupCountRet);
 
     return result;
 }

source/loader/ur_ldrddi.cpp

@@ -7571,7 +7571,13 @@ __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
-    uint32_t *pGroupCountRet    ///< [out] pointer to maximum number of groups
+    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,
+    ///< that will be used when the kernel is launched
+    uint32_t *pGroupCountRet ///< [out] pointer to maximum number of groups
 ) {
     ur_result_t result = UR_RESULT_SUCCESS;
 
@@ -7587,7 +7593,8 @@ __urdlllocal ur_result_t UR_APICALL urKernelSuggestMaxCooperativeGroupCountExp(
     hKernel = reinterpret_cast<ur_kernel_object_t *>(hKernel)->handle;
 
     // forward to device-platform
-    result = pfnSuggestMaxCooperativeGroupCountExp(hKernel, pGroupCountRet);
+    result = pfnSuggestMaxCooperativeGroupCountExp(
+        hKernel, localWorkSize, dynamicSharedMemorySize, pGroupCountRet);
 
     return result;
 }