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]
This commit is contained in:
Wen-Heng (Jack) Chung
2023-12-08 17:46:53 -06:00
committed by GitHub
parent 22730bdc4d
commit 529f72ceaa
4 changed files with 32 additions and 9 deletions
@@ -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());