From 8b8179a689a68d95dfccad0111e1c8096b5d2664 Mon Sep 17 00:00:00 2001 From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com> Date: Sun, 21 Jan 2024 19:00:50 -0800 Subject: [PATCH] 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 [ROCm/rccl commit: 7e25d5bc5562e13dea9688c8c4796b45279663e6] --- projects/rccl/CMakeLists.txt | 1 - projects/rccl/src/enqueue.cc | 13 ++--- projects/rccl/src/include/strongstream.h | 10 ++-- projects/rccl/src/misc/strongstream.cc | 68 +++++++++--------------- 4 files changed, 37 insertions(+), 55 deletions(-) diff --git a/projects/rccl/CMakeLists.txt b/projects/rccl/CMakeLists.txt index de18c52b4b..abf4da4133 100644 --- a/projects/rccl/CMakeLists.txt +++ b/projects/rccl/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 diff --git a/projects/rccl/src/enqueue.cc b/projects/rccl/src/enqueue.cc index 04fc6bb220..a7dae3534c 100644 --- a/projects/rccl/src/enqueue.cc +++ b/projects/rccl/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; diff --git a/projects/rccl/src/include/strongstream.h b/projects/rccl/src/include/strongstream.h index 0984dfe574..c063f37980 100644 --- a/projects/rccl/src/include/strongstream.h +++ b/projects/rccl/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; diff --git a/projects/rccl/src/misc/strongstream.cc b/projects/rccl/src/misc/strongstream.cc index faeec4bca3..7913fcb0ea 100644 --- a/projects/rccl/src/misc/strongstream.cc +++ b/projects/rccl/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