Merge remote-tracking branch 'nccl/master' into develop

src/include/alloc.h

@@ -11,28 +11,40 @@
 #include "nccl.h"
 #include "checks.h"
 #include "align.h"
+#include "utils.h"
 #include <sys/mman.h>
 #include <unistd.h>
 #include <stdlib.h>
 #include <string.h>
 #include "rccl_vars.h"
 
+uint64_t clockNano(); // from utils.h with which we have a circular dependency
+
 template <typename T>
-static ncclResult_t ncclCudaHostCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
-  CUDACHECK(hipHostMalloc(ptr, nelem*sizeof(T), hipHostMallocMapped));
+ncclResult_t ncclCudaHostCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
+  ncclResult_t result = ncclSuccess;
+  uint64_t time = 0;
+  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  *ptr = nullptr;
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  time = clockNano();
+  CUDACHECKGOTO(hipHostMalloc(ptr, nelem*sizeof(T), hipHostMallocMapped), result, finish);
+  time = clockNano() - time;
   memset(*ptr, 0, nelem*sizeof(T));
-  INFO(NCCL_ALLOC, "%s:%d Cuda Host Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
-  return ncclSuccess;
+  INFO(NCCL_ALLOC, "%s:%d Cuda Host Alloc Size %ld pointer %p seconds: hipHostAlloc=%g", filefunc, line, nelem*sizeof(T), *ptr, double(time)/1.e9);
+finish:
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  return result;
 }
 #define ncclCudaHostCalloc(...) ncclCudaHostCallocDebug(__VA_ARGS__, __FILE__, __LINE__)
 
-static inline ncclResult_t ncclCudaHostFree(void* ptr) {
+inline ncclResult_t ncclCudaHostFree(void* ptr) {
   CUDACHECK(hipHostFree(ptr));
   return ncclSuccess;
 }
 
 template <typename T>
-static ncclResult_t ncclCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
+ncclResult_t ncclCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
   void* p = malloc(nelem*sizeof(T));
   if (p == NULL) {
     WARN("Failed to malloc %ld bytes", nelem*sizeof(T));
@@ -46,7 +58,7 @@ static ncclResult_t ncclCallocDebug(T** ptr, size_t nelem, const char *filefunc,
 #define ncclCalloc(...) ncclCallocDebug(__VA_ARGS__, __FILE__, __LINE__)
 
 template <typename T>
-static ncclResult_t ncclRealloc(T** ptr, size_t oldNelem, size_t nelem) {
+ncclResult_t ncclRealloc(T** ptr, size_t oldNelem, size_t nelem) {
   if (nelem < oldNelem) return ncclInternalError;
   if (nelem == oldNelem) return ncclSuccess;
 
@@ -78,54 +90,126 @@ static_assert(sizeof(struct allocationTracker) == 64, "allocationTracker must be
 extern struct allocationTracker allocTracker[];
 
 template <typename T>
-static ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, bool isFineGrain = false) {
-
-  // Need async stream for P2P pre-connect + CUDA Graph
-  static bool streamCreated = false;
-  static hipStream_t stream;
-  if (rcclParamEnableHipGraph() && !streamCreated)
-  {
-    // Create stream only once to avoid performance penalty
-    CUDACHECK(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
-    streamCreated = true;
-  }
-
+ncclResult_t ncclCudaMallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, bool isFineGrain = false) {
+  ncclResult_t result = ncclSuccess;
+  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  *ptr = nullptr;
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  uint64_t time = clockNano();
   if (isFineGrain)
-    CUDACHECK(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained));
+    CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
   else
-    CUDACHECK(hipMalloc(ptr, nelem*sizeof(T)));
+    CUDACHECKGOTO(hipMalloc(ptr, nelem*sizeof(T)), result, finish);
+  time = clockNano() - time;
+finish:
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p seconds: hipMalloc=%g", filefunc, line, nelem*sizeof(T), *ptr, double(time)/1.e9);
+  return result;
+}
+#define ncclCudaMalloc(...) ncclCudaMallocDebug( __FILE__, __LINE__, __VA_ARGS__)
 
-  if (rcclParamEnableHipGraph()) {
-    CUDACHECK(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream));
-    CUDACHECK(hipStreamSynchronize(stream));
-    // NOTE: Currently the re-used stream is not destroyed
-    //CUDACHECK(hipStreamDestroy(stream));
-  } else {
-    CUDACHECK(hipMemset(*ptr, 0, nelem*sizeof(T)));
-    CUDACHECK(hipStreamSynchronize(NULL));
-  }
-
-  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
+template <typename T>
+ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, bool isFineGrain = false) {
+  ncclResult_t result = ncclSuccess;
+  uint64_t time0=0, time1=0, time2=0;
+  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  *ptr = nullptr;
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  // Need a side stream so as not to interfere with graph capture.
+  hipStream_t stream;
+  time0 = clockNano();
+  CUDACHECK(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
+  time1 = clockNano();
+  if (isFineGrain)
+    CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
+  else
+    CUDACHECKGOTO(hipMalloc(ptr, nelem*sizeof(T)), result, finish);
+  time2 = clockNano();
+  CUDACHECKGOTO(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
+  CUDACHECKGOTO(hipStreamSynchronize(stream), result, finish);
+  CUDACHECKGOTO(hipStreamDestroy(stream), result, finish);
   int dev;
   CUDACHECK(hipGetDevice(&dev));
   if (dev < MAX_ALLOC_TRACK_NGPU) {
-    __atomic_fetch_add(&allocTracker[dev].totalAlloc, 1, __ATOMIC_SEQ_CST);
-    __atomic_fetch_add(&allocTracker[dev].totalAllocSize, nelem*sizeof(T), __ATOMIC_SEQ_CST);
+    __atomic_fetch_add(&allocTracker[dev].totalAlloc, 1, __ATOMIC_RELAXED);
+    __atomic_fetch_add(&allocTracker[dev].totalAllocSize, nelem*sizeof(T), __ATOMIC_RELAXED);
   }
-  return ncclSuccess;
+  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p seconds: hipStreamCreateWithFlags=%g hipMalloc=%g", filefunc, line, nelem*sizeof(T), *ptr, double(time1-time0)/1.e9, double(time2-time1)/1.e9);
+finish:
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  return result;
 }
 #define ncclCudaCalloc(...) ncclCudaCallocDebug(__FILE__, __LINE__, __VA_ARGS__)
 
 template <typename T>
-static ncclResult_t ncclCudaMemcpy(T* dst, T* src, size_t nelem) {
-  CUDACHECK(hipMemcpy(dst, src, nelem*sizeof(T), hipMemcpyDefault));
-  return ncclSuccess;
+ncclResult_t ncclCudaCallocAsyncDebug(const char *filefunc, int line, T** ptr, size_t nelem, hipStream_t stream, bool isFineGrain = false) {
+  ncclResult_t result = ncclSuccess;
+  uint64_t time = 0;
+  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  *ptr = nullptr;
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  time = clockNano();
+  if (isFineGrain)
+    CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
+  else
+    CUDACHECKGOTO(hipMalloc(ptr, nelem*sizeof(T)), result, finish);
+  time = clockNano() - time;
+  CUDACHECKGOTO(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
+  int dev;
+  CUDACHECK(hipGetDevice(&dev));
+  if (dev < MAX_ALLOC_TRACK_NGPU) {
+    __atomic_fetch_add(&allocTracker[dev].totalAlloc, 1, __ATOMIC_RELAXED);
+    __atomic_fetch_add(&allocTracker[dev].totalAllocSize, nelem*sizeof(T), __ATOMIC_RELAXED);
+  }
+  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p seconds: hipMalloc=%g", filefunc, line, nelem*sizeof(T), *ptr, double(time)/1.e9);
+finish:
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  return result;
+}
+#define ncclCudaCallocAsync(...) ncclCudaCallocAsyncDebug(__FILE__, __LINE__, __VA_ARGS__)
+
+template <typename T>
+ncclResult_t ncclCudaMemcpy(T* dst, T* src, size_t nelem) {
+  ncclResult_t result = ncclSuccess;
+  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  // Need a side stream so as not to interfere with graph capture.
+  hipStream_t stream;
+  CUDACHECKGOTO(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking), result, finish);
+  NCCLCHECKGOTO(ncclCudaMemcpyAsync(dst, src, nelem, stream), result, finish);
+  CUDACHECKGOTO(hipStreamSynchronize(stream), result, finish);
+  CUDACHECKGOTO(hipStreamDestroy(stream), result, finish);
+finish:
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  return result;
+}
+
+template <typename T>
+ncclResult_t ncclCudaMemcpyAsync(T* dst, T* src, size_t nelem, hipStream_t stream) {
+  ncclResult_t result = ncclSuccess;
+  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECKGOTO(hipMemcpyAsync(dst, src, nelem*sizeof(T), hipMemcpyDefault, stream), result, finish);
+finish:
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  return result;
+}
+
+template <typename T>
+ncclResult_t ncclCudaFree(T* ptr) {
+  ncclResult_t result = ncclSuccess;
+  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECKGOTO(hipFree(ptr), result, finish);
+finish:
+  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  return result;
 }
 
 // Allocate memory to be potentially ibv_reg_mr'd. This needs to be
 // allocated on separate pages as those pages will be marked DONTFORK
 // and if they are shared, that could cause a crash in a child process
-static ncclResult_t ncclIbMallocDebug(void** ptr, size_t size, const char *filefunc, int line) {
+inline ncclResult_t ncclIbMallocDebug(void** ptr, size_t size, const char *filefunc, int line) {
   size_t page_size = sysconf(_SC_PAGESIZE);
   void* p;
   int size_aligned = ROUNDUP(size, page_size);

src/include/bootstrap.h

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/

src/include/channel.h

@@ -31,7 +31,8 @@ static ncclResult_t ncclChannelComputeBase(struct ncclComm* comm, int peer, int
 }
 
 static ncclResult_t ncclChannelComputeFromBase(struct ncclComm* comm, int base, int channelInc, int*channelId) {
-  *channelId = (base+comm->p2pChannels[channelInc]) % comm->p2pnChannels;
+  //*channelId = (base+comm->p2pChannels[channelInc]) % comm->p2pnChannels;
+  *channelId = (comm->p2pChannels[base%comm->p2pnChannels]+channelInc) % comm->p2pnChannels;
   return ncclSuccess;
 }
 

src/include/checks.h

@@ -10,7 +10,7 @@
 
 #include "debug.h"
 
-// Check CUDA calls
+// Check CUDA RT calls
 #define CUDACHECK(cmd) do {                                 \
     hipError_t err = cmd;                                    \
     if( err != hipSuccess ) {                                \
@@ -143,9 +143,9 @@
   if (tmpAbortFlag) NEQCHECKGOTO(*tmpAbortFlag, 0, res, label); \
 } while (!(cond));
 
-#define NCCLCHECKTHREAD(a) do { \
-  if ((args->ret = (a)) != ncclSuccess) { \
-    INFO(NCCL_INIT,"%s:%d -> %d [Async thread]", __FILE__, __LINE__, args->ret); \
+#define NCCLCHECKTHREAD(a, args) do { \
+  if (((args)->ret = (a)) != ncclSuccess) { \
+    INFO(NCCL_INIT,"%s:%d -> %d [Async thread]", __FILE__, __LINE__, (args)->ret); \
     return args; \
   } \
 } while(0)

src/include/coll_net.h

@@ -10,25 +10,26 @@
 #include "nccl.h"
 #include "nccl_net.h"
 
-extern ncclCollNet_t* ncclCollNet;
 typedef char collNetHandle_t[NCCL_NET_HANDLE_MAXSIZE];
 
 // Translation to external API
-static const char* collNetName() { return ncclCollNet->name; }
-static ncclResult_t collNetDevices(int* ndev) { NCCLCHECK(ncclCollNet->devices(ndev)); return ncclSuccess; }
-static ncclResult_t collNetGetProperties(int dev, ncclNetProperties_t* props) { NCCLCHECK(ncclCollNet->getProperties(dev, props)); return ncclSuccess; }
-static ncclResult_t collNetListen(int dev, void* handle, void** listenComm) { NCCLCHECK(ncclCollNet->listen(dev, handle, listenComm)); return ncclSuccess; }
-static ncclResult_t collNetConnect(void* handles[], int nranks, int rank, void* listenComm, void** collComm) { NCCLCHECK(ncclCollNet->connect(handles, nranks, rank, listenComm, collComm)); return ncclSuccess; }
-static ncclResult_t collNetReduceSupport(ncclDataType_t dataType, ncclRedOp_t redOp, int* supported) { NCCLCHECK(ncclCollNet->reduceSupport(dataType, redOp, supported)); return ncclSuccess; }
-static ncclResult_t collNetRegMr(void* comm, void* data, int size, int type, void** mhandle) { NCCLCHECK(ncclCollNet->regMr(comm, data, size, type, mhandle)); return ncclSuccess; }
-static ncclResult_t collNetDeregMr(void* comm, void* mhandle) { NCCLCHECK(ncclCollNet->deregMr(comm, mhandle)); return ncclSuccess; }
-static ncclResult_t collNetIallreduce(void* collComm, void* sendData, void* recvData, int count, ncclDataType_t dataType, ncclRedOp_t redOp, void* sendMhandle, void* recvMhandle,  void** request) {
-  NCCLCHECK(ncclCollNet->iallreduce(collComm, sendData, recvData, count, dataType, redOp, sendMhandle, recvMhandle, request)); return ncclSuccess; }
-static ncclResult_t collNetIflush(void* collComm, void* data, int size, void* mhandle, void** request) { NCCLCHECK(ncclCollNet->iflush(collComm, data, size, mhandle, request)); return ncclSuccess; }
-static ncclResult_t collNetTest(void* request, int* done, int* size) { NCCLCHECK(ncclCollNet->test(request, done, size)); return ncclSuccess; }
-static ncclResult_t collNetCloseColl(void* collComm) { NCCLCHECK(ncclCollNet->closeColl(collComm)); return ncclSuccess; }
-static ncclResult_t collNetCloseListen(void* listenComm) { NCCLCHECK(ncclCollNet->closeListen(listenComm)); return ncclSuccess; }
+static const char* collNetName(struct ncclComm* comm) { return comm->ncclCollNet->name; }
+static ncclResult_t collNetDevices(struct ncclComm* comm, int* ndev) { NCCLCHECK(comm->ncclCollNet->devices(ndev)); return ncclSuccess; }
+static ncclResult_t collNetGetProperties(struct ncclComm* comm, int dev, ncclNetProperties_t* props) { NCCLCHECK(comm->ncclCollNet->getProperties(dev, props)); return ncclSuccess; }
+static ncclResult_t collNetListen(struct ncclComm* comm, int dev, void* handle, void** listenComm) { NCCLCHECK(comm->ncclCollNet->listen(dev, handle, listenComm)); return ncclSuccess; }
+static ncclResult_t collNetConnect(struct ncclComm* comm, void* handles[], int nranks, int rank, void* listenComm, void** collComm) { NCCLCHECK(comm->ncclCollNet->connect(handles, nranks, rank, listenComm, collComm)); return ncclSuccess; }
+static ncclResult_t collNetReduceSupport(struct ncclComm* comm, ncclDataType_t dataType, ncclRedOp_t redOp, int* supported) { NCCLCHECK(comm->ncclCollNet->reduceSupport(dataType, redOp, supported)); return ncclSuccess; }
+static ncclResult_t collNetRegMr(struct ncclComm* comm, void* collComm, void* data, int size, int type, void** mhandle) { NCCLCHECK(comm->ncclCollNet->regMr(collComm, data, size, type, mhandle)); return ncclSuccess; }
+/* DMA-BUF support */
+static ncclResult_t collNetRegMrDmaBuf(struct ncclComm* comm, void* collComm, void* data, int size, int type, uint64_t offset, int fd, void** mhandle) { NCCLCHECK(comm->ncclCollNet->regMrDmaBuf(collComm, data, size, type, offset, fd, mhandle)); return ncclSuccess; }
+static ncclResult_t collNetDeregMr(struct ncclComm* comm, void* collComm, void* mhandle) { NCCLCHECK(comm->ncclCollNet->deregMr(collComm, mhandle)); return ncclSuccess; }
+static ncclResult_t collNetIallreduce(struct ncclComm* comm, void* collComm, void* sendData, void* recvData, int count, ncclDataType_t dataType, ncclRedOp_t redOp, void* sendMhandle, void* recvMhandle,  void** request) {
+  NCCLCHECK(comm->ncclCollNet->iallreduce(collComm, sendData, recvData, count, dataType, redOp, sendMhandle, recvMhandle, request)); return ncclSuccess; }
+static ncclResult_t collNetIflush(struct ncclComm* comm, void* collComm, void* data, int size, void* mhandle, void** request) { NCCLCHECK(comm->ncclCollNet->iflush(collComm, data, size, mhandle, request)); return ncclSuccess; }
+static ncclResult_t collNetTest(struct ncclComm* comm, void* request, int* done, int* size) { NCCLCHECK(comm->ncclCollNet->test(request, done, size)); return ncclSuccess; }
+static ncclResult_t collNetCloseColl(struct ncclComm* comm, void* collComm) { NCCLCHECK(comm->ncclCollNet->closeColl(collComm)); return ncclSuccess; }
+static ncclResult_t collNetCloseListen(struct ncclComm* comm, void* listenComm) { NCCLCHECK(comm->ncclCollNet->closeListen(listenComm)); return ncclSuccess; }
 
-static int collNetSupport() { return ncclCollNet != nullptr ? 1 : 0; }
+static int collNetSupport(struct ncclComm* comm) { return comm->ncclCollNet != nullptr ? 1 : 0; }
 
 #endif

src/include/collectives.h

@@ -47,10 +47,10 @@ struct ncclDevRedOpFull {
 /* Declare all collective operations */
 #define DECL5(func, algo, proto, devredop, type) \
   extern __device__ __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, devredop, type)(); \
-  extern __global__ void NCCL_KERN_NAME(func, algo, proto, devredop, type)(struct ncclDevComm* comm); \
-  extern __global__ void NCCL_KERN_NAME_DEBUG(func, algo, proto, devredop, type)(struct ncclDevComm* comm); \
-  extern __global__ void NCCL_KERN_NAME_LL128(func, algo, proto, devredop, type)(struct ncclDevComm* comm); \
-  extern __global__ void NCCL_KERN_NAME_LL128_DEBUG(func, algo, proto, devredop, type)(struct ncclDevComm* comm);
+  extern __global__ void NCCL_KERN_NAME(func, algo, proto, devredop, type)(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead); \
+  extern __global__ void NCCL_KERN_NAME_DEBUG(func, algo, proto, devredop, type)(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead); \
+  extern __global__ void NCCL_KERN_NAME_LL128(func, algo, proto, devredop, type)(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead); \
+  extern __global__ void NCCL_KERN_NAME_LL128_DEBUG(func, algo, proto, devredop, type)(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead);
 
 #define CONCAT(a,b) a##b
 #define MACRO_IF(cond, t, f) CONCAT(MACRO_IF_, cond)(t, f)

src/include/comm.h

@@ -10,25 +10,13 @@
 
 #include "transport.h"
 #include "p2p.h"
-// [RCCL]
-//#include "clique/CliqueManager.h"
-// [/RCCL]
-
-// Convert volatile access to atomic
-#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-#define LOAD(VAR) __atomic_load_n((VAR), __ATOMIC_SEQ_CST)
-#define STORE(DST, SRC) __atomic_store_n((DST), (SRC), __ATOMIC_SEQ_CST)
-#else
-#define LOAD(VAR) *(VAR)
-#define STORE(DST, SRC) *(DST) = (SRC)
-#endif
-
-
-#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-  #define HIPRT_CB
-#else
 #include "collectives.h"
+#include "proxy.h"
+#include "strongstream.h"
 
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+#define HIPRT_CB
+#else
 #if CUDART_VERSION < 9000
 struct cudaLaunchParams {
   void *func;
@@ -77,8 +65,6 @@ struct ncclRecvMem {
   };
 };
 
-typedef hipError_t(*pfn_cuMemGetAddressRange_t)(void**, size_t*, void*);
-
 enum helperThreadState {ThreadStart, ThreadStop};
 
 #define NCCL_IPC_POOL_SIZE (2*NCCL_MAX_LOCAL_RANKS*NCCL_MAX_OPS)
@@ -104,15 +90,87 @@ struct ncclNodeRanks {
   int* localRankToRank;
 };
 
-struct ncclComm {
-  struct ncclChannel channels[MAXCHANNELS];
+struct ncclDestructor {
+  struct ncclDestructor* next;
+  void* obj;
+  ncclResult_t(*fn)(struct ncclDestructor* me);
+};
 
+struct ncclCommCallback {
+  struct ncclCommCallback* next;
+  ncclResult_t(*fn)(struct ncclComm* comm, struct ncclCommCallback* cb);
+};
+
+struct ncclChannel {
+  struct ncclChannelPeer* peers;
+  struct ncclDevChannelPeer* devPeers;
+  struct ncclRing ring;
+  int* devRingUserRanks;
+  struct ncclTree tree;
+  struct ncclDirect collTree;
+  int id; // index of this channel
+  uint32_t workFifoSent; // last used work index+1
+  uint64_t p2pOpCount;
+};
+
+struct ncclWorkList {
+  struct ncclWorkList* next;
+  struct ncclWork work;
+};
+
+struct ncclPointerList {
+  struct ncclPointerList* next;
+  void *ptr;
+};
+
+struct ncclKernelPlan {
+  // A kernel plan is also a callback that reclaims itself. Hence this must
+  // be the first member.
+  struct ncclCommCallback reclaimer;
+  struct ncclMemoryPool memPool_ncclProxyOp; // memory to return to comm in cleanup
+
+  struct ncclComm* comm;
+  struct ncclKernelPlan* next;
+
+  bool persistent; // aka captured in a graph
+  void *kernelFn;
+  int channelUbound; // only channels c < channelUbound are present
+  int channelCount; // number of channels present
+  uint64_t channelMask; // which channels are present, channelCount == popcount(channelMask)
+  bool hasProxyOps; // does any channel have a non-empty proxyOpQueue
+  int threadPerBlock;
+  // workHeap fields are null until uploadWorkFifo() or preparePersistentKernel()
+  struct ncclWork* workHead;
+
+  int collOpCount; // zero based for this plan
+
+  struct ncclIntruQueue<struct ncclPointerList, &ncclPointerList::next> ipcMemQueue;
+
+  struct Channel {
+    int nWork;
+    union {
+      int nWorkElem; // used for coll and reg coll
+      int p2pTailElem[2]; // used for p2p, indexed by ncclWorkElemP2pType-1
+    };
+    size_t collBytes;
+    struct ncclIntruQueue<struct ncclWorkList, &ncclWorkList::next> workQueue;
+    struct ncclIntruQueue<struct ncclProxyOp, &ncclProxyOp::enqNext> proxyOpQueue;
+  } channels[MAXCHANNELS];
+};
+
+struct ncclComm {
+  struct ncclMemoryStack memPermanent, memScoped;
+  // List of destructors to run when comm is destructed
+  struct ncclDestructor* destructorHead;
+
+  struct ncclChannel channels[MAXCHANNELS];
   struct ncclPeerInfo* peerInfo;
   struct ncclTopoSystem* topo;
 
+  ncclNet_t* ncclNet;
+  ncclCollNet_t* ncclCollNet;
   void* bootstrap;
   // Bitmasks for ncclTransportP2pSetup
-  int connect[NCCL_MAX_CONNS];
   uint32_t* connectSend;
   uint32_t* connectRecv;
 
@@ -135,19 +193,13 @@ struct ncclComm {
   // localRanks and localRanktoRank for all nodes
   struct ncclNodeRanks* nodeRanks;
 
-  enum { GROUP, PARALLEL, GROUP_GRAPH } launchMode;
-  hipStream_t userStream;
-  bool userStreamSet;
-  hipEvent_t doneEvent;
-  hipEvent_t intDoneEvent;
   bool checkPointers;
+  bool dmaBufSupport;
 
   // Counter for tracking CUDA launches (P2P and collectives included)
   uint64_t opCount;
   // Collective operation counter
   uint64_t collOpCount;
-  // P2P operation counter
-  uint64_t p2pOpCount;
 
   // Channels for collectives
   int nChannels;
@@ -165,10 +217,6 @@ struct ncclComm {
   float bandwidths[NCCL_NUM_FUNCTIONS][NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
   int maxThreads[NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
 
-  // An internal CUDA stream for NCCL kernel CGMD launches
-  int groupCudaStream;
-  hipStream_t groupStream;
-
   // Whether there has been a fatal error in this communicator.
   ncclResult_t fatalError;
 
@@ -178,26 +226,33 @@ struct ncclComm {
   // Flags for enable P2P NET
   uint32_t p2pNet;
   uint32_t useIntraNet;
+  bool hasFineGrain;
 
-  // Device side of the communicator
-  struct ncclDevComm *devComm;
-  // Host copy of the devComm (to free CUDA allocs)
-  struct ncclDevComm hostDevComm;
+  // Device side of the communicator (for cudaFree's)
+  struct ncclDevComm* devComm; // actually = &ncclDevCommAndChannels::comm
+
+  // Operation pool.
+  int workFifoDepth; // size of workFifoHeap[], power of 2
+  struct ncclWork* workFifoHeap;
+  struct ncclWork* devWorkFifoHeap;
+  void* workFifoHeapGdrHandle;
+
+  // Work completion notificaion
+  uint32_t* workFifoDone/*[MAXCHANNELS]*/; // in cudaHost memory
+  uint32_t workFifoSent; // Monotonic (mod 1<<32) index of next unused fifo slot.
+  uint32_t workFifoAckdMin; // Monotonic index of least unprocessed fifo slot over all channels.
 
   // Intra-process sync
+  struct ncclComm* intraComm0; // leader of intra-process comms (self possible)
+  struct ncclComm* intraNext; // next of intra-process comms, intraComm0 is head
+  int intraRefs; // reference count from intra-process comms (zero if not leader else intraRanks)
   int intraRank;
   int intraRanks;
-  int* intraBarrier;
-  int intraPhase;
-
-  // Storage for deferred intra-process launch
-  hipLaunchParams * intraParams;
-  hipLaunchParams *myParams;
-  pthread_t* intraThreads;
-  int* intraCudaDevs;
-  int* intraCGMode; // Whether we can use CUDA9 CGMD or not
-  int* intraCC; // Only to check all have the same ComputeCap and disable CGMode if not
-  void* argsptrs[1];
+  uint32_t intraBarrierPhase;
+  char intraPad1[64 - sizeof(uint64_t)];
+  uint64_t intraBarrierCounter; // only used if this is intraComm0
+  char intraPad2[64 - sizeof(uint64_t)];
+  uint64_t intraBarrierGate; // only used if this is intraComm0
 
   struct ncclProxyState proxyState;
 
@@ -205,44 +260,108 @@ struct ncclComm {
   int collNetSupport;
   int intraHighestTransportType;
 
-  // Store info of async operations
-  struct ncclInfo* asyncOps;
-  int asyncOpCount;
-  size_t asyncTotalSize;
-  ssize_t channelSize;
-  int lastChannel;
-  enum { ROUND_ROBIN, SHORTEST_QUEUE } asyncAllocMode;
+  size_t channelSize; // User requested work size (bytes) for channel partitions
 
-  //list of async p2p operation queued in a group semantics
-  ncclP2Plist** p2pSends;
-  ncclP2Plist** p2pRecvs;
-  int p2pSendCount;
-  int p2pRecvCount;
+  // Internal streams
+  struct ncclStrongStream deviceStream, hostStream;
 
-  // [RCCL]
-  //CliqueManager* cliqueManager;    // CliqueManager handles pointer collection / distribution for clique-based kernels
-  //int rootPid;                     // Process ID of root
-  // [/RCCL]
-
-  // Store info for cudaGraph
-  int usingCudaGraph; // Only use it during capture time, not launch time
-  struct ncclQueueInfo* enqueueInfo;
-  int nQueueInfoCreated;
-  int nQueueInfoDestroyed;
-  hipGraphNode_t lastSetupNode;
-  unsigned long long lastCudaGraphId;
-  int driverVersion;
-  pfn_cuMemGetAddressRange_t pfnCuMemGetAddressRange;
-  pthread_t graphHelperThread;
-  struct ncclGraphHelperResources* graphHelperResources;
-  int disableGraphHelper;
-  int graphRegister;
+  // pools backed by comm->memPermanent
+  struct ncclMemoryPool memPool_ncclProxyOp;
+  struct ncclMemoryPool memPool_ncclKernelPlan;
+  struct ncclMemoryPool memPool_ncclPointerList;
+  // Next comm in this thread's active ncclGroup[Start|End](). Holds "0x1" when
+  // this comm is not yet in a group.
+  struct ncclComm* groupNext;
+  // Subset of those in groupNext list. Holds 0x1 if not needing preconnect.
+  struct ncclComm* preconnectNext;
+  int persistentRefs; // number of persistent plan-lists capturing this comm
+  struct ncclTasks tasks;
 
   // user-created reduction ops
   int userRedOpCapacity, userRedOpFreeHead;
   ncclUserRedOp *userRedOps;
+
+  // Queue of things for the main thread to do
+  struct ncclIntruQueueMpsc<struct ncclCommCallback, &ncclCommCallback::next> callbackQueue;
+
+  // List of kernel plans built form tasks.
+  struct ncclIntruQueue<struct ncclKernelPlan, &ncclKernelPlan::next> planQueue;
+  // First of the unlaunched kernels in `planQueue`
+  struct ncclKernelPlan* unlaunchedPlansHead;
+
+  hipEvent_t doneEvent;
+  hipStream_t lastStream;
+
+#ifdef ENABLE_COLLTRACE
+  struct ncclCollTrace* collTrace;
+  volatile uint32_t *collTraceTail;
+  pthread_t collTraceThread;
+  volatile bool collTraceExit;
+#endif
 };
 
+// Set to true during an `atexit()` handler. We use this to intentionally leak
+// unfreed CUDA resources when cleaning up after return of `main()` to avoid
+// CUDA calls after CUDA runtime teardown.
+extern bool ncclMainExited;
+
+enum ncclLaunchMode {
+  ncclLaunchModeInvalid=0,
+  ncclLaunchModeParallel,
+  ncclLaunchModeGroup
+};
+extern enum ncclLaunchMode ncclParamLaunchMode;
+
+void ncclCommPushFree(struct ncclComm* comm, void* buf);
+void ncclCommPushCudaFree(struct ncclComm* comm, void* buf);
+void ncclCommPushCudaHostFree(struct ncclComm* comm, void* buf);
+void ncclCommPushCudaGdrFree(struct ncclComm* comm, void* handle);
+
+inline ncclResult_t ncclCommPollCallbacks(struct ncclComm* comm) {
+  struct ncclCommCallback* cb = ncclIntruQueueMpscDequeueAll(&comm->callbackQueue, /*waitSome=*/false);
+  while (cb != nullptr) {
+    struct ncclCommCallback* next = cb->next;
+    NCCLCHECK(cb->fn(comm, cb)); // may reclaim memory of cb
+    cb = next;
+  }
+  return ncclSuccess;
+}
+
+inline void ncclCommIntraBarrierIn(struct ncclComm* comm, uint32_t x) {
+  int phase = comm->intraBarrierPhase;
+  if (comm->intraRanks == 1) {
+    // Release everyone (just me).
+    comm->intraBarrierGate = (uint64_t(x)<<32) | (phase^1);
+  } else {
+    struct ncclComm* comm0 = comm->intraComm0;
+    uint64_t count = __atomic_add_fetch(&comm0->intraBarrierCounter, (uint64_t(x)<<32) + 1, __ATOMIC_RELEASE);
+    if (uint32_t(count) == uint32_t(comm->intraRanks)) {
+      // Reset.
+      __atomic_store_n(&comm0->intraBarrierCounter, 0, __ATOMIC_RELAXED);
+      // Release everyone.
+      __atomic_store_n(&comm0->intraBarrierGate, (count>>32<<32) | (phase^1), __ATOMIC_RELEASE);
+    }
+  }
+}
+
+// returns sum of x values contributed to ncclCommIntraBarrierIn(comm, x)
+inline uint32_t ncclCommIntraBarrierOut(struct ncclComm* comm) {
+  struct ncclComm* comm0 = comm->intraComm0;
+  comm->intraBarrierPhase ^= 1;
+  uint32_t phase = comm->intraBarrierPhase;
+  uint64_t gate = __atomic_load_n(&comm0->intraBarrierGate, __ATOMIC_RELAXED);
+  if ((gate & 1) != phase) {
+    uint64_t t0 = clockNano();
+    do {
+      // Spin vigorously for first 5us.
+      if (clockNano()-t0 >= 5*1000) sched_yield();
+      gate = __atomic_load_n(&comm0->intraBarrierGate, __ATOMIC_RELAXED);
+    } while ((gate & 1) != phase);
+  }
+  if (comm->intraRanks != 1) __atomic_thread_fence(__ATOMIC_ACQUIRE);
+  return gate>>32;
+}
+
 // Scrambles the bits of non-builtin values of ncclRedOp_t according to the
 // communicator memory address. Used to catch bugs so that integer handles
 // associated with this communicator won't collide with handles of other

src/include/core.h

@@ -37,7 +37,9 @@ static __inline__ int ncclTypeSize(ncclDataType_t type) {
     case ncclUint8:
       return 1;
     case ncclFloat16:
+#if defined(RCCL_BFLOAT16)
     case ncclBfloat16:
+#endif
       return 2;
     case ncclInt32:
     case ncclUint32:
@@ -54,6 +56,7 @@ static __inline__ int ncclTypeSize(ncclDataType_t type) {
 
 #include "debug.h"
 #include "checks.h"
+#include "rocmwrap.h"
 #include "alloc.h"
 #include "utils.h"
 #include "param.h"

src/include/cudawrap.h

@@ -0,0 +1,88 @@
+/*************************************************************************
+ * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef NCCL_CUDAWRAP_H_
+#define NCCL_CUDAWRAP_H_
+
+#include <cuda.h>
+
+#if CUDART_VERSION >= 11030
+#include <cudaTypedefs.h>
+#else
+typedef CUresult (CUDAAPI *PFN_cuInit)(unsigned int Flags);
+typedef CUresult (CUDAAPI *PFN_cuDriverGetVersion)(int *driverVersion);
+typedef CUresult (CUDAAPI *PFN_cuGetProcAddress)(const char *symbol, void **pfn, int driverVersion, cuuint64_t flags);
+#endif
+
+#define CUPFN(symbol) pfn_##symbol
+
+// Check CUDA PFN driver calls
+#define CUCHECK(cmd) do {				      \
+    CUresult err = pfn_##cmd;				      \
+    if( err != CUDA_SUCCESS ) {				      \
+      const char *errStr;				      \
+      (void) pfn_cuGetErrorString(err, &errStr);	      \
+      WARN("Cuda failure '%s'", errStr);		      \
+      return ncclUnhandledCudaError;			      \
+    }							      \
+} while(false)
+
+#define CUCHECKGOTO(cmd, res, label) do {		      \
+    CUresult err = pfn_##cmd;				      \
+    if( err != CUDA_SUCCESS ) {				      \
+      const char *errStr;				      \
+      (void) pfn_cuGetErrorString(err, &errStr);	      \
+      WARN("Cuda failure '%s'", errStr);		      \
+      res = ncclUnhandledCudaError;			      \
+      goto label;					      \
+    }							      \
+} while(false)
+
+// Report failure but clear error and continue
+#define CUCHECKIGNORE(cmd) do {						\
+    CUresult err = pfn_##cmd;						\
+    if( err != CUDA_SUCCESS ) {						\
+      const char *errStr;						\
+      (void) pfn_cuGetErrorString(err, &errStr);			\
+      INFO(NCCL_ALL,"%s:%d Cuda failure '%s'", __FILE__, __LINE__, errStr);	\
+    }									\
+} while(false)
+
+#define CUCHECKTHREAD(cmd, args) do {					\
+    CUresult err = pfn_##cmd;						\
+    if (err != CUDA_SUCCESS) {						\
+      INFO(NCCL_INIT,"%s:%d -> %d [Async thread]", __FILE__, __LINE__, err); \
+      args->ret = ncclUnhandledCudaError;				\
+      return args;							\
+    }									\
+} while(0)
+
+#define DECLARE_CUDA_PFN_EXTERN(symbol) extern PFN_##symbol pfn_##symbol
+
+#if CUDART_VERSION >= 11030
+/* CUDA Driver functions loaded with cuGetProcAddress for versioning */
+DECLARE_CUDA_PFN_EXTERN(cuDeviceGet);
+DECLARE_CUDA_PFN_EXTERN(cuDeviceGetAttribute);
+DECLARE_CUDA_PFN_EXTERN(cuGetErrorString);
+DECLARE_CUDA_PFN_EXTERN(cuGetErrorName);
+DECLARE_CUDA_PFN_EXTERN(cuMemGetAddressRange);
+DECLARE_CUDA_PFN_EXTERN(cuCtxCreate_v3020);
+DECLARE_CUDA_PFN_EXTERN(cuCtxDestroy);
+DECLARE_CUDA_PFN_EXTERN(cuCtxSetCurrent);
+#if CUDA_VERSION >= 11070
+DECLARE_CUDA_PFN_EXTERN(cuMemGetHandleForAddressRange); // DMA-BUF support
+#endif
+#endif
+
+/* CUDA Driver functions loaded with dlsym() */
+DECLARE_CUDA_PFN_EXTERN(cuInit);
+DECLARE_CUDA_PFN_EXTERN(cuDriverGetVersion);
+DECLARE_CUDA_PFN_EXTERN(cuGetProcAddress);
+
+
+ncclResult_t cudaLibraryInit(void);
+
+#endif

src/include/debug.h

@@ -10,8 +10,8 @@
 #include "nccl_net.h"
 #include <stdio.h>
 #include <chrono>
+#include <type_traits>
 
-#include <sys/syscall.h>
 #include <limits.h>
 #include <string.h>
 #include <pthread.h>
@@ -21,7 +21,7 @@
 
 extern int ncclDebugLevel;
 extern uint64_t ncclDebugMask;
-extern pthread_mutex_t ncclDebugOutputLock;
+extern pthread_mutex_t ncclDebugLock;
 extern FILE *ncclDebugFile;
 extern ncclResult_t getHostName(char* hostname, int maxlen, const char delim);
 
@@ -29,13 +29,15 @@ void ncclDebugLog(ncclDebugLogLevel level, unsigned long flags, const char *file
 
 // Let code temporarily downgrade WARN into INFO
 extern thread_local int ncclDebugNoWarn;
+extern char ncclLastError[];
 
 #define WARN(...) ncclDebugLog(NCCL_LOG_WARN, NCCL_ALL, __FILE__, __LINE__, __VA_ARGS__)
 #define INFO(FLAGS, ...) ncclDebugLog(NCCL_LOG_INFO, (FLAGS), __func__, __LINE__, __VA_ARGS__)
+#define TRACE_CALL(...) ncclDebugLog(NCCL_LOG_TRACE, NCCL_CALL, __func__, __LINE__, __VA_ARGS__)
 
 #ifdef ENABLE_TRACE
 #define TRACE(FLAGS, ...) ncclDebugLog(NCCL_LOG_TRACE, (FLAGS), __func__, __LINE__, __VA_ARGS__)
-extern std::chrono::high_resolution_clock::time_point ncclEpoch;
+extern std::chrono::steady_clock::time_point ncclEpoch;
 #else
 #define TRACE(...)
 #endif

src/include/devcomm.h

@@ -15,9 +15,6 @@
 #include "npkit/npkit_struct.h"
 #endif
 #include <stdint.h>
-// [RCCL] Support for clique-based kernels
-//#include "clique/CliqueCommon.h"
-// [/RCCL]
 
 
 #define NCCL_NUM_FUNCTIONS 5 // SendRecv and AllToAllPivot not included for now
@@ -33,7 +30,6 @@ extern const char* ncclAlgoStr[NCCL_NUM_ALGORITHMS];
 #define NCCL_NUM_PROTOCOLS 3 // Simple/LL/LL128
 #define NCCL_PROTO_LL 0
 #define NCCL_PROTO_LL128 1
-#define NCCL_PROTO_CLIQUE 1  // [RCCL] Clique takes up same protocol as unused LL128
 #define NCCL_PROTO_SIMPLE 2
 extern const char* ncclProtoStr[NCCL_NUM_PROTOCOLS];
 
@@ -83,10 +79,6 @@ static_assert(NCCL_LL_CLEAN_MASK % NCCL_STEPS == 0, "Invalid NCCL_LL_CLEAN_MASK
 #define NCCL_LL128_MAX_NTHREADS 256
 #define NCCL_LL128_ELEMS_PER_THREAD 28
 
-// Receiving from up to 3 sources is more compute intensive than sending
-// to 3 dests. Use 70% for reduce and 30% for bcast.
-#define NCCL_LL128_SPLIT(nt) ((nt*7/(10*32))*32)
-
 #define NCCL_LL128_SHMEM_ELEMS_PER_THREAD 4
 #define NCCL_LL128_SHMEM_SIZE (NCCL_LL128_SHMEM_ELEMS_PER_THREAD*NCCL_LL128_MAX_NTHREADS)
 
@@ -145,7 +137,6 @@ struct ncclRing {
   // since we need to know how the user expects data to be ordered across
   // devices. Ordered from current device.
   int* userRanks;
-  int* devUserRanks;
 
   int index; // This rank's index in the ring
 };
@@ -171,7 +162,7 @@ struct ncclDirect {
 
 #define NCCL_CONN_IDX_P2P_NET 2
 #define NCCL_MAX_CONNS 3
-struct ncclPeer {
+struct ncclChannelPeer {
   struct ncclConnector send[NCCL_MAX_CONNS];
   struct ncclConnector recv[NCCL_MAX_CONNS];
 };
@@ -185,31 +176,43 @@ struct ncclDevComm;
 /* Make sure to adjust padding at the end of ncclWorkElem. */
 #define NCCL_WORK_SIZE 256
 
-enum ncclWorkElemType : uint8_t {
+enum ncclWorkType : uint8_t {
    ncclWorkTypeUnused=0,
    ncclWorkTypeColl=1,
    ncclWorkTypeP2p=2,
    ncclWorkTypeRegColl=3
 };
-enum ncclWorkElemSubType : uint8_t {
-  ncclWorkSubTypeUnused =0,
-  ncclWorkSubTypeSend,
-  ncclWorkSubTypeRecv
+enum ncclWorkP2PType : uint8_t {
+  ncclWorkP2pTypeUnused=0,
+  ncclWorkP2pTypeSend,
+  ncclWorkP2pTypeRecv
 };
 
-struct ncclWorkElemHeader {
+struct ncclWorkHeader {
+  union {
+    int32_t workNext;  // when isLast=0: Offset from kernel argument workHead
+    uint32_t doneAcks; // when isLast=1: Monotonic (mod 1<<32) ack value to send back.
+  };
   uint16_t funcIndex;
-  enum ncclWorkElemType type;
-  uint8_t nWarps:5;
-  uint8_t isLast:1;
+  uint8_t isLast:1; // last work for this kernel
+  uint8_t inFifo:1; // is this work in the fifo
+  enum ncclWorkType type;
 };
 
 struct ncclWorkElem {
-  struct ncclWorkElemHeader header;
-  uint8_t regUsed;
+  union {
+    uint8_t flagBits;
+    struct {
+      uint8_t isUsed:1, redOpArgIsPtr:1, regUsed:1, pad_0:1, nWarps:4;
+    };
+  };
   uint8_t direct;
-  uint8_t redOpArgIsPtr;
-  uint8_t pad_0;
+  uint8_t bid;
+  uint8_t nChannels;
+  struct {
+    uint32_t root:30;
+    uint32_t connIndex:2;
+  };
 
   const void * sendbuff;
   void * recvbuff;
@@ -221,29 +224,40 @@ struct ncclWorkElem {
     // Instead, it needs the number of bidirectional rings.
     size_t pivotA2ANumBiRings;
   };
-  uint32_t root;
-  uint8_t bid;
-  uint8_t nChannels;
-  uint16_t connIndex;
   uint64_t redOpArg;
   uint64_t opCount;
 };
-static_assert(NCCL_WORK_SIZE % sizeof(struct ncclWorkElem) == 0, "ncclWorkElem size must be a multiple of ncclWork size");
+
+static_assert((NCCL_WORK_SIZE - alignUp(sizeof(ncclWorkHeader), alignof(ncclWorkElem)))/sizeof(ncclWorkElem) == 4, "Sanity check: NCCL_MAX_WORK_ELEMENTS == 4");
+#define NCCL_MAX_WORK_ELEMENTS 1
 
 struct ncclWorkElemP2p {
-  struct ncclWorkElemHeader header;
-  int32_t peer;
-  void* buff;
-  size_t count;
+  struct {
+    int32_t peer:30;
+    uint32_t connIndex:2;
+  };
+  union {
+    uint16_t flagBits;
+    struct {
+      enum ncclWorkP2PType p2pType:4;
+      uint16_t nWarps:4;
+      uint16_t warpStart:4;
+      uint16_t ngroups:4;
+    };
+  };
+  uint16_t opCount;
+  // Important not to use any fields with greater than 4-byte alignment since
+  // we need sizeof(ncclWorkElemP2p)==28, but that would be padded up to 32 if
+  // there were 8-byte fields.
+  //void* buff;
+  uint32_t buffHi32, buffLo32; // buff = buffHi32<<32 | buffLo32;
+  //size_t count;
+  uint32_t countHi32, countLo32; // count = countHi32<<32 | countLo32;
   int chunkSize;
-  uint8_t ngroups:4;
-  uint8_t warpStart:4;
-  uint8_t nWarps:4;
-  enum ncclWorkElemSubType subType:4;
-  uint16_t opCount:12;
-  uint16_t connIndex:4;
 };
-static_assert(NCCL_WORK_SIZE % sizeof(struct ncclWorkElemP2p) == 0, "ncclWorkElemP2p size must be a multiple of ncclWork size");
+
+static_assert(((NCCL_WORK_SIZE - alignUp(sizeof(ncclWorkHeader), alignof(ncclWorkElemP2p)))/sizeof(ncclWorkElemP2p)) == 8, "Sanity check: NCCL_MAX_WORK_ELEMENTS_P2P == 8");
+#define NCCL_MAX_WORK_ELEMENTS_P2P 2
 
 struct ncclWorkElemReg {
   struct ncclWorkElem elem;
@@ -251,56 +265,31 @@ struct ncclWorkElemReg {
   void* dnOutputs[NCCL_MAX_DIRECT_ARITY+1];
   void* upOutputs[NCCL_MAX_DIRECT_ARITY+1];
 };
-static_assert(NCCL_WORK_SIZE % sizeof(struct ncclWorkElemReg) == 0, "ncclWork size must be a multiple of ncclWorkElemReg size");
-static_assert(sizeof(struct ncclWorkElemReg) % sizeof(struct ncclWorkElem) == 0, "ncclWorkElemReg size must be a multiple of ncclWorkElem size");
 
-#define NCCL_MAX_WORK_ELEMENTS 1
-#define NCCL_MAX_WORK_ELEMENTS_P2P 2
-#define NCCL_MAX_WORK_ELEMENTS_REG (NCCL_WORK_SIZE/sizeof(struct ncclWorkElemReg))
+#define NCCL_MAX_WORK_ELEMENTS_REG ((NCCL_WORK_SIZE - alignUp(sizeof(ncclWorkHeader), alignof(ncclWorkElemReg)))/sizeof(ncclWorkElemReg))
+static_assert(NCCL_MAX_WORK_ELEMENTS_REG == 1, "Sanity check: NCCL_MAX_WORK_ELEMENTS_REG == 1");
+
 // Number of named barriers supported by CUDA
 #define NCCL_MAX_GROUPS (NCCL_MAX_NTHREADS/WARP_SIZE)
 
 struct ncclWork {
+  struct ncclWorkHeader header;
   union {
-    char pad[NCCL_WORK_SIZE];
-    struct ncclWorkElemHeader header;
+    char pad[NCCL_WORK_SIZE - sizeof(struct ncclWorkHeader)];
     struct ncclWorkElem elems[NCCL_MAX_WORK_ELEMENTS];
     struct ncclWorkElemP2p p2pElems[NCCL_MAX_WORK_ELEMENTS_P2P];
     struct ncclWorkElemReg regElems[NCCL_MAX_WORK_ELEMENTS_REG];
   };
 };
+static_assert(sizeof(struct ncclWork) == NCCL_WORK_SIZE, "Sanity check: sizeof(struct ncclWork) == NCCL_WORK_SIZE");
+static_assert(sizeof(struct ncclWork)%16 == 0, "Sanity check: sizeof(struct ncclWork)%16 == 0");
 
-static_assert(sizeof(struct ncclWork) == NCCL_WORK_SIZE, "ncclWork size needs to be well aligned");
-
-struct ncclChannel {
-  union {
-    struct {
-      struct ncclRing ring;
-      struct ncclTree tree;
-      struct ncclDirect collTree;
-
-      int id;
-
-      // Communication structures
-      struct ncclPeer* peers;
-      struct ncclPeer* devPeers;
-
-      // Operation list for aggregation
-      struct ncclWork* workFifo;
-      int workCount;
-      size_t totalSize;
-      uint64_t workFifoTail; // Only used by CPU
-      uint16_t index;        // Only used by GPU
-
-      // GDRCOPY support
-      struct ncclWork* workFifoGdr;
-      struct ncclWork* workFifoDev;
-      void* gdrMemDesc;
-    };
-    int data[0x80];
-  };
+struct ncclDevChannelPeer {
+  // Stripped version of ncclChannelPeer where we only keep the ncclConnInfo
+  // instead of the full ncclConnector.
+  struct ncclConnInfo send[NCCL_MAX_CONNS];
+  struct ncclConnInfo recv[NCCL_MAX_CONNS];
 };
-static_assert(sizeof(struct ncclChannel) == 0x80*sizeof(int), "ncclChannel must have a pow2 size");
 #pragma pack(pop)   /* restore original alignment from stack */
 
 #ifdef ENABLE_PROFILING
@@ -361,38 +350,48 @@ static_assert(sizeof(struct ncclCollTrace) == 8*sizeof(int), "ncclCollTrace must
 #define COLLTRACE_NUM_ITEMS 8192
 #endif
 
+struct alignas(16) ncclDevChannel {
+  struct ncclDevChannelPeer *peers;
+  struct ncclRing ring;
+  struct ncclTree tree;
+  struct ncclDirect collTree;
+  uint32_t* workFifoDone; // Location of done counter, device writes index+1 of last work processed
+};
+
 struct ncclDevComm {
   int rank;
   int nRanks;
   int buffSizes[NCCL_NUM_PROTOCOLS];
 
+  // Operation list for aggregation
+  int workFifoDepth;
+  struct ncclWork* workFifoHeap; // may be cudaHost or GDR memory
+
   // Flag to ask NCCL kernels to abort
-  volatile uint32_t *abortFlag;
+  volatile uint32_t* abortFlag;
 
   // Channels, device side
-  struct ncclChannel* channels;
+  struct ncclDevChannel* channels/*[MAXCHANNELS]*/;
 
 #if defined(ENABLE_NPKIT)
   NpKitEventCollectContext* npKitEventCollectContexts;
   uint64_t* cpuTimestamp;
 #endif
 
-#ifdef ENABLE_PROFILING
-  // Profiling counters
-  struct ncclProf* devProf;
-#endif
-
 #ifdef ENABLE_COLLTRACE
   struct ncclCollTrace* collTrace;
-  uint32_t collTraceHead, *collTraceTail;
+  volatile uint32_t *collTraceTail;
   pthread_t collTraceThread;
-  bool collTraceExit;
+#endif
+
+#ifdef ENABLE_PROFILING
+  struct ncclProf* devProf;
 #endif
 };
 
-struct ncclDevCommAndChannels {
-  ncclDevComm comm;
-  ncclChannel channels[MAXCHANNELS];
+struct alignas(16) ncclDevCommAndChannels {
+  struct ncclDevComm comm;
+  struct ncclDevChannel channels[MAXCHANNELS];
 };
 
 #endif

src/include/enqueue.h

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -11,6 +10,7 @@
 #include "comm.h"
 #include "group.h"
 #include "collectives.h"
+#include "utils.h"
 
 #define NCCL_MIN_CHANNEL_SIZE (NCCL_LL_THREAD_THRESHOLD*64)
 #define NCCL_AGG_CHANNEL_SIZE (1LL << 21) /* 2 MiB, ideal per-channel size to fully utilize bandwidth */
@@ -19,117 +19,10 @@ size_t ncclKernMaxLocalSize();
 size_t ncclKernLocalSize(int i);
 ncclResult_t ncclKernSetSharedMemoryCarveout(int carveOut);
 ncclResult_t ncclEnqueueCheck(struct ncclInfo* info);
-ncclResult_t ncclCpuBarrierIn(struct ncclComm* comm, int* isLast);
-ncclResult_t ncclCpuBarrierLast(struct ncclComm* comm);
-ncclResult_t ncclCpuBarrierOut(struct ncclComm* comm);
-ncclResult_t ncclLaunchBarrier(struct ncclComm* comm);
-ncclResult_t ncclLaunchKernel(ncclComm_t comm);
-ncclResult_t ncclRecordEvents(struct ncclComm* comm);
-ncclResult_t ncclLaunchReset(ncclComm_t comm);
-ncclResult_t ncclSetupP2pKernel(struct ncclInfo* info);
-ncclResult_t ncclSetupAsyncKernels(struct ncclComm* comm);
-template<int USING_CUDA_GRAPH>
-void HIPRT_CB ncclEnqueueHostSetup(void* arg);
-ncclResult_t ncclGetCudaGraph(ncclComm_t comm, hipGraph_t* graph);
-ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, hipGraph_t graph);
+ncclResult_t ncclLaunchPrepare(struct ncclComm* comm);
+ncclResult_t ncclLaunchKernelBefore_NoUncapturedCuda(struct ncclComm* comm, struct ncclKernelPlan* plan);
+ncclResult_t ncclLaunchKernel(struct ncclComm* comm, struct ncclKernelPlan* plan);
+ncclResult_t ncclLaunchKernelAfter_NoCuda(struct ncclComm* comm, struct ncclKernelPlan* plan);
+ncclResult_t ncclLaunchFinish(struct ncclComm* comm);
 
-struct ncclBuffRegInfo {
-  void* sendbuffsBase[NCCL_MAX_LOCAL_RANKS];
-  void* recvbuffsBase[NCCL_MAX_LOCAL_RANKS];
-  void* sendbuffs[NCCL_MAX_LOCAL_RANKS];
-  void* recvbuffs[NCCL_MAX_LOCAL_RANKS];
-  int nBuffs;
-};
-
-// Enqueue information (for kernel and proxy) for each operation
-struct ncclQueueElem {
-  struct ncclWork work;
-  struct ncclProxyOp proxyOp;
-  struct ncclBuffRegInfo buffRegInfo;
-};
-
-typedef ncclRecyclableList<struct ncclQueueElem> ncclQueueElemList;
-
-// Structure passed to CUDA graph
-struct ncclQueueInfo {
-  ncclComm_t comm;
-  int maxChannels;    // Dynamic version of gridDim
-  ncclResult_t ret;   // Return value of host setup call
-  int nRegBuffs;
-  ncclQueueElemList* elemList;
-};
-
-static ncclResult_t ncclCreateQueueInfo(struct ncclQueueInfo** eqInfo, ncclComm_t comm) {
-  NCCLCHECK(ncclCalloc(eqInfo, 1));
-  (*eqInfo)->comm = comm;
-  (*eqInfo)->elemList = new ncclQueueElemList();
-  (*eqInfo)->comm->nQueueInfoCreated++;
-  return ncclSuccess;
-}
-
-// Reset element queue
-static ncclResult_t ncclResetQueueInfo(struct ncclQueueInfo* eqInfo) {
-  if (eqInfo == NULL) return ncclInternalError;
-  eqInfo->maxChannels = 0;
-  eqInfo->ret = ncclSuccess;
-  eqInfo->nRegBuffs = 0;
-  eqInfo->elemList->recycle();
-  return ncclSuccess;
-}
-
-// Destroy enqueue info space
-// used by both CUDA graph and non CUDA graph
-static void ncclDestroyQueueInfo(void* ptr) {
-  if (ptr == NULL) return;
-  struct ncclQueueInfo* eqInfo = (struct ncclQueueInfo*)ptr;
-  struct ncclComm* comm = eqInfo->comm;
-  // Close IPC mem handles for registered buffers
-  struct ncclQueueElem* eqElem = eqInfo->elemList->begin();
-#if 0
-  // Ideally, the deregistration should happen here
-  // but currently the destroy function of CUDA objects does not allow CUDA API calls
-  while (eqElem != NULL) {
-    for (int i=0; i<eqElem->buffRegInfo.nBuffs; i++) {
-      if (i == eqInfo->comm->localRank) continue;
-      CUDACHECKIGNORE(cudaIpcCloseMemHandle(eqElem->buffRegInfo.sendbuffsBase[i]));
-      CUDACHECKIGNORE(cudaIpcCloseMemHandle(eqElem->buffRegInfo.recvbuffsBase[i]));
-    }
-    eqElem = eqInfo->elemList->getNext();
-  }
-#else
-  // Instead, we push these pointers to a pool owned by ncclComm
-  // and asks a helper thread to close mem handles
-  struct ncclGraphHelperResources* res = comm->graphHelperResources;
-  int ipcTailOld = 0;
-  if (res == NULL || (!comm->graphHelperThread) || eqInfo->nRegBuffs == 0) goto skip;
-
-  pthread_mutex_lock(&res->threadLock);
-  ipcTailOld = res->ipcTail;
-  while (eqElem != NULL) {
-    for (int i=0; i<eqElem->buffRegInfo.nBuffs; i++) {
-      if (eqElem->buffRegInfo.sendbuffsBase[i] != NULL) {
-        res->ipcBases[res->ipcTail] = eqElem->buffRegInfo.sendbuffsBase[i];
-        res->ipcTail = (res->ipcTail+1)%NCCL_IPC_POOL_SIZE;
-      }
-      if (eqElem->buffRegInfo.recvbuffsBase[i] != NULL) {
-        res->ipcBases[res->ipcTail] = eqElem->buffRegInfo.recvbuffsBase[i];
-        res->ipcTail = (res->ipcTail+1)%NCCL_IPC_POOL_SIZE;
-      }
-    }
-    eqElem = eqInfo->elemList->getNext();
-  }
-  if (res->ipcTail != ipcTailOld) {
-    res->threadState = ThreadStart;
-    TRACE(NCCL_COLL, "CUDA Graph destroy function signaling helper thread with %d IPC handles", res->ipcTail-ipcTailOld);
-    pthread_cond_signal(&res->threadCond);
-  }
-  pthread_mutex_unlock(&res->threadLock);
-#endif
-
-skip:
-  delete eqInfo->elemList;
-  free(eqInfo);
-  comm->nQueueInfoDestroyed++;
-  return;
-}
 #endif // End include guard

src/include/gdrwrap.h

@@ -263,7 +263,7 @@ static ncclResult_t ncclGdrCudaFree(void* gdrHandle) {
   gdr_mem_desc_t *md = (gdr_mem_desc_t*)gdrHandle;
   NCCLCHECK(wrap_gdr_unmap(ncclGdrCopy, md->gdrMh, md->gdrMap, md->gdrMapSize));
   NCCLCHECK(wrap_gdr_unpin_buffer(ncclGdrCopy, md->gdrMh));
-  CUDACHECK(hipFree(md->gdrDevMem));
+  CUDACHECK(cudaFree(md->gdrDevMem));
   free(md);
 
   return ncclSuccess;

src/include/graph.h

@@ -24,7 +24,7 @@ ncclResult_t ncclTopoGetSystem(struct ncclComm* comm, struct ncclTopoSystem** sy
 ncclResult_t ncclTopoSortSystem(struct ncclTopoSystem* system);
 ncclResult_t ncclTopoPrint(struct ncclTopoSystem* system);
 
-ncclResult_t ncclTopoComputePaths(struct ncclTopoSystem* system, struct ncclPeerInfo* info);
+ncclResult_t ncclTopoComputePaths(struct ncclTopoSystem* system, struct ncclComm* comm);
 void ncclTopoFree(struct ncclTopoSystem* system);
 ncclResult_t ncclTopoTrimSystem(struct ncclTopoSystem* system, struct ncclComm* comm);
 ncclResult_t ncclTopoComputeP2pChannels(struct ncclComm* comm);
@@ -37,7 +37,7 @@ ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* topo, int64_t busId, int ne
 #define MAX_XGMI_INTER_GPUS 4
 ncclResult_t ncclTopoGetIntraNetDev(struct ncclTopoSystem* system, int rank, struct ncclTopoGraph* graph, int channelId, int type, int* dev);
 ncclResult_t ncclTopoGetLinkType(struct ncclTopoSystem* system, int cudaDev1, int cudaDev2, bool* isXGMI, int maxInter=MAX_XGMI_INTER_GPUS, int nInter=0, int *inter=nullptr);
-int ncclPxnDisable();
+int ncclPxnDisable(struct ncclComm* comm);
 ncclResult_t ncclTopoGetPxnRanks(struct ncclComm* comm, int** intermediateRanks, int* nranks);
 ncclResult_t ncclTopoGetLocalRank(struct ncclTopoSystem* system, int rank, int* localRank);
 

src/include/group.h

@@ -11,15 +11,78 @@
 #include "nccl.h"
 #include "comm.h"
 
-bool ncclAsyncMode();
-ncclResult_t ncclAsyncErrCheck(ncclResult_t ret);
+ncclResult_t ncclGroupErrCheck(ncclResult_t ret);
+void ncclGroupCommJoin(struct ncclComm* comm);
+void ncclGroupCommPreconnect(struct ncclComm* comm);
+void ncclGroupCommLeave(struct ncclComm* comm);
 
-typedef ncclResult_t(*ncclInitFunc_t)(ncclComm_t* newcomm, int ndev, ncclUniqueId commId, int myrank, int cudaDev, int virtualId);
+struct ncclAsyncJob {
+  struct ncclAsyncJob* next;
+  pthread_t thread;
+  ncclResult_t result;
+  ncclResult_t(*func)(struct ncclAsyncJob*);
+  void(*undo)(struct ncclAsyncJob*);
+  void(*destructor)(void*);
+};
 
-ncclResult_t ncclAsyncInit(ncclInitFunc_t func, ncclComm_t* newcomm, int ndev, ncclUniqueId commId, int myrank, int cudaDev, int virtualId);
+ncclResult_t ncclAsyncLaunch(
+  struct ncclAsyncJob* job,
+  ncclResult_t(*func)(struct ncclAsyncJob*),
+  void(*undo)(struct ncclAsyncJob*),
+  void(*destructor)(void*)
+);
 
-typedef ncclResult_t(*ncclCollFunc_t)(const void* sendbuff, void* recvbuff, size_t count,
-    ncclDataType_t type, ncclRedOp_t op, int root, ncclComm_t comm, hipStream_t stream);
+ncclResult_t ncclGroupStartInternal();
+ncclResult_t ncclGroupEndInternal();
+
+////////////////////////////////////////////////////////////////////////////////
+
+extern __thread int ncclGroupDepth; // depth of ncclGroupStart nesting
+extern __thread ncclResult_t ncclGroupError;
+extern __thread struct ncclComm* ncclGroupCommHead;
+extern __thread struct ncclComm* ncclGroupCommPreconnectHead;
+
+inline ncclResult_t ncclGroupStartInternal() {
+  ncclGroupDepth++;
+  return ncclSuccess;
+}
+
+inline ncclResult_t ncclGroupErrCheck(ncclResult_t ret) {
+  if (ncclGroupDepth > 0) {
+    if (ncclGroupError == ncclSuccess || ret != ncclSuccess) ncclGroupError = ret;
+  }
+  return ret;
+}
+
+// Add comm to this thread's group
+inline void ncclGroupCommJoin(struct ncclComm* comm) {
+  if (comm->groupNext == reinterpret_cast<struct ncclComm*>(0x1)) {
+    // Insert comm into ncclGroupCommHead adjacent to sibling comms. This preserves
+    // the users program order yet insures siblings occur consecutively. This
+    // is required by doLaunches() in "group.cc".
+    struct ncclComm** pp = &ncclGroupCommHead;
+    while (*pp != nullptr && comm->intraComm0 != (*pp)->intraComm0)
+      pp = &(*pp)->groupNext;
+    comm->groupNext = *pp;
+    *pp = comm;
+    // Comms gets a new memory stack scope upon joining. Each task batched for
+    // this comm is allocated there.
+    ncclMemoryStackPush(&comm->memScoped);
+  }
+}
+
+// Add comm to this thread's group needing preconnect
+inline void ncclGroupCommPreconnect(struct ncclComm* comm) {
+  if (comm->preconnectNext == reinterpret_cast<struct ncclComm*>(0x1)) {
+    comm->preconnectNext = ncclGroupCommPreconnectHead;
+    ncclGroupCommPreconnectHead = comm;
+  }
+}
+
+// Comm has left group
+inline void ncclGroupCommLeave(struct ncclComm* comm) {
+  comm->groupNext = reinterpret_cast<struct ncclComm*>(0x1);
+  ncclMemoryStackPop(&comm->memScoped);
+}
 
-ncclResult_t ncclAsyncColl(ncclComm_t comm);
 #endif

src/include/ibvwrap.h

@@ -1067,6 +1067,9 @@ ncclResult_t wrap_ibv_dealloc_pd(struct ibv_pd *pd);
 ncclResult_t wrap_ibv_reg_mr(struct ibv_mr **ret, struct ibv_pd *pd, void *addr, size_t length, int access);
 struct ibv_mr * wrap_direct_ibv_reg_mr(struct ibv_pd *pd, void *addr, size_t length, int access);
 ncclResult_t wrap_ibv_reg_mr_iova2(struct ibv_mr **ret, struct ibv_pd *pd, void *addr, size_t length, uint64_t iova, int access);
+/* DMA-BUF support */
+ncclResult_t wrap_ibv_reg_dmabuf_mr(struct ibv_mr **ret, struct ibv_pd *pd, uint64_t offset, size_t length, uint64_t iova, int fd, int access);
+struct ibv_mr * wrap_direct_ibv_reg_dmabuf_mr(struct ibv_pd *pd, uint64_t offset, size_t length, uint64_t iova, int fd, int access);
 ncclResult_t wrap_ibv_dereg_mr(struct ibv_mr *mr);
 ncclResult_t wrap_ibv_create_comp_channel(struct ibv_comp_channel **ret, struct ibv_context *context);
 ncclResult_t wrap_ibv_destroy_comp_channel(struct ibv_comp_channel *channel);

src/include/info.h

@@ -11,6 +11,9 @@
 #include "nccl.h"
 #include "devcomm.h"
 #include "collectives.h"
+#include "core.h"
+#include "utils.h"
+#include "strongstream.h"
 
 typedef enum : uint8_t {
   ncclPatternRing,
@@ -53,8 +56,66 @@ struct ncclInfo {
   int nchunksPerLoop;
   int chunkSize;
   int channelId;
-  uint16_t connIndex;
-  uint64_t opCount;
+};
+
+inline ncclResult_t ncclInfoSetDerived(struct ncclInfo* info, int nRanks) {
+  info->nBytes = info->count * ncclTypeSize(info->datatype);
+  if (info->coll == ncclFuncAllGather || info->coll == ncclFuncBroadcast || info->coll == ncclFuncAllToAllPivot) {
+    info->count = info->nBytes;
+    info->datatype = ncclInt8;
+  }
+  if (info->coll == ncclFuncAllGather || info->coll == ncclFuncReduceScatter) info->nBytes *= nRanks; // count is per rank
+  return ncclSuccess;
+}
+
+struct ncclTaskColl {
+  struct ncclTaskColl* next;
+  ncclFunc_t func;
+  void const* sendbuff;
+  void* recvbuff;
+  size_t count;
+  int root;
+  ncclDataType_t datatype;
+  ncclDevRedOpFull op;
+  int chunkSteps, sliceSteps;
+};
+struct ncclTaskP2p {
+  ncclTaskP2p *next;
+  void *buff;
+  size_t bytes;
+  // Stateful chunk index. If a p2p gets "cut" over two plans this keeps track
+  // of where it left off.
+  int chunk;
+};
+
+struct ncclCudaStreamList {
+  struct ncclCudaStreamList *next;
+  hipStream_t stream;
+};
+
+struct ncclTasks {
+  struct Peer {
+    bool sendSeen, recvSeen;
+    struct ncclIntruQueue<struct ncclTaskP2p, &ncclTaskP2p::next> sendQueue;
+    struct ncclIntruQueue<struct ncclTaskP2p, &ncclTaskP2p::next> recvQueue;
+  };
+  struct ncclIntruQueue<ncclTaskColl, &ncclTaskColl::next> collQueue;
+  size_t collBytesTotal;
+  struct Peer* peers/*[nRanks]*/;
+  int *p2pSendOrder/*[nRanks]*/, *p2pRecvOrder/*[nRanks]*/;
+  int nTasksColl, nTasksP2p;
+
+  // The list of user streams aggregated over all tasks present.
+  struct ncclCudaStreamList* streams;
+  // Keep track of the number of user streams
+  int numStreams;
+  // The most recent user stream. Ignored if streams==nullptr
+  hipStream_t streamRecent;
+  // The graph capturing all user streams or invalid if none. Thus we restrict the
+  // user that all streams must be captured in the same graph or not captured
+  // at all. Technically we could probably relax this, but that would mean
+  // collecting a different `ncclTasks` per graph and one for non-graph.
+  struct ncclCudaGraph capturingGraph;
 };
 
 #endif

src/include/nccl_net.h

@@ -14,12 +14,13 @@
 
 #define NCCL_PTR_HOST 0x1
 #define NCCL_PTR_CUDA 0x2
+#define NCCL_PTR_DMABUF 0x4
 
 // Maximum number of requests per comm object
 #define NCCL_NET_MAX_REQUESTS 8
 
 typedef enum {NCCL_LOG_NONE=0, NCCL_LOG_VERSION=1, NCCL_LOG_WARN=2, NCCL_LOG_INFO=3, NCCL_LOG_ABORT=4, NCCL_LOG_TRACE=5} ncclDebugLogLevel;
-typedef enum {NCCL_INIT=1, NCCL_COLL=2, NCCL_P2P=4, NCCL_SHM=8, NCCL_NET=16, NCCL_GRAPH=32, NCCL_TUNING=64, NCCL_ENV=128, NCCL_ALLOC=256, NCCL_ALL=~0} ncclDebugLogSubSys;
+typedef enum {NCCL_INIT=1, NCCL_COLL=2, NCCL_P2P=4, NCCL_SHM=8, NCCL_NET=16, NCCL_GRAPH=32, NCCL_TUNING=64, NCCL_ENV=128, NCCL_ALLOC=256, NCCL_CALL=512, NCCL_ALL=~0} ncclDebugLogSubSys;
 
 typedef void (*ncclDebugLogger_t)(ncclDebugLogLevel level, unsigned long flags, const char *file, int line, const char *fmt, ...);
 
@@ -28,15 +29,15 @@ typedef struct {
   char* pciPath;  // Path to the PCI device in /sys.
   uint64_t guid;  // Unique identifier for the NIC chip. Important for
                   // cards with multiple PCI functions (Physical or virtual).
-  int ptrSupport; // NCCL_PTR_HOST or NCCL_PTR_HOST|NCCL_PTR_CUDA
+  int ptrSupport; // [NCCL_PTR_HOST|NCCL_PTR_CUDA|NCCL_PTR_DMABUF]
   int speed;      // Port speed in Mbps.
   int port;       // Port number.
   float latency;  // Network latency
   int maxComms;   // Maximum number of comms we can create
   int maxRecvs;   // Maximum number of grouped receives.
-}ncclNetProperties_v5_t;
+}ncclNetProperties_v6_t;
 
-typedef ncclNetProperties_v5_t ncclNetProperties_t;
+typedef ncclNetProperties_v6_t ncclNetProperties_t;
 
 typedef struct {
   // Name of the network (mainly for logs)
@@ -46,7 +47,103 @@ typedef struct {
   // Return the number of adapters.
   ncclResult_t (*devices)(int* ndev);
   // Get various device properties.
-  ncclResult_t (*getProperties)(int dev, ncclNetProperties_v5_t* props);
+  ncclResult_t (*getProperties)(int dev, ncclNetProperties_v6_t* props);
+  // Create a receiving object and provide a handle to connect to it. The
+  // handle can be up to NCCL_NET_HANDLE_MAXSIZE bytes and will be exchanged
+  // between ranks to create a connection.
+  ncclResult_t (*listen)(int dev, void* handle, void** listenComm);
+  // Connect to a handle and return a sending comm object for that peer.
+  // This call must not block for the connection to be established, and instead
+  // should return successfully with sendComm == NULL with the expectation that
+  // it will be called again until sendComm != NULL.
+  ncclResult_t (*connect)(int dev, void* handle, void** sendComm);
+  // Finalize connection establishment after remote peer has called connect.
+  // This call must not block for the connection to be established, and instead
+  // should return successfully with recvComm == NULL with the expectation that
+  // it will be called again until recvComm != NULL.
+  ncclResult_t (*accept)(void* listenComm, void** recvComm);
+  // Register/Deregister memory. Comm can be either a sendComm or a recvComm.
+  // Type is either NCCL_PTR_HOST or NCCL_PTR_CUDA.
+  ncclResult_t (*regMr)(void* comm, void* data, int size, int type, void** mhandle);
+  /* DMA-BUF support */
+  ncclResult_t (*regMrDmaBuf)(void* comm, void* data, size_t size, int type, uint64_t offset, int fd, void** mhandle);
+  ncclResult_t (*deregMr)(void* comm, void* mhandle);
+  // Asynchronous send to a peer.
+  // May return request == NULL if the call cannot be performed (or would block)
+  ncclResult_t (*isend)(void* sendComm, void* data, int size, int tag, void* mhandle, void** request);
+  // Asynchronous recv from a peer.
+  // May return request == NULL if the call cannot be performed (or would block)
+  ncclResult_t (*irecv)(void* recvComm, int n, void** data, int* sizes, int* tags, void** mhandles, void** request);
+  // Perform a flush/fence to make sure all data received with NCCL_PTR_CUDA is
+  // visible to the GPU
+  ncclResult_t (*iflush)(void* recvComm, int n, void** data, int* sizes, void** mhandles, void** request);
+  // Test whether a request is complete. If size is not NULL, it returns the
+  // number of bytes sent/received.
+  ncclResult_t (*test)(void* request, int* done, int* sizes);
+  // Close and free send/recv comm objects
+  ncclResult_t (*closeSend)(void* sendComm);
+  ncclResult_t (*closeRecv)(void* recvComm);
+  ncclResult_t (*closeListen)(void* listenComm);
+} ncclNet_v6_t;
+
+typedef ncclNet_v6_t ncclNet_t;
+
+#define NCCL_PLUGIN_SYMBOL ncclNetPlugin_v6
+
+typedef struct {
+  // Name of the collective network (mainly for logs)
+  const char* name;
+  // Initialize the collective network.
+  ncclResult_t (*init)(ncclDebugLogger_t logFunction);
+  // Return the number of adapters capable of doing collective operations.
+  // If ndev returns 0, all other functions might be set to NULL.
+  ncclResult_t (*devices)(int* ndev);
+  // Get various device properties.
+  ncclResult_t (*getProperties)(int dev, ncclNetProperties_v6_t* props);
+  // Create a receiving object and provide a handle to connect to it. The
+  // handle can be up to NCCL_NET_HANDLE_MAXSIZE bytes and will be exchanged
+  // between ranks to create connections.
+  ncclResult_t (*listen)(int dev, void* handle, void** listenComm);
+  // Create a group for collective operations. handles have been created
+  // using listen() above. rank indicates caller's rank in the collective network.
+  ncclResult_t (*connect)(void* handles[], int nranks, int rank, void* listenComm, void** collComm);
+  // Returns whether a reduction operation on a data type is supported.
+  // 1 for supported, 0 otherwise.
+  ncclResult_t (*reduceSupport)(ncclDataType_t dataType, ncclRedOp_t redOp, int* supported);
+  // Register/Deregister memory. Type is either NCCL_PTR_HOST or NCCL_PTR_CUDA.
+  ncclResult_t (*regMr)(void* collComm, void* data, int size, int type, void** mhandle);
+  /* DMA-BUF support */
+  ncclResult_t (*regMrDmaBuf)(void* collComm, void* data, size_t size, int type, uint64_t offset, int fd, void** mhandle);
+  ncclResult_t (*deregMr)(void* collComm, void* mhandle);
+  // Performs an asynchronous allreduce operation on the collective group.
+  // May return request == NULL if the call cannot be performed (or would block).
+  ncclResult_t (*iallreduce)(void* collComm, void* sendData, void* recvData, int count,
+      ncclDataType_t dataType, ncclRedOp_t redOp, void* sendMhandle, void* recvMhandle, void** request);
+  // Perform a flush/fence to make sure all data received with NCCL_PTR_CUDA is
+  // visible to the GPU
+  ncclResult_t (*iflush)(void* collComm, void* data, int size, void* mhandle, void** request);
+  // Test whether a request is complete. If size is not NULL, it returns the
+  // number of bytes sent/received.
+  ncclResult_t (*test)(void* request, int* done, int* size);
+  // Close and free collective comm objects
+  ncclResult_t (*closeColl)(void* collComm);
+  ncclResult_t (*closeListen)(void* listenComm);
+} ncclCollNet_v6_t;
+
+typedef ncclCollNet_v6_t ncclCollNet_t;
+
+#define NCCL_COLLNET_PLUGIN_SYMBOL ncclCollNetPlugin_v6
+
+// v5 struct for backwards compatibility
+typedef struct {
+  // Name of the network (mainly for logs)
+  const char* name;
+  // Initialize the network.
+  ncclResult_t (*init)(ncclDebugLogger_t logFunction);
+  // Return the number of adapters.
+  ncclResult_t (*devices)(int* ndev);
+  // Get various device properties.
+  ncclResult_t (*getProperties)(int dev, ncclNetProperties_v6_t* props);
   // Create a receiving object and provide a handle to connect to it. The
   // handle can be up to NCCL_NET_HANDLE_MAXSIZE bytes and will be exchanged
   // between ranks to create a connection.
@@ -83,10 +180,7 @@ typedef struct {
   ncclResult_t (*closeListen)(void* listenComm);
 } ncclNet_v5_t;
 
-typedef ncclNet_v5_t ncclNet_t;
-
-#define NCCL_PLUGIN_SYMBOL ncclNetPlugin_v5
-
+// v5 struct for backwards compatibility
 typedef struct {
   // Name of the collective network (mainly for logs)
   const char* name;
@@ -96,7 +190,7 @@ typedef struct {
   // If ndev returns 0, all other functions might be set to NULL.
   ncclResult_t (*devices)(int* ndev);
   // Get various device properties.
-  ncclResult_t (*getProperties)(int dev, ncclNetProperties_v5_t* props);
+  ncclResult_t (*getProperties)(int dev, ncclNetProperties_v6_t* props);
   // Create a receiving object and provide a handle to connect to it. The
   // handle can be up to NCCL_NET_HANDLE_MAXSIZE bytes and will be exchanged
   // between ranks to create connections.
@@ -125,10 +219,7 @@ typedef struct {
   ncclResult_t (*closeListen)(void* listenComm);
 } ncclCollNet_v5_t;
 
-typedef ncclCollNet_v5_t ncclCollNet_t;
-
-#define NCCL_COLLNET_PLUGIN_SYMBOL ncclCollNetPlugin_v5
-
+// v4 struct for backwards compatibility
 typedef struct {
   char* name;     // Used mostly for logging.
   char* pciPath;  // Path to the PCI device in /sys.
@@ -140,6 +231,7 @@ typedef struct {
   int maxComms;   // Maximum number of comms we can create
 } ncclNetProperties_v4_t;
 
+// v4 struct for backwards compatibility
 typedef struct {
   // Name of the network (mainly for logs)
   const char* name;
@@ -179,6 +271,7 @@ typedef struct {
   ncclResult_t (*closeListen)(void* listenComm);
 } ncclNet_v4_t;
 
+// v4 struct for backwards compatibility
 typedef struct {
   // Name of the collective network (mainly for logs)
   const char* name;

src/include/net.h

@@ -9,33 +9,36 @@
 
 #include "nccl.h"
 #include "nccl_net.h"
+#include "comm.h"
 #include "checks.h"
 
-extern ncclNet_t* ncclNet;
 typedef char ncclNetHandle_t[NCCL_NET_HANDLE_MAXSIZE];
 
-ncclResult_t ncclNetInit();
-int ncclNetVersion();
+ncclResult_t ncclNetPluginInit();
+ncclResult_t ncclNetInit(struct ncclComm* comm);
+int ncclNetVersion(struct ncclComm* comm);
 
 // Translation to external API
-static const char* ncclNetName() { return ncclNet->name; }
-static ncclResult_t ncclNetDevices(int* ndev) { NCCLCHECK(ncclNet->devices(ndev)); return ncclSuccess; }
-static ncclResult_t ncclNetGetProperties(int dev, ncclNetProperties_t* props) { NCCLCHECK(ncclNet->getProperties(dev, props)); return ncclSuccess; }
-static ncclResult_t ncclNetListen(int dev, void* handle, void** listenComm) { NCCLCHECK(ncclNet->listen(dev, handle, listenComm)); return ncclSuccess; }
-static ncclResult_t ncclNetConnect(int dev, void* handle, void** sendComm) { NCCLCHECK(ncclNet->connect(dev, handle, sendComm)); return ncclSuccess; }
-static ncclResult_t ncclNetAccept(void* listenComm, void** recvComm) { NCCLCHECK(ncclNet->accept(listenComm, recvComm)); return ncclSuccess; }
-static ncclResult_t ncclNetRegMr(void* comm, void* data, int size, int type, void** mhandle) { NCCLCHECK(ncclNet->regMr(comm, data, size, type, mhandle)); return ncclSuccess; }
-static ncclResult_t ncclNetDeregMr(void* comm, void* mhandle) { NCCLCHECK(ncclNet->deregMr(comm, mhandle)); return ncclSuccess; }
-static ncclResult_t ncclNetIsend(void* sendComm, void* data, int size, int tag, void* mhandle, void** request) { NCCLCHECK(ncclNet->isend(sendComm, data, size, tag, mhandle, request)); return ncclSuccess; }
-static ncclResult_t ncclNetIrecv(void* recvComm, int n, void** data, int* sizes, int* tags, void** mhandles, void** request) { NCCLCHECK(ncclNet->irecv(recvComm, n, data, sizes, tags, mhandles, request)); return ncclSuccess; }
-static ncclResult_t ncclNetIflush(void* recvComm, int n, void** data, int* sizes, void** mhandles, void** request) { NCCLCHECK(ncclNet->iflush(recvComm, n, data, sizes, mhandles, request)); return ncclSuccess; }
-static ncclResult_t ncclNetTest(void* request, int* done, int* sizes) { NCCLCHECK(ncclNet->test(request, done, sizes)); return ncclSuccess; }
-static ncclResult_t ncclNetCloseSend(void* sendComm) { NCCLCHECK(ncclNet->closeSend(sendComm)); return ncclSuccess; }
-static ncclResult_t ncclNetCloseRecv(void* recvComm) { NCCLCHECK(ncclNet->closeRecv(recvComm)); return ncclSuccess; }
-static ncclResult_t ncclNetCloseListen(void* listenComm) { NCCLCHECK(ncclNet->closeListen(listenComm)); return ncclSuccess; }
+static const char* ncclNetName(struct ncclComm* comm) { return comm->ncclNet->name; }
+static ncclResult_t ncclNetDevices(struct ncclComm* comm, int* ndev) { NCCLCHECK(comm->ncclNet->devices(ndev)); return ncclSuccess; }
+static ncclResult_t ncclNetGetProperties(struct ncclComm* comm, int dev, ncclNetProperties_t* props) { NCCLCHECK(comm->ncclNet->getProperties(dev, props)); return ncclSuccess; }
+static ncclResult_t ncclNetListen(struct ncclComm* comm, int dev, void* handle, void** listenComm) { NCCLCHECK(comm->ncclNet->listen(dev, handle, listenComm)); return ncclSuccess; }
+static ncclResult_t ncclNetConnect(struct ncclComm* comm, int dev, void* handle, void** sendComm) { NCCLCHECK(comm->ncclNet->connect(dev, handle, sendComm)); return ncclSuccess; }
+static ncclResult_t ncclNetAccept(struct ncclComm* comm, void* listenComm, void** recvComm) { NCCLCHECK(comm->ncclNet->accept(listenComm, recvComm)); return ncclSuccess; }
+static ncclResult_t ncclNetRegMr(struct ncclComm* comm, void* netComm, void* data, int size, int type, void** mhandle) { NCCLCHECK(comm->ncclNet->regMr(netComm, data, size, type, mhandle)); return ncclSuccess; }
+/* DMA-BUF support */
+static ncclResult_t ncclNetRegMrDmaBuf(struct ncclComm* comm, void* netComm, void* data, size_t size, int type, uint64_t offset, int fd, void** mhandle) { NCCLCHECK(comm->ncclNet->regMrDmaBuf(netComm, data, size, type, offset, fd, mhandle)); return ncclSuccess; }
+static ncclResult_t ncclNetDeregMr(struct ncclComm* comm, void* netComm, void* mhandle) { NCCLCHECK(comm->ncclNet->deregMr(netComm, mhandle)); return ncclSuccess; }
+static ncclResult_t ncclNetIsend(struct ncclComm* comm, void* sendComm, void* data, int size, int tag, void* mhandle, void** request) { NCCLCHECK(comm->ncclNet->isend(sendComm, data, size, tag, mhandle, request)); return ncclSuccess; }
+static ncclResult_t ncclNetIrecv(struct ncclComm* comm, void* recvComm, int n, void** data, int* sizes, int* tags, void** mhandles, void** request) { NCCLCHECK(comm->ncclNet->irecv(recvComm, n, data, sizes, tags, mhandles, request)); return ncclSuccess; }
+static ncclResult_t ncclNetIflush(struct ncclComm* comm, void* recvComm, int n, void** data, int* sizes, void** mhandles, void** request) { NCCLCHECK(comm->ncclNet->iflush(recvComm, n, data, sizes, mhandles, request)); return ncclSuccess; }
+static ncclResult_t ncclNetTest(struct ncclComm* comm, void* request, int* done, int* sizes) { NCCLCHECK(comm->ncclNet->test(request, done, sizes)); return ncclSuccess; }
+static ncclResult_t ncclNetCloseSend(struct ncclComm* comm, void* sendComm) { NCCLCHECK(comm->ncclNet->closeSend(sendComm)); return ncclSuccess; }
+static ncclResult_t ncclNetCloseRecv(struct ncclComm* comm, void* recvComm) { NCCLCHECK(comm->ncclNet->closeRecv(recvComm)); return ncclSuccess; }
+static ncclResult_t ncclNetCloseListen(struct ncclComm* comm, void* listenComm) { NCCLCHECK(comm->ncclNet->closeListen(listenComm)); return ncclSuccess; }
 
 // Test whether the current GPU support GPU Direct RDMA.
-ncclResult_t ncclGpuGdrSupport(int* gdrSupport);
+ncclResult_t ncclGpuGdrSupport(struct ncclComm* comm, int* gdrSupport);
 
 extern ncclNet_t ncclNetIb;
 extern ncclNet_t ncclNetSocket;

src/include/nvtx3/nvToolsExtCuda.h

@@ -8,7 +8,7 @@
 
 #include "nvToolsExt.h"
 
-#include "cuda.h"
+#include "hip/hip_runtime.h"
 
 #ifndef NVTOOLSEXT_CUDA_V3
 #define NVTOOLSEXT_CUDA_V3
@@ -42,10 +42,10 @@ extern "C" {
 */
 typedef enum nvtxResourceCUDAType_t
 {
-    NVTX_RESOURCE_TYPE_CUDA_DEVICE = NVTX_RESOURCE_MAKE_TYPE(CUDA, 1), /* CUdevice */
-    NVTX_RESOURCE_TYPE_CUDA_CONTEXT = NVTX_RESOURCE_MAKE_TYPE(CUDA, 2), /* CUcontext */
-    NVTX_RESOURCE_TYPE_CUDA_STREAM = NVTX_RESOURCE_MAKE_TYPE(CUDA, 3), /* CUstream */
-    NVTX_RESOURCE_TYPE_CUDA_EVENT = NVTX_RESOURCE_MAKE_TYPE(CUDA, 4), /* CUevent */
+    NVTX_RESOURCE_TYPE_CUDA_DEVICE = NVTX_RESOURCE_MAKE_TYPE(CUDA, 1), /* hipDevice_t */
+    NVTX_RESOURCE_TYPE_CUDA_CONTEXT = NVTX_RESOURCE_MAKE_TYPE(CUDA, 2), /* hipCtx_t */
+    NVTX_RESOURCE_TYPE_CUDA_STREAM = NVTX_RESOURCE_MAKE_TYPE(CUDA, 3), /* hipStream_t */
+    NVTX_RESOURCE_TYPE_CUDA_EVENT = NVTX_RESOURCE_MAKE_TYPE(CUDA, 4), /* hipEvent_t */
 } nvtxResourceCUDAType_t;
 
 
@@ -59,8 +59,8 @@ typedef enum nvtxResourceCUDAType_t
  *
  * \version \NVTX_VERSION_1
  * @{ */
-NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceA(CUdevice device, const char* name);
-NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceW(CUdevice device, const wchar_t* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceA(hipDevice_t device, const char* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceW(hipDevice_t device, const wchar_t* name);
 /** @} */
 
 /* ------------------------------------------------------------------------- */
@@ -73,16 +73,16 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceW(CUdevice device, const wchar_t* na
  *
  * \par Example:
  * \code
- * CUresult status = cuCtxCreate( &cuContext, 0, cuDevice );
- * if ( CUDA_SUCCESS != status )
+ * hipError_t status = hipCtxCreate( &cuContext, 0, cuDevice );
+ * if ( hipSuccess != status )
  *     goto Error;
  * nvtxNameCuContext(cuContext, "CTX_NAME");
  * \endcode
  *
  * \version \NVTX_VERSION_1
  * @{ */
-NVTX_DECLSPEC void NVTX_API nvtxNameCuContextA(CUcontext context, const char* name);
-NVTX_DECLSPEC void NVTX_API nvtxNameCuContextW(CUcontext context, const wchar_t* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCuContextA(hipCtx_t context, const char* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCuContextW(hipCtx_t context, const wchar_t* name);
 /** @} */
 
 /* ------------------------------------------------------------------------- */
@@ -95,8 +95,8 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuContextW(CUcontext context, const wchar_t*
  *
  * \version \NVTX_VERSION_1
  * @{ */
-NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamA(CUstream stream, const char* name);
-NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamW(CUstream stream, const wchar_t* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamA(hipStream_t stream, const char* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamW(hipStream_t stream, const wchar_t* name);
 /** @} */
 
 /* ------------------------------------------------------------------------- */
@@ -109,8 +109,8 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamW(CUstream stream, const wchar_t* na
  *
  * \version \NVTX_VERSION_1
  * @{ */
-NVTX_DECLSPEC void NVTX_API nvtxNameCuEventA(CUevent event, const char* name);
-NVTX_DECLSPEC void NVTX_API nvtxNameCuEventW(CUevent event, const wchar_t* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCuEventA(hipEvent_t event, const char* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCuEventW(hipEvent_t event, const wchar_t* name);
 /** @} */
 
 /** @} */ /* END RESOURCE_NAMING */

src/include/nvtx3/nvToolsExtCudaRt.h

@@ -8,8 +8,8 @@
 
 #include "nvToolsExt.h"
 
-#include "cuda.h"
-#include "driver_types.h"
+#include "hip/hip_runtime.h"
+#include "hip/driver_types.h"
 
 #ifndef NVTOOLSEXT_CUDART_V3
 #define NVTOOLSEXT_CUDART_V3
@@ -44,8 +44,8 @@ extern "C" {
 typedef enum nvtxResourceCUDARTType_t
 {
     NVTX_RESOURCE_TYPE_CUDART_DEVICE = NVTX_RESOURCE_MAKE_TYPE(CUDART, 0), /* int device */
-    NVTX_RESOURCE_TYPE_CUDART_STREAM = NVTX_RESOURCE_MAKE_TYPE(CUDART, 1), /* cudaStream_t */
-    NVTX_RESOURCE_TYPE_CUDART_EVENT = NVTX_RESOURCE_MAKE_TYPE(CUDART, 2), /* cudaEvent_t */
+    NVTX_RESOURCE_TYPE_CUDART_STREAM = NVTX_RESOURCE_MAKE_TYPE(CUDART, 1), /* hipStream_t */
+    NVTX_RESOURCE_TYPE_CUDART_EVENT = NVTX_RESOURCE_MAKE_TYPE(CUDART, 2), /* hipEvent_t */
 } nvtxResourceCUDARTType_t;
 
 
@@ -73,8 +73,8 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCudaDeviceW(int device, const wchar_t* name)
  *
  * \version \NVTX_VERSION_1
  * @{ */
-NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamA(cudaStream_t stream, const char* name);
-NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamW(cudaStream_t stream, const wchar_t* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamA(hipStream_t stream, const char* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamW(hipStream_t stream, const wchar_t* name);
 /** @} */
 
 /* ------------------------------------------------------------------------- */
@@ -87,8 +87,8 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamW(cudaStream_t stream, const wchar
  *
  * \version \NVTX_VERSION_1
  * @{ */
-NVTX_DECLSPEC void NVTX_API nvtxNameCudaEventA(cudaEvent_t event, const char* name);
-NVTX_DECLSPEC void NVTX_API nvtxNameCudaEventW(cudaEvent_t event, const wchar_t* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCudaEventA(hipEvent_t event, const char* name);
+NVTX_DECLSPEC void NVTX_API nvtxNameCudaEventW(hipEvent_t event, const wchar_t* name);
 /** @} */
 
 /** @} */ /* END RESOURCE_NAMING */

src/include/nvtx3/nvtxDetail/nvtxImplCudaRt_v3.h

@@ -16,10 +16,10 @@ extern "C" {
 
 typedef void (NVTX_API * nvtxNameCudaDeviceA_impl_fntype)(int device, const char* name);
 typedef void (NVTX_API * nvtxNameCudaDeviceW_impl_fntype)(int device, const wchar_t* name);
-typedef void (NVTX_API * nvtxNameCudaStreamA_impl_fntype)(cudaStream_t stream, const char* name);
-typedef void (NVTX_API * nvtxNameCudaStreamW_impl_fntype)(cudaStream_t stream, const wchar_t* name);
-typedef void (NVTX_API * nvtxNameCudaEventA_impl_fntype)(cudaEvent_t event, const char* name);
-typedef void (NVTX_API * nvtxNameCudaEventW_impl_fntype)(cudaEvent_t event, const wchar_t* name);
+typedef void (NVTX_API * nvtxNameCudaStreamA_impl_fntype)(hipStream_t stream, const char* name);
+typedef void (NVTX_API * nvtxNameCudaStreamW_impl_fntype)(hipStream_t stream, const wchar_t* name);
+typedef void (NVTX_API * nvtxNameCudaEventA_impl_fntype)(hipEvent_t event, const char* name);
+typedef void (NVTX_API * nvtxNameCudaEventW_impl_fntype)(hipEvent_t event, const wchar_t* name);
 
 NVTX_DECLSPEC void NVTX_API nvtxNameCudaDeviceA(int device, const char* name)
 {
@@ -39,7 +39,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCudaDeviceW(int device, const wchar_t* name)
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamA(cudaStream_t stream, const char* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamA(hipStream_t stream, const char* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCudaStreamA_impl_fntype local = (nvtxNameCudaStreamA_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCudaStreamA_impl_fnptr;
@@ -48,7 +48,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamA(cudaStream_t stream, const char*
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamW(cudaStream_t stream, const wchar_t* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamW(hipStream_t stream, const wchar_t* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCudaStreamW_impl_fntype local = (nvtxNameCudaStreamW_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCudaStreamW_impl_fnptr;
@@ -57,7 +57,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCudaStreamW(cudaStream_t stream, const wchar
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCudaEventA(cudaEvent_t event, const char* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCudaEventA(hipEvent_t event, const char* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCudaEventA_impl_fntype local = (nvtxNameCudaEventA_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCudaEventA_impl_fnptr;
@@ -66,7 +66,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCudaEventA(cudaEvent_t event, const char* na
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCudaEventW(cudaEvent_t event, const wchar_t* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCudaEventW(hipEvent_t event, const wchar_t* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCudaEventW_impl_fntype local = (nvtxNameCudaEventW_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCudaEventW_impl_fnptr;

src/include/nvtx3/nvtxDetail/nvtxImplCuda_v3.h

@@ -15,16 +15,16 @@
 extern "C" {
 #endif /* __cplusplus */
 
-typedef void (NVTX_API * nvtxNameCuDeviceA_impl_fntype)(CUdevice device, const char* name);
-typedef void (NVTX_API * nvtxNameCuDeviceW_impl_fntype)(CUdevice device, const wchar_t* name);
-typedef void (NVTX_API * nvtxNameCuContextA_impl_fntype)(CUcontext context, const char* name);
-typedef void (NVTX_API * nvtxNameCuContextW_impl_fntype)(CUcontext context, const wchar_t* name);
-typedef void (NVTX_API * nvtxNameCuStreamA_impl_fntype)(CUstream stream, const char* name);
-typedef void (NVTX_API * nvtxNameCuStreamW_impl_fntype)(CUstream stream, const wchar_t* name);
-typedef void (NVTX_API * nvtxNameCuEventA_impl_fntype)(CUevent event, const char* name);
-typedef void (NVTX_API * nvtxNameCuEventW_impl_fntype)(CUevent event, const wchar_t* name);
+typedef void (NVTX_API * nvtxNameCuDeviceA_impl_fntype)(hipDevice_t device, const char* name);
+typedef void (NVTX_API * nvtxNameCuDeviceW_impl_fntype)(hipDevice_t device, const wchar_t* name);
+typedef void (NVTX_API * nvtxNameCuContextA_impl_fntype)(hipCtx_t context, const char* name);
+typedef void (NVTX_API * nvtxNameCuContextW_impl_fntype)(hipCtx_t context, const wchar_t* name);
+typedef void (NVTX_API * nvtxNameCuStreamA_impl_fntype)(hipStream_t stream, const char* name);
+typedef void (NVTX_API * nvtxNameCuStreamW_impl_fntype)(hipStream_t stream, const wchar_t* name);
+typedef void (NVTX_API * nvtxNameCuEventA_impl_fntype)(hipEvent_t event, const char* name);
+typedef void (NVTX_API * nvtxNameCuEventW_impl_fntype)(hipEvent_t event, const wchar_t* name);
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceA(CUdevice device, const char* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceA(hipDevice_t device, const char* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCuDeviceA_impl_fntype local = (nvtxNameCuDeviceA_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCuDeviceA_impl_fnptr;
@@ -33,7 +33,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceA(CUdevice device, const char* name)
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceW(CUdevice device, const wchar_t* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceW(hipDevice_t device, const wchar_t* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCuDeviceW_impl_fntype local = (nvtxNameCuDeviceW_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCuDeviceW_impl_fnptr;
@@ -42,7 +42,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuDeviceW(CUdevice device, const wchar_t* na
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCuContextA(CUcontext context, const char* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCuContextA(hipCtx_t context, const char* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCuContextA_impl_fntype local = (nvtxNameCuContextA_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCuContextA_impl_fnptr;
@@ -51,7 +51,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuContextA(CUcontext context, const char* na
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCuContextW(CUcontext context, const wchar_t* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCuContextW(hipCtx_t context, const wchar_t* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCuContextW_impl_fntype local = (nvtxNameCuContextW_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCuContextW_impl_fnptr;
@@ -60,7 +60,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuContextW(CUcontext context, const wchar_t*
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamA(CUstream stream, const char* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamA(hipStream_t stream, const char* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCuStreamA_impl_fntype local = (nvtxNameCuStreamA_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCuStreamA_impl_fnptr;
@@ -69,7 +69,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamA(CUstream stream, const char* name)
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamW(CUstream stream, const wchar_t* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamW(hipStream_t stream, const wchar_t* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCuStreamW_impl_fntype local = (nvtxNameCuStreamW_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCuStreamW_impl_fnptr;
@@ -78,7 +78,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuStreamW(CUstream stream, const wchar_t* na
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCuEventA(CUevent event, const char* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCuEventA(hipEvent_t event, const char* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCuEventA_impl_fntype local = (nvtxNameCuEventA_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCuEventA_impl_fnptr;
@@ -87,7 +87,7 @@ NVTX_DECLSPEC void NVTX_API nvtxNameCuEventA(CUevent event, const char* name)
 #endif /*NVTX_DISABLE*/
 }
 
-NVTX_DECLSPEC void NVTX_API nvtxNameCuEventW(CUevent event, const wchar_t* name)
+NVTX_DECLSPEC void NVTX_API nvtxNameCuEventW(hipEvent_t event, const wchar_t* name)
 {
 #ifndef NVTX_DISABLE
     nvtxNameCuEventW_impl_fntype local = (nvtxNameCuEventW_impl_fntype)NVTX_VERSIONED_IDENTIFIER(nvtxGlobals).nvtxNameCuEventW_impl_fnptr;

src/include/nvtx3/nvtxDetail/nvtxTypes.h

@@ -18,7 +18,7 @@
 /* ------ Dependency-free types binary-compatible with real types ------- */
 
 /* In order to avoid having the NVTX core API headers depend on non-NVTX
-*  headers like cuda.h, NVTX defines binary-compatible types to use for
+*  headers like hip/hip_runtime.h, NVTX defines binary-compatible types to use for
 *  safely making the initialization versions of all NVTX functions without
 *  needing to have definitions for the real types. */
 

src/include/p2p.h

@@ -9,21 +9,4 @@
 #ifndef NCCL_P2P_H_
 #define NCCL_P2P_H_
 
-struct ncclP2Pinfo {
-  void* buff;
-  ssize_t nbytes;
-  uint64_t opCount;
-};
-
-typedef ncclRecyclableList<struct ncclP2Pinfo> ncclP2Plist;
-
-static ncclResult_t ncclSaveP2pInfo(ncclP2Plist* &p2p, void* buff, ssize_t nBytes, uint64_t opCount) {
-  if (p2p == NULL) p2p = new ncclP2Plist();
-  struct ncclP2Pinfo* next;
-  NCCLCHECK(p2p->getNewElem(&next));
-  next->buff = buff;
-  next->nbytes = nBytes;
-  next->opCount = opCount;
-  return ncclSuccess;
-}
 #endif

src/include/proxy.h

@@ -26,18 +26,26 @@ struct ncclProxyOp {
   int channelId;
   int nsteps;
   ssize_t nbytes;
-  int root;
+  struct {
+    int root:30;
+    uint32_t connIndex:2;
+  };
   int next;
 
   uint64_t opCount;
   int sliceSteps;
   int chunkSteps;
   int chunkSize;
-  ncclDataType_t dtype;
-  ncclRedOp_t redOp;
-  ncclPattern_t pattern; // uint8_t
+  uint8_t /*ncclDataType_t*/ dtype;
+  uint8_t /*ncclDevRedOp_t*/ redOp;
+  uint8_t /*ncclPattern_t*/ pattern;
   uint8_t protocol;
-  uint16_t connIndex;
+
+  union {
+    uint64_t unused;
+    // For use by enqueue.cc
+    struct ncclProxyOp *enqNext;
+  };
 };
 static_assert(sizeof(struct ncclProxyOp) == 64, "Keep ProxyOp aligned with cache lines for effective prefetch");
 
@@ -73,9 +81,9 @@ struct ncclProxyArgs {
   int sliceSteps;
   int chunkSteps;
   int chunkSize;
-  ncclDataType_t dtype;
-  ncclRedOp_t redOp;
-  ncclPattern_t pattern;
+  uint8_t /*ncclDataType_t*/ dtype;
+  uint8_t /*ncclDevRedOp_t*/ redOp;
+  uint8_t /*ncclPattern_t*/ pattern;
   uint8_t protocol;
   int state;
   char* sharedBuff[NCCL_STEPS];
@@ -164,6 +172,7 @@ struct ncclProxyState {
   pthread_t thread;
   struct ncclSocket* listenSock;
   int stop;
+  hipCtx_t cudaCtx;
 
   // Used by main thread
   union ncclSocketAddress* peerAddresses;
@@ -193,9 +202,8 @@ enum proxyMode {
   proxyTo = 2
 };
 
-ncclResult_t ncclProxySaveColl(struct ncclComm* comm, struct ncclProxyOp* proxyOp, int nranks);
+ncclResult_t ncclProxySaveOp(struct ncclComm* comm, struct ncclProxyOp* proxyOp, bool *justInquire);
 ncclResult_t ncclProxyComputeP2p(struct ncclInfo* info, struct ncclProxyOp* proxyOp);
-ncclResult_t ncclProxySaveP2p(struct ncclComm* comm, struct ncclProxyOp* proxyOp);
 ncclResult_t ncclProxyStart(struct ncclComm* comm);
 ncclResult_t ncclProxyInit(struct ncclComm* comm, struct ncclSocket* sock, union ncclSocketAddress* peerAddresses);
 ncclResult_t ncclProxyCreate(struct ncclComm* comm);

src/include/rocmwrap.h

@@ -0,0 +1,73 @@
+/*************************************************************************
+ * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef NCCL_ROCMWRAP_H_
+#define NCCL_ROCMWRAP_H_
+
+#include <hsa/hsa.h>
+
+typedef hsa_status_t (*PFN_hsa_init)();
+typedef hsa_status_t (*PFN_hsa_system_get_info)(hsa_system_info_t attribute, void* value);
+typedef hsa_status_t (*PFN_hsa_status_string)(hsa_status_t status, const char ** status_string);
+typedef hsa_status_t (*PFN_hsa_amd_portable_export_dmabuf)(const void* ptr, size_t size, int* dmabuf, uint64_t* offset);
+
+
+#define CUPFN(symbol) pfn_##symbol
+
+// Check CUDA PFN driver calls
+#define CUCHECK(cmd) do {				      \
+    hsa_status_t err = pfn_##cmd;				      \
+    if( err != HSA_STATUS_SUCCESS ) {				      \
+      const char *errStr;				      \
+      pfn_hsa_status_string(err, &errStr);	      \
+      WARN("ROCr failure '%s'", errStr);		      \
+      return ncclUnhandledCudaError;			      \
+    }							      \
+} while(false)
+
+#define CUCHECKGOTO(cmd, res, label) do {		      \
+    hsa_status_t err = pfn_##cmd;				      \
+    if( err != HSA_STATUS_SUCCESS ) {				      \
+      const char *errStr;				      \
+      pfn_hsa_status_string(err, &errStr);	      \
+      WARN("ROCr failure '%s'", errStr);		      \
+      res = ncclUnhandledCudaError;			      \
+      goto label;					      \
+    }							      \
+} while(false)
+
+// Report failure but clear error and continue
+#define CUCHECKIGNORE(cmd) do {						\
+    hsa_status_t err = pfn_##cmd;						\
+    if( err != HSA_STATUS_SUCCESS ) {						\
+      const char *errStr;						\
+      pfn_hsa_status_string(err, &errStr);			\
+      INFO(NCCL_ALL,"%s:%d ROCr failure '%s'", __FILE__, __LINE__, errStr);	\
+    }									\
+} while(false)
+
+#define CUCHECKTHREAD(cmd, args) do {					\
+    hsa_status_t err = pfn_##cmd;						\
+    if (err != HSA_STATUS_SUCCESS) {						\
+      INFO(NCCL_INIT,"%s:%d -> %d [Async thread]", __FILE__, __LINE__, err); \
+      args->ret = ncclUnhandledCudaError;				\
+      return args;							\
+    }									\
+} while(0)
+
+#define DECLARE_ROCM_PFN_EXTERN(symbol) extern PFN_##symbol pfn_##symbol
+
+DECLARE_ROCM_PFN_EXTERN(hsa_amd_portable_export_dmabuf); // DMA-BUF support
+
+/* ROCr Driver functions loaded with dlsym() */
+DECLARE_ROCM_PFN_EXTERN(hsa_init);
+DECLARE_ROCM_PFN_EXTERN(hsa_system_get_info);
+DECLARE_ROCM_PFN_EXTERN(hsa_status_string);
+
+ncclResult_t rocmLibraryInit(void);
+
+#endif

src/include/strongstream.h

@@ -0,0 +1,142 @@
+/*************************************************************************
+ * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef NCCL_STRONGSTREAM_H_
+#define NCCL_STRONGSTREAM_H_
+
+#include "nccl.h"
+#include "checks.h"
+
+#include <stdint.h>
+
+/* ncclCudaGraph: Wraps a cudaGraph_t so that we can support pre-graph CUDA runtimes
+ * easily.
+ */
+struct ncclCudaGraph {
+#if CUDART_VERSION >= 11030
+  cudaGraph_t graph;
+  uint64_t graphId;
+#endif
+};
+
+inline struct ncclCudaGraph ncclCudaGraphNull() {
+  struct ncclCudaGraph tmp;
+  #if CUDART_VERSION >= 11030
+    tmp.graph = nullptr;
+    tmp.graphId = ULLONG_MAX;
+  #endif
+  return tmp;
+}
+
+inline bool ncclCudaGraphValid(struct ncclCudaGraph graph) {
+  #if CUDART_VERSION >= 11030
+    return graph.graph != nullptr;
+  #else
+    return false;
+  #endif
+}
+
+inline bool ncclCudaGraphSame(struct ncclCudaGraph a, struct ncclCudaGraph b) {
+  #if CUDART_VERSION >= 11030
+    return a.graphId == b.graphId;
+  #else
+    return true;
+  #endif
+}
+
+ncclResult_t ncclCudaGetCapturingGraph(struct ncclCudaGraph* graph, hipStream_t stream);
+ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, hipHostFn_t fn, void* arg);
+
+
+/* ncclStrongStream: An abstraction over CUDA streams that do not lose their
+ * identity while being captured. Regular streams have the deficiency that the
+ * captured form of a stream in one graph launch has no relation to the
+ * uncaptured stream or to the captured form in other graph launches. This makes
+ * streams unfit for the use of serializing access to a persistent resource.
+ * Strong streams have been introduced to address this need.
+ *
+ * Constraints of using strong streams:
+ *
+ * - Operations that enqueue work to the strong stream need to be enclosed by
+ *   ncclStrongStream[Acquire/Release] pairs. Acquire/release act like fences,
+ *   the strong stream is not stateful so there is no harm in redundant acquire
+ *   or releases.
+ *
+ * - An {Acquire; ...; Release} sequence must not be concurrent with any
+ *   other operations against the strong stream including graph launches which
+ *   reference this stream.
+ *
+ * - All strong stream functions take a "graph" parameter which must reference
+ *   the currently capturing graph, or null if none.
+ */
+struct ncclStrongStream;
+
+ncclResult_t ncclStrongStreamConstruct(struct ncclStrongStream* ss);
+ncclResult_t ncclStrongStreamDestruct(struct ncclStrongStream* ss);
+
+// Has this strong stream ever been captured in a graph.
+bool ncclStrongStreamEverCaptured(struct ncclStrongStream* ss);
+
+// Acquire-fence the strong stream.
+ncclResult_t ncclStrongStreamAcquire(
+  struct ncclCudaGraph graph, struct ncclStrongStream* ss
+);
+
+// Acquire-fence the strong stream assuming no graph is capturing. This permits
+// the caller to enqueue directly to the `ss->stream` member using native CUDA
+// calls. Strong stream must be released via:
+//   ncclStrongStreamRelease(ncclCudaGraphNull(), graphRefs, ss);
+ncclResult_t ncclStrongStreamAcquireUncaptured(struct ncclStrongStream* ss);
+
+// Release-fence of the strong stream.
+ncclResult_t ncclStrongStreamRelease(struct ncclCudaGraph graph, struct ncclStrongStream* ss);
+
+// Add a host launch to the stream.
+ncclResult_t ncclStrongStreamLaunchHost(
+  struct ncclCudaGraph graph, struct ncclStrongStream* ss,
+  hipHostFn_t fn, void* arg
+);
+// Add a kernel launch to the stream.
+ncclResult_t ncclStrongStreamLaunchKernel(
+  struct ncclCudaGraph graph, struct ncclStrongStream* ss,
+  void* fn, dim3 grid, dim3 block, void** args, size_t sharedMemBytes
+);
+// Cause `a` to wait for the current state `b`. Both `a` and `b` must be acquired.
+ncclResult_t ncclStrongStreamWaitStream(
+  struct ncclCudaGraph graph, struct ncclStrongStream* a, struct ncclStrongStream* b
+);
+// `b` must be capturing within `graph`.
+ncclResult_t ncclStrongStreamWaitStream(
+  struct ncclCudaGraph graph, struct ncclStrongStream* a, hipStream_t b
+);
+// `a` must be capturing within `graph`.
+ncclResult_t ncclStrongStreamWaitStream(
+  struct ncclCudaGraph graph, hipStream_t a, struct ncclStrongStream* b
+);
+
+// Synchrnoization does not need the strong stream to be acquired.
+ncclResult_t ncclStrongStreamSynchronize(struct ncclStrongStream* ss);
+
+////////////////////////////////////////////////////////////////////////////////
+
+struct ncclStrongStream {
+  hipStream_t stream;
+  hipEvent_t event;
+  #if CUDART_VERSION >= 11030
+  cudaGraphNode_t node; // null if never captured, otherwise never null again
+  uint64_t graphId:63, eventIsLagging:1;
+  #endif
+};
+
+inline bool ncclStrongStreamEverCaptured(struct ncclStrongStream* ss) {
+  #if CUDART_VERSION >= 11030
+    return ss->node != nullptr;
+  #else
+    return false;
+  #endif
+}
+
+#endif

src/include/transport.h

@@ -21,7 +21,12 @@
 
 #include "proxy.h"
 
-extern struct ncclTransport ncclTransports[];
+extern struct ncclTransport p2pTransport;
+extern struct ncclTransport shmTransport;
+extern struct ncclTransport netTransport;
+extern struct ncclTransport collNetTransport;
+
+extern struct ncclTransport* ncclTransports[];
 
 // Forward declarations
 struct ncclRing;
@@ -66,7 +71,7 @@ struct ncclTransport {
   struct ncclTransportComm recv;
 };
 
-ncclResult_t ncclTransportP2pConnect(struct ncclComm* comm, struct ncclChannel* channel, int nrecv, int* peerRecv, int nsend, int* peerSend, int connIndex);
+ncclResult_t ncclTransportP2pConnect(struct ncclComm* comm, int channelId, int nrecv, int* peerRecv, int nsend, int* peerSend, int connIndex);
 ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType=NULL);
 
 enum { collNetRecv=0, collNetSend=1 };

src/include/utils.h

@@ -8,8 +8,12 @@
 #define NCCL_UTILS_H_
 
 #include "nccl.h"
+#include "alloc.h"
 #include "checks.h"
 #include <stdint.h>
+#include <time.h>
+#include <sched.h>
+#include <new>
 
 int ncclCudaCompCap();
 
@@ -38,81 +42,446 @@ static long log2i(long n) {
  return l;
 }
 
-// Recyclable list that avoids frequent malloc/free
+inline uint64_t clockNano() {
+  struct timespec ts;
+  clock_gettime(CLOCK_MONOTONIC, &ts);
+  return uint64_t(ts.tv_sec)*1000*1000*1000 + ts.tv_nsec;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+template<typename Int>
+inline void ncclAtomicRefCountIncrement(Int* refs) {
+  __atomic_fetch_add(refs, 1, __ATOMIC_RELAXED);
+}
+
+template<typename Int>
+inline Int ncclAtomicRefCountDecrement(Int* refs) {
+  return __atomic_sub_fetch(refs, 1, __ATOMIC_ACQ_REL);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/* ncclMemoryStack: Pools memory for fast LIFO ordered allocation. Note that
+ * granularity of LIFO is not per object, instead frames containing many objects
+ * are pushed and popped. Therefor deallocation is extremely cheap since its
+ * done at the frame granularity.
+ *
+ * The initial state of the stack is with one frame, the "nil" frame, which
+ * cannot be popped. Therefor objects allocated in the nil frame cannot be
+ * deallocated sooner than stack destruction.
+ */
+struct ncclMemoryStack;
+
+void ncclMemoryStackConstruct(struct ncclMemoryStack* me);
+void ncclMemoryStackDestruct(struct ncclMemoryStack* me);
+void ncclMemoryStackPush(struct ncclMemoryStack* me);
+void ncclMemoryStackPop(struct ncclMemoryStack* me);
 template<typename T>
-struct ncclListElem {
-  T data;
-  struct ncclListElem* next;
+T* ncclMemoryStackAlloc(struct ncclMemoryStack* me, size_t n=1);
+
+////////////////////////////////////////////////////////////////////////////////
+/* ncclMemoryPool: A free-list of same-sized allocations. It is an invalid for
+ * a pool instance to ever hold objects whose type have differing
+ * (sizeof(T), alignof(T)) pairs. The underlying memory is supplied by
+ * a backing `ncclMemoryStack` passed during Alloc(). If memory
+ * backing any currently held object is deallocated then it is an error to do
+ * anything other than reconstruct it, after which it is a valid empty pool.
+ */
+struct ncclMemoryPool;
+
+// Equivalent to zero-initialization
+void ncclMemoryPoolConstruct(struct ncclMemoryPool* me);
+template<typename T>
+T* ncclMemoryPoolAlloc(struct ncclMemoryPool* me, struct ncclMemoryStack* backing);
+template<typename T>
+void ncclMemoryPoolFree(struct ncclMemoryPool* me, T* obj);
+void ncclMemoryPoolTakeAll(struct ncclMemoryPool* me, struct ncclMemoryPool* from);
+
+////////////////////////////////////////////////////////////////////////////////
+/* ncclIntruQueue: A singly-linked list queue where the per-object next pointer
+ * field is given via the `next` template argument.
+ *
+ * Example:
+ *   struct Foo {
+ *     struct Foo *next1, *next2; // can be a member of two lists at once
+ *   };
+ *   ncclIntruQueue<Foo, &Foo::next1> list1;
+ *   ncclIntruQueue<Foo, &Foo::next2> list2;
+ */
+template<typename T, T *T::*next>
+struct ncclIntruQueue;
+
+template<typename T, T *T::*next>
+void ncclIntruQueueConstruct(ncclIntruQueue<T,next> *me);
+template<typename T, T *T::*next>
+bool ncclIntruQueueEmpty(ncclIntruQueue<T,next> *me);
+template<typename T, T *T::*next>
+T* ncclIntruQueueHead(ncclIntruQueue<T,next> *me);
+template<typename T, T *T::*next>
+void ncclIntruQueueEnqueue(ncclIntruQueue<T,next> *me, T *x);
+template<typename T, T *T::*next>
+T* ncclIntruQueueDequeue(ncclIntruQueue<T,next> *me);
+template<typename T, T *T::*next>
+T* ncclIntruQueueTryDequeue(ncclIntruQueue<T,next> *me);
+template<typename T, T *T::*next>
+void ncclIntruQueueFreeAll(ncclIntruQueue<T,next> *me, ncclMemoryPool *memPool);
+
+////////////////////////////////////////////////////////////////////////////////
+/* ncclThreadSignal: Couples a pthread mutex and cond together. The "mutex"
+ * and "cond" fields are part of the public interface.
+ */
+struct ncclThreadSignal {
+  pthread_mutex_t mutex;
+  pthread_cond_t cond;
 };
 
-template<typename T>
-class ncclRecyclableList {
- private:
-  struct ncclListElem<T>* head;
-  struct ncclListElem<T>* tail;
-  struct ncclListElem<T>* cursor;
-  int n;
+// returns {PTHREAD_MUTEX_INITIALIZER, PTHREAD_COND_INITIALIZER}
+constexpr ncclThreadSignal ncclThreadSignalStaticInitializer();
 
- public:
-  ncclRecyclableList() {
-    tail = cursor = head = NULL;
-    n = 0;
-  }
+void ncclThreadSignalConstruct(struct ncclThreadSignal* me);
+void ncclThreadSignalDestruct(struct ncclThreadSignal* me);
 
-  int count() const { return n; }
+// A convenience instance per-thread.
+extern __thread struct ncclThreadSignal ncclThreadSignalLocalInstance;
 
-  // Get a new element from the list and return pointer
-  ncclResult_t getNewElem(T** dataOut) {
-    if (tail != NULL) {
-      *dataOut = &tail->data;
-      memset(*dataOut, 0, sizeof(T));
-    } else {
-      NCCLCHECK(ncclCalloc(&tail, 1));
-      *dataOut = &tail->data;
-      cursor = head = tail;
-    }
-    if (tail->next == NULL) {
-      NCCLCHECK(ncclCalloc(&tail->next, 1));
-    }
-    tail = tail->next;
-    n += 1;
-    return ncclSuccess;
-  }
+////////////////////////////////////////////////////////////////////////////////
 
-  T* begin() {
-    if (head == NULL || head == tail) return NULL;
-    cursor = head->next;
-    return &head->data;
-  }
+template<typename T, T *T::*next>
+struct ncclIntruQueueMpsc;
 
-  // Get next element from the list during an iteration
-  T* getNext() {
-    // tail always points to the next element to be enqueued
-    // hence does not contain valid data
-    if (cursor == NULL || cursor == tail) return NULL;
-    T* rv = &cursor->data;
-    cursor = cursor->next;
-    return rv;
-  }
+template<typename T, T *T::*next>
+void ncclIntruQueueMpscConstruct(struct ncclIntruQueueMpsc<T,next>* me);
+template<typename T, T *T::*next>
+bool ncclIntruQueueMpscEmpty(struct ncclIntruQueueMpsc<T,next>* me);
+// Enqueue element. Returns true if queue is not abandoned. Even if queue is
+// abandoned the element enqueued, so the caller needs to make arrangements for
+// the queue to be tended.
+template<typename T, T *T::*next>
+bool ncclIntruQueueMpscEnqueue(struct ncclIntruQueueMpsc<T,next>* me, T* x);
+// Dequeue all elements at a glance. If there aren't any and `waitSome` is
+// true then this call will wait until it can return a non empty list.
+template<typename T, T *T::*next>
+T* ncclIntruQueueMpscDequeueAll(struct ncclIntruQueueMpsc<T,next>* me, bool waitSome);
+// Dequeue all elements and set queue to abandoned state.
+template<typename T, T *T::*next>
+T* ncclIntruQueueMpscAbandon(struct ncclIntruQueueMpsc<T,next>* me);
 
-  T* peakNext() {
-    if (cursor == NULL || cursor == tail) return NULL;
-    return &cursor->data;
-  }
+////////////////////////////////////////////////////////////////////////////////
 
-  // Recycle the list without freeing the space
-  void recycle() {
-    tail = cursor = head;
-    n = 0;
-  }
+struct ncclMemoryStack {
+  struct Hunk {
+    struct Hunk* above; // reverse stack pointer
+    size_t size; // size of this allocation (including this header struct)
+  };
+  struct Unhunk { // proxy header for objects allocated out-of-hunk
+    struct Unhunk* next;
+    void* obj;
+  };
+  struct Frame {
+    struct Hunk* hunk; // top of non-empty hunks
+    uintptr_t bumper, end; // points into top hunk
+    struct Unhunk* unhunks;
+    struct Frame* below;
+  };
 
-  ~ncclRecyclableList() {
-    while (head != NULL) {
-      struct ncclListElem<T>* temp = head;
-      head = head->next;
-      free(temp);
-    }
-  }
+  static void* allocateSpilled(struct ncclMemoryStack* me, size_t size, size_t align);
+  static void* allocate(struct ncclMemoryStack* me, size_t size, size_t align);
+
+  struct Hunk stub;
+  struct Frame topFrame;
 };
 
+inline void ncclMemoryStackConstruct(struct ncclMemoryStack* me) {
+  me->stub.above = nullptr;
+  me->stub.size = 0;
+  me->topFrame.hunk = &me->stub;
+  me->topFrame.bumper = 0;
+  me->topFrame.end = 0;
+  me->topFrame.unhunks = nullptr;
+  me->topFrame.below = nullptr;
+}
+
+inline void* ncclMemoryStack::allocate(struct ncclMemoryStack* me, size_t size, size_t align) {
+  uintptr_t o = (me->topFrame.bumper + align-1) & -uintptr_t(align);
+  void* obj;
+  if (__builtin_expect(o + size <= me->topFrame.end, true)) {
+    me->topFrame.bumper = o + size;
+    obj = reinterpret_cast<void*>(o);
+  } else {
+    obj = allocateSpilled(me, size, align);
+  }
+  return obj;
+}
+
+template<typename T>
+inline T* ncclMemoryStackAlloc(struct ncclMemoryStack* me, size_t n) {
+  void *obj = ncclMemoryStack::allocate(me, n*sizeof(T), alignof(T));
+  memset(obj, 0, n*sizeof(T));
+  return (T*)obj;
+}
+
+inline void ncclMemoryStackPush(struct ncclMemoryStack* me) {
+  using Frame = ncclMemoryStack::Frame;
+  Frame tmp = me->topFrame;
+  Frame* snapshot = (Frame*)ncclMemoryStack::allocate(me, sizeof(Frame), alignof(Frame));
+  *snapshot = tmp; // C++ struct assignment
+  me->topFrame.unhunks = nullptr;
+  me->topFrame.below = snapshot;
+}
+
+inline void ncclMemoryStackPop(struct ncclMemoryStack* me) {
+  ncclMemoryStack::Unhunk* un = me->topFrame.unhunks;
+  while (un != nullptr) {
+    free(un->obj);
+    un = un->next;
+  }
+  me->topFrame = *me->topFrame.below; // C++ struct assignment
+}
+
+
+////////////////////////////////////////////////////////////////////////////////
+
+struct ncclMemoryPool {
+  struct Cell {
+    Cell *next;
+  };
+  template<int Size, int Align>
+  union CellSized {
+    Cell cell;
+    alignas(Align) char space[Size];
+  };
+  struct Cell* head;
+  struct Cell* tail; // meaningful only when head != nullptr
+};
+
+inline void ncclMemoryPoolConstruct(struct ncclMemoryPool* me) {
+  me->head = nullptr;
+}
+
+template<typename T>
+inline T* ncclMemoryPoolAlloc(struct ncclMemoryPool* me, struct ncclMemoryStack* backing) {
+  using Cell = ncclMemoryPool::Cell;
+  using CellSized = ncclMemoryPool::CellSized<sizeof(T), alignof(T)>;
+  Cell* cell;
+  if (__builtin_expect(me->head != nullptr, true)) {
+    cell = me->head;
+    me->head = cell->next;
+  } else {
+    // Use the internal allocate() since it doesn't memset to 0 yet.
+    cell = (Cell*)ncclMemoryStack::allocate(backing, sizeof(CellSized), alignof(CellSized));
+  }
+  memset(cell, 0, sizeof(T));
+  return reinterpret_cast<T*>(cell);
+}
+
+template<typename T>
+inline void ncclMemoryPoolFree(struct ncclMemoryPool* me, T* obj) {
+  using Cell = ncclMemoryPool::Cell;
+  Cell* cell = reinterpret_cast<Cell*>(obj);
+  cell->next = me->head;
+  if (me->head == nullptr) me->tail = cell;
+  me->head = cell;
+}
+
+inline void ncclMemoryPoolTakeAll(struct ncclMemoryPool* me, struct ncclMemoryPool* from) {
+  if (from->head != nullptr) {
+    from->tail->next = me->head;
+    if (me->head == nullptr) me->tail = from->tail;
+    me->head = from->head;
+    from->head = nullptr;
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, T *T::*next>
+struct ncclIntruQueue {
+  T *head, *tail;
+};
+
+template<typename T, T *T::*next>
+inline void ncclIntruQueueConstruct(ncclIntruQueue<T,next> *me) {
+  me->head = nullptr;
+  me->tail = nullptr;
+}
+
+template<typename T, T *T::*next>
+inline bool ncclIntruQueueEmpty(ncclIntruQueue<T,next> *me) {
+  return me->head == nullptr;
+}
+
+template<typename T, T *T::*next>
+inline T* ncclIntruQueueHead(ncclIntruQueue<T,next> *me) {
+  return me->head;
+}
+
+template<typename T, T *T::*next>
+inline T* ncclIntruQueueTail(ncclIntruQueue<T,next> *me) {
+  return me->tail;
+}
+
+template<typename T, T *T::*next>
+inline void ncclIntruQueueEnqueue(ncclIntruQueue<T,next> *me, T *x) {
+  x->*next = nullptr;
+  (me->head ? me->tail->*next : me->head) = x;
+  me->tail = x;
+}
+
+template<typename T, T *T::*next>
+inline T* ncclIntruQueueDequeue(ncclIntruQueue<T,next> *me) {
+  T *ans = me->head;
+  me->head = ans->*next;
+  if (me->head == nullptr) me->tail = nullptr;
+  return ans;
+}
+
+template<typename T, T *T::*next>
+inline T* ncclIntruQueueTryDequeue(ncclIntruQueue<T,next> *me) {
+  T *ans = me->head;
+  if (ans != nullptr) {
+    me->head = ans->*next;
+    if (me->head == nullptr) me->tail = nullptr;
+  }
+  return ans;
+}
+
+template<typename T, T *T::*next>
+void ncclIntruQueueFreeAll(ncclIntruQueue<T,next> *me, ncclMemoryPool *pool) {
+  T *head = me->head;
+  me->head = nullptr;
+  me->tail = nullptr;
+  while (head != nullptr) {
+    T *tmp = head->*next;
+    ncclMemoryPoolFree(pool, tmp);
+    head = tmp;
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+constexpr ncclThreadSignal ncclThreadSignalStaticInitializer() {
+  return {PTHREAD_MUTEX_INITIALIZER, PTHREAD_COND_INITIALIZER};
+}
+
+inline void ncclThreadSignalConstruct(struct ncclThreadSignal* me) {
+  pthread_mutex_init(&me->mutex, nullptr);
+  pthread_cond_init(&me->cond, nullptr);
+}
+
+inline void ncclThreadSignalDestruct(struct ncclThreadSignal* me) {
+  pthread_mutex_destroy(&me->mutex);
+  pthread_cond_destroy(&me->cond);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, T *T::*next>
+struct ncclIntruQueueMpsc {
+  T* head;
+  uintptr_t tail;
+  struct ncclThreadSignal* waiting;
+};
+
+template<typename T, T *T::*next>
+void ncclIntruQueueMpscConstruct(struct ncclIntruQueueMpsc<T,next>* me) {
+  me->head = nullptr;
+  me->tail = 0x0;
+  me->waiting = nullptr;
+}
+
+template<typename T, T *T::*next>
+bool ncclIntruQueueMpscEmpty(struct ncclIntruQueueMpsc<T,next>* me) {
+  return __atomic_load_n(&me->tail, __ATOMIC_RELAXED) <= 0x2;
+}
+
+template<typename T, T *T::*next>
+bool ncclIntruQueueMpscEnqueue(ncclIntruQueueMpsc<T,next>* me, T* x) {
+  __atomic_store_n(&(x->*next), nullptr, __ATOMIC_RELAXED);
+  uintptr_t utail = __atomic_exchange_n(&me->tail, reinterpret_cast<uintptr_t>(x), __ATOMIC_ACQ_REL);
+  T* prev = reinterpret_cast<T*>(utail);
+  T** prevNext = utail <= 0x2 ? &me->head : &(prev->*next);
+  __atomic_store_n(prevNext, x, __ATOMIC_RELAXED);
+  if (utail == 0x1) { // waiting
+    __atomic_thread_fence(__ATOMIC_ACQUIRE); // to see me->waiting
+    // This lock/unlock is essential to ensure we don't race ahead of the consumer
+    // and signal the cond before they begin waiting on it.
+    struct ncclThreadSignal* waiting = me->waiting;
+    pthread_mutex_lock(&waiting->mutex);
+    pthread_mutex_unlock(&waiting->mutex);
+    pthread_cond_broadcast(&waiting->cond);
+  }
+  return utail != 0x2; // not abandoned
+}
+
+template<typename T, T *T::*next>
+T* ncclIntruQueueMpscDequeueAll(ncclIntruQueueMpsc<T,next>* me, bool waitSome) {
+  T* head = __atomic_load_n(&me->head, __ATOMIC_RELAXED);
+  if (head == nullptr) {
+    if (!waitSome) return nullptr;
+    uint64_t t0 = clockNano();
+    bool sleeping = false;
+    do {
+      if (clockNano()-t0 >= 10*1000) { // spin for first 10us
+        struct ncclThreadSignal* waitSignal = &ncclThreadSignalLocalInstance;
+        pthread_mutex_lock(&waitSignal->mutex);
+        uintptr_t expected = sleeping ? 0x1 : 0x0;
+        uintptr_t desired = 0x1;
+        me->waiting = waitSignal; // release done by successful compare exchange
+        if (__atomic_compare_exchange_n(&me->tail, &expected, desired, /*weak=*/true, __ATOMIC_RELEASE, __ATOMIC_RELAXED)) {
+          sleeping = true;
+          pthread_cond_wait(&waitSignal->cond, &waitSignal->mutex);
+        }
+        pthread_mutex_unlock(&waitSignal->mutex);
+      }
+      head = __atomic_load_n(&me->head, __ATOMIC_RELAXED);
+    } while (head == nullptr);
+  }
+
+  __atomic_store_n(&me->head, nullptr, __ATOMIC_RELAXED);
+  uintptr_t utail = __atomic_exchange_n(&me->tail, 0x0, __ATOMIC_ACQ_REL);
+  T* tail = utail <= 0x2 ? nullptr : reinterpret_cast<T*>(utail);
+  T *x = head;
+  while (x != tail) {
+    T *x1;
+    int spins = 0;
+    while (true) {
+      x1 = __atomic_load_n(&(x->*next), __ATOMIC_RELAXED);
+      if (x1 != nullptr) break;
+      if (++spins == 1024) { spins = 1024-1; sched_yield(); }
+    }
+    x = x1;
+  }
+  return head;
+}
+
+template<typename T, T *T::*next>
+T* ncclIntruQueueMpscAbandon(ncclIntruQueueMpsc<T,next>* me) {
+  uintptr_t expected = 0x0;
+  if (__atomic_compare_exchange_n(&me->tail, &expected, /*desired=*/0x2, /*weak=*/true, __ATOMIC_RELAXED, __ATOMIC_RELAXED)) {
+    return nullptr;
+  } else {
+    int spins = 0;
+    T* head;
+    while (true) {
+      head = __atomic_load_n(&me->head, __ATOMIC_RELAXED);
+      if (head != nullptr) break;
+      if (++spins == 1024) { spins = 1024-1; sched_yield(); }
+    }
+    __atomic_store_n(&me->head, nullptr, __ATOMIC_RELAXED);
+    uintptr_t utail = __atomic_exchange_n(&me->tail, 0x2, __ATOMIC_ACQ_REL);
+    T* tail = utail <= 0x2 ? nullptr : reinterpret_cast<T*>(utail);
+    T *x = head;
+    while (x != tail) {
+      T *x1;
+      spins = 0;
+      while (true) {
+        x1 = __atomic_load_n(&(x->*next), __ATOMIC_RELAXED);
+        if (x1 != nullptr) break;
+        if (++spins == 1024) { spins = 1024-1; sched_yield(); }
+      }
+      x = x1;
+    }
+    return head;
+  }
+}
 #endif