Removing redundant LOAD and STORE on primitives plus adding some atomics (#585)

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

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

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]) :

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

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;