MSCCL: Improve executor and integrate scheduler (#694)

* MSCCL: improve executor and add scheduler for testing

* Use external scheduler

* Fix cmake error

* Address comments

* Fix thread safe issue

* Make MSCCL lifecycle APIs thread safe

* Make MSCCL internal scheduler aware of topology hint

* Revise error message

[ROCm/rccl commit: e3b2342f39]

projects/rccl/src/misc/msccl/msccl_lifecycle.cc

@@ -4,116 +4,192 @@
  ************************************************************************/
 
 #include <atomic>
+#include <map>
+#include <mutex>
+#include <set>
+
+#include <dirent.h>
+#include <dlfcn.h>
+#include <error.h>
+#include <link.h>
 
 #include "alloc.h"
 #include "checks.h"
+#include "graph/topo.h"
 
 #include "msccl/msccl_lifecycle.h"
 #include "msccl/msccl_parser.h"
 #include "msccl/msccl_setup.h"
 #include "msccl/msccl_status.h"
 
-RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 0);
+RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 1);
 static const char* mscclAlgoFilePathEnv = "MSCCL_ALGO_FILE_PATH";
 static std::atomic<bool> mscclInitialized;
 static bool mscclSchedulerTriedLoadAlgo = false;
+static std::mutex mscclLifecycleMutex;
 
 bool mscclEnabled() {
   return rcclParamMscclEnabled();
 }
 
-static bool mscclIsCallerFlag = false;
-
 void mscclSetIsCallerFlag() {
-  mscclIsCallerFlag = true;
+  mscclGetThreadLocalStatus().mscclIsCallerFlag = true;
 }
 
 void mscclClearIsCallerFlag() {
-  mscclIsCallerFlag = false;
+  mscclGetThreadLocalStatus().mscclIsCallerFlag = false;
 }
 
 bool mscclIsCaller() {
-  return mscclIsCallerFlag;
+  return mscclGetThreadLocalStatus().mscclIsCallerFlag;
 }
 
 bool mscclAvailable() {
   return mscclEnabled() && mscclInitialized.load(std::memory_order_acquire);
 }
 
-ncclResult_t mscclInit(ncclComm_t comm) {
-  if (comm->intraRanks > 1) {
-    mscclInitialized.store(false, std::memory_order_release);
-    INFO(NCCL_INIT, "MSCCL doesn't support multiple GPUs in one process and is not available");
-    return ncclSuccess;
+static bool mscclCommCompatible(ncclComm_t comm) {
+  std::map<uint64_t, std::set<uint64_t>> hostHashToPidHashes;
+  for (int i = 0; i < comm->nRanks; i++) {
+    uint64_t hostHash = comm->peerInfo[i].hostHash;
+    uint64_t pidHash = comm->peerInfo[i].pidHash;
+    if (hostHashToPidHashes.find(hostHash) != hostHashToPidHashes.end()) {
+      auto& pidHashSet = hostHashToPidHashes[hostHash];
+      if (pidHashSet.find(pidHash) != pidHashSet.end()) {
+        return false;
+      }
+    }
+    hostHashToPidHashes[hostHash].insert(pidHash);
+  }
+  return true;
+}
+
+static const char* mscclSchedulerPathEnv = "MSCCL_SCHEDULER";
+static const char* mscclSchedulerDefaultPath = "libmsccl-scheduler.so";
+static const char* mscclAlgoDirEnv = "MSCCL_ALGO_DIR";
+static const char* mscclAlgoDefaultDir = "msccl-algorithms";
+extern "C" bool mscclUnitTestMode() __attribute__((__weak__));
+static const char* mscclUnitTestAlgoDefaultDir = "msccl-unit-test-algorithms";
+
+static ncclResult_t mscclInternalSchedulerInit() {
+  mscclStatus& status = mscclGetStatus();
+  const char* mscclAlgoDir = getenv(mscclAlgoDirEnv);
+  std::string mscclAlgoDirStr;
+  if (mscclAlgoDir == nullptr) {
+    // Try to find default algorithm directory based on librccl.so path
+    Dl_info dl_info;
+    struct link_map *link_map_ptr = nullptr;
+    if (!dladdr1((void *)mscclInternalSchedulerInit, &dl_info, (void **)&link_map_ptr, RTLD_DL_LINKMAP)) {
+      WARN("MSCCL Internal Scheduler: dladdr1 failed");
+      return ncclInvalidUsage;
+    }
+    std::string selfLibPath = link_map_ptr->l_name;
+    mscclAlgoDirStr = selfLibPath.substr(0, selfLibPath.find_last_of("/\\") + 1);
+    mscclAlgoDirStr += (mscclUnitTestMode && mscclUnitTestMode()) ? mscclUnitTestAlgoDefaultDir : mscclAlgoDefaultDir;
+    mscclAlgoDir = mscclAlgoDirStr.c_str();
+  }
+  struct dirent *entry = nullptr;
+  DIR *dp = nullptr;
+  dp = opendir(mscclAlgoDir);
+  if (dp == nullptr) {
+    WARN("MSCCL Internal Scheduler: open algorithm directory %s failed", mscclAlgoDir);
+    return ncclInvalidUsage;
+  }
+  while ((entry = readdir(dp))) {
+    if (entry->d_type != DT_LNK && entry->d_type != DT_REG) {
+      continue;
+    }
+    status.algoMetas.emplace_back();
+    std::string fullPath = mscclAlgoDir;
+    fullPath += "/";
+    fullPath += entry->d_name;
+    NCCLCHECK(mscclGetAlgoMetaFromXmlFile(fullPath.c_str(), &(status.algoMetas.back())));
+  }
+  if (closedir(dp)) {
+    WARN("MSCCL Internal Scheduler: closedir failed, error %d", errno);
+    return ncclInvalidUsage;
+  }
+  status.rankToAlgoHandles.resize(status.algoMetas.size());
+  return ncclSuccess;
+}
+
+static ncclResult_t mscclSchedulerInit() {
+  mscclStatus& status = mscclGetStatus();
+  bool useInternalScheduler = false;
+
+  const char* mscclSchedulerPath = getenv(mscclSchedulerPathEnv);
+  if (mscclSchedulerPath) {
+    status.mscclSchedulerLib = dlopen(mscclSchedulerPath, RTLD_NOW | RTLD_LOCAL);
   } else {
+    status.mscclSchedulerLib = dlopen(mscclSchedulerDefaultPath, RTLD_NOW | RTLD_LOCAL);
+  }
+  if (status.mscclSchedulerLib == nullptr) {
+    INFO(NCCL_INIT, "MSCCL: No external scheduler found, using internal implementation");
+    useInternalScheduler = true;
+  } else {
+    status.mscclSchedulerPtr = (mscclSchedulerInterface *)dlsym(status.mscclSchedulerLib, "mscclScheduler");
+    if (status.mscclSchedulerPtr == nullptr) {
+      INFO(NCCL_INIT, "MSCCL: Failed to find mscclScheduler symbol, using internal implementation");
+      useInternalScheduler = true;
+    }
+  }
+  if (useInternalScheduler) {
+    NCCLCHECK(mscclInternalSchedulerInit());
+  } else {
+    NCCLCHECK(status.mscclSchedulerPtr->init());
+  }
+  return ncclSuccess;
+}
+
+ncclResult_t mscclInit(ncclComm_t comm) {
+  // Always initialize thread local status
+  mscclThreadLocalStatus threadLocalStatus = mscclGetThreadLocalStatus();
+  threadLocalStatus.groupStatus = mscclNoGroup;
+  threadLocalStatus.groupDepth = 0;
+  comm->mscclCompatible = mscclCommCompatible(comm);
+
+  {
+    std::lock_guard<std::mutex> lock(mscclLifecycleMutex);
+
+    if (mscclInitialized.load(std::memory_order_acquire)) {
+      return ncclSuccess;
+    }
+
+    mscclStatus& status = mscclGetStatus();
+    status.scratchBuffer = nullptr;
+    status.scratchBufferSize = 0;
+    status.workIndex = 1;
+    status.freeAlgoHandles.resize(MSCCL_MAX_NUM_ALGOS);
+    for (int i = 0; i < MSCCL_MAX_NUM_ALGOS; i++) {
+      status.freeAlgoHandles[i] = MSCCL_MAX_NUM_ALGOS - i - 1;
+    }
+    NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS));
+    status.lastStream = nullptr;
+    mscclSchedulerTriedLoadAlgo = false;
+
+    NCCLCHECK(mscclSchedulerInit());
+
     mscclInitialized.store(true, std::memory_order_release);
   }
 
-  mscclStatus& status = mscclGetStatus();
-  status.scratchBuffer = nullptr;
-  status.scratchBufferSize = 0;
-  status.rank = comm->rank;
-  status.workIndex = 1;
-  status.freeAlgoHandles.resize(MSCCL_MAX_NUM_ALGOS);
-  for (int i = 0; i < MSCCL_MAX_NUM_ALGOS; i++) {
-    status.freeAlgoHandles[i] = MSCCL_MAX_NUM_ALGOS - i - 1;
-  }
-  NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS));
-  status.groupStatus = mscclNoGroup;
-  status.groupDepth = 0;
-  mscclSchedulerTriedLoadAlgo = false;
+  INFO(NCCL_INIT, "MSCCL: Initialization finished");
   return ncclSuccess;
 }
 
 ncclResult_t mscclGroupStart() {
-  mscclStatus& status = mscclGetStatus();
-  status.groupDepth++;
-  if (status.groupStatus == mscclNoGroup) {
-    status.groupStatus = mscclGroupSupportedOp;
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
+  threadLocalStatus.groupDepth++;
+  if (threadLocalStatus.groupStatus == mscclNoGroup) {
+    threadLocalStatus.groupStatus = mscclGroupSupportedOp;
   }
   return ncclSuccess;
 }
 
-static ncclResult_t mscclScheduler(struct mscclSchedulerParam* param) {
-  static bool algoAvailable = false;
-  static mscclAlgoHandle_t loadedAlgoHandle;
-  static mscclAlgo* loadedHostAlgo = nullptr;
-
+static ncclResult_t mscclInternalSchedulerSelectAlgo(struct mscclSchedulerParam* param) {
+  mscclStatus& status = mscclGetStatus();
   param->scheduled = false;
 
-  if (!mscclSchedulerTriedLoadAlgo) {
-    mscclSchedulerTriedLoadAlgo = true;
-    const char* mscclAlgoFilePath = getenv(mscclAlgoFilePathEnv);
-    if (mscclAlgoFilePath != nullptr) {
-      NCCLCHECK(mscclLoadAlgo(mscclAlgoFilePath, &loadedAlgoHandle));
-      mscclStatus& status = mscclGetStatus();
-      loadedHostAlgo = status.hostAlgos[loadedAlgoHandle];
-      algoAvailable = true;
-    }
-  }
-  if (!algoAvailable) {
-    return ncclSuccess;
-  }
-
-  bool mscclAlgoFuncIsValid = loadedHostAlgo->func == param->func;
-  if (!mscclAlgoFuncIsValid) {
-    return ncclSuccess;
-  }
-
-  bool numGpusIsValid = loadedHostAlgo->nRanks == param->comm->nRanks;
-  if (!numGpusIsValid) {
-    return ncclSuccess;
-  }
-
-  size_t nBytes = param->count * ncclTypeSize(param->dataType) * loadedHostAlgo->sizeMultiplier;
-  bool msgSizeIsValid =
-    param->count > 0 && (param->count % loadedHostAlgo->nChunksPerLoop) == 0 &&
-    nBytes >= loadedHostAlgo->minBytes &&
-    (loadedHostAlgo->maxBytes == 0 || nBytes <= loadedHostAlgo->maxBytes);
-  if (!msgSizeIsValid) {
-    return ncclSuccess;
-  }
-
+  // Whether the algorithm is in-place
   bool isInPlace = false;
   if (param->func == mscclFuncReduce ||
       param->func == mscclFuncBroadcast ||
@@ -123,120 +199,153 @@ static ncclResult_t mscclScheduler(struct mscclSchedulerParam* param) {
     isInPlace = param->sendBuff == param->recvBuff;
   } else if (param->func == mscclFuncAllGather ||
              param->func == mscclFuncGather) {
-    isInPlace = (char*)param->sendBuff == (char*)param->recvBuff + param->comm->rank * param->count * ncclTypeSize(param->dataType);
+    isInPlace = (char*)param->sendBuff == (char*)param->recvBuff + param->rank * param->count * ncclTypeSize(param->dataType);
   } else if (param->func == mscclFuncReduceScatter ||
              param->func == mscclFuncScatter) {
-    isInPlace = (char*)param->recvBuff == (char*)param->sendBuff + param->comm->rank * param->count * ncclTypeSize(param->dataType);
-  }
-  bool inPlaceOutOfPlaceIsValid = isInPlace ? loadedHostAlgo->inPlace : loadedHostAlgo->outOfPlace;
-  if (!inPlaceOutOfPlaceIsValid) {
-    return ncclSuccess;
+    isInPlace = (char*)param->recvBuff == (char*)param->sendBuff + param->rank * param->count * ncclTypeSize(param->dataType);
+  }
+
+  // Search suitable algorithms
+  for (size_t i = 0; i < status.algoMetas.size(); i++) {
+    auto &m = status.algoMetas[i];
+    size_t nBytes = param->count * ncclTypeSize(param->dataType) * m.sizeMultiplier;
+    bool msgSizeIsValid =
+      param->count > 0 && (param->count % m.nChunksPerLoop) == 0 &&
+      nBytes >= m.minBytes && (m.maxBytes == 0 || nBytes <= m.maxBytes);
+    if (msgSizeIsValid &&
+        m.nRanks == param->nRanks &&
+        m.func == param->func &&
+        (isInPlace ? m.inPlace : m.outOfPlace)) {
+      // If not loaded for current rank, load it
+      if (status.rankToAlgoHandles[i].find(param->rank) == status.rankToAlgoHandles[i].end()) {
+        mscclAlgoHandle_t algoHandle;
+        NCCLCHECK(mscclLoadAlgo(m.filePath.c_str(), &algoHandle, param->rank));
+        status.rankToAlgoHandles[i][param->rank] = algoHandle;
+      }
+      param->handle = status.rankToAlgoHandles[i][param->rank];
+      param->scheduled = true;
+      return ncclSuccess;
+    }
   }
 
-  param->handle = loadedAlgoHandle;
-  param->scheduled = true;
   return ncclSuccess;
 }
 
-static ncclResult_t mscclSetSchedulerParam(
+static ncclResult_t mscclSchedulerSelectAlgo(struct mscclSavedSchedulerParam* param) {
+  mscclStatus& status = mscclGetStatus();
+  if (status.mscclSchedulerPtr) {
+    NCCLCHECK(status.mscclSchedulerPtr->selectAlgo(&(param->p)));
+  } else {
+    if (param->comm->topo->mscclEnabled) {
+      NCCLCHECK(mscclInternalSchedulerSelectAlgo(&(param->p)));
+    } else {
+      param->p.scheduled = false;
+    }
+  }
+  return ncclSuccess;
+}
+
+static ncclResult_t mscclSetSavedSchedulerParam(
   const void* sendBuff, const size_t sendCounts[], const size_t sDisPls[],
   void* recvBuff, const size_t recvCounts[], const size_t rDisPls[],
   size_t count, ncclDataType_t dataType, int root, int peer, ncclRedOp_t op,
   mscclFunc_t func, ncclComm_t comm, hipStream_t stream,
-  struct mscclSchedulerParam* param) {
-  param->sendBuff = sendBuff;
-  param->sendCounts = sendCounts;
-  param->sDisPls = sDisPls;
-  param->recvBuff = recvBuff;
-  param->recvCounts = recvCounts;
-  param->rDisPls = rDisPls;
-  param->count = count;
-  param->dataType = dataType;
-  param->root = root;
-  param->peer = peer;
-  param->op = op;
-  param->func = func;
+  struct mscclSavedSchedulerParam* param) {
+  param->p.sendBuff = sendBuff;
+  param->p.sendCounts = sendCounts;
+  param->p.sDisPls = sDisPls;
+  param->p.recvBuff = recvBuff;
+  param->p.recvCounts = recvCounts;
+  param->p.rDisPls = rDisPls;
+  param->p.count = count;
+  param->p.dataType = dataType;
+  param->p.root = root;
+  param->p.peer = peer;
+  param->p.op = op;
+  param->p.func = func;
+  param->p.rank = comm->rank;
+  param->p.nRanks = comm->nRanks;
   param->comm = comm;
   param->stream = stream;
   return ncclSuccess;
 }
 
-static ncclResult_t mscclSaveCountsAndDispls(struct mscclSchedulerParam* param) {
-  if (param->sendCounts) {
-    param->savedSendCounts.assign(param->sendCounts, param->sendCounts + param->comm->nRanks);
-    param->sendCounts = param->savedSendCounts.data();
-    param->savedSDisPls.assign(param->sDisPls, param->sDisPls + param->comm->nRanks);
-    param->sDisPls = param->savedSDisPls.data();
-    param->savedRecvCounts.assign(param->recvCounts, param->recvCounts + param->comm->nRanks);
-    param->recvCounts = param->savedRecvCounts.data();
-    param->savedRDisPls.assign(param->rDisPls, param->rDisPls + param->comm->nRanks);
-    param->rDisPls = param->savedRDisPls.data();
+static ncclResult_t mscclSaveCountsAndDispls(struct mscclSavedSchedulerParam* param) {
+  if (param->p.sendCounts) {
+    param->savedSendCounts.assign(param->p.sendCounts, param->p.sendCounts + param->p.nRanks);
+    param->p.sendCounts = param->savedSendCounts.data();
+    param->savedSDisPls.assign(param->p.sDisPls, param->p.sDisPls + param->p.nRanks);
+    param->p.sDisPls = param->savedSDisPls.data();
+    param->savedRecvCounts.assign(param->p.recvCounts, param->p.recvCounts + param->p.nRanks);
+    param->p.recvCounts = param->savedRecvCounts.data();
+    param->savedRDisPls.assign(param->p.rDisPls, param->p.rDisPls + param->p.nRanks);
+    param->p.rDisPls = param->savedRDisPls.data();
   }
   return ncclSuccess;
 }
 
 static ncclResult_t mscclRunSavedParams() {
-  mscclStatus& status = mscclGetStatus();
-  for (auto& param : status.savedSchedulerParams) {
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
+  for (auto& param : threadLocalStatus.savedSchedulerParams) {
     NCCLCHECK(mscclRunAlgo(
-      param.sendBuff, param.sendCounts, param.sDisPls,
-      param.recvBuff, param.recvCounts, param.rDisPls,
-      param.count, param.dataType, param.root, param.peer, param.op, param.handle, param.comm, param.stream));
+      param.p.sendBuff, param.p.sendCounts, param.p.sDisPls,
+      param.p.recvBuff, param.p.recvCounts, param.p.rDisPls,
+      param.p.count, param.p.dataType, param.p.root, param.p.peer, param.p.op, param.p.handle, param.comm, param.stream));
   }
-  status.savedSchedulerParams.clear();
+  threadLocalStatus.savedSchedulerParams.clear();
   return ncclSuccess;
 }
 
 static ncclResult_t mscclFallBackSavedParams() {
-  mscclStatus& status = mscclGetStatus();
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
   mscclSetIsCallerFlag();
-  for (auto& param : status.savedSchedulerParams) {
-    switch (param.func) {
+  for (auto& param : threadLocalStatus.savedSchedulerParams) {
+    switch (param.p.func) {
       case mscclFuncReduce:
-        NCCLCHECK(ncclReduce(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.op, param.root, param.comm, param.stream));
+        NCCLCHECK(ncclReduce(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.op, param.p.root, param.comm, param.stream));
         break;
       case mscclFuncBroadcast:
-        NCCLCHECK(ncclBroadcast(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.root, param.comm, param.stream));
+        NCCLCHECK(ncclBroadcast(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.root, param.comm, param.stream));
         break;
       case mscclFuncAllReduce:
-        NCCLCHECK(ncclAllReduce(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.op, param.comm, param.stream));
+        NCCLCHECK(ncclAllReduce(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.op, param.comm, param.stream));
         break;
       case mscclFuncReduceScatter:
-        NCCLCHECK(ncclReduceScatter(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.op, param.comm, param.stream));
+        NCCLCHECK(ncclReduceScatter(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.op, param.comm, param.stream));
         break;
       case mscclFuncAllGather:
-        NCCLCHECK(ncclAllGather(param.sendBuff, param.recvBuff, param.count, param.dataType,
+        NCCLCHECK(ncclAllGather(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
           param.comm, param.stream));
         break;
       case mscclFuncSend:
-        NCCLCHECK(ncclSend(param.sendBuff, param.count, param.dataType,
-          param.peer, param.comm, param.stream));
+        NCCLCHECK(ncclSend(param.p.sendBuff, param.p.count, param.p.dataType,
+          param.p.peer, param.comm, param.stream));
         break;
       case mscclFuncRecv:
-        NCCLCHECK(ncclRecv(param.recvBuff, param.count, param.dataType,
-          param.peer, param.comm, param.stream));
+        NCCLCHECK(ncclRecv(param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.peer, param.comm, param.stream));
         break;
       case mscclFuncGather:
-        NCCLCHECK(ncclGather(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.root, param.comm, param.stream));
+        NCCLCHECK(ncclGather(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.root, param.comm, param.stream));
         break;
       case mscclFuncScatter:
-        NCCLCHECK(ncclScatter(param.sendBuff, param.recvBuff, param.count, param.dataType,
-          param.root, param.comm, param.stream));
+        NCCLCHECK(ncclScatter(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
+          param.p.root, param.comm, param.stream));
         break;
       case mscclFuncAllToAll:
-        NCCLCHECK(ncclAllToAll(param.sendBuff, param.recvBuff, param.count, param.dataType,
+        NCCLCHECK(ncclAllToAll(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
           param.comm, param.stream));
         break;
       case mscclFuncAllToAllv:
         NCCLCHECK(ncclAllToAllv(
-          param.sendBuff, param.sendCounts, param.sDisPls,
-          param.recvBuff, param.recvCounts, param.rDisPls,
-          param.dataType, param.comm, param.stream));
+          param.p.sendBuff, param.p.sendCounts, param.p.sDisPls,
+          param.p.recvBuff, param.p.recvCounts, param.p.rDisPls,
+          param.p.dataType, param.comm, param.stream));
         break;
       default:
         WARN("Invalid MSCCL function type in saved parameter");
@@ -244,7 +353,7 @@ static ncclResult_t mscclFallBackSavedParams() {
     }
   }
   mscclClearIsCallerFlag();
-  status.savedSchedulerParams.clear();
+  threadLocalStatus.savedSchedulerParams.clear();
   return ncclSuccess;
 }
 
@@ -253,40 +362,43 @@ ncclResult_t mscclEnqueueCheck(
     void* recvBuff, const size_t recvCounts[], const size_t rDisPls[],
     size_t count, ncclDataType_t dataType, int root, int peer, ncclRedOp_t op,
     mscclFunc_t func, ncclComm_t comm, hipStream_t stream) {
-  mscclStatus& status = mscclGetStatus();
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
   hipStreamCaptureStatus captureStatus;
   unsigned long long pid;
 
-  status.savedSchedulerParams.push_back({});
-  NCCLCHECK(mscclSetSchedulerParam(
+  threadLocalStatus.savedSchedulerParams.push_back({});
+  NCCLCHECK(mscclSetSavedSchedulerParam(
     sendBuff, sendCounts, sDisPls, recvBuff, recvCounts, rDisPls,
     count, dataType, root, peer, op, func, comm, stream,
-    &status.savedSchedulerParams.back()));
+    &threadLocalStatus.savedSchedulerParams.back()));
 
-  switch (status.groupStatus) {
+  switch (threadLocalStatus.groupStatus) {
     case mscclNoGroup:
-      CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
-      if (captureStatus == hipStreamCaptureStatusNone) {
-        NCCLCHECK(mscclScheduler(&status.savedSchedulerParams.back()));
-        if (status.savedSchedulerParams.back().scheduled) {
-          NCCLCHECK(mscclRunSavedParams());
-          break;
+      if (comm->mscclCompatible) {
+        CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
+        if (captureStatus == hipStreamCaptureStatusNone) {
+          NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
+          if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {
+            NCCLCHECK(mscclRunSavedParams());
+            break;
+          }
         }
       }
       NCCLCHECK(mscclFallBackSavedParams());
       break;
     case mscclGroupSupportedOp:
-      CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
-      if (captureStatus == hipStreamCaptureStatusNone) {
-        NCCLCHECK(mscclScheduler(&status.savedSchedulerParams.back()));
-        if (status.savedSchedulerParams.back().scheduled) {
-          // Only save counts and displs when there is suitable MSCCL algorithm for this
-          NCCLCHECK(mscclSaveCountsAndDispls(&status.savedSchedulerParams.back()));
-          break;
+      if (comm->mscclCompatible) {
+        CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
+        if (captureStatus == hipStreamCaptureStatusNone) {
+          NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
+          if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {
+            // Only save counts and displs when there is suitable MSCCL algorithm for this
+            NCCLCHECK(mscclSaveCountsAndDispls(&threadLocalStatus.savedSchedulerParams.back()));
+            break;
+          }
         }
       }
-      NCCLCHECK(mscclFallBackSavedParams());
-      break;
+      threadLocalStatus.groupStatus = mscclGroupUnsupportedOp;
     case mscclGroupUnsupportedOp:
       NCCLCHECK(mscclFallBackSavedParams());
       break;
@@ -297,37 +409,71 @@ ncclResult_t mscclEnqueueCheck(
 }
 
 ncclResult_t mscclGroupEnd() {
-  mscclStatus& status = mscclGetStatus();
-  status.groupDepth--;
-  if (status.groupDepth == 0) {
-    if (status.groupStatus == mscclGroupSupportedOp) {
+  mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
+  threadLocalStatus.groupDepth--;
+  if (threadLocalStatus.groupDepth == 0) {
+    if (threadLocalStatus.groupStatus == mscclGroupSupportedOp) {
       NCCLCHECK(mscclRunSavedParams());
     }
-    status.groupStatus = mscclNoGroup;
+    threadLocalStatus.groupStatus = mscclNoGroup;
   }
   return ncclSuccess;
 }
 
-ncclResult_t mscclTeardown() {
-  if (!mscclInitialized.load(std::memory_order_acquire)) {
-    return ncclSuccess;
-  }
+static ncclResult_t mscclInternalSchedulerTeardown() {
+  ncclResult_t ret = ncclSuccess, tmpRet = ncclSuccess;
   mscclStatus& status = mscclGetStatus();
-  for (auto &p : status.hostAlgos) {
-    free(p.second);
-    status.freeAlgoHandles.push_back(p.first);
+  for (auto &m : status.rankToAlgoHandles) {
+    for (auto &p : m) {
+      tmpRet = mscclUnloadAlgo(p.second);
+      if (ret == ncclSuccess) {
+        ret = tmpRet;
+      }
+    }
   }
-  for (auto &p : status.devAlgos) {
-    CUDACHECK(hipFree(p.second));
+  status.algoMetas.clear();
+  status.rankToAlgoHandles.clear();
+  return ret;
+}
+
+ncclResult_t mscclTeardown() {
+  // Always teardown thread local status
+  mscclThreadLocalStatus threadLocalStatus = mscclGetThreadLocalStatus();
+  threadLocalStatus.savedSchedulerParams.clear();
+
+  {
+    std::lock_guard<std::mutex> lock(mscclLifecycleMutex);
+
+    if (!mscclInitialized.load(std::memory_order_acquire)) {
+      return ncclSuccess;
+    }
+    mscclStatus& status = mscclGetStatus();
+    for (auto &p : status.hostAlgos) {
+      free(p.second);
+      status.freeAlgoHandles.push_back(p.first);
+    }
+    for (auto &p : status.devAlgos) {
+      CUDACHECK(hipFree(p.second));
+    }
+    CUDACHECK(hipFree(status.scratchBuffer));
+    CUDACHECK(hipFree(status.syncFlags));
+    status.hostAlgos.clear();
+    status.devAlgos.clear();
+    status.freeAlgoHandles.clear();
+    status.scratchBuffer = nullptr;
+    status.scratchBufferSize = 0;
+    status.connectedAlgos.clear();
+    if (status.mscclSchedulerPtr) {
+      NCCLCHECK(status.mscclSchedulerPtr->teardown());
+      status.mscclSchedulerPtr = nullptr;
+      dlclose(status.mscclSchedulerLib);
+      status.mscclSchedulerLib = nullptr;
+    } else {
+      NCCLCHECK(mscclInternalSchedulerTeardown());
+    }
+    mscclInitialized.store(false, std::memory_order_release);
   }
-  CUDACHECK(hipFree(status.scratchBuffer));
-  CUDACHECK(hipFree(status.syncFlags));
-  status.hostAlgos.clear();
-  status.devAlgos.clear();
-  status.freeAlgoHandles.clear();
-  status.scratchBuffer = nullptr;
-  status.scratchBufferSize = 0;
-  status.workIndex = 1;
-  mscclInitialized.store(false, std::memory_order_release);
+
+  INFO(NCCL_INIT, "MSCCL: Teardown finished");
   return ncclSuccess;
 }

projects/rccl/src/misc/msccl/msccl_parser.cc

@@ -701,3 +701,88 @@ ncclResult_t mscclGetAlgoFromXmlFile(const char* str, struct mscclAlgo* algo, in
   free(xml);
   return ncclSuccess;
 }
+
+ncclResult_t mscclXmlLoadSingleNode(FILE* file, struct mscclXmlNode* node) {
+  memset(node, 0, sizeof(struct mscclXmlNode));
+  return mscclXmlGetNode(file, node);
+}
+
+ncclResult_t mscclAlgoMetaXmlLoad(const char* xmlFilePath, struct mscclXmlNode* node) {
+  ncclResult_t ret = ncclSuccess;
+  FILE* file = fopen(xmlFilePath, "r");
+  if (file == NULL) {
+    fprintf(stderr, "Could not open MSCCL XML algorithm file %s : %s", xmlFilePath, strerror(errno));
+    return ncclSystemError;
+  }
+  NCCLCHECK(mscclXmlLoadSingleNode(file, node));
+  fclose(file);
+  return ncclSuccess;
+}
+
+ncclResult_t mscclGetAlgoMetaFromXmlFile(const char* str, struct mscclAlgoMeta* algoMeta) {
+  ncclResult_t ret = ncclSuccess;
+  struct mscclXmlNode* node;
+  node = (struct mscclXmlNode *)malloc(sizeof(struct mscclXmlNode));
+  NCCLCHECK(mscclAlgoMetaXmlLoad(str, node));
+
+  algoMeta->filePath = str;
+
+  int nChunksPerLoop;
+  NCCLCHECK(mscclXmlGetAttrInt(node, "nchunksperloop", &nChunksPerLoop));
+  algoMeta->nChunksPerLoop  = nChunksPerLoop;
+
+  int nGpus;
+  NCCLCHECK(mscclXmlGetAttrInt(node, "ngpus", &nGpus));
+  algoMeta->nRanks = nGpus;
+
+  const char* coll;
+  NCCLCHECK(mscclXmlGetAttrStr(node, "coll", &coll));
+  algoMeta->sizeMultiplier = 1;
+  if (strcmp(coll, "reduce") == 0) {
+    algoMeta->func = mscclFuncReduce;
+  } else if (strcmp(coll, "broadcast") == 0) {
+    algoMeta->func = mscclFuncBroadcast;
+  } else if (strcmp(coll, "allreduce") == 0) {
+    algoMeta->func = mscclFuncAllReduce;
+  } else if (strcmp(coll, "reducescatter") == 0) {
+    algoMeta->sizeMultiplier = nGpus;
+    algoMeta->func = mscclFuncReduceScatter;
+  } else if (strcmp(coll, "allgather") == 0) {
+    algoMeta->sizeMultiplier = nGpus;
+    algoMeta->func = mscclFuncAllGather;
+  } else if (strcmp(coll, "send") == 0) {
+    algoMeta->func = mscclFuncSend;
+  } else if (strcmp(coll, "recv") == 0) {
+    algoMeta->func = mscclFuncRecv;
+  } else if (strcmp(coll, "gather") == 0) {
+    algoMeta->func = mscclFuncGather;
+  } else if (strcmp(coll, "scatter") == 0) {
+    algoMeta->func = mscclFuncScatter;
+  } else if (strcmp(coll, "alltoall") == 0) {
+    algoMeta->sizeMultiplier = nGpus;
+    algoMeta->func = mscclFuncAllToAll;
+  } else if (strcmp(coll, "alltoallv") == 0) {
+    algoMeta->func = mscclFuncAllToAllv;
+  } else {
+    return ncclInvalidUsage;
+  }
+
+  int64_t minBytes;
+  NCCLCHECK(mscclXmlGetAttrInt64(node, "minBytes", &minBytes));
+  algoMeta->minBytes = minBytes;
+
+  int64_t maxBytes;
+  NCCLCHECK(mscclXmlGetAttrInt64(node, "maxBytes", &maxBytes));
+  algoMeta->maxBytes = maxBytes;
+
+  int inplace;
+  NCCLCHECK(mscclXmlGetAttrInt(node, "inplace", &inplace));
+  algoMeta->inPlace = (bool)inplace;
+
+  int outofplace;
+  NCCLCHECK(mscclXmlGetAttrInt(node, "outofplace", &outofplace));
+  algoMeta->outOfPlace = (bool)outofplace;
+
+  free(node);
+  return ncclSuccess;
+}

projects/rccl/src/misc/msccl/msccl_setup.cc

@@ -258,6 +258,11 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
     ncclDataType_t dataType, ncclRedOp_t op, struct mscclAlgo* hostAlgo, struct mscclAlgo* devAlgo,
     ncclComm_t comm, hipStream_t stream) {
   mscclStatus& status = mscclGetStatus();
+
+  if (status.lastStream != stream && status.lastStream != nullptr) {
+    CUDACHECK(hipStreamWaitEvent(stream, comm->doneEvent, 0));
+  }
+
   dim3 grid = {(uint32_t)hostAlgo->nBlocks, 1, 1};
   dim3 block = {NCCL_MAX_NTHREADS, 1, 1};
   ncclDevRedOpFull opFull;
@@ -278,7 +283,8 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
 
   void *args[3] = {&comm->devComm, &devAlgo, &work};
   void *func = mscclKernelEntries[(opFull.op * ncclNumTypes + dataType) * NCCL_NUM_PROTOCOLS + hostAlgo->protocol];
-  CUDACHECK(hipExtLaunchKernel(func, grid, block, args, 0, stream, NULL, NULL,0));
+  CUDACHECK(hipExtLaunchKernel(func, grid, block, args, 0, stream, NULL, comm->doneEvent, 0));
   status.workIndex++;
+  status.lastStream = stream;
   return ncclSuccess;
 }

projects/rccl/src/misc/msccl/msccl_status.cc

@@ -9,3 +9,8 @@ mscclStatus& mscclGetStatus() {
   static mscclStatus status;
   return status;
 }
+
+mscclThreadLocalStatus& mscclGetThreadLocalStatus() {
+  static thread_local mscclThreadLocalStatus threadLocalStatus;
+  return threadLocalStatus;
+}