[hip] add initial implementation for hipLaunchCooperativeKernel API (#1339)

* [hip] add initial implementation for hipLaunchCooperativeKernel API

* [hip] use total number of work groups to initialize the GWS resource

* [hip] use only one argument for init_gws kernel

* [hip] use the device associated with the stream for checking the device properties


[ROCm/clr commit: 32ce882d6e]

projects/clr/hipamd/include/hip/hcc_detail/device_library_decls.h

@@ -59,6 +59,10 @@ extern "C" __device__ __attribute__((const)) float __ocml_trunc_f32(float);
 extern "C" __device__ __attribute__((const)) float __ocml_fmin_f32(float, float);
 extern "C" __device__ __attribute__((const)) float __ocml_fmax_f32(float, float);
 
+extern "C" __device__ __attribute__((convergent)) void __ockl_gws_init(uint nwm1, uint rid);
+extern "C" __device__ __attribute__((convergent)) void __ockl_gws_barrier(uint nwm1, uint rid);
+
+
 // Introduce local address space
 #define __local __attribute__((address_space(3)))
 

projects/clr/hipamd/src/hip_module.cpp

@@ -386,6 +386,100 @@ hipError_t hipExtLaunchMultiKernelMultiDevice(hipLaunchParams* launchParamsList,
     return ihipLogStatus(result);
 }
 
+namespace {
+// kernel for initializing GWS
+// nwm1 is the total number of work groups minus 1
+__global__ void init_gws(uint nwm1) {
+    __ockl_gws_init(nwm1, 0);
+}
+}
+
+hipError_t hipLaunchCooperativeKernel(const void* f, dim3 gridDim,
+        dim3 blockDimX, void** kernelParams, unsigned int sharedMemBytes,
+        hipStream_t stream) {
+
+    HIP_INIT_API(hipLaunchCooperativeKernel, f, gridDim, blockDimX, kernelParams, sharedMemBytes, stream);
+    hipError_t result;
+
+
+    if ((f == nullptr) || (stream == nullptr) || (kernelParams == nullptr)) {
+        return ihipLogStatus(hipErrorNotInitialized);
+    }
+
+    if (!stream->getDevice()->_props.cooperativeLaunch) {
+        return ihipLogStatus(hipErrorInvalidConfiguration);
+    }
+
+    // Prepare the kernel descriptor for initializing the GWS
+    hipFunction_t gwsKD = hip_impl::get_program_state().kernel_descriptor(
+            reinterpret_cast<std::uintptr_t>(&init_gws),
+            hip_impl::target_agent(stream));
+
+    if (gwsKD == nullptr) {
+        return ihipLogStatus(hipErrorInvalidValue);
+    }
+    hip_impl::kernargs_size_align gwsKargs =
+            hip_impl::get_program_state().get_kernargs_size_align(
+                    reinterpret_cast<std::uintptr_t>(&init_gws));
+
+    gwsKD->_kernarg_layout = *reinterpret_cast<const std::vector<
+            std::pair<std::size_t, std::size_t>>*>(gwsKargs.getHandle());
+
+    // Prepare the kernel descriptor for the main kernel
+    hipFunction_t kd = hip_impl::get_program_state().kernel_descriptor(
+            reinterpret_cast<std::uintptr_t>(f),
+            hip_impl::target_agent(stream));
+    if (kd == nullptr) {
+        return ihipLogStatus(hipErrorInvalidValue);
+    }
+    hip_impl::kernargs_size_align kargs =
+            hip_impl::get_program_state().get_kernargs_size_align(
+                    reinterpret_cast<std::uintptr_t>(f));
+
+    kd->_kernarg_layout = *reinterpret_cast<const std::vector<
+            std::pair<std::size_t, std::size_t>>*>(kargs.getHandle());
+
+
+    void *gwsKernelParam[1];
+    // calculate total number of work groups minus 1 for the main kernel
+    uint nwm1 = (gridDim.x * gridDim.y * gridDim.z) - 1;
+    gwsKernelParam[0] = &nwm1;
+
+    LockedAccessor_StreamCrit_t streamCrit(stream->criticalData(), false);
+#if (__hcc_workweek__ >= 19213)
+    streamCrit->_av.acquire_locked_hsa_queue();
+#endif
+
+    // launch the init_gws kernel to initialize the GWS
+    result = ihipModuleLaunchKernel(tls, gwsKD, 1, 1, 1, 1, 1, 1,
+             0, stream, gwsKernelParam, nullptr, nullptr, nullptr, 0, true);
+
+    if (result != hipSuccess) {
+        stream->criticalData().unlock();
+#if (__hcc_workweek__ >= 19213)
+        stream->criticalData()._av.release_locked_hsa_queue();
+#endif
+
+        return ihipLogStatus(hipErrorLaunchFailure);
+    }
+
+    // launch the main kernel
+    result = ihipModuleLaunchKernel(tls, kd,
+            gridDim.x * blockDimX.x,
+            gridDim.y * blockDimX.y,
+            gridDim.z * blockDimX.z,
+            blockDimX.x, blockDimX.y, blockDimX.z,
+            sharedMemBytes, stream, kernelParams, nullptr, nullptr,
+            nullptr, 0, true);
+
+    stream->criticalData().unlock();
+#if (__hcc_workweek__ >= 19213)
+    stream->criticalData()._av.release_locked_hsa_queue();
+#endif
+
+    return ihipLogStatus(result);
+}
+
 namespace hip_impl {
     hsa_executable_t executable_for(hipModule_t hmod) {
         return hmod->executable;