From abf265a9115d7abf5cbc506a15aec87005f20e28 Mon Sep 17 00:00:00 2001 From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com> Date: Thu, 4 Jan 2024 08:18:48 -0800 Subject: [PATCH] Rework barriers and adjust scope of atomics (#1019) --- src/collectives/device/msccl_kernel_impl.h | 10 ++++--- src/collectives/device/prims_ll.h | 35 +++++----------------- src/collectives/device/prims_ll128.h | 6 ++-- src/collectives/device/prims_simple.h | 12 ++------ 4 files changed, 19 insertions(+), 44 deletions(-) diff --git a/src/collectives/device/msccl_kernel_impl.h b/src/collectives/device/msccl_kernel_impl.h index 1622af5c6c..236782d899 100644 --- a/src/collectives/device/msccl_kernel_impl.h +++ b/src/collectives/device/msccl_kernel_impl.h @@ -338,8 +338,9 @@ __device__ __forceinline__ void mscclRunInterpreter( if (tid == 0) { NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_SEND_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP()); } -#endif - prims.template send<1>(srcOffset, thisNelem); // LL.send is the only situation where there is no barrier at the end. + +#endif + prims.send(srcOffset, thisNelem); // LL.send is the only situation where there is no barrier at the end. #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_SEND_EXIT) if (tid == 0) { @@ -352,8 +353,9 @@ __device__ __forceinline__ void mscclRunInterpreter( if (tid == 0) { NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RECV_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP()); } -#endif - prims.template recv<1>(dstOffset, thisNelem); +#endif + prims.recv(dstOffset, thisNelem); + #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_RECV_EXIT) if (tid == 0) { NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RECV_EXIT, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP()); diff --git a/src/collectives/device/prims_ll.h b/src/collectives/device/prims_ll.h index e6b7fe343d..96286f6be9 100644 --- a/src/collectives/device/prims_ll.h +++ b/src/collectives/device/prims_ll.h @@ -89,13 +89,6 @@ private: #endif } - inline __device__ void msccl_barrier() { -#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__) - __builtin_amdgcn_s_barrier(); -#else -#endif - } - uint32_t abort = 0; inline __device__ int checkAbort(int &spins, int send) { @@ -107,7 +100,6 @@ private: return abort; } - template inline __device__ void waitSend(int nbytes) { #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_WAIT_SEND_ENTRY) if (tid == 0) { @@ -125,15 +117,11 @@ private: __asm__ __volatile__("s_wakeup"); if (sendConnFifoPtr) { int size = ((sendConnHead & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) ? stepLines*sizeof(union ncclLLFifoLine) : nbytes; - __atomic_store_n(sendConnFifoPtr+sendConnHead%NCCL_STEPS, (size), __ATOMIC_SEQ_CST); + __atomic_store_n(sendConnFifoPtr+sendConnHead%NCCL_STEPS, (size), __ATOMIC_RELAXED); } sendConnHead += 1; } - if (MSCCL) { - msccl_barrier(); - } else { - barrier(); - } + barrier(); #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_WAIT_SEND_EXIT) if (tid == 0) { NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_LL_WAIT_SEND_EXIT, nbytes, 0, NPKIT_GET_GPU_TIMESTAMP(), @@ -145,13 +133,8 @@ private: inline __device__ void incRecv(int i) { recvStep[i] += 1; } - template inline __device__ void postRecv() { - if (MSCCL) { - msccl_barrier(); - } else { - barrier(); - } + barrier(); if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1); } @@ -402,7 +385,7 @@ private: } } - template + template __device__ void LLGenericOp(intptr_t srcIx, intptr_t dstIx, int nelem, bool postOp) { constexpr int SRC = SrcBuf != -1 ? 1 : 0; constexpr int DST = DstBuf != -1 ? 1 : 0; @@ -411,7 +394,7 @@ private: // Always waitSend in case of cleanup nelem = nelem < 0 ? 0 : nelem; - if (SEND) waitSend(divUp(nelem, EltPerLine)*sizeof(ncclLLFifoLine)); + if (SEND) waitSend(divUp(nelem, EltPerLine)*sizeof(ncclLLFifoLine)); #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) if (tid == 0) { @@ -494,7 +477,7 @@ private: if (RECV) { for (int i=0; i < MaxRecv; i++) incRecv(i); - postRecv(); + postRecv(); } if (SEND) { for (int i=1; i < MaxSend && i < fan.nsend(); i++) @@ -656,7 +639,6 @@ private: userBufs[Output] += delta; } - template __device__ void send(intptr_t inpIx, int eltN) { #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_ENTRY) if (tid == 0) { @@ -664,7 +646,7 @@ private: ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx); } #endif - LLGenericOp<0, 1, Input, -1, MSCCL>(inpIx, -1, eltN, false); + LLGenericOp<0, 1, Input, -1>(inpIx, -1, eltN, false); #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_EXIT) if (tid == 0) { NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(), @@ -687,7 +669,6 @@ private: } #endif } - template __device__ void recv(intptr_t outIx, int eltN, bool postOp=false) { #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_ENTRY) if (tid == 0) { @@ -695,7 +676,7 @@ private: ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx); } #endif - LLGenericOp<1, 0, -1, Output, MSCCL>(-1, outIx, eltN, postOp); + LLGenericOp<1, 0, -1, Output>(-1, outIx, eltN, postOp); #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_EXIT) if (tid == 0) { NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(), diff --git a/src/collectives/device/prims_ll128.h b/src/collectives/device/prims_ll128.h index 561a490bd1..10bc6753b8 100644 --- a/src/collectives/device/prims_ll128.h +++ b/src/collectives/device/prims_ll128.h @@ -99,12 +99,12 @@ private: int spins = 0; while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) { __builtin_amdgcn_s_sleep(1); - sendConnHeadCache = atomicAdd_system((unsigned long long *)sendConnHeadPtr, 0); + sendConnHeadCache = __atomic_load_n(sendConnHeadPtr, __ATOMIC_RELAXED); if (checkAbort(spins, wid, 1)) break; } __asm__ __volatile__("s_wakeup"); if (sendConnFifoPtr) { - __atomic_store_n(sendConnFifoPtr+sendStep[wid]%NCCL_STEPS, nbytes, __ATOMIC_SEQ_CST); + __atomic_store_n(sendConnFifoPtr+sendStep[wid]%NCCL_STEPS, nbytes, __ATOMIC_RELAXED); } sendConnHead += 1; } @@ -546,14 +546,12 @@ public: userBufs[Output] += delta; } - template __device__ void send(intptr_t inpIx, int eltN) { return GenericOp<0, 1, Input, -1>(inpIx, -1, eltN, false); } __device__ void sendFromOutput(intptr_t outIx, int eltN) { return GenericOp<0, 1, Output, -1>(outIx, -1, eltN, false); } - template __device__ void recv(intptr_t outIx, int eltN, bool postOp=false) { return GenericOp<1, 0, -1, Output>(-1, outIx, eltN, postOp); } diff --git a/src/collectives/device/prims_simple.h b/src/collectives/device/prims_simple.h index 5337675fd2..2b175621a8 100644 --- a/src/collectives/device/prims_simple.h +++ b/src/collectives/device/prims_simple.h @@ -81,7 +81,7 @@ private: inline __device__ bool checkAbort(int &spins) { spins++; if (!(flags & Aborted) && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) { - if (atomicAdd_system((unsigned int *)ncclShmem.comm.abortFlag, 0)) { + if (__atomic_load_n(ncclShmem.comm.abortFlag, __ATOMIC_SEQ_CST)) { flags |= Aborted; ncclShmem.aborted = 1; } @@ -100,11 +100,7 @@ private: #endif // volatile is faster than acquire but not as correct. Make sure reduceCopy // loads data using volatile so it doesn't see stale data in L1. -#ifdef __GFX9__ - return atomicAdd((unsigned long long *)ptr, 0); -#else - return __atomic_load_n(ptr, __ATOMIC_SEQ_CST); -#endif + return __atomic_load_n(ptr, __ATOMIC_RELAXED); } template @@ -127,7 +123,7 @@ private: if (flags & (Recv*RoleWaitRecv | Send*RoleWaitSend)) { if (isSendNotRecv && (flags & SizesFifoEnabled)) - __atomic_store_n(connSizesFifoPtr+step%NCCL_STEPS, nelts*sizeof(T), __ATOMIC_SEQ_CST); + __atomic_store_n(connSizesFifoPtr+step%NCCL_STEPS, nelts*sizeof(T), __ATOMIC_RELAXED); void **ptrs = isSendNotRecv ? (ncclShmem.groups[group].dsts + Dst) : (ncclShmem.groups[group].srcs + Src); @@ -710,7 +706,6 @@ private: if (flags & RoleOutput) userBuff = (T*)outputBuf; } - template __device__ __forceinline__ void send(intptr_t inpIx, int eltN) { genericOp<0, 0, 0, 1, Input, -1>(inpIx, -1, eltN, false); } @@ -724,7 +719,6 @@ private: genericOp<0, 1, 0, 1, Output, -1>(outIx, outIx, eltN, false); } - template __device__ __forceinline__ void recv(intptr_t outIx, int eltN, bool postOp=false) { genericOp<0, 0, 1, 0, -1, Output>(-1, outIx, eltN, postOp); }