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

[ROCm/rccl commit: 3a919c1f49]

projects/rccl/src/enqueue.cc

@@ -12,81 +12,61 @@
 #include <hip/hip_runtime.h>
 #include <hip/hip_ext.h>
 #include "gdrwrap.h"
+#include "bootstrap.h"
+#include <cstring>
+
+#include <cstring> // std::memcpy
 
 // Only generate inline kernels for LL
-#define NCCL_FUNC5(func, algo, redop, dtype) \
-  NCCL_KERN_NAME(func, algo, LL, redop, dtype), \
-  NCCL_KERN_NAME(func, algo, LL, redop, dtype), \
-  NCCL_KERN_NAME(func, algo, LL, redop, dtype)
+#define NCCL_FUNC5(func, algo, devredop, dtype) \
+  (void*)NCCL_KERN_NAME(func, algo, LL, devredop, dtype), \
+  (void*)NCCL_KERN_NAME(func, algo, LL, devredop, dtype), \
+  (void*)NCCL_KERN_NAME(func, algo, LL, devredop, dtype)
 
-#define NCCL_FUNC4(func, redop, type) \
-  NCCL_FUNC5(func, TREE,    redop, type), \
-  NCCL_FUNC5(func, RING,    redop, type), \
-  NCCL_FUNC5(func, COLLNET, redop, type)
+#define NCCL_FUNC4(func, devredop, type) \
+  (void*)NCCL_FUNC5(func, TREE,    devredop, type), \
+  (void*)NCCL_FUNC5(func, RING,    devredop, type), \
+  (void*)NCCL_FUNC5(func, COLLNET, devredop, type)
 
-#if defined(__CUDA_BF16_TYPES_EXIST__)
 // Must be consistent with ncclDataType_t
-#define NCCL_FUNCS3A(func, redop) \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, uint8_t), \
-  (void*)NCCL_FUNC4(func, redop, int32_t), \
-  (void*)NCCL_FUNC4(func, redop, uint32_t), \
-  (void*)NCCL_FUNC4(func, redop, int64_t), \
-  (void*)NCCL_FUNC4(func, redop, uint64_t), \
-  (void*)NCCL_FUNC4(func, redop, half), \
-  (void*)NCCL_FUNC4(func, redop, float), \
-  (void*)NCCL_FUNC4(func, redop, double), \
-  (void*)NCCL_FUNC4(func, redop, __nv_bfloat16)
-#define NCCL_FUNCS3B(func, redop) \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, int8_t), \
-  (void*)NCCL_FUNC4(func, redop, int8_t)
-#else
-// Must be consistent with ncclDataType_t
-#define NCCL_FUNCS3A(func, redop) \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, uint8_t), \
-  NCCL_FUNC4(func, redop, int32_t), \
-  NCCL_FUNC4(func, redop, uint32_t), \
-  NCCL_FUNC4(func, redop, int64_t), \
-  NCCL_FUNC4(func, redop, uint64_t), \
-  NCCL_FUNC4(func, redop, half), \
-  NCCL_FUNC4(func, redop, float), \
-  NCCL_FUNC4(func, redop, double), \
-  NCCL_FUNC4(func, redop, rccl_bfloat16)
-#define NCCL_FUNCS3B(func, redop) \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, int8_t), \
-  NCCL_FUNC4(func, redop, int8_t)
-#endif
+#define NCCL_FUNCS3A(func, devredop) \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, uint8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int32_t), \
+  (void*)NCCL_FUNC4(func, devredop, uint32_t), \
+  (void*)NCCL_FUNC4(func, devredop, int64_t), \
+  (void*)NCCL_FUNC4(func, devredop, uint64_t), \
+  (void*)NCCL_FUNC4(func, devredop, half), \
+  (void*)NCCL_FUNC4(func, devredop, float), \
+  (void*)NCCL_FUNC4(func, devredop, double), \
+  (void*)NCCL_FUNC4(func, devredop, rccl_bfloat16)
+#define NCCL_FUNCS3B(func, devredop) \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int8_t), \
+  (void*)NCCL_FUNC4(func, devredop, int8_t)
 
-// Must be consistent with ncclRedOp_t -- but we only generate kernel for sums.
+// Must be consistent with ncclDevRedOp_t -- but we only generate kernel for sums.
 #define NCCL_FUNCS2A(func) \
-  NCCL_FUNCS3A(func, Sum), \
-  NCCL_FUNCS3A(func, Sum), \
-  NCCL_FUNCS3A(func, Sum), \
-  NCCL_FUNCS3A(func, Sum), \
-  NCCL_FUNCS3A(func, Sum)
+  NCCL_FUNCS3A(func, Sum), /*Sum*/ \
+  NCCL_FUNCS3A(func, Sum), /*Prod*/ \
+  NCCL_FUNCS3A(func, Sum), /*Max*/ \
+  NCCL_FUNCS3A(func, Sum), /*Min*/ \
+  NCCL_FUNCS3A(func, Sum), /*PreMulSum*/ \
+  NCCL_FUNCS3A(func, Sum)  /*SumPostDiv*/
 #define NCCL_FUNCS2B(func) \
-  NCCL_FUNCS3B(func, Sum), \
-  NCCL_FUNCS3B(func, Sum), \
-  NCCL_FUNCS3A(func, Sum), \
-  NCCL_FUNCS3B(func, Sum), \
-  NCCL_FUNCS3B(func, Sum)
+  NCCL_FUNCS3B(func, Sum), /*Sum*/ \
+  NCCL_FUNCS3B(func, Sum), /*Prod*/ \
+  NCCL_FUNCS3B(func, Sum), /*Max*/ \
+  NCCL_FUNCS3B(func, Sum), /*Min*/ \
+  NCCL_FUNCS3B(func, Sum), /*PreMulSum*/ \
+  NCCL_FUNCS3B(func, Sum)  /*SumPostDiv*/
 
 typedef void(*ncclKern_t)(struct ncclWorkElem first);
 // Must be consistent with the ncclFuncSet enum
@@ -145,7 +125,6 @@ static ncclResult_t getNextOp(struct ncclChannel* channel, struct ncclWork** wor
   // Initialize with work elem if provided
   if (base) memcpy(e, base, sizeof(struct ncclWorkElem));
   e->active = 1;
-  e->index = opIndex;
   channel->workFifoTail++;
   channel->workCount++;
   if (work) *work = w;
@@ -183,10 +162,11 @@ static ncclResult_t setupLaunch(struct ncclQueueInfo* eqInfo, int usingCudaGraph
 
     if (c == 0) {
       // As we inline the first coll directly, we can free it immediately.
-      // Except P2P or aggregation cases
+      // Except P2P or aggregation or registration cases
       struct ncclWork* work = channel->workFifo+((channel->workFifoTail-channel->workCount)%NCCL_MAX_OPS);
       struct ncclWorkElem* elem = work->elems;
-      if (elem->funcIndex != FUNC_INDEX_P2P && eqInfo->elemList->count() == 1) elem->active = 0;
+      if (elem->funcIndex != FUNC_INDEX_P2P && eqInfo->elemList->count() == 1 && elem->regUsed == 0)
+        elem->active = 0;
     }
 
     if (channel->gdrMemDesc) {
@@ -371,7 +351,7 @@ RCCL_PARAM(SharpThreshold, "SHARP_THRESHOLD", 16384);
 
 static inline ncclResult_t getCollNetSupport(struct ncclInfo* info, int* collNetTypeSupport) {
   if (info->comm->collNetSupport > 0 && info->nBytes < rcclParamSharpThreshold()) {
-    ncclRedOp_t netOp = info->op == ncclAvg ? ncclSum : info->op;
+    ncclRedOp_t netOp = info->op == ncclAvg || info->op >= ncclNumOps ? ncclSum : info->op;
     NCCLCHECK(collNetReduceSupport(info->datatype, netOp, collNetTypeSupport));
   } else {
     *collNetTypeSupport = 0;
@@ -381,30 +361,35 @@ static inline ncclResult_t getCollNetSupport(struct ncclInfo* info, int* collNet
 
 static ncclResult_t getAlgoInfo(struct ncclInfo* info, int collNetTypeSupport, int numPipeOps) {
   struct ncclComm* comm = info->comm;
-  float minTime = 3600000000.0; // Hopefully no operation will take an hour to complete.
-  // Find algorithm / protocol.
-  info->algorithm = -1;
-  info->protocol = -1;
-  if (comm->nRanks == 1) return ncclSuccess;
-  int nAlgos = NCCL_NUM_ALGORITHMS;
-  for (int a=0; a<nAlgos; a++) {
-    if (a == NCCL_ALGO_COLLNET && collNetTypeSupport != 1) continue;
-    for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
-      float time;
-      NCCLCHECK(ncclTopoGetAlgoTime(info, a, p, numPipeOps, &time));
-      if (time >= 0 && time < minTime) {
-        info->algorithm = a;
-        info->protocol = p;
-        minTime = time;
+  if (comm->nRanks == 1) {
+    info->algorithm = NCCL_ALGO_RING;
+    info->protocol = NCCL_PROTO_SIMPLE;
+  }
+  else {
+    float minTime = 3600000000.0; // Hopefully no operation will take an hour to complete.
+    // Find algorithm / protocol.
+    info->algorithm = -1;
+    info->protocol = -1;
+    int nAlgos = NCCL_NUM_ALGORITHMS;
+    for (int a=0; a<nAlgos; a++) {
+      if (a == NCCL_ALGO_COLLNET && collNetTypeSupport != 1) continue;
+      for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
+        float time;
+        NCCLCHECK(ncclTopoGetAlgoTime(info, a, p, numPipeOps, &time));
+        if (time >= 0 && time < minTime) {
+          info->algorithm = a;
+          info->protocol = p;
+          minTime = time;
+        }
       }
     }
+    if (info->algorithm == -1 || info->protocol == -1) {
+      WARN("Error : no algorithm/protocol available");
+      return ncclInternalError;
+    }
+    //if (comm->rank == 0) INFO(NCCL_TUNING, "%ld Bytes -> Algo %d proto %d time %f", info->nBytes, info->algorithm, info->protocol, minTime);
+    TRACE(NCCL_COLL, "%ld Bytes -> Algo %d proto %d time %f", info->nBytes, info->algorithm, info->protocol, minTime);
   }
-  if (info->algorithm == -1 || info->protocol == -1) {
-    WARN("Error : no algorithm/protocol available");
-    return ncclInternalError;
-  }
-  //if (comm->rank == 0) INFO(NCCL_TUNING, "%ld Bytes -> Algo %d proto %d time %f", info->nBytes, info->algorithm, info->protocol, minTime);
-  TRACE(NCCL_COLL, "%ld Bytes -> Algo %d proto %d time %f", info->nBytes, info->algorithm, info->protocol, minTime);
 
   int nc = (info->nChannels > 0) ? info->nChannels : comm->nChannels;
   int nt = comm->maxThreads[info->algorithm][info->protocol];
@@ -498,15 +483,23 @@ comp_next:
   NCCLCHECK(getPatternInfo(info));
   NCCLCHECK(getLoopInfo(info));
 
-  work->op.opCount = info->comm->collOpCount;
+  work->coll.opCount = info->comm->collOpCount;
   work->sendbuff = info->sendbuff;
   work->recvbuff = info->recvbuff;
   work->coll.root = info->root;
   work->coll.count = info->count;
   work->coll.nChannels = info->nChannels;
   work->nThreads = info->nThreads;
+  work->coll.redOpArg = info->opFull.scalarArg;
+  work->redOpArgIsPtr = info->opFull.scalarArgIsPtr;
 
-  work->funcIndex = FUNC_INDEX(info->coll, info->op, info->datatype, info->algorithm, info->protocol);
+  if (info->comm->nRanks == 1) {
+    // one-rank reduce index
+    work->funcIndex = FUNC_INDEX_P2P - ncclNumTypes + int(info->datatype);
+    return ncclSuccess;
+  }
+
+  work->funcIndex = FUNC_INDEX(info->coll, info->opFull.op, info->datatype, info->algorithm, info->protocol);
 
   work->coll.connIndex = 0;
   proxyArgs->connIndex = 0;
@@ -533,7 +526,7 @@ comp_next:
                                                                info->comm->nChannels,
                                                                &work->clique.nChannels));
       work->clique.count = info->count;
-      work->funcIndex = FUNC_INDEX(info->coll, info->op, info->datatype, info->algorithm, info->protocol);
+      work->funcIndex = FUNC_INDEX(info->coll, info->opFull.op, info->datatype, info->algorithm, info->protocol);
 
       // Setup pointers to where all the input/output pointers will be
       NCCLCHECK(info->comm->cliqueManager->WaitForPointers(work));
@@ -563,6 +556,8 @@ comp_next:
     while (info->nBytes / (info->nChannels*info->comm->channels[0].collTree.nHeads*chunkSize) < info->comm->channels[0].collTree.depth*8 && chunkSize > 32768) chunkSize /= 2;
     // Use lastChunkSize as chunkSize
     work->coll.lastChunkSize = chunkSize / ncclTypeSize(info->datatype);
+    // Set direct direction for broadcast-gather (read or write)
+    work->direct = (info->nBytes / info->nChannels <= 1024*1024) ? NCCL_DIRECT_WRITE : NCCL_DIRECT_READ;
   } else if (info->protocol == NCCL_PROTO_LL) {
     const ssize_t sliceSize = stepSize*sizeof(uint64_t)/sizeof(union ncclLLFifoLine);
     const ssize_t loopSize = info->nChannels*info->nchunksPerLoop*(ssize_t)sliceSize;
@@ -592,7 +587,7 @@ comp_next:
   proxyArgs->protocol = info->protocol;
   proxyArgs->dtype = info->datatype;
   proxyArgs->redOp = info->algorithm != NCCL_ALGO_COLLNET ? ncclNumOps : // Only set redOp when using CollNet
-                     info->op == ncclAvg ? ncclSum : // Network sees avg as sum
+                     info->opFull.op==ncclDevPreMulSum || info->opFull.op==ncclDevSumPostDiv ? ncclSum : // Network sees avg as sum
                      info->op;
   proxyArgs->pattern = info->pattern;
   proxyArgs->root = info->root;
@@ -618,12 +613,61 @@ static ncclResult_t checkSetStream(struct ncclInfo* info) {
   return ncclSuccess;
 }
 
+struct ncclBuffRegHandle {
+  hipIpcMemHandle_t sendBuffIpc;
+  hipIpcMemHandle_t recvBuffIpc;
+  ssize_t sendBuffOffset;
+  ssize_t recvBuffOffset;
+};
+
+// Register input and output buffers
+// Exchange with ranks on the same host
+static ncclResult_t ncclRegBuffAndExchange(struct ncclInfo* info, struct ncclBuffRegInfo* regInfo) {
+  ncclComm_t comm = info->comm;
+  if (comm->localRanks == 1) return ncclSuccess;
+  if (comm->pfnCuMemGetAddressRange == NULL) return ncclSuccess;  // CUDA toolkit or driver version too old
+
+  struct ncclBuffRegHandle regHandles[NCCL_MAX_INTRA_RANKS];
+  // Get IPC handles
+  // Note: the handle only corresponds to the base address of the allocation
+  CUDACHECK(hipIpcGetMemHandle(&regHandles[comm->intraNodeRank].sendBuffIpc, (void*)info->sendbuff));
+  CUDACHECK(hipIpcGetMemHandle(&regHandles[comm->intraNodeRank].recvBuffIpc, (void*)info->recvbuff));
+  // Get offset of user buffer within allocation
+  void* baseAddr;
+  size_t size;
+  CUDACHECK(comm->pfnCuMemGetAddressRange(&baseAddr, &size, (void*)info->sendbuff));
+  regHandles[comm->intraNodeRank].sendBuffOffset = (char*)info->sendbuff - (char*)baseAddr;
+  CUDACHECK(comm->pfnCuMemGetAddressRange(&baseAddr, &size, (void*)info->recvbuff));
+  regHandles[comm->intraNodeRank].recvBuffOffset = (char*)info->recvbuff - (char*)baseAddr;
+  TRACE(NCCL_COLL, "Base %p size %lu offset %ld", baseAddr, size, regHandles[comm->intraNodeRank].recvBuffOffset);
+
+  // Exchange handles within node
+  NCCLCHECK(bootstrapIntraNodeAllGather(comm->bootstrap, comm->intraNodeGlobalRanks, comm->intraNodeRank, comm->localRanks, regHandles, sizeof(struct ncclBuffRegHandle)));
+  // Open handles at local process
+  for (int i=0; i<comm->localRanks; i++) {
+    if (i == comm->intraNodeRank) {
+      regInfo->sendbuffsBase[i] = regInfo->recvbuffsBase[i] = NULL;
+      continue;
+    }
+    CUDACHECK(hipIpcOpenMemHandle(regInfo->sendbuffsBase+i, regHandles[i].sendBuffIpc, hipIpcMemLazyEnablePeerAccess));
+    CUDACHECK(hipIpcOpenMemHandle(regInfo->recvbuffsBase+i, regHandles[i].recvBuffIpc, hipIpcMemLazyEnablePeerAccess));
+    // Get real address of buffer
+    regInfo->sendbuffs[i] = (char*)regInfo->sendbuffsBase[i] + regHandles[i].sendBuffOffset;
+    regInfo->recvbuffs[i] = (char*)regInfo->recvbuffsBase[i] + regHandles[i].recvBuffOffset;
+  }
+  regInfo->nBuffs = comm->localRanks;
+  TRACE(NCCL_COLL, "Rank %d exchanged %d buffers", comm->rank, regInfo->nBuffs);
+  return ncclSuccess;
+}
+
 // Compute enqueue element, save it in list
 // Compute CUDA launch parameters
 // Capture time code in view of CUDA graph
 static ncclResult_t ncclSetupCollKernel(struct ncclInfo* info) {
   ncclComm_t comm = info->comm;
-  if (comm->nRanks == 1) {
+  if (comm->nRanks == 1 &&
+      // User-defined reduction ops may need alter the data even for unitary reductions
+      info->op < ncclNumOps) {
     if (info->sendbuff != info->recvbuff)
       CUDACHECK(hipMemcpyAsync(info->recvbuff, info->sendbuff, info->nBytes, hipMemcpyDeviceToDevice, info->stream));
     return ncclSuccess;
@@ -651,6 +695,19 @@ static ncclResult_t ncclSetupCollKernel(struct ncclInfo* info) {
     comm->args.active = 2;    // I am so far the last element; may be changed later in aggregation mode
   }
 
+  // Register and exchange input and output buffers
+  if (comm->usingCudaGraph &&                   // only in CUDA graph mode
+      comm->graphRegister == 1 &&               // when registration is enabled
+      info->algorithm == NCCL_ALGO_COLLNET &&   // limited to CollNet for now
+      comm->intraHighestTransportType == TRANSPORT_P2P && // only when all ranks can p2p each other
+      comm->intraRanks == 1) {                  // only in multi-process mode
+    NCCLCHECK(ncclRegBuffAndExchange(info, &eqElem->buffRegInfo));
+    // Disable inline argument because we need kernel to copy the entire ncclWork from workFifo
+    // because the registered addresses are in ncclWork
+    if (eqElem->buffRegInfo.nBuffs > 0) comm->args.active = 0;
+    comm->enqueueInfo->nRegBuffs += eqElem->buffRegInfo.nBuffs;
+  }
+
   return ncclSuccess;
 }
 
@@ -667,41 +724,15 @@ static inline int findShortestChannel(ncclComm_t comm) {
   return minC;
 }
 
-static inline ncclResult_t getNextChannel(ncclComm_t comm, int* nextChannel) {
-  if (comm->asyncAllocMode == ncclComm::SHORTEST_QUEUE) {
-    *nextChannel = findShortestChannel(comm);
+static inline int getNextChannel(ncclComm_t comm, int aggMode) {
+  int nextChannel = 0;
+  if (aggMode && comm->asyncAllocMode == ncclComm::SHORTEST_QUEUE) {
+    nextChannel = findShortestChannel(comm);
   } else {
-    *nextChannel = comm->lastChannel % comm->nChannels;
+    nextChannel = comm->lastChannel % comm->nChannels;
     comm->lastChannel++;
   }
-  return ncclSuccess;
-}
-
-// Dynamic enqueue code
-static ncclResult_t ncclEnqueueCollKernel(ncclComm_t comm, struct ncclQueueElem* eqElem) {
-  struct ncclWorkElem* work = &eqElem->work;
-  struct ncclProxyArgs* proxyArgs = &eqElem->proxyArgs;
-
-  int nChannels = work->coll.nChannels;
-  for (int bid=0; bid<nChannels; bid++) {
-    int channelId = comm->lastChannel % comm->nChannels;
-    struct ncclChannel* channel = comm->channels+channelId;
-
-    // Proxy
-    proxyArgs->subs[0].channel = channel;
-    proxyArgs->opCount = comm->collOpCount;
-    proxyArgs->commOpCount = comm->opCount;
-
-    if (proxyArgs->subs[0].nsteps) NCCLCHECK(ncclProxySaveColl(proxyArgs, comm->nRanks));
-
-    comm->lastChannel++;
-    work->coll.bid = bid % nChannels;
-    NCCLCHECK(getNextOp(channel, NULL, work));
-    //INFO(NCCL_COLL, "Host enqueue: bid %d channel %d index %ld nThreads %d funcIndex %d count %ld nChannels %d",
-    //      work->coll.bid, channelId, channel->workFifoTail, work->nThreads, work->funcIndex, work->coll.count, work->coll.nChannels);
-  }
-  comm->collOpCount++;
-  return ncclSuccess;
+  return nextChannel;
 }
 
 ncclResult_t ncclSetupAsyncKernels(ncclComm_t comm) {
@@ -733,7 +764,7 @@ ncclResult_t ncclSetupAsyncKernels(ncclComm_t comm) {
       channelUsed += info->nChannels;
       // We can use fast path if all collectives are the same
       homogeneous &= info->coll == comm->asyncOps[0].coll &&
-                     info->op == comm->asyncOps[0].op &&
+                     info->opFull.op == comm->asyncOps[0].opFull.op &&
                      info->datatype == comm->asyncOps[0].datatype;
       if (allCollNetSupport > 0) NCCLCHECK(getCollNetSupport(info, &allCollNetSupport));
     }
@@ -818,13 +849,22 @@ static ncclResult_t ncclSaveP2p(struct ncclInfo* info) {
   return ncclSuccess;
 }
 
-enum { COLL_SEGMENT=0, P2P_SEGMENT=1 };
+enum { RingTree_Segment=0, P2P_Segment=1, CollNet_Segment=2 };
 static int getSegment(int type, int delta, struct ncclWork* work) {
-  if (type == P2P_SEGMENT) {  // P2P
+  // Current ncclWork is full
+  if (work->elems[NCCL_MAX_WORK_ELEMENTS-1].active != 0) return -1;
+
+  if (type == P2P_Segment) {  // P2P
+    // Do not mix P2P and collective ops
+    if (work->elems[0].funcIndex != FUNC_INDEX_P2P) return -1;
     for (int s=0; s<NCCL_MAX_WORK_ELEMENTS && work->elems[s].p2p.delta != delta; s++) {
       if (work->elems[s].active == 0) return s;
     }
-  } else { // aggregation
+  } else if (type == CollNet_Segment) { // CollNet
+    for (int s=0; s<NCCL_MAX_WORK_ELEMENTS; s+=NCCL_REG_ELEM_FACTOR) {
+      if (work->elems[s].active == 0) return s;
+    }
+  } else {  // Ring or Tree
     for (int s=0; s<NCCL_MAX_WORK_ELEMENTS; s++) {
       if (work->elems[s].active == 0) return s;
     }
@@ -838,7 +878,7 @@ static ncclResult_t computeP2pWorkElem(struct ncclInfo* info /* input */, struct
   elem->nThreads = NCCL_MAX_NTHREADS;
   elem->sendbuff = info->sendbuff;
   elem->recvbuff = info->recvbuff;
-  elem->op.opCount = info->comm->p2pOpCount;
+  elem->p2p.opCount = info->comm->p2pOpCount;
   elem->p2p.sendCount = info->sendbytes;
   elem->p2p.recvCount = info->recvbytes;
   elem->p2p.sendChunkSize = info->sendChunkSize;
@@ -847,13 +887,14 @@ static ncclResult_t computeP2pWorkElem(struct ncclInfo* info /* input */, struct
   return ncclSuccess;
 }
 
-static ncclResult_t enqueueSegOp(int type, struct ncclWorkElem* elem /* input */, struct ncclWork* work, int s) {
+static ncclResult_t enqueueSegOp(int type, struct ncclWorkElem* elem /* input */, struct ncclWork* work, int s,
+    struct ncclBuffRegInfo* regInfo, struct ncclChannel* channel, struct ncclComm* comm) {
   // Copy element into corresponding segment of ncclWork
   memcpy(work->elems+s, elem, sizeof(struct ncclWorkElem));
   work->elems[s].active = 1;
 
   // Determine nThreads at dynamic time
-  if (type == P2P_SEGMENT) {
+  if (type == P2P_Segment) {
     const int nsegments = s+1;
     int nThreads = 512;
     while (nsegments*nThreads > 256) nThreads /= 2;
@@ -861,6 +902,33 @@ static ncclResult_t enqueueSegOp(int type, struct ncclWorkElem* elem /* input */
     for (int i=0; i<nsegments; i++) work->elems[i].p2p.nThreads = nThreads;
   }
 
+  // Copy registered buffer addresses into ncclWork
+  if (regInfo->nBuffs > 0) {
+    struct ncclWorkRegElem* regElem = (struct ncclWorkRegElem*)(work->elems+s);
+    // For CollNet
+    for (int i=0; i<NCCL_MAX_DIRECT_ARITY; i++) {
+      int peer = channel->collTree.down[i];
+      if (peer == -1) break;
+      int j = comm->rankToIntraNodeRank[peer];
+      if (j < 0) {
+        WARN("Invalid intra-node rank %d for peer %d", j, peer);
+        return ncclInternalError;
+      }
+      regElem->dnInputs[i] = regInfo->sendbuffs[j];
+      regElem->dnOutputs[i] = regInfo->recvbuffs[j];
+    }
+    for (int i=0; i<NCCL_MAX_DIRECT_ARITY; i++) {
+      int peer = channel->collTree.up[i];
+      if (peer == -1) break;
+      int j = comm->rankToIntraNodeRank[peer];
+      if (j < 0) {
+        WARN("Invalid intra-node rank %d for peer %d", j, peer);
+        return ncclInternalError;
+      }
+      regElem->upOutputs[i] = regInfo->recvbuffs[j];
+    }
+    work->elems[s].regUsed = 1;
+  }
   return ncclSuccess;
 }
 
@@ -873,9 +941,9 @@ ncclResult_t ncclEnqueueP2pKernel(struct ncclComm* comm, struct ncclQueueElem* e
   int opIndex = (channel->workFifoTail-1+NCCL_MAX_OPS)%NCCL_MAX_OPS;
   struct ncclWork* w = channel->workFifo+opIndex;
   int segment = -1;
-  if (channel->workCount && w->elems[0].funcIndex == FUNC_INDEX_P2P && w->elems[NCCL_MAX_WORK_ELEMENTS-1].active == 0) {
+  if (channel->workCount) {
     // Try to pack more segments into a single operation
-    segment = getSegment(P2P_SEGMENT, workElem->p2p.delta, w);
+    segment = getSegment(P2P_Segment, workElem->p2p.delta, w);
   }
   if (segment == -1) {
     NCCLCHECK(getNextOp(channel, &w, NULL));
@@ -884,7 +952,7 @@ ncclResult_t ncclEnqueueP2pKernel(struct ncclComm* comm, struct ncclQueueElem* e
 
   // store work element into FIFO
   NCCLCHECK(ncclProxySaveP2p(comm, proxyArgs));
-  NCCLCHECK(enqueueSegOp(P2P_SEGMENT, workElem, w, segment));
+  NCCLCHECK(enqueueSegOp(P2P_Segment, workElem, w, segment, &eqElem->buffRegInfo, channel, comm));
   comm->p2pOpCount++;
   return ncclSuccess;
 }
@@ -920,15 +988,18 @@ ncclResult_t ncclSetupP2pKernel(struct ncclInfo* info) {
   return ncclSuccess;
 }
 
-ncclResult_t ncclEnqueueAsyncKernel(struct ncclComm* comm, struct ncclQueueElem* eqElem) {
+// Dynamic enqueue function for collective kernels
+// Supports both aggregated and non-aggregated modes
+ncclResult_t ncclEnqueueCollKernel(struct ncclComm* comm, struct ncclQueueElem* eqElem, int aggMode) {
   struct ncclWorkElem* work = &eqElem->work;
   struct ncclProxyArgs* proxyArgs = &eqElem->proxyArgs;
 
   int nChannels = work->coll.nChannels;
   size_t channelSize = work->coll.count*ncclTypeSize(proxyArgs->dtype)/work->coll.nChannels;
+  int segmentType = proxyArgs->redOp == ncclNumOps ? RingTree_Segment : CollNet_Segment;  // redOp is only set when using CollNet
+
   for (int bid=0; bid<nChannels; bid++) {
-    int channelId;
-    NCCLCHECK(getNextChannel(comm, &channelId));
+    int channelId = getNextChannel(comm, aggMode);
     struct ncclChannel* channel = comm->channels+channelId;
 
     // Proxy
@@ -937,18 +1008,19 @@ ncclResult_t ncclEnqueueAsyncKernel(struct ncclComm* comm, struct ncclQueueElem*
     proxyArgs->commOpCount = comm->opCount;
     if (proxyArgs->subs[0].nsteps) NCCLCHECK(ncclProxySaveColl(proxyArgs, comm->nRanks));
 
-    // Try to reuse last work if not full yet
     work->coll.bid = bid % nChannels;
-    int opIndex = (channel->workFifoTail-1+NCCL_MAX_OPS)%NCCL_MAX_OPS;
-    struct ncclWork* w = channel->workFifo+opIndex;
+    struct ncclWork* w = NULL;
     int segment = -1;
-    if (channel->workCount && w->elems[NCCL_MAX_WORK_ELEMENTS-1].active == 0 &&
-        // All elems in work must have same (funcIndex,nThreads),
-        // see "src/collectives/device/common.h"
-        w->elems[0].funcIndex == work->funcIndex &&
-        w->elems[0].nThreads == work->nThreads) {
+    if (aggMode && channel->workCount) {
       // Try to pack more segments into a single operation
-      segment = getSegment(COLL_SEGMENT, 0, w);
+      int opIndex = (channel->workFifoTail-1+NCCL_MAX_OPS)%NCCL_MAX_OPS;
+      w = channel->workFifo+opIndex;
+      // All elems in work must have same (funcIndex,nThreads),
+      // see "src/collectives/device/common.h"
+      if (w->elems[0].funcIndex == work->funcIndex &&
+          w->elems[0].nThreads == work->nThreads) {
+        segment = getSegment(segmentType, 0, w);
+      }
     }
     if (segment == -1) {
       NCCLCHECK(getNextOp(channel, &w, NULL));
@@ -956,7 +1028,7 @@ ncclResult_t ncclEnqueueAsyncKernel(struct ncclComm* comm, struct ncclQueueElem*
     }
 
     // store work element into FIFO
-    NCCLCHECK(enqueueSegOp(COLL_SEGMENT, work, w, segment));
+    NCCLCHECK(enqueueSegOp(segmentType, work, w, segment, &eqElem->buffRegInfo, channel, comm));
     channel->totalSize += channelSize;
   }
   comm->collOpCount++;
@@ -968,17 +1040,15 @@ void HIPRT_CB ncclEnqueueHostSetup(void* arg) {
   ncclResult_t ret;
   struct ncclQueueInfo* eqInfo = (struct ncclQueueInfo*)arg;
   ncclComm_t comm = eqInfo->comm;
+  int aggMode = eqInfo->elemList->count() > 1 ? 1 : 0;
 
   // Iterate through the element list
   struct ncclQueueElem* eqElem = eqInfo->elemList->begin();
   while (eqElem != NULL) {
     if (eqElem->work.funcIndex == FUNC_INDEX_P2P) {
       NCCLCHECKGOTO(ncclEnqueueP2pKernel(comm, eqElem), ret, cb_end);
-    } else if (eqInfo->elemList->count() > 1) {
-      // We have more than one operation, hence aggregating
-      NCCLCHECKGOTO(ncclEnqueueAsyncKernel(comm, eqElem), ret, cb_end);
     } else {
-      NCCLCHECKGOTO(ncclEnqueueCollKernel(comm, eqElem), ret, cb_end);
+      NCCLCHECKGOTO(ncclEnqueueCollKernel(comm, eqElem, aggMode), ret, cb_end);
     }
     eqElem = eqInfo->elemList->getNext();
   }
@@ -996,51 +1066,95 @@ cb_end:
 template void HIPRT_CB ncclEnqueueHostSetup<0>(void*);
 template void HIPRT_CB ncclEnqueueHostSetup<1>(void*);
 
+void* graphHelperFunc(void *args) {
+  struct ncclGraphHelperResources* res = (struct ncclGraphHelperResources*)args;
+  if (res == NULL) {
+    WARN("CUDA Graph helper resource is null");
+    return NULL;
+  }
+  int dev = res->comm->cudaDev;
+  CUDACHECKIGNORE(hipSetDevice(dev));
+  INFO(NCCL_COLL, "CUDA Graph helper thread created for device %d", dev);
+
+  volatile enum helperThreadState* state = &res->threadState;
+  volatile int* ipcTail = &res->ipcTail;
+  while (1) {
+    int ipcTailMark = *ipcTail;
+    int ipcCount = 0;
+    while (res->ipcHead != ipcTailMark) {
+      if (res->ipcBases[res->ipcHead] != NULL)
+        CUDACHECKIGNORE(hipIpcCloseMemHandle(res->ipcBases[res->ipcHead]));
+      res->ipcBases[res->ipcHead] = NULL;
+      res->ipcHead = (res->ipcHead+1)%NCCL_IPC_POOL_SIZE;
+      ipcCount++;
+    }
+    TRACE(NCCL_COLL, "CUDA Graph helper thread closed %d IPC handles", ipcCount);
+    pthread_mutex_lock(&res->threadLock);
+    while (res->ipcHead == *ipcTail && *state != ThreadStop) {
+      pthread_cond_wait(&res->threadCond, &res->threadLock);
+    }
+    pthread_mutex_unlock(&res->threadLock);
+    if (*state == ThreadStop) {
+      INFO(NCCL_COLL, "CUDA Graph helper thread for device %d returning", dev);
+      return NULL;
+    }
+  }
+}
+
 ncclResult_t ncclGetCudaGraph(ncclComm_t comm, cudaGraph_t* graph) {
   comm->usingCudaGraph = 0;
 #if CUDART_VERSION >= 11030
-  cudaStreamCaptureStatus captureStatus;
-  unsigned long long cudaGraphId;
+  hipStreamCaptureStatus captureStatus;
+  unsigned long long hipGraphId;
   if (comm->driverVersion < 11030) {
-    CUDACHECK(cudaStreamIsCapturing(comm->userStream, &captureStatus));
-    if (captureStatus != cudaStreamCaptureStatusNone) {
+    CUDACHECK(hipStreamIsCapturing(comm->userStream, &captureStatus));
+    if (captureStatus != hipStreamCaptureStatusNone) {
       WARN("The installed CUDA driver is older than the minimum version (R465) required for NCCL's CUDA Graphs support");
       return ncclInvalidUsage;
     }
     return ncclSuccess;
   }
-  CUDACHECK(cudaStreamGetCaptureInfo_v2(comm->userStream, &captureStatus, &cudaGraphId, graph, NULL, NULL));
-  if (captureStatus == cudaStreamCaptureStatusActive) {
-    if (cudaGraphId != comm->lastCudaGraphId) {
-      INFO(NCCL_COLL, "stream is being captured by a new graph, id %llu", cudaGraphId);
+  CUDACHECK(hipStreamGetCaptureInfo_v2(comm->userStream, &captureStatus, &hipGraphId, graph, NULL, NULL));
+  if (captureStatus == hipStreamCaptureStatusActive) {
+    if (hipGraphId != comm->lastCudaGraphId) {
+      INFO(NCCL_COLL, "stream is being captured by a new graph, id %llu", hipGraphId);
       // We are in a new graph, hence need to forget the last setup node so that
       // the first setup node in the new graph will not have a dependency
-      comm->lastCudaGraphId = cudaGraphId;
+      comm->lastCudaGraphId = hipGraphId;
       comm->lastSetupNode = NULL;
     }
     if (comm->launchMode == ncclComm::GROUP) comm->launchMode = ncclComm::GROUP_GRAPH;
     comm->usingCudaGraph = 1;
+
+    // Create helper thread that closes IPC handles during graph destruction
+    // Only create this thread when buffer registration is enabled
+    if ((!comm->graphHelperThread) && comm->graphRegister == 1 && comm->disableGraphHelper == 0) {
+      pthread_mutex_init(&comm->graphHelperResources->threadLock, NULL);
+      pthread_cond_init(&comm->graphHelperResources->threadCond, NULL);
+      comm->graphHelperResources->threadState = ThreadStart;
+      pthread_create(&comm->graphHelperThread, NULL, graphHelperFunc, comm->graphHelperResources);
+    }
   }
 #endif
   return ncclSuccess;
 }
 
-ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, cudaGraph_t graph) {
+ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, hipGraph_t graph) {
 #if CUDART_VERSION >= 11030
   struct ncclQueueInfo* eqInfo = comm->enqueueInfo;
   // Create a CUDA object to wrap around the argument space
   // which CUDA graph would manage lifetime of
-  cudaUserObject_t object;
-  CUDACHECK(cudaUserObjectCreate(&object, eqInfo, ncclDestroyQueueInfo, 1/*initialRefcount*/, cudaUserObjectNoDestructorSync));
-  CUDACHECK(cudaGraphRetainUserObject(graph, object, 1, cudaGraphUserObjectMove));
+  hipUserObject_t object;
+  CUDACHECK(hipUserObjectCreate(&object, eqInfo, ncclDestroyQueueInfo, 1/*initialRefcount*/, hipUserObjectNoDestructorSync));
+  CUDACHECK(hipGraphRetainUserObject(graph, object, 1, hipGraphUserObjectMove));
 
-  cudaHostFn_t fn = ncclEnqueueHostSetup<1>;
+  hipHostFn_t fn = ncclEnqueueHostSetup<1>;
   // Add a CPU node to the graph
-  cudaGraphNode_t setupNode;
-  cudaHostNodeParams setupNodeParams = {fn, eqInfo};
+  hipGraphNode_t setupNode;
+  hipHostNodeParams setupNodeParams = {fn, eqInfo};
   int numDependencies = comm->lastSetupNode == NULL ? 0 : 1;
-  CUDACHECK(cudaGraphAddHostNode(&setupNode, graph, &comm->lastSetupNode, numDependencies, &setupNodeParams));
-  CUDACHECK(cudaStreamUpdateCaptureDependencies(comm->userStream, &setupNode, 1, cudaStreamAddCaptureDependencies));
+  CUDACHECK(hipGraphAddHostNode(&setupNode, graph, &comm->lastSetupNode, numDependencies, &setupNodeParams));
+  CUDACHECK(hipStreamUpdateCaptureDependencies(comm->userStream, &setupNode, 1, hipStreamAddCaptureDependencies));
   comm->lastSetupNode = setupNode;
   return ncclSuccess;
 #else
@@ -1049,6 +1163,74 @@ ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, cudaGraph_t graph) {
 #endif
 }
 
+static ncclResult_t hostToDevRedOp(
+    ncclDevRedOpFull *opFull, ncclRedOp_t op, ncclDataType_t datatype, ncclComm *comm
+  ) {
+  union {
+    int8_t i8;
+    uint8_t u8;
+    int32_t i32;
+    uint32_t u32;
+    int64_t i64;
+    uint64_t u64;
+    half f16;
+    #if defined(RCCL_BFLOAT16)
+      rccl_bfloat16 bf16;
+    #endif
+    float f32;
+    double f64;
+    void *ptr;
+  };
+  u64 = 0;
+  opFull->scalarArgIsPtr = false;
+  switch (int(op)) {
+  case ncclSum:  opFull->op = ncclDevSum;  break;
+  case ncclProd: opFull->op = ncclDevProd; break;
+  case ncclMax:  opFull->op = ncclDevMax;  break;
+  case ncclMin:  opFull->op = ncclDevMin;  break;
+  case ncclAvg:
+    switch ((int)datatype) {
+    case ncclInt8:  case ncclInt32:  case ncclInt64:
+    case ncclUint8: case ncclUint32: case ncclUint64:
+      opFull->op = ncclDevSumPostDiv;
+      u64 = comm->nRanks;
+      break;
+    case ncclFloat16:
+      opFull->op = ncclDevPreMulSum;
+      f16 = __float2half(float(1.0/comm->nRanks)); // __double2half not supported pre CUDA 11.x
+      break;
+    #if defined(RCCL_BFLOAT16)
+    case ncclBfloat16:
+      opFull->op = ncclDevPreMulSum;
+      bf16 = (rccl_bfloat16)(float(1.0/comm->nRanks));
+      break;
+    #endif
+    case ncclFloat32:
+      opFull->op = ncclDevPreMulSum;
+      f32 = float(1.0/comm->nRanks);
+      break;
+    case ncclFloat64:
+      opFull->op = ncclDevPreMulSum;
+      f64 = 1.0/comm->nRanks;
+      break;
+    }
+    opFull->scalarArgIsPtr = false;
+    opFull->scalarArg = u64;
+    break;
+  default: // user created
+    int ix = int(ncclUserRedOpMangle(comm, op)) - int(ncclNumOps);
+    ncclUserRedOp *user = &comm->userRedOps[ix];
+    if (datatype != user->datatype) {
+      WARN("Data type supplied to user-created ncclRedOp_t does not match type "
+           "given to reduction operation");
+      return ncclInvalidArgument;
+    }
+    *opFull = user->opFull;
+    break;
+  }
+  return ncclSuccess;
+}
+
 ncclResult_t ncclEnqueueCheck(struct ncclInfo* info) {
   // [RCCL] Check for clique-based kernel support
   {
@@ -1064,40 +1246,39 @@ ncclResult_t ncclEnqueueCheck(struct ncclInfo* info) {
   }
   // [/RCCL]
 
+  ncclResult_t ret = ncclSuccess;
+  bool isAsync = ncclAsyncMode();
+  int savedDev = -1;
+  // Check arguments
+  NCCLCHECK(PtrCheck(info->comm, info->opName, "comm"));
+  if (isAsync && info->comm->checkPointers) {
+    CUDACHECKGOTO(hipGetDevice(&savedDev), ret, end);
+    CUDACHECKGOTO(hipSetDevice(info->comm->cudaDev), ret, end);
+  }
+  NCCLCHECKGOTO(ArgsCheck(info), ret, end);
+
+  // Copy reduction op state from op handle into info struct here since the
+  // op handle may be destroyed before ncclGroupEnd().
+  NCCLCHECKGOTO(hostToDevRedOp(&info->opFull, info->op, info->datatype, info->comm), ret, end);
+
   // Launch asynchronously if needed
-  if (ncclAsyncMode()) {
-    ncclResult_t ret = ncclSuccess;
-    int savedDev = -1;
-    // Check arguments
-    NCCLCHECK(PtrCheck(info->comm, info->opName, "comm"));
-    if (info->comm->checkPointers) {
-      CUDACHECKGOTO(hipGetDevice(&savedDev), ret, end);
-      CUDACHECKGOTO(hipSetDevice(info->comm->cudaDev), ret, end);
-    }
-    NCCLCHECKGOTO(ArgsCheck(info), ret, end);
+  if (isAsync) {
     // Always register comm even in case of error to make sure ncclGroupEnd
     // cleans it up.
     NCCLCHECKGOTO(ncclAsyncColl(info->comm), ret, end);
     NCCLCHECKGOTO(checkSetStream(info), ret, end);
 
-    INFO(NCCL_COLL,"%s: opCount %lx sendbuff %p recvbuff %p count %zi datatype %d op %d root %d comm %p [nranks=%d] stream %p",
+    INFO(NCCL_COLL,"%s: opCount %lx sendbuff %p recvbuff %p count %zi datatype %d op %d root %d comm %p [nranks=%d] stream %p devRedOp %d isPtr %d scaler %lx",
         info->opName, info->coll == ncclFuncSendRecv ? info->comm->p2pOpCount : info->comm->collOpCount, info->sendbuff, info->recvbuff, info->count,
-        info->datatype, info->op, info->root, info->comm, info->comm->nRanks, info->stream);
+        info->datatype, info->op, info->root, info->comm, info->comm->nRanks, info->stream, info->opFull.op, info->opFull.scalarArgIsPtr, info->opFull.scalarArg);
 
     if (info->coll == ncclFuncSendRecv) { //p2p stored separately
       NCCLCHECKGOTO(ncclSaveP2p(info), ret, end);
     } else {
       NCCLCHECKGOTO(ncclSaveAsyncColl(info), ret, end);
     }
-
-end:
-    if (savedDev != -1) CUDACHECK(hipSetDevice(savedDev));
-    ncclAsyncErrCheck(ret);
-    return ret;
   } else {
-    NCCLCHECK(PtrCheck(info->comm, info->opName, "comm"));
-    NCCLCHECK(ArgsCheck(info));
-    NCCLCHECK(checkSetStream(info));
+    NCCLCHECKGOTO(checkSetStream(info), ret, end);
 
     INFO(NCCL_COLL,"%s: opCount %lx sendbuff %p recvbuff %p count %zi datatype %d op %d root %d comm %p [nranks=%d] stream %p",
         info->opName, info->comm->collOpCount, info->sendbuff, info->recvbuff, info->count,
@@ -1106,24 +1287,82 @@ end:
     // Check whether we are in cuda graph mode
     cudaGraph_t graph;
     ncclComm_t comm = info->comm;
-    NCCLCHECK(ncclGetCudaGraph(comm, &graph));
+    NCCLCHECKGOTO(ncclGetCudaGraph(comm, &graph), ret, end);
 
     // Common part between graph mode and non-graph mode
-    NCCLCHECK(ncclSetupCollKernel(info));
+    NCCLCHECKGOTO(ncclSetupCollKernel(info), ret, end);
 
     // Host setup
     if (comm->usingCudaGraph) {
-      NCCLCHECK(ncclCudaGraphHostSetup(comm, graph));
+      NCCLCHECKGOTO(ncclCudaGraphHostSetup(comm, graph), ret, end);
     } else {
       ncclEnqueueHostSetup<0>(comm->enqueueInfo);
-      NCCLCHECK(comm->enqueueInfo->ret);
+      NCCLCHECKGOTO(comm->enqueueInfo->ret, ret, end);
     }
 
     // Common part between graph mode and non-graph mode
-    NCCLCHECK(ncclLaunchBarrier(comm));
-    NCCLCHECK(ncclLaunchKernel(comm));
-    NCCLCHECK(ncclRecordEvents(comm));
-    NCCLCHECK(ncclLaunchReset(comm));
-    return ncclSuccess;
+    NCCLCHECKGOTO(ncclLaunchBarrier(comm), ret, end);
+    NCCLCHECKGOTO(ncclLaunchKernel(comm), ret, end);
+    NCCLCHECKGOTO(ncclRecordEvents(comm), ret, end);
+    NCCLCHECKGOTO(ncclLaunchReset(comm), ret, end);
   }
+end:
+  if (isAsync && savedDev != -1) CUDACHECK(hipSetDevice(savedDev));
+  if (isAsync) ncclAsyncErrCheck(ret);
+  return ret;
+}
+
+NCCL_API(ncclResult_t, ncclRedOpCreatePreMulSum, ncclRedOp_t *op, void *scalar, ncclDataType_t datatype, ncclScalarResidence_t residence, ncclComm_t comm);
+ncclResult_t ncclRedOpCreatePreMulSum(ncclRedOp_t *op, void *scalar, ncclDataType_t datatype, ncclScalarResidence_t residence, ncclComm_t comm) {
+  if (comm->userRedOpFreeHead == comm->userRedOpCapacity) {
+    // double capacity and resize
+    int cap = 2*comm->userRedOpCapacity;
+    if (cap < 4) cap = 4;
+    ncclUserRedOp *ops = new ncclUserRedOp[cap];
+    std::memcpy(ops, comm->userRedOps, comm->userRedOpCapacity*sizeof(ncclUserRedOp));
+    for(int ix=comm->userRedOpCapacity; ix < cap; ix++)
+      ops[ix].freeNext = ix + 1;
+    delete[] comm->userRedOps;
+    comm->userRedOps = ops;
+    comm->userRedOpCapacity = cap;
+  }
+  // pop from free list
+  int ix = comm->userRedOpFreeHead;
+  ncclUserRedOp *user = &comm->userRedOps[ix];
+  comm->userRedOpFreeHead = user->freeNext;
+
+  user->freeNext = -1; // allocated
+  user->datatype = datatype;
+  user->opFull.op = ncclDevPreMulSum;
+  if (residence == ncclScalarHostImmediate) {
+    user->opFull.scalarArgIsPtr = false;
+    std::memcpy(&user->opFull.scalarArg, scalar, ncclTypeSize(datatype));
+  } else {
+    user->opFull.scalarArgIsPtr = true;
+    user->opFull.scalarArg = reinterpret_cast<uint64_t>(scalar);
+  }
+  *op = ncclRedOp_t(int(ncclNumOps) + ix);
+  *op = ncclUserRedOpMangle(comm, *op);
+  return ncclSuccess;
+}
+
+NCCL_API(ncclResult_t, ncclRedOpDestroy, ncclRedOp_t op, ncclComm_t comm);
+ncclResult_t ncclRedOpDestroy(ncclRedOp_t op, ncclComm_t comm) {
+  if (0 <= int(op) && int(op) < int(ncclNumOps)) {
+    WARN("ncclRedOpDestroy : operator is a NCCL builtin.");
+    return ncclInvalidArgument;
+  }
+  if (int(op) < 0 || int(ncclMaxRedOp) < int(op)) {
+    WARN("ncclRedOpDestroy :  operator is garbage.");
+    return ncclInvalidArgument;
+  }
+  int ix = int(ncclUserRedOpMangle(comm, op)) - int(ncclNumOps);
+  if (comm->userRedOpCapacity <= ix || comm->userRedOps[ix].freeNext != -1) {
+    WARN("ncclRedOpDestroy : operator unknown to this communicator.");
+    return ncclInvalidArgument;
+  }
+  // push to free list
+  comm->userRedOps[ix].freeNext = comm->userRedOpFreeHead;
+  comm->userRedOpFreeHead = ix;
+  return ncclSuccess;
 }