Merge remote-tracking branch 'nccl/master' into develop
This commit is contained in:
Wenkai Du
+212
-396
@@ -11,446 +11,262 @@
|
||||
#include "transport.h"
|
||||
#include "channel.h"
|
||||
|
||||
#define MAX_ASYNC_OPS 128
|
||||
thread_local pthread_t ncclGroupThreads[MAX_ASYNC_OPS];
|
||||
thread_local int ncclGroupIndex = 0;
|
||||
thread_local int ncclGroupMode = 0;
|
||||
thread_local ncclResult_t ncclGroupError = ncclSuccess;
|
||||
extern struct allocationTracker allocTracker[];
|
||||
__thread int ncclGroupDepth = 0; // depth of ncclGroupStart nesting
|
||||
__thread ncclResult_t ncclGroupError = ncclSuccess;
|
||||
__thread struct ncclComm* ncclGroupCommHead = nullptr;
|
||||
__thread struct ncclComm* ncclGroupCommPreconnectHead = nullptr;
|
||||
__thread struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> ncclAsyncJobs;
|
||||
|
||||
bool ncclAsyncMode() {
|
||||
return ncclGroupMode > 0;
|
||||
}
|
||||
|
||||
ncclResult_t ncclAsyncErrCheck(ncclResult_t ret) {
|
||||
if (ncclGroupError == ncclSuccess || ret != ncclSuccess) ncclGroupError = ret;
|
||||
return ret;
|
||||
}
|
||||
|
||||
struct ncclInitArgs {
|
||||
ncclInitFunc_t func;
|
||||
int cudaDev;
|
||||
ncclComm_t* newcomm;
|
||||
int ndev;
|
||||
ncclUniqueId commId;
|
||||
int myrank;
|
||||
int virtualId;
|
||||
};
|
||||
struct ncclCollArgs {
|
||||
ncclComm_t comm;
|
||||
uint16_t connIndex;
|
||||
};
|
||||
|
||||
enum ncclAsyncFuncType {
|
||||
ASYNC_FUNC_INVALID = 0,
|
||||
ASYNC_FUNC_INIT = 1,
|
||||
ASYNC_FUNC_COLL = 2,
|
||||
};
|
||||
struct ncclAsyncArgs {
|
||||
ncclResult_t ret;
|
||||
enum ncclAsyncFuncType funcType;
|
||||
union {
|
||||
ncclCollArgs coll;
|
||||
ncclInitArgs init;
|
||||
};
|
||||
};
|
||||
|
||||
thread_local struct ncclAsyncArgs ncclGroupArgs[MAX_ASYNC_OPS];
|
||||
|
||||
void* ncclAsyncThreadMain(void* args_) {
|
||||
struct ncclAsyncArgs* args = (struct ncclAsyncArgs*)args_;
|
||||
NCCLCHECKTHREAD(args->init.func(args->init.newcomm, args->init.ndev, args->init.commId, args->init.myrank,
|
||||
args->init.cudaDev, args->init.virtualId));
|
||||
return args;
|
||||
}
|
||||
|
||||
ncclResult_t ncclAsyncInit(ncclInitFunc_t func, ncclComm_t* newcomm, int ndev, ncclUniqueId commId, int myrank, int cudaDev, int virtualId) {
|
||||
if (ncclGroupIndex >= MAX_ASYNC_OPS) {
|
||||
WARN("Too many async operations in progress, max is %d", MAX_ASYNC_OPS);
|
||||
return ncclAsyncErrCheck(ncclInvalidUsage);
|
||||
ncclResult_t ncclAsyncLaunch(
|
||||
struct ncclAsyncJob* job,
|
||||
ncclResult_t(*func)(struct ncclAsyncJob*),
|
||||
void(*undo)(struct ncclAsyncJob*),
|
||||
void(*destructor)(void*)
|
||||
) {
|
||||
if (0 == ncclGroupDepth) {
|
||||
ncclResult_t res = func(job);
|
||||
if (res != ncclSuccess && undo) undo(job);
|
||||
if (destructor) destructor(job);
|
||||
return res;
|
||||
} else {
|
||||
job->func = func;
|
||||
job->undo = undo;
|
||||
job->destructor = destructor;
|
||||
ncclIntruQueueEnqueue(&ncclAsyncJobs, job);
|
||||
return ncclSuccess;
|
||||
}
|
||||
int index = ncclGroupIndex++;
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+index;
|
||||
args->funcType = ASYNC_FUNC_INIT;
|
||||
args->init.func = func;
|
||||
args->init.cudaDev = cudaDev;
|
||||
args->init.newcomm = newcomm;
|
||||
args->init.ndev = ndev;
|
||||
memcpy(&args->init.commId, &commId, sizeof(commId));
|
||||
args->init.myrank = myrank;
|
||||
args->init.virtualId = virtualId;
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
ncclResult_t ncclAsyncColl(ncclComm_t comm) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs;
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
if (args->coll.comm == comm) return ncclSuccess;
|
||||
args++;
|
||||
void* ncclAsyncJobMain(void* arg) {
|
||||
struct ncclAsyncJob* job = (struct ncclAsyncJob*)arg;
|
||||
job->result = job->func(job);
|
||||
if (job->result != ncclSuccess) {
|
||||
INFO(NCCL_INIT,"%s:%d -> %d [Async thread]", __FILE__, __LINE__, job->result);
|
||||
}
|
||||
if (ncclGroupIndex >= MAX_ASYNC_OPS) {
|
||||
WARN("Too many async operations in progress, max is %d", MAX_ASYNC_OPS);
|
||||
return ncclAsyncErrCheck(ncclInvalidUsage);
|
||||
}
|
||||
ncclGroupIndex++;
|
||||
args->funcType = ASYNC_FUNC_COLL;
|
||||
args->coll.comm = comm;
|
||||
return ncclSuccess;
|
||||
return arg;
|
||||
}
|
||||
|
||||
NCCL_API(ncclResult_t, ncclGroupStart);
|
||||
ncclResult_t ncclGroupStart() {
|
||||
NVTX3_FUNC_RANGE_IN(nccl_domain);
|
||||
if (ncclGroupMode == 0) {
|
||||
memset(ncclGroupArgs, 0, sizeof(struct ncclAsyncArgs)*MAX_ASYNC_OPS);
|
||||
}
|
||||
ncclGroupMode++;
|
||||
NCCLCHECK(ncclGroupStartInternal());
|
||||
TRACE_CALL("ncclGroupStart()");
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t scheduleSend(struct ncclComm* comm, int peer, int channelId, size_t count, void* buff, uint64_t opCount, uint16_t connIndex) {
|
||||
struct ncclInfo info = { ncclFuncSend, "Send",
|
||||
NULL, buff, count, ncclInt8, ncclSum, peer, comm, comm->userStream, /* Args */
|
||||
1, 1 };
|
||||
info.channelId = channelId;
|
||||
info.opCount = opCount;
|
||||
info.connIndex = connIndex;
|
||||
NCCLCHECK(ncclSetupP2pKernel(&info));
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t scheduleRecv(struct ncclComm* comm, int peer, int channelId, size_t count, void* buff, uint64_t opCount, uint16_t connIndex) {
|
||||
struct ncclInfo info = { ncclFuncRecv, "Recv",
|
||||
NULL, buff, count, ncclInt8, ncclSum, peer, comm, comm->userStream, /* Args */
|
||||
1, 1 };
|
||||
info.channelId = channelId;
|
||||
info.opCount = opCount;
|
||||
info.connIndex = connIndex;
|
||||
NCCLCHECK(ncclSetupP2pKernel(&info));
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
void* ncclAsyncThreadPreconnect(void* args_) {
|
||||
struct ncclAsyncArgs* args = (struct ncclAsyncArgs*)args_;
|
||||
struct ncclComm* comm = args->coll.comm;
|
||||
CUDACHECKTHREAD(hipSetDevice(comm->cudaDev));
|
||||
if (CPU_COUNT(&comm->cpuAffinity)) sched_setaffinity(0, sizeof(cpu_set_t), &comm->cpuAffinity);
|
||||
NCCLCHECKTHREAD(ncclTransportP2pSetup(comm, NULL, args->coll.connIndex));
|
||||
return args;
|
||||
}
|
||||
|
||||
static size_t getP2pChunkSize(size_t totalSize, int minChannels, int maxChannels, size_t minSize, size_t maxSize) {
|
||||
size_t size = std::max(minSize, DIVUP(totalSize, minChannels));
|
||||
int nChannels = minChannels;
|
||||
while (size > maxSize && nChannels <= maxChannels/2) {
|
||||
nChannels *= 2;
|
||||
size = DIVUP(totalSize, nChannels);
|
||||
}
|
||||
ALIGN_SIZE(size, minSize);
|
||||
return size;
|
||||
}
|
||||
|
||||
RCCL_PARAM(P2pNetThreshold, "P2P_NET_THRESHOLD", 131072);
|
||||
|
||||
NCCL_API(ncclResult_t, ncclGroupEnd);
|
||||
ncclResult_t ncclGroupEnd() {
|
||||
NVTX3_FUNC_RANGE_IN(nccl_domain);
|
||||
if (ncclGroupMode == 0) {
|
||||
NCCLCHECK(ncclGroupEndInternal());
|
||||
TRACE_CALL("ncclGroupEnd()");
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
struct ncclPreconnectJob {
|
||||
struct ncclAsyncJob base;
|
||||
struct ncclComm* comm;
|
||||
};
|
||||
ncclResult_t ncclPreconnectFunc(struct ncclAsyncJob* job_) {
|
||||
struct ncclPreconnectJob* job = (struct ncclPreconnectJob*)job_;
|
||||
struct ncclComm* comm = job->comm;
|
||||
CUDACHECK(hipSetDevice(comm->cudaDev));
|
||||
if (CPU_COUNT(&comm->cpuAffinity)) sched_setaffinity(0, sizeof(cpu_set_t), &comm->cpuAffinity);
|
||||
NCCLCHECK(ncclTransportP2pSetup(comm, NULL, 1));
|
||||
if (comm->p2pNet) NCCLCHECK(ncclTransportP2pSetup(comm, NULL, NCCL_CONN_IDX_P2P_NET));
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t doLaunches(struct ncclComm* head) {
|
||||
ncclResult_t result = ncclSuccess;
|
||||
struct ncclComm* cliqueComm0 = head->intraComm0;
|
||||
struct ncclComm* cliqueHead = head;
|
||||
struct ncclComm* cliqueNextHead;
|
||||
bool useBarrier = ncclParamLaunchMode == ncclLaunchModeGroup;
|
||||
// This outer loop iterates over cliques of comms which are siblings of the
|
||||
// same global entity. We calculate a clique as all comms which have the same
|
||||
// `intraComm0` value.
|
||||
do {
|
||||
struct ncclComm* comm = cliqueHead;
|
||||
bool capturingYes = false, capturingNo = false;
|
||||
do {
|
||||
(ncclCudaGraphValid(comm->tasks.capturingGraph) ? capturingYes : capturingNo) = true;
|
||||
CUDACHECKGOTO(hipSetDevice(comm->cudaDev), result, failure);
|
||||
NCCLCHECKGOTO(ncclLaunchPrepare(comm), result, failure);
|
||||
if (useBarrier) ncclCommIntraBarrierIn(comm, 1);
|
||||
comm = comm->groupNext;
|
||||
} while (comm != nullptr && comm->intraComm0 == cliqueComm0);
|
||||
cliqueNextHead = comm;
|
||||
|
||||
if (capturingYes && capturingNo) {
|
||||
// We have entered barriers but are aborting without leaving them. Thus
|
||||
// these comms are permanently trashed. We need a good mechanism for
|
||||
// tracking and reporting that.
|
||||
WARN("Either none or all communicators in a ncclGroup() can be CUDA graph captured.");
|
||||
result = ncclInvalidUsage;
|
||||
goto failure;
|
||||
}
|
||||
|
||||
while (true) { // Iterate rounds of launches for clique.
|
||||
bool moreRounds;
|
||||
comm = cliqueHead;
|
||||
do { // Iterate clique members.
|
||||
struct ncclComm* next = comm->groupNext;
|
||||
if (useBarrier) {
|
||||
// Barrier reduction result tells us if this was the final round.
|
||||
moreRounds = 0 != ncclCommIntraBarrierOut(comm);
|
||||
} else {
|
||||
moreRounds = comm->unlaunchedPlansHead != nullptr;
|
||||
}
|
||||
if (moreRounds) {
|
||||
// Pop next unlaunched kernel
|
||||
struct ncclKernelPlan* plan = comm->unlaunchedPlansHead;
|
||||
if (plan != nullptr) {
|
||||
comm->unlaunchedPlansHead = plan->next;
|
||||
CUDACHECKGOTO(hipSetDevice(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 (plan != nullptr) {
|
||||
NCCLCHECKGOTO(ncclLaunchKernelAfter_NoCuda(comm, plan), result, failure);
|
||||
}
|
||||
} else { // Final round.
|
||||
CUDACHECKGOTO(hipSetDevice(comm->cudaDev), result, failure);
|
||||
NCCLCHECKGOTO(ncclLaunchFinish(comm), result, failure);
|
||||
}
|
||||
comm = next;
|
||||
} while (comm != cliqueNextHead);
|
||||
if (!moreRounds) break;
|
||||
}
|
||||
cliqueHead = cliqueNextHead;
|
||||
} while (cliqueHead != nullptr);
|
||||
failure:
|
||||
return result;
|
||||
}
|
||||
|
||||
ncclResult_t ncclGroupEndInternal() {
|
||||
if (ncclGroupDepth == 0) {
|
||||
WARN("ncclGroupEnd: not in a group call.");
|
||||
return ncclInvalidUsage;
|
||||
}
|
||||
ncclGroupMode--;
|
||||
if (ncclGroupMode > 0) return ncclSuccess;
|
||||
ncclGroupDepth--;
|
||||
if (ncclGroupDepth > 0) return ncclSuccess;
|
||||
|
||||
int savedDev;
|
||||
CUDACHECK(hipGetDevice(&savedDev));
|
||||
int activeThreads = 0;
|
||||
int doneArray[MAX_ASYNC_OPS];
|
||||
for (int i=0; i<ncclGroupIndex; i++) doneArray[i] = 1;
|
||||
|
||||
ncclResult_t ret = ncclGroupError;
|
||||
int usingCudaGraphAll = -1;
|
||||
hipGraph_t* graphs = NULL;
|
||||
if (ret != ncclSuccess) goto group_cleanup;
|
||||
bool jobsDone = false;
|
||||
if (ret != ncclSuccess) goto failure;
|
||||
|
||||
/* Launch async ncclCommInitRank */
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_INIT) {
|
||||
pthread_create(ncclGroupThreads+i, NULL, ncclAsyncThreadMain, args);
|
||||
activeThreads++;
|
||||
doneArray[i] = 0;
|
||||
}
|
||||
}
|
||||
/* For init, since we use threads, we just wait for threads to complete */
|
||||
while (activeThreads) {
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_INIT && doneArray[i] == 0) {
|
||||
int err = pthread_tryjoin_np(ncclGroupThreads[i], NULL);
|
||||
if (err == EBUSY) continue;
|
||||
if (err != 0) ret = ncclSystemError;
|
||||
if (args->ret != ncclSuccess) ret = args->ret;
|
||||
doneArray[i] = 1;
|
||||
activeThreads--;
|
||||
}
|
||||
}
|
||||
if (ncclGroupCommPreconnectHead != nullptr) {
|
||||
struct ncclComm* comm = ncclGroupCommPreconnectHead;
|
||||
do {
|
||||
struct ncclPreconnectJob* job;
|
||||
NCCLCHECK(ncclCalloc(&job, 1));
|
||||
job->base.func = ncclPreconnectFunc;
|
||||
job->base.undo = nullptr;
|
||||
job->base.destructor = free;
|
||||
job->comm = comm;
|
||||
ncclIntruQueueEnqueue(&ncclAsyncJobs, &job->base);
|
||||
|
||||
struct ncclComm* next = comm->preconnectNext;
|
||||
comm->preconnectNext = reinterpret_cast<struct ncclComm*>(0x1);
|
||||
comm = next;
|
||||
} while (comm != nullptr);
|
||||
}
|
||||
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL && args->coll.comm->connect[1]) {
|
||||
args->coll.connIndex = 1;
|
||||
pthread_create(ncclGroupThreads+i, NULL, ncclAsyncThreadPreconnect, args);
|
||||
}
|
||||
}
|
||||
if (!ncclIntruQueueEmpty(&ncclAsyncJobs)) {
|
||||
struct ncclAsyncJob* job = ncclIntruQueueHead(&ncclAsyncJobs);
|
||||
do {
|
||||
pthread_create(&job->thread, nullptr, ncclAsyncJobMain, job);
|
||||
job = job->next;
|
||||
} while (job != nullptr);
|
||||
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL && args->coll.comm->connect[1]) {
|
||||
int err = pthread_join(ncclGroupThreads[i], NULL);
|
||||
job = ncclIntruQueueHead(&ncclAsyncJobs);
|
||||
do {
|
||||
int err = pthread_join(job->thread, nullptr);
|
||||
if (err != 0) {
|
||||
WARN("Error waiting for pthread_join : %s", strerror(errno));
|
||||
return ncclSystemError;
|
||||
ret = ncclSystemError;
|
||||
}
|
||||
INFO(NCCL_INIT, "comm %p rank %d total %ld bytes - P2P preconnect COMPLETE", args->coll.comm, args->coll.comm->rank, allocTracker[args->coll.comm->cudaDev].totalAllocSize);
|
||||
NCCLCHECKGOTO(args->ret, ret, end);
|
||||
args->coll.comm->connect[1] = 0;
|
||||
}
|
||||
if (ret == ncclSuccess && job->result != ncclSuccess) ret = job->result;
|
||||
job = job->next;
|
||||
} while (job != nullptr);
|
||||
|
||||
jobsDone = true;
|
||||
if (ret != ncclSuccess) goto failure;
|
||||
}
|
||||
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL && args->coll.comm->connect[NCCL_CONN_IDX_P2P_NET]) {
|
||||
args->coll.connIndex = NCCL_CONN_IDX_P2P_NET;
|
||||
pthread_create(ncclGroupThreads+i, NULL, ncclAsyncThreadPreconnect, args);
|
||||
}
|
||||
if (ncclGroupCommHead != nullptr) {
|
||||
NCCLCHECKGOTO(doLaunches(ncclGroupCommHead), ret, failure);
|
||||
do {
|
||||
struct ncclComm* comm = ncclGroupCommHead;
|
||||
struct ncclComm* next = comm->groupNext;
|
||||
ncclGroupCommLeave(comm);
|
||||
ncclGroupCommHead = next;
|
||||
} while (ncclGroupCommHead != nullptr);
|
||||
}
|
||||
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL && args->coll.comm->connect[NCCL_CONN_IDX_P2P_NET]) {
|
||||
int err = pthread_join(ncclGroupThreads[i], NULL);
|
||||
if (err != 0) {
|
||||
WARN("Error waiting for pthread_join : %s", strerror(errno));
|
||||
return ncclSystemError;
|
||||
if (false) {
|
||||
failure:
|
||||
struct ncclComm* comm = ncclGroupCommHead;
|
||||
while (comm != nullptr) {
|
||||
struct ncclComm* next = comm->groupNext;
|
||||
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++) {
|
||||
comm->tasks.peers[i].sendSeen = false;
|
||||
comm->tasks.peers[i].recvSeen = false;
|
||||
comm->connectSend[i] = 0;
|
||||
comm->connectRecv[i] = 0;
|
||||
comm->connectSend[i+comm->nRanks*NCCL_CONN_IDX_P2P_NET] = 0;
|
||||
comm->connectRecv[i+comm->nRanks*NCCL_CONN_IDX_P2P_NET] = 0;
|
||||
}
|
||||
INFO(NCCL_INIT, "comm %p rank %d total %ld bytes - P2P NET preconnect COMPLETE", args->coll.comm, args->coll.comm->rank, allocTracker[args->coll.comm->cudaDev].totalAllocSize);
|
||||
NCCLCHECKGOTO(args->ret, ret, end);
|
||||
args->coll.comm->connect[NCCL_CONN_IDX_P2P_NET] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL) {
|
||||
struct ncclComm* comm = args->coll.comm;
|
||||
int node = comm->node;
|
||||
int nNodes = comm->nNodes;
|
||||
int localRank = comm->localRank;
|
||||
|
||||
// Compute how much to split operations
|
||||
// Natural step size matching buffer steps.
|
||||
ssize_t stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / NCCL_STEPS;
|
||||
// Try to use all channels
|
||||
int nChannelsMax = comm->p2pnChannelsPerPeer;
|
||||
int nChannelsMin = nChannelsMax;
|
||||
// Try to use all channels, but one channel per operation.
|
||||
//while (nChannelsMin*comm->nRanks > std::max(comm->nChannels, comm->p2pnChannels) && nChannelsMin > 1) nChannelsMin /= 2;
|
||||
// Avoid overloading channels with 8+ operations as we loose the sync warp, hence a bit of bandwidth.
|
||||
//while (nChannelsMax*comm->nRanks > std::max(comm->nChannels, comm->p2pnChannels)*4 && nChannelsMax > 1) nChannelsMax /= 2;
|
||||
|
||||
while (comm->p2pSendCount > 0 || comm->p2pRecvCount > 0) {
|
||||
// schedule delta 0, +1, -1, +2, -2, ...
|
||||
// also make sure we don't do 0 twice, nor +n/2 and -n/2 if n is even.
|
||||
for (int d=0; d<=nNodes/4; d++) {
|
||||
int deltas[4] = { d, (nNodes-d)%nNodes, nNodes/2-d, (nNodes-(nNodes/2-d))%nNodes };
|
||||
int index = 0;
|
||||
int delta = deltas[index];
|
||||
sched_delta:
|
||||
uint32_t recvNode = (node+nNodes-delta)%nNodes;
|
||||
uint32_t sendNode = (node+delta)%nNodes;
|
||||
int steps = comm->maxLocalRanks;
|
||||
for (int s=0; s<steps; s++) {
|
||||
int recvIndex = (localRank-s+steps)%steps;
|
||||
int recvPeer = recvIndex<comm->nodeRanks[recvNode].localRanks ? comm->nodeRanks[recvNode].localRankToRank[recvIndex] : -1;
|
||||
int sendIndex = (localRank+s)%steps;
|
||||
int sendPeer = sendIndex<comm->nodeRanks[sendNode].localRanks ? comm->nodeRanks[sendNode].localRankToRank[sendIndex] : -1;
|
||||
struct ncclP2Pinfo* recv = recvPeer != -1 && comm->p2pRecvs[recvPeer] ? comm->p2pRecvs[recvPeer]->getNext() : NULL;
|
||||
struct ncclP2Pinfo* send = sendPeer != -1 && comm->p2pSends[sendPeer] ? comm->p2pSends[sendPeer]->getNext() : NULL;
|
||||
if (recv != NULL || send != NULL) {
|
||||
ssize_t totRecvBytes = -1, totSendBytes = -1;
|
||||
if (recv != NULL) totRecvBytes = recv->nbytes;
|
||||
if (send != NULL) totSendBytes = send->nbytes;
|
||||
if (recv) comm->p2pRecvCount--;
|
||||
if (send) comm->p2pSendCount--;
|
||||
if (recvPeer == comm->rank) { // Check self send/recv
|
||||
if (sendPeer != comm->rank) { WARN("Sendrecv schedule not aligned for self"); ret = ncclInternalError; goto group_cleanup; }
|
||||
if (send && recv == NULL) { WARN("Trying to send to self without a matching recv"); ret = ncclInvalidUsage; goto group_cleanup; }
|
||||
if (send == NULL && recv) { WARN("Trying to recv to self without a matching send"); ret = ncclInvalidUsage; goto group_cleanup; }
|
||||
}
|
||||
void* recvBuff = recv ? recv->buff : NULL;
|
||||
void* sendBuff = send ? send->buff : NULL;
|
||||
// After we recycle p2pSend/Recv, we're no longer allowed to dereference send or recv, only use them as boolean NULL/not NULL.
|
||||
if (recv && comm->p2pRecvs[recvPeer]->peakNext() == NULL) comm->p2pRecvs[recvPeer]->recycle();
|
||||
if (send && comm->p2pSends[sendPeer]->peakNext() == NULL) comm->p2pSends[sendPeer]->recycle();
|
||||
|
||||
ssize_t recvChunkSize = getP2pChunkSize(totRecvBytes, nChannelsMin, nChannelsMax, stepSize, SENDRECV_SLICEFACTOR*stepSize);
|
||||
ssize_t sendChunkSize = getP2pChunkSize(totSendBytes, nChannelsMin, nChannelsMax, stepSize, SENDRECV_SLICEFACTOR*stepSize);
|
||||
|
||||
uint16_t sendIdx = 1, recvIdx = 1;
|
||||
if(comm->p2pNet && totSendBytes > rcclParamP2pNetThreshold())
|
||||
sendIdx = NCCL_CONN_IDX_P2P_NET;
|
||||
if(comm->p2pNet && totRecvBytes > rcclParamP2pNetThreshold())
|
||||
recvIdx = NCCL_CONN_IDX_P2P_NET;
|
||||
|
||||
ssize_t sendOffset = 0;
|
||||
ssize_t recvOffset = 0;
|
||||
int sendRemaining = 1, recvRemaining = 1;
|
||||
int chunk = 0;
|
||||
do {
|
||||
int channelId;
|
||||
// Shuffle channels with s intra-node, and delta inter-node. Inter-node, make sure
|
||||
// to use multiple channels to guarantee progress on all ranks from the same node.
|
||||
ssize_t recvbytes = totRecvBytes-recvOffset;
|
||||
ssize_t sendbytes = totSendBytes-sendOffset;
|
||||
if (recvbytes > recvChunkSize) { recvbytes = recvChunkSize; } else { recvRemaining = 0; }
|
||||
if (sendbytes > sendChunkSize) { sendbytes = sendChunkSize; } else { sendRemaining = 0; }
|
||||
// 0-bytes send/recv are considered as syncs. Make sure we only add syncs when requested
|
||||
// (total size == 0), otherwise set size to -1.
|
||||
if (sendbytes < 0 || (sendbytes == 0 && totSendBytes != 0)) send = NULL;
|
||||
if (recvbytes < 0 || (recvbytes == 0 && totRecvBytes != 0)) recv = NULL;
|
||||
if (send || recv) {
|
||||
if (recv) {
|
||||
NCCLCHECK(ncclChannelCompute(comm, recvPeer, chunk%comm->p2pnChannelsPerPeer, ncclFuncRecv, &channelId));
|
||||
}
|
||||
else
|
||||
recvPeer = -1;
|
||||
if (send) {
|
||||
NCCLCHECK(ncclChannelCompute(comm, sendPeer, chunk%comm->p2pnChannelsPerPeer, ncclFuncSend, &channelId));
|
||||
}
|
||||
else
|
||||
sendPeer = -1;
|
||||
NCCLCHECKGOTO(scheduleRecv(comm, recvPeer, channelId, recvbytes, recv ? ((char*)recvBuff)+recvOffset : NULL, recv ? recv->opCount : 0, recvIdx), ret, group_cleanup);
|
||||
NCCLCHECKGOTO(scheduleSend(comm, sendPeer, channelId, sendbytes, send ? ((char*)sendBuff)+sendOffset : NULL, send ? send->opCount : 0, sendIdx), ret, group_cleanup);
|
||||
}
|
||||
recvOffset += recvChunkSize;
|
||||
sendOffset += sendChunkSize;
|
||||
chunk++;
|
||||
} while (sendRemaining || recvRemaining);
|
||||
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);
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
index++;
|
||||
if (index == 1 && deltas[1] == deltas[0]) index++;
|
||||
if (index == 2 && deltas[2] == deltas[0]) index++;
|
||||
if (index == 3 && deltas[3] == deltas[2]) index++;
|
||||
if (index == 3 && deltas[3] == deltas[1]) index++;
|
||||
if (index < 4) {
|
||||
delta = deltas[index];
|
||||
goto sched_delta;
|
||||
}
|
||||
ncclMemoryPoolFree(&comm->memPool_ncclKernelPlan, plan);
|
||||
}
|
||||
}
|
||||
// Reset comm->tasks to empty.
|
||||
comm->tasks.nTasksColl = 0;
|
||||
comm->tasks.nTasksP2p = 0;
|
||||
comm->tasks.streams = nullptr;
|
||||
ncclIntruQueueConstruct(&comm->tasks.collQueue);
|
||||
comm->tasks.collBytesTotal = 0;
|
||||
for (int i=0; i < comm->nRanks; i++) {
|
||||
ncclIntruQueueConstruct(&comm->tasks.peers[i].sendQueue);
|
||||
ncclIntruQueueConstruct(&comm->tasks.peers[i].recvQueue);
|
||||
}
|
||||
comm = next;
|
||||
}
|
||||
}
|
||||
|
||||
/* Collectives are done in three steps :
|
||||
* 0. Save kernels previously enqueued. Compute channel, algo, proto, etc.
|
||||
* 1. Barrier Check In. Only the last call may call cudaLaunchKernel[cooperative]
|
||||
* 2. Barrier Wait. No CUDA call is permitted
|
||||
* 3. Enqueue Events. CUDA event wait/enqueue.
|
||||
* This is needed because step 2 cannot call any CUDA primitive, otherwise if
|
||||
* cudaFree happens between 1 and 3, it could block that CUDA call and
|
||||
* prevent some ranks from launching their network threads, which would
|
||||
* prevent the NCCL call from completing, blocking the cudaFree call.
|
||||
*/
|
||||
|
||||
// Check whether we are in cuda graph mode
|
||||
NCCLCHECK(ncclCalloc(&graphs, ncclGroupIndex));
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL) {
|
||||
ncclComm_t comm = args->coll.comm;
|
||||
NCCLCHECKGOTO(ncclGetCudaGraph(comm, graphs+i), ret, group_cleanup);
|
||||
if (usingCudaGraphAll == -1) {
|
||||
usingCudaGraphAll = comm->usingCudaGraph;
|
||||
} else if (usingCudaGraphAll != comm->usingCudaGraph) {
|
||||
WARN("Illegal to have some communicators in graph mode while others not");
|
||||
ret = ncclInvalidUsage;
|
||||
goto group_cleanup;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL) {
|
||||
ncclComm_t comm = args->coll.comm;
|
||||
NCCLCHECKGOTO(ncclSetupAsyncKernels(comm), ret, group_cleanup);
|
||||
}
|
||||
}
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL) {
|
||||
if (args->coll.comm->userStream == hipStreamDefault/* ||
|
||||
args->coll.comm->userStream == hipStreamPerThread ||
|
||||
args->coll.comm->userStream == hipStreamLegacy*/)
|
||||
CUDACHECKGOTO(hipSetDevice(args->coll.comm->cudaDev), ret, end);
|
||||
if (usingCudaGraphAll == 1) {
|
||||
NCCLCHECKGOTO(ncclCudaGraphHostSetup(args->coll.comm, graphs[i]), ret, end);
|
||||
} else {
|
||||
ncclEnqueueHostSetup<0>(args->coll.comm->enqueueInfo);
|
||||
}
|
||||
NCCLCHECKGOTO(ncclLaunchBarrier(args->coll.comm), ret, end);
|
||||
}
|
||||
}
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL) {
|
||||
CUDACHECKGOTO(hipSetDevice(args->coll.comm->cudaDev), ret, end);
|
||||
NCCLCHECKGOTO(ncclLaunchKernel(args->coll.comm), ret, end);
|
||||
}
|
||||
}
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_COLL) {
|
||||
if (args->coll.comm->userStream == hipStreamDefault/* ||
|
||||
args->coll.comm->userStream == hipStreamPerThread ||
|
||||
args->coll.comm->userStream == hipStreamLegacy*/)
|
||||
CUDACHECKGOTO(hipSetDevice(args->coll.comm->cudaDev), ret, end);
|
||||
NCCLCHECKGOTO(ncclRecordEvents(args->coll.comm), ret, end);
|
||||
NCCLCHECKGOTO(ncclLaunchReset(args->coll.comm), ret, end);
|
||||
}
|
||||
while (!ncclIntruQueueEmpty(&ncclAsyncJobs)) {
|
||||
struct ncclAsyncJob* job = ncclIntruQueueDequeue(&ncclAsyncJobs);
|
||||
if (ret != ncclSuccess && jobsDone && job->undo) job->undo(job);
|
||||
if (job->destructor) job->destructor((void*)job);
|
||||
}
|
||||
|
||||
goto end;
|
||||
group_cleanup:
|
||||
if (ret != ncclSuccess) {
|
||||
// At least one call in the group failed. Since we want to make that group
|
||||
// an atomic operation, we need to cancel all operations.
|
||||
for (int i=0; i<ncclGroupIndex; i++) {
|
||||
struct ncclAsyncArgs* args = ncclGroupArgs+i;
|
||||
if (args->funcType == ASYNC_FUNC_INIT) {
|
||||
if (args->init.newcomm) ncclCommDestroy(*args->init.newcomm);
|
||||
*args->init.newcomm = NULL;
|
||||
} else {
|
||||
struct ncclComm* comm = args->coll.comm;
|
||||
// Reset aggregation counters
|
||||
comm->asyncOpCount = 0;
|
||||
comm->asyncTotalSize = 0;
|
||||
// Dequeue p2p lists
|
||||
if (comm->p2pSendCount > 0 || comm->p2pRecvCount > 0) {
|
||||
for (int peer=0; peer<comm->nRanks; peer++) {
|
||||
if (comm->p2pSends[peer]) comm->p2pSends[peer]->recycle();
|
||||
if (comm->p2pRecvs[peer]) comm->p2pRecvs[peer]->recycle();
|
||||
}
|
||||
comm->p2pSendCount = comm->p2pRecvCount = 0;
|
||||
}
|
||||
ncclLaunchReset(comm);
|
||||
}
|
||||
}
|
||||
}
|
||||
end:
|
||||
ncclGroupError = ncclSuccess;
|
||||
ncclGroupIndex = 0;
|
||||
ncclGroupCommHead = nullptr;
|
||||
ncclGroupCommPreconnectHead = nullptr;
|
||||
CUDACHECK(hipSetDevice(savedDev)); // do other clean-ups first before calling hipSetDevice, because this call can fail too
|
||||
if (graphs) free(graphs);
|
||||
return ret;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user