Merge branch 'oneapi-src:main' into victor/fix_ur_return

Author: Victor Lomuller

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

@@ -245,13 +245,14 @@ setKernelParams(const ur_context_handle_t Context,
           return UR_RESULT_SUCCESS;
         };
 
-        size_t KernelLocalWorkGroupSize = 0;
+        size_t KernelLocalWorkGroupSize = 1;
         for (size_t Dim = 0; Dim < WorkDim; Dim++) {
           auto Err = IsValid(Dim);
           if (Err != UR_RESULT_SUCCESS)
             return Err;
-          // If no error then sum the total local work size per dim.
-          KernelLocalWorkGroupSize += LocalWorkSize[Dim];
+          // If no error then compute the total local work size as a product of
+          // all dims.
+          KernelLocalWorkGroupSize *= LocalWorkSize[Dim];
         }
 
         if (hasExceededMaxRegistersPerBlock(Device, Kernel,
@@ -493,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;
 }