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

[ROCm/rccl commit: 3a919c1f49]

projects/rccl/src/include/alloc.h

@@ -12,6 +12,9 @@
 #include "checks.h"
 #include "align.h"
 #include <sys/mman.h>
+#include <unistd.h>
+#include <stdlib.h>
+#include <string.h>
 
 template <typename T>
 static ncclResult_t ncclCudaHostCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
@@ -28,7 +31,7 @@ static inline ncclResult_t ncclCudaHostFree(void* ptr) {
 }
 
 template <typename T>
-static ncclResult_t ncclCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
+static ncclResult_t ncclCalloc(T** ptr, size_t nelem) {
   void* p = malloc(nelem*sizeof(T));
   if (p == NULL) {
     WARN("Failed to malloc %ld bytes", nelem*sizeof(T));
@@ -36,10 +39,8 @@ static ncclResult_t ncclCallocDebug(T** ptr, size_t nelem, const char *filefunc,
   }
   memset(p, 0, nelem*sizeof(T));
   *ptr = (T*)p;
-  INFO(NCCL_ALLOC, "%s:%d Mem Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
   return ncclSuccess;
 }
-#define ncclCalloc(...) ncclCallocDebug(__VA_ARGS__, __FILE__, __LINE__)
 
 struct __attribute__ ((aligned(64))) allocationTracker {
   union {

projects/rccl/src/include/bootstrap.h

@@ -17,7 +17,8 @@ ncclResult_t bootstrapInit(ncclUniqueId* id, int rank, int nranks, void** commSt
 ncclResult_t bootstrapAllGather(void* commState, void* allData, int size);
 ncclResult_t bootstrapSend(void* commState, int peer, int tag, void* data, int size);
 ncclResult_t bootstrapRecv(void* commState, int peer, int tag, void* data, int size);
-ncclResult_t bootstrapBarrier(void* commState, int *ranks, int tag, int rank, int nranks);
+ncclResult_t bootstrapBarrier(void* commState, int *ranks, int rank, int nranks, int tag);
+ncclResult_t bootstrapIntraNodeAllGather(void* commState, int *ranks, int rank, int nranks, void* allData, int size);
 ncclResult_t bootstrapRemAlloc(size_t size, int rank, void* commState, int* id, hipIpcMemHandle_t* ipc, void** ptr);
 ncclResult_t bootstrapRemFree(int id, int rank, void* commState);
 ncclResult_t bootstrapClose(void* commState);

projects/rccl/src/include/collectives.h

@@ -8,78 +8,104 @@
 #ifndef NCCL_COLLECTIVES_H_
 #define NCCL_COLLECTIVES_H_
 
-#define FUNC_INDEX_P2P (NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumOps)
-#define FUNC_INDEX(func, redop, ncclType, al, pr) ((((((func)*ncclNumOps + (redop))*ncclNumTypes) + (ncclType))*NCCL_NUM_ALGORITHMS+(al))*NCCL_NUM_PROTOCOLS+(pr))
+enum ncclDevRedOp_t {
+  ncclDevSum, ncclDevProd, ncclDevMax, ncclDevMin,
+  ncclDevPreMulSum, ncclDevSumPostDiv,
+  ncclNumDevRedOps
+};
+struct ncclDevRedOpFull {
+  ncclDevRedOp_t op;
+  bool scalarArgIsPtr;
+  uint64_t scalarArg;
+};
 
-#define NCCL_FUNC_NAME(func, algo, proto, redop, type) \
-  ncclFunction_##func##_##algo##_##proto##_##redop##_##type
+#define FUNC_INDEX_P2P (ncclNumTypes+NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumDevRedOps)
+#define FUNC_INDEX(func, devredop, ncclType, al, pr) ((((((func)*ncclNumDevRedOps + (devredop))*ncclNumTypes) + (ncclType))*NCCL_NUM_ALGORITHMS+(al))*NCCL_NUM_PROTOCOLS+(pr))
 
-#define NCCL_KERN_NAME(func, algo, proto, redop, type) \
-  ncclKernel_##func##_##algo##_##proto##_##redop##_##type
+#define NCCL_FUNC_NAME(func, algo, proto, devredop, type) \
+  ncclFunction_##func##_##algo##_##proto##_##devredop##_##type
+
+#define NCCL_ONERANK_REDUCE_NAME(devredop, type) \
+  ncclFunction_OneRankReduce_##devredop##_##type
+
+#define NCCL_KERN_NAME(func, algo, proto, devredop, type) \
+  ncclKernel_##func##_##algo##_##proto##_##devredop##_##type
 
 #define NCCL_IMPL_NAME(func, algo, proto) \
   nccl##func##algo##proto
 
 /* Declare all collective operations */
-#define DECL5(func, algo, proto, redop, type) \
-  extern __device__ __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, redop, type)(struct ncclWorkElem* args); \
-  extern __global__ void NCCL_KERN_NAME(func, algo, proto, redop, type)(struct ncclWorkElem first); \
+#define DECL5(func, algo, proto, devredop, type) \
+  extern __device__ __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, devredop, type)(struct ncclWorkElem* args); \
+  extern __global__ void NCCL_KERN_NAME(func, algo, proto, devredop, type)(ncclWorkElem c); \
 
-  //extern __device__ void NCCL_FUNC_NAME(func, algo, proto, redop, type)(); \
-  //extern __global__ void NCCL_KERN_NAME(func, algo, proto, redop, type)(ncclWorkElem c); \
+#define CONCAT(a,b) a##b
+#define MACRO_IF(cond, t, f) CONCAT(MACRO_IF_, cond)(t, f)
+#define MACRO_IF_0(t, f) f
+#define MACRO_IF_1(t, f) t
 
-#define DECL4(func, algo, redop, type) \
-  DECL5(func, algo, SIMPLE, redop, type) \
-  DECL5(func, algo, LL,     redop, type) \
-  DECL5(func, algo, LL128,  redop, type)
+#define DECL4(func, algo, devredop, type, undef) \
+  MACRO_IF(undef, /*undefined*/, DECL5(func, algo, SIMPLE, devredop, type)) \
+  MACRO_IF(undef, /*undefined*/, DECL5(func, algo, LL,     devredop, type)) \
+  MACRO_IF(undef, /*undefined*/, DECL5(func, algo, LL128,  devredop, type))
 
-#define DECL3(func, redop, type) \
-  DECL4(func, RING,    redop, type) \
-  DECL4(func, TREE,    redop, type) \
-  DECL4(func, COLLNET, redop, type)
+#define DECL3(func, devredop, type, undef) \
+  DECL4(func, RING,    devredop, type, undef) \
+  DECL4(func, TREE,    devredop, type, undef) \
+  DECL4(func, COLLNET, devredop, type, undef)
 
-#if defined(__CUDA_BF16_TYPES_EXIST__)
-#define DECL2(func, redop) \
-  DECL3(func, redop, int8_t) \
-  DECL3(func, redop, uint8_t) \
-  DECL3(func, redop, int32_t) \
-  DECL3(func, redop, uint32_t) \
-  DECL3(func, redop, int64_t) \
-  DECL3(func, redop, uint64_t) \
-  DECL3(func, redop, half) \
-  DECL3(func, redop, float) \
-  DECL3(func, redop, double) \
-  DECL3(func, redop, __nv_bfloat16)
+#if defined(RCCL_BFLOAT16)
+#define DECL2(func, devredop, undefForFloat) \
+  DECL3(func, devredop, int8_t, /*undef=*/0) \
+  DECL3(func, devredop, uint8_t, /*undef=*/0) \
+  DECL3(func, devredop, int32_t, /*undef=*/0) \
+  DECL3(func, devredop, uint32_t, /*undef=*/0) \
+  DECL3(func, devredop, int64_t, /*undef=*/0) \
+  DECL3(func, devredop, uint64_t, /*undef=*/0) \
+  DECL3(func, devredop, half, /*undef=*/undefForFloat) \
+  DECL3(func, devredop, float, /*undef=*/undefForFloat) \
+  DECL3(func, devredop, double, /*undef=*/undefForFloat) \
+  DECL3(func, devredop, rccl_bfloat16, /*undef=*/undefForFloat)
 #else
-#define DECL2(func, redop) \
-  DECL3(func, redop, int8_t) \
-  DECL3(func, redop, uint8_t) \
-  DECL3(func, redop, int32_t) \
-  DECL3(func, redop, uint32_t) \
-  DECL3(func, redop, int64_t) \
-  DECL3(func, redop, uint64_t) \
-  DECL3(func, redop, half) \
-  DECL3(func, redop, float) \
-  DECL3(func, redop, double) \
-  DECL3(func, redop, rccl_bfloat16)
+#define DECL2(func, devredop, undefForFloat) \
+  DECL3(func, devredop, int8_t, /*undef=*/0) \
+  DECL3(func, devredop, uint8_t, /*undef=*/0) \
+  DECL3(func, devredop, int32_t, /*undef=*/0) \
+  DECL3(func, devredop, uint32_t, /*undef=*/0) \
+  DECL3(func, devredop, int64_t, /*undef=*/0) \
+  DECL3(func, devredop, uint64_t, /*undef=*/0) \
+  DECL3(func, devredop, half, /*undef=*/undefForFloat) \
+  DECL3(func, devredop, float, /*undef=*/undefForFloat) \
+  DECL3(func, devredop, double, /*undef=*/undefForFloat)
 #endif
 
 #define DECL(func) \
-  DECL2(func, Sum) \
-  DECL2(func, Prod) \
-  DECL2(func, Min) \
-  DECL2(func, Max) \
-  DECL2(func, Avg) \
+  DECL2(func, Sum, /*undefForFloat=*/0) \
+  DECL2(func, Prod, /*undefForFloat=*/0) \
+  DECL2(func, Min, /*undefForFloat=*/0) \
+  DECL2(func, Max, /*undefForFloat=*/0) \
+  DECL2(func, PreMulSum, /*undefForFloat=*/0) \
+  DECL2(func, SumPostDiv, /*undefForFloat=*/1)
 
-#define DECL_ALL \
-  DECL2(Broadcast, Sum) \
-  DECL(Reduce) \
-  DECL2(AllGather, Sum) \
-  DECL(ReduceScatter) \
-  DECL(AllReduce) \
-  DECL5(SendRecv, RING, SIMPLE, Sum, int8_t) \
+DECL2(Broadcast, Sum, /*undefForFloat=*/0)
+DECL(Reduce)
+DECL2(AllGather, Sum, /*undefForFloat=*/0)
+DECL(ReduceScatter)
+DECL(AllReduce)
+DECL5(SendRecv, RING, SIMPLE, Sum, int8_t)
 
-DECL_ALL
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, int8_t)(struct ncclWorkElem* args);
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, uint8_t)(struct ncclWorkElem* args);
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, int32_t)(struct ncclWorkElem* args);
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, uint32_t)(struct ncclWorkElem* args);
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, int64_t)(struct ncclWorkElem* args);
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, uint64_t)(struct ncclWorkElem* args);
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, half)(struct ncclWorkElem* args);
+#if defined(RCCL_BFLOAT16)
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, rccl_bfloat16)(struct ncclWorkElem* args);
+#endif
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, float)(struct ncclWorkElem* args);
+extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, double)(struct ncclWorkElem* args);
 
 // CHUNKSIZE must be a multiple of SLICESIZE
 //#define ALLREDUCE_SLICESTEPS (NCCL_STEPS/4)
@@ -99,5 +125,6 @@ DECL_ALL
 #define REDUCE_SLICESTEPS 1
 #define REDUCE_CHUNKSTEPS 1
 #define SENDRECV_SLICEFACTOR 1
+#define NCCL_MAX_SLICE_PER_CHUNK 2  // max value for CHUNKSTEPS/SLICESTEPS, must accord with above
 
 #endif

projects/rccl/src/include/comm.h

@@ -19,6 +19,8 @@
   typedef void *cudaGraphNode_t;
   #define HIPRT_CB
 #else
+#include "collectives.h"
+
 #if CUDART_VERSION < 9000
 struct cudaLaunchParams {
   void *func;
@@ -40,13 +42,16 @@ struct cudaLaunchParams {
 #define NCCL_LL128_THREAD_THRESHOLD 8
 #define NCCL_SIMPLE_THREAD_THRESHOLD 64
 
+#define NCCL_MAX_INTRA_RANKS 32
+
 struct ncclSendMem {
   union {
     struct {
       uint64_t head;
       char pad1[CACHE_LINE_SIZE-sizeof(uint64_t)];
       void* ptrExchange;
-      char pad2[CACHE_LINE_SIZE-sizeof(void*)];
+      uint64_t redOpArgExchange[2];
+      char pad2[CACHE_LINE_SIZE-sizeof(void*)-2*sizeof(uint64_t)];
     };
     char pad3[MEM_ALIGN];
   };
@@ -66,6 +71,28 @@ struct ncclRecvMem {
   char buff[1]; // Actually larger than that
 };
 
+typedef hipError_t(*pfn_cuMemGetAddressRange_t)(void**, size_t*, void*);
+
+enum helperThreadState {ThreadStart, ThreadStop};
+
+#define NCCL_IPC_POOL_SIZE (2*NCCL_MAX_INTRA_RANKS*NCCL_MAX_OPS)
+
+struct ncclGraphHelperResources {
+  ncclComm* comm;
+  pthread_mutex_t threadLock;
+  pthread_cond_t  threadCond;
+  enum helperThreadState threadState;
+  void* ipcBases[NCCL_IPC_POOL_SIZE];
+  int ipcTail;
+  int ipcHead;
+};
+
+struct ncclUserRedOp {
+  int freeNext; // -1=allocated, otherwise index of next free entry in array
+  ncclDataType_t datatype;
+  ncclDevRedOpFull opFull;
+};
+
 struct ncclComm {
   struct ncclChannel channels[MAXCHANNELS];
 
@@ -86,7 +113,12 @@ struct ncclComm {
 
   int node;
   int nNodes;
+
+  // Intra-node rank info
+  int intraNodeGlobalRanks[NCCL_MAX_INTRA_RANKS];
   int localRanks;
+  int intraNodeRank;
+  int8_t* rankToIntraNodeRank;
 
   enum { GROUP, PARALLEL, GROUP_GRAPH } launchMode;
   hipStream_t userStream;
@@ -158,6 +190,7 @@ struct ncclComm {
 
   // Whether this communicator uses collNet
   int collNetSupport;
+  int intraHighestTransportType;
 
   // Store info of async operations
   struct ncclInfo* asyncOps;
@@ -181,9 +214,38 @@ struct ncclComm {
   // Store info for cudaGraph
   int usingCudaGraph; // Only use it during capture time, not launch time
   struct ncclQueueInfo* enqueueInfo;
+  int nQueueInfoCreated;
+  int nQueueInfoDestroyed;
   cudaGraphNode_t lastSetupNode;
   unsigned long long lastCudaGraphId;
   int driverVersion;
+  pfn_cuMemGetAddressRange_t pfnCuMemGetAddressRange;
+  pthread_t graphHelperThread;
+  struct ncclGraphHelperResources* graphHelperResources;
+  int disableGraphHelper;
+  int graphRegister;
+
+  // user-created reduction ops
+  int userRedOpCapacity, userRedOpFreeHead;
+  ncclUserRedOp *userRedOps;
 };
 
+// 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
+// communicatrs. This function is its own inverse.
+static inline ncclRedOp_t ncclUserRedOpMangle(ncclComm *comm, ncclRedOp_t op) {
+  // Preserve the built-in values.
+  if(int(op) < int(ncclNumOps))
+    return op;
+  uint64_t h = reinterpret_cast<uint64_t>(comm);
+  h ^= h >> 32;
+  h *= 0x9e3779b97f4a7c13u; // Knuth's 64-bit magical hash constant
+  h >>= 32; // h is now an excellent 32-bit hash of the comm pointer
+  h &= int(ncclMaxRedOp); // ncclMaxRedOp is a power of 2 minus 1
+  int op1 = int(h) ^ int(op);
+  // Since builtin values are preserved, we also have to preserve their preimage.
+  return op1 < int(ncclNumOps) ? op : ncclRedOp_t(op1);
+}
+
 #endif

projects/rccl/src/include/core.h

@@ -1,10 +1,6 @@
 /*************************************************************************
-<<<<<<< HEAD
- * Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
-=======
  * Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
->>>>>>> nccl/master
+ * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/

projects/rccl/src/include/debug.h

@@ -7,17 +7,14 @@
 #ifndef NCCL_DEBUG_H_
 #define NCCL_DEBUG_H_
 
-#include "core.h"
-
+#include "nccl_net.h"
 #include <stdio.h>
 #include <chrono>
 
 #include <sys/syscall.h>
 #include <limits.h>
 #include <string.h>
-#include "nccl_net.h"
-
-#define gettid() (pid_t) syscall(SYS_gettid)
+#include <pthread.h>
 
 extern int ncclDebugLevel;
 extern uint64_t ncclDebugMask;

projects/rccl/src/include/devcomm.h

@@ -96,8 +96,11 @@ static_assert(NCCL_LL_CLEAN_MASK % NCCL_STEPS == 0, "Invalid NCCL_LL_CLEAN_MASK
 #define NCCL_LL128_SHMEM_ELEMS_PER_THREAD 2
 #define NCCL_LL128_SHMEM_SIZE (NCCL_LL128_SHMEM_ELEMS_PER_THREAD*NCCL_LL128_MAX_NTHREADS)
 
-#define NCCL_DIRECT_GPU 0x01
-#define NCCL_DIRECT_NIC 0x10
+#define NCCL_DIRECT_WRITE 0x01
+#define NCCL_DIRECT_READ  0x02
+#define NCCL_DIRECT_NIC   0x04
+#define NCCL_IPC_WRITE    0x08
+#define NCCL_IPC_READ     0x10
 
 struct ncclConnInfo {
   // Regular comm mechanism
@@ -108,6 +111,7 @@ struct ncclConnInfo {
   int direct;         // Direct communication
   int shared;         // Buffers are shared
   void **ptrExchange; // Pointer exchange for direct communication
+  uint64_t* redOpArgExchange; // PreOp scaler exchange for direct pull case
 
   int *sizesFifo;     // Sizes fifo from GPU to proxy
   void* *ptrsFifo;      // Buffer fifo from proxy to GPU
@@ -175,7 +179,7 @@ struct ncclPeer {
 struct ncclDevComm;
 
 #pragma pack(push)  /* push current alignment to stack */
-#pragma pack(8)     /* set alignment to 4 bytes boundary */
+#pragma pack(8)     /* set alignment to 8 bytes boundary */
 #define NCCL_MAX_WORK_ELEMENTS 1
 #define NCCL_MAX_GROUPS (NCCL_MAX_NTHREADS/WARP_SIZE)
 
@@ -187,8 +191,9 @@ struct ncclWorkElem {
   struct ncclDevComm* comm;
   uint16_t nThreads;
   uint16_t funcIndex;
-  uint16_t index;
-  uint16_t active;
+  uint8_t regUsed;
+  uint8_t direct;
+  uint8_t active, redOpArgIsPtr;
 
   const void * sendbuff;
   void * recvbuff;
@@ -198,10 +203,12 @@ struct ncclWorkElem {
     struct {
       size_t count;
       size_t lastChunkSize;
-      uint32_t root;
+      uint64_t redOpArg;
+      uint16_t root;
       uint8_t bid;
       uint8_t nChannels;
-      uint8_t connIndex;
+      uint16_t connIndex;
+      uint16_t opCount;
     } coll;
     struct {
       size_t sendCount;
@@ -217,18 +224,16 @@ struct ncclWorkElem {
         };
         uint16_t padding;
       };
-    } p2p;
-    struct {
-      uint16_t padding[15];
       uint16_t opCount;
-    } op;
+    } p2p;
     // [RCCL] Clique-based arguments
     //        NOTE: Follows same field structure as coll
     //              because nChannels is accessed from "coll" struct.
     struct {
       size_t count;
       cliqueDevicePtrs_t* ptrs;
-      uint32_t unused;
+      uint64_t unused_1;
+      uint16_t unused_2;
       uint8_t bid;
       uint8_t nChannels;
     } clique;
@@ -236,11 +241,24 @@ struct ncclWorkElem {
     uint64_t align[4];
   };
 };
-struct ncclWork {
-  struct ncclWorkElem elems[NCCL_MAX_WORK_ELEMENTS];
-};
 static_assert(sizeof(struct ncclWorkElem) == (0x10*sizeof(int)), "ncclWorkElem must have a pow2 size");
 
+struct ncclWorkRegElem {
+  struct ncclWorkElem elem;
+  void* dnInputs[NCCL_MAX_DIRECT_ARITY+1];
+  void* dnOutputs[NCCL_MAX_DIRECT_ARITY+1];
+  void* upOutputs[NCCL_MAX_DIRECT_ARITY+1];
+};
+#define NCCL_REG_ELEM_FACTOR 4
+static_assert(sizeof(struct ncclWorkRegElem) == (NCCL_REG_ELEM_FACTOR*sizeof(struct ncclWorkElem)), "ncclWorkRegElem size must be pow2 times ncclWorkElem size");
+
+struct ncclWork {
+  union {
+    struct ncclWorkElem elems[NCCL_MAX_WORK_ELEMENTS];
+    struct ncclWorkRegElem regElems[NCCL_MAX_WORK_ELEMENTS/NCCL_REG_ELEM_FACTOR];
+  };
+};
+
 struct ncclChannel {
   union {
     struct {
@@ -330,6 +348,7 @@ struct ncclProf {
 typedef enum {
   ncclCollTraceNotReady,
   ncclCollTraceKernelLaunchType,
+  ncclCollTraceKernelEndType,
   ncclCollTraceCollEndType,
   ncclCollTraceAbortType,
   ncclCollTraceDataType

projects/rccl/src/include/enqueue.h

@@ -1,5 +1,6 @@
 /*************************************************************************
  * Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -30,10 +31,19 @@ void HIPRT_CB ncclEnqueueHostSetup(void* arg);
 ncclResult_t ncclGetCudaGraph(ncclComm_t comm, cudaGraph_t* graph);
 ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, cudaGraph_t graph);
 
+struct ncclBuffRegInfo {
+  void* sendbuffsBase[NCCL_MAX_INTRA_RANKS];
+  void* recvbuffsBase[NCCL_MAX_INTRA_RANKS];
+  void* sendbuffs[NCCL_MAX_INTRA_RANKS];
+  void* recvbuffs[NCCL_MAX_INTRA_RANKS];
+  int nBuffs;
+};
+
 // Enqueue information (for kernel and proxy) for each operation
 struct ncclQueueElem {
   struct ncclWorkElem work;
   struct ncclProxyArgs proxyArgs;
+  struct ncclBuffRegInfo buffRegInfo;
 };
 
 typedef ncclRecyclableList<struct ncclQueueElem> ncclQueueElemList;
@@ -43,6 +53,7 @@ struct ncclQueueInfo {
   ncclComm_t comm;
   int maxChannels;    // Dynamic version of gridDim
   ncclResult_t ret;   // Return value of host setup call
+  int nRegBuffs;
   ncclQueueElemList* elemList;
 };
 
@@ -50,6 +61,7 @@ static ncclResult_t ncclCreateQueueInfo(struct ncclQueueInfo** eqInfo, ncclComm_
   NCCLCHECK(ncclCalloc(eqInfo, 1));
   (*eqInfo)->comm = comm;
   (*eqInfo)->elemList = new ncclQueueElemList();
+  (*eqInfo)->comm->nQueueInfoCreated++;
   return ncclSuccess;
 }
 
@@ -58,6 +70,7 @@ 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;
 }
@@ -67,7 +80,54 @@ static ncclResult_t ncclResetQueueInfo(struct ncclQueueInfo* eqInfo) {
 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->intraNodeRank) 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

projects/rccl/src/include/gdrwrap.h

@@ -11,6 +11,7 @@
 #include "nccl.h"
 #include <stdint.h> // for standard [u]intX_t types
 #include <stdio.h>
+#include <stdlib.h>
 
 // These can be used if the GDR library isn't thread safe
 #include <pthread.h>

projects/rccl/src/include/info.h

@@ -10,6 +10,7 @@
 
 #include "nccl.h"
 #include "devcomm.h"
+#include "collectives.h"
 
 typedef enum {
   ncclPatternRing,
@@ -39,6 +40,7 @@ struct ncclInfo {
   int chunkSteps;
   int sliceSteps;
   // Computed later
+  ncclDevRedOpFull opFull;
   int algorithm;
   int protocol;
   ncclPattern_t pattern;

projects/rccl/src/include/param.h

@@ -8,6 +8,7 @@
 #ifndef NCCL_PARAM_H_
 #define NCCL_PARAM_H_
 
+#include <stdio.h>
 #include <stdlib.h>
 #include <unistd.h>
 #include <sys/types.h>

projects/rccl/src/include/transport.h

@@ -55,7 +55,7 @@ struct ncclTransport {
 };
 
 ncclResult_t ncclTransportP2pConnect(struct ncclComm* comm, struct ncclChannel* channel, int nrecv, int* peerRecv, int nsend, int* peerSend, int connIndex);
-ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex);
+ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType=NULL);
 
 enum { collNetRecv=0, collNetSend=1 };
 int ncclTransportCollNetSetup(struct ncclComm* comm, struct ncclTopoGraph* collNetGraph, struct ncclChannel* channel, int masterRank, int masterPeer, int collNetGraphChannelId, int type);