[Device] Adjust threadblock size for gfx950 to increase LL64/Simple performance for AR, RS and AG (#1978)

* Add initial commit to increase tb size to 512
* Fix LL perf issue when subset of NCCL_MAX_NTHREADS is used
Adding a constant to barrier_generic logic from using fallback logic when nthreads < NCCL_MAX_NTHREADS and nthreads == blockDim.X
* Adjust nthreads for LL
* Opt threads for reduce_scatter upper small range
* Add macro for single node
* Restrict MSCCL to 256 threads to prevent mem access fault
* Support pre-MI350 compatibility
* Partially refactor threadblock size override
* Use const macros instead of numerals
* opt out of unused function

[ROCm/rccl commit: 12f51ba8bf]

projects/rccl/src/device/msccl_kernel_impl.h

@@ -28,7 +28,7 @@ extern __shared__ struct mscclShmemData mscclShmem;
 
 inline __device__ static void barrier(int nthreads) {
 #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
-  assert(nthreads == NCCL_MAX_NTHREADS);
+  assert(nthreads == MSCCL_MAX_NTHREADS);
   #ifdef __GFX12__
     __asm__ __volatile__("s_waitcnt vmcnt(0) lgkmcnt(0)\ns_barrier_signal -1\ns_barrier_wait -1");
   #else
@@ -96,7 +96,7 @@ __device__ __forceinline__ void mscclRunInterpreter(
   struct ncclDevComm* comm, struct mscclAlgo* algo, struct mscclWork* work) {
   const int tid = threadIdx.x;
   const int bid = blockIdx.x;
-  const int nthreads = NCCL_MAX_NTHREADS;
+  const int nthreads = MSCCL_MAX_NTHREADS;
 
 #if defined(ENABLE_NPKIT)
   uint64_t timestamp_entry = 0;

projects/rccl/src/device/primitives.h

@@ -16,9 +16,8 @@
 #include "common.h"
 
 #define NCCL_SPINS_BEFORE_CHECK_ABORT 10000
-
 #define barrier_generic(__THREAD_FENCE, NWORKERS, BARRIER_NEXT, BARRIERS_PTR) do { \
-  if (nthreads == NCCL_MAX_NTHREADS) { \
+  if (nthreads == threadsPerBlock) { \
     __THREAD_FENCE; __builtin_amdgcn_s_barrier(); \
   } else { \
     /**const int w = threadIdx.x/WARP_SIZE //unused variable - compiler warning**/;\

projects/rccl/src/device/prims_ll.h

@@ -25,6 +25,7 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p, isNetOffload, Metadata, Pi
   const int nthreads;
   const int wid;
   const int group;
+  const int threadsPerBlock;
   const int stepLines;
   Fan fan;
   T *userBufs[3];
@@ -149,7 +150,7 @@ private:
   __device__ uint64_t readLL(int offset, int i) {
     union ncclLLFifoLine* src = recvPtr(i) + offset;
     uint32_t flag = recvFlag(i);
-    uint32_t data1, flag1, data2, flag2; 
+    uint32_t data1, flag1, data2, flag2;
     (void)data1; (void)flag1; (void)data2; (void)flag2; // unused variable - compiler warning
     int spins = 0;
 
@@ -651,7 +652,7 @@ public:
       bool ipcReg = false, bool netReg = false, int stepSize_ = 0
     ):
     redOp(redOpArg),
-    tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), group(group),
+    tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), group(group), threadsPerBlock(blockDim.x),
     stepLines(ncclShmem.comm.buffSizes[NCCL_PROTO_LL]/NCCL_STEPS/sizeof(ncclLLFifoLine)) {
     auto *channel = &ncclShmem.channel;
     barriers = &ncclShmem.groups[group].barrier;

projects/rccl/src/device/prims_ll128.h

@@ -35,6 +35,7 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload, Metadata,
   const int warpInBlock; // warp index in thread block
   const bool flagThread;
   const int group;
+  const int threadsPerBlock;
   Fan fan;
   T *userBufs[3];
   struct ncclConnInfo* recvConn = NULL;
@@ -138,8 +139,8 @@ private:
     if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1);
   }
   inline __device__ void postSend() {
-    __atomic_signal_fence(__ATOMIC_SEQ_CST);	  
-    asm volatile("s_waitcnt lgkmcnt(0) vmcnt(0)"); 	  
+    __atomic_signal_fence(__ATOMIC_SEQ_CST);
+    asm volatile("s_waitcnt lgkmcnt(0) vmcnt(0)");
     __atomic_signal_fence(__ATOMIC_SEQ_CST);
 
     if (sendConnTailPtr) {
@@ -576,8 +577,8 @@ public:
     ):
     redOp(redOpArg),
     tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE),                                /*compiler warnings*/
-    stepSize(ncclShmem.comm.buffSizes[NCCL_PROTO_LL128]/NCCL_STEPS/sizeof(uint64_t)), 
-    warp(tid/WARP_SIZE), warpInBlock(threadIdx.x/WARP_SIZE), flagThread((tid%4)==3), group(group){
+    stepSize(ncclShmem.comm.buffSizes[NCCL_PROTO_LL128]/NCCL_STEPS/sizeof(uint64_t)),
+    warp(tid/WARP_SIZE), warpInBlock(threadIdx.x/WARP_SIZE), flagThread((tid%4)==3), group(group), threadsPerBlock(blockDim.x){
     auto *channel = &ncclShmem.channel;
     barriers = &ncclShmem.groups[group].barrier;
     int nrecv=0, nsend=0;

projects/rccl/src/device/prims_simple.h

@@ -51,6 +51,7 @@ class Primitives<
   int index; // Peer index I'm responsible for
   int flags;
   const int group;
+  const int threadsPerBlock;
   uint64_t step;
   struct ncclConnInfo* conn = NULL;
   struct ncclConnFifo* connFifo = NULL;
@@ -757,7 +758,7 @@ public:
       struct ncclDevWorkP2p* p2pWork = nullptr, int stepSize_ = 0, int mode = primsModeDefault
     ):
     tid(tid), tidInBlock(threadIdx.x), nthreads(nthreads), /*compiler warnings*/
-    stepSize(stepSize_ == 0 ? ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/sizeof(T) : stepSize_), group(group) {
+    stepSize(stepSize_ == 0 ? ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/sizeof(T) : stepSize_), group(group), threadsPerBlock(blockDim.x){
 
     barriers = &ncclShmem.groups[group].barrier;
     // PAT uses the same barrier for each group