Use separate threads for send and receive

[ROCm/rccl commit: b90735c935]

projects/rccl/src/collectives/device/all_to_all.h

@@ -48,16 +48,40 @@ __device__ void ncclAllToAllKernel(struct CollectiveArgs* args) {
           ReduceOrCopyMulti<UNROLL, FUNC, T, 1, 1, 1, 1>(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<UNROLL, ALLTOALL_CHUNKSTEPS/ALLTOALL_SLICESTEPS, ALLTOALL_SLICESTEPS, T, 1, 1, 0, FUNC>
-          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<UNROLL, ALLTOALL_CHUNKSTEPS/ALLTOALL_SLICESTEPS, ALLTOALL_SLICESTEPS, T, 2, 1, 0, FUNC>
+          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<UNROLL, ALLTOALL_CHUNKSTEPS/ALLTOALL_SLICESTEPS, ALLTOALL_SLICESTEPS, T, 1, 2, 0, FUNC>
+          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);
+        }
       }
     }
   }

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; \

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 <typename T, int NRECV>
 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();
 

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;