AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Adding opt-in hipGraph support for RCCL via RCCL_ENABLE_HIPGRAPH (#608)

Browse Source 
Adding opt-in hipGraph support via RCCL_ENABLE_HIPGRAPH
This commit is contained in:
gilbertlee-amd
2022-09-06 10:29:46 -06:00
committed by GitHub
parent 06bce9d0c9
commit 47b2fc3a30
6 changed files with 99 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/enqueue.cc
+29 -14
View File
@@ -14,7 +14,7 @@
#include "gdrwrap.h"
#include "bootstrap.h"
#include "channel.h"
#include "rccl_vars.h"
#include <cstring> // std::memcpy
// Only generate inline kernels for LL
@@ -293,7 +293,7 @@ ncclResult_t ncclLaunchKernel(ncclComm_t comm) {
if (comm->launchMode == ncclComm::GROUP) {
NCCLCHECK(ncclCpuBarrierOut(comm));
} else {
if (!comm->usingCudaGraph)
if (!rcclParamEnableHipGraph() || !comm->usingCudaGraph)
CUDACHECK(hipExtLaunchKernel(params->func, params->gridDim, params->blockDim, params->args, params->sharedMem, params->stream, NULL, comm->doneEvent, 0));
else
CUDACHECK(hipLaunchKernel(params->func, params->gridDim, params->blockDim, params->args, params->sharedMem, params->stream));
@@ -328,7 +328,10 @@ ncclResult_t ncclRecordEvents(ncclComm_t comm) {
// Enqueue event after NCCL kernel (only in non-graph mode)
// [RCCL] move event record into hipExtLaunchKernel
// if (!comm->usingCudaGraph) CUDACHECK(hipEventRecord(comm->doneEvent, params->stream));
if (rcclParamEnableHipGraph())
{
if (!comm->usingCudaGraph) CUDACHECK(hipEventRecord(comm->doneEvent, params->stream));
}
// Use internal NCCL stream for CGMD/GROUP launch if required or if the user stream is NULL
if (comm->launchMode == ncclComm::GROUP &&
(comm->groupCudaStream ||
@@ -1231,33 +1234,40 @@ void* graphHelperFunc(void *args) {
}
}
}
RCCL_PARAM(EnableHipGraph, "ENABLE_HIPGRAPH", 0);
// Check if we are in CUDA Graph capture mode
ncclResult_t ncclGetCudaGraph(ncclComm_t comm, hipGraph_t* graph) {
comm->usingCudaGraph = 0;
// Feature requires CUDA 11.3/R465 or above
#if CUDART_VERSION >= 11030
cudaStreamCaptureStatus captureStatus;
#if HIP_VERSION >= 50322000
hipStreamCaptureStatus captureStatus;
unsigned long long cudaGraphId;
ncclResult_t ret = ncclSuccess;
if (comm->driverVersion < 11030) {
if (comm->driverVersion < 50322000) {
// Runtime driver version older than compiler version
// Enhanced compat fallback
goto enh_compat_end;
}
// Get CUDA Graph handle
CUDACHECKGOTO(cudaStreamGetCaptureInfo_v2(comm->userStream, &captureStatus, &cudaGraphId, graph, NULL, NULL), ret, enh_compat_end);
if (captureStatus == cudaStreamCaptureStatusActive) {
CUDACHECKGOTO(hipStreamGetCaptureInfo_v2(comm->userStream, &captureStatus, &cudaGraphId, graph, NULL, NULL), ret, enh_compat_end);
if (captureStatus == hipStreamCaptureStatusActive) {
if (cudaGraphId != comm->lastCudaGraphId) {
INFO(NCCL_COLL, "stream is being captured by a new graph, id %llu", cudaGraphId);
// We are in a new graph, hence need to forget the last setup node so that
// the first setup node in the new graph will not have a dependency
comm->lastCudaGraphId = hipGraphId;
comm->lastCudaGraphId = cudaGraphId;
comm->lastSetupNode = NULL;
}
if (comm->launchMode == ncclComm::GROUP) comm->launchMode = ncclComm::GROUP_GRAPH;
comm->usingCudaGraph = 1;
if (!rcclParamEnableHipGraph())
{
WARN("RCCL_ENABLE_HIPGRAPH must be set to non-zero in order to support hipGraph usage");
return ncclInvalidUsage;
}
// Create helper thread that closes IPC handles during graph destruction
// Only create this thread when buffer registration is enabled
if ((!comm->graphHelperThread) && comm->graphRegister == 1 && comm->disableGraphHelper == 0) {
@@ -1276,9 +1286,9 @@ ncclResult_t ncclGetCudaGraph(ncclComm_t comm, hipGraph_t* graph) {
enh_compat_end: // Enhanced compat fallback
(void)ret;
CUDACHECK(cudaStreamIsCapturing(comm->userStream, &captureStatus));
if (captureStatus != cudaStreamCaptureStatusNone) {
WARN("The installed CUDA driver is older than the minimum version (R465) required for NCCL's CUDA Graphs support");
CUDACHECK(hipStreamIsCapturing(comm->userStream, &captureStatus));
if (captureStatus != hipStreamCaptureStatusNone) {
WARN("The installed ROCm driver is older than the minimum version (50322000) required for RCCL's HIP Graphs support");
return ncclInvalidUsage;
}
// If we are not in capture mode, we can ignore the driver being lower
@@ -1288,7 +1298,12 @@ enh_compat_end: // Enhanced compat fallback
// Create host setup node in CUDA Graph
ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, hipGraph_t graph) {
#if CUDART_VERSION >= 11030
#if HIP_VERSION >= 50322000
if (!rcclParamEnableHipGraph())
{
WARN("RCCL_ENABLE_HIPGRAPH must be set to non-zero to enable HIP graph feature");
return ncclInternalError;
}
struct ncclQueueInfo* eqInfo = comm->enqueueInfo;
// Create a CUDA object to wrap around the argument space
// which CUDA graph would manage lifetime of
@@ -1309,7 +1324,7 @@ ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, hipGraph_t graph) {
comm->lastSetupNode = setupNode;
return ncclSuccess;
#else
WARN("NCCL does not support this CUDA version for CUDA graph feature");
WARN("RCCL does not support this ROCm version for HIP graph feature");
return ncclInternalError;
#endif
}