From 3c674c9c1fc507f7d2362daaae65d7d56820cc72 Mon Sep 17 00:00:00 2001 From: Wenkai Du Date: Mon, 29 Jun 2020 08:47:15 -0700 Subject: [PATCH] Use separate threads for send and receive [ROCm/rccl commit: b90735c9350a24f1738b1460d035ba834a105a3a] --- .../rccl/src/collectives/device/all_to_all.h | 42 +++++++++++++++---- projects/rccl/src/collectives/device/common.h | 10 ++++- .../rccl/src/collectives/device/primitives.h | 19 ++++++++- projects/rccl/src/include/devcomm.h | 5 +++ 4 files changed, 65 insertions(+), 11 deletions(-) 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/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;