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

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