Integrated RCCL with MSCCL++ for small message sizes (#1231)

src/misc/msccl/msccl_lifecycle.cc

@@ -22,6 +22,10 @@
 #include "msccl/msccl_setup.h"
 #include "msccl/msccl_status.h"
 
+#ifdef ENABLE_MSCCLPP
+#include "mscclpp/mscclpp_nccl.h"
+#endif
+
 RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 1);
 RCCL_PARAM(MscclForceEnabled, "MSCCL_FORCE_ENABLE", 0);
 static const char* mscclAlgoFilePathEnv = "MSCCL_ALGO_FILE_PATH";
@@ -448,18 +452,58 @@ ncclResult_t mscclEnqueueCheck(
     count, dataType, root, peer, op, func, comm, stream,
     &threadLocalStatus.savedSchedulerParams.back()));
 
+  size_t nBytes = count * ncclTypeSize(dataType);
+
   switch (threadLocalStatus.groupStatus) {
     case mscclNoGroup:
+#ifdef ENABLE_MSCCLPP
+      if (comm->mscclppCompatible) {
+        /* check if one rank per GPU and graph mode is enabled */
+        if ((nBytes >= 32) && (threadLocalStatus.captureStatus != mscclNoCapture) && comm->mscclCompatible) {
+          if (func == mscclFuncAllReduce && nBytes <= comm->mscclpp_threshold) {
+            INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclAllReduce (groupStatus=mscclNoGroup)");
+            NCCLCHECK(mscclpp_ncclAllReduce(sendBuff, recvBuff, count, dataType, op, comm->mscclpp_comm, stream));
+            threadLocalStatus.savedSchedulerParams.clear();
+            break;
+          }
+          else if (func == mscclFuncAllGather && nBytes * comm->nRanks <= comm->mscclpp_threshold) {
+            INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclAllGather (groupStatus=mscclNoGroup)");
+            NCCLCHECK(mscclpp_ncclAllGather(sendBuff, recvBuff, count, dataType, comm->mscclpp_comm, stream));
+            threadLocalStatus.savedSchedulerParams.clear();
+            break;
+          }
+        }
+      }
+#endif
       if (comm->mscclCompatible) {
           NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
           if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {
             NCCLCHECK(mscclRunSavedParams());
             break;
           }
-        }
+      }
       NCCLCHECK(mscclFallBackSavedParams());
       break;
     case mscclGroupSupportedOp:
+#ifdef ENABLE_MSCCLPP
+      if (comm->mscclppCompatible) {
+        /* check if one rank per GPU and graph mode is enabled */
+        if ((nBytes >= 32) && (threadLocalStatus.captureStatus != mscclNoCapture) && comm->mscclCompatible) {
+          if (func == mscclFuncAllReduce && nBytes <= comm->mscclpp_threshold) {
+            INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclAllReduce (groupStatus=mscclGroupSupportedOp)");
+            NCCLCHECK(mscclpp_ncclAllReduce(sendBuff, recvBuff, count, dataType, op, comm->mscclpp_comm, stream));
+            threadLocalStatus.savedSchedulerParams.clear();
+            break;
+          }
+          else if (func == mscclFuncAllGather && nBytes * comm->nRanks <= comm->mscclpp_threshold) {
+            INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclAllGather (groupStatus=mscclGroupSupportedOp)");
+            NCCLCHECK(mscclpp_ncclAllGather(sendBuff, recvBuff, count, dataType, comm->mscclpp_comm, stream));
+            threadLocalStatus.savedSchedulerParams.clear();
+            break;
+          }
+        }
+      }
+#endif
       if (comm->mscclCompatible) {
           NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
           if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {

src/misc/mscclpp/mscclpp_nccl.cc

@@ -0,0 +1,46 @@
+/*************************************************************************
+ * Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt and NOTICES.txt for license information
+ ************************************************************************/
+
+#include "mscclpp/mscclpp_nccl.h"
+#include "debug.h"
+#include <dlfcn.h>
+#include <unordered_map>
+
+#define MSCCLPP_DECLARE(X) decltype(mscclpp_##X) mscclpp_##X = nullptr
+#define MSCCLPP_LOAD(HANDLE, X) do {                            \
+  (mscclpp_##X) = (decltype(mscclpp_##X))dlsym((HANDLE), (#X)); \
+  const char* error;                                            \
+  if ((error = dlerror()) != nullptr) {                         \
+    WARN("MSCCL++: failed to load %s : %s", (#X), error);       \
+    return false;                                               \
+  }                                                             \
+} while (false)
+
+static const char mscclpp_nccl_lib_name[] = "libmscclpp_nccl.so";
+
+MSCCLPP_DECLARE(ncclGetUniqueId);
+MSCCLPP_DECLARE(ncclCommInitRank);
+MSCCLPP_DECLARE(ncclCommDestroy);
+MSCCLPP_DECLARE(ncclAllReduce);
+MSCCLPP_DECLARE(ncclAllGather);
+
+bool mscclpp_init() {
+  void* handle = dlopen(mscclpp_nccl_lib_name, RTLD_LAZY);
+  if (!handle) {
+    WARN("MSCCL++: failed to access %s : %s", mscclpp_nccl_lib_name, dlerror());
+    return false;
+  }
+  dlerror(); // Clear any errors.
+
+  MSCCLPP_LOAD(handle, ncclGetUniqueId);
+  MSCCLPP_LOAD(handle, ncclCommInitRank);
+  MSCCLPP_LOAD(handle, ncclCommDestroy);
+  MSCCLPP_LOAD(handle, ncclAllReduce);
+  MSCCLPP_LOAD(handle, ncclAllGather);
+  return true;
+}
+
+std::unordered_map<ncclUniqueId, mscclpp_ncclUniqueId> mscclpp_uniqueIdMap;