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

2022-10-20 15:40:03 +00:00
@@ -23,17 +23,15 @@ uint64_t clockNano(); // from utils.h with which we have a circular dependency
 template <typename T>
 ncclResult_t ncclCudaHostCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
  ncclResult_t result = ncclSuccess;
-  uint64_t time = 0;
  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
  *ptr = nullptr;
  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
-  time = clockNano();
  CUDACHECKGOTO(hipHostMalloc(ptr, nelem*sizeof(T), hipHostMallocMapped), result, finish);
-  time = clockNano() - time;
  memset(*ptr, 0, nelem*sizeof(T));
-  INFO(NCCL_ALLOC, "%s:%d Cuda Host Alloc Size %ld pointer %p seconds: hipHostAlloc=%g", filefunc, line, nelem*sizeof(T), *ptr, double(time)/1.e9);
 finish:
  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  if (*ptr == nullptr) WARN("Failed to CUDA host alloc %ld bytes", nelem*sizeof(T));
+  INFO(NCCL_ALLOC, "%s:%d Cuda Host Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
  return result;
 }
 #define ncclCudaHostCalloc(...) ncclCudaHostCallocDebug(__VA_ARGS__, __FILE__, __LINE__)
@@ -95,15 +93,14 @@ ncclResult_t ncclCudaMallocDebug(const char *filefunc, int line, T** ptr, size_t
  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
  *ptr = nullptr;
  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
-  uint64_t time = clockNano();
  if (isFineGrain)
    CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
  else
    CUDACHECKGOTO(hipMalloc(ptr, nelem*sizeof(T)), result, finish);
-  time = clockNano() - time;
 finish:
  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
-  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p seconds: hipMalloc=%g", filefunc, line, nelem*sizeof(T), *ptr, double(time)/1.e9);
+  if (*ptr == nullptr) WARN("Failed to CUDA malloc %ld bytes", nelem*sizeof(T));
+  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
  return result;
 }
 #define ncclCudaMalloc(...) ncclCudaMallocDebug( __FILE__, __LINE__, __VA_ARGS__)
@@ -111,21 +108,17 @@ finish:
 template <typename T>
 ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, hipStream_t sideStream = nullptr, bool isFineGrain = false) {
  ncclResult_t result = ncclSuccess;
-  uint64_t time0=0, time1=0, time2=0;
  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
  *ptr = nullptr;
  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
  // Need a side stream so as not to interfere with graph capture.
  hipStream_t stream = sideStream;
-  time0 = clockNano();
  if (stream == nullptr)
    CUDACHECK(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
-  time1 = clockNano();
  if (isFineGrain)
    CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
  else
    CUDACHECKGOTO(hipMalloc(ptr, nelem*sizeof(T)), result, finish);
-  time2 = clockNano();
  CUDACHECKGOTO(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
  CUDACHECKGOTO(hipStreamSynchronize(stream), result, finish);
  if (sideStream == nullptr)
@@ -136,9 +129,10 @@ ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t
    __atomic_fetch_add(&allocTracker[dev].totalAlloc, 1, __ATOMIC_RELAXED);
    __atomic_fetch_add(&allocTracker[dev].totalAllocSize, nelem*sizeof(T), __ATOMIC_RELAXED);
  }
-  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p seconds: hipStreamCreateWithFlags=%g hipMalloc=%g", filefunc, line, nelem*sizeof(T), *ptr, double(time1-time0)/1.e9, double(time2-time1)/1.e9);
 finish:
+  if (*ptr == nullptr) WARN("Failed to CUDA calloc %ld bytes", nelem*sizeof(T));
  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
  return result;
 }
 #define ncclCudaCalloc(...) ncclCudaCallocDebug(__FILE__, __LINE__, __VA_ARGS__)
@@ -146,16 +140,13 @@ finish:
 template <typename T>
 ncclResult_t ncclCudaCallocAsyncDebug(const char *filefunc, int line, T** ptr, size_t nelem, hipStream_t stream, bool isFineGrain = false) {
  ncclResult_t result = ncclSuccess;
-  uint64_t time = 0;
  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
  *ptr = nullptr;
  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
-  time = clockNano();
  if (isFineGrain)
    CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
  else
    CUDACHECKGOTO(hipMalloc(ptr, nelem*sizeof(T)), result, finish);
-  time = clockNano() - time;
  CUDACHECKGOTO(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
  int dev;
  CUDACHECK(hipGetDevice(&dev));
@@ -163,9 +154,10 @@ ncclResult_t ncclCudaCallocAsyncDebug(const char *filefunc, int line, T** ptr, s
    __atomic_fetch_add(&allocTracker[dev].totalAlloc, 1, __ATOMIC_RELAXED);
    __atomic_fetch_add(&allocTracker[dev].totalAllocSize, nelem*sizeof(T), __ATOMIC_RELAXED);
  }
-  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p seconds: hipMalloc=%g", filefunc, line, nelem*sizeof(T), *ptr, double(time)/1.e9);
 finish:
  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  if (*ptr == nullptr) WARN("Failed to CUDA calloc async %ld bytes", nelem*sizeof(T));
+  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
  return result;
 }
 #define ncclCudaCallocAsync(...) ncclCudaCallocAsyncDebug(__FILE__, __LINE__, __VA_ARGS__)
@@ -107,7 +107,7 @@
 // Propagate errors up
 #define NCCLCHECK(call) do { \
  ncclResult_t res = call; \
-  if (res != ncclSuccess) { \
+  if (res != ncclSuccess && res != ncclInProgress) { \
    /* Print the back trace*/ \
    if (ncclDebugNoWarn == 0) INFO(NCCL_ALL,"%s:%d -> %d", __FILE__, __LINE__, res);    \
    return res; \
@@ -116,7 +116,7 @@

 #define NCCLCHECKGOTO(call, res, label) do { \
  res = call; \
-  if (res != ncclSuccess) { \
+  if (res != ncclSuccess && res != ncclInProgress) { \
    /* Print the back trace*/ \
    if (ncclDebugNoWarn == 0) INFO(NCCL_ALL,"%s:%d -> %d", __FILE__, __LINE__, res);    \
    goto label; \
@@ -126,7 +126,7 @@
 #define NCCLWAIT(call, cond, abortFlagPtr) do {         \
  volatile uint32_t* tmpAbortFlag = (abortFlagPtr);     \
  ncclResult_t res = call;                \
-  if (res != ncclSuccess) {               \
+  if (res != ncclSuccess && res != ncclInProgress) {               \
    if (ncclDebugNoWarn == 0) INFO(NCCL_ALL,"%s:%d -> %d", __FILE__, __LINE__, res);    \
    return ncclInternalError;             \
  }                                       \
@@ -136,7 +136,7 @@
 #define NCCLWAITGOTO(call, cond, abortFlagPtr, res, label) do { \
  volatile uint32_t* tmpAbortFlag = (abortFlagPtr);             \
  res = call;                             \
-  if (res != ncclSuccess) {               \
+  if (res != ncclSuccess && res != ncclInProgress) {               \
    if (ncclDebugNoWarn == 0) INFO(NCCL_ALL,"%s:%d -> %d", __FILE__, __LINE__, res);    \
    goto label;                           \
  }                                       \
@@ -144,7 +144,7 @@
 } while (!(cond));

 #define NCCLCHECKTHREAD(a, args) do { \
-  if (((args)->ret = (a)) != ncclSuccess) { \
+  if (((args)->ret = (a)) != ncclSuccess && (args)->ret != ncclInProgress) { \
    INFO(NCCL_INIT,"%s:%d -> %d [Async thread]", __FILE__, __LINE__, (args)->ret); \
    return args; \
  } \
@@ -52,8 +52,9 @@ struct ncclDevRedOpFull {
  extern __global__ void NCCL_KERN_NAME_LL128(func, algo, proto, devredop, type)(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead); \
  extern __global__ void NCCL_KERN_NAME_LL128_DEBUG(func, algo, proto, devredop, type)(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead);

+#define SINGLE_ARG(...) __VA_ARGS__
 #define CONCAT(a,b) a##b
-#define MACRO_IF(cond, t, f) CONCAT(MACRO_IF_, cond)(t, f)
+#define MACRO_IF(cond, t, f) CONCAT(MACRO_IF_, cond)(SINGLE_ARG(t), SINGLE_ARG(f))
 #define MACRO_IF_0(t, f) f
 #define MACRO_IF_1(t, f) t

@@ -65,7 +66,8 @@ struct ncclDevRedOpFull {
 #define DECL3(func, devredop, type, undef) \
  DECL4(func, RING,    devredop, type, undef) \
  DECL4(func, TREE,    devredop, type, undef) \
-  DECL4(func, COLLNET, devredop, type, undef)
+  DECL4(func, COLLNET_DIRECT, devredop, type, undef) \
+  DECL4(func, COLLNET_CHAIN, devredop, type, undef)

 #if defined(RCCL_BFLOAT16)
 #define DECL2(func, devredop, undefForFloat) \
@@ -132,7 +134,6 @@ extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, double)();
 #define BROADCAST_CHUNKSTEPS 1
 #define REDUCE_SLICESTEPS 1
 #define REDUCE_CHUNKSTEPS 1
-#define SENDRECV_SLICEFACTOR 4
 #define NCCL_MAX_SLICE_PER_CHUNK 2  // max value for CHUNKSTEPS/SLICESTEPS, must accord with above
 #define ALLTOALL_PIVOT_SLICESTEPS 2
 #define ALLTOALL_PIVOT_CHUNKSTEPS 4
@@ -107,8 +107,9 @@ struct ncclChannel {
  struct ncclRing ring;
  int* devRingUserRanks;
  struct ncclTree tree;
+  struct ncclTree collnetChain;
+  struct ncclDirect collnetDirect;
  struct ncclTree binTree;
-  struct ncclDirect collTree;
  int id; // index of this channel
  uint32_t workFifoSent; // last used work index+1
  uint64_t p2pOpCount;
@@ -134,6 +135,7 @@ struct ncclKernelPlan {
  struct ncclKernelPlan* next;

  bool persistent; // aka captured in a graph
+  bool kernelSpecialized;
  void *kernelFn;
  int channelUbound; // only channels c < channelUbound are present
  int channelCount; // number of channels present
@@ -209,8 +211,12 @@ struct ncclComm {
  int p2pnChannelsPerPeer;
  int p2pChannels[MAXCHANNELS];

+  // Should this comm allocate LL buffers for network P2P connections?
+  bool allocP2pNetLLBuffers;
+
  // Buffer sizes
  int buffSizes[NCCL_NUM_PROTOCOLS];
+  int p2pNetChunkSize;

  // Algorithm/Protocols thresholds
  ssize_t threadThresholds[NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
@@ -218,8 +224,9 @@ struct ncclComm {
  float bandwidths[NCCL_NUM_FUNCTIONS][NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
  int maxThreads[NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];

-  // Whether there has been a fatal error in this communicator.
-  ncclResult_t fatalError;
+  /* This attribute can indicate the states of communicators and return code of
+   * asynchronous NCCL operations. */
+  ncclResult_t asyncResult;

  // Flag to ask NCCL kernels to abort
  volatile uint32_t *abortFlag;
@@ -301,12 +308,16 @@ struct ncclComm {
  pthread_t collTraceThread;
  volatile bool collTraceExit;
 #endif
-};

-// Set to true during an `atexit()` handler. We use this to intentionally leak
-// unfreed CUDA resources when cleaning up after return of `main()` to avoid
-// CUDA calls after CUDA runtime teardown.
-extern bool ncclMainExited;
+  // communicator mode
+  int blocking;
+  // initState is to more conveniently reclaim resources when errors happen.
+  ncclResult_t initState;
+  // flag to indicate if ncclCommFinalize() is called
+  bool finalizeCalled;
+  // shared structures for finalization
+  int finalizeRankCnt;
+};

 enum ncclLaunchMode {
  ncclLaunchModeInvalid=0,
@@ -320,13 +331,16 @@ void ncclCommPushCudaFree(struct ncclComm* comm, void* buf);
 void ncclCommPushCudaHostFree(struct ncclComm* comm, void* buf);
 void ncclCommPushCudaGdrFree(struct ncclComm* comm, void* handle);

-inline ncclResult_t ncclCommPollCallbacks(struct ncclComm* comm) {
-  struct ncclCommCallback* cb = ncclIntruQueueMpscDequeueAll(&comm->callbackQueue, /*waitSome=*/false);
+inline ncclResult_t ncclCommPollCallbacks(struct ncclComm* comm, bool waitSome) {
+  ncclResult_t result = ncclSuccess;
+  struct ncclCommCallback* cb = ncclIntruQueueMpscDequeueAll(&comm->callbackQueue, waitSome);
  while (cb != nullptr) {
    struct ncclCommCallback* next = cb->next;
-    NCCLCHECK(cb->fn(comm, cb)); // may reclaim memory of cb
+    ncclResult_t res1 = cb->fn(comm, cb); // may reclaim memory of cb
+    if (res1 != ncclSuccess) result = res1;
    cb = next;
  }
+  NCCLCHECK(result);
  return ncclSuccess;
 }

@@ -383,4 +397,7 @@ static inline ncclRedOp_t ncclUserRedOpMangle(ncclComm *comm, ncclRedOp_t op) {
  return op1 < int(ncclNumOps) ? op : ncclRedOp_t(op1);
 }

+ncclResult_t ncclCommEnsureReady(ncclComm_t comm);
+ncclResult_t ncclCommSetAsyncError(ncclComm_t comm, ncclResult_t nextState);
+
 #endif
@@ -12,9 +12,9 @@
 #if CUDART_VERSION >= 11030
 #include <cudaTypedefs.h>
 #else
-typedef CUresult (CUDAAPI *PFN_cuInit)(unsigned int Flags);
-typedef CUresult (CUDAAPI *PFN_cuDriverGetVersion)(int *driverVersion);
-typedef CUresult (CUDAAPI *PFN_cuGetProcAddress)(const char *symbol, void **pfn, int driverVersion, cuuint64_t flags);
+typedef CUresult (CUDAAPI *PFN_cuInit_v2000)(unsigned int Flags);
+typedef CUresult (CUDAAPI *PFN_cuDriverGetVersion_v2020)(int *driverVersion);
+typedef CUresult (CUDAAPI *PFN_cuGetProcAddress_v11030)(const char *symbol, void **pfn, int driverVersion, cuuint64_t flags);
 #endif

 #define CUPFN(symbol) pfn_##symbol
@@ -60,27 +60,27 @@ typedef CUresult (CUDAAPI *PFN_cuGetProcAddress)(const char *symbol, void **pfn,
    }									\
 } while(0)

-#define DECLARE_CUDA_PFN_EXTERN(symbol) extern PFN_##symbol pfn_##symbol
+#define DECLARE_CUDA_PFN_EXTERN(symbol,version) extern PFN_##symbol##_v##version pfn_##symbol

 #if CUDART_VERSION >= 11030
 /* CUDA Driver functions loaded with cuGetProcAddress for versioning */
-DECLARE_CUDA_PFN_EXTERN(cuDeviceGet);
-DECLARE_CUDA_PFN_EXTERN(cuDeviceGetAttribute);
-DECLARE_CUDA_PFN_EXTERN(cuGetErrorString);
-DECLARE_CUDA_PFN_EXTERN(cuGetErrorName);
-DECLARE_CUDA_PFN_EXTERN(cuMemGetAddressRange);
-DECLARE_CUDA_PFN_EXTERN(cuCtxCreate_v3020);
-DECLARE_CUDA_PFN_EXTERN(cuCtxDestroy);
-DECLARE_CUDA_PFN_EXTERN(cuCtxSetCurrent);
+DECLARE_CUDA_PFN_EXTERN(cuDeviceGet, 2000);
+DECLARE_CUDA_PFN_EXTERN(cuDeviceGetAttribute, 2000);
+DECLARE_CUDA_PFN_EXTERN(cuGetErrorString, 6000);
+DECLARE_CUDA_PFN_EXTERN(cuGetErrorName, 6000);
+DECLARE_CUDA_PFN_EXTERN(cuMemGetAddressRange, 3020);
+DECLARE_CUDA_PFN_EXTERN(cuCtxCreate, 3020);
+DECLARE_CUDA_PFN_EXTERN(cuCtxDestroy, 4000);
+DECLARE_CUDA_PFN_EXTERN(cuCtxSetCurrent, 4000);
 #if CUDA_VERSION >= 11070
-DECLARE_CUDA_PFN_EXTERN(cuMemGetHandleForAddressRange); // DMA-BUF support
+DECLARE_CUDA_PFN_EXTERN(cuMemGetHandleForAddressRange, 11070); // DMA-BUF support
 #endif
 #endif

 /* CUDA Driver functions loaded with dlsym() */
-DECLARE_CUDA_PFN_EXTERN(cuInit);
-DECLARE_CUDA_PFN_EXTERN(cuDriverGetVersion);
-DECLARE_CUDA_PFN_EXTERN(cuGetProcAddress);
+DECLARE_CUDA_PFN_EXTERN(cuInit, 2000);
+DECLARE_CUDA_PFN_EXTERN(cuDriverGetVersion, 2020);
+DECLARE_CUDA_PFN_EXTERN(cuGetProcAddress, 11030);


 ncclResult_t cudaLibraryInit(void);
@@ -21,10 +21,11 @@
 typedef enum { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv, ncclFuncSend, ncclFuncRecv, ncclFuncAllToAllPivot, ncclNumFuncs} ncclFunc_t;
 extern const char* ncclFuncStr[NCCL_NUM_FUNCTIONS+2];

-#define NCCL_NUM_ALGORITHMS 3 // Tree/Ring/CollNet
+#define NCCL_NUM_ALGORITHMS 4 // Tree/Ring/CollNet*
 #define NCCL_ALGO_TREE 0
 #define NCCL_ALGO_RING 1
-#define NCCL_ALGO_COLLNET 2
+#define NCCL_ALGO_COLLNET_DIRECT 2
+#define NCCL_ALGO_COLLNET_CHAIN 3
 extern const char* ncclAlgoStr[NCCL_NUM_ALGORITHMS];

 #define NCCL_NUM_PROTOCOLS 3 // Simple/LL/LL128
@@ -234,8 +235,9 @@ static_assert((NCCL_WORK_SIZE - alignUp(sizeof(ncclWorkHeader), alignof(ncclWork

 struct ncclWorkElemP2p {
  struct {
-    int32_t peer:30;
+    int32_t peer:28;
    uint32_t connIndex:2;
+    int32_t proto:2;
  };
  union {
    uint16_t flagBits;
@@ -356,8 +358,9 @@ struct alignas(16) ncclDevChannel {
  struct ncclDevChannelPeer *peers;
  struct ncclRing ring;
  struct ncclTree tree;
+  struct ncclTree collnetChain;
+  struct ncclDirect collnetDirect;
  struct ncclTree binTree;
-  struct ncclDirect collTree;
  uint32_t* workFifoDone; // Location of done counter, device writes index+1 of last work processed
 };

@@ -37,6 +37,7 @@ ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* topo, int64_t busId, int ne
 #define MAX_XGMI_INTER_GPUS 4
 ncclResult_t ncclTopoGetIntraNetDev(struct ncclTopoSystem* system, int rank, struct ncclTopoGraph* graph, int channelId, int type, int* dev);
 ncclResult_t ncclTopoGetLinkType(struct ncclTopoSystem* system, int cudaDev1, int cudaDev2, bool* isXGMI, int maxInter=MAX_XGMI_INTER_GPUS, int nInter=0, int *inter=nullptr);
+ncclResult_t ncclTopoCheckNet(struct ncclTopoSystem* system, int64_t id1, int64_t id2, int* net);
 int ncclPxnDisable(struct ncclComm* comm);
 ncclResult_t ncclTopoGetPxnRanks(struct ncclComm* comm, int** intermediateRanks, int* nranks);
 ncclResult_t ncclTopoGetLocalRank(struct ncclTopoSystem* system, int rank, int* localRank);
@@ -57,6 +58,7 @@ ncclResult_t ncclTopoGetCpuAffinity(struct ncclTopoSystem* system, int rank, cpu
 #define NCCL_TOPO_CPU_TYPE_YONGFENG 1
 ncclResult_t ncclTopoCpuType(struct ncclTopoSystem* system, int* arch, int* vendor, int* model);
 ncclResult_t ncclTopoGetNetCount(struct ncclTopoSystem* system, int* count);
+ncclResult_t ncclTopoGetNvsCount(struct ncclTopoSystem* system, int* count);
 ncclResult_t ncclTopoGetLocalNet(struct ncclTopoSystem* system, int rank, int* id);

 #define NCCL_TOPO_MAX_NODES 256
@@ -78,8 +80,8 @@ struct ncclTopoGraph {
  int maxChannels;
  // Output
  int nChannels;
-  float speedIntra;
-  float speedInter;
+  float bwIntra;
+  float bwInter;
  float latencyInter;
  int typeIntra;
  int typeInter;
@@ -1,6 +1,5 @@
 /*************************************************************************
 * Copyright (c) 2015-2017, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
 *
 * See LICENSE.txt for license information
 ************************************************************************/
@@ -14,7 +13,18 @@
 ncclResult_t ncclGroupErrCheck(ncclResult_t ret);
 void ncclGroupCommJoin(struct ncclComm* comm);
 void ncclGroupCommPreconnect(struct ncclComm* comm);
-void ncclGroupCommLeave(struct ncclComm* comm);
+ncclResult_t ncclGroupCommLeave(struct ncclComm* comm);
+void ncclGroupJobAbort();
+
+typedef ncclResult_t(*ncclInitFunc_t)(ncclComm_t* newcomm, int ndev, ncclUniqueId commId, int myrank, int cudaDev);
+
+ncclResult_t ncclAsyncInit(ncclInitFunc_t func, ncclComm_t* newcomm, int ndev, ncclUniqueId commId, int myrank, int cudaDev);
+
+typedef enum ncclGroupJobState {
+  ncclGroupJobRunning = 0,
+  ncclGroupJobDone    = 1,
+  ncclGroupJobJoined  = 2,
+} ncclGroupJobState_t;

 struct ncclAsyncJob {
  struct ncclAsyncJob* next;
@@ -23,17 +33,31 @@ struct ncclAsyncJob {
  ncclResult_t(*func)(struct ncclAsyncJob*);
  void(*undo)(struct ncclAsyncJob*);
  void(*destructor)(void*);
+  ncclGroupJobState_t state;
+  volatile uint32_t *abortFlag; /* point to comm abortFlag */
+  ncclComm_t comm;
 };

 ncclResult_t ncclAsyncLaunch(
  struct ncclAsyncJob* job,
  ncclResult_t(*func)(struct ncclAsyncJob*),
  void(*undo)(struct ncclAsyncJob*),
-  void(*destructor)(void*)
+  void(*destructor)(void*), ncclComm_t comm
 );

+struct ncclGroupJob {
+  struct ncclAsyncJob base;
+  struct ncclComm **groupCommHeadPtr;
+  struct ncclComm **groupCommPreconnectHeadPtr;
+  ncclResult_t *groupErrorPtr;
+  volatile bool *abortFlagPtr;
+  struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> *asyncJobsPtr;
+  bool doneFlag;
+};
+
 ncclResult_t ncclGroupStartInternal();
 ncclResult_t ncclGroupEndInternal();
+ncclResult_t ncclAsyncJobComplete(struct ncclAsyncJob* job);

 ////////////////////////////////////////////////////////////////////////////////

@@ -41,6 +65,7 @@ extern __thread int ncclGroupDepth; // depth of ncclGroupStart nesting
 extern __thread ncclResult_t ncclGroupError;
 extern __thread struct ncclComm* ncclGroupCommHead;
 extern __thread struct ncclComm* ncclGroupCommPreconnectHead;
+extern __thread int ncclGroupBlocking;

 inline ncclResult_t ncclGroupStartInternal() {
  ncclGroupDepth++;
@@ -49,7 +74,7 @@ inline ncclResult_t ncclGroupStartInternal() {

 inline ncclResult_t ncclGroupErrCheck(ncclResult_t ret) {
  if (ncclGroupDepth > 0) {
-    if (ncclGroupError == ncclSuccess || ret != ncclSuccess) ncclGroupError = ret;
+    if (ret != ncclSuccess && ret != ncclInProgress) ncclGroupError = ret;
  }
  return ret;
 }
@@ -69,6 +94,8 @@ inline void ncclGroupCommJoin(struct ncclComm* comm) {
    // this comm is allocated there.
    ncclMemoryStackPush(&comm->memScoped);
  }
+
+  ncclGroupBlocking = comm->blocking;
 }

 // Add comm to this thread's group needing preconnect
@@ -80,9 +107,10 @@ inline void ncclGroupCommPreconnect(struct ncclComm* comm) {
 }

 // Comm has left group
-inline void ncclGroupCommLeave(struct ncclComm* comm) {
+inline ncclResult_t ncclGroupCommLeave(struct ncclComm* comm) {
  comm->groupNext = reinterpret_cast<struct ncclComm*>(0x1);
  ncclMemoryStackPop(&comm->memScoped);
+  return ncclSuccess;
 }

 #endif
@@ -23,7 +23,8 @@ typedef enum : uint8_t {
  ncclPatternTreeUp,
  ncclPatternTreeDown,
  ncclPatternTreeUpDown,
-  ncclPatternCollTreeUpDown,
+  ncclPatternCollnetChain,
+  ncclPatternCollnetDirect,
  ncclPatternSend,
  ncclPatternRecv
 } ncclPattern_t;
@@ -173,6 +173,7 @@ struct ncclProxyState {
  struct ncclSocket* listenSock;
  int stop;
  hipCtx_t cudaCtx;
+  int safeAbortFlag;

  // Used by main thread
  union ncclSocketAddress* peerAddresses;
@@ -192,6 +193,7 @@ struct ncclProxyConnection {
  struct ncclProxyArgs *proxyAppend;
  struct ncclProxyArgs **proxyAppendPtr;
  void* transportResources;
+  bool initFlag;
 };

 typedef ncclResult_t (*threadFunc_t)(struct ncclProxyArgs*);