nccl/rccl integration

apps/nccl/src/nccl.cu

@@ -16,6 +16,8 @@
 #if defined(ENABLE_NPKIT)
 #include <mscclpp/npkit/npkit.hpp>
 #endif
+#include <dlfcn.h>
+
 #include "allgather.hpp"
 #include "allreduce.hpp"
 #include "broadcast.hpp"
@@ -35,6 +37,101 @@
 
 #define NUM_CHANNELS_PER_CONNECTION 64
 
+typedef enum mscclppNcclDlopenErr {
+  dlopenSuccess = 0,
+  dlopenError = 1,
+} mscclppNcclDlopenErr_t;
+
+typedef struct _mscclppNcclOps_t {
+  ncclResult_t (*CommInitRank)(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
+  ncclResult_t (*GetUniqueId)(ncclUniqueId* uniqueId);
+  ncclResult_t (*CommDestroy)(ncclComm_t comm);
+  ncclResult_t (*CommUserRank)(const ncclComm_t, int* rank);
+  ncclResult_t (*AllReduce)(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, ncclRedOp_t op,
+                            ncclComm_t comm, cudaStream_t stream);
+  ncclResult_t (*AllGather)(const void* sendbuff, void* recvbuff, size_t sendcount, ncclDataType_t datatype,
+                            ncclComm_t comm, cudaStream_t stream);
+  ncclResult_t (*Broadcast)(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, int root,
+                            ncclComm_t comm, cudaStream_t stream);
+  ncclResult_t (*ReduceScatter)(const void* sendbuff, void* recvbuff, size_t recvcount, ncclDataType_t datatype,
+                                ncclRedOp_t op, ncclComm_t comm, cudaStream_t stream);
+} mscclppNcclOps_t;
+
+mscclppNcclOps_t mscclppNcclOps;
+void* mscclppNcclDlHandle = NULL;
+bool mscclppNcclDlopenSharedLib = false;
+
+#define QUOTE(symbol) #symbol
+
+#define NCCL_DLSYM(_struct_, _handle_, _prefix_, _function_, _type_)                               \
+  do {                                                                                             \
+    _struct_._function_ = (_type_)dlsym((_handle_), QUOTE(_prefix_##_function_));                  \
+    if (_struct_._function_ == NULL) {                                                             \
+      printf("Failed: dlsym error: Cannot open %s: %s\n", QUOTE(_prefix_##_function_), dlerror()); \
+      exit(dlopenError);                                                                           \
+    }                                                                                              \
+  } while (0)
+
+static inline int mscclppNcclDlopenInit() {
+  const char* ncclLibPath = mscclpp::env()->ncclSharedLibPath.c_str();
+  if (ncclLibPath != nullptr && ncclLibPath[0] != '\0') {
+    if (std::filesystem::is_directory(ncclLibPath)) {
+      WARN("The value of the environment variable %s is a directory", ncclLibPath);
+      return dlopenError;
+    }
+
+    mscclppNcclDlHandle = dlopen(ncclLibPath, RTLD_LAZY | RTLD_NODELETE);
+    if (!mscclppNcclDlHandle) {
+      WARN("Cannot open the shared library specified by MSCCLPP_NCCL_LIB_PATH: %s\n", dlerror());
+      return dlopenError;
+    }
+  } else {
+    WARN("The value of MSCCLPP_NCCL_LIB_PATH is empty!\n");
+    return dlopenError;
+  }
+
+  NCCL_DLSYM(mscclppNcclOps, mscclppNcclDlHandle, nccl, CommInitRank,
+             ncclResult_t(*)(ncclComm_t*, int, ncclUniqueId, int));
+  NCCL_DLSYM(mscclppNcclOps, mscclppNcclDlHandle, nccl, GetUniqueId, ncclResult_t(*)(ncclUniqueId*));
+  NCCL_DLSYM(mscclppNcclOps, mscclppNcclDlHandle, nccl, CommDestroy, ncclResult_t(*)(ncclComm_t));
+  NCCL_DLSYM(mscclppNcclOps, mscclppNcclDlHandle, nccl, CommUserRank, ncclResult_t(*)(ncclComm_t, int*));
+  NCCL_DLSYM(mscclppNcclOps, mscclppNcclDlHandle, nccl, AllReduce,
+             ncclResult_t(*)(const void*, void*, size_t, ncclDataType_t, ncclRedOp_t, ncclComm_t, cudaStream_t));
+  NCCL_DLSYM(mscclppNcclOps, mscclppNcclDlHandle, nccl, AllGather,
+             ncclResult_t(*)(const void*, void*, size_t, ncclDataType_t, ncclComm_t, cudaStream_t));
+  NCCL_DLSYM(mscclppNcclOps, mscclppNcclDlHandle, nccl, Broadcast,
+             ncclResult_t(*)(const void*, void*, size_t, ncclDataType_t, int, ncclComm_t, cudaStream_t));
+  NCCL_DLSYM(mscclppNcclOps, mscclppNcclDlHandle, nccl, ReduceScatter,
+             ncclResult_t(*)(const void*, void*, size_t, ncclDataType_t, ncclRedOp_t, ncclComm_t, cudaStream_t));
+
+  return dlopenSuccess;
+}
+
+static inline void mscclppNcclDlopenFinalize() {
+  if (mscclppNcclDlHandle) {
+    dlclose(mscclppNcclDlHandle);
+  }
+}
+
+static inline int mscclppNcclInFallbackList(const char* collOps, const char* fallbackList) {
+  if (fallbackList == nullptr || fallbackList[0] == '\0' || strcmp(fallbackList, "all") == 0) {
+    return 1;
+  }
+
+  char* fallbackListCopy = strdup(fallbackList);
+  char* token = strtok(fallbackListCopy, ",");
+  while (token != NULL) {
+    if (strcmp(collOps, token) == 0) {
+      free(fallbackListCopy);
+      return 1;
+    }
+    token = strtok(NULL, ",");
+  }
+
+  free(fallbackListCopy);
+  return 0;
+}
+
 // static const mscclpp::Transport IBs[] = {mscclpp::Transport::IB0, mscclpp::Transport::IB1, mscclpp::Transport::IB2,
 //                             mscclpp::Transport::IB3, mscclpp::Transport::IB4, mscclpp::Transport::IB5,
 //                             mscclpp::Transport::IB6, mscclpp::Transport::IB7};
@@ -96,6 +193,8 @@ struct ncclComm {
 
   uint32_t numScratchBuff;
   uint32_t buffFlag;
+
+  void* mscclppNcclComm;
 };
 
 static size_t ncclTypeSize(ncclDataType_t type) {
@@ -501,6 +600,34 @@ NCCL_API ncclResult_t ncclCommInitRank(ncclComm_t* comm, int nranks, ncclUniqueI
     NpKit::Init(rank);
   }
 #endif
+
+  const bool mscclppEnableNcclFallback = mscclpp::env()->enableNcclFallback;
+  if (mscclppEnableNcclFallback == true && mscclppNcclDlHandle == NULL) {
+    int dlopenStatus = mscclppNcclDlopenInit();
+    if (dlopenStatus == dlopenSuccess) {
+      mscclppNcclDlopenSharedLib = true;
+    } else {
+      return ncclInternalError;
+    }
+  }
+
+  if (mscclppNcclDlopenSharedLib == true) {
+    ncclUniqueId mscclppNcclUniqueId;
+    if (rank == 0) {
+      mscclppNcclOps.GetUniqueId(&mscclppNcclUniqueId);
+    }
+    // After broadcast, mscclppNcclUniqueId on each rank has the same ncclUniqueId
+    bootstrap->broadcast(&mscclppNcclUniqueId, sizeof(ncclUniqueId), 0);
+
+    commPtr->mscclppNcclComm = new ncclComm_t();
+    if (commPtr->mscclppNcclComm == nullptr) {
+      WARN("Failed to allocate memory for mscclppNcclComm");
+      return ncclInternalError;
+    }
+    mscclppNcclOps.CommInitRank(reinterpret_cast<ncclComm_t*>(commPtr->mscclppNcclComm), nranks, mscclppNcclUniqueId,
+                                rank);
+  }
+
   return ncclSuccess;
 }
 
@@ -527,6 +654,13 @@ NCCL_API ncclResult_t ncclCommDestroy(ncclComm_t comm) {
     NpKit::Shutdown();
   }
 #endif
+
+  if (mscclppNcclDlopenSharedLib == true) {
+    mscclppNcclOps.CommDestroy(*reinterpret_cast<ncclComm_t*>(comm->mscclppNcclComm));
+    mscclppNcclDlopenFinalize();
+    delete static_cast<ncclComm_t*>(comm->mscclppNcclComm);
+  }
+
   delete comm;
   return ncclSuccess;
 }
@@ -628,6 +762,11 @@ NCCL_API ncclResult_t ncclCommUserRank(const ncclComm_t comm, int* rank) {
     WARN("comm is nullptr or rank is nullptr");
     return ncclInvalidArgument;
   }
+
+  if (mscclppNcclDlopenSharedLib == true) {
+    return mscclppNcclOps.CommUserRank(*reinterpret_cast<ncclComm_t*>(comm->mscclppNcclComm), rank);
+  }
+
   *rank = comm->comm->bootstrap()->getRank();
   return ncclSuccess;
 }
@@ -715,6 +854,12 @@ NCCL_API ncclResult_t ncclBroadcast(const void* sendbuff, void* recvbuff, size_t
     return ncclInvalidArgument;
   }
 
+  const char* fallbackList = mscclpp::env()->forceNcclFallbackOperation.c_str();
+  if (mscclppNcclDlopenSharedLib == true && mscclppNcclInFallbackList("broadcast", fallbackList)) {
+    return mscclppNcclOps.Broadcast(sendbuff, recvbuff, count, datatype, root,
+                                    *reinterpret_cast<ncclComm_t*>(comm->mscclppNcclComm), stream);
+  }
+
   int rank = comm->comm->bootstrap()->getRank();
 
   std::vector<executionPlanInstance>& plans = comm->executionPlans["broadcast"];
@@ -767,6 +912,12 @@ NCCL_API ncclResult_t ncclAllReduce(const void* sendbuff, void* recvbuff, size_t
     return ncclInvalidArgument;
   }
 
+  const char* fallbackList = mscclpp::env()->forceNcclFallbackOperation.c_str();
+  if (mscclppNcclDlopenSharedLib == true && mscclppNcclInFallbackList("allreduce", fallbackList)) {
+    return mscclppNcclOps.AllReduce(sendbuff, recvbuff, count, datatype, reductionOperation,
+                                    *reinterpret_cast<ncclComm_t*>(comm->mscclppNcclComm), stream);
+  }
+
   // Declarating variables
   size_t bytes = count * ncclTypeSize(datatype);
   int rank = comm->comm->bootstrap()->getRank();
@@ -811,7 +962,7 @@ NCCL_API ncclResult_t ncclAllReduce(const void* sendbuff, void* recvbuff, size_t
 }
 
 NCCL_API ncclResult_t ncclReduceScatter(const void* sendbuff, void* recvbuff, size_t recvcount, ncclDataType_t datatype,
-                                        ncclRedOp_t, ncclComm_t comm, cudaStream_t stream) {
+                                        ncclRedOp_t op, ncclComm_t comm, cudaStream_t stream) {
   size_t bytes = recvcount * ncclTypeSize(datatype);
   if (sendbuff == nullptr || recvbuff == nullptr || bytes == 0 || comm == nullptr) {
     WARN(
@@ -820,6 +971,12 @@ NCCL_API ncclResult_t ncclReduceScatter(const void* sendbuff, void* recvbuff, si
     return ncclInvalidArgument;
   }
 
+  const char* fallbackList = mscclpp::env()->forceNcclFallbackOperation.c_str();
+  if (mscclppNcclDlopenSharedLib == true && mscclppNcclInFallbackList("reducescatter", fallbackList)) {
+    return mscclppNcclOps.ReduceScatter(sendbuff, recvbuff, recvcount, datatype, op,
+                                        *reinterpret_cast<ncclComm_t*>(comm->mscclppNcclComm), stream);
+  }
+
   int rank = comm->comm->bootstrap()->getRank();
   int nRank = comm->comm->bootstrap()->getNranks();
 
@@ -876,6 +1033,12 @@ NCCL_API ncclResult_t ncclAllGather(const void* sendbuff, void* recvbuff, size_t
     return ncclInvalidArgument;
   }
 
+  const char* fallbackList = mscclpp::env()->forceNcclFallbackOperation.c_str();
+  if (mscclppNcclDlopenSharedLib == true && mscclppNcclInFallbackList("allgather", fallbackList)) {
+    return mscclppNcclOps.AllGather(sendbuff, recvbuff, sendcount, datatype,
+                                    *reinterpret_cast<ncclComm_t*>(comm->mscclppNcclComm), stream);
+  }
+
   int rank = comm->comm->bootstrap()->getRank();
   int nRank = comm->comm->bootstrap()->getNranks();
 

include/mscclpp/core.hpp

@@ -114,6 +114,17 @@ class TcpBootstrap : public Bootstrap {
   /// @param size The size of the data each rank sends.
   void allGather(void* allData, int size) override;
 
+  /// Broadcast data from the root process to all processes using a ring-based algorithm.
+  ///
+  /// When called by the root rank, this sends the `size` bytes starting at memory location `data` to all other
+  /// ranks. Non-root ranks receive these bytes into their own `data` buffer, overwriting its previous contents.
+  /// The data propagates sequentially through a logical ring of processes until all ranks have received it.
+  ///
+  /// @param data Pointer to the send buffer (root) or receive buffer (non-root)
+  /// @param size Number of bytes to broadcast
+  /// @param root Rank initiating the broadcast
+  void broadcast(void* data, int size, int root);
+
   /// Synchronize all processes.
   void barrier() override;
 

include/mscclpp/env.hpp

@@ -30,6 +30,9 @@ class Env {
   const std::string executionPlanDir;
   const std::string npkitDumpDir;
   const bool cudaIpcUseDefaultStream;
+  const std::string ncclSharedLibPath;
+  const std::string forceNcclFallbackOperation;
+  const bool enableNcclFallback;
   const bool disableChannelCache;
 
  private:

src/bootstrap/bootstrap.cc

@@ -83,6 +83,7 @@ class TcpBootstrap::Impl {
   int getNranks();
   int getNranksPerNode();
   void allGather(void* allData, int size);
+  void broadcast(void* data, int size, int root);
   void send(void* data, int size, int peer, int tag);
   void recv(void* data, int size, int peer, int tag);
   void barrier();
@@ -465,6 +466,40 @@ void TcpBootstrap::Impl::allGather(void* allData, int size) {
   TRACE(MSCCLPP_INIT, "rank %d nranks %d size %d - DONE", rank, nRanks, size);
 }
 
+void TcpBootstrap::Impl::broadcast(void* data, int size, int root) {
+  int rank = rank_;
+  int nRanks = nRanks_;
+
+  if (nRanks == 1) return;
+
+  TRACE(MSCCLPP_INIT, "rank %d nranks %d root %d size %d", rank, nRanks, root, size);
+
+  /*Ring-based broadcast, propagate data in nRanks-1 steps*/
+
+  // For ring-based broadcast, propagate data in nRanks-1 steps
+  for (int step = 0; step < nRanks - 1; step++) {
+    if (rank == root) {
+      // Root sends data to next rank in first step
+      if (step == 0) {
+        netSend(ringSendSocket_.get(), data, size);
+      }
+    } else {
+      // Calculate when this rank should receive data
+      int receiveStep = (rank - root - 1 + nRanks) % nRanks;
+      if (step == receiveStep) {
+        // Receive from previous rank
+        netRecv(ringRecvSocket_.get(), data, size);
+        // Forward to next rank (if not last step)
+        if (step < nRanks - 2) {
+          netSend(ringSendSocket_.get(), data, size);
+        }
+      }
+    }
+  }
+
+  TRACE(MSCCLPP_INIT, "rank %d nranks %d root %d size %d - DONE", rank, nRanks, root, size);
+}
+
 std::shared_ptr<Socket> TcpBootstrap::Impl::getPeerSendSocket(int peer, int tag) {
   auto it = peerSendSockets_.find(std::make_pair(peer, tag));
   if (it != peerSendSockets_.end()) {
@@ -555,6 +590,8 @@ MSCCLPP_API_CPP void TcpBootstrap::recv(void* data, int size, int peer, int tag)
 
 MSCCLPP_API_CPP void TcpBootstrap::allGather(void* allData, int size) { pimpl_->allGather(allData, size); }
 
+MSCCLPP_API_CPP void TcpBootstrap::broadcast(void* data, int size, int root) { pimpl_->broadcast(data, size, root); }
+
 MSCCLPP_API_CPP void TcpBootstrap::initialize(UniqueId uniqueId, int64_t timeoutSec) {
   pimpl_->initialize(uniqueId, timeoutSec);
 }

src/env.cpp

@@ -62,6 +62,9 @@ Env::Env()
       executionPlanDir(readEnv<std::string>("MSCCLPP_EXECUTION_PLAN_DIR", "")),
       npkitDumpDir(readEnv<std::string>("MSCCLPP_NPKIT_DUMP_DIR", "")),
       cudaIpcUseDefaultStream(readEnv<bool>("MSCCLPP_CUDAIPC_USE_DEFAULT_STREAM", false)),
+      ncclSharedLibPath(readEnv<std::string>("MSCCLPP_NCCL_LIB_PATH", "")),
+      forceNcclFallbackOperation(readEnv<std::string>("MSCCLPP_FORCE_NCCL_FALLBACK_OPERATION", "")),
+      enableNcclFallback(readEnv<bool>("MSCCLPP_ENABLE_NCCL_FALLBACK", false)),
       disableChannelCache(readEnv<bool>("MSCCLPP_DISABLE_CHANNEL_CACHE", false)) {}
 
 std::shared_ptr<Env> env() {
@@ -81,6 +84,9 @@ std::shared_ptr<Env> env() {
     logEnv("MSCCLPP_EXECUTION_PLAN_DIR", globalEnv->executionPlanDir);
     logEnv("MSCCLPP_NPKIT_DUMP_DIR", globalEnv->npkitDumpDir);
     logEnv("MSCCLPP_CUDAIPC_USE_DEFAULT_STREAM", globalEnv->cudaIpcUseDefaultStream);
+    logEnv("MSCCLPP_NCCL_LIB_PATH", globalEnv->ncclSharedLibPath);
+    logEnv("MSCCLPP_FORCE_NCCL_FALLBACK_OPERATION", globalEnv->forceNcclFallbackOperation);
+    logEnv("MSCCLPP_ENABLE_NCCL_FALLBACK", globalEnv->enableNcclFallback);
     logEnv("MSCCLPP_DISABLE_CHANNEL_CACHE", globalEnv->disableChannelCache);
   }
   return globalEnv;