Use new HIP graph API compatible with CUDA 11030 (#991)

* Use new HIP graph API compatible with CUDA 11030

* Update dependency to ROCm 6.1

* Fix single stream use case

CMakeLists.txt

@@ -316,7 +316,6 @@ set(SRC_FILES
   src/collectives/scatter.cc
   src/collectives/sendrecv.cc
   src/debug.cc
-  src/enhcompat.cc
   src/enqueue.cc
   src/graph/connect.cc
   src/graph/paths.cc

src/enqueue.cc

@@ -944,14 +944,14 @@ ncclResult_t ncclLaunchPrepare(struct ncclComm* comm) {
     cudaStream_t launchStream = tasks->streams->stream;
     NCCLCHECKGOTO(ncclStrongStreamAcquire(tasks->capturingGraph, &comm->sharedRes->deviceStream), result, failure);
 
-    if (tasks->numStreams != 1) {
+    if (tasks->numStreams != 1 || persistent) {
       // Create dependency for device stream on user streams. First from extra user
       // streams to deviceStream. Then deviceStream to first user stream.
       for (struct ncclCudaStreamList* l=tasks->streams->next; l != nullptr; l = l->next) {
         NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, &comm->sharedRes->deviceStream, l->stream), result, failure);
       }
       NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, launchStream, &comm->sharedRes->deviceStream), result, failure);
-    } else if (tasks->streams->stream != comm->lastStream && comm->lastStream != nullptr) {
+    } else if (tasks->streams->stream != comm->lastStream && comm->lastStream != nullptr && !persistent) {
       // Stream changed from last call, create dependency against last NCCL kernel launch
       CUDACHECK(hipStreamWaitEvent(tasks->streams->stream, comm->doneEvent, 0));
     }
@@ -971,7 +971,7 @@ ncclResult_t ncclLaunchPrepare(struct ncclComm* comm) {
       }
       if (acquired) {
         // Make to-be-launched kernels dependent on just-launched host stream tasks.
-        if (tasks->numStreams != 1) NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, launchStream, &comm->sharedRes->hostStream), result, failure);
+        NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, launchStream, &comm->sharedRes->hostStream), result, failure);
         NCCLCHECKGOTO(ncclStrongStreamRelease(tasks->capturingGraph, &comm->sharedRes->hostStream), result, failure);
       }
     }
@@ -1010,7 +1010,7 @@ ncclResult_t ncclLaunchKernel(struct ncclComm* comm, struct ncclKernelPlan* plan
   dim3 block = {(unsigned)plan->threadPerBlock, 1, 1};
   size_t smem = ncclShmemDynamicSize(comm->cudaArch);
   void *args[3] = {&comm->devComm, &plan->channelMask, &plan->workHead};
-  if (tasks->numStreams == 1) {
+  if (tasks->numStreams == 1 && !plan->persistent) {
     CUDACHECK(hipExtLaunchKernel(plan->kernelFn, grid, block, args, 0, tasks->streams->stream, NULL, comm->doneEvent, 0));
     comm->lastStream = tasks->streams->stream;
     return ncclSuccess;
@@ -1079,6 +1079,7 @@ ncclResult_t ncclLaunchKernelAfter_NoCuda(struct ncclComm* comm, struct ncclKern
 ncclResult_t ncclLaunchFinish(struct ncclComm* comm) {
   ncclResult_t result = ncclSuccess;
   struct ncclTasks* tasks = &comm->tasks;
+  bool persistent = ncclCudaGraphValid(tasks->capturingGraph);
   tasks->collBytesTotal = 0; // Just in case subtraction during scheduleCollTasksToPlan() doesn't get to 0
 
   // Deallocate ncclWork's. This frame exists so long as ncclLaunchPrepare
@@ -1093,14 +1094,14 @@ ncclResult_t ncclLaunchFinish(struct ncclComm* comm) {
     // Create dependency for deviceStream on launchStream. We know that deviceStream
     // hasn't been modified since launchStream waited on it (in ncclLaunchPrepare),
     // so we can say that launchStream subsumes it.
-    if (tasks->numStreams != 1) NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, &comm->sharedRes->deviceStream, launchStream, /*b_subsumes_a=*/true), result, resume1);
+    if (persistent || tasks->numStreams != 1) NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, &comm->sharedRes->deviceStream, launchStream, /*b_subsumes_a=*/true), result, resume1);
   resume1:
     // Create dependency for other user streams (skip launch stream) on deviceStream.
     // Again, the user streams haven't been touched since deviceStream waited on them
     // so we can say they are subsumed by deviceStream.
     struct ncclCudaStreamList* sl = tasks->streams->next;
     tasks->streams = nullptr; // Reset comm->tasks.streams to empty.
-    while (sl != nullptr && tasks->numStreams != 1) {
+    while (sl != nullptr && (tasks->numStreams != 1 || persistent)) {
       NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, sl->stream, &comm->sharedRes->deviceStream, /*b_subsumes_a=*/true), result, resume2);
     resume2:
       sl = sl->next;

src/include/strongstream.h

@@ -16,7 +16,7 @@
  * easily.
  */
 struct ncclCudaGraph {
-#if CUDART_VERSION >= 11030
+#if ROCM_VERSION >= 60100
   cudaGraph_t graph;
   unsigned long long graphId;
 #endif
@@ -24,7 +24,7 @@ struct ncclCudaGraph {
 
 inline struct ncclCudaGraph ncclCudaGraphNone() {
   struct ncclCudaGraph tmp;
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     tmp.graph = nullptr;
     tmp.graphId = ULLONG_MAX;
   #endif
@@ -32,7 +32,7 @@ inline struct ncclCudaGraph ncclCudaGraphNone() {
 }
 
 inline bool ncclCudaGraphValid(struct ncclCudaGraph graph) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     return graph.graph != nullptr;
   #else
     return false;
@@ -40,7 +40,7 @@ inline bool ncclCudaGraphValid(struct ncclCudaGraph graph) {
 }
 
 inline bool ncclCudaGraphSame(struct ncclCudaGraph a, struct ncclCudaGraph b) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     return a.graphId == b.graphId;
   #else
     return true;
@@ -123,7 +123,7 @@ struct ncclStrongStreamGraph; // internal to ncclStrongStream
 struct ncclStrongStream {
   // Used when not graph capturing.
   cudaStream_t cudaStream;
-#if CUDART_VERSION >= 11030
+#if ROCM_VERSION >= 60100
   // The event used to establish order between graphs and streams. During acquire
   // this event is waited on, during release it is recorded to.
   cudaEvent_t serialEvent;

src/misc/strongstream.cc

@@ -38,39 +38,21 @@ static void ncclStrongStreamGraphDelete(struct ncclStrongStreamGraph* g) {
 ncclResult_t ncclCudaGetCapturingGraph(
     struct ncclCudaGraph* graph, cudaStream_t stream
   ) {
-  #if CUDART_VERSION >= 10000 // cudaStreamGetCaptureInfo
-    int driver;
-    NCCLCHECK(ncclCudaDriverVersion(&driver));
-    if (CUDART_VERSION < 11030 || driver < 11030) {
-      cudaStreamCaptureStatus status;
-      unsigned long long gid;
-      CUDACHECK(cudaStreamGetCaptureInfo(stream, &status, &gid));
-      #if CUDART_VERSION >= 11030
-        graph->graph = nullptr;
-        graph->graphId = ULLONG_MAX;
-      #endif
-      if (status != cudaStreamCaptureStatusNone) {
-        WARN("NCCL cannot be captured in a graph if either it wasn't built with CUDA runtime >= 11.3 or if the installed CUDA driver < R465.");
-        return ncclInvalidUsage;
-      }
-    } else {
-      #if CUDART_VERSION >= 11030
-        cudaStreamCaptureStatus status;
-        unsigned long long gid;
-        CUDACHECK(cudaStreamGetCaptureInfo_v2(stream, &status, &gid, &graph->graph, nullptr, nullptr));
-        if (status != cudaStreamCaptureStatusActive) {
-          graph->graph = nullptr;
-          gid = ULLONG_MAX;
-        }
-        graph->graphId = gid;
-      #endif
-    }
-  #endif
+#if ROCM_VERSION >= 60100
+  hipStreamCaptureStatus status;
+  unsigned long long gid;
+  CUDACHECK(hipStreamGetCaptureInfo_v2(stream, &status, &gid, &graph->graph, nullptr, nullptr));
+  if (status != hipStreamCaptureStatusActive) {
+    graph->graph = nullptr;
+    gid = ULLONG_MAX;
+  }
+  graph->graphId = gid;
+#endif
   return ncclSuccess;
 }
 
 ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, cudaHostFn_t fn, void* arg) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     cudaUserObject_t object;
     CUDACHECK(cudaUserObjectCreate(
       &object, arg, fn, /*initialRefcount=*/1, cudaUserObjectNoDestructorSync
@@ -87,7 +69,7 @@ ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, cudaHostFn_t
 
 ncclResult_t ncclStrongStreamConstruct(struct ncclStrongStream* ss) {
   CUDACHECK(cudaStreamCreateWithFlags(&ss->cudaStream, cudaStreamNonBlocking));
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     CUDACHECK(cudaEventCreateWithFlags(&ss->serialEvent, cudaEventDisableTiming));
     ss->everCaptured = false;
     ss->serialEventNeedsRecord = false;
@@ -108,7 +90,7 @@ static void graphDestructor(void* arg) {
 
 ncclResult_t ncclStrongStreamDestruct(struct ncclStrongStream* ss) {
   CUDACHECK(cudaStreamDestroy(ss->cudaStream));
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     CUDACHECK(cudaEventDestroy(ss->serialEvent));
     // Delete list of per-graph chains.
     struct ncclStrongStreamGraph* g = ss->graphHead;
@@ -126,7 +108,7 @@ ncclResult_t ncclStrongStreamDestruct(struct ncclStrongStream* ss) {
   return ncclSuccess;
 }
 
-NCCL_PARAM(GraphMixingSupport, "GRAPH_MIXING_SUPPORT", 1)
+NCCL_PARAM(GraphMixingSupport, "GRAPH_MIXING_SUPPORT", 0)
 
 static void ensureTips(struct ncclStrongStreamGraph* g, int n) {
   if (g->tipCapacity < n) {
@@ -138,7 +120,7 @@ static void ensureTips(struct ncclStrongStreamGraph* g, int n) {
 ncclResult_t ncclStrongStreamAcquire(
     struct ncclCudaGraph graph, struct ncclStrongStream* ss
   ) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     bool mixing = ncclParamGraphMixingSupport();
     if (graph.graph == nullptr) {
       if (mixing && ss->everCaptured) {
@@ -200,7 +182,7 @@ ncclResult_t ncclStrongStreamAcquire(
 }
 
 ncclResult_t ncclStrongStreamAcquireUncaptured(struct ncclStrongStream* ss) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     bool mixing = ncclParamGraphMixingSupport();
     if (mixing && ss->everCaptured) {
       CUDACHECK(cudaStreamWaitEvent(ss->cudaStream, ss->serialEvent, 0));
@@ -219,7 +201,7 @@ static ncclResult_t checkGraphId(struct ncclStrongStreamGraph* g, unsigned long
 }
 
 ncclResult_t ncclStrongStreamRelease(struct ncclCudaGraph graph, struct ncclStrongStream* ss) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     bool mixing = ncclParamGraphMixingSupport();
     if (mixing && ss->serialEventNeedsRecord) {
       if (graph.graph == nullptr) {
@@ -243,7 +225,7 @@ ncclResult_t ncclStrongStreamRelease(struct ncclCudaGraph graph, struct ncclStro
 ncclResult_t ncclStrongStreamLaunchHost(
     struct ncclCudaGraph graph, struct ncclStrongStream* ss, cudaHostFn_t fn, void* arg
   ) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     if (graph.graph == nullptr) {
       CUDACHECK(cudaLaunchHostFunc(ss->cudaStream, fn, arg));
     } else {
@@ -267,7 +249,7 @@ ncclResult_t ncclStrongStreamLaunchKernel(
     struct ncclCudaGraph graph, struct ncclStrongStream* ss,
     void* fn, dim3 grid, dim3 block, void* args[], size_t sharedMemBytes
   ) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     if (graph.graph == nullptr) {
       CUDACHECK(cudaLaunchKernel(fn, grid, block, args, sharedMemBytes, ss->cudaStream));
     } else {
@@ -308,7 +290,7 @@ ncclResult_t ncclStrongStreamWaitStream(
     struct ncclCudaGraph graph, struct ncclStrongStream* a, struct ncclStrongStream* b,
     bool b_subsumes_a
   ) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     if (graph.graph == nullptr) {
       if (b->serialEventNeedsRecord) {
         b->serialEventNeedsRecord = false;
@@ -335,7 +317,7 @@ ncclResult_t ncclStrongStreamWaitStream(
     struct ncclCudaGraph graph, struct ncclStrongStream* a, cudaStream_t b,
     bool b_subsumes_a
   ) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     if (graph.graph == nullptr) {
       // It is ok to use a->serialEvent to record b since we'll be setting
       // a->serialEventNeedsRecord so the event won't be considered accurate
@@ -347,7 +329,7 @@ ncclResult_t ncclStrongStreamWaitStream(
       unsigned long long bGraphId;
       cudaGraphNode_t const* bNodes;
       size_t bCount = 0;
-      CUDACHECK(cudaStreamGetCaptureInfo_v2(b, &status, &bGraphId, nullptr, &bNodes, &bCount));
+      CUDACHECK(hipStreamGetCaptureInfo_v2(b, &status, &bGraphId, nullptr, &bNodes, &bCount));
       if (status != cudaStreamCaptureStatusActive || graph.graphId != bGraphId) {
         WARN("Stream is not being captured by the expected graph.");
         return ncclInvalidUsage;
@@ -369,7 +351,7 @@ ncclResult_t ncclStrongStreamWaitStream(
     struct ncclCudaGraph graph, cudaStream_t a, struct ncclStrongStream* b,
     bool b_subsumes_a
   ) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     if (graph.graph == nullptr) {
       if (b->serialEventNeedsRecord) {
         b->serialEventNeedsRecord = false;
@@ -379,7 +361,7 @@ ncclResult_t ncclStrongStreamWaitStream(
     } else {
       struct ncclStrongStreamGraph* bg = b->graphHead;
       NCCLCHECK(checkGraphId(bg, graph.graphId));
-      CUDACHECK(cudaStreamUpdateCaptureDependencies(a, bg->tipNodes, bg->tipCount,
+      CUDACHECK(hipStreamUpdateCaptureDependencies(a, bg->tipNodes, bg->tipCount,
         b_subsumes_a ? cudaStreamSetCaptureDependencies : cudaStreamAddCaptureDependencies
       ));
     }
@@ -391,7 +373,7 @@ ncclResult_t ncclStrongStreamWaitStream(
 }
 
 ncclResult_t ncclStrongStreamSynchronize(struct ncclStrongStream* ss) {
-  #if CUDART_VERSION >= 11030
+  #if ROCM_VERSION >= 60100
     CUDACHECK(cudaStreamWaitEvent(ss->cudaStream, ss->serialEvent, 0));
     ss->serialEventNeedsRecord = false;
   #endif