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

src/include/alloc.h

@@ -19,6 +19,11 @@
 #include <string.h>
 #include "rccl_vars.h"
 
+#if CUDART_VERSION >= 11030
+#include <cuda.h>
+#include "cudawrap.h"
+#endif
+
 uint64_t clockNano(); // from utils.h with which we have a circular dependency
 
 template<typename T>
@@ -26,6 +31,81 @@ constexpr size_t ncclSizeOfT() { return sizeof(T); }
 template<>
 constexpr size_t ncclSizeOfT<void>() { return 1; }
 
+#if CUDART_VERSION >= 12020
+
+static inline ncclResult_t ncclCuMemHostAlloc(void** ptr, CUmemGenericAllocationHandle *handlep, size_t size) {
+  ncclResult_t result = ncclSuccess;
+  size_t granularity = 0;
+  CUdevice currentDev;
+  CUmemAllocationProp prop = {};
+  CUmemAccessDesc accessDesc = {};
+  CUmemGenericAllocationHandle handle;
+  int cudaDev;
+  int cpuNumaNodeId = -1;
+  CUmemAllocationHandleType type = ncclCuMemHandleType;
+
+  CUDACHECK(cudaGetDevice(&cudaDev));
+  CUCHECK(cuDeviceGet(&currentDev, cudaDev));
+  CUCHECK(cuDeviceGetAttribute(&cpuNumaNodeId, CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID, currentDev));
+  if (cpuNumaNodeId < 0) cpuNumaNodeId = 0;
+  prop.location.type = CU_MEM_LOCATION_TYPE_HOST_NUMA;
+  prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
+  prop.requestedHandleTypes = type; // So it can be exported
+  prop.location.id = cpuNumaNodeId;
+  CUCHECK(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
+  ALIGN_SIZE(size, granularity);
+  /* Allocate the physical memory on the device */
+  CUCHECK(cuMemCreate(&handle, size, &prop, 0));
+  /* Reserve a virtual address range */
+  CUCHECK(cuMemAddressReserve((CUdeviceptr*)ptr, size, granularity, 0, 0));
+  /* Map the virtual address range to the physical allocation */
+  CUCHECK(cuMemMap((CUdeviceptr)*ptr, size, 0, handle, 0));
+  /* Now allow RW access to the newly mapped memory for local GPU */
+  accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+  accessDesc.location.id = cudaDev;
+  accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+  CUCHECK(cuMemSetAccess((CUdeviceptr)*ptr, size, &accessDesc, 1));
+
+  /* Now allow RW access to the newly mapped memory from the CPU */
+  accessDesc.location.type = CU_MEM_LOCATION_TYPE_HOST_NUMA;
+  accessDesc.location.id = cpuNumaNodeId;
+  accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+  CUCHECK(cuMemSetAccess((CUdeviceptr)*ptr, size, &accessDesc, 1));
+
+  if (handlep) *handlep = handle;
+  INFO(NCCL_ALLOC, "CUMEM Host Alloc Size %zi pointer %p handle %llx numa %d dev %d granularity %ld", size, *ptr, handle, cpuNumaNodeId, cudaDev, granularity);
+  return result;
+}
+
+static inline ncclResult_t ncclCuMemHostFree(void* ptr) {
+  if (ptr == NULL) return ncclSuccess;
+  ncclResult_t result = ncclSuccess;
+  CUmemGenericAllocationHandle handle;
+  size_t size = 0;
+  CUCHECK(cuMemRetainAllocationHandle(&handle, ptr));
+  CUCHECK(cuMemRelease(handle));
+  CUCHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr));
+  TRACE(NCCL_ALLOC, "CUMEM Host Free Size %zi pointer %p handle 0x%llx", size, ptr, handle);
+  CUCHECK(cuMemUnmap((CUdeviceptr)ptr, size));
+  CUCHECK(cuMemRelease(handle));
+  CUCHECK(cuMemAddressFree((CUdeviceptr)ptr, size));
+  return result;
+}
+
+#else /* CUDART_VERSION >= 12020 */
+
+static inline ncclResult_t ncclCuMemHostAlloc(void** ptr, void* handlep, size_t size) {
+  WARN("CUMEM Host is not supported prior to CUDA 12.2");
+  return ncclInternalError;
+}
+
+static inline ncclResult_t ncclCuMemHostFree(void* ptr) {
+  WARN("CUMEM Host is not supported prior to CUDA 12.2");
+  return ncclInternalError;
+}
+
+#endif  /* CUDART_VERSION >= 12020 */
+
 template <typename T>
 ncclResult_t ncclCudaHostCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
   ncclResult_t result = ncclSuccess;
@@ -51,24 +131,25 @@ finish:
   INFO(NCCL_ALLOC, "%s:%d Cuda Host Alloc Size %ld pointer %p", filefunc, line, nelem*ncclSizeOfT<T>(), *ptr);
   return result;
 }
-#define ncclCudaHostCalloc(...) ncclCudaHostCallocDebug(__VA_ARGS__, __FILE__, __LINE__)
 
-inline ncclResult_t ncclCudaHostFree(void* ptr) {
+static inline ncclResult_t ncclCudaHostFree(void* ptr) {
   CUDACHECK(cudaFreeHost(ptr));
   return ncclSuccess;
 }
 
+#define ncclCudaHostCalloc(...) ncclCudaHostCallocDebug(__VA_ARGS__, __FILE__, __LINE__)
+
 template <typename T>
 ncclResult_t ncclCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
   if (nelem > 0) {
-    void* p = malloc(nelem*ncclSizeOfT<T>());
+    T* p = (T*)malloc(nelem*ncclSizeOfT<T>());
     if (p == NULL) {
       WARN("Failed to malloc %ld bytes", nelem*ncclSizeOfT<T>());
       return ncclSystemError;
     }
     //INFO(NCCL_ALLOC, "%s:%d malloc Size %ld pointer %p", filefunc, line, nelem*ncclSizeOfT<T>(), p);
     memset(p, 0, nelem*ncclSizeOfT<T>());
-    *ptr = (T*)p;
+    *ptr = p;
   } else {
     *ptr = NULL;
   }
@@ -78,17 +159,17 @@ ncclResult_t ncclCallocDebug(T** ptr, size_t nelem, const char *filefunc, int li
 
 template <typename T>
 ncclResult_t ncclRealloc(T** ptr, size_t oldNelem, size_t nelem) {
-  if (nelem < oldNelem) return ncclInternalError;
+  T* oldp = *ptr;
+  if (nelem < oldNelem || (oldp == NULL && oldNelem > 0)) return ncclInternalError;
   if (nelem == oldNelem) return ncclSuccess;
 
-  T* oldp = *ptr;
   T* p = (T*)malloc(nelem*ncclSizeOfT<T>());
   if (p == NULL) {
     WARN("Failed to malloc %ld bytes", nelem*ncclSizeOfT<T>());
     return ncclSystemError;
   }
-  memcpy(p, oldp, oldNelem*ncclSizeOfT<T>());
-  free(oldp);
+  if (oldp && oldNelem) memcpy(p, oldp, oldNelem * ncclSizeOfT<T>());
+  if (oldp) free(oldp);
   memset(p+oldNelem, 0, (nelem-oldNelem)*ncclSizeOfT<T>());
   *ptr = (T*)p;
   INFO(NCCL_ALLOC, "Mem Realloc old size %ld, new size %ld pointer %p", oldNelem*ncclSizeOfT<T>(), nelem*ncclSizeOfT<T>(), *ptr);
@@ -113,6 +194,40 @@ extern struct allocationTracker allocTracker[];
 #include <cuda.h>
 #include "cudawrap.h"
 
+// ncclCuMemAllocAddr takes memory handle and size and returns the mapped address pointer
+static inline ncclResult_t ncclCuMemAllocAddr(void **ptr, CUmemGenericAllocationHandle *handleIn, size_t size) {
+  ncclResult_t result = ncclSuccess;
+  size_t granularity = 0;
+  CUmemAllocationProp prop = {};
+  CUmemAccessDesc accessDesc = {};
+  int cudaDev;
+  CUDACHECK(cudaGetDevice(&cudaDev));
+  CUCHECK(cuMemGetAllocationPropertiesFromHandle(&prop, *handleIn));
+  CUCHECK(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
+  ALIGN_SIZE(size, granularity);
+  /* Reserve a virtual address range */
+  CUCHECK(cuMemAddressReserve((CUdeviceptr *)ptr, size, granularity, 0, 0));
+  /* Map the virtual address range to the physical allocation */
+  CUCHECK(cuMemMap((CUdeviceptr)*ptr, size, 0, *handleIn, 0));
+  /* Now allow RW access to the newly mapped memory */
+  accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+  accessDesc.location.id = cudaDev;
+  accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+  CUCHECK(cuMemSetAccess((CUdeviceptr)*ptr, size, &accessDesc, 1));
+  TRACE(NCCL_ALLOC, "CuMem Map Size %zu pointer %p handle %llx", size, *ptr, *handleIn);
+  return result;
+}
+
+static inline ncclResult_t ncclCuMemFreeAddr(void *ptr) {
+  if (ptr == NULL) return ncclSuccess;
+  ncclResult_t result = ncclSuccess;
+  size_t size = 0;
+  CUCHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr));
+  CUCHECK(cuMemUnmap((CUdeviceptr)ptr, size));
+  CUCHECK(cuMemAddressFree((CUdeviceptr)ptr, size));
+  return result;
+}
+
 static inline ncclResult_t ncclCuMemAlloc(void **ptr, CUmemGenericAllocationHandle *handlep, size_t size) {
   ncclResult_t result = ncclSuccess;
   size_t granularity = 0;
@@ -130,7 +245,7 @@ static inline ncclResult_t ncclCuMemAlloc(void **ptr, CUmemGenericAllocationHand
   prop.requestedHandleTypes = type;
   prop.location.id = currentDev;
   // Query device to see if RDMA support is available
-  CUCHECK(cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_SUPPORTED, currentDev));
+  CUCHECK(cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, currentDev));
   if (flag) prop.allocFlags.gpuDirectRDMACapable = 1;
   CUCHECK(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
   ALIGN_SIZE(size, granularity);
@@ -178,6 +293,15 @@ static inline ncclResult_t ncclCuMemFree(void *ptr) {
   return ncclInternalError;
 }
 
+static inline ncclResult_t ncclCuMemAllocAddr(void **ptr, CUmemGenericAllocationHandle *handleIn, size_t size) {
+  WARN("CUMEM not supported prior to CUDA 11.3");
+  return ncclInternalError;
+}
+
+static inline ncclResult_t ncclCuMemFreeAddr(void *ptr) {
+  WARN("CUMEM not supported prior to CUDA 11.3");
+  return ncclInternalError;
+}
 #endif
 
 template <typename T>
@@ -297,7 +421,8 @@ finish:
 // and if they are shared, that could cause a crash in a child process
 inline ncclResult_t ncclIbMallocDebug(void** ptr, size_t size, const char *filefunc, int line) {
   if (size > 0) {
-    size_t page_size = sysconf(_SC_PAGESIZE);
+    long page_size = sysconf(_SC_PAGESIZE);
+    if (page_size < 0) return ncclSystemError;
     void* p;
     int size_aligned = ROUNDUP(size, page_size);
     int ret = posix_memalign(&p, page_size, size_aligned);

src/include/bitops.h

@@ -185,6 +185,8 @@ inline __host__ __device__ Int pow2Up(Int x) {
 
 template<typename Int>
 inline __host__ __device__ Int pow2Down(Int x) {
+  // True, log2Down can return -1, but we don't normally pass 0 as an argument...
+  // coverity[negative_shift]
   return Int(1)<<log2Down(x);
 }
 
@@ -274,4 +276,13 @@ inline __host__ __device__ uint32_t u32fp8Decode(uint8_t x) {
   return u32fpDecode(x, 3);
 }
 
+inline __host__ __device__ uint64_t getHash(const char* string, int n) {
+  // Based on DJB2a, result = result * 33 ^ char
+  uint64_t result = 5381;
+  for (int c = 0; c < n; c++) {
+    result = ((result << 5) + result) ^ string[c];
+  }
+  return result;
+}
+
 #endif

src/include/bootstrap.h

@@ -19,8 +19,8 @@ static_assert(sizeof(struct ncclBootstrapHandle) <= sizeof(ncclUniqueId), "Boots
 ncclResult_t bootstrapNetInit();
 ncclResult_t bootstrapCreateRoot(struct ncclBootstrapHandle* handle, bool idFromEnv);
 ncclResult_t bootstrapGetUniqueId(struct ncclBootstrapHandle* handle);
-ncclResult_t bootstrapInit(struct ncclBootstrapHandle* handle, struct ncclComm* comm);
-ncclResult_t bootstrapSplit(struct ncclBootstrapHandle* handle, struct ncclComm* comm, struct ncclComm* parent, int color, int key, int* parentRanks);
+ncclResult_t bootstrapInit(int nHandles, void* handle, struct ncclComm* comm);
+ncclResult_t bootstrapSplit(uint64_t magic, struct ncclComm* comm, struct ncclComm* parent, int color, int key, int* parentRanks);
 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);

src/include/checks.h

@@ -38,21 +38,17 @@
 
 #include <errno.h>
 // Check system calls
-#define SYSCHECK(call, name) do { \
+#define SYSCHECK(statement, name) do { \
   int retval; \
-  SYSCHECKVAL(call, name, retval); \
-} while (false)
-
-#define SYSCHECKVAL(call, name, retval) do { \
-  SYSCHECKSYNC(call, name, retval); \
+  SYSCHECKSYNC((statement), name, retval); \
   if (retval == -1) { \
-    WARN("Call to " name " failed : %s", strerror(errno)); \
+    WARN("Call to " name " failed: %s", strerror(errno)); \
     return ncclSystemError; \
   } \
 } while (false)
 
-#define SYSCHECKSYNC(call, name, retval) do { \
-  retval = call; \
+#define SYSCHECKSYNC(statement, name, retval) do { \
+  retval = (statement); \
   if (retval == -1 && (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN)) { \
     INFO(NCCL_ALL,"Call to " name " returned %s, retrying", strerror(errno)); \
   } else { \
@@ -60,14 +56,33 @@
   } \
 } while(true)
 
-#define SYSCHECKGOTO(statement, RES, label) do { \
-  if ((statement) == -1) {    \
-    /* Print the back trace*/ \
-    RES = ncclSystemError;    \
-    INFO(NCCL_ALL,"%s:%d -> %d (%s)", __FILE__, __LINE__, RES, strerror(errno));    \
+#define SYSCHECKGOTO(statement, name, RES, label) do { \
+  int retval; \
+  SYSCHECKSYNC((statement), name, retval); \
+  if (retval == -1) { \
+    WARN("Call to " name " failed: %s", strerror(errno)); \
+    RES = ncclSystemError; \
     goto label; \
   } \
-} while (0);
+} while (0)
+
+// Pthread calls don't set errno and never return EINTR.
+#define PTHREADCHECK(statement, name) do { \
+  int retval = (statement); \
+  if (retval != 0) { \
+    WARN("Call to " name " failed: %s", strerror(retval)); \
+    return ncclSystemError; \
+  } \
+} while (0)
+
+#define PTHREADCHECKGOTO(statement, name, RES, label) do { \
+  int retval = (statement); \
+  if (retval != 0) { \
+    WARN("Call to " name " failed: %s", strerror(retval)); \
+    RES = ncclSystemError; \
+    goto label; \
+  } \
+} while (0)
 
 #define NEQCHECK(statement, value) do {   \
   if ((statement) != value) {             \
@@ -75,7 +90,7 @@
     INFO(NCCL_ALL,"%s:%d -> %d (%s)", __FILE__, __LINE__, ncclSystemError, strerror(errno));    \
     return ncclSystemError;     \
   }                             \
-} while (0);
+} while (0)
 
 #define NEQCHECKGOTO(statement, value, RES, label) do { \
   if ((statement) != value) { \
@@ -84,7 +99,7 @@
     INFO(NCCL_ALL,"%s:%d -> %d (%s)", __FILE__, __LINE__, RES, strerror(errno));    \
     goto label; \
   } \
-} while (0);
+} while (0)
 
 #define EQCHECK(statement, value) do {    \
   if ((statement) == value) {             \
@@ -92,7 +107,7 @@
     INFO(NCCL_ALL,"%s:%d -> %d (%s)", __FILE__, __LINE__, ncclSystemError, strerror(errno));    \
     return ncclSystemError;     \
   }                             \
-} while (0);
+} while (0)
 
 #define EQCHECKGOTO(statement, value, RES, label) do { \
   if ((statement) == value) { \
@@ -101,7 +116,7 @@
     INFO(NCCL_ALL,"%s:%d -> %d (%s)", __FILE__, __LINE__, RES, strerror(errno));    \
     goto label; \
   } \
-} while (0);
+} while (0)
 
 // Propagate errors up
 #define NCCLCHECK(call) do { \
@@ -111,7 +126,7 @@
     if (ncclDebugNoWarn == 0) INFO(NCCL_ALL,"%s:%d -> %d", __FILE__, __LINE__, RES);    \
     return RES; \
   } \
-} while (0);
+} while (0)
 
 #define NCCLCHECKGOTO(call, RES, label) do { \
   RES = call; \
@@ -120,7 +135,7 @@
     if (ncclDebugNoWarn == 0) INFO(NCCL_ALL,"%s:%d -> %d", __FILE__, __LINE__, RES);    \
     goto label; \
   } \
-} while (0);
+} while (0)
 
 #define NCCLWAIT(call, cond, abortFlagPtr) do {         \
   uint32_t* tmpAbortFlag = (abortFlagPtr);     \
@@ -130,7 +145,7 @@
     return ncclInternalError;             \
   }                                       \
   if (__atomic_load(tmpAbortFlag, __ATOMIC_ACQUIRE)) NEQCHECK(*tmpAbortFlag, 0); \
-} while (!(cond));
+} while (!(cond))
 
 #define NCCLWAITGOTO(call, cond, abortFlagPtr, RES, label) do { \
   uint32_t* tmpAbortFlag = (abortFlagPtr);             \
@@ -140,7 +155,7 @@
     goto label;                           \
   }                                       \
   if (__atomic_load(tmpAbortFlag, __ATOMIC_ACQUIRE)) NEQCHECKGOTO(*tmpAbortFlag, 0, RES, label); \
-} while (!(cond));
+} while (!(cond))
 
 #define NCCLCHECKTHREAD(a, args) do { \
   if (((args)->ret = (a)) != ncclSuccess && (args)->ret != ncclInProgress) { \

src/include/collectives.h

@@ -72,4 +72,490 @@ struct ncclConnFifo {
   ssize_t size;
   void* ptr;
 };
+
+#include <stdio.h>
+
+template<typename T>
+class PatRSAlgorithm{
+  size_t offset;
+  size_t end;
+  size_t count;
+  int chunkCount;
+  int nelem;
+  int rank;
+  int nranks;
+  int nrPow2;
+  int postFreq;
+  int lastA;
+
+  int aggFactor;
+  int as; // aggregated steps
+  int a; // step inside aggregated step
+  int sendSkipped; // number of skipped steps during aggregation
+  int recvSkipped; // number of skipped steps during aggregation
+  int phase2recv;  // receive offset for phase 2
+  int aggDelta;
+  int scale;
+  int phase;
+
+  __device__ __host__ int min(int a, int b) {
+    return (a<b)?a:b;
+  }
+
+  __device__ __host__ int getNelem() {
+    return min(chunkCount, end-offset);
+  }
+
+  __device__ __host__ int mirrorInvert(int i, int max) {
+    int ret = 0;
+    for (int mask=1, imask=max/2; mask<max; mask<<=1, imask>>=1) {
+      if ((i&mask) == 0) ret += imask;
+    }
+    return ret;
+  }
+
+  __device__ __host__ int firstBitSet(int i, int max) {
+    int ffs =
+#ifdef __CUDA_ARCH__
+      __ffs(i);
+#else
+      __builtin_ffs(i);
+#endif
+    return ffs ? ffs-1 : max;
+  }
+
+  __device__ __host__ void resetA() {
+    a = 0;
+    sendSkipped = recvSkipped = 0;
+    lastA = aggFactor;
+    if (phase >= 2) lastA /= 2*scale;
+  }
+
+  __device__ __host__ void reset() {
+    nelem = getNelem();
+    phase = 0;
+    scale = 1;
+    phase2recv = 0;
+    as = aggDelta - 1;
+    resetA();
+  }
+
+  __device__ __host__ int nBitsSet(int i) {
+    int nbits =
+#ifdef __CUDA_ARCH__
+      __popc(i);
+#else
+      __builtin_popcount(i);
+#endif
+    return nbits;
+  }
+
+  // Return 1 when only upper bits are set. For example, if nrpow2==16 we'll return 1 for 8, 12, 14, 15.
+  // A number being in the form of 1111000 implies that the complementary is 0000111 meaning it's a power of 2 minus 1.
+  __device__ __host__ int newPeer(int i, int pow2) {
+    //printf("New peer %d/%d -> %d\n", i, pow2, nBitsSet((i ^ (pow2-1)) + 1) == 1 ? 1 : 0);
+    return nBitsSet((i ^ (pow2-1)) + 1) == 1 ? 1 : 0;
+  }
+
+public:
+   __device__ __host__ PatRSAlgorithm(int stepSize, int stepDepth, size_t offset, size_t end, size_t count, int chunkCount, int rank, int nranks):
+     offset(offset), end(end), count(count), chunkCount(chunkCount), rank(rank), nranks(nranks) {
+    aggDelta = nrPow2 = (1<<log2Up(nranks));
+
+    aggFactor = 1;
+    size_t channelSize = end-offset;
+    while (stepSize / (channelSize*sizeof(T)*aggFactor) >= 2 && aggFactor < nranks/2) {
+      aggFactor *= 2;
+      aggDelta /= 2;
+    }
+    postFreq = aggFactor;
+    int d = stepDepth;
+    while (d > 1 && aggFactor < nranks/2) {
+      d /= 2;
+      aggFactor *= 2;
+      aggDelta /= 2;
+    }
+
+    reset();
+  }
+
+  __device__ __host__ void getNextOp(int &recvDim, int &sendDim, size_t &inpIx, size_t &outIx, int &recvOffset, int &sendOffset, int &sendStepOffset, int &nelemOut, int &postRecv, int &postSend, int &last) {
+restart:
+    last = 0;
+    nelemOut = nelem;
+    outIx = offset;
+    int skip = 0;
+    //printf("Phase %d as %d/%d a %d/%d scale %d\n", phase, as, aggDelta, a, lastA, scale);
+    if (phase == 0) {
+      int s = mirrorInvert(a, lastA)*aggDelta + as;
+      if (s >= nranks) skip = 1;
+      int sendDataRank = (rank + s) % nranks;
+      inpIx = sendDataRank * count + offset;
+      recvDim = -1;
+      sendDim = 0;
+      outIx = 0;
+      recvOffset = -1;
+      sendOffset = ((a - sendSkipped)%postFreq) * nelem;
+      sendStepOffset = 0;
+      if ((((a - sendSkipped)%postFreq) + 1 >= postFreq) || (a == lastA-1)) {
+        postSend = 1;
+      } else {
+        postSend = 0;
+      }
+      postRecv = 0;
+      if (skip) sendSkipped++;
+      if (++a == lastA) {
+        phase = as == 1 ? (aggFactor > 1 ? 2 : 4) : 1; // If as == 1, switch to phase 2
+        resetA();
+      }
+      if (skip == 0) return;
+    } else if (phase == 1) {
+      int s = mirrorInvert(a, lastA)*aggDelta + as;
+      if (s >= nranks) skip = 1;
+      recvDim = firstBitSet(s, nrPow2);
+      sendOffset = ((a - sendSkipped)%postFreq)*nelem;
+      recvOffset = ((a - recvSkipped)%postFreq)*nelem;
+      postSend = 0;
+      if (recvDim == 0) {
+        if ((((a - sendSkipped)%postFreq) + 1 >= postFreq) || (a == lastA-1)) postSend = 1;
+        sendStepOffset = 0;
+      } else {
+        sendStepOffset = (a - sendSkipped)/postFreq;
+      }
+      if ((((a - recvSkipped)%postFreq) + 1 >= postFreq) || (a == lastA-1)) {
+        postRecv = 1;
+      } else {
+        postRecv = 0;
+      }
+      s -= (1<<recvDim);
+      int recvDataRank = (rank + nranks + s) % nranks;
+      inpIx = recvDataRank * count + offset;
+      sendDim = s ? firstBitSet(s, nrPow2) : -1;
+      if (sendDim == -1) {
+        sendOffset = -1;
+        sendStepOffset = 0;
+      } else if (as - (1<<recvDim) == 0) {
+        if (newPeer(a, aggFactor)) sendSkipped = a;
+        int foffset = a - sendSkipped;
+        sendStepOffset = recvDim == 0 ? 0 : foffset/postFreq;
+        sendOffset = (foffset%postFreq)*nelem;
+      }
+      if (s < nranks && skip) {
+        recvDim = -1;
+        recvOffset = -1;
+        postRecv = 0;
+        skip = 0;
+      }
+      if (skip || recvDim == -1) recvSkipped++;
+      if (skip) sendSkipped++;
+      if (++a == lastA) {
+        as--;
+        phase = as % 2 == 1 ? 0 : 1;
+        resetA();
+      }
+      if (skip == 0) return;
+    } else if (phase == 2) {
+      int s = (2*mirrorInvert(a, lastA)+1)*scale*aggDelta + 1;
+      postRecv = 0;
+      if (s >= nranks) skip = 1;
+      recvDim = 0;
+      postSend = a == lastA-1 ? 1 : 0;
+      s -= 1;
+      if (s < nranks && skip) {
+        recvDim = -1;
+        recvOffset = -1;
+        skip = 0;
+      } else if (!skip) {
+        int foffset = phase2recv;
+        phase2recv++;
+        postRecv |= ((foffset+1)%postFreq) == 0 ? 1 : 0;
+        recvOffset = (foffset%postFreq) * nelem;
+      }
+      int recvDataRank = (rank + nranks + s) % nranks;
+      inpIx = recvDataRank * count + offset;
+      sendDim = s ? firstBitSet(s, nrPow2) : -1;
+      int foffset = a - sendSkipped;
+      postSend |= ((foffset+1)%postFreq) == 0 ? 1 : 0;
+      sendStepOffset = 0;
+      sendOffset = (foffset%postFreq) * nelem;
+      if (skip || sendDim == -1) sendSkipped++;
+      if (++a == lastA) {
+        phase = 3;
+        resetA();
+      }
+      if (skip == 0) return;
+    } else if (phase == 3) {
+      int s = (2*mirrorInvert(a, lastA)+1)*scale*aggDelta;
+      postRecv = a == lastA-1 ? 1 : 0;
+      if (s >= nranks) skip = 1;
+      recvDim = firstBitSet(s, nrPow2);
+      postSend = 0;
+      s -= (1<<recvDim);
+      int foffset = a - recvSkipped;
+      postRecv |= (foffset+1)%postFreq == 0 ? 1 : 0;
+      recvOffset = (foffset%postFreq) * nelem;
+      int recvDataRank = (rank + nranks + s) % nranks;
+      inpIx = recvDataRank * count + offset;
+      sendDim = s ? firstBitSet(s, nrPow2) : -1;
+      if (s < nranks && skip) {
+        recvDim = -1;
+        recvOffset = -1;
+        postRecv = 0;
+        skip = 0;
+      }
+      if (newPeer(a, aggFactor/(2*scale))) sendSkipped = a;
+      foffset = a - sendSkipped;
+      sendStepOffset = foffset / postFreq; // Accumulate on next steps
+      sendOffset = sendDim >= 0 ? (foffset%postFreq) * nelem : -1;
+      if (skip || recvDim == -1) recvSkipped++;
+      if (skip) sendSkipped++;
+      if (++a == lastA) {
+        scale *= 2;
+        phase = scale < aggFactor ? 2 : 4;
+        resetA();
+      }
+      if (skip == 0) return;
+    } else if (phase == 4) {
+      recvDim = 0;
+      sendDim = -1;
+      inpIx = rank * count + offset;
+      recvOffset = (phase2recv%postFreq) * nelem;
+      sendStepOffset = 0;
+      sendOffset = -1;
+      postRecv = 1;
+      postSend = 0;
+      offset += chunkCount;
+      if (offset >= end) {
+        last = 1;
+      } else {
+        reset();
+      }
+      return;
+    }
+    goto restart;
+  }
+};
+
+template<typename T>
+class PatAGAlgorithm{
+  size_t offset;
+  size_t end;
+  size_t count;
+  int chunkCount;
+  int nelem;
+  int rank;
+  int nranks;
+  int nrPow2;
+  int postFreq;
+  int lastA;
+
+  int aggFactor;
+  int as; // aggregated steps
+  int a; // step inside aggregated step
+  int aggDelta;
+
+  int scale;
+
+  int phase;
+
+  // AS computation
+  int asDim;
+  int v;
+  int bitCount[32];
+  int bitZeroStep[32];
+
+  __device__ __host__ int min(int a, int b) {
+    return (a<b)?a:b;
+  }
+
+  __device__ __host__ int getNelem() {
+    return min(chunkCount, end-offset);
+  }
+
+  __device__ __host__ int mirror(int i, int max) {
+    int ret = 0;
+    for (int mask=1, imask=max/2; mask<max; mask<<=1, imask>>=1) {
+      if ((i&mask)) ret += imask;
+    }
+    return ret;
+  }
+
+  __device__ __host__ int firstBitSet(int i, int max) {
+    int ffs =
+#ifdef __CUDA_ARCH__
+      __ffs(i);
+#else
+      __builtin_ffs(i);
+#endif
+    return ffs ? ffs-1 : max;
+  }
+
+  __device__ __host__ void resetA() {
+    a = 0;
+    lastA = aggFactor;
+    if (phase >= 2) lastA /= 2*scale;
+  }
+
+  __device__ __host__ void reset() {
+    nelem = getNelem();
+    scale = aggFactor/2;
+    phase = scale ? 2 : 1;
+    v = 0;
+    for (int i = 0; i<asDim; i++) {
+      bitCount[i] = asDim-i;
+      bitZeroStep[i] = 1;
+    }
+    as = nextAs();
+    resetA();
+  }
+
+  __device__ __host__ int nextAs() {
+    for (int d=0; d<asDim; d++) {
+      int p = 1<<d;
+      bitCount[d]--;
+      if (bitCount[d] == 0) {
+        v ^= p;
+        bitCount[d] = p;
+        if ((v&p) == 0) {
+          bitCount[d] += firstBitSet(bitZeroStep[d], asDim) - 1;
+          if (bitCount[d] == 0) {
+            v ^= p;
+            bitCount[d] = p;
+          }
+          bitZeroStep[d]++;
+        }
+      }
+    }
+    return v;
+  }
+
+
+public:
+   __device__ __host__ PatAGAlgorithm(int stepSize, int stepDepth, size_t offset, size_t end, size_t count, int chunkCount, int rank, int nranks):
+     offset(offset), end(end), count(count), chunkCount(chunkCount), rank(rank), nranks(nranks) {
+    aggDelta = nrPow2 = (1<<log2Up(nranks));
+
+    aggFactor = 1;
+    size_t channelSize = end-offset;
+    while (stepSize / (channelSize*sizeof(T)*aggFactor) >= 2 && aggFactor < nranks/2) {
+      aggFactor *= 2;
+      aggDelta /= 2;
+    }
+    postFreq = aggFactor;
+    int d = stepDepth;
+    while (d > 1 && aggFactor < nranks/2) {
+      d /= 2;
+      aggFactor *= 2;
+      aggDelta /= 2;
+    }
+    //printf("AggFactor %d PostFreq %d AggDelta %d\n", aggFactor, postFreq, aggDelta);
+
+    asDim = log2Up(aggDelta);
+    reset();
+  }
+
+  __device__ __host__ void getNextOp(int &recvDim, int &sendDim, size_t &inpIx, size_t &outIx, int &recvOffset, int &sendOffset, int &recvStepOffset, int &nelemOut, int &postRecv, int &postSend, int &last) {
+restart:
+    //printf("Phase %d as %d/%d a %d/%d scale %d\n", phase, as, aggDelta, a, lastA, scale);
+    last = 0;
+    nelemOut = nelem;
+    inpIx = offset;
+    int skip = 0;
+    if (phase == 0) {
+      int s = a*aggDelta + as;
+      if (s >= nranks) skip = 1;
+      int nextSkip = (a+1)*aggDelta + as >= nranks ? 1 : 0;
+      int recvDataRank = (rank + s) % nranks;
+      outIx = recvDataRank * count + offset;
+      sendDim = -1;
+      recvDim = 0;
+      inpIx = 0;
+      sendOffset = -1;
+      recvOffset = (a % postFreq) * nelem;
+      recvStepOffset = 0;
+      postRecv = (a % postFreq == postFreq-1) || ((a+1)*aggDelta+as >= nranks) ? 1 : 0;
+      postSend = 0;
+      a++;
+      if (nextSkip) {
+        as = nextAs();
+        if (as == aggDelta/2) {
+          offset += chunkCount;
+          if (offset >= end) {
+            last = 1;
+          } else {
+            reset();
+          }
+          return;
+        }
+        phase = 1;
+        resetA();
+      }
+      if (skip == 0) return;
+   } else if (phase == 1) {
+      int s = a*aggDelta + as;
+      if (s >= nranks) skip = 1;
+      sendDim = firstBitSet(s, nrPow2);
+      s -= (1<<sendDim);
+      int sendDataRank = (rank + nranks + s) % nranks;
+      outIx = sendDataRank * count + offset;
+      recvDim = s ? firstBitSet(s, nrPow2) : -1;
+      sendOffset = recvOffset = (a % postFreq) * nelem;
+      postSend = (a % postFreq == postFreq-1) || ((a+1)*aggDelta+as >= nranks) ? 1 : 0;
+      postRecv = (sendDim == 0) && ((a % postFreq == postFreq-1) || ((a+1)*aggDelta+as-1 >= nranks)) ? 1 : 0;
+      recvStepOffset = (sendDim == 0) ? 0 : a/postFreq;
+      if (recvDim == -1) {
+        recvOffset = -1;
+        postRecv = 0;
+      } else if (as - (1<<sendDim) == 0) {
+        int foffset = (a*aggDelta) >> (recvDim+1);
+        recvOffset = (foffset%postFreq)*nelem;
+        postRecv = (sendDim == 0) && ((foffset % postFreq == postFreq-1) || ((((foffset+1)*2)+1)<<recvDim) >= nranks) ? 1 : 0;
+        recvStepOffset = (sendDim == 0) ? 0 : foffset/postFreq;
+      }
+      if (s < nranks && sendDim == 0 && skip) {
+        // Don't forget to receive at least once even if we don't send afterwards
+        sendDim = -1;
+        sendOffset = -1;
+        postSend = 0;
+        skip = 0;
+      }
+      if (++a == lastA) {
+        if (as % 2 == 1) {
+          phase = 0;
+        } else {
+          as = nextAs();
+        }
+        resetA();
+      }
+      if (skip == 0) return;
+    } else if (phase == 2) {
+      int s = (2*a+1)*scale*aggDelta;
+      postSend = (a % postFreq == postFreq-1) || ((2*(a+1)+1)*scale*aggDelta >= nranks) ? 1 : 0;
+      postRecv = 0;
+      if (s >= nranks) skip = 1;
+      sendDim = firstBitSet(s, nrPow2);
+      s -= (1<<sendDim);
+      sendOffset = (a%postFreq) * nelem;
+      recvStepOffset = a / postFreq;
+      int sendDataRank = (rank + nranks + s) % nranks;
+      outIx = sendDataRank * count + offset;
+      recvDim = s ? firstBitSet(s, nrPow2) : -1;
+      s -= (1<<recvDim);
+      if (recvDim == -1) {
+        recvOffset = -1;
+      } else {
+        int foffset = (a*2*scale*aggDelta) >> (recvDim+1);
+        recvOffset = (foffset%postFreq)*nelem;
+        recvStepOffset = foffset / postFreq;
+      }
+      if (++a == lastA) {
+        scale /= 2;
+        phase = scale ? 2 : 1;
+        resetA();
+      }
+      if (skip == 0) return;
+    }
+    goto restart;
+  }
+};
 #endif

src/include/comm.h

@@ -18,6 +18,7 @@
 #include "register.h"
 #include "graph.h"
 #include "nvmlwrap.h"
+#include "profiler.h"
 
 #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
 #define HIPRT_CB
@@ -110,6 +111,11 @@ struct ncclCommCallback {
   struct ncclCommCallback* next;
   ncclResult_t(*fn)(struct ncclComm* comm, struct ncclCommCallback* cb);
 };
+struct ncclCommEventCallback {
+  struct ncclCommEventCallback* next;
+  cudaEvent_t event;
+  ncclResult_t(*fn)(struct ncclComm* comm, struct ncclCommEventCallback* cb);
+};
 
 struct ncclSharedResources {
   int refCount;
@@ -179,6 +185,56 @@ struct ncclCollnetHandleList {
   struct ncclProxyConnector* proxyconn;
 };
 
+struct ncclTaskColl {
+  struct ncclTaskColl* next;
+  ncclFunc_t func;
+  void const* sendbuff;
+  void* recvbuff;
+  size_t count;
+  int root;
+  ncclDataType_t datatype;
+  ncclRedOp_t opHost;
+  struct ncclDevRedOpFull opDev;
+  int chunkSteps, sliceSteps;
+  // Computed later:
+  size_t trafficBytes;
+  int32_t nMaxChannels:8;
+  int32_t nWarps:8;
+  int32_t algorithm:8, protocol:8;
+  uint32_t isCollnet:1, isNvls:1;
+  uint32_t devFuncId:30;
+  enum ncclRegBufferType regBufType;
+  uint64_t opCount;
+  // number of elements in planner->ipcMemQueue associated with this collective
+  int nCleanupQueueElts;
+
+  void* sendMhandle;
+  void* recvMhandle;
+  // index for IPC record lookup
+  uintptr_t sendbuffOffset;
+  uintptr_t recvbuffOffset;
+  uintptr_t* sendbuffRmtAddrs;
+  uintptr_t* recvbuffRmtAddrs;
+
+  // Profiler plugin
+  int eActivationMask;
+  void* eventHandle;
+};
+struct ncclTaskP2p {
+  struct ncclTaskP2p* next;
+  ncclFunc_t func;
+  void* buff;
+  size_t count;
+  ncclDataType_t datatype;
+  int root;
+  size_t bytes;
+  uint64_t opCount;
+
+  // Profiler plugin
+  int eActivationMask;
+  void* eventHandle;
+};
+
 struct ncclKernelPlan {
   // A kernel plan is also a callback that reclaims itself. Hence this must
   // be the first member.
@@ -204,42 +260,12 @@ struct ncclKernelPlan {
   struct ncclIntruQueue<struct ncclCommCallback, &ncclCommCallback::next> cleanupQueue;
   void* workBufPersistent;
 
+  struct ncclIntruQueue<struct ncclTaskP2p, &ncclTaskP2p::next> p2pTaskQueue;
+  struct ncclIntruQueue<struct ncclTaskColl, &ncclTaskColl::next> collTaskQueue;
   struct ncclIntruQueue<struct ncclProxyOp, &ncclProxyOp::enqNext> proxyOpQueue;
-};
 
-////////////////////////////////////////////////////////////////////////////////
-
-struct ncclTaskColl {
-  struct ncclTaskColl* next;
-  ncclFunc_t func;
-  void const* sendbuff;
-  void* recvbuff;
-  size_t count;
-  int root;
-  ncclDataType_t datatype;
-  ncclRedOp_t opHost;
-  struct ncclDevRedOpFull opDev;
-  int chunkSteps, sliceSteps;
-  // Computed later:
-  size_t trafficBytes;
-  int32_t nMaxChannels:8;
-  int32_t nWarps:8;
-  int32_t algorithm:8, protocol:8;
-  uint32_t isCollnet:1, isNvls:1;
-  uint32_t devFuncId:30;
-  enum ncclRegBufferType regBufType;
-  // number of elements in planner->ipcMemQueue associated with this collective
-  int nCleanupQueueElts;
-
-  void* sendMhandle;
-  void* recvMhandle;
-  uint64_t opCount;
-};
-struct ncclTaskP2p {
-  struct ncclTaskP2p* next;
-  void* buff;
-  size_t bytes;
-  uint64_t opCount;
+  // Profiler plugin
+  void* groupEventHandle;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -395,6 +421,7 @@ struct ncclPeerInfo {
   // MNNVL support
   nvmlGpuFabricInfoV_t fabricInfo;
   int cuMemSupport;
+  int version;
 };
 
 struct ncclComm {
@@ -410,6 +437,8 @@ struct ncclComm {
   struct ncclChannel channels[MAXCHANNELS];
   struct ncclPeerInfo* peerInfo;
   struct ncclTopoSystem* topo;
+  struct ncclProxyConnector* gproxyConn;
+  struct ncclIntruQueue<struct ncclCommCallback, &ncclCommCallback::next> legacyRegCleanupQueue;
 
   int netPluginLoaded;
   ncclNet_t* ncclNet;
@@ -422,10 +451,12 @@ struct ncclComm {
   struct ncclTopoGraph graphs[NCCL_NUM_ALGORITHMS];
   bool initAlgoChannels[NCCL_NUM_ALGORITHMS];
   bool runtimeConn; // if dynamic connection is supported
+  bool directMode;
   int cuMemSupport;
 
   uint64_t magic; // Magic number for all network communication. Not a security key -- only goal is to detect mismatches.
 
+  const char* commName;
   uint64_t commHash;
   int rank;    // my rank in the communicator
   int nRanks;  // number of GPUs in communicator
@@ -488,7 +519,7 @@ struct ncclComm {
   int maxThreads[NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
 
   /* This attribute can indicate the states of communicators and return code of
-   * asynchronous NCCL operations. */
+  * asynchronous NCCL operations. */
   ncclResult_t asyncResult;
 
   // Flag to ask NCCL kernels to abort
@@ -537,7 +568,7 @@ struct ncclComm {
   int collNetSupport;
   bool collNetRegSupport;
   uint8_t collNetSupportMatrix[4/*sum,prod,max,min*/][ncclNumTypes];
-  int intraHighestTransportType;
+  bool intraNodeP2pSupport;
   int* collNetHeads;
   int collNetHeadsNum;
   int* collNetDenseToUserRank;
@@ -552,6 +583,8 @@ struct ncclComm {
   struct ncclNvlsSharedRes* nvlsResources;
 
   // pools backed by comm->memPermanent
+  struct ncclMemoryPool memPool_ncclTaskColl;
+  struct ncclMemoryPool memPool_ncclTaskP2p;
   struct ncclMemoryPool memPool_ncclProxyOp;
   struct ncclMemoryPool memPool_ncclKernelPlan;
 
@@ -566,6 +599,13 @@ struct ncclComm {
   struct ncclKernelPlanner planner;
 
   hipStream_t sideStream; // [RCCL] Cached non-captured stream
+  
+  cudaMemPool_t memPool;
+  // Queue of events and associated callbacks for cleaning up asynchronous work.
+  // Using this is preferable to using CUDA host callbacks because host callbacks
+  // won't allow the work following the callback to run until the callback completes,
+  // which comes at expense to perf.
+  struct ncclIntruQueue<struct ncclCommEventCallback, &ncclCommEventCallback::next> eventCallbackQueue;
 
   // user-created reduction ops
   int userRedOpCapacity, userRedOpFreeHead;
@@ -614,6 +654,11 @@ struct ncclComm {
   int tunerPluginLoaded;
   ncclTuner_t* tuner;
   void *tunerContext;
+
+  // Profiler plugin
+  void* profilerContext;
+  uint64_t seqNumber[NCCL_NUM_FUNCTIONS];
+
   // buffer registration cache
   struct ncclRegCache regCache;
   uint64_t endMagic;
@@ -647,6 +692,27 @@ inline ncclResult_t ncclCommPollCallbacks(struct ncclComm* comm, bool waitSome)
   return ncclSuccess;
 }
 
+inline ncclResult_t ncclCommPollEventCallbacks(struct ncclComm *comm) {
+  ncclResult_t result = ncclSuccess;
+  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  while (true) {
+    struct ncclCommEventCallback* cb = ncclIntruQueueHead(&comm->eventCallbackQueue);
+    if (cb == nullptr) break;
+    cudaError_t ok = cudaEventSynchronize(cb->event);
+    if (ok == cudaErrorNotReady) break;
+    ncclIntruQueueDequeue(&comm->eventCallbackQueue);
+    if (ok == cudaSuccess) {
+      NCCLCHECKGOTO(cb->fn(comm, cb), result, finish);
+    } else {
+      CUDACHECKGOTO(ok, result, finish);
+    }
+  }
+finish:
+  cudaThreadExchangeStreamCaptureMode(&mode);
+  return ncclSuccess;
+}
+
 inline void ncclCommIntraBarrierIn(struct ncclComm* comm, uint32_t x) {
   int phase = comm->intraBarrierPhase;
   if (comm->intraRanks == 1) {

src/include/cudawrap.h

@@ -13,6 +13,7 @@
 
 // Is cuMem API usage enabled
 extern int ncclCuMemEnable();
+extern int ncclCuMemHostEnable();
 
 #if CUDART_VERSION >= 11030
 #include <cudaTypedefs.h>
@@ -96,6 +97,7 @@ DECLARE_CUDA_PFN_EXTERN(cuMemRelease);
 DECLARE_CUDA_PFN_EXTERN(cuMemRetainAllocationHandle);
 DECLARE_CUDA_PFN_EXTERN(cuMemSetAccess);
 DECLARE_CUDA_PFN_EXTERN(cuMemUnmap);
+DECLARE_CUDA_PFN_EXTERN(cuMemGetAllocationPropertiesFromHandle);
 #if CUDA_VERSION >= 11070
 DECLARE_CUDA_PFN_EXTERN(cuMemGetHandleForAddressRange); // DMA-BUF support
 #endif

src/include/device.h

@@ -54,9 +54,9 @@ struct ncclDevRedOpFull {
 
 union ncclLLFifoLine {
   /* Flags have to be *after* data, because otherwise, an incomplete receive
-     from the network may receive the flag but not the data.
-     Note this is assuming that either we receive contiguous chunks of data
-     (sockets) or data is written with an atomicity of 8 bytes (IB/RDMA). */
+    from the network may receive the flag but not the data.
+    Note this is assuming that either we receive contiguous chunks of data
+    (sockets) or data is written with an atomicity of 8 bytes (IB/RDMA). */
   struct {
     uint32_t data1;
     uint32_t flag1;
@@ -144,6 +144,8 @@ struct ncclConnInfo {
 };
 
 struct ncclProxyConnector {
+  bool initialized;
+  int rank;
   int tpRank;
   int tpLocalRank;
   int sameProcess;
@@ -157,6 +159,8 @@ struct ncclConnector {
   struct ncclTransportComm* transportComm;
   void* transportResources;
   struct ncclConnInfo conn;
+  int sendMemSameProcess;
+  int recvMemSameProcess;
 };
 
 struct ncclRing {
@@ -247,6 +251,8 @@ struct alignas(16) ncclDevWorkP2p {
   uint8_t sendProtoLL:1, recvProtoLL:1;
   uint8_t sendRegistered:1, recvRegistered:1;
 
+  uint8_t sendIpcReg:1, recvIpcReg:1;
+
   uint8_t sendConnIndex:2, recvConnIndex:2;
 };
 
@@ -298,6 +304,10 @@ struct alignas(16) ncclDevWorkColl {
   uint16_t pivotA2ANumBiRings;
   void* recvbuff;
   void* sendbuff;
+  uintptr_t sendbuffOffset;
+  uintptr_t recvbuffOffset;
+  uintptr_t* sendbuffRmtAddrs;
+  uintptr_t* recvbuffRmtAddrs;
   union {
     // Continuous-byte-distribution scheduling. The lo and hi channels are of
     // different size than the channels in the middle.
@@ -319,9 +329,9 @@ struct alignas(16) ncclDevWorkColl {
 
 __host__ __device__ constexpr int ncclProtoGrainSize(int proto) {
   return proto == NCCL_PROTO_LL ? 16 :
-         proto == NCCL_PROTO_LL128 ? WARP_SIZE*NCCL_LL128_SHMEM_ELEMS_PER_THREAD/NCCL_LL128_LINEELEMS*NCCL_LL128_DATAELEMS*sizeof(uint64_t) :
-         proto == NCCL_PROTO_SIMPLE ? 512 :
-         -1;
+        proto == NCCL_PROTO_LL128 ? WARP_SIZE*NCCL_LL128_SHMEM_ELEMS_PER_THREAD/NCCL_LL128_LINEELEMS*NCCL_LL128_DATAELEMS*sizeof(uint64_t) :
+        proto == NCCL_PROTO_SIMPLE ? 512 :
+        -1;
 }
 
 template<typename Int>
@@ -367,7 +377,7 @@ enum ncclDevWorkType: uint8_t {
 
 constexpr size_t ncclDevWorkSize(enum ncclDevWorkType type) {
   return type == ncclDevWorkTypeP2p ? sizeof(ncclDevWorkP2p) :
-         type == ncclDevWorkTypeColl ? sizeof(ncclDevWorkColl) : sizeof(ncclDevWorkCollReg);
+        type == ncclDevWorkTypeColl ? sizeof(ncclDevWorkColl) : sizeof(ncclDevWorkCollReg);
 }
 
 #define NCCL_MAX_DEV_WORK_BATCH_BYTES 128
@@ -493,6 +503,7 @@ struct ncclDevComm {
   int nNodes;
   int buffSizes[NCCL_NUM_PROTOCOLS];
   int p2pChunkSize;
+  int isNvlink;
   int p2pnChannelsPerPeer;
 
   // Work fifo return credits
@@ -506,6 +517,8 @@ struct ncclDevComm {
   // Channels, device side
   struct ncclDevChannel* channels/*[MAXCHANNELS]*/;
 
+  int* rankToLocalRank;
+
 #if defined(ENABLE_NPKIT)
   NpKitEventCollectContext* npKitEventCollectContexts;
   uint64_t* cpuTimestamp;
@@ -686,7 +699,7 @@ inline int ncclDevFuncId(int coll, int devRedOp, int type, int algo, int proto)
       row += (((algo * NCCL_NUM_PROTOCOLS + proto) * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - NCCL_NUM_FLOATS * (algo * NCCL_NUM_PROTOCOLS + proto);
       break;
     }
-    row += (NCCL_NUM_ALGORITHMS - 4) * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
+    row += (NCCL_NUM_ALGORITHMS - 5) * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - NCCL_NUM_FLOATS);
 
     // RING / SIMPLE / Sum / int8_t
     if (coll == ncclFuncAllToAllPivot) break;

src/include/graph.h

@@ -34,16 +34,17 @@ ncclResult_t ncclTopoComputeCommCPU(struct ncclComm* comm);
 
 // Query topology
 ncclResult_t ncclTopoGetNetDev(struct ncclComm* comm, int rank, struct ncclTopoGraph* graph, int channelId, int peerRank, int64_t* id, int* dev, int* proxyRank);
-ncclResult_t ncclTopoCheckP2p(struct ncclTopoSystem* system, int64_t id1, int64_t id2, int* p2p, int *read, int* intermediateRank);
+ncclResult_t ncclTopoCheckP2p(struct ncclTopoSystem* system, int rank1, int rank2, int* p2p, int *read, int* intermediateRank);
 ncclResult_t ncclTopoCheckMNNVL(struct ncclTopoSystem* system, struct ncclPeerInfo* info1, struct ncclPeerInfo* info2, int* ret);
 ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* topo, int64_t busId, int64_t netId, int read, int* useGdr);
 #define MAX_XGMI_INTER_GPUS 4
 ncclResult_t ncclTopoGetIntraNetDev(struct ncclTopoSystem* system, int rank, struct ncclTopoGraph* graph, int channelId, int type, int64_t* id, 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);
 ncclResult_t ncclTopoNeedFlush(struct ncclTopoSystem* system, int64_t busId, int* flush);
-ncclResult_t ncclTopoCheckNet(struct ncclTopoSystem* system, int64_t id1, int64_t id2, int* net);
+ncclResult_t ncclTopoCheckNet(struct ncclTopoSystem* system, int rank1, int rank2, int* net);
 int ncclPxnDisable(struct ncclComm* comm);
 ncclResult_t ncclTopoGetPxnRanks(struct ncclComm* comm, int** intermediateRanks, int* nranks);
+ncclResult_t ncclGetLocalCpu(struct ncclTopoSystem* system, int gpu, int* retCpu);
 
 // Find CPU affinity
 ncclResult_t ncclTopoGetCpuAffinity(struct ncclTopoSystem* system, int rank, cpu_set_t* affinity);
@@ -82,7 +83,7 @@ ncclResult_t ncclTopoSearchInit(struct ncclTopoSystem* system);
 #define NCCL_TOPO_PATTERN_COLLNET_DIRECT 6  // Collnet Direct
 struct ncclTopoGraph {
   // Input / output
-  int id; // ring : 0, tree : 1, collnet : 2
+  int id; // ring : 0, tree : 1, collnet : 2, nvls : 3, collnetDirect : 4
   int pattern;
   int crossNic;
   int collNet;

src/include/nccl_common.h

@@ -55,7 +55,7 @@ typedef enum {
   ncclNumFuncs = 9
 } ncclFunc_t;
 
-#define NCCL_NUM_ALGORITHMS 6 // Tree/Ring/CollNet*
+#define NCCL_NUM_ALGORITHMS 7 // Tree/Ring/CollNet*
 #define NCCL_ALGO_UNDEF -1
 #define NCCL_ALGO_TREE 0
 #define NCCL_ALGO_RING 1
@@ -63,6 +63,7 @@ typedef enum {
 #define NCCL_ALGO_COLLNET_CHAIN 3
 #define NCCL_ALGO_NVLS 4
 #define NCCL_ALGO_NVLS_TREE 5
+#define NCCL_ALGO_PAT 6
 
 #define NCCL_NUM_PROTOCOLS 3 // Simple/LL/LL128
 #define NCCL_PROTO_UNDEF -1

src/include/nccl_profiler.h

@@ -0,0 +1,150 @@
+/*************************************************************************
+ * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef NCCL_PROFILER_H_
+#define NCCL_PROFILER_H_
+
+#include <cstdint>
+
+enum {
+  ncclProfileGroup     = (1 << 0),  // group event type
+  ncclProfileColl      = (1 << 1),  // host collective call event type
+  ncclProfileP2p       = (1 << 2),  // host point-to-point call event type
+  ncclProfileProxyOp   = (1 << 3),  // proxy operation event type
+  ncclProfileProxyStep = (1 << 4),  // proxy step event type
+  ncclProfileProxyCtrl = (1 << 5),  // proxy control event type
+  ncclProfileNumEvents = (     6),
+};
+
+typedef struct {
+  uint8_t type;                 // event type descriptor: ncclProfileColl, ...
+  void* parentObj;              // pointer to the profiler parent object (for coll is the group)
+  int rank;                     // originating rank
+  union {
+    struct {
+      const char* name;
+      uint64_t commHash;
+      uint64_t seqNumber;
+      uint8_t func;
+      void const* sendBuff;
+      void* recvBuff;
+      size_t count;
+      int root;
+      uint8_t datatype;
+      uint32_t op;
+      size_t trafficBytes;
+      uint8_t nMaxChannels;
+      uint8_t nWarps;
+      uint8_t algo;
+      uint8_t proto;
+      int isCollnet;
+      int isNvls;
+    } coll;
+
+    struct {
+      const char* name;
+      uint64_t commHash;
+      uint8_t func;
+      void* buff;
+      uint8_t datatype;
+      size_t count;
+      int peer;
+    } p2p;
+
+    struct {
+      pid_t pid;                // pid of the originating process
+      uint8_t channelId;        // channel id for this proxy operation
+      int peer;                 // remote rank for send/recv
+      int nSteps;               // number of steps for this proxy operation
+      int chunkSize;            // amount of data transferred by this proxy operation
+      int isSend;
+    } proxyOp;
+
+    struct {
+      int step;
+    } proxyStep;
+  };
+} ncclProfilerEventDescr_v1_t;
+
+typedef enum {
+  ncclProfilerProxyOpSendPosted,
+  ncclProfilerProxyOpSendRemFifoWait,
+  ncclProfilerProxyOpSendTransmitted,
+  ncclProfilerProxyOpSendDone,
+  ncclProfilerProxyOpRecvPosted,
+  ncclProfilerProxyOpRecvReceived,
+  ncclProfilerProxyOpRecvTransmitted,
+  ncclProfilerProxyOpRecvDone,
+
+  /* Legacy proxy profiler states */
+  ncclProfilerProxyStepSendGPUWait,
+  ncclProfilerProxyStepSendWait,
+  ncclProfilerProxyStepRecvWait,
+  ncclProfilerProxyStepRecvFlushWait,
+  ncclProfilerProxyStepRecvGPUWait,
+
+  /* Legacy proxy control states */
+  ncclProfilerProxyCtrlIdle,
+  ncclProfilerProxyCtrlActive,
+  ncclProfilerProxyCtrlSleep,
+  ncclProfilerProxyCtrlWakeup,
+  ncclProfilerProxyCtrlAppend,
+  ncclProfilerProxyCtrlAppendEnd,
+} ncclProfilerEventState_v1_t;
+
+typedef union {
+  struct {
+    size_t transSize;
+    int steps;
+  } proxyOp;
+
+  struct {
+    int appendedProxyOps;
+  } proxyCtrl;
+} ncclProfilerEventStateArgs_v1_t;
+
+typedef struct {
+  const char* name;
+
+  // init - initialize the profiler plugin
+  // Input
+  //  - context        : opaque profiler context object for separating profiler behavior across comms
+  // Output
+  //  - eActivationMask: bitmask of active events set by the plugin
+  ncclResult_t (*init)(void** context, int* eActivationMask);
+
+  // startEvent - initialize and start a new event for the supplied event descriptor inside the eventset
+  // Input
+  //  - context: opaque profiler context object
+  //  - eDescr : pointer to ncclProfilerEventDescr_t object
+  // Output
+  //  - eHandle: return event handle for supplied event descriptor object
+  ncclResult_t (*startEvent)(void* context, void** eHandle, ncclProfilerEventDescr_v1_t* eDescr);
+
+  // stopEvent - stop/finalize an event inside and event set
+  // Input
+  //  - eHandle: handle to event object
+  ncclResult_t (*stopEvent)(void* eHandle);
+
+  // recordEventState - record event state transitions and event attribute updates
+  // Input
+  //  - eHandle   : handle to event object created through startEvent
+  //  - eStateArgs: optional argument used to capture event attribute updates associated with the state transition
+  //  - eState    : event state transition
+  ncclResult_t (*recordEventState)(void* eHandle, ncclProfilerEventState_v1_t eState, ncclProfilerEventStateArgs_v1_t* eStateArgs);
+
+  // finalize - finalize the profiler plugin
+  // Input
+  //  - context: opaque profiler context object
+  ncclResult_t (*finalize)(void* context);
+} ncclProfiler_v1_t;
+
+typedef ncclProfilerEventDescr_v1_t ncclProfilerEventDescr_t;
+typedef ncclProfilerEventState_v1_t ncclProfilerEventState_t;
+typedef ncclProfilerEventStateArgs_v1_t ncclProfilerEventStateArgs_t;
+typedef ncclProfiler_v1_t ncclProfiler_t;
+
+#endif

src/include/nvtx.h

@@ -17,32 +17,35 @@
 #endif
 
 // Define all NCCL-provided static schema IDs here (avoid duplicates).
-#define NVTX_SID_CommInitRank  0
-#define NVTX_SID_CommInitAll   1
-#define NVTX_SID_CommDestroy   2 // same schema as NVTX_SID_CommInitRank
-#define NVTX_SID_CommAbort     3 // same schema as NVTX_SID_CommInitRank
-#define NVTX_SID_AllGather     4
-#define NVTX_SID_AllReduce     5
-#define NVTX_SID_AllToAll      6
-#define NVTX_SID_AllToAllv     7
-#define NVTX_SID_Broadcast     8
-#define NVTX_SID_Gather        9
-#define NVTX_SID_MSCCL         10
-#define NVTX_SID_ReduceScatter 11
-#define NVTX_SID_Reduce        12
-#define NVTX_SID_Scatter       13
-#define NVTX_SID_Send          14
-#define NVTX_SID_Recv          15
+#define NVTX_SID_CommInitRank         0
+#define NVTX_SID_CommInitAll          1
+#define NVTX_SID_CommDestroy          2 // same schema as NVTX_SID_CommInitRank
+#define NVTX_SID_CommAbort            3 // same schema as NVTX_SID_CommInitRank
+#define NVTX_SID_AllGather            4
+#define NVTX_SID_AllReduce            5
+#define NVTX_SID_AllToAll             6
+#define NVTX_SID_AllToAllv            7
+#define NVTX_SID_Broadcast            8
+#define NVTX_SID_Gather               9
+#define NVTX_SID_MSCCL                10
+#define NVTX_SID_ReduceScatter        11
+#define NVTX_SID_Reduce               12
+#define NVTX_SID_Scatter              13
+#define NVTX_SID_Send                 14
+#define NVTX_SID_Recv                 15
+#define NVTX_SID_CommInitRankConfig   16 // same schema as NVTX_SID_CommInitRank
+#define NVTX_SID_CommInitRankScalable 17 // same schema as NVTX_SID_CommInitRank
+#define NVTX_SID_CommSplit            18
 
 // Define static schema ID for the reduction operation.
-#define NVTX_PAYLOAD_ENTRY_NCCL_REDOP 11 + NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START
+#define NVTX_PAYLOAD_ENTRY_NCCL_REDOP 14 + NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START
 
 extern const nvtxDomainHandle_t ncclNvtxDomainHandle;
 
 struct nccl_domain{static constexpr char const* name{"NCCL"};};
 
 class payload_schema {
- public:
+public:
   explicit payload_schema(const nvtxPayloadSchemaEntry_t entries[], size_t numEntries, const uint64_t schemaId, const char* schemaName = nullptr) noexcept
   {
     schema_attr.name = schemaName;
@@ -59,7 +62,7 @@ class payload_schema {
   payload_schema(payload_schema&&) = default;
   payload_schema& operator=(payload_schema&&) = default;
 
- private:
+private:
   nvtxPayloadSchemaAttr_t schema_attr{
     NVTX_PAYLOAD_SCHEMA_ATTR_TYPE |
     NVTX_PAYLOAD_SCHEMA_ATTR_ENTRIES |

src/include/p2p.h

@@ -34,11 +34,36 @@ typedef union {
   // Legacy CUDA IPC
   cudaIpcMemHandle_t devIpc;
   // cuMem API support
-  ncclCuDesc cuDesc;
+  struct {
+    ncclCuDesc cuDesc;
+    CUmemGenericAllocationHandle memHandle;
+  };
 } ncclIpcDesc;
 
-ncclResult_t ncclP2pAllocateShareableBuffer(size_t size, ncclIpcDesc *ipcDesc, void **ptr);
+enum ncclIpcRegType {
+  NCCL_IPC_SENDRECV = 0,
+  NCCL_IPC_COLLECTIVE = 1
+};
+
+struct ncclIpcImpInfo {
+  void* rmtRegAddr;
+  bool legacyIpcCap;
+  uintptr_t offset;
+};
+
+struct ncclIpcRegInfo {
+  int peerRank;
+  void* baseAddr;
+  struct ncclProxyConnector* ipcProxyconn;
+  struct ncclIpcImpInfo impInfo;
+};
+
+ncclResult_t ncclP2pAllocateShareableBuffer(size_t size, int directMap, ncclIpcDesc *ipcDesc, void **ptr);
 ncclResult_t ncclP2pFreeShareableBuffer(ncclIpcDesc *ipcDesc);
-ncclResult_t ncclP2pImportShareableBuffer(struct ncclComm *comm, int tpPeer, size_t size, ncclIpcDesc *ipcDesc, void **devMemPtr);
+ncclResult_t ncclP2pImportShareableBuffer(struct ncclComm *comm, int peer, size_t size, ncclIpcDesc *ipcDesc, void **devMemPtr);
+ncclResult_t ncclIpcLocalRegisterBuffer(ncclComm* comm, const void* userbuff, size_t buffSize, int* peerRanks, int nPeers, ncclIpcRegType type, int* regBufFlag, uintptr_t* offsetOut, uintptr_t** peerRmtAddrsOut);
+ncclResult_t ncclIpcGraphRegisterBuffer(ncclComm* comm, const void* userbuff, size_t buffSize, int* peerRanks, int nPeers, ncclIpcRegType type, int* regBufFlag, uintptr_t* offsetOut, uintptr_t** peerRmtAddrsOut, void* cleanupQueuePtr, int* nCleanupQueueElts);
+
+ncclResult_t ncclIpcDeregBuffer(struct ncclComm* comm, struct ncclIpcRegInfo* regInfo);
 
 #endif

src/include/profiler.h

@@ -4,34 +4,52 @@
  * See LICENSE.txt for license information
  ************************************************************************/
 
-#ifndef NCCL_PROFILER_H_
-#define NCCL_PROFILER_H_
+#ifndef PROFILER_H_
+#define PROFILER_H_
 
-#include "proxy.h"
+#include <cuda_runtime.h>
+#include "nccl_profiler.h"
 
-enum ncclProxyProfileState {
-  ncclProxyProfileBegin = 0,
+struct ncclProxyArgs;
+struct ncclKernelPlan;
+struct ncclTaskColl;
+struct ncclTaskP2p;
+struct ncclInfo;
+struct ncclComm;
+struct ncclProxyOp;
 
-  ncclProxyProfileSendGPUWait = 1,
-  ncclProxyProfileSendWait = 2,
+// Plugin Init/Finalize Wrappers
+ncclResult_t ncclProfilerPluginInit(struct ncclComm* comm);
+ncclResult_t ncclProfilerPluginFinalize(struct ncclComm* comm);
 
-  ncclProxyProfileRecvWait = 1,
-  ncclProxyProfileRecvFlushWait = 2,
-  ncclProxyProfileRecvGPUWait = 3,
+// Profiler Start/Stop Group Wrappers
+ncclResult_t ncclProfilerStartGroupEvent(struct ncclKernelPlan* plan);
+ncclResult_t ncclProfilerStopGroupEvent(struct ncclKernelPlan* plan);
 
-  ncclProxyProfileEnd = 4,
+// Profiler Start/Stop Task Events Wrappers
+ncclResult_t ncclProfilerStartTaskEvents(struct ncclKernelPlan* plan);
+ncclResult_t ncclProfilerStopTaskEvents(struct ncclKernelPlan* plan);
 
-  ncclProxyProfileSleep = 8,
-  ncclProxyProfileWakeup = 9,
+// Proxy Op Start/Stop Event Wrappers
+ncclResult_t ncclProfilerStartSendProxyOpEvent(int sub, struct ncclProxyArgs* args);
+ncclResult_t ncclProfilerStartRecvProxyOpEvent(int sub, struct ncclProxyArgs* args);
+ncclResult_t ncclProfilerStopProxyOpEvent(int sub, struct ncclProxyArgs* args);
 
-  ncclProxyProfileIdle = 16,
-  ncclProxyProfileActive = 17,
+// Proxy Step Start/Stop Event Wrappers
+ncclResult_t ncclProfilerStartSendProxyStepEvents(int sub, struct ncclProxyArgs* args, uint64_t stepLo, uint64_t stepHi);
+ncclResult_t ncclProfilerStartRecvProxyStepEvents(int sub, struct ncclProxyArgs* args, uint64_t stepLo, uint64_t stepHi);
+ncclResult_t ncclProfilerStopProxyStepEvents(int sub, struct ncclProxyArgs* args, uint64_t stepLo, uint64_t stepHi);
 
-  ncclProxyProfileAppend = 24,
-  ncclProxyProfileAppendEnd = 25
-};
+// Proxy Control Start/Stop Events Wrappers
+ncclResult_t ncclProfilerStartProxyCtrlEvent(void* profilerContext, void** eHandle);
+ncclResult_t ncclProfilerStopProxyCtrlEvent(void* eHandle);
 
-ncclResult_t ncclProfilingRecord(struct ncclProxyArgs* args, int sub, int step, int state);
-void ncclProfilingDump();
+// Record Event Wrappers
+ncclResult_t ncclProfilerRecordProxyOpEventState(int sub, struct ncclProxyArgs* args, int steps, size_t transSize, ncclProfilerEventState_t eState);
+ncclResult_t ncclProfilerRecordProxyStepEventStates(int sub, struct ncclProxyArgs* args, uint64_t stepLo, uint64_t stepHi, ncclProfilerEventState_t eState);
+ncclResult_t ncclProfilerRecordProxyCtrlEventState(void*eHandle, int appended, ncclProfilerEventState_t eState);
+
+// Profiler utility functions
+ncclResult_t ncclProfilerAddPidToProxyOp(struct ncclProxyOp* op);
 
 #endif

src/include/proxy.h

@@ -15,7 +15,7 @@
 #include "ipcsocket.h"
 #include "nccl_net.h"
 #include <pthread.h>
-#include "shm.h"
+#include "shmutils.h"
 #include "p2p.h"
 
 typedef enum : uint8_t {
@@ -30,6 +30,8 @@ typedef enum : uint8_t {
   ncclPatternCollnetDirect,
   ncclPatternNvls,
   ncclPatternNvlsTree,
+  ncclPatternPatUp,
+  ncclPatternPatDown,
   ncclPatternSend,
   ncclPatternRecv
 } ncclPattern_t;
@@ -79,6 +81,19 @@ struct ncclProxyOp {
 
   union ncclProxyOpSpecifics specifics;
 
+  // Profiler plugin
+  union {
+    struct ncclTaskColl* coll;
+    struct ncclTaskP2p* p2p;
+  } task;
+
+  int eActivationMask;
+  void* taskEventHandle;
+  int rank;
+  int peer;
+  pid_t pid;
+  void* profilerContext;
+
   struct ncclProxyOp *enqNext;
 };
 
@@ -107,7 +122,15 @@ struct ncclProxySubArgs {
   uint64_t done;
   uint64_t end;
   void* requests[NCCL_STEPS];
-  void* profilingEvents[NCCL_STEPS];
+
+  // Profiler plugin
+  int eActivationMask;
+  int rank;
+  void* taskEventHandle;
+  void* opEventHandle;
+  void* stepEventHandles[NCCL_STEPS];
+  size_t transSize;
+
   void* recvRequestsCache[NCCL_STEPS];
   int recvRequestsSubCount;
 
@@ -142,6 +165,10 @@ struct ncclProxyArgs {
   int idle;
   uint64_t hdp_flushed;
 
+  // Profiler plugin
+  pid_t pid;
+  void* profilerContext;
+
   // Element linking
   struct ncclProxyArgs* next;
   struct ncclProxyArgs* nextPeer;
@@ -279,6 +306,7 @@ struct ncclProxyState {
   ncclNet_t* ncclNet;
   ncclCollNet_t* ncclCollNet;
   uint32_t* abortFlag;
+  bool directMode;
   // Service threads
   pthread_t thread;
   pthread_t threadUDS;
@@ -299,6 +327,9 @@ struct ncclProxyState {
   // Progress thread
   struct ncclProxyProgressState progressState;
 
+  // Profiler plugin
+  void* profilerContext;
+
   // Queue of expected responses from the proxy
   struct ncclExpectedProxyResponse* expectedResponses;
 };
@@ -350,8 +381,9 @@ enum ncclProxyMsgType {
   ncclProxyMsgAbort = 7,
   ncclProxyMsgStop = 8,
   ncclProxyMsgGetFd = 9, // cuMem API support (UDS)
-  ncclProxyMsgRegister = 10,
-  ncclProxyMsgDeregister = 11
+  ncclProxyMsgQueryFd = 10,
+  ncclProxyMsgRegister = 11,
+  ncclProxyMsgDeregister = 12
 };
 
 // This function is called by a client of the proxy that needs to invoke any of the non-progress proxyOp types
@@ -365,6 +397,7 @@ ncclResult_t ncclPollProxyResponse(struct ncclComm* comm, struct ncclProxyConnec
 
 // UDS support
 ncclResult_t ncclProxyClientGetFdBlocking(struct ncclComm* comm, int rank, void *handle, int* convertedFd);
+ncclResult_t ncclProxyClientQueryFdBlocking(struct ncclComm* comm, struct ncclProxyConnector* proxyConn, int localFd, int* rmtFd);
 
 ncclResult_t ncclProxyStop(struct ncclComm* comm);
 ncclResult_t ncclProxyShmUnlink(struct ncclComm* comm);

src/include/register.h

@@ -11,7 +11,13 @@ enum {
   NVLS_REG_COMPLETE = 0x02,
   NVLS_REG_POSSIBLE = 0x04,
   NVLS_REG_NO_SUPPORT = 0x08,
-  COLLNET_REG_COMPLETE = 0x10
+  COLLNET_REG_COMPLETE = 0x10,
+  IPC_REG_COMPLETE = 0x20
+};
+
+struct ncclPeerRegIpcAddr {
+  uintptr_t* devPeerRmtAddrs;
+  uintptr_t* hostPeerRmtAddrs;
 };
 
 struct ncclReg {
@@ -34,7 +40,10 @@ struct ncclReg {
   uintptr_t caddrs[NCCL_MAX_LOCAL_RANKS]; /* use to check if NVLS buffers match among intra-node ranks */
   // collnet reg
   void* collnetHandle;
-  struct ncclProxyConnector* proxyconn;
+  struct ncclProxyConnector* collnetProxyconn;
+  // general ipc reg
+  struct ncclPeerRegIpcAddr regIpcAddrs;
+  struct ncclIpcRegInfo* ipcInfos[NCCL_MAX_LOCAL_RANKS];
 };
 
 struct ncclRegCache {

src/include/rocmwrap.h

@@ -15,7 +15,6 @@ typedef hsa_status_t (*PFN_hsa_system_get_info)(hsa_system_info_t attribute, voi
 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
@@ -68,6 +67,9 @@ DECLARE_ROCM_PFN_EXTERN(hsa_init);
 DECLARE_ROCM_PFN_EXTERN(hsa_system_get_info);
 DECLARE_ROCM_PFN_EXTERN(hsa_status_string);
 
+extern int ncclCuMemEnable();
+extern int ncclCuMemHostEnable();
+
 ncclResult_t rocmLibraryInit(void);
 
 extern bool ncclCudaLaunchBlocking; // initialized by ncclCudaLibraryInit()

src/include/shm.h

@@ -1,26 +1,37 @@
-/*************************************************************************
- * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
- *
- * See LICENSE.txt for license information
- ************************************************************************/
-
 #ifndef NCCL_SHM_H_
 #define NCCL_SHM_H_
 
-#include "nccl.h"
+#include "comm.h"
 
-typedef void* ncclShmHandle_t;
-ncclResult_t ncclShmOpen(char* shmPath, size_t shmSize, void** shmPtr, void** devShmPtr, int refcount, ncclShmHandle_t* handle);
-ncclResult_t ncclShmClose(ncclShmHandle_t handle);
-ncclResult_t ncclShmUnlink(ncclShmHandle_t handle);
-
-struct ncclShmemCollBuff {
-  volatile size_t *cnt[2];
-  volatile void *ptr[2];
-  int round;
-  size_t maxTypeSize;
+struct shmLegacyIpc {
+  char shmSuffix[7];
+  ncclShmHandle_t handle;
+  size_t shmSize;
 };
 
-ncclResult_t ncclShmemAllgather(struct ncclComm *comm, struct ncclShmemCollBuff *shmem, void *sendbuff, void *recvbuff, size_t typeSize);
+struct shmCuIpc {
+  union {
+    CUmemFabricHandle handle;
+    CUmemGenericAllocationHandle data;
+  };
+  int tpProxyRank;
+  void *ptr;
+  size_t size;
+};
+
+struct shmIpcDesc {
+  union
+  {
+    struct shmLegacyIpc shmli;
+    struct shmCuIpc shmci;
+  };
+  bool legacy;
+};
+
+typedef struct shmIpcDesc ncclShmIpcDesc_t;
+
+ncclResult_t ncclShmAllocateShareableBuffer(int tpProxyRank, size_t size, bool legacy, ncclShmIpcDesc_t *descOut, void **hptr, void **dptr);
+ncclResult_t ncclShmImportShareableBuffer(struct ncclComm *comm, ncclShmIpcDesc_t *desc, void **hptr, void **dptr, ncclShmIpcDesc_t *descOut);
+ncclResult_t ncclShmIpcClose(ncclShmIpcDesc_t *desc);
 
 #endif

src/include/shmutils.h

@@ -0,0 +1,26 @@
+/*************************************************************************
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef NCCL_SHMUTILS_H_
+#define NCCL_SHMUTILS_H_
+
+#include "nccl.h"
+
+typedef void* ncclShmHandle_t;
+ncclResult_t ncclShmOpen(char* shmPath, size_t shmSize, void** shmPtr, void** devShmPtr, int refcount, ncclShmHandle_t* handle);
+ncclResult_t ncclShmClose(ncclShmHandle_t handle);
+ncclResult_t ncclShmUnlink(ncclShmHandle_t handle);
+
+struct ncclShmemCollBuff {
+  volatile size_t *cnt[2];
+  volatile void *ptr[2];
+  int round;
+  size_t maxTypeSize;
+};
+
+ncclResult_t ncclShmemAllgather(struct ncclComm *comm, struct ncclShmemCollBuff *shmem, void *sendbuff, void *recvbuff, size_t typeSize);
+
+#endif

src/include/timer.h

@@ -33,15 +33,15 @@ static double startTimes[8];
 #define TIME_START(index) do { \
   counts[index]++; \
   startTimes[index] = gettime(); \
-} while (0);
+} while (0)
 
 #define TIME_STOP(index) do { \
   times[index] += gettime() - startTimes[index]; \
-} while (0);
+} while (0)
 
 #define TIME_CANCEL(index) do { \
   counts[index]--; \
-} while (0);
+} while (0)
 
 #define TIME_PRINT(name) do { \
   printf("%s stats", name); \
@@ -50,11 +50,11 @@ static double startTimes[8];
     counts[i] = 0; \
   } \
   printf("\n"); \
-} while (0);
+} while (0)
 #else
-#define TIME_START(index) while(0);
-#define TIME_STOP(index) while(0);
-#define TIME_CANCEL(index) while(0);
+#define TIME_START(index) do {} while(0)
+#define TIME_STOP(index) do {} while(0)
+#define TIME_CANCEL(index) do {} while(0)
 #define TIME_PRINT(name)
 #endif
 #endif

src/include/transport.h

@@ -36,7 +36,7 @@ struct ncclConnector;
 struct ncclComm;
 
 #define CHANNEL_MASK_OFFSET(nranks, connIndex) (nranks * (connIndex == NCCL_CONN_IDX_P2P_NET ? NCCL_CONN_IDX_P2P_NET : 0))
-#define CONNECT_SIZE 128
+#define CONNECT_SIZE 256
 struct ncclConnect {
   char data[CONNECT_SIZE];
 };
@@ -77,7 +77,6 @@ struct ncclCollNetSharedRes {
   void* resources;
   int nChannels;
   size_t buffSize;
-  int intraHighestTransportType;
 };
 
 struct ncclTransportComm {
@@ -95,13 +94,14 @@ struct ncclTransportComm {
 
 struct ncclTransport {
   const char name[8];
-  ncclResult_t (*canConnect)(int*, struct ncclTopoSystem* topo, struct ncclTopoGraph* graph, struct ncclPeerInfo*, struct ncclPeerInfo*);
+  ncclResult_t (*canConnect)(int*, struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo*, struct ncclPeerInfo*);
   struct ncclTransportComm send;
   struct ncclTransportComm recv;
 };
 
 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, bool* needsProxy=NULL);
+ncclResult_t ncclTransportCheckP2pType(struct ncclComm* comm, bool* intraNodeP2pSupport, bool* directMode);
 
 ncclResult_t ncclNvlsInit(struct ncclComm* comm);
 ncclResult_t ncclNvlsSetup(struct ncclComm* comm, struct ncclComm* parent);
@@ -113,7 +113,7 @@ ncclResult_t ncclNvlsDeregBuffer(CUmemGenericAllocationHandle *mcHandler, CUdevi
 ncclResult_t ncclNvlsFree(struct ncclComm* comm);
 
 enum { collNetRecv=0, collNetSend=1 };
-int ncclTransportCollNetSetup(struct ncclComm* comm, struct ncclTopoGraph* collNetGraph, struct ncclChannel* channel, int masterRank, int masterPeer, int collNetGraphChannelId, int type, ncclConnect* connect);
+bool ncclTransportCollNetSetup(struct ncclComm* comm, struct ncclTopoGraph* collNetGraph, struct ncclChannel* channel, int masterRank, int masterPeer, int collNetGraphChannelId, int type, ncclConnect* connect);
 ncclResult_t ncclTransportCollNetCheck(struct ncclComm* comm, int collNetSetupFail);
 ncclResult_t ncclTransportCollNetFree(struct ncclComm* comm);
 ncclResult_t ncclCollnetLocalRegisterBuffer(struct ncclComm* comm, const void* userbuff, size_t buffSize, int type, int* outRegBufUsed, void** outHandle);
@@ -122,6 +122,7 @@ ncclResult_t ncclCollnetDeregBuffer(struct ncclComm* comm, struct ncclProxyConne
 
 ncclResult_t ncclTransportRingConnect(struct ncclComm* comm);
 ncclResult_t ncclTransportTreeConnect(struct ncclComm* comm);
+ncclResult_t ncclTransportPatConnect(struct ncclComm* comm);
 
 ncclResult_t ncclCollNetSetup(ncclComm_t comm, ncclComm_t parent, struct ncclTopoGraph* graphs[]);
 ncclResult_t ncclCollNetChainBufferSetup(ncclComm_t comm);

src/include/utils.h

@@ -27,7 +27,6 @@ ncclResult_t busIdToInt64(const char* busId, int64_t* id);
 ncclResult_t getBusId(int cudaDev, int64_t *busId);
 
 ncclResult_t getHostName(char* hostname, int maxlen, const char delim);
-uint64_t getHash(const char* string, int n);
 uint64_t getHostHash();
 uint64_t getPidHash();
 ncclResult_t getRandomData(void* buffer, size_t bytes);