From baadda4bd8828cb16b0acb6ec0ac39b2d70b86aa Mon Sep 17 00:00:00 2001 From: "Wen-Heng (Jack) Chung" Date: Fri, 8 Dec 2023 17:46:53 -0600 Subject: [PATCH] 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. --- src/collectives/device/msccl_kernel_impl.h | 4 +-- src/collectives/device/prims_ll.h | 33 +++++++++++++++++----- src/collectives/device/prims_ll128.h | 2 ++ src/collectives/device/prims_simple.h | 2 ++ 4 files changed, 32 insertions(+), 9 deletions(-) diff --git a/src/collectives/device/msccl_kernel_impl.h b/src/collectives/device/msccl_kernel_impl.h index 0c537349ba..c211c86921 100644 --- a/src/collectives/device/msccl_kernel_impl.h +++ b/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()); diff --git a/src/collectives/device/prims_ll.h b/src/collectives/device/prims_ll.h index 681a84a448..e6b7fe343d 100644 --- a/src/collectives/device/prims_ll.h +++ b/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 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 inline __device__ void postRecv() { - barrier(); + if (MSCCL) { + msccl_barrier(); + } else { + barrier(); + } if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1); } @@ -385,7 +402,7 @@ private: } } - template + template __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(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(); } if (SEND) { for (int i=1; i < MaxSend && i < fan.nsend(); i++) @@ -639,6 +656,7 @@ private: userBufs[Output] += delta; } + template __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 __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(), diff --git a/src/collectives/device/prims_ll128.h b/src/collectives/device/prims_ll128.h index bb236c0679..858ae5ca71 100644 --- a/src/collectives/device/prims_ll128.h +++ b/src/collectives/device/prims_ll128.h @@ -546,12 +546,14 @@ public: userBufs[Output] += delta; } + template __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 __device__ void recv(intptr_t outIx, int eltN, bool postOp=false) { return GenericOp<1, 0, -1, Output>(-1, outIx, eltN, postOp); } diff --git a/src/collectives/device/prims_simple.h b/src/collectives/device/prims_simple.h index 9bcac9e88e..70e6333ca0 100644 --- a/src/collectives/device/prims_simple.h +++ b/src/collectives/device/prims_simple.h @@ -710,6 +710,7 @@ private: if (flags & RoleOutput) userBuff = (T*)outputBuf; } + template __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 __device__ __forceinline__ void recv(intptr_t outIx, int eltN, bool postOp=false) { genericOp<0, 0, 1, 0, -1, Output>(-1, outIx, eltN, postOp); }