From a523b37ac7a95fd788c9908860e932ba09cafc7d Mon Sep 17 00:00:00 2001 From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com> Date: Wed, 28 Sep 2022 16:49:29 -0700 Subject: [PATCH] Another threadfence and flags rework (#629) --- src/collectives/device/common.h | 8 ++++++++ src/collectives/device/prims_ll.h | 4 ++-- src/collectives/device/prims_ll128.h | 4 ++-- src/collectives/device/prims_simple.h | 28 ++++++++++++++++----------- 4 files changed, 29 insertions(+), 15 deletions(-) diff --git a/src/collectives/device/common.h b/src/collectives/device/common.h index cb16565ab0..bb96806460 100644 --- a/src/collectives/device/common.h +++ b/src/collectives/device/common.h @@ -19,6 +19,14 @@ #define __synclds() \ asm volatile("s_waitcnt lgkmcnt(0) \n s_barrier"); +#if defined(__gfx90a__) +#define STORE(DST, SRC) \ + { __threadfence_block(); atomicExch((unsigned long long *)(DST), (SRC)); } +#else +#define STORE(DST, SRC) \ + { __atomic_store_n((DST), (SRC), __ATOMIC_SEQ_CST); } +#endif + #define NCCL_FUNC5(func, algo, devredop, type, nullify) \ MACRO_IF(nullify, nullptr, NCCL_FUNC_NAME(func, algo, LL, devredop, type)), \ MACRO_IF(nullify, nullptr, NCCL_FUNC_NAME(func, algo, LL, devredop, type)), \ diff --git a/src/collectives/device/prims_ll.h b/src/collectives/device/prims_ll.h index 0cf7fc8cb6..aac108bb0f 100644 --- a/src/collectives/device/prims_ll.h +++ b/src/collectives/device/prims_ll.h @@ -106,7 +106,7 @@ 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_SEQ_CST); } sendConnHead += 1; } @@ -124,7 +124,7 @@ private: } inline __device__ void postRecv() { barrier(); - if (recvConnHeadPtr) atomicExch_system((unsigned long long *)recvConnHeadPtr, recvConnHead += 1); + if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1); } inline __device__ void incSend(int i, int offset) { diff --git a/src/collectives/device/prims_ll128.h b/src/collectives/device/prims_ll128.h index 0ccea1548c..b7ef95dd0f 100644 --- a/src/collectives/device/prims_ll128.h +++ b/src/collectives/device/prims_ll128.h @@ -101,10 +101,10 @@ private: } inline __device__ void postRecv() { - if (recvConnHeadPtr) atomicExch_system((unsigned long long *)recvConnHeadPtr, recvConnHead += 1); + if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1); } inline __device__ void postSend() { - if (sendConnTailPtr) { __threadfence(); atomicExch_system((unsigned long long *)sendConnTailPtr, sendConnTail += 1); } + if (sendConnTailPtr) { __threadfence(); STORE((unsigned long long *)sendConnTailPtr, sendConnTail += 1); } } template diff --git a/src/collectives/device/prims_simple.h b/src/collectives/device/prims_simple.h index bc8904f0c0..230b145c91 100644 --- a/src/collectives/device/prims_simple.h +++ b/src/collectives/device/prims_simple.h @@ -116,7 +116,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_SEQ_CST); void **ptrs = isSendNotRecv ? (ncclShmem.groups[group].dsts + Dst) : (ncclShmem.groups[group].srcs + Src); @@ -149,15 +149,11 @@ private: template inline __device__ void postPeer() { if ((flags & Send*RolePostSend) && next_hdp_reg) - atomicExch_system(next_hdp_reg, 0x1); + STORE((unsigned int *)next_hdp_reg, 0x1); if (flags & (Recv*RolePostRecv | Send*RolePostSend)) { step += StepPerSlice; -#if defined(__gfx90a__) - atomicExch_system((unsigned long long *)connStepPtr, step); -#else - __atomic_store_n(connStepPtr, step, __ATOMIC_SEQ_CST); -#endif + STORE(connStepPtr, step); } } @@ -323,11 +319,16 @@ private: } barrier(); // This barrier has a counterpart in following loop #if defined(__gfx90a__) - if ((MaxSend == 0 || MaxRecv == 0) && Send && (flags & RolePostSend) && index == 0) __threadfence_system(); + if (Send && (flags & RolePostSend) && index == 0) { + if (MaxSend == 0 || MaxRecv == 0) + __threadfence_system(); + else + __asm__ __volatile__("s_waitcnt vmcnt(0) lgkmcnt(0); buffer_wbinvl1_vol"); + } #else if (Send && (flags & RolePostSend) && index == 0) __threadfence_system(); #endif - __syncwarp(); + __syncwarp(); postPeer(); offset += sliceSize; slice += 1; @@ -346,7 +347,12 @@ private: } barrier(); // Has couterpart in preceding worker-only loop. #if defined(__gfx90a__) - if ((MaxSend == 0 || MaxRecv == 0) && Send && (flags & RolePostSend) && sliceSize > 0 && index == 0) __threadfence_system(); + if (Send && (flags & RolePostSend) && sliceSize > 0 && index == 0) { + if (MaxSend == 0 || MaxRecv == 0) + __threadfence_system(); + else + __asm__ __volatile__("s_waitcnt vmcnt(0) lgkmcnt(0); buffer_wbinvl1_vol"); + } #else if (Send && (flags & RolePostSend) && sliceSize > 0 && index == 0) __threadfence_system(); #endif @@ -433,7 +439,7 @@ private: step = roundUp(step, SlicePerChunk*StepPerSlice); if (flags & RolePostRecv) { connStepPtr = conn->head; - atomicExch_system((unsigned long long *)connStepPtr, step); // Return credits in case we rounded up. + STORE(connStepPtr, step); // Return credits in case we rounded up. } if (flags & RoleWaitRecv) { ncclShmem.groups[group].recvConns[index] = conn; // WaitRecv role saves since that's who needs it in setDataPtrs()