Template unroll for RCCL kernels (#1250)

* Template unroll for RCCL kernels

* Adding unroll template arg during CMake hipification

* Reduce linking parallel jobs to avoid OOM in CI

* Workaround issues with UT tests

SWDEV-469533: register spill fix is needed for mainline build
LWPCOMMLIBS-369: cannot enable 112 channels with 80 CUs
Use -parallel-jobs=8 for linking

* CI: do not use -j 16 when building

* CI: use -j 8 when building

* Only reduce parallel linking job for CI extended

* Restore original jenkins command. Change parallel linking jobs in cmake

* Disable MSCCLPP

---------

Co-authored-by: gilbertlee-amd <gilbert.lee@amd.com>

src/device/common.cu

@@ -18,11 +18,17 @@ struct RunWorkNop {
 };
 
 __launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
-  ncclKernelMain<-1, RunWorkNop, false>(comm, channelMask, workHead);
+  ncclKernelMain<-1, RunWorkNop, false, 2>(comm, channelMask, workHead);
+}
+__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_4(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
+  ncclKernelMain<-1, RunWorkNop, false, 4>(comm, channelMask, workHead);
 }
 #ifdef ENABLE_COLLTRACE
 __launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
-  ncclKernelMain<-1, RunWorkNop, true>(comm, channelMask, workHead);
+  ncclKernelMain<-1, RunWorkNop, true, 2>(comm, channelMask, workHead);
+}
+__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_4(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
+  ncclKernelMain<-1, RunWorkNop, true, 4>(comm, channelMask, workHead);
 }
 #endif
 

src/device/common.h

@@ -15,12 +15,6 @@
 #include "network/unpack/unpack_defs.h"
 #include "comm.h"
 
-#if defined(__gfx908__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
-#define COLL_UNROLL 2
-#else
-#define COLL_UNROLL 4
-#endif
-
 #define NCCL_MAX_DEV_ARITY (NCCL_MAX_TREE_ARITY-1)  // Using balanced tree instead of split tree
 
 #define __syncwarp()
@@ -185,14 +179,14 @@ inline __device__ void copyToShmem16(int tid, void* dst, void const* src, int by
   }
 }
 
-template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto>
+template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int COLL_UNROLL>
 struct RunWorkElement {
   __device__ void run(ncclWorkElem*) {
     // Put NOT IMPLEMENTED behavior here.
   }
 };
 
-template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto>
+template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int COLL_UNROLL>
 struct RunWork {
   // This __forceinline__ is necessary. The compiler was inserting a function call
   // here from the LL ncclKernel.
@@ -203,7 +197,7 @@ struct RunWork {
     #pragma unroll 1
     while ((char*)we + stride <= (char*)(w+1) && we->isUsed) {
       if (wid < we->nWarps) {
-        RunWorkElement<Fn, T, RedOp, Algo, Proto>().run(we);
+        RunWorkElement<Fn, T, RedOp, Algo, Proto, COLL_UNROLL>().run(we);
       }
       we = (ncclWorkElem*)((char*)we + stride);
     }
@@ -231,7 +225,7 @@ static __forceinline__ __device__ void ncclRedopPtrDeref(struct ncclWorkElem* we
   }
 }
 
-template<int SpecializedFnId, typename SpecializedRunWork, bool COLLTRACE>
+template<int SpecializedFnId, typename SpecializedRunWork, bool COLLTRACE, int COLL_UNROLL>
 __forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) {
   const int tid = threadIdx.x;
   int x = tid;
@@ -350,9 +344,15 @@ __forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, struct
       SpecializedRunWork().run(&ncclShmem.work);
     } else {
 #ifdef USE_INDIRECT_FUNCTION_CALL
-      ncclDevFuncTable[ncclShmem.work.header.funcIndex]();
+      if (COLL_UNROLL == 4)
+        ncclDevFuncTable_4[ncclShmem.work.header.funcIndex]();
+      else
+        ncclDevFuncTable[ncclShmem.work.header.funcIndex]();
 #else
-      NCCL_CALL_FUNCTIONS(ncclShmem.work.header.funcIndex);
+      if (COLL_UNROLL == 4)
+        NCCL_CALL_FUNCTIONS_4(ncclShmem.work.header.funcIndex);
+      else
+        NCCL_CALL_FUNCTIONS(ncclShmem.work.header.funcIndex);
 #endif
     }
 
@@ -383,35 +383,27 @@ __forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, struct
 }
 
 __global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
+__global__ void ncclDevKernel_Generic_4(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
 #ifdef ENABLE_COLLTRACE
 __global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
-#endif
-
-#ifdef ENABLE_COLLTRACE
-#define DEFINE_ncclDevKernel(suffix, coll, redop, ty, algo, proto, specializedFnId) \
-  __global__ void ncclDevKernel_##suffix(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) { \
-    ncclKernelMain<specializedFnId, RunWork<coll, ty, redop<ty>, algo, proto>, false>(comm, channelMask, workHead); \
-  } \
-  \
-  __global__ void ncclDevKernelDebug_##suffix(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) { \
-    ncclKernelMain<specializedFnId, RunWork<coll, ty, redop<ty>, algo, proto>, true>(comm, channelMask, workHead); \
-  }
-#else
-#define DEFINE_ncclDevKernel(suffix, coll, redop, ty, algo, proto, specializedFnId) \
-  __global__ void ncclDevKernel_##suffix(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead) { \
-    ncclKernelMain<specializedFnId, RunWork<coll, ty, redop<ty>, algo, proto>, false>(comm, channelMask, workHead); \
-  }
+__global__ void ncclDevKernelDebug_Generic_4(struct ncclDevComm* comm, struct channelMasks channelMask, struct ncclWork* workHead);
 #endif
 
 #ifdef USE_INDIRECT_FUNCTION_CALL
 #define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto) \
   __device__ void ncclDevFunc_##suffix() { \
-    RunWork<coll, ty, redop<ty>, algo, proto>().run(&ncclShmem.work); \
+    RunWork<coll, ty, redop<ty>, algo, proto, 2>().run(&ncclShmem.work); \
+  } \
+  __device__ void ncclDevFunc_##suffix##_4() { \
+    RunWork<coll, ty, redop<ty>, algo, proto, 4>().run(&ncclShmem.work); \
   }
 #else
 #define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto) \
   __device__ __attribute__((noinline)) void ncclDevFunc_##suffix() { \
-    RunWork<coll, ty, redop<ty>, algo, proto>().run(&ncclShmem.work); \
+    RunWork<coll, ty, redop<ty>, algo, proto, 2>().run(&ncclShmem.work); \
+  } \
+  __device__ __attribute__((noinline)) void ncclDevFunc_##suffix##_4() { \
+    RunWork<coll, ty, redop<ty>, algo, proto, 4>().run(&ncclShmem.work); \
   }
 #endif
 

src/device/msccl_kernel_impl.h

@@ -404,7 +404,7 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, LL128, fullOps)(struct n
   mscclRunInterpreter<type, Func##devredop<type>, ProtoLL128, fullOps>(comm, algo, work); \
 } \
 __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, Simple, fullOps)(struct ncclDevComm* comm, struct mscclAlgo* algo, struct mscclWork* work) { \
-  mscclRunInterpreter<type, Func##devredop<type>, ProtoSimple<MSCCL_CHUNKSTEPS/MSCCL_SLICESTEPS, MSCCL_SLICESTEPS>, fullOps>(comm, algo, work); \
+  mscclRunInterpreter<type, Func##devredop<type>, ProtoSimple<MSCCL_CHUNKSTEPS/MSCCL_SLICESTEPS, MSCCL_SLICESTEPS, 2>, fullOps>(comm, algo, work); \
 }
 
 #define MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(devredop, fullOps) \

src/device/onerank.cu

@@ -11,6 +11,12 @@
 #include "common.h"
 #include <cuda_runtime.h>
 
+#if defined(__gfx908__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+#define COLL_UNROLL 2
+#else
+#define COLL_UNROLL 4
+#endif
+
 namespace {
   template<typename RedOp>
   __global__ __launch_bounds__(512, 1)

src/device/primitives.h

@@ -55,7 +55,7 @@
  * to how that protocol operates with a consistent interface so that our
  * algorithm code can operate protocol parametrically.
  */
-template<int SlicePerChunk_1, int StepPerSlice_1, int Unroll_1 = COLL_UNROLL, int MultimemSrcs_1 = 0, int MultimemDsts_1 = 0>
+template<int SlicePerChunk_1, int StepPerSlice_1, int Unroll_1, int MultimemSrcs_1 = 0, int MultimemDsts_1 = 0>
 struct ProtoSimple {
   static constexpr int Id = NCCL_PROTO_SIMPLE;
   static constexpr int SlicePerChunk = SlicePerChunk_1;

src/enqueue.cc

@@ -30,13 +30,16 @@ struct ncclKernelMatch {
 };
 
 #ifdef ENABLE_COLLTRACE
-static ncclKernelMatch const ncclKerns[2] = {
+static ncclKernelMatch const ncclKerns[4] = {
   {(void *)ncclDevKernel_Generic, true},
+  {(void *)ncclDevKernel_Generic_4, true},
   {(void *)ncclDevKernelDebug_Generic, true},
+  {(void *)ncclDevKernelDebug_Generic_4, true},
 };
 #else
-static ncclKernelMatch const ncclKerns[1] = {
-  {(void*)ncclDevKernel_Generic, true}
+static ncclKernelMatch const ncclKerns[2] = {
+  {(void*)ncclDevKernel_Generic, true},
+  {(void*)ncclDevKernel_Generic_4, true},
 };
 #endif
 
@@ -57,6 +60,21 @@ static ncclResult_t getPatternInfo(struct ncclInfo* collInfo);
 static ncclResult_t getLoopInfo(struct ncclInfo* collInfo);
 static ncclResult_t getCollNetSupport(struct ncclInfo* info, int* collNetSupport);
 
+int ncclGetKernelIndex(struct ncclComm* comm) {
+#if ENABLE_COLLTRACE
+  int start_idx = comm->collTraceThread ? 2 : 0;
+#else
+  int start_idx = 0;
+#endif
+  hipDeviceProp_t devProp;
+  CUDACHECK(hipGetDeviceProperties(&devProp, comm->cudaDev));
+  if(IsArchMatch(devProp.gcnArchName, "gfx908") || (IsArchMatch(devProp.gcnArchName, "gfx94")
+    && devProp.multiProcessorCount > 80))
+    return start_idx;
+  else
+    return start_idx + 1;
+}
+
 // Returns maximum kernel stack size of all CUDA kernels
 ncclResult_t ncclInitKernelsForDevice(int cudaArch, size_t* maxStackSize) {
   constexpr int KernelCount = sizeof(ncclKerns)/sizeof(ncclKerns[0]);

src/graph/connect.cc

@@ -625,10 +625,11 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa
   NCCLCHECK(connectTrees(comm, treeToParent, treeToChild0, treeToChild1, treePatterns));
 
   // Only use full MAXCHANNELS for gfx94x
-  int maxChannels = IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx94") ? MAXCHANNELS : 2*CHANNEL_LIMIT;
+  int maxChannels = IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx94") ?
+    (comm->topo->nodes[GPU].nodes[0].gpu.cu == 80 ? 80 : MAXCHANNELS) : 2*CHANNEL_LIMIT;
 
   if (graphs[NCCL_ALGO_RING]->nIntraChannels > 0 || comm->nNodes > 1) {
-        maxChannels = std::min(64, maxChannels);
+    maxChannels = std::min(64, maxChannels);
   }
 
   // Duplicate ringPrev/ringNext for ncclBuildRing
@@ -680,6 +681,10 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa
     minNchannels = 2;
     WARN("NCCL_MIN_NCHANNELS set by environment is ignored due to less than 8 GPUs.");
   }
+  if (minNchannels > maxChannels) {
+    minNchannels = 2;
+    WARN("NCCL_MIN_NCHANNELS set by environment is ignored due to greater than max allowed %d channels.", maxChannels);
+  }
 
   if (mscclEnabled() && (comm->topo->mscclEnabled || mscclForceEnabled())) {
     int mscclNumChannelsRequired = maxNchannels;

src/graph/topo.cc

@@ -379,6 +379,8 @@ ncclResult_t ncclTopoAddGpu(struct ncclXmlNode* xmlGpu, struct ncclTopoSystem* s
 #if defined(__HIP_PLATFORM_AMD__) || defined(__HCC__) || defined(__HIPCC__)
   // There is no direct mapping between CUDA SM to HIP GFX. Use SM60 as compatibility level.
   gpu->gpu.cudaCompCap = 60;
+  // Repurpose previously unused "sm" as CU counts
+  NCCLCHECK(xmlGetAttrInt(xmlGpu, "sm", &gpu->gpu.cu));
 #else
   NCCLCHECK(xmlGetAttrInt(xmlGpu, "sm", &gpu->gpu.cudaCompCap));
 #endif

src/graph/topo.h

@@ -128,6 +128,7 @@ struct ncclTopoNode {
       int gdrSupport;
       const char* gcn;
       hipDeviceArch_t arch;
+      int cu;
     }gpu;
     struct {
       uint64_t asic;
@@ -240,12 +241,6 @@ static float ncclTopoXGMISpeed(const char* gcn) {
     return VEGA_XGMI_WIDTH;
 }
 
-#if ENABLE_COLLTRACE
-  #define ncclGetKernelIndex(p_comm) ((p_comm)->collTraceThread ? 1 : 0)
-#else
-  #define ncclGetKernelIndex(p_comm) (0)
-#endif
-
 // Returns NVLink bw in GB/s
 static float ncclTopoNVLinkBw(int cudaCompCap) {
   return

src/graph/xml.cc

@@ -651,11 +651,9 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, uint32_t rocmDev
 
   NCCLCHECK(xmlGetAttrIndex(gpuNode, "sm", &index));
   if (index == -1) {
-    int cudaMajor, cudaMinor;
     cudaDeviceProp devProp;
     CUDACHECK(cudaGetDeviceProperties(&devProp, 0));
-    cudaMajor = devProp.major; cudaMinor = devProp.minor;
-    NCCLCHECK(xmlSetAttrInt(gpuNode, "sm", cudaMajor*10+cudaMinor));
+    NCCLCHECK(xmlSetAttrInt(gpuNode, "sm", devProp.multiProcessorCount));
   }
   int sm;
   NCCLCHECK(xmlGetAttrInt(gpuNode, "sm", &sm));