AI/rocm-systems
2
0
Форкнуть 0
Код Задачи Запросы на слияние Действия Пакеты Проекты Релизы Вики Активность

Rework barriers and adjust scope of atomics (#1019)

Просмотр исходного кода
Этот коммит содержится в:
Wenkai Du
2024-01-04 08:18:48 -08:00
коммит произвёл GitHub
родитель 0a53077c9c
Коммит abf265a911
4 изменённых файлов: 19 добавлений и 44 удалений
Скачать Patch файл Скачать Diff файл Показать все файлы Свернуть все файлы
src/collectives/device/msccl_kernel_impl.h
+6 -4
Просмотреть файл
@@ -338,8 +338,9 @@ __device__ __forceinline__ void mscclRunInterpreter(
if (tid == 0) {
NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_SEND_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
}
#endif
prims.template send<1>(srcOffset, thisNelem); // LL.send is the only situation where there is no barrier at the end.
#endif
prims.send(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) {
@@ -352,8 +353,9 @@ __device__ __forceinline__ void mscclRunInterpreter(
if (tid == 0) {
NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RECV_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
}
#endif
prims.template recv<1>(dstOffset, thisNelem);
#endif
prims.recv(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());
src/collectives/device/prims_ll.h
+8 -27
Просмотреть файл
@@ -89,13 +89,6 @@ 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) {
@@ -107,7 +100,6 @@ 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) {
@@ -125,15 +117,11 @@ 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_RELAXED);
}
sendConnHead += 1;
}
if (MSCCL) {
msccl_barrier();
} else {
barrier();
}
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(),
@@ -145,13 +133,8 @@ private:
inline __device__ void incRecv(int i) {
recvStep[i] += 1;
}
template<int MSCCL = 0>
inline __device__ void postRecv() {
if (MSCCL) {
msccl_barrier();
} else {
barrier();
}
barrier();
if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1);
}
@@ -402,7 +385,7 @@ private:
}
}
template <int RECV, int SEND, int SrcBuf, int DstBuf, int MSCCL = 0>
template <int RECV, int SEND, int SrcBuf, int DstBuf>
__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;
@@ -411,7 +394,7 @@ private:
// Always waitSend in case of cleanup
nelem = nelem < 0 ? 0 : nelem;
if (SEND) waitSend<MSCCL>(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) {
@@ -494,7 +477,7 @@ private:
if (RECV) {
for (int i=0; i < MaxRecv; i++) incRecv(i);
postRecv<MSCCL>();
postRecv();
}
if (SEND) {
for (int i=1; i < MaxSend && i < fan.nsend(); i++)
@@ -656,7 +639,6 @@ 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) {
@@ -664,7 +646,7 @@ private:
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
#endif
LLGenericOp<0, 1, Input, -1, MSCCL>(inpIx, -1, eltN, false);
LLGenericOp<0, 1, Input, -1>(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(),
@@ -687,7 +669,6 @@ 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) {
@@ -695,7 +676,7 @@ private:
ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
}
#endif
LLGenericOp<1, 0, -1, Output, MSCCL>(-1, outIx, eltN, postOp);
LLGenericOp<1, 0, -1, Output>(-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(),
src/collectives/device/prims_ll128.h
+2 -4
Просмотреть файл
@@ -99,12 +99,12 @@ private:
int spins = 0;
while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) {
__builtin_amdgcn_s_sleep(1);
sendConnHeadCache = atomicAdd_system((unsigned long long *)sendConnHeadPtr, 0);
sendConnHeadCache = __atomic_load_n(sendConnHeadPtr, __ATOMIC_RELAXED);
if (checkAbort(spins, wid, 1)) break;
}
__asm__ __volatile__("s_wakeup");
if (sendConnFifoPtr) {
__atomic_store_n(sendConnFifoPtr+sendStep[wid]%NCCL_STEPS, nbytes, __ATOMIC_SEQ_CST);
__atomic_store_n(sendConnFifoPtr+sendStep[wid]%NCCL_STEPS, nbytes, __ATOMIC_RELAXED);
}
sendConnHead += 1;
}
@@ -546,14 +546,12 @@ 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);
}
src/collectives/device/prims_simple.h
+3 -9
Просмотреть файл
@@ -81,7 +81,7 @@ private:
inline __device__ bool checkAbort(int &spins) {
spins++;
if (!(flags & Aborted) && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) {
if (atomicAdd_system((unsigned int *)ncclShmem.comm.abortFlag, 0)) {
if (__atomic_load_n(ncclShmem.comm.abortFlag, __ATOMIC_SEQ_CST)) {
flags |= Aborted;
ncclShmem.aborted = 1;
}
@@ -100,11 +100,7 @@ private:
#endif
// volatile is faster than acquire but not as correct. Make sure reduceCopy
// loads data using volatile so it doesn't see stale data in L1.
#ifdef __GFX9__
return atomicAdd((unsigned long long *)ptr, 0);
#else
return __atomic_load_n(ptr, __ATOMIC_SEQ_CST);
#endif
return __atomic_load_n(ptr, __ATOMIC_RELAXED);
}
template <int DirectRecv, int DirectSend, int Recv, int Send, int Src, int Dst>
@@ -127,7 +123,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_RELAXED);
void **ptrs = isSendNotRecv ? (ncclShmem.groups[group].dsts + Dst)
: (ncclShmem.groups[group].srcs + Src);
@@ -710,7 +706,6 @@ 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);
}
@@ -724,7 +719,6 @@ 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);
}