AI/rocm-systems
2
0
Форкнуть 0
Код Задачи Запросы на слияние Действия Пакеты Проекты Релизы Вики Активность

Fix P2P scheduling (#560)

Просмотр исходного кода
Этот коммит содержится в:
Wenkai Du
2022-06-06 13:32:28 -07:00
коммит произвёл GitHub
родитель 7a6c6927ae
Коммит 5cb2aca3d9
4 изменённых файлов: 95 добавлений и 68 удалений
Скачать Patch файл Скачать Diff файл Показать все файлы Свернуть все файлы
src/collectives/device/prims_simple.h
+13 -3
Просмотреть файл
@@ -47,6 +47,7 @@ class Primitives<
uint64_t* barriers;
uint64_t* barrier_next;
const uint64_t opCount;
uint32_t* next_hdp_reg;
// Don't use barrier 0 as it's used by the final sync
inline __device__ void barrier() {
@@ -96,6 +97,7 @@ class Primitives<
connStepCache = LOAD(connStepPtr);
if (checkAbort(spins)) break;
//if (spins == 0) printf("r=%d b=%d t=%d SPUN OUT got=%d want=%d\n", ncclShmem->comm.rank, blockIdx.x, threadIdx.x, int(connStepCache + (isSendNotRecv ? NCCL_STEPS : 0)), int(step+StepPerSlice));
if (spins == 0) traceData(__LINE__, threadIdx.x, int(connStepCache + (isSendNotRecv ? NCCL_STEPS : 0)), int(step+StepPerSlice));
}
__asm__ __volatile__("s_wakeup");
}
@@ -134,6 +136,9 @@ class Primitives<
template<int Recv, int Send>
inline __device__ void postPeer() {
if ((flags & Send*RolePostSend) && next_hdp_reg)
atomicExch_system(next_hdp_reg, 0x1);
if (flags & (Recv*RolePostRecv | Send*RolePostSend)) {
step += StepPerSlice;
atomicExch_system((unsigned long long *)connStepPtr, step);
@@ -225,8 +230,10 @@ class Primitives<
sliceSize);
}
barrier(); // This barrier has a counterpart in following loop
//if (Send && (flags & RolePostSend) && index == 0) __threadfence_system();
__syncwarp();
#if defined(__gfx1030__)
if (Send && (flags & RolePostSend) && index == 0) __threadfence_system();
#endif
__syncwarp();
postPeer<Recv, Send>();
offset += sliceSize;
slice += 1;
@@ -245,7 +252,9 @@ class Primitives<
waitPeer<DirectRecv, DirectSend, Recv, Send, Src, Dst>(0, 0, 0, 0);
}
barrier(); // Has couterpart in preceding worker-only loop.
//if (Send && (flags & RolePostSend) && sliceSize > 0 && index == 0) __threadfence_system();
#if defined(__gfx1030__)
if (Send && (flags & RolePostSend) && sliceSize > 0 && index == 0) __threadfence_system();
#endif
__syncwarp();
postPeer<Recv, Send>();
offset += sliceSize;
@@ -371,6 +380,7 @@ class Primitives<
step = roundUp(step, SlicePerChunk*StepPerSlice);
if (flags & RolePostSend) {
connStepPtr = conn->tail;
next_hdp_reg = conn->next_hdp_reg;
}
if (flags & RoleWaitSend) {
ncclShmem->groups[group].sendConns[index] = conn; // WaitSend role saves since that's who needs it in setDataPtrs()
src/enqueue.cc
+5 -5
Просмотреть файл
@@ -956,7 +956,7 @@ static int getSegment(enum ncclWorkElemType type, enum ncclWorkElemSubType subTy
if (type == ncclWorkTypeP2p) { // P2P
int start = subType == ncclWorkSubTypeRecv ? 0 : 1;
for (int s=start; s<NCCL_MAX_WORK_ELEMENTS_P2P && s<NCCL_MAX_NTHREADS/comm->WarpSize; s+=2) {
if (work->p2pElems[s].peer == -1) return s;
if (work->p2pElems[s].peer == -2) return s;
// Do not aggregate multiple sends to the same peer (or receives from the same peer)
if (work->p2pElems[s].peer == peer) return -1;
}
@@ -1072,14 +1072,14 @@ ncclResult_t ncclEnqueueP2pKernel(struct ncclComm* comm, struct ncclQueueElem* e
if (segment == -1) {
NCCLCHECK(getNextOp(channel, &w, NULL));
segment = workElem->subType == ncclWorkSubTypeRecv ? 0 : 1;
// Initialize work as P2P, set peer=-1 to designate the p2p elem is not used.
// Initialize work as P2P, set peer=-2 to designate the p2p elem is not used.
w->header.type = ncclWorkTypeP2p;
for (int i=0; i<NCCL_MAX_WORK_ELEMENTS_P2P && i<NCCL_MAX_NTHREADS/comm->WarpSize; i++) w->p2pElems[i].peer = -1;
for (int i=0; i<NCCL_MAX_WORK_ELEMENTS_P2P && i<NCCL_MAX_NTHREADS/comm->WarpSize; i++) w->p2pElems[i].peer = -2;
}
//printf("%s to %d -> Channel %d OpCount %ld Segment %d\n", workElem->subType == ncclWorkSubTypeRecv ? "Recv" : "Send", proxyOp->root, channel->id, channel->workFifoTail-1, segment);
//INFO(NCCL_COLL, "%s to %d -> Channel %d OpCount %ld Segment %d", workElem->subType == ncclWorkSubTypeRecv ? "Recv" : "Send", workElem->peer, channel->id, channel->workFifoTail-1, segment);
// store work element into FIFO
NCCLCHECK(ncclProxySaveP2p(comm, proxyOp));
if (workElem->peer != -1) NCCLCHECK(ncclProxySaveP2p(comm, proxyOp));
NCCLCHECK(enqueueSegOp(ncclWorkTypeP2p, &eqElem->work, w, segment, &eqElem->buffRegInfo, channel, comm));
return ncclSuccess;
}
src/group.cc
+75 -60
Просмотреть файл
@@ -105,12 +105,10 @@ ncclResult_t ncclGroupStart() {
return ncclSuccess;
}
static ncclResult_t scheduleSend(struct ncclComm* comm, int peer, int chunk, size_t count, void* buff, uint64_t opCount, uint16_t connIndex) {
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 };
int channelId;
NCCLCHECK(ncclChannelCompute(comm, peer, chunk%comm->p2pnChannelsPerPeer, ncclFuncSend, &channelId));
info.channelId = channelId;
info.opCount = opCount;
info.connIndex = connIndex;
@@ -118,12 +116,10 @@ static ncclResult_t scheduleSend(struct ncclComm* comm, int peer, int chunk, siz
return ncclSuccess;
}
static ncclResult_t scheduleRecv(struct ncclComm* comm, int peer, int chunk, size_t count, void* buff, uint64_t opCount, uint16_t connIndex) {
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 };
int channelId;
NCCLCHECK(ncclChannelCompute(comm, peer, chunk%comm->p2pnChannelsPerPeer, ncclFuncRecv, &channelId));
info.channelId = channelId;
info.opCount = opCount;
info.connIndex = connIndex;
@@ -244,8 +240,9 @@ ncclResult_t ncclGroupEnd() {
struct ncclAsyncArgs* args = ncclGroupArgs+i;
if (args->funcType == ASYNC_FUNC_COLL) {
struct ncclComm* comm = args->coll.comm;
int rank = comm->rank;
int nRanks = comm->nRanks;
int node = comm->node;
int nNodes = comm->nNodes;
int localRank = comm->localRank;
// Compute how much to split operations
// Natural step size matching buffer steps.
@@ -254,71 +251,89 @@ ncclResult_t ncclGroupEnd() {
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;
//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 (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<=nRanks/4; d++) {
int deltas[4] = { d, (nRanks-d)%nRanks, nRanks/2-d, (nRanks-(nRanks/2-d))%nRanks };
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 recvPeer = (rank+nRanks-delta)%nRanks;
uint32_t sendPeer = (rank+delta)%nRanks;
struct ncclP2Pinfo* recv = comm->p2pRecvs[recvPeer] ? comm->p2pRecvs[recvPeer]->getNext() : NULL;
struct ncclP2Pinfo* send = 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();
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);
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;
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 {
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.
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 (recv) {
NCCLCHECKGOTO(scheduleRecv(comm, recvPeer, chunk, recvbytes, ((char*)(recv->buff))+recvOffset, recv->opCount, recvIdx), ret, group_cleanup);
}
if (send) {
NCCLCHECKGOTO(scheduleSend(comm, sendPeer, chunk, sendbytes, ((char*)(send->buff))+sendOffset, send->opCount, sendIdx), ret, group_cleanup);
}
recvOffset += recvChunkSize;
sendOffset += sendChunkSize;
chunk++;
} while (sendRemaining || recvRemaining);
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);
}
}
index++;
if (index == 1 && deltas[1] == deltas[0]) index++;
src/proxy.cc
+2
Просмотреть файл
@@ -412,6 +412,8 @@ ncclResult_t ncclProxyComputeP2p(struct ncclInfo* info, struct ncclProxyOp* op)
info->chunkSize = stepSize;
op->root = info->root;
op->nbytes = info->count;
if (info->root == -1) return ncclSuccess;
struct ncclPeer* peer = channel->peers + op->root;
if (info->coll == ncclFuncSend) {