diff --git a/src/device/common.h b/src/device/common.h index a1e5a8958b..a9ecd732e6 100644 --- a/src/device/common.h +++ b/src/device/common.h @@ -114,7 +114,6 @@ struct ncclShmemGroup { void* srcs[NCCL_MAX_ARITY+1]; void* dsts[NCCL_MAX_ARITY+1]; uint64_t barrier; - uint64_t barrier_next[NCCL_MAX_GROUPS]; union { unpackGroupShmem unpack; } devicePlugin; @@ -473,8 +472,6 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a ncclShmem.groups[tid-WARP_SIZE].barrier = 0; break; case 2: - if (tid < 2*WARP_SIZE + NCCL_MAX_GROUPS*NCCL_MAX_GROUPS) - ncclShmem.groups[(tid-2*WARP_SIZE)/NCCL_MAX_GROUPS].barrier_next[(tid-2*WARP_SIZE)%NCCL_MAX_GROUPS] = 0; break; case 3: /* set abort flag to 0 */ @@ -551,6 +548,15 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a if (ncclShmem.nextBatchIx == -1) break; int batchIx = ncclShmem.nextBatchIx; __synclds(); + switch (tid/WARP_SIZE) { + case 1: + if (tid < WARP_SIZE + NCCL_MAX_GROUPS) + ncclShmem.groups[tid-WARP_SIZE].barrier = 0; + break; + default: + break; + } + __synclds(); loadWorkBatchToShmem(tid%WARP_SIZE, tn, args, batchIx); // Check whether the last operation was aborted and make sure all threads exit diff --git a/src/device/primitives.h b/src/device/primitives.h index 0ecd48ff42..10c8710c61 100644 --- a/src/device/primitives.h +++ b/src/device/primitives.h @@ -28,12 +28,12 @@ const int w = threadIdx.x/WARP_SIZE; \ const int wid = threadIdx.x%WARP_SIZE; \ if (wid == 0) { \ - barrier_next[w] += nthreads/WARP_SIZE; \ + barrier_next += nthreads/WARP_SIZE; \ __hip_atomic_fetch_add(barriers, 1, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_WORKGROUP); \ int spins = 0; \ int rate_limit = 50; \ __THREAD_FENCE; \ - while (__hip_atomic_load(barriers, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) < barrier_next[w]) { \ + while (__hip_atomic_load(barriers, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP) < barrier_next) { \ spins++; \ if (spins == NCCL_SPINS_BEFORE_CHECK_ABORT) { \ if (__atomic_load_n(ncclShmem.comm.abortFlag, __ATOMIC_SEQ_CST)) { \ @@ -44,7 +44,7 @@ } \ if (spins == 0 && rate_limit > 0) { \ rate_limit --; \ - traceData(__LINE__, threadIdx.x, __hip_atomic_load(barriers, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP), barrier_next[w]); \ + traceData(__LINE__, threadIdx.x, __hip_atomic_load(barriers, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_WORKGROUP), barrier_next); \ } \ __builtin_amdgcn_s_sleep(1); \ } \ diff --git a/src/device/prims_ll.h b/src/device/prims_ll.h index dec5bf80dc..d561a69971 100644 --- a/src/device/prims_ll.h +++ b/src/device/prims_ll.h @@ -66,7 +66,7 @@ private: inline __device__ uint32_t sendFlag(int i) { return NCCL_LL_FLAG(sendStep[i]+1); } uint64_t* barriers; - uint64_t* barrier_next; + uint64_t barrier_next = 0; inline __device__ void barrier() { #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) @@ -622,7 +622,6 @@ private: stepLines(ncclShmem.comm.buffSizes[NCCL_PROTO_LL]/NCCL_STEPS/sizeof(ncclLLFifoLine)) { auto *channel = &ncclShmem.channel; barriers = &ncclShmem.groups[group].barrier; - barrier_next = ncclShmem.groups[group].barrier_next; // If we are going to support oneshot collNet + LL, then we would need to add connector index here int nrecv=0, nsend=0; // We compare with Fan::MaxRecv here because this->MaxRecv is always at least 1 diff --git a/src/device/prims_ll128.h b/src/device/prims_ll128.h index de29528e8a..771ee82a25 100644 --- a/src/device/prims_ll128.h +++ b/src/device/prims_ll128.h @@ -62,7 +62,7 @@ class Primitives: inline __device__ uint64_t sendFlag(int i) { return sendStep[i]+1; } uint64_t* barriers; - uint64_t* barrier_next; + uint64_t barrier_next = 0; #if defined(ENABLE_NPKIT) public: @@ -508,7 +508,6 @@ public: stepSize(ncclShmem.comm.buffSizes[NCCL_PROTO_LL128]/NCCL_STEPS/sizeof(uint64_t)) { auto *channel = &ncclShmem.channel; barriers = &ncclShmem.groups[group].barrier; - barrier_next = ncclShmem.groups[group].barrier_next; int nrecv=0, nsend=0; while (nrecv < MaxRecv && recvPeers[nrecv] >= 0) { loadRecvConn(&channel->peers[recvPeers[nrecv]]->recv[connIndexRecv], nrecv); diff --git a/src/device/prims_simple.h b/src/device/prims_simple.h index 1887ea154b..d8083bea8f 100644 --- a/src/device/prims_simple.h +++ b/src/device/prims_simple.h @@ -54,7 +54,7 @@ class Primitives< void* netDeviceHandle; uint32_t* next_hdp_reg; uint64_t* barriers; - uint64_t* barrier_next; + uint64_t barrier_next = 0; int repeat; #if defined(ENABLE_NPKIT) @@ -671,7 +671,6 @@ private: // For send operations, we need an extra warp to overlap the threadfence and the copy barriers = &ncclShmem.groups[group].barrier; - barrier_next = ncclShmem.groups[group].barrier_next; this->nworkers = nthreads; int nrecv=0, nsend=0;