Merge remote-tracking branch 'nccl/master' into develop
This commit is contained in:
BertanDogancay
+245
-115
@@ -19,16 +19,14 @@
|
||||
|
||||
using namespace rccl;
|
||||
|
||||
#define GROUP_MAX_RECLAIM_STEPS 10
|
||||
|
||||
__thread int ncclGroupDepth = 0; // depth of ncclGroupStart nesting
|
||||
__thread ncclResult_t ncclGroupError = ncclSuccess;
|
||||
__thread struct ncclComm* ncclGroupCommHead = nullptr;
|
||||
__thread struct ncclComm* ncclGroupCommHead[ncclGroupTaskTypeNum] = {nullptr};
|
||||
__thread struct ncclComm* ncclGroupCommPreconnectHead = nullptr;
|
||||
__thread struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> ncclAsyncJobs;
|
||||
__thread struct ncclGroupJob *ncclGroupJobMainPtr = NULL;
|
||||
__thread struct ncclGroupJob ncclGroupJobMain;
|
||||
__thread int ncclGroupBlocking = -1; /* default mode */
|
||||
__thread bool ncclGroupJobAbortFlag = false;
|
||||
|
||||
void* ncclAsyncJobMain(void* arg);
|
||||
|
||||
ncclResult_t ncclAsyncLaunch(
|
||||
@@ -219,6 +217,66 @@ fail:
|
||||
goto exit;
|
||||
}
|
||||
|
||||
struct ncclGroupSymmetricJob {
|
||||
struct ncclAsyncJob base;
|
||||
struct ncclComm* comm;
|
||||
};
|
||||
|
||||
NCCL_PARAM(WinStride, "WIN_STRIDE", -1);
|
||||
|
||||
ncclResult_t ncclCommGroupRegisterSymmetric(struct ncclAsyncJob* job_) {
|
||||
struct ncclGroupSymmetricJob* job = (struct ncclGroupSymmetricJob*)job_;
|
||||
struct ncclComm* comm = job->comm;
|
||||
ncclResult_t ret = ncclSuccess;
|
||||
|
||||
CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), ret, fail);
|
||||
if (comm->baseStride == 0) {
|
||||
cudaStream_t hostStream;
|
||||
// first time to allocate symmetric VA space.
|
||||
// calling into this function means symmetric is supported.
|
||||
struct ncclSymDevBase* symBase = NULL;
|
||||
size_t size = ncclSymDevBase::size(comm->localRanks);
|
||||
if (ncclParamWinStride() != -1) {
|
||||
comm->baseStride = ncclParamWinStride();
|
||||
} else {
|
||||
size_t maxStride = 0;
|
||||
for (int r = 0; r < comm->nRanks; ++r)
|
||||
if (comm->peerInfo[r].totalGlobalMem > maxStride) maxStride = comm->peerInfo[r].totalGlobalMem;
|
||||
comm->baseStride = maxStride;
|
||||
}
|
||||
INFO(NCCL_INIT, "rank %d base stride %zuGB total VM %zuGB", comm->rank, comm->baseStride >> 30, (comm->baseStride * comm->localRanks) >> 30);
|
||||
NCCLCHECKGOTO(ncclIpcSymmetricInit(comm), ret, fail);
|
||||
NCCLCHECKGOTO(ncclNvlsSymmetricInit(comm), ret, fail);
|
||||
comm->symAllocHead = 0;
|
||||
|
||||
// Allocate symmetric memory for NCCL internal usage
|
||||
NCCLCHECKGOTO(ncclCommSymmetricAllocInternal(comm, size, alignof(struct ncclSymDevBase), (void**)&symBase), ret, fail);
|
||||
assert((void*)symBase == (void*)(comm->baseUCSymPtr + comm->localRank * comm->baseStride));
|
||||
NCCLCHECKGOTO(ncclStrongStreamAcquire(ncclCudaGraphNone(), &comm->sharedRes->hostStream, /*concurrent=*/false, &hostStream), ret, fail);
|
||||
CUDACHECKGOTO(cudaMemsetAsync(symBase, 0, size, hostStream), ret, fail);
|
||||
CUDACHECKGOTO(cudaStreamSynchronize(hostStream), ret, fail);
|
||||
NCCLCHECKGOTO(ncclStrongStreamRelease(ncclCudaGraphNone(), &comm->sharedRes->hostStream, /*concurrent=*/false), ret, fail);
|
||||
|
||||
comm->symDevComm.base = (struct ncclSymDevBase*)(comm->baseUCSymPtr + comm->localRank * comm->baseStride);
|
||||
comm->symDevComm.baseMc = (struct ncclSymDevBase*)comm->baseMCSymPtr;
|
||||
comm->symDevComm.nRanks = comm->localRanks;
|
||||
comm->symDevComm.nRanks_rcp32 = idivRcp32(comm->localRanks);
|
||||
comm->symDevComm.rank = comm->localRank;
|
||||
comm->symDevComm.stride4G = comm->baseStride >> 32;
|
||||
}
|
||||
|
||||
while (!ncclIntruQueueEmpty(&comm->symRegTaskQueue)) {
|
||||
struct ncclSymRegTask* task = ncclIntruQueueDequeue(&comm->symRegTaskQueue);
|
||||
NCCLCHECKGOTO(ncclCommSymmetricRegisterInternal(comm, task->buff, task->baseSize, task->alignment, task->memHandle, task->regHandle), ret, fail);
|
||||
free(task);
|
||||
}
|
||||
|
||||
exit:
|
||||
return ret;
|
||||
fail:
|
||||
goto exit;
|
||||
}
|
||||
|
||||
static ncclResult_t doLaunches(struct ncclComm* head) {
|
||||
ncclResult_t result = ncclSuccess;
|
||||
struct ncclComm* cliqueHead = head;
|
||||
@@ -235,7 +293,7 @@ static ncclResult_t doLaunches(struct ncclComm* head) {
|
||||
CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
|
||||
NCCLCHECKGOTO(ncclLaunchPrepare(comm), result, failure);
|
||||
if (useBarrier) ncclCommIntraBarrierIn(comm, 1);
|
||||
comm = comm->groupNext;
|
||||
comm = comm->groupNext[ncclGroupTaskTypeCollective];
|
||||
} while (comm != nullptr && comm->intraComm0 == cliqueHead->intraComm0);
|
||||
cliqueNextHead = comm;
|
||||
|
||||
@@ -252,7 +310,7 @@ static ncclResult_t doLaunches(struct ncclComm* head) {
|
||||
bool moreRounds = false;
|
||||
comm = cliqueHead;
|
||||
do { // Iterate clique members.
|
||||
struct ncclComm* next = comm->groupNext;
|
||||
struct ncclComm* next = comm->groupNext[ncclGroupTaskTypeCollective];
|
||||
if (useBarrier) {
|
||||
// Barrier reduction result tells us if this was the final round.
|
||||
moreRounds = 0 != ncclCommIntraBarrierOut(comm);
|
||||
@@ -287,66 +345,62 @@ failure:
|
||||
return result;
|
||||
}
|
||||
|
||||
static inline void groupResetJobState(struct ncclGroupJob* job) {
|
||||
if (job) {
|
||||
if (job->groupBlockingPtr) *job->groupBlockingPtr = -1;
|
||||
if (job->abortFlagPtr) *job->abortFlagPtr = false;
|
||||
if (job->groupErrorPtr) *job->groupErrorPtr = ncclSuccess;
|
||||
if (job->groupCommHeadPtr) *job->groupCommHeadPtr = NULL;
|
||||
if (job->groupCommPreconnectHeadPtr) *job->groupCommPreconnectHeadPtr = NULL;
|
||||
memset(job, 0, sizeof(struct ncclGroupJob));
|
||||
}
|
||||
static inline void groupLocalResetJobState() {
|
||||
ncclGroupError = ncclSuccess;
|
||||
for (int type = 0; type < ncclGroupTaskTypeNum; ++type) ncclGroupCommHead[type] = NULL;
|
||||
ncclGroupCommPreconnectHead = NULL;
|
||||
ncclGroupBlocking = -1;
|
||||
ncclIntruQueueConstruct(&ncclAsyncJobs);
|
||||
return;
|
||||
}
|
||||
|
||||
static void groupCleanup(struct ncclComm** groupCommHeadPtr, struct ncclComm** groupCommPreconnectHeadPtr, struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next>* asyncJobsPtr, ncclResult_t* groupErrorPtr, int* groupBlockingPtr, volatile bool* groupJobAbortFlagPtr, ncclResult_t error) {
|
||||
struct ncclComm* comm = *groupCommHeadPtr;
|
||||
|
||||
/* reset all thread local variables */
|
||||
*groupCommHeadPtr = NULL;
|
||||
*groupCommPreconnectHeadPtr = NULL;
|
||||
*groupErrorPtr = ncclSuccess;
|
||||
*groupBlockingPtr = -1;
|
||||
*groupJobAbortFlagPtr = false;
|
||||
|
||||
while (comm != nullptr) {
|
||||
struct ncclComm* next = comm->groupNext;
|
||||
(void) ncclGroupCommLeave(comm); // overwrites comm->groupNext
|
||||
// We don't know if preconnect succeeded or happened at all, so clear
|
||||
// the flags that let `taskAppend()` skip over checking if preconnect
|
||||
// is needed.
|
||||
comm->preconnectNext = reinterpret_cast<struct ncclComm*>(0x1);
|
||||
for (int i = 0; i < comm->nRanks; i++) {
|
||||
for (int j = 0; j < MAXCHANNELS/64; j++) {
|
||||
comm->connectSend[i].masks[j] = 0UL;
|
||||
comm->connectRecv[i].masks[j] = 0UL;
|
||||
}
|
||||
}
|
||||
// Reclaim abandoned kernel plan memory. Note ncclWork structs were already
|
||||
// reclaimed by a `ncclMemoryStackPop(&comm->memScoped)` during `ncclGroupCommLeave()`.
|
||||
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) {
|
||||
while (!ncclIntruQueueEmpty(&plan->proxyOpQueue)) {
|
||||
struct ncclProxyOp* pxop = ncclIntruQueueDequeue(&plan->proxyOpQueue);
|
||||
ncclMemoryPoolFree(&comm->memPool_ncclProxyOp, pxop);
|
||||
static void groupCleanup(struct ncclComm** groupCommHeadPtr, struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next>* asyncJobsPtr, ncclResult_t error) {
|
||||
struct ncclComm* comm;
|
||||
for (int type = 0; type < ncclGroupTaskTypeNum; ++type) {
|
||||
comm = groupCommHeadPtr[type];
|
||||
// reset groupCommHeadPtr[type]
|
||||
groupCommHeadPtr[type] = nullptr;
|
||||
while (comm != nullptr) {
|
||||
struct ncclComm* next = comm->groupNext[type];
|
||||
(void)ncclGroupCommLeave(comm, type); // overwrites comm->groupNext
|
||||
// We don't know if preconnect succeeded or happened at all, so clear
|
||||
// the flags that let `taskAppend()` skip over checking if preconnect
|
||||
// is needed.
|
||||
if (type == ncclGroupTaskTypeCollective) {
|
||||
comm->preconnectNext = reinterpret_cast<struct ncclComm*>(0x1);
|
||||
for (int i = 0; i < comm->nRanks; i++) {
|
||||
for (int j = 0; j < MAXCHANNELS/64; j++) {
|
||||
comm->connectSend[i].masks[j] = 0UL;
|
||||
comm->connectRecv[i].masks[j] = 0UL;
|
||||
}
|
||||
}
|
||||
// Reclaim abandoned kernel plan memory. Note ncclWork structs were already
|
||||
// reclaimed by a `ncclMemoryStackPop(&comm->memScoped)` during `ncclGroupCommLeave()`.
|
||||
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) {
|
||||
while (!ncclIntruQueueEmpty(&plan->proxyOpQueue)) {
|
||||
struct ncclProxyOp* pxop = ncclIntruQueueDequeue(&plan->proxyOpQueue);
|
||||
ncclMemoryPoolFree(&comm->memPool_ncclProxyOp, pxop);
|
||||
}
|
||||
ncclMemoryPoolFree(&comm->memPool_ncclKernelPlan, plan);
|
||||
}
|
||||
}
|
||||
|
||||
{ // Reset comm->planner to empty.
|
||||
ncclKernelPlanner::Peer* tmp = comm->planner.peers;
|
||||
memset(&comm->planner, 0, sizeof(comm->planner));
|
||||
comm->planner.peers = tmp;
|
||||
if (comm->planner.peers != NULL) memset(comm->planner.peers, 0, comm->nRanks * sizeof(comm->planner.peers[0]));
|
||||
}
|
||||
ncclMemoryPoolFree(&comm->memPool_ncclKernelPlan, plan);
|
||||
}
|
||||
}
|
||||
|
||||
{ // Reset comm->planner to empty.
|
||||
ncclKernelPlanner::Peer* tmp = comm->planner.peers;
|
||||
memset(&comm->planner, 0, sizeof(comm->planner));
|
||||
comm->planner.peers = tmp;
|
||||
if (comm->planner.peers != NULL) memset(comm->planner.peers, 0, comm->nRanks*sizeof(comm->planner.peers[0]));
|
||||
if (!comm->config.blocking)
|
||||
(void)ncclCommSetAsyncError(comm, error);
|
||||
comm = next;
|
||||
}
|
||||
|
||||
if (!comm->config.blocking)
|
||||
(void) ncclCommSetAsyncError(comm, error);
|
||||
comm = next;
|
||||
}
|
||||
|
||||
/* reset everything */
|
||||
@@ -423,11 +477,10 @@ fail:
|
||||
static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInfo = NULL) {
|
||||
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;
|
||||
struct ncclComm **groupCommHeadMain = gjob->groupCommHead;
|
||||
struct ncclComm *groupCommPreconnectHeadMain = gjob->groupCommPreconnectHead;
|
||||
struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> *asyncJobsMain = &gjob->asyncJobs;
|
||||
bool *groupAbortFlag = &gjob->abortFlag;
|
||||
|
||||
if (!simInfo && groupCommPreconnectHeadMain != nullptr) {
|
||||
struct ncclComm* comm = groupCommPreconnectHeadMain;
|
||||
@@ -451,9 +504,41 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInf
|
||||
|
||||
NCCLCHECKGOTO(asyncJobLaunch(asyncJobsMain, groupAbortFlag), ret, fail);
|
||||
|
||||
// only loop through sym alloc and register tasks
|
||||
for (int type = ncclGroupTaskTypeSymRegister; type <= ncclGroupTaskTypeSymRegister; ++type) {
|
||||
if (groupCommHeadMain[type]) {
|
||||
struct ncclComm* cliqueHead = groupCommHeadMain[type];
|
||||
struct ncclComm* comm = NULL;
|
||||
struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> asyncSymJobs;
|
||||
ncclIntruQueueConstruct(&asyncSymJobs);
|
||||
do {
|
||||
comm = cliqueHead;
|
||||
do {
|
||||
struct ncclGroupSymmetricJob* job;
|
||||
NCCLCHECKGOTO(ncclCalloc(&job, 1), ret, fail);
|
||||
job->base.func = ncclCommGroupRegisterSymmetric;
|
||||
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(&asyncSymJobs, (struct ncclAsyncJob*)job);
|
||||
comm = comm->groupNext[type];
|
||||
} while (comm != nullptr && comm->intraComm0 == cliqueHead->intraComm0);
|
||||
NCCLCHECKGOTO(asyncJobLaunch(&asyncSymJobs, groupAbortFlag), ret, fail);
|
||||
while (!ncclIntruQueueEmpty(&asyncSymJobs)) {
|
||||
struct ncclAsyncJob* job = ncclIntruQueueDequeue(&asyncSymJobs);
|
||||
if (job->destructor) job->destructor((void*)job);
|
||||
}
|
||||
cliqueHead = comm;
|
||||
} while (cliqueHead != nullptr);
|
||||
}
|
||||
}
|
||||
|
||||
/* Connect channels at runtime if cumem is supported */
|
||||
if (groupCommHeadMain != nullptr) {
|
||||
struct ncclComm* cliqueHead = groupCommHeadMain;
|
||||
if (groupCommHeadMain[ncclGroupTaskTypeCollective] != nullptr) {
|
||||
struct ncclComm* cliqueHead = groupCommHeadMain[ncclGroupTaskTypeCollective];
|
||||
struct ncclComm* comm = NULL;
|
||||
struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> asyncCollJobs;
|
||||
ncclIntruQueueConstruct(&asyncCollJobs);
|
||||
@@ -484,7 +569,7 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInf
|
||||
memcpy(job->algoNeedConnect, algoNeedConnect, sizeof(bool) * NCCL_NUM_ALGORITHMS);
|
||||
ncclIntruQueueEnqueue(&asyncCollJobs, &job->base);
|
||||
}
|
||||
comm = comm->groupNext;
|
||||
comm = comm->groupNext[ncclGroupTaskTypeCollective];
|
||||
} while (comm != nullptr && comm->intraComm0 == cliqueHead->intraComm0);
|
||||
// connect
|
||||
NCCLCHECKGOTO(asyncJobLaunch(&asyncCollJobs, groupAbortFlag), ret, fail);
|
||||
@@ -496,42 +581,49 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInf
|
||||
} while (cliqueHead != nullptr);
|
||||
|
||||
// done with all buffer allocation, start registration and enqueue
|
||||
comm = groupCommHeadMain;
|
||||
comm = groupCommHeadMain[ncclGroupTaskTypeCollective];
|
||||
do {
|
||||
CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), ret, fail);
|
||||
NCCLCHECKGOTO(ncclTasksRegAndEnqueue(comm), ret, fail);
|
||||
comm = comm->groupNext;
|
||||
comm = comm->groupNext[ncclGroupTaskTypeCollective];
|
||||
} while (comm);
|
||||
}
|
||||
|
||||
if ((!simInfo) && (groupCommHeadMain != nullptr)) {
|
||||
NCCLCHECKGOTO(doLaunches(groupCommHeadMain), ret, fail);
|
||||
if ((!simInfo) && (groupCommHeadMain[ncclGroupTaskTypeCollective] != nullptr)) {
|
||||
NCCLCHECKGOTO(doLaunches(groupCommHeadMain[ncclGroupTaskTypeCollective]), ret, fail);
|
||||
}
|
||||
|
||||
while (!ncclIntruQueueEmpty(asyncJobsMain)) {
|
||||
struct ncclAsyncJob* job = ncclIntruQueueDequeue(asyncJobsMain);
|
||||
if (!job->destroyFlag && job->comm && !job->comm->config.blocking)
|
||||
if (!job->destroyFlag && job->comm && !job->comm->config.blocking && groupCommHeadMain[ncclGroupTaskTypeCollective] == nullptr)
|
||||
(void) ncclCommSetAsyncError(job->comm, ret);
|
||||
if (job->destructor) job->destructor((void*)job);
|
||||
}
|
||||
|
||||
while (groupCommHeadMain != nullptr) {
|
||||
struct ncclComm* comm = groupCommHeadMain;
|
||||
struct ncclComm* next = comm->groupNext;
|
||||
// Poll for callbacks sent to us from other threads. Typically these free
|
||||
// resources from to our memory pools and UB
|
||||
NCCLCHECKGOTO(ncclCommPollCallbacks(comm, /*waitSome=*/false), ret, fail);
|
||||
(void) ncclGroupCommLeave(comm);
|
||||
if (!comm->config.blocking) {
|
||||
(void) ncclCommSetAsyncError(comm, ret);
|
||||
for (int type = 0; type < ncclGroupTaskTypeNum; ++type) {
|
||||
while (groupCommHeadMain[type] != nullptr) {
|
||||
struct ncclComm* comm = groupCommHeadMain[type];
|
||||
struct ncclComm* next = comm->groupNext[type];
|
||||
// Poll for callbacks sent to us from other threads. Typically these free
|
||||
// resources from to our memory pools and UB
|
||||
if (comm->reclaimSteps == GROUP_MAX_RECLAIM_STEPS) {
|
||||
NCCLCHECKGOTO(ncclCommPollCallbacks(comm, /*waitSome=*/false), ret, fail);
|
||||
comm->reclaimSteps = 0;
|
||||
} else {
|
||||
comm->reclaimSteps++;
|
||||
}
|
||||
(void)ncclGroupCommLeave(comm, type);
|
||||
if (!comm->config.blocking) {
|
||||
(void)ncclCommSetAsyncError(comm, ret);
|
||||
}
|
||||
groupCommHeadMain[type] = next;
|
||||
}
|
||||
groupCommHeadMain = next;
|
||||
}
|
||||
|
||||
exit:
|
||||
return ret;
|
||||
fail:
|
||||
groupCleanup(gjob->groupCommHeadPtr, gjob->groupCommPreconnectHeadPtr, gjob->asyncJobsPtr, gjob->groupErrorPtr, gjob->groupBlockingPtr, gjob->abortFlagPtr, ret);
|
||||
groupCleanup(gjob->groupCommHead, &gjob->asyncJobs, ret);
|
||||
goto exit;
|
||||
}
|
||||
|
||||
@@ -544,6 +636,8 @@ ncclResult_t ncclGroupEndInternal(ncclSimInfo_t* simInfo) {
|
||||
ncclSimInfo_t internalSimInfo = NCCL_SIM_INFO_INITIALIZER;
|
||||
ncclSimInfo_t* internalSimInfoPtr = NULL;
|
||||
size_t realSize = 0;
|
||||
bool hasCommHead = false;
|
||||
ncclGroupJob* groupJob = NULL;
|
||||
|
||||
internalSimInfo.magic = 0;
|
||||
|
||||
@@ -573,72 +667,108 @@ ncclResult_t ncclGroupEndInternal(ncclSimInfo_t* simInfo) {
|
||||
internalSimInfoPtr = &internalSimInfo;
|
||||
}
|
||||
|
||||
if (ncclGroupCommHead != nullptr || !ncclIntruQueueEmpty(&ncclAsyncJobs) || ncclGroupCommPreconnectHead != nullptr) {
|
||||
ncclGroupJobMain.groupCommHeadPtr = &ncclGroupCommHead;
|
||||
ncclGroupJobMain.groupCommPreconnectHeadPtr = &ncclGroupCommPreconnectHead;
|
||||
ncclGroupJobMain.groupErrorPtr = &ncclGroupError;
|
||||
ncclGroupJobMain.asyncJobsPtr = &ncclAsyncJobs;
|
||||
ncclGroupJobMain.abortFlagPtr = &ncclGroupJobAbortFlag;
|
||||
ncclGroupJobMain.groupBlockingPtr = &ncclGroupBlocking;
|
||||
ncclGroupJobMain.initialized = true;
|
||||
ncclGroupJobMainPtr = &ncclGroupJobMain;
|
||||
for (int type = 0; type < ncclGroupTaskTypeNum; ++type) {
|
||||
if (ncclGroupCommHead[type]) {
|
||||
hasCommHead = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
NCCLCHECKGOTO(ncclCalloc(&groupJob, 1), ret, fail);
|
||||
ncclIntruQueueConstruct(&groupJob->asyncJobs);
|
||||
groupJob->groupRefCount = 0;
|
||||
groupJob->nonBlockingInit = false;
|
||||
memcpy(groupJob->groupCommHead, ncclGroupCommHead, sizeof(ncclGroupCommHead));
|
||||
groupJob->groupCommPreconnectHead = ncclGroupCommPreconnectHead;
|
||||
groupJob->groupError = ncclSuccess;
|
||||
groupJob->abortFlag = false;
|
||||
groupJob->joined = false;
|
||||
ncclIntruQueueTransfer(&groupJob->asyncJobs, &ncclAsyncJobs);
|
||||
|
||||
if (hasCommHead || !ncclIntruQueueEmpty(&groupJob->asyncJobs) || ncclGroupCommPreconnectHead != nullptr) {
|
||||
/* make sure ncclGroupBlocking has been set. */
|
||||
assert(ncclGroupBlocking == 0 || ncclGroupBlocking == 1);
|
||||
if (ncclGroupBlocking == 0 && (ncclGroupCommPreconnectHead != nullptr || !ncclIntruQueueEmpty(&ncclAsyncJobs))) {
|
||||
/* nonblocking group */
|
||||
if (!ncclIntruQueueEmpty(&ncclAsyncJobs)) {
|
||||
ncclAsyncJob* job = ncclIntruQueueHead(&ncclAsyncJobs);
|
||||
if (!ncclIntruQueueEmpty(&groupJob->asyncJobs)) {
|
||||
ncclAsyncJob* job = ncclIntruQueueHead(&groupJob->asyncJobs);
|
||||
do {
|
||||
NCCLCHECKGOTO(ncclCommSetAsyncError(job->comm, ncclInProgress), ret, fail);
|
||||
job->comm->groupJob = ncclGroupJobMainPtr;
|
||||
if (job->comm->groupJob == NULL) {
|
||||
job->comm->groupJob = groupJob;
|
||||
groupJob->groupRefCount++;
|
||||
}
|
||||
job = job->next;
|
||||
} while (job);
|
||||
}
|
||||
|
||||
if (ncclGroupCommHead) {
|
||||
ncclComm_t comm = ncclGroupCommHead;
|
||||
do {
|
||||
NCCLCHECKGOTO(ncclCommSetAsyncError(comm, ncclInProgress), ret, fail);
|
||||
/* link group job to communicators. */
|
||||
comm->groupJob = ncclGroupJobMainPtr;
|
||||
comm = comm->groupNext;
|
||||
} while (comm);
|
||||
for (int type = 0; type < ncclGroupTaskTypeNum; ++type) {
|
||||
if (ncclGroupCommHead[type]) {
|
||||
ncclComm_t comm = ncclGroupCommHead[type];
|
||||
do {
|
||||
NCCLCHECKGOTO(ncclCommSetAsyncError(comm, ncclInProgress), ret, fail);
|
||||
/* link group job to communicators. */
|
||||
if (comm->groupJob == NULL) {
|
||||
comm->groupJob = groupJob;
|
||||
groupJob->groupRefCount++;
|
||||
}
|
||||
comm = comm->groupNext[type];
|
||||
} while (comm);
|
||||
}
|
||||
}
|
||||
|
||||
ncclGroupJobMainPtr->base.func = groupLaunchNonBlocking;
|
||||
PTHREADCHECKGOTO(pthread_create(&ncclGroupJobMainPtr->base.thread, NULL, ncclAsyncJobMain, (void*)&ncclGroupJobMainPtr->base), "pthread_create", ret, fail);
|
||||
groupJob->base.func = groupLaunchNonBlocking;
|
||||
PTHREADCHECKGOTO(pthread_create(&groupJob->base.thread, NULL, ncclAsyncJobMain, (void*)&groupJob->base), "pthread_create", ret, fail);
|
||||
groupJob->nonBlockingInit = true;
|
||||
ret = ncclInProgress;
|
||||
} else {
|
||||
/* blocking group */
|
||||
int savedDev;
|
||||
CUDACHECKGOTO(cudaGetDevice(&savedDev), ret, fail);
|
||||
NCCLCHECKGOTO(groupLaunch(&ncclGroupJobMainPtr->base, internalSimInfoPtr), ret, fail);
|
||||
NCCLCHECKGOTO(groupLaunch(&groupJob->base, internalSimInfoPtr), ret, fail);
|
||||
CUDACHECKGOTO(cudaSetDevice(savedDev), ret, fail);
|
||||
if (simInfo) memcpy((void*)simInfo, (void*)internalSimInfoPtr, realSize);
|
||||
groupResetJobState(ncclGroupJobMainPtr);
|
||||
free(groupJob);
|
||||
}
|
||||
}
|
||||
/* Reset the job state for the next group call. */
|
||||
groupLocalResetJobState();
|
||||
|
||||
exit:
|
||||
return ret;
|
||||
fail:
|
||||
groupCleanup(&ncclGroupCommHead, &ncclGroupCommPreconnectHead, &ncclAsyncJobs, &ncclGroupError, &ncclGroupBlocking, &ncclGroupJobAbortFlag, ret);
|
||||
if (groupJob) {
|
||||
groupCleanup(groupJob->groupCommHead, &groupJob->asyncJobs, ret);
|
||||
free(groupJob);
|
||||
} else {
|
||||
groupCleanup(ncclGroupCommHead, &ncclAsyncJobs, ret);
|
||||
}
|
||||
groupLocalResetJobState();
|
||||
goto exit;
|
||||
}
|
||||
|
||||
ncclResult_t ncclGroupJobComplete(struct ncclGroupJob* groupJob) {
|
||||
ncclResult_t ret = ncclSuccess;
|
||||
if (groupJob && groupJob->initialized) {
|
||||
ret = ncclAsyncJobComplete(&groupJob->base);
|
||||
groupResetJobState(groupJob);
|
||||
if (groupJob && groupJob->nonBlockingInit) {
|
||||
if (!__atomic_exchange_n(&groupJob->joined, true, __ATOMIC_ACQ_REL)) {
|
||||
ret = ncclAsyncJobComplete(&groupJob->base);
|
||||
}
|
||||
if (ncclAtomicRefCountDecrement(&groupJob->groupRefCount) == 0) {
|
||||
free(groupJob);
|
||||
}
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
ncclResult_t ncclGroupJobAbort(struct ncclGroupJob* groupJob) {
|
||||
if (groupJob && groupJob->initialized) {
|
||||
__atomic_store_n(groupJob->abortFlagPtr, true, __ATOMIC_RELEASE);
|
||||
NCCLCHECK(ncclGroupJobComplete(groupJob));
|
||||
if (groupJob && groupJob->nonBlockingInit) {
|
||||
if (!__atomic_exchange_n(&groupJob->joined, true, __ATOMIC_ACQ_REL)) {
|
||||
__atomic_store_n(&groupJob->abortFlag, true, __ATOMIC_RELAXED);
|
||||
ncclAsyncJobComplete(&groupJob->base);
|
||||
}
|
||||
if (ncclAtomicRefCountDecrement(&groupJob->groupRefCount) == 0) {
|
||||
free(groupJob);
|
||||
}
|
||||
}
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user