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

src/include/comm.h

@@ -393,6 +393,13 @@ struct ncclComm {
   // shared structures for finalization
   int finalizeRankCnt;
 
+#if defined(ENABLE_MSCCLPP)
+  // Whether this comm is compatible with MSCCLPP
+  bool mscclppCompatible;
+  struct mscclpp_ncclComm* mscclpp_comm;
+  size_t mscclpp_threshold;
+#endif
+
   // Whether this comm is compatible with MSCCL
   bool mscclCompatible;
   // group job to support multi-thread FT

src/include/mscclpp/mscclpp_nccl.h

@@ -0,0 +1,48 @@
+/*************************************************************************
+ * Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt and NOTICES.txt for license information
+ ************************************************************************/
+
+#ifndef MSCCLPP_NCCL_H_
+#define MSCCLPP_NCCL_H_
+
+#include "nccl.h"
+#include <unordered_map>
+
+typedef struct mscclpp_ncclComm* mscclpp_ncclComm_t;
+
+typedef struct { char internal[NCCL_UNIQUE_ID_BYTES]; } mscclpp_ncclUniqueId;
+
+bool mscclpp_init();
+
+/* A ncclUniqueId and a mscclpp_ncclUniqueId will always be created together and used alternatively. This maps between them. */
+extern std::unordered_map<ncclUniqueId, mscclpp_ncclUniqueId> mscclpp_uniqueIdMap;
+
+/* See ncclGetUniqueId. */
+extern ncclResult_t  (*mscclpp_ncclGetUniqueId)(mscclpp_ncclUniqueId* uniqueId);
+
+/* See ncclCommInitRank. */
+extern ncclResult_t  (*mscclpp_ncclCommInitRank)(mscclpp_ncclComm_t* comm, int nranks, mscclpp_ncclUniqueId commId, int rank);
+
+/* See ncclCommDestroy. */
+extern ncclResult_t  (*mscclpp_ncclCommDestroy)(mscclpp_ncclComm_t comm);
+
+/* See ncclAllReduce. */
+extern ncclResult_t  (*mscclpp_ncclAllReduce)(const void* sendbuff, void* recvbuff, size_t count,
+    ncclDataType_t datatype, ncclRedOp_t op, mscclpp_ncclComm_t comm, hipStream_t stream);
+
+/* See ncclAllGather. */
+extern ncclResult_t  (*mscclpp_ncclAllGather)(const void* sendbuff, void* recvbuff, size_t sendcount,
+    ncclDataType_t datatype, mscclpp_ncclComm_t comm, hipStream_t stream);
+
+namespace std {
+  template <>
+  struct hash<ncclUniqueId> {
+    size_t operator ()(const ncclUniqueId& uniqueId) const noexcept;
+  };
+}
+
+bool operator ==(const ncclUniqueId& a, const ncclUniqueId& b);
+
+#endif

src/init.cc

@@ -45,6 +45,9 @@
 #include "hip_rocm_version_info.h"
 //#include "clique/CliqueManager.h"
 //#include <hsa/hsa_ext_amd.h>
+#ifdef ENABLE_MSCCLPP
+#include "mscclpp/mscclpp_nccl.h"
+#endif
 // [/RCCL]
 
 #include "msccl/msccl_lifecycle.h"
@@ -89,6 +92,16 @@ static uint64_t hashUniqueId(ncclUniqueId const &id) {
   return h;
 }
 
+#ifdef ENABLE_MSCCLPP
+size_t std::hash<ncclUniqueId>::operator ()(const ncclUniqueId& uniqueId) const noexcept {
+  return (size_t)hashUniqueId(uniqueId);
+}
+
+bool operator ==(const ncclUniqueId& a, const ncclUniqueId& b) {
+  return memcmp(a.internal, b.internal, NCCL_UNIQUE_ID_BYTES) == 0;
+}
+#endif
+
 // GDRCOPY support: Off by default
 NCCL_PARAM(GdrCopyEnable, "GDRCOPY_ENABLE", 0);
 
@@ -149,6 +162,11 @@ static ncclResult_t ncclInit() {
     }
 #ifndef NVTX_NO_IMPL
     initNvtxRegisteredEnums();
+#endif
+#ifdef ENABLE_MSCCLPP
+    if (!mscclpp_init()) {
+      return ncclSystemError;
+    }
 #endif
     __atomic_store_n(&initialized, true, __ATOMIC_RELEASE);
   }
@@ -163,12 +181,32 @@ ncclResult_t ncclGetVersion(int* version) {
   return ncclSuccess;
 }
 
+#ifdef ENABLE_MSCCLPP
+RCCL_PARAM(EnableMscclpp, "ENABLE_MSCCLPP", 0);
+RCCL_PARAM(MscclppThreshold, "MSCCLPP_THRESHOLD", (size_t)(1024*1024));
+#endif
+
 NCCL_API(ncclResult_t, ncclGetUniqueId, ncclUniqueId* out);
 ncclResult_t ncclGetUniqueId(ncclUniqueId* out) {
   NCCLCHECK(ncclInit());
   NCCLCHECK(PtrCheck(out, "GetUniqueId", "out"));
   ncclResult_t res = bootstrapGetUniqueId((struct ncclBootstrapHandle*)out);
   TRACE_CALL("ncclGetUniqueId(0x%llx)", (unsigned long long)hashUniqueId(*out));
+#ifdef ENABLE_MSCCLPP
+  if (rcclParamEnableMscclpp()) {
+    NCCLCHECK(res);
+    int dev;
+    CUDACHECK(cudaGetDevice(&dev));
+    hipDeviceProp_t devProp;
+    CUDACHECK(hipGetDeviceProperties(&devProp, dev));
+    if (IsArchMatch(devProp.gcnArchName, "gfx94")) {
+      INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclGetUniqueId");
+      res = mscclpp_ncclGetUniqueId(&(mscclpp_uniqueIdMap[*out]));
+    } else {
+      WARN("MSCCL++: Cannot enable MSCCL++ on %s architecture", devProp.gcnArchName);
+    }
+  }
+#endif
   return res;
 }
 
@@ -1930,6 +1968,24 @@ static ncclResult_t ncclCommInitRankFunc(struct ncclAsyncJob* job_) {
 
   NCCLCHECKGOTO(initTransportsRank(comm, job->parent), res, fail);
 
+#ifdef ENABLE_MSCCLPP
+  if (rcclParamEnableMscclpp()) {
+    hipDeviceProp_t devProp;
+    CUDACHECK(hipGetDeviceProperties(&devProp, cudaDev));
+    comm->mscclppCompatible = IsArchMatch(devProp.gcnArchName, "gfx94");
+    if (comm->mscclppCompatible) {
+      NCCLCHECKGOTO(bootstrapIntraNodeBroadcast(comm->bootstrap, comm->localRankToRank, comm->localRank, comm->localRanks, 0, &(mscclpp_uniqueIdMap[job->commId]), sizeof(mscclpp_ncclUniqueId)), res, fail);
+      INFO(NCCL_INIT, "MSCCL++: Broadcast mscclpp_ncclUniqueId to %d ranks", (comm->localRanks - 1));
+      comm->mscclpp_threshold = rcclParamMscclppThreshold();
+      INFO(NCCL_INIT, "MSCCL++: Enabled! Msg size threshold=%zu", comm->mscclpp_threshold);
+      INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclCommInitRank (nranks=%d)", job->nranks);
+      NCCLCHECKGOTO(mscclpp_ncclCommInitRank(&(comm->mscclpp_comm), job->nranks, mscclpp_uniqueIdMap[job->commId], job->myrank), res, fail);
+    } else {
+      WARN("MSCCL++: Cannot enable MSCCL++ on %s architecture", devProp.gcnArchName);
+    }
+  }
+#endif
+
   NCCLCHECKGOTO(ncclLoadTunerPlugin(&comm->tuner), res, fail);
   if (comm->tuner) {
     NCCLCHECK(comm->tuner->init(comm->nRanks, comm->nNodes, ncclDebugLog));
@@ -2528,6 +2584,18 @@ ncclResult_t ncclCommDestroy(ncclComm_t comm) {
     return ncclSuccess;
   }
 
+#ifdef ENABLE_MSCCLPP
+  if (comm->mscclppCompatible) {
+    INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclCommDestroy");
+    ncclResult_t res = mscclpp_ncclCommDestroy(comm->mscclpp_comm);
+    if (res != ncclSuccess) {
+      WARN("MSCCL++: mscclpp_ncclCommDestroy failed (%s)", ncclGetErrorString(res));
+    }
+    comm->mscclppCompatible = false;
+    comm->mscclpp_comm = nullptr;
+  }
+#endif
+
   int rank = comm->rank, nranks = comm->nRanks, cudaDev = comm->cudaDev;
 
   NvtxParamsCommInitRank payload{rank, nranks, cudaDev};

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;