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

Merge remote-tracking branch 'nccl/master' into develop

Browse Source 
[ROCm/rccl commit: 36343be84f]
This commit is contained in:
BertanDogancay
2025-01-23 11:48:18 -06:00
committed by Corey Derochie
parent e3a29f5eab 5ca1b6c160
commit 1b000665df
122 changed files with 9009 additions and 4702 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rccl/src/group.cc
+195 -66
View File
@@ -13,6 +13,7 @@
#include "channel.h"
#include "api_trace.h"
#include <assert.h>
#include "bootstrap.h"
#include "msccl/msccl_lifecycle.h"
@@ -36,6 +37,7 @@ ncclResult_t ncclAsyncLaunch(
) {
ncclResult_t ret = ncclSuccess;
job->destroyFlag = comm->destroyFlag;
if (ncclGroupDepth == 0) {
ret = func(job);
if (ret != ncclSuccess && undo) undo(job);
@@ -45,11 +47,15 @@ ncclResult_t ncclAsyncLaunch(
job->undo = undo;
job->destructor = destructor;
job->abortFlag = comm->abortFlag;
job->abortFlagDev = comm->abortFlagDev;
job->childAbortFlag = comm->childAbortFlag;
job->childAbortFlagDev = comm->childAbortFlagDev;
job->state = ncclGroupJobRunning;
job->comm = comm;
/* check if there are blocking and nonblocking comms at the same time in group. */
if (ncclGroupBlocking == -1) {
if (comm->destroyFlag) {
ncclGroupBlocking = 1;
} else if (ncclGroupBlocking == -1) {
/* first met communicator */
ncclGroupBlocking = comm->config.blocking;
} else if (ncclGroupBlocking != comm->config.blocking) {
@@ -111,11 +117,23 @@ exit:
return ret;
}
NCCL_API(ncclResult_t, ncclGroupSimulateEnd, ncclSimInfo_t* simInfo);
ncclResult_t ncclGroupSimulateEnd(ncclSimInfo_t* simInfo) {
ncclResult_t ret = ncclSuccess;
NVTX3_FUNC_RANGE_IN(nccl_domain);
NCCLCHECKGOTO(ncclGroupEndInternal(simInfo), ret, exit);
TRACE_CALL("ncclGroupSimulateEnd()");
exit:
return ret;
}
struct ncclPreconnectJob {
struct ncclAsyncJob base;
struct ncclComm* comm;
bool* algoNeedConnect;
};
ncclResult_t ncclPreconnectFunc(struct ncclAsyncJob* job_) {
ncclResult_t ncclP2PPreconnectFunc(struct ncclAsyncJob* job_) {
struct ncclPreconnectJob* job = (struct ncclPreconnectJob*)job_;
struct ncclComm* comm = job->comm;
CUDACHECK(cudaSetDevice(comm->cudaDev));
@@ -125,6 +143,57 @@ ncclResult_t ncclPreconnectFunc(struct ncclAsyncJob* job_) {
return ncclSuccess;
}
ncclResult_t ncclCollPreconnectFunc(struct ncclAsyncJob* job_) {
struct ncclPreconnectJob* job = (struct ncclPreconnectJob*)job_;
struct ncclComm* comm = job->comm;
ncclResult_t ret = ncclSuccess;
CUDACHECK(cudaSetDevice(comm->cudaDev));
if (CPU_COUNT(&comm->cpuAffinity)) sched_setaffinity(0, sizeof(cpu_set_t), &comm->cpuAffinity);
for (int i = 0; i < NCCL_NUM_ALGORITHMS; ++i) {
if (job->algoNeedConnect[i]) {
switch (i) {
case NCCL_ALGO_RING: {
NCCLCHECKGOTO(ncclTransportRingConnect(comm), ret, fail);
break;
}
case NCCL_ALGO_TREE: {
NCCLCHECKGOTO(ncclTransportTreeConnect(comm), ret, fail);
break;
}
case NCCL_ALGO_NVLS: {
/* If we are using NVLS_TREE algo, we must mark NVLS algo to set up
* NVLS intra-node buffer */
NCCLCHECKGOTO(ncclNvlsBufferSetup(comm), ret, fail);
break;
}
case NCCL_ALGO_NVLS_TREE: {
NCCLCHECKGOTO(ncclNvlsTreeConnect(comm), ret, fail);
break;
}
case NCCL_ALGO_COLLNET_CHAIN: {
NCCLCHECKGOTO(ncclCollNetChainBufferSetup(comm), ret, fail);
break;
}
case NCCL_ALGO_COLLNET_DIRECT: {
NCCLCHECKGOTO(ncclCollNetDirectBufferSetup(comm), ret, fail);
break;
}
default: {
ret = ncclInternalError;
goto fail;
}
}
}
}
exit:
free(job->algoNeedConnect);
return ret;
fail:
goto exit;
}
static ncclResult_t doLaunches(struct ncclComm* head) {
ncclResult_t result = ncclSuccess;
struct ncclComm* cliqueComm0 = head->intraComm0;
@@ -138,7 +207,7 @@ static ncclResult_t doLaunches(struct ncclComm* head) {
struct ncclComm* comm = cliqueHead;
bool capturingYes = false, capturingNo = false;
do {
(ncclCudaGraphValid(comm->tasks.capturingGraph) ? capturingYes : capturingNo) = true;
(ncclCudaGraphValid(comm->planner.capturingGraph) ? capturingYes : capturingNo) = true;
CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
NCCLCHECKGOTO(ncclLaunchPrepare(comm), result, failure);
if (useBarrier) ncclCommIntraBarrierIn(comm, 1);
@@ -164,19 +233,19 @@ static ncclResult_t doLaunches(struct ncclComm* head) {
// Barrier reduction result tells us if this was the final round.
moreRounds = 0 != ncclCommIntraBarrierOut(comm);
} else {
moreRounds |= comm->unlaunchedPlansHead != nullptr;
moreRounds |= comm->planner.unlaunchedPlansHead != nullptr;
}
if (moreRounds) {
// Pop next unlaunched kernel
struct ncclKernelPlan* plan = comm->unlaunchedPlansHead;
struct ncclKernelPlan* plan = comm->planner.unlaunchedPlansHead;
if (plan != nullptr) {
comm->unlaunchedPlansHead = plan->next;
comm->planner.unlaunchedPlansHead = plan->next;
CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
NCCLCHECKGOTO(ncclLaunchKernelBefore_NoUncapturedCuda(comm, plan), result, failure);
NCCLCHECKGOTO(ncclLaunchKernel(comm, plan), result, failure);
}
// Barrier reduction input indicates if we require further rounds.
if (useBarrier) ncclCommIntraBarrierIn(comm, comm->unlaunchedPlansHead != nullptr ? 1 : 0);
if (useBarrier) ncclCommIntraBarrierIn(comm, comm->planner.unlaunchedPlansHead != nullptr ? 1 : 0);
if (plan != nullptr) {
NCCLCHECKGOTO(ncclLaunchKernelAfter_NoCuda(comm, plan), result, failure);
}
@@ -224,39 +293,31 @@ static void groupCleanup(struct ncclComm** groupCommHeadPtr, struct ncclComm** g
// is needed.
comm->preconnectNext = reinterpret_cast<struct ncclComm*>(0x1);
for (int i = 0; i < comm->nRanks; i++) {
comm->tasks.peers[i].sendSeen = false;
comm->tasks.peers[i].recvSeen = false;
for (int j = 0; j < MAXCHANNELS/64; j++) {
comm->connectSend[i].masks[j] = 0UL;
comm->connectRecv[i].masks[j] = 0UL;
}
}
comm->unlaunchedPlansHead = nullptr;
// Reclaim abandoned kernel plan memory. Note ncclWork structs were already
// reclaimed by a `ncclMemoryStackPop(&comm->memScoped)` during `ncclGroupCommLeave()`.
while (!ncclIntruQueueEmpty(&comm->planQueue)) {
struct ncclKernelPlan* plan = ncclIntruQueueDequeue(&comm->planQueue);
while (!ncclIntruQueueEmpty(&comm->planner.planQueue)) {
struct ncclKernelPlan* plan = ncclIntruQueueDequeue(&comm->planner.planQueue);
// Persistent plans will be reclaimed via the callbackQueue when the
// graph drops its UserObject reference.
if (!plan->persistent) {
for (int c = 0; c < MAXCHANNELS; c++) {
while (!ncclIntruQueueEmpty(&plan->channels[c].proxyOpQueue)) {
struct ncclProxyOp* pxop = ncclIntruQueueDequeue(&plan->channels[c].proxyOpQueue);
ncclMemoryPoolFree(&comm->memPool_ncclProxyOp, pxop);
}
while (!ncclIntruQueueEmpty(&plan->proxyOpQueue)) {
struct ncclProxyOp* pxop = ncclIntruQueueDequeue(&plan->proxyOpQueue);
ncclMemoryPoolFree(&comm->memPool_ncclProxyOp, pxop);
}
ncclMemoryPoolFree(&comm->memPool_ncclKernelPlan, plan);
}
}
// Reset comm->tasks to empty.
comm->tasks.nTasksColl = 0;
comm->tasks.nTasksP2p = 0;
comm->tasks.workBytesTotal = 0;
comm->tasks.streams = nullptr;
ncclIntruQueueConstruct(&comm->tasks.collQueue);
for (int i = 0; i < comm->nRanks; i++) {
ncclIntruQueueConstruct(&comm->tasks.peers[i].sendQueue);
ncclIntruQueueConstruct(&comm->tasks.peers[i].recvQueue);
{ // Reset comm->planner to empty.
ncclKernelPlanner::Peer* tmp = comm->planner.peers;
memset(&comm->planner, 0, sizeof(comm->planner));
comm->planner.peers = tmp;
memset(comm->planner.peers, 0, comm->nRanks*sizeof(comm->planner.peers[0]));
}
if (!comm->config.blocking)
@@ -276,37 +337,10 @@ static void groupCleanup(struct ncclComm** groupCommHeadPtr, struct ncclComm** g
return;
}
static ncclResult_t groupLaunch(struct ncclAsyncJob *job_) {
int savedDev;
static ncclResult_t asyncJobLaunch(struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> *asyncJobsMain, volatile bool *groupAbortFlag) {
ncclResult_t ret = ncclSuccess;
bool jobsDone = false;
bool errorJobAbortFlag = false;
struct ncclGroupJob *gjob = (struct ncclGroupJob*) job_;
struct ncclComm *groupCommHeadMain = *gjob->groupCommHeadPtr;
struct ncclComm *groupCommPreconnectHeadMain = *gjob->groupCommPreconnectHeadPtr;
struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> *asyncJobsMain = gjob->asyncJobsPtr;
volatile bool *groupAbortFlag = gjob->abortFlagPtr;
CUDACHECKGOTO(cudaGetDevice(&savedDev), ret, fail);
if (groupCommPreconnectHeadMain != nullptr) {
struct ncclComm* comm = groupCommPreconnectHeadMain;
do {
struct ncclPreconnectJob* job;
NCCLCHECKGOTO(ncclCalloc(&job, 1), ret, fail);
job->base.func = ncclPreconnectFunc;
job->base.undo = nullptr;
job->base.destructor = free;
job->base.state = ncclGroupJobRunning;
job->base.abortFlag = comm->abortFlag;
job->comm = comm;
ncclIntruQueueEnqueue(asyncJobsMain, &job->base);
struct ncclComm* next = comm->preconnectNext;
comm->preconnectNext = reinterpret_cast<struct ncclComm*>(0x1);
comm = next;
} while (comm != nullptr);
}
if (!ncclIntruQueueEmpty(asyncJobsMain)) {
struct ncclAsyncJob* job = ncclIntruQueueHead(asyncJobsMain);
@@ -337,9 +371,13 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_) {
assert(state == ncclGroupJobJoined);
}
if (__atomic_load_n(groupAbortFlag, __ATOMIC_RELAXED) || errorJobAbortFlag == true) {
__atomic_store_n(job->abortFlag, 1, __ATOMIC_RELAXED);
if (job->childAbortFlag) __atomic_store_n(job->childAbortFlag, 1, __ATOMIC_RELAXED);
if (!job->destroyFlag && (__atomic_load_n(groupAbortFlag, __ATOMIC_ACQUIRE) || errorJobAbortFlag == true)) {
__atomic_store_n(job->abortFlag, 1, __ATOMIC_RELEASE);
__atomic_store_n(job->abortFlagDev, 1, __ATOMIC_RELEASE);
if (job->childAbortFlag) {
__atomic_store_n(job->childAbortFlag, 1, __ATOMIC_RELEASE);
__atomic_store_n(job->childAbortFlagDev, 1, __ATOMIC_RELEASE);
}
}
job = job->next;
@@ -351,17 +389,86 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_) {
if (ret != ncclSuccess) goto fail;
}
if (groupCommHeadMain != nullptr) {
NCCLCHECKGOTO(doLaunches(groupCommHeadMain), ret, fail);
}
while (!ncclIntruQueueEmpty(asyncJobsMain)) {
struct ncclAsyncJob* job = ncclIntruQueueDequeue(asyncJobsMain);
if (job->comm && !job->comm->config.blocking)
if (!job->destroyFlag && job->comm && !job->comm->config.blocking)
(void) ncclCommSetAsyncError(job->comm, ret);
if (job->destructor) job->destructor((void*)job);
}
exit:
return ret;
fail:
goto exit;
}
static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInfo = NULL) {
int savedDev;
ncclResult_t ret = ncclSuccess;
struct ncclGroupJob *gjob = (struct ncclGroupJob*) job_;
struct ncclComm *groupCommHeadMain = *gjob->groupCommHeadPtr;
struct ncclComm *groupCommPreconnectHeadMain = *gjob->groupCommPreconnectHeadPtr;
struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> *asyncJobsMain = gjob->asyncJobsPtr;
bool *groupAbortFlag = gjob->abortFlagPtr;
CUDACHECKGOTO(cudaGetDevice(&savedDev), ret, fail);
if (!simInfo && groupCommPreconnectHeadMain != nullptr) {
struct ncclComm* comm = groupCommPreconnectHeadMain;
do {
struct ncclPreconnectJob* job;
NCCLCHECKGOTO(ncclCalloc(&job, 1), ret, fail);
job->base.func = ncclP2PPreconnectFunc;
job->base.undo = nullptr;
job->base.destructor = free;
job->base.state = ncclGroupJobRunning;
job->base.abortFlag = comm->abortFlag;
job->base.abortFlagDev = comm->abortFlagDev;
job->comm = comm;
ncclIntruQueueEnqueue(asyncJobsMain, &job->base);
struct ncclComm* next = comm->preconnectNext;
comm->preconnectNext = reinterpret_cast<struct ncclComm*>(0x1);
comm = next;
} while (comm != nullptr);
}
NCCLCHECKGOTO(asyncJobLaunch(asyncJobsMain, groupAbortFlag), ret, fail);
/* Connect channels at runtime if cumem is supported */
if (groupCommHeadMain != nullptr) {
struct ncclComm* comm = groupCommHeadMain;
do {
bool needConnect = false;
bool algoNeedConnect[NCCL_NUM_ALGORITHMS];
memset(algoNeedConnect, 0, sizeof(bool) * NCCL_NUM_ALGORITHMS);
NCCLCHECKGOTO(ncclPrepareTasks(comm, algoNeedConnect, &needConnect, simInfo), ret, fail);
if (comm->cuMemSupport && needConnect) {
struct ncclPreconnectJob* job;
NCCLCHECKGOTO(ncclCalloc(&job, 1), ret, fail);
job->base.func = ncclCollPreconnectFunc;
job->base.undo = nullptr;
job->base.destructor = free;
job->base.state = ncclGroupJobRunning;
job->base.abortFlag = comm->abortFlag;
job->comm = comm;
NCCLCHECKGOTO(ncclCalloc(&job->algoNeedConnect, NCCL_NUM_ALGORITHMS), ret, fail);
memcpy(job->algoNeedConnect, algoNeedConnect, sizeof(bool) * NCCL_NUM_ALGORITHMS);
ncclIntruQueueEnqueue(asyncJobsMain, &job->base);
}
comm = comm->groupNext;
} while (comm);
NCCLCHECKGOTO(asyncJobLaunch(asyncJobsMain, groupAbortFlag), ret, fail);
}
if ((!simInfo) && (groupCommHeadMain != nullptr)) {
NCCLCHECKGOTO(doLaunches(groupCommHeadMain), ret, fail);
}
while (groupCommHeadMain != nullptr) {
struct ncclComm* comm = groupCommHeadMain;
struct ncclComm* next = comm->groupNext;
@@ -381,8 +488,17 @@ fail:
goto exit;
}
ncclResult_t ncclGroupEndInternal() {
static ncclResult_t groupLaunchNonBlocking(struct ncclAsyncJob *job_) {
return groupLaunch(job_ /* estimatedTime = NULL */);
}
ncclResult_t ncclGroupEndInternal(ncclSimInfo_t* simInfo) {
ncclResult_t ret = ncclSuccess;
ncclSimInfo_t internalSimInfo = NCCL_SIM_INFO_INITIALIZER;
ncclSimInfo_t* internalSimInfoPtr = NULL;
size_t realSize = 0;
internalSimInfo.magic = 0;
if (ncclGroupDepth == 0) {
WARN("ncclGroupEnd: not in a group call.");
@@ -398,6 +514,18 @@ ncclResult_t ncclGroupEndInternal() {
if ((ret = ncclGroupError) != ncclSuccess) goto fail;
if (simInfo) {
memcpy((void*)&realSize, (void*)&simInfo->size, sizeof(size_t));
realSize = realSize > sizeof(ncclSimInfo_t) ? sizeof(ncclSimInfo_t) : realSize;
memcpy((void*)&internalSimInfo, (void*)simInfo, realSize);
if (internalSimInfo.magic != 0x74685283) {
WARN("ncclSimInfo_t argument not initialized via NCCL_SIM_INFO_INITIALIZER");
ret = ncclInvalidArgument;
goto fail;
}
internalSimInfoPtr = &internalSimInfo;
}
if (ncclGroupCommHead != nullptr || !ncclIntruQueueEmpty(&ncclAsyncJobs) || ncclGroupCommPreconnectHead != nullptr) {
ncclGroupJobMain.groupCommHeadPtr = &ncclGroupCommHead;
ncclGroupJobMain.groupCommPreconnectHeadPtr = &ncclGroupCommPreconnectHead;
@@ -430,12 +558,13 @@ ncclResult_t ncclGroupEndInternal() {
} while (comm);
}
ncclGroupJobMainPtr->base.func = groupLaunch;
ncclGroupJobMainPtr->base.func = groupLaunchNonBlocking;
SYSCHECKGOTO(pthread_create(&ncclGroupJobMainPtr->base.thread, NULL, ncclAsyncJobMain, (void*)&ncclGroupJobMainPtr->base), ret, fail);
ret = ncclInProgress;
} else {
/* blocking group */
NCCLCHECKGOTO(groupLaunch(&ncclGroupJobMainPtr->base), ret, fail);
NCCLCHECKGOTO(groupLaunch(&ncclGroupJobMainPtr->base, internalSimInfoPtr), ret, fail);
if (simInfo) memcpy((void*)simInfo, (void*)internalSimInfoPtr, realSize);
groupResetJobState(ncclGroupJobMainPtr);
}
}
@@ -458,7 +587,7 @@ ncclResult_t ncclGroupJobComplete(struct ncclGroupJob* groupJob) {
ncclResult_t ncclGroupJobAbort(struct ncclGroupJob* groupJob) {
if (groupJob && groupJob->initialized) {
__atomic_store_n(groupJob->abortFlagPtr, true, __ATOMIC_RELAXED);
__atomic_store_n(groupJob->abortFlagPtr, true, __ATOMIC_RELEASE);
NCCLCHECK(ncclGroupJobComplete(groupJob));
}
return ncclSuccess;