Use hipExtLaunchMultiKernelMultiDevice API (#100)

Depends on HIP version with this pull request:
https://github.com/ROCm-Developer-Tools/HIP/pull/1232

src/include/core.h

@@ -20,17 +20,6 @@
 #include <hip/hip_runtime_api.h>
 #include <hip/hip_runtime.h>
 
-#if CUDART_VERSION < 9000
-struct cudaLaunchParams {
-  void (*func)(struct ncclColl);
-  dim3 gridDim;
-  dim3 blockDim;
-  struct ncclColl **args;
-  size_t sharedMem;
-  hipStream_t stream;
-};
-#endif
-
 #define MAXRINGS 16
 #define MAXTHREADS 256
 #define DEFAULT_BUFFER_SIZE_BYTES (1LL << 22) /* 4MiB */
@@ -248,8 +237,8 @@ struct ncclComm {
   int intraPhase;
 
   // Storage for deferred intra-process launch
-  struct cudaLaunchParams * intraParams;
-  struct cudaLaunchParams *myParams;
+  hipLaunchParams* intraParams;
+  hipLaunchParams* myParams;
   int* intraCudaDevs;
   int* intraCGMode; // Whether we can use CUDA9 CGMD or not
   int* intraCC; // Only to check all have the same ComputeCap and disable CGMode if not

src/init.cu

@@ -455,8 +455,8 @@ void* waitForNonNullPtr(void* p) {
 }
 
 ncclResult_t initParams(struct ncclComm* comm) {
-  struct cudaLaunchParams* params = comm->myParams = comm->intraParams+comm->intraRank;
-  params->args = &comm->argsptr;
+  hipLaunchParams* params = comm->myParams = comm->intraParams+comm->intraRank;
+  params->args = (void **)&comm->argsptr;
   params->stream = NULL;
   params->sharedMem = 0;
   params->blockDim.x = 0; params->blockDim.y = params->blockDim.z = 1;
@@ -489,7 +489,7 @@ ncclResult_t ncclCommSetIntra(struct ncclComm* comm, int rank, int ranks, struct
     comm->intraCC = CC;
   } else {
     comm->intraBarrier = (int*)waitForNonNullPtr(&comm0->intraBarrier);
-    comm->intraParams = (struct cudaLaunchParams*)waitForNonNullPtr(&comm0->intraParams);
+    comm->intraParams = (hipLaunchParams*)waitForNonNullPtr(&comm0->intraParams);
     comm->intraCudaDevs = (int*)waitForNonNullPtr(&comm0->intraCudaDevs);
     comm->intraCGMode = (int*)waitForNonNullPtr(&comm0->intraCGMode);
     comm->intraCC = (int*)waitForNonNullPtr(&comm0->intraCC);
@@ -497,7 +497,7 @@ ncclResult_t ncclCommSetIntra(struct ncclComm* comm, int rank, int ranks, struct
   comm->intraCudaDevs[comm->intraRank] = comm->cudaDev;
   NCCLCHECK(initParams(comm));
 
-  int cgMdLaunch = 0;
+  int cgMdLaunch = 1;
 
   // Set CG Mode
   comm->launchMode = ncclComm::GROUP;

src/misc/enqueue.cu

@@ -61,27 +61,23 @@ static ncclKern_t const ncclKerns[ncclCollCount*ncclNumOps*ncclNumTypes*2] = {
   NCCL_FUNCS2A(ncclAllReduce)
 };
 
-ncclResult_t ncclLaunchCooperativeKernelMultiDevice(struct cudaLaunchParams *paramsList, int* cudaDevs, int numDevices, int cgMode) {
-#if CUDART_VERSION >= 9000
+ncclResult_t ncclLaunchCooperativeKernelMultiDevice(hipLaunchParams *paramsList, int* cudaDevs, int numDevices, int cgMode) {
   if (cgMode & 0x01) {
-    CUDACHECK(cudaLaunchCooperativeKernelMultiDevice(paramsList, numDevices,
-            // These flags are to reduce the latency of using this API
-            cudaCooperativeLaunchMultiDeviceNoPreSync|cudaCooperativeLaunchMultiDeviceNoPostSync));
+    CUDACHECK(hipExtLaunchMultiKernelMultiDevice(paramsList, numDevices, 0));
     return ncclSuccess;
   }
-#endif
   int savedDev;
   CUDACHECK(hipGetDevice(&savedDev));
   for (int i = 0; i < numDevices; i++) {
-    struct cudaLaunchParams* params = paramsList+i;
+    hipLaunchParams* params = paramsList+i;
     CUDACHECK(hipSetDevice(cudaDevs[i]));
-    hipLaunchKernelGGL(params->func, params->gridDim, params->blockDim, params->sharedMem, params->stream, **params->args);
+    hipLaunchKernelGGL((void (*)(struct ncclColl))params->func, params->gridDim, params->blockDim, params->sharedMem, params->stream, **((struct ncclColl **)(params->args)));
   }
   CUDACHECK(hipSetDevice(savedDev));
   return ncclSuccess;
 }
 
-ncclResult_t setupLaunch(struct ncclComm* comm, struct cudaLaunchParams* params) {
+ncclResult_t setupLaunch(struct ncclComm* comm, hipLaunchParams* params) {
   params->gridDim.x = std::min((int) params->gridDim.x, comm->nRings);
 
   // Set active = 2 for the last operation
@@ -97,7 +93,7 @@ ncclResult_t setupLaunch(struct ncclComm* comm, struct cudaLaunchParams* params)
   // As we pass that coll directly, we can free it immediately.
   STORE(&coll->active, 0);
 
-  params->func = ncclKerns[coll->funcIndex];
+  params->func = (void *)ncclKerns[coll->funcIndex];
   return ncclSuccess;
 }
 
@@ -142,7 +138,7 @@ ncclResult_t ncclCpuBarrierOut(struct ncclComm* comm) {
 
 ncclResult_t ncclBarrierEnqueue(struct ncclComm* comm) {
   if (comm->nRanks == 1) return ncclSuccess;
-  struct cudaLaunchParams* params = comm->myParams;
+  hipLaunchParams* params = comm->myParams;
 
   NCCLCHECK(setupLaunch(comm, params));
 
@@ -187,9 +183,9 @@ ncclResult_t ncclBarrierEnqueueWait(ncclComm_t comm) {
 
   NCCLCHECK(ncclCpuBarrierOut(comm));
 
-  struct cudaLaunchParams *params = comm->myParams;
+  hipLaunchParams *params = comm->myParams;
   if (comm->launchMode == ncclComm::PARALLEL) {
-    hipLaunchKernelGGL(params->func, params->gridDim, params->blockDim, params->sharedMem, params->stream, **params->args);
+    hipLaunchKernelGGL((void (*)(struct ncclColl))params->func, params->gridDim, params->blockDim, params->sharedMem, params->stream, **((struct ncclColl **)(params->args)));
   }
   // Start the network proxies as soon as the kernel has been launched. We can't
   // perform any CUDA call between the two or having a hipFree between the CUDA
@@ -207,7 +203,7 @@ ncclResult_t ncclBarrierEnqueueWait(ncclComm_t comm) {
 }
 
 ncclResult_t ncclEnqueueEvents(ncclComm_t comm) {
-  struct cudaLaunchParams *params = comm->myParams;
+  hipLaunchParams *params = comm->myParams;
   // Enqueue event after NCCL kernel
   CUDACHECK(hipEventRecord(comm->doneEvent, params->stream));
   // Use internal NCCL stream for CGMD/GROUP launch if required or if the user stream is NULL