Force enable proto and/or algo after model selection (#1799)

* Force enable proto or algo

* Remove inc nccl_common.h

* Move logic and add error checks

* Fix topo_expl compatibility

* Allow algo/proto overrides

* Remove extra function decl

* Clarify warning message

* Move algo/proto overrides into separate functions

* Update CHANGELOG.md

src/enqueue.cc

@@ -2103,6 +2103,8 @@ static ncclResult_t topoGetAlgoInfo(
   if (info->algorithm == NCCL_ALGO_TREE) nt = NCCL_MAX_NTHREADS; // Tree now uses all threads always.
   if (info->algorithm == NCCL_ALGO_PAT) nt = NCCL_MAX_NTHREADS;
   info->nWarps = nt/comm->WarpSize;
+  rcclOverrideAlgorithm(ncclAlgoStr, table, info);
+  rcclOverrideProtocol(ncclProtoStr, table, info);
   return ncclSuccess;
 }
 

src/include/rccl_common.h

@@ -81,6 +81,9 @@ inline size_t rcclGetSizePerRank(ncclFunc_t const& func, size_t const& nBytes, i
   // For AR, this is the send/recv size per rank
   return (func == ncclFuncReduceScatter || func == ncclFuncAllGather || func == ncclFuncBroadcast) ? nBytes / nRanks : nBytes;
 }
+ncclResult_t rcclGetAlgoProtoIndex(const char *envStr, const char* algoProtoString[], int nEntries, int& result);
+ncclResult_t rcclOverrideProtocol(const char* ncclProtoStr[], float table[][NCCL_NUM_PROTOCOLS], struct ncclTaskColl* info);
+ncclResult_t rcclOverrideAlgorithm(const char* ncclAlgoStr[], float table[][NCCL_NUM_PROTOCOLS], struct ncclTaskColl* info);
 void rcclUpdateCollectiveProtocol(struct ncclComm* comm, size_t const& nBytes, struct ncclTaskColl* info);
 void rcclUpdateThreadThreshold(struct ncclComm* comm, size_t const& nBytes, struct ncclTaskColl* info, int& threadThreshold);
 void rcclSetPipelining(struct ncclComm* comm, size_t const& nBytes, struct ncclTaskColl* info);

src/rccl_wrap.cc

@@ -85,6 +85,74 @@ void rcclUpdateCollectiveProtocol(struct ncclComm* comm, size_t const& nBytes, s
   }
 }
 
+ncclResult_t rcclGetAlgoProtoIndex(const char *envStr, const char* algoProtoString[], int nEntries, int& result) {
+  if(envStr) {
+    for (int i = 0; i < nEntries; ++i) {
+      if (strcasecmp(envStr, algoProtoString[i]) == 0) {
+        result = i;
+        return ncclSuccess;
+      }
+    }
+    static bool failedProtoWarn = false;
+    if (!failedProtoWarn) {
+      WARN("Invalid algo or protocol string passed %s", envStr);
+      failedProtoWarn = true;
+      return ncclInvalidUsage;
+    }
+  }
+  return ncclInvalidUsage;
+}
+
+ncclResult_t rcclOverrideProtocol(const char* ncclProtoStr[], float table[][NCCL_NUM_PROTOCOLS], struct ncclTaskColl* info) {
+  static const char* protoOverrideEnv = ncclGetEnv("RCCL_OVERRIDE_PROTO");
+  static bool validInput = true;
+  if (!validInput) return ncclInvalidUsage;
+
+  if (protoOverrideEnv) {
+    static int protoVal = NCCL_PROTO_UNDEF;
+    if (protoVal == NCCL_PROTO_UNDEF) {
+      if (rcclGetAlgoProtoIndex(protoOverrideEnv, ncclProtoStr, NCCL_NUM_PROTOCOLS, protoVal) != ncclSuccess) {
+        validInput = false;
+        return ncclInvalidUsage;
+      }
+    }
+    if (protoVal > NCCL_PROTO_UNDEF) {
+      if (table[info->algorithm][protoVal] == NCCL_ALGO_PROTO_IGNORE) {
+        WARN("Failed to force unsupported protocol %s for function %s with datatype %s", protoOverrideEnv, ncclFuncToString(info->func), ncclDatatypeToString(info->datatype));
+        return ncclInternalError;
+      } else {
+        info->protocol = protoVal;
+      }
+    }
+  }
+  return ncclSuccess;
+}
+
+ncclResult_t rcclOverrideAlgorithm(const char* ncclAlgoStr[], float table[][NCCL_NUM_PROTOCOLS], struct ncclTaskColl* info) {
+  static const char* algoOverrideEnv = ncclGetEnv("RCCL_OVERRIDE_ALGO");
+  static bool validInput = true;
+  if (!validInput) return ncclInvalidUsage;
+
+  if (algoOverrideEnv) {
+    static int algoVal = NCCL_ALGO_UNDEF;
+    if (algoVal == NCCL_ALGO_UNDEF) {
+      if (rcclGetAlgoProtoIndex(algoOverrideEnv, ncclAlgoStr, NCCL_NUM_ALGORITHMS, algoVal) != ncclSuccess) {
+        validInput = false;
+        return ncclInvalidUsage;
+      }
+    }
+    if (algoVal > NCCL_ALGO_UNDEF) {
+      if (table[algoVal][info->protocol] == NCCL_ALGO_PROTO_IGNORE) {
+        WARN("Failed to force unsupported algorithm %s for function %s with datatype %s", algoOverrideEnv, ncclFuncToString(info->func), ncclDatatypeToString(info->datatype));
+        return ncclInternalError;
+      } else {
+        info->algorithm = algoVal;
+      }
+    }
+  }
+  return ncclSuccess;
+}
+
 void rcclUpdateThreadThreshold(struct ncclComm* comm, size_t const& nBytes, struct ncclTaskColl* info, int& threadThreshold) {
   // Honor user input for thread thresholds
   static int userChannelControlInput = -2;