From 6ce074d92d46fd7057411fd5c7535243709f5101 Mon Sep 17 00:00:00 2001 From: Ziyue Yang Date: Thu, 16 Nov 2023 08:43:28 +0800 Subject: [PATCH] Fix MSCCL work FIFO allocation with HIP graph enabled (#967) [ROCm/rccl commit: 7fc891bc8d0a55d160e7051922299d8c5b8b2127] --- projects/rccl/src/collectives/msccl.cc | 2 +- projects/rccl/src/include/msccl/msccl_setup.h | 4 + .../rccl/src/include/msccl/msccl_struct.h | 21 +++-- .../rccl/src/misc/msccl/msccl_lifecycle.cc | 15 ++-- projects/rccl/src/misc/msccl/msccl_setup.cc | 78 ++++++++++++++----- 5 files changed, 80 insertions(+), 40 deletions(-) diff --git a/projects/rccl/src/collectives/msccl.cc b/projects/rccl/src/collectives/msccl.cc index 0d8197b731..1fae57c27c 100644 --- a/projects/rccl/src/collectives/msccl.cc +++ b/projects/rccl/src/collectives/msccl.cc @@ -27,7 +27,7 @@ ncclResult_t mscclLoadAlgo(const char *mscclAlgoFilePath, mscclAlgoHandle_t *msc status.hostAlgos[*mscclAlgoHandle] = hostAlgo; struct mscclAlgo* devAlgo; - NCCLCHECK(ncclCudaCalloc(&devAlgo, 1)); + NCCLCHECK(ncclCudaMalloc(&devAlgo, 1)); CUDACHECK(hipMemcpy(devAlgo, hostAlgo, sizeof(struct mscclAlgo), hipMemcpyHostToDevice)); status.devAlgos[*mscclAlgoHandle] = devAlgo; diff --git a/projects/rccl/src/include/msccl/msccl_setup.h b/projects/rccl/src/include/msccl/msccl_setup.h index 893a2f5b7f..7a326236ce 100644 --- a/projects/rccl/src/include/msccl/msccl_setup.h +++ b/projects/rccl/src/include/msccl/msccl_setup.h @@ -27,4 +27,8 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count ncclDataType_t dataType, ncclRedOp_t op, struct mscclAlgo* hostAlgo, struct mscclAlgo* devAlgo, ncclComm_t comm, hipStream_t stream); +ncclResult_t mscclInitWorkFifoStatus(mscclWorkFifoStatus* workFifoStatus); + +ncclResult_t mscclDestroyWorkFifoStatus(mscclWorkFifoStatus* workFifoStatus); + #endif diff --git a/projects/rccl/src/include/msccl/msccl_struct.h b/projects/rccl/src/include/msccl/msccl_struct.h index bd288de4ce..4ce0b3e6df 100644 --- a/projects/rccl/src/include/msccl/msccl_struct.h +++ b/projects/rccl/src/include/msccl/msccl_struct.h @@ -36,8 +36,6 @@ #define MSCCL_LOCAL_COPY 6 #define MSCCL_REDUCE 7 -#define MSCCL_WORK_FIFO_DEPTH (64 << 10) - struct mscclTransmission { int16_t dependencePointer; // index to the first dependence int16_t numDependencies; // dependencePointer+numDependencies indicate the last dependence @@ -189,6 +187,17 @@ struct mscclThreadLocalStatus { hipGraph_t graph; }; +struct mscclWorkFifoStatus { + uint64_t workFifoDepth; + struct mscclWork* workFifo; + uint32_t* workFifoDone; + uint32_t workFifoSent; + uint32_t workFifoSentPerThreadBlock[MSCCL_MAX_NUM_THREAD_BLOCKS]; + uint32_t workFifoAckdMin; +}; + +typedef std::map mscclSavedGraphWorkFifoStatus; + struct mscclStatus { std::vector freeAlgoHandles; std::map hostAlgos; @@ -214,12 +223,8 @@ struct mscclStatus { bool graphEnabled; bool graphFirstKernel; bool needsProxy; - uint64_t workFifoDepth; - struct mscclWork* workFifo; - uint32_t* workFifoDone; - uint32_t workFifoSent; - uint32_t workFifoSentPerThreadBlock[MSCCL_MAX_NUM_THREAD_BLOCKS]; - uint32_t workFifoAckdMin; + mscclWorkFifoStatus defaultWorkFifoStatus; + mscclSavedGraphWorkFifoStatus graphWorkFifoStatus; }; #pragma pack(push) diff --git a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc b/projects/rccl/src/misc/msccl/msccl_lifecycle.cc index e402fb019d..01d605c945 100644 --- a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc +++ b/projects/rccl/src/misc/msccl/msccl_lifecycle.cc @@ -190,14 +190,7 @@ ncclResult_t mscclInit(ncclComm_t comm) { NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS)); status.lastStream = nullptr; status.needsProxy = false; - status.workFifoDepth = MSCCL_WORK_FIFO_DEPTH; - NCCLCHECK(ncclCudaCalloc(&status.workFifo, status.workFifoDepth, nullptr, true)); - NCCLCHECK(ncclCudaHostCalloc(&status.workFifoDone, MSCCL_MAX_NUM_THREAD_BLOCKS)); - status.workFifoSent = 0; - for (int i = 0; i < MSCCL_MAX_NUM_THREAD_BLOCKS; i++) { - status.workFifoSentPerThreadBlock[i] = 0; - } - status.workFifoAckdMin = 0; + NCCLCHECK(mscclInitWorkFifoStatus(&(status.defaultWorkFifoStatus))); mscclSchedulerTriedLoadAlgo = false; NCCLCHECK(mscclSchedulerInit()); @@ -521,8 +514,10 @@ ncclResult_t mscclTeardown() { } else { NCCLCHECK(mscclInternalSchedulerTeardown()); } - NCCLCHECK(ncclCudaFree(status.workFifo)); - NCCLCHECK(ncclCudaHostFree(status.workFifoDone)); + NCCLCHECK(mscclDestroyWorkFifoStatus(&(status.defaultWorkFifoStatus))); + for (auto &p : status.graphWorkFifoStatus) { + NCCLCHECK(mscclDestroyWorkFifoStatus(&(p.second))); + } mscclInitialized.store(false, std::memory_order_release); } diff --git a/projects/rccl/src/misc/msccl/msccl_setup.cc b/projects/rccl/src/misc/msccl/msccl_setup.cc index 4c9a12565a..199c911a3d 100644 --- a/projects/rccl/src/misc/msccl/msccl_setup.cc +++ b/projects/rccl/src/misc/msccl/msccl_setup.cc @@ -20,7 +20,10 @@ RCCL_PARAM(MscclEnableDoneEvent, "MSCCL_ENABLE_DONE_EVENT", 1); #endif +RCCL_PARAM(MscclWorkFifoDepth, "MSCCL_WORK_FIFO_DEPTH", 64<<10); + ncclResult_t mscclGetCaptureStatus(hipStream_t stream) { + mscclStatus& status = mscclGetStatus(); mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus(); mscclSavedProxyArgs& savedProxyArgs = mscclGetSavedProxyArgs(); cudaStreamCaptureStatus captureStatus; @@ -30,6 +33,7 @@ ncclResult_t mscclGetCaptureStatus(hipStream_t stream) { if (savedProxyArgs.count(captureId) == 0) { threadLocalStatus.captureStatus = mscclNewCapture; savedProxyArgs[captureId] = std::vector(); + NCCLCHECK(mscclInitWorkFifoStatus(&(status.graphWorkFifoStatus[captureId]))); } else { INFO(NCCL_NET,"mscclGetCaptureStatus: captureId %llu is same with the previous one\n", captureId); threadLocalStatus.captureStatus = mscclExistingCapture; @@ -69,7 +73,7 @@ ncclResult_t mscclSetupScratch(struct mscclAlgo* hostAlgo, hipStream_t stream) { size_t sizeNeeded = (status.nBytes * (size_t)(hostAlgo->nScratchChunks)) / (size_t)(hostAlgo->nChunksPerLoop); if (sizeNeeded > status.scratchBufferSize){ NCCLCHECK(ncclCudaFree(status.scratchBuffer)); - NCCLCHECK(ncclCudaCalloc((char**)&status.scratchBuffer, sizeNeeded)); + NCCLCHECK(ncclCudaMalloc((char**)&status.scratchBuffer, sizeNeeded)); status.scratchBufferSize = sizeNeeded; } return ncclSuccess; @@ -333,12 +337,11 @@ static inline uint32_t rollingMin32(uint32_t a, uint32_t b) { return (b-a <= PositiveMax) ? a : b; } -static void mscclWaitWorkFifoAvailable(uint32_t desiredSent) { - mscclStatus& status = mscclGetStatus(); - if (__builtin_expect(rollingLess32(status.workFifoAckdMin + status.workFifoDepth, desiredSent), false)) { +static void mscclWaitWorkFifoAvailable(uint32_t desiredSent, mscclWorkFifoStatus* status) { + if (__builtin_expect(rollingLess32(status->workFifoAckdMin + status->workFifoDepth, desiredSent), false)) { while (1) { // We have to poll for notifications from device. - uint32_t* doneLive = status.workFifoDone; + uint32_t* doneLive = status->workFifoDone; uint32_t ackd[MSCCL_MAX_NUM_THREAD_BLOCKS]; for (int c=0; c < MSCCL_MAX_NUM_THREAD_BLOCKS; c++) { ackd[c] = __atomic_load_n(&doneLive[c], __ATOMIC_RELAXED); @@ -346,10 +349,10 @@ static void mscclWaitWorkFifoAvailable(uint32_t desiredSent) { // Compiler-only fence to prevent fusion of loops to encourage dense loads. __atomic_signal_fence(__ATOMIC_SEQ_CST); - uint32_t ackdAll = status.workFifoSent; + uint32_t ackdAll = status->workFifoSent; for (int c=0; c < MSCCL_MAX_NUM_THREAD_BLOCKS; c++) { // ackdAll is min over all non-quiesced channels - if (ackd[c] != status.workFifoSentPerThreadBlock[c]) + if (ackd[c] != status->workFifoSentPerThreadBlock[c]) ackdAll = rollingMin32(ackdAll, ackd[c]); } @@ -359,15 +362,15 @@ static void mscclWaitWorkFifoAvailable(uint32_t desiredSent) { for (int c=0; c < MSCCL_MAX_NUM_THREAD_BLOCKS; c++) { // Advance counter on quiesced channels so they don't lag behind // too far where they could get lost in 32-bit wraparound. - if (ackd[c] == status.workFifoSentPerThreadBlock[c]) { - status.workFifoSentPerThreadBlock[c] = ackdAll; + if (ackd[c] == status->workFifoSentPerThreadBlock[c]) { + status->workFifoSentPerThreadBlock[c] = ackdAll; __atomic_store_n(&doneLive[c], ackdAll, __ATOMIC_RELAXED); } } - status.workFifoAckdMin = ackdAll; + status->workFifoAckdMin = ackdAll; // See if that was enough. - if (!rollingLess32(status.workFifoAckdMin + status.workFifoDepth, desiredSent)) break; + if (!rollingLess32(status->workFifoAckdMin + status->workFifoDepth, desiredSent)) break; sched_yield(); } } @@ -379,6 +382,7 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count ncclDataType_t dataType, ncclRedOp_t op, struct mscclAlgo* hostAlgo, struct mscclAlgo* devAlgo, ncclComm_t comm, hipStream_t stream) { mscclStatus& status = mscclGetStatus(); + mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus(); bool enableDoneEvent = #ifndef HIP_EVENT_DISABLE_FENCE @@ -421,9 +425,22 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count work.redOpArgIsPtr = opFull.scalarArgIsPtr; work.fnIndex = fnIndex; INFO(NCCL_COLL, "MSCCL: typeMask %x fnIndex %d Setup Kernel finished", hostAlgo->typeMask, fnIndex); - - uint32_t workFifoIdxMask = status.workFifoDepth - 1; - uint32_t workFifoSent = status.workFifoSent; + + mscclWorkFifoStatus* workFifoStatus = nullptr; + if (threadLocalStatus.captureStatus == mscclNoCapture) { + workFifoStatus = &(status.defaultWorkFifoStatus); + } else { + workFifoStatus = &(status.graphWorkFifoStatus[threadLocalStatus.captureId]); + } + + uint32_t workFifoIdxMask = workFifoStatus->workFifoDepth - 1; + uint32_t workFifoSent = workFifoStatus->workFifoSent; + + if (threadLocalStatus.captureStatus != mscclNoCapture && workFifoSent + numBlocks > workFifoStatus->workFifoDepth) { + WARN("MSCCL: number of captured works (%u) > max limit (%lu)", workFifoSent + numBlocks, workFifoStatus->workFifoDepth); + return ncclInternalError; + } + // First work for a channel has to be at workHeap+blockIdx.x which means // we cannot tolerate fifo wraparound. So round up to the wrap boundary // if not doing so would incur crossing it. @@ -432,18 +449,19 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count // Need to update workFifoSent so waitWorkFifoAvailable() knows we've // skipped those elements. Consider if all the channels report quiesced, // this way the skipped slots will be considered consumed as well. - status.workFifoSent = workFifoSent; + workFifoStatus->workFifoSent = workFifoSent; } - mscclWaitWorkFifoAvailable(workFifoSent + numBlocks); + mscclWaitWorkFifoAvailable(workFifoSent + numBlocks, workFifoStatus); for (int i = 0; i < numBlocks; i++) { work.workFifoDoneAck = workFifoSent + i; - work.workFifoDone = status.workFifoDone + i; - status.workFifoSentPerThreadBlock[i] = workFifoSent + i; - status.workFifo[(workFifoSent + i) & workFifoIdxMask] = work; + work.workFifoDone = workFifoStatus->workFifoDone + i; + workFifoStatus->workFifoSentPerThreadBlock[i] = workFifoSent + i; + workFifoStatus->workFifo[(workFifoSent + i) & workFifoIdxMask] = work; } - status.workFifoSent = workFifoSent + numBlocks; - struct mscclWork *workPtr = status.workFifo + (workFifoSent & workFifoIdxMask); + struct mscclWork *workPtr = workFifoStatus->workFifo + (workFifoSent & workFifoIdxMask); + workFifoStatus->workFifoSent = workFifoSent + numBlocks; + void *args[3] = {&comm->devComm, &devAlgo, &workPtr}; void *func = mscclKernelEntries[fnIndex]; if (enableDoneEvent) { @@ -472,3 +490,21 @@ size_t mscclKernMaxLocalSize() { error: return (res != ncclSuccess) ? 0 : max; } + +ncclResult_t mscclInitWorkFifoStatus(mscclWorkFifoStatus* workFifoStatus) { + workFifoStatus->workFifoDepth = rcclParamMscclWorkFifoDepth(); + NCCLCHECK(ncclCudaMalloc(&(workFifoStatus->workFifo), workFifoStatus->workFifoDepth, true)); + NCCLCHECK(ncclCudaHostCalloc(&(workFifoStatus->workFifoDone), MSCCL_MAX_NUM_THREAD_BLOCKS)); + workFifoStatus->workFifoSent = 0; + for (int i = 0; i < MSCCL_MAX_NUM_THREAD_BLOCKS; i++) { + workFifoStatus->workFifoSentPerThreadBlock[i] = 0; + } + workFifoStatus->workFifoAckdMin = 0; + return ncclSuccess; +} + +ncclResult_t mscclDestroyWorkFifoStatus(mscclWorkFifoStatus* workFifoStatus) { + NCCLCHECK(ncclCudaFree(workFifoStatus->workFifo)); + NCCLCHECK(ncclCudaHostFree(workFifoStatus->workFifoDone)); + return ncclSuccess; +}