From 451c287aa6e40039ac71051177013ef89b5bc430 Mon Sep 17 00:00:00 2001 From: akolliasAMD <99202231+akolliasAMD@users.noreply.github.com> Date: Thu, 21 Jul 2022 13:04:57 -0600 Subject: [PATCH] Removing redundant LOAD and STORE on primitives plus adding some atomics (#585) --- src/collectives/device/common_kernel.h | 2 +- src/collectives/device/prims_ll.h | 4 ++-- src/collectives/device/prims_simple.h | 18 +++++++++--------- src/include/comm.h | 10 ++++++++++ src/include/devcomm.h | 9 --------- 5 files changed, 22 insertions(+), 21 deletions(-) diff --git a/src/collectives/device/common_kernel.h b/src/collectives/device/common_kernel.h index e117f70561..35c5e34ccb 100644 --- a/src/collectives/device/common_kernel.h +++ b/src/collectives/device/common_kernel.h @@ -20,7 +20,7 @@ static __device__ int min(int a, ssize_t b) { return (a < b) ? a : b; } inline __device__ int loadInt(int* ptr) { int v; #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__) - v = LOAD(ptr); + v = atomicAdd_system((unsigned long long *)ptr, 0); #else asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(v) : "l"(ptr)); diff --git a/src/collectives/device/prims_ll.h b/src/collectives/device/prims_ll.h index 8ff56b5caa..9c7ec8107d 100644 --- a/src/collectives/device/prims_ll.h +++ b/src/collectives/device/prims_ll.h @@ -83,7 +83,7 @@ private: inline __device__ int checkAbort(int &spins, int send) { spins++; if (abort == 0 && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) { - abort = LOAD(ncclShmem->comm.abortFlag); + abort = __atomic_load_n((ncclShmem->comm.abortFlag), __ATOMIC_SEQ_CST); spins = 0; } return abort; @@ -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; - STORE(sendConnFifoPtr+sendConnHead%NCCL_STEPS, size); + __atomic_store_n((sendConnFifoPtr+sendConnHead%NCCL_STEPS), (size), __ATOMIC_SEQ_CST); } sendConnHead += 1; } diff --git a/src/collectives/device/prims_simple.h b/src/collectives/device/prims_simple.h index edc0ac9a65..33fa274531 100644 --- a/src/collectives/device/prims_simple.h +++ b/src/collectives/device/prims_simple.h @@ -117,7 +117,7 @@ private: if (flags & (Recv*RoleWaitRecv | Send*RoleWaitSend)) { if (isSendNotRecv && (flags & SizesFifoEnabled)) - STORE(connSizesFifoPtr+step%NCCL_STEPS, nelts*sizeof(T)); + atomicExch_system((unsigned long long *)(connSizesFifoPtr+step%NCCL_STEPS), nelts*sizeof(T)); void **ptrs = isSendNotRecv ? (ncclShmem->groups[group].dsts + Dst) : (ncclShmem->groups[group].srcs + Src); @@ -440,12 +440,12 @@ private: step = roundUp(step, SlicePerChunk*StepPerSlice); if (flags & RolePostRecv) { connStepPtr = conn->head; - STORE(connStepPtr, step); // Return credits in case we rounded up. + atomicExch_system((unsigned long long *)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() connStepPtr = conn->tail; - connStepCache = LOAD(connStepPtr); + connStepCache = atomicAdd_system((unsigned long long *)connStepPtr, 0); flags |= (conn->offsFifo != nullptr) ? OffsFifoEnabled : 0; if (Direct) { // User buffers have been registered @@ -485,7 +485,7 @@ private: if (flags & RoleWaitSend) { ncclShmem->groups[group].sendConns[index] = conn; // WaitSend role saves since that's who needs it in setDataPtrs() connStepPtr = conn->head; - connStepCache = LOAD(connStepPtr); + connStepCache = atomicAdd_system((unsigned long long *)connStepPtr, 0); flags |= (conn->offsFifo != nullptr) ? OffsFifoEnabled : 0; if (flags & OffsFifoEnabled) connOffsFifoPtr = conn->offsFifo; @@ -575,7 +575,7 @@ private: // Save steps for the next operation if (flags & (RolePostSend|RolePostRecv)) { auto *conns = (flags & RolePostSend) ? ncclShmem->groups[group].sendConns : ncclShmem->groups[group].recvConns; - STORE(&conns[index]->step, step); + conns[index]->step = step; } // Make sure all threads are done writing back conn->step and done using // ncclShmem->groups[group] @@ -598,7 +598,7 @@ private: int spins = 0; void *volatile *slot = ncclShmem->groups[group].recvConns[index]->ptrExchange; // Wait for consumer to consume previous value before trampling it. - while (LOAD(slot) != nullptr && !checkAbort(spins)); + while ((void *)atomicAdd_system((unsigned long long *) slot,0) != nullptr && !checkAbort(spins)); directBuff = (T*)outputBuf; // Encode pointer by XOR'ing against some address they definitely wouldn't send // since we want to allow them sending us nullptr while not colliding with @@ -610,7 +610,7 @@ private: void *volatile *slot = ncclShmem->groups[group].sendConns[index]->ptrExchange; void *ptr; while (true) { - ptr = LOAD(slot); + ptr = (void *)atomicAdd_system((unsigned long long *) slot,0); if (ptr != nullptr || checkAbort(spins)) break; } directBuff = regUsed ? (T*)(e->dnOutputs[index]) : @@ -623,7 +623,7 @@ private: volatile uint64_t* argSlot0 = ncclShmem->groups[group].sendConns[index]->redOpArgExchange; volatile uint64_t* argSlot1 = ncclShmem->groups[group].sendConns[index]->redOpArgExchange+1; // Wait for consumer to consume previous value before trampling it. - while ((*slot != nullptr || *argSlot0 != 0 || *argSlot1 !=0) && !checkAbort(spins)); + while (((void *)atomicAdd_system((unsigned long long *) slot,0) != nullptr || *argSlot0 != 0 || *argSlot1 !=0) && !checkAbort(spins)); // If there is no recv, then we are directly pulling from input buffer (e.g. directScatter) // Otherwise, we are pulling from output buffer (e.g. recvCopyDirectSend) directBuff = MaxRecv == 0 ? (T*)inputBuf : (T*)outputBuf; @@ -642,7 +642,7 @@ private: volatile uint64_t* argSlot1 = ncclShmem->groups[group].recvConns[index]->redOpArgExchange+1; void *ptr; while (true) { - ptr = *slot; + ptr = (void *)atomicAdd_system((unsigned long long *) slot,0); if (ptr != nullptr || checkAbort(spins)) break; } directBuff = regUsed ? (T*)(MaxSend == 0 ? e->upOutputs[index] : e->dnInputs[index]) : diff --git a/src/include/comm.h b/src/include/comm.h index 3360982963..359b7666b1 100644 --- a/src/include/comm.h +++ b/src/include/comm.h @@ -14,6 +14,16 @@ //#include "clique/CliqueManager.h" // [/RCCL] +// Convert volatile access to atomic +#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__) +#define LOAD(VAR) __atomic_load_n((VAR), __ATOMIC_SEQ_CST) +#define STORE(DST, SRC) __atomic_store_n((DST), (SRC), __ATOMIC_SEQ_CST) +#else +#define LOAD(VAR) *(VAR) +#define STORE(DST, SRC) *(DST) = (SRC) +#endif + + #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__) #define HIPRT_CB #else diff --git a/src/include/devcomm.h b/src/include/devcomm.h index 855923cfb1..8e4e97d11e 100644 --- a/src/include/devcomm.h +++ b/src/include/devcomm.h @@ -19,15 +19,6 @@ //#include "clique/CliqueCommon.h" // [/RCCL] -// Convert volatile access to atomic -#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__) -#define LOAD(VAR) __atomic_load_n((VAR), __ATOMIC_SEQ_CST) -#define STORE(DST, SRC) __atomic_store_n((DST), (SRC), __ATOMIC_SEQ_CST) -#else -#define LOAD(VAR) *(VAR) -#define STORE(DST, SRC) *(DST) = (SRC) -#endif - #define NCCL_NUM_FUNCTIONS 5 // SendRecv and AllToAllPivot not included for now typedef enum { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv, ncclFuncSend, ncclFuncRecv, ncclFuncAllToAllPivot, ncclNumFuncs} ncclFunc_t;