diff --git a/projects/rccl/src/collectives/device/all_to_all.h b/projects/rccl/src/collectives/device/all_to_all.h index e21579f127..b85338635c 100644 --- a/projects/rccl/src/collectives/device/all_to_all.h +++ b/projects/rccl/src/collectives/device/all_to_all.h @@ -48,16 +48,40 @@ __device__ void ncclAllToAllKernel(struct CollectiveArgs* args) { ReduceOrCopyMulti(tid, nthreads, 1, &sendbuff, 1, &recvbuff, nelem); } } - else { - int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks; - int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks; - ncclPrimitives - prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm, args->opCount); + } + } - ssize_t send_offset = chunkOffset + peerSend*size; - ssize_t recv_offset = chunkOffset + peerRecv*size; - prims.send(thisInput+send_offset, nelem); - prims.recv(thisOutput+recv_offset, nelem); + for (int i = 0; i < peersPerChan; i++) { + if ((peersPerChan == 1 && blockIdx.x >= (nChannels/nranks)*nranks) || + (peersPerChan > 1 && blockIdx.x*peersPerChan+i >= nranks)) + continue; + if ((blockIdx.x*peersPerChan+i)%nranks != 0) { + int nthreadsSplit = nthreads/2; + int peerNone[2] = {-1,-1}; + if (tid < nthreadsSplit ) { + int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks; + ncclPrimitives + prims(tid, nthreadsSplit, peerNone, &peerSend, NULL, stepSize, channel, comm, args->opCount); + for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) { + int realChunkSize = min(chunkSize, DIVUP(size-gridOffset, (peersPerChan == 1 ? (nChannels/nranks) : 1))); + ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T)); + ssize_t chunkOffset = gridOffset + (peersPerChan == 1 ? (bid/nranks)*realChunkSize : 0); + int nelem = min(realChunkSize, size-chunkOffset); + ssize_t send_offset = chunkOffset + peerSend*size; + prims.send(thisInput+send_offset, nelem); + } + } else { + int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks; + ncclPrimitives + prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &peerRecv, peerNone, NULL, stepSize, channel, comm, args->opCount); + for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) { + int realChunkSize = min(chunkSize, DIVUP(size-gridOffset, (peersPerChan == 1 ? (nChannels/nranks) : 1))); + ALIGN_SIZE(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T)); + ssize_t chunkOffset = gridOffset + (peersPerChan == 1 ? (bid/nranks)*realChunkSize : 0); + int nelem = min(realChunkSize, size-chunkOffset); + ssize_t recv_offset = chunkOffset + peerRecv*size; + prims.recv(thisOutput+recv_offset, nelem); + } } } } diff --git a/projects/rccl/src/collectives/device/common.h b/projects/rccl/src/collectives/device/common.h index da67970e45..76bcb68d5c 100644 --- a/projects/rccl/src/collectives/device/common.h +++ b/projects/rccl/src/collectives/device/common.h @@ -246,11 +246,19 @@ __global__ void NCCL_KERN_NAME(coll, op, dtype)(struct ncclDevComm* comm) { \ ALLOCATE_SHMEM; \ __shared__ struct ncclColl localColl; \ __shared__ uint32_t abortCount; \ + __shared__ uint64_t barrier[MAXBARRIERS]; \ + __shared__ uint64_t barrier_next[MAXBARRIERS*MAXWARPS]; \ if (tid == 0) abortCount = 0; \ __syncthreads(); \ \ struct ncclChannel* channel = comm->channels+bid; \ - channel->sync = sync; \ + if (tid == 0) { \ + channel->sync = sync; \ + channel->barrier = barrier; \ + channel->barrier_next = barrier_next; \ + for (auto i = 0; i < MAXBARRIERS; i++) barrier[i] = 0; \ + for (auto i = 0; i < MAXBARRIERS*MAXWARPS; i++) barrier_next[i] = 0; \ + } \ if (!load_coll(&localColl, channel->collectives+channel->collFifoHead, tid, comm, &abortCount)) { \ if (tid == 0) traceAbort(-1); \ return; \ diff --git a/projects/rccl/src/collectives/device/primitives.h b/projects/rccl/src/collectives/device/primitives.h index ce5c83c3e9..b5c6d24bfb 100644 --- a/projects/rccl/src/collectives/device/primitives.h +++ b/projects/rccl/src/collectives/device/primitives.h @@ -32,6 +32,13 @@ } \ } while (0) +#define barrier_by_id(id) do { \ + const int w = threadIdx.x/WARP_SIZE; \ + barrier_next[id*MAXWARPS+w] += nthreads/WARP_SIZE; \ + __atomic_fetch_add(barriers+id, 1, __ATOMIC_SEQ_CST); \ + while (LOAD(barriers+id) < barrier_next[id*MAXWARPS+w]) /* spin */; \ +} while (0) + template class ncclPrimitivesRecvData { public: @@ -90,9 +97,17 @@ class ncclPrimitives { inline __device__ const T* recvPtr(int i) { return ((const T*)r.recvBuff[i])+recvOffset(i); } inline __device__ T* sendPtr(int i) { return ((T*)s.sendBuff[i])+sendOffset(i); } + uint64_t* barriers; + uint64_t* barrier_next; + inline __device__ void barrier() { #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__) - __syncthreads(); + if (nthreads == gridDim.x) { + __syncthreads(); + } else if (wid == 0) { + if (NRECV < NSEND) barrier_by_id(0); + else barrier_by_id(1); + } #else asm volatile ("bar.sync 1, %0;" :: "r"(nthreads+WARP_SIZE)); #endif @@ -359,6 +374,8 @@ inline __device__ int directSendInc(int i, int directInc, int sliceInc) { } inline __device__ void init(int* recvPeers, int* sendPeers, struct ncclChannel* channel) { + barriers = channel->barrier; + barrier_next = channel->barrier_next; // Make sure step is updated before we read it. barrier(); diff --git a/projects/rccl/src/collectives/device/sendrecv.h b/projects/rccl/src/collectives/device/sendrecv.h index 1331244b15..60eebc6960 100644 --- a/projects/rccl/src/collectives/device/sendrecv.h +++ b/projects/rccl/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/projects/rccl/src/include/collectives.h b/projects/rccl/src/include/collectives.h index 7f4b30c45d..56c268ba19 100644 --- a/projects/rccl/src/include/collectives.h +++ b/projects/rccl/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/projects/rccl/src/include/devcomm.h b/projects/rccl/src/include/devcomm.h index e5c5f43727..4fd7b46d94 100644 --- a/projects/rccl/src/include/devcomm.h +++ b/projects/rccl/src/include/devcomm.h @@ -89,6 +89,9 @@ static_assert(NCCL_LL_CLEAN_MASK % NCCL_STEPS == 0, "Invalid NCCL_LL_CLEAN_MASK #define NCCL_DIRECT_GPU 0x01 #define NCCL_DIRECT_NIC 0x10 +#define MAXBARRIERS 2 +#define MAXWARPS (NCCL_MAX_NTHREADS/WARP_SIZE) + struct ncclConnInfo { // Regular comm mechanism char *buffs[NCCL_NUM_PROTOCOLS]; // Local for recv, remote for send @@ -220,6 +223,8 @@ struct ncclChannel { int collFifoTail; // Only used by CPU uint32_t* sync; + uint64_t* barrier; + uint64_t* barrier_next; #ifdef ENABLE_PROFILING struct timeval tvs; uint64_t sizes; diff --git a/projects/rccl/src/proxy.cc b/projects/rccl/src/proxy.cc index 9985953938..636cb4d9a9 100644 --- a/projects/rccl/src/proxy.cc +++ b/projects/rccl/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)); } diff --git a/projects/rccl/test/test_AllGather.cpp b/projects/rccl/test/test_AllGather.cpp index 558e927ffb..1e7b9441c0 100644 --- a/projects/rccl/test/test_AllGather.cpp +++ b/projects/rccl/test/test_AllGather.cpp @@ -105,9 +105,9 @@ namespace CorrectnessTests ncclFloat64, ncclBfloat16), // Number of elements - testing::Values(3072, 3145728), + testing::Values(2520, 3026520), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false, true), testing::Values(""))); diff --git a/projects/rccl/test/test_AllReduce.cpp b/projects/rccl/test/test_AllReduce.cpp index 41f1af5a8b..1452ccdfef 100644 --- a/projects/rccl/test/test_AllReduce.cpp +++ b/projects/rccl/test/test_AllReduce.cpp @@ -55,7 +55,7 @@ namespace CorrectnessTests // Number of elements testing::Values(1024, 1048576), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false, true), testing::Values(""))); diff --git a/projects/rccl/test/test_AllToAll.cpp b/projects/rccl/test/test_AllToAll.cpp index 2c1d29da1b..05d5e40475 100644 --- a/projects/rccl/test/test_AllToAll.cpp +++ b/projects/rccl/test/test_AllToAll.cpp @@ -59,7 +59,7 @@ namespace CorrectnessTests // Number of elements testing::Values(1024, 1048576), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false), testing::Values("RCCL_ALLTOALL_KERNEL_DISABLE=0", "RCCL_ALLTOALL_KERNEL_DISABLE=1"))); diff --git a/projects/rccl/test/test_Broadcast.cpp b/projects/rccl/test/test_Broadcast.cpp index 85b087a163..52e4403c75 100644 --- a/projects/rccl/test/test_Broadcast.cpp +++ b/projects/rccl/test/test_Broadcast.cpp @@ -63,7 +63,7 @@ namespace CorrectnessTests // Number of elements testing::Values(1024, 1048576), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false, true), testing::Values(""))); diff --git a/projects/rccl/test/test_CombinedCalls.cpp b/projects/rccl/test/test_CombinedCalls.cpp index 35d61e82f3..1fe86a4125 100644 --- a/projects/rccl/test/test_CombinedCalls.cpp +++ b/projects/rccl/test/test_CombinedCalls.cpp @@ -91,9 +91,9 @@ namespace CorrectnessTests ncclFloat64, ncclBfloat16), // Number of elements - testing::Values(3072, 3145728), + testing::Values(2520, 3026520), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false, true), testing::Values(""))); diff --git a/projects/rccl/test/test_Gather.cpp b/projects/rccl/test/test_Gather.cpp index 4324434e9e..95dc35c461 100644 --- a/projects/rccl/test/test_Gather.cpp +++ b/projects/rccl/test/test_Gather.cpp @@ -63,7 +63,7 @@ namespace CorrectnessTests // Number of elements testing::Values(1024, 1048576), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false), testing::Values("RCCL_ALLTOALL_KERNEL_DISABLE=0", "RCCL_ALLTOALL_KERNEL_DISABLE=1"))); diff --git a/projects/rccl/test/test_GroupCalls.cpp b/projects/rccl/test/test_GroupCalls.cpp index cc2bd7dcd2..2f652e96e4 100644 --- a/projects/rccl/test/test_GroupCalls.cpp +++ b/projects/rccl/test/test_GroupCalls.cpp @@ -111,9 +111,9 @@ namespace CorrectnessTests ncclFloat64, ncclBfloat16), // Number of elements - testing::Values(3072, 3145728), + testing::Values(2520, 3026520), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false, true), testing::Values(""))); diff --git a/projects/rccl/test/test_Reduce.cpp b/projects/rccl/test/test_Reduce.cpp index c181ec90ac..59dac111f9 100644 --- a/projects/rccl/test/test_Reduce.cpp +++ b/projects/rccl/test/test_Reduce.cpp @@ -63,7 +63,7 @@ namespace CorrectnessTests // Number of elements testing::Values(1024, 1048576), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false, true), testing::Values(""))); diff --git a/projects/rccl/test/test_ReduceScatter.cpp b/projects/rccl/test/test_ReduceScatter.cpp index 014ecd0ba4..ab4dc08860 100644 --- a/projects/rccl/test/test_ReduceScatter.cpp +++ b/projects/rccl/test/test_ReduceScatter.cpp @@ -59,9 +59,9 @@ namespace CorrectnessTests ncclFloat64, ncclBfloat16), // Number of elements - testing::Values(3072, 3145728), + testing::Values(2520, 3026520), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false, true), testing::Values(""))); diff --git a/projects/rccl/test/test_Scatter.cpp b/projects/rccl/test/test_Scatter.cpp index dca86b1ff3..80ac6e534c 100644 --- a/projects/rccl/test/test_Scatter.cpp +++ b/projects/rccl/test/test_Scatter.cpp @@ -63,7 +63,7 @@ namespace CorrectnessTests // Number of elements testing::Values(1024, 1048576), // Number of devices - testing::Values(2,3,4), + testing::Values(2,3,4,5,6,7,8), // In-place or not testing::Values(false), testing::Values("RCCL_ALLTOALL_KERNEL_DISABLE=0", "RCCL_ALLTOALL_KERNEL_DISABLE=1")));