[DEVICE] Adding ability to choose unroll factor at runtime (#1734)

* Adding runtime unroll factor selection via RCCL_UNROLL_FACTOR
* [BUILD] Add support for user-defined UNROLL for debugging
* Update CHANGELOG.md
* Fix COLLTRACE errors in CI
* Add debug statements for unroll and resolve warnings
* Incorporate UNROLL into ONLY_FUNCS for debugging

---------

Signed-off-by: nileshnegi <Nilesh.Negi@amd.com>
Co-authored-by: gilbertlee-amd <44450918+gilbertlee-amd@users.noreply.github.com>
Co-authored-by: Jeffrey Novotny <jnovotny@amd.com>

src/enqueue.cc

@@ -28,30 +28,31 @@
 
 using namespace rccl;
 
-struct ncclKernelMatch {
-  void* kernelFn;
-  bool specialized;
-};
+/* [RCCL] Determine which GPU kernel to execute */
+void* rcclGetKernelIndex(int unroll, bool useCollTrace, struct ncclTaskColl* task = NULL)
+{
+  // At this time, unroll factor is controlled only by passed in unroll argument
+  // After more investigation, this may be further tuned by the actual task being processed
 
 #ifdef ENABLE_COLLTRACE
-#define ncclGetKernelIndex(p_comm) ((p_comm)->unroll + ((p_comm)->collTraceEnabled ? 3 : 0))
-static ncclKernelMatch const ncclKerns[6] = {
-  {(void *)ncclDevKernel_Generic_1, true},
-  {(void *)ncclDevKernel_Generic_2, true},
-  {(void *)ncclDevKernel_Generic_4, true},
-  {(void *)ncclDevKernelDebug_Generic_1, true},
-  {(void *)ncclDevKernelDebug_Generic_2, true},
-  {(void *)ncclDevKernelDebug_Generic_4, true}
-};
+  int numKernels = sizeof(rcclKernelTable) / sizeof(rcclKernelTable[0]) / 2;
+  int firstKernel = useCollTrace ? numKernels : 0;
 #else
-#define ncclGetKernelIndex(p_comm) ((p_comm)->unroll)
-static ncclKernelMatch const ncclKerns[3] = {
-  {(void*)ncclDevKernel_Generic_1, true},
-  {(void*)ncclDevKernel_Generic_2, true},
-  {(void*)ncclDevKernel_Generic_4, true}
-};
+  int numKernels = sizeof(rcclKernelTable) / sizeof(rcclKernelTable[0]);
+  int firstKernel = 0;
 #endif
 
+  // Check if the requested unroll exists
+  for (int kernelIdx = 0; kernelIdx < numKernels; kernelIdx++) {
+    if (rcclKernelTable[firstKernel + kernelIdx].unroll == unroll) {
+      return rcclKernelTable[firstKernel + kernelIdx].funcPtr;
+    }
+  }
+  // Fall back to default unroll
+  WARN("Requested RCCL_UNROLL_FACTOR: %d does not exist in `rcclKernelTable`. Falling back to default unroll: %d", unroll, rcclKernelTable[firstKernel].unroll);
+  return rcclKernelTable[firstKernel].funcPtr;
+}
+
 static int rcclProtoGrainSize(int proto, ncclComm *comm){
   switch (proto) {
     case NCCL_PROTO_LL: return 16;
@@ -81,7 +82,7 @@ NCCL_PARAM(L1SharedMemoryCarveout, "L1_SHARED_MEMORY_CARVEOUT", 0);
 
 // Returns maximum kernel stack size of all CUDA kernels
 ncclResult_t ncclInitKernelsForDevice(int cudaArch, int maxSharedMem, size_t* maxStackSize) {
-  constexpr int KernelCount = sizeof(ncclKerns)/sizeof(ncclKerns[0]);
+  constexpr int KernelCount = sizeof(rcclKernelTable)/sizeof(rcclKernelTable[0]);
   ncclResult_t result = ncclSuccess;
   int print = 0;
 
@@ -95,7 +96,7 @@ ncclResult_t ncclInitKernelsForDevice(int cudaArch, int maxSharedMem, size_t* ma
   int ncclMaxSharedMem = rcclShmemDynamicSize(cudaArch, WarpSize);
 
   for (int k=0; k < KernelCount; k++) {
-    void* fn = ncclKerns[k].kernelFn;
+    void* fn = rcclKernelTable[k].funcPtr;
     cudaFuncAttributes attr = {0};
     if (fn == nullptr) continue;
 
@@ -783,8 +784,12 @@ static ncclResult_t scheduleCollTasksToPlan(
     //plan->channelMask.masks[channelId/64] |= (2ull<<devWork->channelHi) - (1ull<<devWork->channelLo);
     plan->threadPerBlock = std::max(plan->threadPerBlock, 192 /* 3*WARP_SIZE */);
     if (!plan->kernelSpecialized) {
-      plan->kernelFn = ncclKerns[ncclGetKernelIndex(comm)].kernelFn;
-      plan->kernelSpecialized = ncclKerns[ncclGetKernelIndex(comm)].specialized;
+#ifdef ENABLE_COLLTRACE
+      plan->kernelFn = rcclGetKernelIndex(comm->unroll, comm->collTraceEnabled);
+#else
+      plan->kernelFn = rcclGetKernelIndex(comm->unroll, false);
+#endif
+      plan->kernelSpecialized = true;
     }
 
     if (comm->rank == 0) {
@@ -1084,8 +1089,12 @@ static ncclResult_t scheduleP2pTasksToPlan(
 
   plan->threadPerBlock = std::max(plan->threadPerBlock, NCCL_MAX_NTHREADS);
   if (!plan->kernelSpecialized) {
-    plan->kernelFn = ncclKerns[ncclGetKernelIndex(comm)].kernelFn;
-    plan->kernelSpecialized = ncclKerns[ncclGetKernelIndex(comm)].specialized;
+#ifdef ENABLE_COLLTRACE
+    plan->kernelFn = rcclGetKernelIndex(comm->unroll, comm->collTraceEnabled);
+#else
+    plan->kernelFn = rcclGetKernelIndex(comm->unroll, false);
+#endif
+    plan->kernelSpecialized = true;
   }
 
   // Compute how much to split operations