Another threadfence and flags rework (#629)

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)), \

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

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<int WordPerThread>

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<int Recv, int Send>
   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<Recv, Send>();
         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()