From 964c4c20617448a7b1bfe1d0ece82390b8698a94 Mon Sep 17 00:00:00 2001 From: Wenkai Du Date: Mon, 29 Jun 2020 08:47:46 -0700 Subject: [PATCH] Merge sendrecv kernel from NCCL 2.7.3 This commit was cherry-picked and modified from https://github.com/NVIDIA/nccl/commit/5949d96f36d050e59d05872f8bbffd2549318e95 --- src/collectives/device/sendrecv.h | 76 +++++++++++++++++++------------ src/include/collectives.h | 2 +- src/proxy.cc | 4 +- 3 files changed, 51 insertions(+), 31 deletions(-) diff --git a/src/collectives/device/sendrecv.h b/src/collectives/device/sendrecv.h index 1331244b15..60eebc6960 100644 --- a/src/collectives/device/sendrecv.h +++ b/src/collectives/device/sendrecv.h @@ -22,8 +22,15 @@ __device__ void ncclSendRecvKernel(struct CollectiveArgs* args) { if (args->p2p.delta == 0) { if (tid < nthreads && sendbuff != recvbuff) { - // local copy - ReduceOrCopyMulti(tid, nthreads, 1, &sendbuff, 1, &recvbuff, args->p2p.sendCount); + // local copy : ReduceOrCopyMulti takes an int as number of elements, + // so we split it in blocks of 1G elements. + int blockSize = 1<<30; + for (size_t offset=0; offsetp2p.sendCount; offset += blockSize) { + size_t remaining = args->p2p.sendCount - offset; + if (remaining < blockSize) blockSize = remaining; + ReduceOrCopyMulti(tid, nthreads, 1, &sendbuff, 1, &recvbuff, blockSize); + sendbuff += blockSize; recvbuff += blockSize; + } } return; } @@ -31,40 +38,53 @@ __device__ void ncclSendRecvKernel(struct CollectiveArgs* args) { struct ncclDevComm* comm = args->comm; struct ncclChannel* channel = comm->channels+blockIdx.x; - const ssize_t sendSize = args->p2p.sendCount; - const ssize_t recvSize = args->p2p.recvCount; - const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS); - const int chunkSize = stepSize; - int peerRecv = recvSize >= 0 ? (comm->rank-(int)args->p2p.delta+comm->nRanks)%comm->nRanks : -1; - int peerSend = sendSize >= 0 ? (comm->rank+(int)args->p2p.delta)%comm->nRanks : -1; + const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE]/(sizeof(T)*NCCL_STEPS)/SENDRECV_SLICEFACTOR; - ncclPrimitives - prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm, args->opCount); - - int maxSize = sendSize-chunkSize>recvSize ? sendSize-chunkSize : recvSize; - - if (sendSize >= 0) { - int realChunkSize = min(chunkSize, sendSize); - ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T)); - int nelem = min(realChunkSize, sendSize); - prims.send(sendbuff, nelem); + int nthreadsSplit; + const int64_t sendCount = static_cast(args->p2p.sendCount); + const int64_t recvCount = static_cast(args->p2p.recvCount); + if (sendCount >= 0 && recvCount >= 0) + nthreadsSplit = nthreads/2; + else { + if (sendCount >= 0) nthreadsSplit = nthreads; + else nthreadsSplit = 0; } + // We set NRECV or NSEND to 2 to use different barriers in primitives for the send threads and + // receive threads, but then we define all peers to -1 since sender threads don't receive and + // receive threads don't send. + int peerNone[2] = {-1,-1}; - for (ssize_t gridOffset = 0; gridOffset < maxSize; gridOffset += chunkSize) { - if (gridOffset+chunkSize < sendSize) { - int realChunkSize = min(chunkSize, sendSize-gridOffset-chunkSize); + if (tid < nthreadsSplit ) { + const ssize_t sendSize = args->p2p.sendCount; + if (sendSize < 0) return; + + int peer = (comm->rank+(int)args->p2p.delta)%comm->nRanks; + ncclPrimitives + prims(tid, nthreadsSplit, peerNone, &peer, recvbuff, stepSize*SENDRECV_SLICEFACTOR, channel, comm, args->opCount); + + if (sendSize == 0) { + prims.send(sendbuff, 0); + } else for (ssize_t offset = 0; offset < sendSize; offset += stepSize) { + int realChunkSize = min(stepSize, sendSize-offset); ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T)); - ssize_t offset = gridOffset + chunkSize; int nelem = min(realChunkSize, sendSize-offset); - prims.send(sendbuff+offset, nelem); + prims.directSend(sendbuff+offset, offset, nelem); } - if (gridOffset < recvSize) { - int realChunkSize = min(chunkSize, recvSize-gridOffset); + } else { + const ssize_t recvSize = args->p2p.recvCount; + if (recvSize < 0) return; + + int peer = (comm->rank-(int)args->p2p.delta+comm->nRanks)%comm->nRanks; + ncclPrimitives + prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &peer, peerNone, recvbuff, stepSize*SENDRECV_SLICEFACTOR, channel, comm, args->opCount); + + if (recvSize == 0) { + prims.recv(recvbuff, 0); + } else for (ssize_t offset = 0; offset < recvSize; offset += stepSize) { + int realChunkSize = min(stepSize, recvSize-offset); ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T)); - ssize_t offset = gridOffset; int nelem = min(realChunkSize, recvSize-offset); - prims.recv(recvbuff+offset, nelem); + prims.directRecv(recvbuff+offset, offset, nelem); } } - if (recvSize == 0) prims.recv(recvbuff,0); } diff --git a/src/include/collectives.h b/src/include/collectives.h index 7f4b30c45d..56c268ba19 100644 --- a/src/include/collectives.h +++ b/src/include/collectives.h @@ -82,7 +82,7 @@ DECL_ALL_COLLS #define BROADCAST_CHUNKSTEPS 1 #define REDUCE_SLICESTEPS 1 #define REDUCE_CHUNKSTEPS 1 -#define SENDRECV_SLICEFACTOR 4 +#define SENDRECV_SLICEFACTOR 1 #define GATHER_SLICESTEPS 4 #define GATHER_CHUNKSTEPS 4 #define SCATTER_SLICESTEPS 4 diff --git a/src/proxy.cc b/src/proxy.cc index 9985953938..636cb4d9a9 100644 --- a/src/proxy.cc +++ b/src/proxy.cc @@ -152,13 +152,13 @@ ncclResult_t ncclProxySaveP2p(struct ncclInfo* info, struct ncclChannel* channel args.dtype = info->datatype; if (info->delta > 0 && info->sendbytes >= 0) { int peersend = (info->comm->rank+info->delta)%info->comm->nRanks; - args.nsteps = DIVUP(info->sendbytes, info->comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS); + args.nsteps = DIVUP(info->sendbytes, info->comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/SENDRECV_SLICEFACTOR); if (args.nsteps == 0) args.nsteps = 1; NCCLCHECK(SaveProxy(peersend, &args)); } if (info->delta > 0 && info->recvbytes >= 0) { int peerrecv = (info->comm->nRanks+info->comm->rank-info->delta)%info->comm->nRanks; - args.nsteps = DIVUP(info->recvbytes, info->comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS); + args.nsteps = DIVUP(info->recvbytes, info->comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/SENDRECV_SLICEFACTOR); if (args.nsteps == 0) args.nsteps = 1; NCCLCHECK(SaveProxy(peerrecv, &args)); }