Dynamically select unroll factor to build for when targeting local arch (#1371)

* Dynamically select unroll factor to build for when targeting local arch only

src/enqueue.cc

@@ -31,16 +31,15 @@ struct ncclKernelMatch {
 };
 
 #ifdef ENABLE_COLLTRACE
-static ncclKernelMatch const ncclKerns[4] = {
+#define ncclGetKernelIndex(p_comm) ((p_comm)->collTraceThread ? 1 : 0)
+static ncclKernelMatch const ncclKerns[2] = {
   {(void *)ncclDevKernel_Generic, true},
-  {(void *)ncclDevKernel_Generic_4, true},
   {(void *)ncclDevKernelDebug_Generic, true},
-  {(void *)ncclDevKernelDebug_Generic_4, true},
 };
 #else
-static ncclKernelMatch const ncclKerns[2] = {
-  {(void*)ncclDevKernel_Generic, true},
-  {(void*)ncclDevKernel_Generic_4, true},
+#define ncclGetKernelIndex(p_comm) (0)
+static ncclKernelMatch const ncclKerns[1] = {
+  {(void*)ncclDevKernel_Generic, true}
 };
 #endif
 
@@ -56,24 +55,19 @@ static ncclResult_t initCollProxyOp(struct ncclInfo* collInfo, int channelId, ui
 static ncclResult_t getTunerInfo(struct ncclInfo* collInfo, int collNetSupport, int nvlsSupport, int numPipeOps);
 static ncclResult_t topoGetAlgoInfo(struct ncclInfo* collInfo, int collNetSupport, int nvlsSupport, int numPipeOps);
 static ncclResult_t getChannnelThreadInfo(struct ncclInfo* collInfo);
-static ncclResult_t computeCollWorkFunc(struct ncclInfo* collInfo);
+static ncclResult_t computeCollWorkFunc(struct ncclInfo* collInfo, int unroll);
 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
+int getUnrollFactor(struct ncclComm* comm) {
   hipDeviceProp_t devProp;
   CUDACHECK(hipGetDeviceProperties(&devProp, comm->cudaDev));
   if(IsArchMatch(devProp.gcnArchName, "gfx908") || (IsArchMatch(devProp.gcnArchName, "gfx94")
     && devProp.multiProcessorCount > 80))
-    return start_idx;
+    return NCCL_UNROLL_2;
   else
-    return start_idx + 1;
+    return NCCL_UNROLL_4;
 }
 
 // Returns maximum kernel stack size of all CUDA kernels
@@ -194,7 +188,7 @@ static ncclResult_t appendWorkElemP2p(
     struct ncclComm* comm, struct ncclKernelPlan* plan, int channelId,
     struct ncclWorkElemP2p const *elem, bool fuseOk
   ) {
-  int funcIndex = ncclDevFuncId_P2p();
+  int funcIndex = ncclDevFuncId_P2p(plan->unroll);
   if (funcIndex < 0) {
     WARN("%s: unsupported collective. Please ensure the collective has been enabled in build.", __func__);
     return ncclInvalidUsage;
@@ -220,7 +214,7 @@ static ncclResult_t appendWorkElemP2p(
   }
   q = ncclMemoryStackAlloc<struct ncclWorkList>(&comm->memScoped);
   q->work.header.type = ncclWorkTypeP2p;
-  q->work.header.funcIndex = ncclDevFuncId_P2p();
+  q->work.header.funcIndex = funcIndex;
   chan->p2pTailElem[ncclWorkP2pTypeRecv-1] = 0;
   chan->p2pTailElem[ncclWorkP2pTypeSend-1] = 1;
   q->work.p2pElems[chan->p2pTailElem[elem->p2pType-1]] = *elem; // C++ struct assignment
@@ -851,7 +845,7 @@ static ncclResult_t scheduleCollTasksToPlan(
         NCCLCHECK(getTunerInfo(aggInfo, collNetSupport, nvlsSupport, 1));
         NCCLCHECK(topoGetAlgoInfo(aggInfo, collNetSupport, nvlsSupport, 1));
         NCCLCHECK(getChannnelThreadInfo(aggInfo));
-        NCCLCHECK(computeCollWorkFunc(aggInfo));
+        NCCLCHECK(computeCollWorkFunc(aggInfo, plan->unroll));
         NCCLCHECK(getPatternInfo(aggInfo));
 
         // Try to assign algo and proto to all possible collectives
@@ -1330,6 +1324,7 @@ ncclResult_t ncclLaunchPrepare(struct ncclComm* comm) {
       plan->comm = comm;
       plan->reclaimer.fn = reclaimPlan;
       plan->persistent = persistent;
+      plan->unroll = getUnrollFactor(comm);
 
       // Non-persistent kernels fill up at most half of our fifo per kernel.
       int nWorkBudget = plan->persistent ? INT_MAX : comm->workFifoDepth/2;
@@ -1762,8 +1757,8 @@ static ncclResult_t getPatternInfo(struct ncclInfo* collInfo) {
 
 RCCL_PARAM(IntraNetThreshold, "INTRANET_THRESHOLD", 8388608);
 
-static ncclResult_t computeCollWorkFunc(struct ncclInfo* collInfo) {
-  collInfo->workFuncIndex = ncclDevFuncId(collInfo->coll, collInfo->opFull.op, collInfo->datatype, collInfo->algorithm, collInfo->protocol);
+static ncclResult_t computeCollWorkFunc(struct ncclInfo* collInfo, int unroll) {
+  collInfo->workFuncIndex = ncclDevFuncId(collInfo->coll, collInfo->opFull.op, collInfo->datatype, collInfo->algorithm, collInfo->protocol, unroll);
   if (collInfo->workFuncIndex < 0) {
     WARN("%s: unsupported collective. Please ensure the collective has been enabled in build.", __func__);
     return ncclInvalidUsage;