Relax workgroup barrier implementation for MSCCL send/recv ops. (#997)

* Trim logic.

* Revert "Trim logic."

This reverts commit 8f2dba6c764108acf2bf5428366b9f41d4d206b9.

* Introduce MSCCL template parameters to send / recv.

* Address review feedbacks.

[ROCm/rccl commit: baadda4bd8]

projects/rccl/src/collectives/device/msccl_kernel_impl.h

@@ -339,7 +339,7 @@ __device__ __forceinline__ void mscclRunInterpreter(
               NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_SEND_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
             }
 #endif	
-          prims.send(srcOffset, thisNelem); // LL.send is the only situation where there is no barrier at the end.
+          prims.template send<1>(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) {
@@ -353,7 +353,7 @@ __device__ __forceinline__ void mscclRunInterpreter(
               NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RECV_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
             }
 #endif	
-          prims.recv(dstOffset, thisNelem);
+          prims.template recv<1>(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());

projects/rccl/src/collectives/device/prims_ll.h

@@ -89,6 +89,13 @@ 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) {
@@ -100,6 +107,7 @@ private:
     return abort;
   }
 
+  template<int MSCCL = 0>
   inline __device__ void waitSend(int nbytes) {
 #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_WAIT_SEND_ENTRY)
     if (tid == 0) {
@@ -121,7 +129,11 @@ private:
       }
       sendConnHead += 1;
     }
-    barrier();
+    if (MSCCL) {
+      msccl_barrier();
+    } else {
+      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(),
@@ -133,8 +145,13 @@ private:
   inline __device__ void incRecv(int i) {
     recvStep[i] += 1;
   }
+  template<int MSCCL = 0>
   inline __device__ void postRecv() {
-    barrier();
+    if (MSCCL) {
+      msccl_barrier();
+    } else {
+      barrier();
+    }
     if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1);
   }
 
@@ -385,7 +402,7 @@ private:
     }
   }
 
-  template <int RECV, int SEND, int SrcBuf, int DstBuf>
+  template <int RECV, int SEND, int SrcBuf, int DstBuf, int MSCCL = 0>
   __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;
@@ -394,7 +411,7 @@ private:
 
     // Always waitSend in case of cleanup
     nelem = nelem < 0 ? 0 : nelem;
-    if (SEND) waitSend(divUp(nelem, EltPerLine)*sizeof(ncclLLFifoLine));
+    if (SEND) waitSend<MSCCL>(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) {
@@ -477,7 +494,7 @@ private:
 
     if (RECV) {
       for (int i=0; i < MaxRecv; i++) incRecv(i);
-      postRecv();
+      postRecv<MSCCL>();
     }
     if (SEND) {
       for (int i=1; i < MaxSend && i < fan.nsend(); i++)
@@ -639,6 +656,7 @@ private:
     userBufs[Output] += delta;
   }
 
+  template<int MSCCL = 0>
   __device__ void send(intptr_t inpIx, int eltN) {
 #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_ENTRY)
     if (tid == 0) {
@@ -646,7 +664,7 @@ private:
           ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
     }
 #endif
-    LLGenericOp<0, 1, Input, -1>(inpIx, -1, eltN, false);
+    LLGenericOp<0, 1, Input, -1, MSCCL>(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(),
@@ -669,6 +687,7 @@ private:
     }
 #endif
   }
+  template<int MSCCL = 0>
   __device__ void recv(intptr_t outIx, int eltN, bool postOp=false) {
 #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_ENTRY)
     if (tid == 0) {
@@ -676,7 +695,7 @@ private:
           ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
     }
 #endif
-    LLGenericOp<1, 0, -1, Output>(-1, outIx, eltN, postOp);
+    LLGenericOp<1, 0, -1, Output, MSCCL>(-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(),

projects/rccl/src/collectives/device/prims_ll128.h

@@ -546,12 +546,14 @@ public:
     userBufs[Output] += delta;
   }
 
+  template<int MSCCL = 0>
   __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<int MSCCL = 0>
   __device__ void recv(intptr_t outIx, int eltN, bool postOp=false) {
     return GenericOp<1, 0, -1, Output>(-1, outIx, eltN, postOp);
   }

projects/rccl/src/collectives/device/prims_simple.h

@@ -710,6 +710,7 @@ private:
     if (flags & RoleOutput) userBuff = (T*)outputBuf;
   }
 
+  template<int MSCCL = 0>
   __device__ __forceinline__ void send(intptr_t inpIx, int eltN) {
     genericOp<0, 0, 0, 1, Input, -1>(inpIx, -1, eltN, false);
   }
@@ -723,6 +724,7 @@ private:
     genericOp<0, 1, 0, 1, Output, -1>(outIx, outIx, eltN, false);
   }
 
+  template<int MSCCL = 0>
   __device__ __forceinline__ void recv(intptr_t outIx, int eltN, bool postOp=false) {
     genericOp<0, 0, 1, 0, -1, Output>(-1, outIx, eltN, postOp);
   }