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

[ROCm/rccl commit: a79d9e3586]
2022-09-09 01:20:52 +00:00
parent fe99249cde c9a50a9ec5
commit 7874a99c75
90 changed files with 5517 additions and 3115 deletions
@@ -45,12 +45,7 @@ ncclResult_t ArgsCheck(struct ncclInfo* info) {
    return ncclInvalidArgument;
  }
  // Type is OK, compute nbytes. Convert Allgather/Broadcast/P2P/AllToAllPivot calls to chars.
-  info->nBytes = info->count * ncclTypeSize(info->datatype);
-  if (info->coll == ncclFuncAllGather || info->coll == ncclFuncBroadcast || info->coll == ncclFuncAllToAllPivot) {
-    info->count = info->nBytes;
-    info->datatype = ncclInt8;
-  }
-  if (info->coll == ncclFuncAllGather || info->coll == ncclFuncReduceScatter) info->nBytes *= info->comm->nRanks; // count is per rank
+  NCCLCHECK(ncclInfoSetDerived(info, info->comm->nRanks));

  if (info->op < 0 || ncclMaxRedOp < info->op) {
    WARN("%s : invalid reduction operation %d", info->opName, info->op);
@@ -0,0 +1,163 @@
+/*************************************************************************
+ * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#include "nccl.h"
+#include "debug.h"
+#include "cudawrap.h"
+
+#include <dlfcn.h>
+
+#define DECLARE_CUDA_PFN(symbol) PFN_##symbol pfn_##symbol = nullptr
+
+#if CUDART_VERSION >= 11030
+/* CUDA Driver functions loaded with cuGetProcAddress for versioning */
+DECLARE_CUDA_PFN(cuDeviceGet);
+DECLARE_CUDA_PFN(cuDeviceGetAttribute);
+DECLARE_CUDA_PFN(cuGetErrorString);
+DECLARE_CUDA_PFN(cuGetErrorName);
+/* enqueue.cc */
+DECLARE_CUDA_PFN(cuMemGetAddressRange);
+/* proxy.cc */
+DECLARE_CUDA_PFN(cuCtxCreate_v3020);
+DECLARE_CUDA_PFN(cuCtxDestroy);
+DECLARE_CUDA_PFN(cuCtxSetCurrent);
+#if CUDA_VERSION >= 11070
+/* transport/collNet.cc/net.cc*/
+DECLARE_CUDA_PFN(cuMemGetHandleForAddressRange); // DMA-BUF support
+#endif
+#endif
+
+/* CUDA Driver functions loaded with dlsym() */
+DECLARE_CUDA_PFN(cuInit);
+DECLARE_CUDA_PFN(cuDriverGetVersion);
+DECLARE_CUDA_PFN(cuGetProcAddress);
+
+static enum { cudaUninitialized, cudaInitializing, cudaInitialized, cudaError } cudaState = cudaUninitialized;
+
+#define CUDA_DRIVER_MIN_VERSION 11030
+
+static void *cudaLib;
+static int cudaDriverVersion;
+
+#if CUDART_VERSION >= 11030
+/*
+  Load the CUDA symbols
+ */
+static int cudaPfnFuncLoader(void) {
+  CUresult res;
+
+#define LOAD_SYM(symbol, ignore) do {                                   \
+    res = pfn_cuGetProcAddress(#symbol, (void **) (&pfn_##symbol), cudaDriverVersion, 0); \
+    if (res != 0) {                                                     \
+      if (!ignore) {                                                    \
+        WARN("Retrieve %s version %d failed with %d", #symbol, cudaDriverVersion, res); \
+        return ncclSystemError; }                                       \
+    } } while(0)
+
+  LOAD_SYM(cuGetErrorString, 0);
+  LOAD_SYM(cuGetErrorName, 0);
+  LOAD_SYM(cuDeviceGet, 0);
+  LOAD_SYM(cuDeviceGetAttribute, 0);
+  LOAD_SYM(cuMemGetAddressRange, 1);
+  LOAD_SYM(cuCtxCreate_v3020, 1);
+  LOAD_SYM(cuCtxDestroy, 1);
+  LOAD_SYM(cuCtxSetCurrent, 1);
+#if CUDA_VERSION >= 11070
+  LOAD_SYM(cuMemGetHandleForAddressRange, 1); // DMA-BUF support
+#endif
+  return ncclSuccess;
+}
+#endif
+
+ncclResult_t cudaLibraryInit(void) {
+  CUresult res;
+
+  if (cudaState == cudaInitialized)
+    return ncclSuccess;
+  if (cudaState == cudaError)
+    return ncclSystemError;
+
+  if (__sync_bool_compare_and_swap(&cudaState, cudaUninitialized, cudaInitializing) == false) {
+    // Another thread raced in front of us. Wait for it to be done.
+    while (cudaState == cudaInitializing) sched_yield();
+    return (cudaState == cudaInitialized) ? ncclSuccess : ncclSystemError;
+  }
+
+  /*
+   * Load CUDA driver library
+   */
+  char path[1024];
+  char *ncclCudaPath = getenv("NCCL_CUDA_PATH");
+  if (ncclCudaPath == NULL)
+    snprintf(path, 1024, "%s", "libcuda.so");
+  else
+    snprintf(path, 1024, "%s%s", ncclCudaPath, "libcuda.so");
+
+  cudaLib = dlopen(path, RTLD_LAZY);
+  if (cudaLib == NULL) {
+    WARN("Failed to find CUDA library in %s (NCCL_CUDA_PATH=%s)", ncclCudaPath, ncclCudaPath);
+    goto error;
+  }
+
+  /*
+   * Load initial CUDA functions
+   */
+
+  pfn_cuInit = (PFN_cuInit) dlsym(cudaLib, "cuInit");
+  if (pfn_cuInit == NULL) {
+    WARN("Failed to load CUDA missing symbol cuInit");
+    goto error;
+  }
+
+  pfn_cuDriverGetVersion = (PFN_cuDriverGetVersion) dlsym(cudaLib, "cuDriverGetVersion");
+  if (pfn_cuDriverGetVersion == NULL) {
+    WARN("Failed to load CUDA missing symbol cuDriverGetVersion");
+    goto error;
+  }
+
+  res = pfn_cuDriverGetVersion(&cudaDriverVersion);
+  if (res != 0) {
+    WARN("cuDriverGetVersion failed with %d", res);
+    goto error;
+  }
+
+  INFO(NCCL_INIT, "cudaDriverVersion %d", cudaDriverVersion);
+
+  if (cudaDriverVersion < CUDA_DRIVER_MIN_VERSION) {
+    // WARN("CUDA Driver version found is %d. Minimum requirement is %d", cudaDriverVersion, CUDA_DRIVER_MIN_VERSION);
+    // Silently ignore version check mismatch for backwards compatibility
+    goto error;
+  }
+
+  pfn_cuGetProcAddress = (PFN_cuGetProcAddress) dlsym(cudaLib, "cuGetProcAddress");
+  if (pfn_cuGetProcAddress == NULL) {
+    WARN("Failed to load CUDA missing symbol cuGetProcAddress");
+    goto error;
+  }
+
+  /*
+   * Required to initialize the CUDA Driver.
+   * Multiple calls of cuInit() will return immediately
+   * without making any relevant change
+   */
+  pfn_cuInit(0);
+
+#if CUDART_VERSION >= 11030
+  if (cudaPfnFuncLoader()) {
+    WARN("CUDA some PFN functions not found in the library");
+    goto error;
+  }
+#endif
+
+  cudaState = cudaInitialized;
+  return ncclSuccess;
+
+error:
+  cudaState = cudaError;
+  return ncclSystemError;
+}
+
+
@@ -57,7 +57,7 @@ ncclResult_t wrap_gdr_symbols(void) {

  if (__sync_bool_compare_and_swap(&gdrState, gdrUninitialized, gdrInitializing) == false) {
    // Another thread raced in front of us. Wait for it to be done.
-    while (gdrState == gdrInitializing) pthread_yield();
+    while (gdrState == gdrInitializing) sched_yield();
    return (gdrState == gdrInitialized) ? ncclSuccess : ncclSystemError;
  }

@@ -30,6 +30,8 @@ struct ibv_pd * (*ibv_internal_alloc_pd)(struct ibv_context *context);
 int (*ibv_internal_dealloc_pd)(struct ibv_pd *pd);
 struct ibv_mr * (*ibv_internal_reg_mr)(struct ibv_pd *pd, void *addr, size_t length, int access);
 struct ibv_mr * (*ibv_internal_reg_mr_iova2)(struct ibv_pd *pd, void *addr, size_t length, uint64_t iova, int access);
+/* DMA-BUF support */
+struct ibv_mr * (*ibv_internal_reg_dmabuf_mr)(struct ibv_pd *pd, uint64_t offset, size_t length, uint64_t iova, int fd, int access);
 int (*ibv_internal_dereg_mr)(struct ibv_mr *mr);
 struct ibv_cq * (*ibv_internal_create_cq)(struct ibv_context *context, int cqe, void *cq_context, struct ibv_comp_channel *channel, int comp_vector);
 int (*ibv_internal_destroy_cq)(struct ibv_cq *cq);
@@ -49,7 +51,7 @@ ncclResult_t wrap_ibv_symbols(void) {

  if (__sync_bool_compare_and_swap(&ibvState, ibvUninitialized, ibvInitializing) == false) {
    // Another thread raced in front of us. Wait for it to be done.
-    while (ibvState == ibvInitializing) pthread_yield();
+    while (ibvState == ibvInitializing) sched_yield();
    return (ibvState == ibvInitialized) ? ncclSuccess : ncclSystemError;
  }

@@ -98,6 +100,8 @@ ncclResult_t wrap_ibv_symbols(void) {
  LOAD_SYM(ibvhandle, "ibv_reg_mr", ibv_internal_reg_mr);
  // Cherry-pick the ibv_reg_mr_iova2 API from IBVERBS 1.8
  LOAD_SYM_VERSION(ibvhandle, "ibv_reg_mr_iova2", ibv_internal_reg_mr_iova2, "IBVERBS_1.8");
+  // Cherry-pick the ibv_reg_dmabuf_mr API from IBVERBS 1.12
+  LOAD_SYM_VERSION(ibvhandle, "ibv_reg_dmabuf_mr", ibv_internal_reg_dmabuf_mr, "IBVERBS_1.12");
  LOAD_SYM(ibvhandle, "ibv_dereg_mr", ibv_internal_dereg_mr);
  LOAD_SYM(ibvhandle, "ibv_create_cq", ibv_internal_create_cq);
  LOAD_SYM(ibvhandle, "ibv_destroy_cq", ibv_internal_destroy_cq);
@@ -126,6 +130,7 @@ teardown:
  ibv_internal_dealloc_pd = NULL;
  ibv_internal_reg_mr = NULL;
  ibv_internal_reg_mr_iova2 = NULL;
+  ibv_internal_reg_dmabuf_mr = NULL;
  ibv_internal_dereg_mr = NULL;
  ibv_internal_create_cq = NULL;
  ibv_internal_destroy_cq = NULL;
@@ -259,7 +264,7 @@ ncclResult_t wrap_ibv_dealloc_pd(struct ibv_pd *pd) { /*returns 0 on success, or
 }

 ncclResult_t wrap_ibv_reg_mr(struct ibv_mr **ret, struct ibv_pd *pd, void *addr, size_t length, int access) {
-  IBV_PTR_CHECK(ibv_internal_reg_mr, ibv_internal_reg_mr(pd, addr, length, access), *ret, NULL, "ibv_reg_mr");
+  IBV_PTR_CHECK_ERRNO(ibv_internal_reg_mr, ibv_internal_reg_mr(pd, addr, length, access), *ret, NULL, "ibv_reg_mr");
 }

 struct ibv_mr * wrap_direct_ibv_reg_mr(struct ibv_pd *pd, void *addr, size_t length, int access) {
@@ -275,7 +280,19 @@ ncclResult_t wrap_ibv_reg_mr_iova2(struct ibv_mr **ret, struct ibv_pd *pd, void
    return ncclInternalError;
  }
  if (ret == NULL) { return ncclSuccess; } // Assume dummy call
-  IBV_PTR_CHECK(ibv_internal_reg_mr_iova2, ibv_internal_reg_mr_iova2(pd, addr, length, iova, access), *ret, NULL, "ibv_reg_mr_iova2");
+  IBV_PTR_CHECK_ERRNO(ibv_internal_reg_mr_iova2, ibv_internal_reg_mr_iova2(pd, addr, length, iova, access), *ret, NULL, "ibv_reg_mr_iova2");
+}
+
+/* DMA-BUF support */
+ncclResult_t wrap_ibv_reg_dmabuf_mr(struct ibv_mr **ret, struct ibv_pd *pd, uint64_t offset, size_t length, uint64_t iova, int fd, int access) {
+  IBV_PTR_CHECK_ERRNO(ibv_internal_reg_dmabuf_mr, ibv_internal_reg_dmabuf_mr(pd, offset, length, iova, fd, access), *ret, NULL, "ibv_reg_dmabuf_mr");
+}
+
+struct ibv_mr * wrap_direct_ibv_reg_dmabuf_mr(struct ibv_pd *pd, uint64_t offset, size_t length, uint64_t iova, int fd, int access) {
+  if (ibv_internal_reg_dmabuf_mr == NULL) {
+    return NULL;
+  }
+  return ibv_internal_reg_dmabuf_mr(pd, offset, length, iova, fd, access);
 }

 ncclResult_t wrap_ibv_dereg_mr(struct ibv_mr *mr) { /*returns 0 on success, or the value of errno on failure (which indicates the failure reason)*/
@@ -0,0 +1,119 @@
+/*************************************************************************
+ * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#include "nccl.h"
+#include "debug.h"
+#include "rocmwrap.h"
+
+#include <dlfcn.h>
+
+#define DECLARE_ROCM_PFN(symbol) PFN_##symbol pfn_##symbol = nullptr
+
+DECLARE_ROCM_PFN(hsa_amd_portable_export_dmabuf); // DMA-BUF support
+
+/* ROCr Driver functions loaded with dlsym() */
+DECLARE_ROCM_PFN(hsa_init);
+DECLARE_ROCM_PFN(hsa_system_get_info);
+DECLARE_ROCM_PFN(hsa_status_string);
+
+static enum { hsaUninitialized, hsaInitializing, hsaInitialized, hsaError } hsaState = hsaUninitialized;
+
+static void *hsaLib;
+static uint16_t version_major, version_minor;
+
+ncclResult_t rocmLibraryInit(void) {
+  hsa_status_t res;
+
+  if (hsaState == hsaInitialized)
+    return ncclSuccess;
+  if (hsaState == hsaError)
+    return ncclSystemError;
+
+  if (__sync_bool_compare_and_swap(&hsaState, hsaUninitialized, hsaInitializing) == false) {
+    // Another thread raced in front of us. Wait for it to be done.
+    while (hsaState == hsaInitializing) sched_yield();
+    return (hsaState == hsaInitialized) ? ncclSuccess : ncclSystemError;
+  }
+
+  /*
+   * Load ROCr driver library
+   */
+  char path[1024];
+  char *ncclCudaPath = getenv("RCCL_ROCR_PATH");
+  if (ncclCudaPath == NULL)
+    snprintf(path, 1024, "%s", "libhsa-runtime64.so");
+  else
+    snprintf(path, 1024, "%s%s", ncclCudaPath, "libhsa-runtime64.so");
+
+  hsaLib = dlopen(path, RTLD_LAZY);
+  if (hsaLib == NULL) {
+    WARN("Failed to find ROCm runtime library in %s (RCCL_ROCR_PATH=%s)", ncclCudaPath, ncclCudaPath);
+    goto error;
+  }
+
+  /*
+   * Load initial ROCr functions
+   */
+
+  pfn_hsa_init = (PFN_hsa_init) dlsym(hsaLib, "hsa_init");
+  if (pfn_hsa_init == NULL) {
+    WARN("Failed to load ROCr missing symbol hsa_init");
+    goto error;
+  }
+
+  pfn_hsa_system_get_info = (PFN_hsa_system_get_info) dlsym(hsaLib, "hsa_system_get_info");
+  if (pfn_hsa_system_get_info == NULL) {
+    WARN("Failed to load ROCr missing symbol hsa_system_get_info");
+    goto error;
+  }
+
+  pfn_hsa_status_string = (PFN_hsa_status_string) dlsym(hsaLib, "hsa_status_string");
+  if (pfn_hsa_status_string == NULL) {
+    WARN("Failed to load ROCr missing symbol hsa_status_string");
+    goto error;
+  }
+
+  res = pfn_hsa_system_get_info(HSA_SYSTEM_INFO_VERSION_MAJOR, &version_major);
+  if (res != 0) {
+    WARN("pfn_hsa_system_get_info failed with %d", res);
+    goto error;
+  }
+  res = pfn_hsa_system_get_info(HSA_SYSTEM_INFO_VERSION_MINOR, &version_minor);
+  if (res != 0) {
+    WARN("pfn_hsa_system_get_info failed with %d", res);
+    goto error;
+  }
+
+  INFO(NCCL_INIT, "ROCr version %d.%d", version_major, version_minor);
+
+  //if (hsaDriverVersion < ROCR_DRIVER_MIN_VERSION) {
+    // WARN("ROCr Driver version found is %d. Minimum requirement is %d", hsaDriverVersion, ROCR_DRIVER_MIN_VERSION);
+    // Silently ignore version check mismatch for backwards compatibility
+    //goto error;
+  //}
+
+  pfn_hsa_amd_portable_export_dmabuf = (PFN_hsa_amd_portable_export_dmabuf) dlsym(hsaLib, "hsa_amd_portable_export_dmabuf");
+  if (pfn_hsa_amd_portable_export_dmabuf == NULL) {
+    WARN("Failed to load ROCr missing symbol hsa_amd_portable_export_dmabuf");
+    goto error;
+  }
+  /*
+   * Required to initialize the ROCr Driver.
+   * Multiple calls of hsa_init() will return immediately
+   * without making any relevant change
+   */
+  pfn_hsa_init();
+
+  hsaState = hsaInitialized;
+  return ncclSuccess;
+
+error:
+  hsaState = hsaError;
+  return ncclSystemError;
+}
+
+
@@ -59,15 +59,15 @@ ncclResult_t ncclShmOpen(char* shmPath, const int shmSize, void** shmPtr, void**

  NCCLCHECKGOTO(ncclShmSetup(shmPath, shmSize, &fd, &ptr, create), res, sysError);
  if (devShmPtr) {
-    CUDACHECKGOTO(hipHostRegister(ptr, shmSize, hipHostRegisterMapped), res, cudaError);
-    CUDACHECKGOTO(hipHostGetDevicePointer(devShmPtr, ptr, 0), res, cudaError);
+    CUDACHECKGOTO(hipHostRegister(ptr, shmSize, hipHostRegisterMapped), res, hipError_t);
+    CUDACHECKGOTO(hipHostGetDevicePointer(devShmPtr, ptr, 0), res, hipError_t);
  }

  *shmPtr = ptr;
  return ncclSuccess;
 sysError:
  WARN("Error while %s shared memory segment %s (size %d)", create ? "creating" : "attaching to", shmPath, shmSize);
-cudaError:
+hipError_t:
  if (fd != -1) close(fd);
  if (create) shm_unlink(shmPath);
  if (ptr != MAP_FAILED) munmap(ptr, shmSize);
@@ -15,6 +15,9 @@
 #include <vector>
 #include <utility>
 #include <unordered_set>
+#include <unistd.h>
+#include <sys/syscall.h>
+
 static std::vector<std::pair<int, std::unordered_set<std::string>>> clientPortPool;

 /* Format a string representation of a (union ncclSocketAddress *) socket address using getnameinfo()
@@ -337,9 +340,10 @@ ncclResult_t ncclSocketListen(struct ncclSocket* sock) {
 #endif
  }

-  /* make all new sockets non-blocking */
-  EQCHECK(flags = fcntl(fd, F_GETFL), -1);
-  SYSCHECK(fcntl(fd, F_SETFL, flags | O_NONBLOCK), "fcntl");
+  if (sock->asyncFlag) {
+    EQCHECK(flags = fcntl(fd, F_GETFL), -1);
+    SYSCHECK(fcntl(fd, F_SETFL, flags | O_NONBLOCK), "fcntl");
+  }

  // addr port should be 0 (Any port)
  SYSCHECK(bind(fd, &sock->addr.sa, salen), "bind");
@@ -378,7 +382,7 @@ static ncclResult_t getFdState(int fd, enum ncclSocketState* state) {
      SYSCHECK(getsockopt(fd, SOL_SOCKET, SO_ERROR, (void*)&ret, &rlen), "getsockopt");
    }

-    if (ret == EINPROGRESS)
+    if (ret == EINPROGRESS || ret == ECONNREFUSED)
      *state = ncclSocketConnecting;
    else if (ret == 0)
      *state = ncclSocketConnected;
@@ -414,10 +418,12 @@ ncclResult_t ncclSocketConnect(struct ncclSocket* sock, int portReuse) {

  const int one = 1;
  SYSCHECK(setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, (char*)&one, sizeof(int)), "setsockopt");
-
+  
  /* support non-blocking socket; by default, the socket is non-blocking */
-  EQCHECK(flags = fcntl(fd, F_GETFL), -1);
-  SYSCHECK(fcntl(fd, F_SETFL, flags | O_NONBLOCK), "fcntl");
+  if (sock->asyncFlag) {
+    EQCHECK(flags = fcntl(fd, F_GETFL), -1);
+    SYSCHECK(fcntl(fd, F_SETFL, flags | O_NONBLOCK), "fcntl");
+  }

  /*  const int bufsize = 128*1024;
    SYSCHECK(setsockopt(fd, SOL_SOCKET, SO_SNDBUF, (char*)&bufsize, sizeof(int)), "setsockopt");
@@ -458,31 +464,26 @@ ncclResult_t ncclSocketConnect(struct ncclSocket* sock, int portReuse) {
  int timedout_retries = 0;
  int refused_retries = 0;
 retry:
-  /* async connect; abort when error happens and abortFlag is present. */
+  /* blocking/non-blocking connect() is determined by asyncFlag. */
  ret = connect(fd, &sock->addr.sa, salen);

-  if (errno == EAGAIN || (errno == ECONNREFUSED && ++refused_retries < RETRY_REFUSED_TIMES) ||
-    (errno == ETIMEDOUT && ++timedout_retries < RETRY_TIMEDOUT_TIMES)) {
-    if (refused_retries % 1000 == 0) INFO(NCCL_ALL, "Call to connect returned %s, retrying", strerror(errno));
+  if (!sock->asyncFlag && (errno == EAGAIN || (errno == ECONNREFUSED && ++refused_retries < RETRY_REFUSED_TIMES) ||
+    (errno == ETIMEDOUT && ++timedout_retries < RETRY_TIMEDOUT_TIMES))) {
+    if (errno == ECONNREFUSED && refused_retries % 1000 == 0) INFO(NCCL_ALL, "Call to connect returned %s, retrying", strerror(errno));
    usleep(SLEEP_INT);
    goto retry;
-  } else if (errno == EINPROGRESS && !sock->asyncFlag) {
-    enum ncclSocketState state;
-    do {
-      if (sock->abortFlag) NEQCHECK(*sock->abortFlag, 0);
-      NCCLCHECK(getFdState(fd, &state));
-    } while (state == ncclSocketConnecting);
-    EQCHECK(state, ncclSocketError);
-    ret = 0;
  }

-  if (ret == 0 || (errno == EINPROGRESS && sock->asyncFlag)) {
+  /* If connect() fails with errno == EAGAIN/EINPROGRESS/ETIMEDOUT, we may want to try connect again.
+   * However, it can return EISCONN instead of success which indicates connection is built up in
+   * background already. No need to call connect() again. */
+  if (ret == 0 || ((errno == EINPROGRESS || errno == ECONNREFUSED) && sock->asyncFlag) || errno == EISCONN) {
    sock->fd = fd;
    return ncclSuccess;
  }

  WARN("Net : Connect to %s failed : %s", ncclSocketToString(&sock->addr, line), strerror(errno));
-  return ncclSystemError;
+  return ncclRemoteError;
 }

 ncclResult_t ncclSocketAccept(struct ncclSocket* sock, struct ncclSocket* listenSocket) {
@@ -535,7 +536,7 @@ static ncclResult_t ncclSocketProgressOpt(int op, struct ncclSocket* sock, void*
    if (bytes == -1) {
      if (errno != EINTR && errno != EWOULDBLOCK && errno != EAGAIN) {
        WARN("Net : Call to recv from %s failed : %s", ncclSocketToString(&sock->addr, line), strerror(errno));
-        return ncclSystemError;
+        return ncclRemoteError;
      } else {
        bytes = 0;
      }
@@ -555,7 +556,7 @@ ncclResult_t ncclSocketProgress(int op, struct ncclSocket* sock, void* ptr, int
  if (closed) {
    char line[SOCKET_NAME_MAXLEN+1];
    WARN("Net : Connection closed by remote peer %s", ncclSocketToString(&sock->addr, line, 0));
-    return ncclSystemError;
+    return ncclRemoteError;
  }
  return ncclSuccess;
 }
@@ -0,0 +1,273 @@
+/*************************************************************************
+ * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#include "strongstream.h"
+#include "checks.h"
+#include "param.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+ncclResult_t ncclCudaGetCapturingGraph(
+    struct ncclCudaGraph* graph, hipStream_t stream
+  ) {
+  #if CUDART_VERSION >= 11030
+    thread_local int driver = -1;
+    if (driver == -1) {
+      CUDACHECK(cudaDriverGetVersion(&driver));
+    }
+    if (driver < 11030) {
+      cudaStreamCaptureStatus status;
+      unsigned long long gid;
+      graph->graph = nullptr;
+      CUDACHECK(cudaStreamGetCaptureInfo(stream, &status, &gid));
+      if (status != cudaStreamCaptureStatusNone) {
+        WARN("The installed CUDA driver is older than the minimum version (R465) required for NCCL's CUDA Graphs support");
+        return ncclInvalidUsage;
+      }
+    } else {
+      cudaStreamCaptureStatus status;
+      unsigned long long gid;
+      CUDACHECK(cudaStreamGetCaptureInfo_v2(stream, &status, &gid, &graph->graph, nullptr, nullptr));
+      if (status != cudaStreamCaptureStatusActive) {
+        graph->graph = nullptr;
+        gid = ULLONG_MAX;
+      }
+      graph->graphId = gid;
+    }
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, hipHostFn_t fn, void* arg) {
+  #if CUDART_VERSION >= 11030
+    cudaUserObject_t object;
+    CUDACHECK(cudaUserObjectCreate(
+      &object, arg, fn, /*initialRefcount=*/1, cudaUserObjectNoDestructorSync
+    ));
+    // Hand over ownership to CUDA Graph
+    CUDACHECK(cudaGraphRetainUserObject(graph.graph, object, 1, cudaGraphUserObjectMove));
+    return ncclSuccess;
+  #else
+    return ncclInvalidUsage;
+  #endif
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+ncclResult_t ncclStrongStreamConstruct(struct ncclStrongStream* ss) {
+  CUDACHECK(hipStreamCreateWithFlags(&ss->stream, hipStreamNonBlocking));
+  CUDACHECK(hipEventCreateWithFlags(&ss->event, hipEventDisableTiming));
+  #if CUDART_VERSION >= 11030
+    ss->node = nullptr;
+    ss->graphId = (1ull<<(8*sizeof(long long)-1))-1;
+    ss->eventIsLagging = 0;
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamDestruct(struct ncclStrongStream* ss) {
+  #if CUDART_VERSION >= 11030
+    CUDACHECK(cudaEventDestroy(ss->event));
+  #endif
+  CUDACHECK(hipStreamDestroy(ss->stream));
+  return ncclSuccess;
+}
+
+NCCL_PARAM(GraphMixingSupport, "GRAPH_MIXING_SUPPORT", 1)
+
+ncclResult_t ncclStrongStreamAcquire(
+    struct ncclCudaGraph graph, struct ncclStrongStream* ss
+  ) {
+  #if CUDART_VERSION >= 11030
+    bool mixing = ncclParamGraphMixingSupport();
+    if (graph.graph == nullptr) {
+      if (mixing && ncclStrongStreamEverCaptured(ss)) {
+        CUDACHECK(cudaStreamWaitEvent(ss->stream, ss->event, 0));
+        ss->eventIsLagging = 0;
+      }
+    } else {
+      if (ss->graphId != graph.graphId) {
+        if (mixing && ss->eventIsLagging) {
+          // Can only be here if previous release was for uncaptured work that
+          // elided updating the event because no capture had yet occurred.
+          CUDACHECK(cudaStreamWaitEvent(ss->stream, ss->event, 0));
+          CUDACHECK(cudaEventRecord(ss->event, ss->stream));
+        }
+        ss->graphId = graph.graphId;
+        ss->eventIsLagging = 0;
+        if (mixing) {
+          CUDACHECK(cudaGraphAddEventWaitNode(&ss->node, graph.graph, nullptr, 0, ss->event));
+        } else {
+          CUDACHECK(cudaGraphAddEmptyNode(&ss->node, graph.graph, nullptr, 0));
+        }
+      }
+    }
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamAcquireUncaptured(struct ncclStrongStream* ss) {
+  #if CUDART_VERSION >= 11030
+    bool mixing = ncclParamGraphMixingSupport();
+    if (mixing && ncclStrongStreamEverCaptured(ss)) {
+      CUDACHECK(cudaStreamWaitEvent(ss->stream, ss->event, 0));
+    }
+    ss->eventIsLagging = 1; // Assume the caller is going to add work to stream.
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamRelease(struct ncclCudaGraph graph, struct ncclStrongStream* ss) {
+  #if CUDART_VERSION >= 11030
+    bool mixing = ncclParamGraphMixingSupport();
+    if (mixing && ss->eventIsLagging) {
+      if (graph.graph == nullptr) {
+        if (ncclStrongStreamEverCaptured(ss)) {
+          CUDACHECK(cudaEventRecord(ss->event, ss->stream));
+          ss->eventIsLagging = 0;
+        }
+      } else {
+        CUDACHECK(cudaGraphAddEventRecordNode(&ss->node, graph.graph, &ss->node, 1, ss->event));
+        ss->eventIsLagging = 0;
+      }
+    }
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamLaunchHost(
+    struct ncclCudaGraph graph, struct ncclStrongStream* ss, hipHostFn_t fn, void* arg
+  ) {
+  #if CUDART_VERSION >= 11030
+    if (graph.graph == nullptr) {
+      CUDACHECK(cudaLaunchHostFunc(ss->stream, fn, arg));
+    } else {
+      cudaHostNodeParams p;
+      p.fn = fn;
+      p.userData = arg;
+      CUDACHECK(cudaGraphAddHostNode(&ss->node, graph.graph, &ss->node, 1, &p));
+    }
+    ss->eventIsLagging = 1;
+  #else
+    //CUDACHECK(hipLaunchHostFunc(ss->stream, fn, arg));
+    CUDACHECK(hipStreamAddCallback(ss->stream, (hipStreamCallback_t)fn, arg, 0));
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamLaunchKernel(
+    struct ncclCudaGraph graph, struct ncclStrongStream* ss,
+    void* fn, dim3 grid, dim3 block, void* args[], size_t sharedMemBytes
+  ) {
+  #if CUDART_VERSION >= 11030
+    if (graph.graph == nullptr) {
+      CUDACHECK(cudaLaunchKernel(fn, grid, block, args, sharedMemBytes, ss->stream));
+    } else {
+      cudaGraphNode_t tip = ss->node;
+      cudaKernelNodeParams p;
+      p.func = fn;
+      p.gridDim = grid;
+      p.blockDim = block;
+      p.kernelParams = args;
+      p.sharedMemBytes = sharedMemBytes;
+      p.extra = nullptr;
+      CUDACHECK(cudaGraphAddKernelNode(&ss->node, graph.graph, &tip, 1, &p));
+    }
+    ss->eventIsLagging = 1;
+  #else
+    CUDACHECK(hipLaunchKernel(fn, grid, block, args, sharedMemBytes, ss->stream));
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamWaitStream(
+    struct ncclCudaGraph graph, struct ncclStrongStream* a, struct ncclStrongStream* b
+  ) {
+  #if CUDART_VERSION >= 11030
+    if (graph.graph == nullptr) {
+      if (b->eventIsLagging) {
+        b->eventIsLagging = 0;
+        CUDACHECK(cudaEventRecord(b->event, b->stream));
+      }
+      CUDACHECK(cudaStreamWaitEvent(a->stream, b->event, 0));
+      a->eventIsLagging = 1;
+    } else {
+      cudaGraphNode_t pair[2] = {a->node, b->node};
+      CUDACHECK(cudaGraphAddEmptyNode(&a->node, graph.graph, pair, 2));
+    }
+  #else
+    CUDACHECK(hipEventRecord(b->event, b->stream));
+    CUDACHECK(hipStreamWaitEvent(a->stream, b->event, 0));
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamWaitStream(
+    struct ncclCudaGraph graph, struct ncclStrongStream* a, hipStream_t b
+  ) {
+  #if CUDART_VERSION >= 11030
+    if (graph.graph == nullptr) {
+      CUDACHECK(cudaEventRecord(a->event, b));
+      CUDACHECK(cudaStreamWaitEvent(a->stream, a->event, 0));
+      // We used a->event to record b so it no longer reflects anything about a.
+      a->eventIsLagging = 1;
+    } else {
+      cudaStreamCaptureStatus status;
+      unsigned long long gid1;
+      cudaGraphNode_t const* deps;
+      size_t depN = 0;
+      CUDACHECK(cudaStreamGetCaptureInfo_v2(b, &status, &gid1, nullptr, &deps, &depN));
+      if (status != cudaStreamCaptureStatusActive || graph.graphId != gid1) {
+        WARN("Stream is not being captured by the expected graph.");
+        return ncclInvalidUsage;
+      }
+      if (depN > 0 && (depN > 1 || deps[0] != a->node)) {
+        cudaGraphNode_t tie;
+        if (depN == 1) {
+          tie = deps[0];
+        } else {
+          CUDACHECK(cudaGraphAddEmptyNode(&tie, graph.graph, deps, depN));
+        }
+        cudaGraphNode_t pair[2] = {a->node, tie};
+        CUDACHECK(cudaGraphAddEmptyNode(&a->node, graph.graph, pair, 2));
+      }
+      // a->eventIsLagging doesn't change since we are just updating the
+      // dependencies of a->node.
+    }
+  #else
+    CUDACHECK(hipEventRecord(a->event, b));
+    CUDACHECK(hipStreamWaitEvent(a->stream, a->event, 0));
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamWaitStream(
+    struct ncclCudaGraph graph, hipStream_t a, struct ncclStrongStream* b
+  ) {
+  #if CUDART_VERSION >= 11030
+    if (graph.graph == nullptr) {
+      if (b->eventIsLagging) {
+        b->eventIsLagging = 0;
+        CUDACHECK(cudaEventRecord(b->event, b->stream));
+      }
+      CUDACHECK(cudaStreamWaitEvent(a, b->event, 0));
+    } else {
+      CUDACHECK(cudaStreamUpdateCaptureDependencies(a, &b->node, 1, cudaStreamAddCaptureDependencies));
+    }
+  #else
+    CUDACHECK(hipEventRecord(b->event, b->stream));
+    CUDACHECK(hipStreamWaitEvent(a, b->event, 0));
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamSynchronize(struct ncclStrongStream* ss) {
+  #if CUDART_VERSION >= 11030
+    CUDACHECK(cudaStreamWaitEvent(ss->stream, ss->event, 0));
+  #endif
+  CUDACHECK(hipStreamSynchronize(ss->stream));
+  return ncclSuccess;
+}
@@ -11,6 +11,8 @@
 #include "nvmlwrap.h"
 #include <hip/hip_runtime.h>

+#include <stdlib.h>
+
 // Get current Compute Capability
 int ncclCudaCompCap() {
  int cudaDev;
@@ -192,3 +194,102 @@ bool matchIfList(const char* string, int port, struct netIf* ifList, int listSiz
  }
  return false;
 }
+
+__thread struct ncclThreadSignal ncclThreadSignalLocalInstance = ncclThreadSignalStaticInitializer();
+
+void* ncclMemoryStack::allocateSpilled(struct ncclMemoryStack* me, size_t size, size_t align) {
+  // `me->hunks` points to the top of the stack non-empty hunks. Hunks above
+  // this (reachable via `->above`) are empty.
+  struct Hunk* top = me->topFrame.hunk;
+  size_t mallocSize = 0;
+
+  // If we have lots of space left in hunk but that wasn't enough then we'll
+  // allocate the object unhunked.
+  if (me->topFrame.end - me->topFrame.bumper >= 8<<10)
+    goto unhunked;
+
+  // If we have another hunk (which must be empty) waiting above this one and
+  // the object fits then use that.
+  if (top && top->above) {
+    struct Hunk* top1 = top->above;
+    uintptr_t uobj = (reinterpret_cast<uintptr_t>(top1) + sizeof(struct Hunk) + align-1) & -uintptr_t(align);
+    if (uobj + size <= reinterpret_cast<uintptr_t>(top1) + top1->size) {
+      me->topFrame.hunk = top1;
+      me->topFrame.bumper = uobj + size;
+      me->topFrame.end = reinterpret_cast<uintptr_t>(top1) + top1->size;
+      return reinterpret_cast<void*>(uobj);
+    }
+  }
+
+  { // If the next hunk we're going to allocate wouldn't be big enough but the
+    // Unhunk proxy fits in the current hunk then go allocate as unhunked.
+    size_t nextSize = (top ? top->size : 0) + (64<<10);
+    constexpr size_t maxAlign = 64;
+    if (nextSize < sizeof(struct Hunk) + maxAlign + size) {
+      uintptr_t uproxy = (me->topFrame.bumper + alignof(Unhunk)-1) & -uintptr_t(alignof(Unhunk));
+      if (uproxy + sizeof(struct Unhunk) <= me->topFrame.end)
+        goto unhunked;
+    }
+
+    // At this point we must need another hunk, either to fit the object
+    // itself or its Unhunk proxy.
+    mallocSize = nextSize;
+    INFO(NCCL_ALLOC, "%s:%d memory stack hunk malloc(%llu)", __FILE__, __LINE__, (unsigned long long)mallocSize);
+    struct Hunk *top1 = (struct Hunk*)malloc(mallocSize);
+    if (top1 == nullptr) goto malloc_exhausted;
+    top1->size = nextSize;
+    top1->above = nullptr;
+    if (top) top->above = top1;
+    top = top1;
+    me->topFrame.hunk = top;
+    me->topFrame.end = reinterpret_cast<uintptr_t>(top) + nextSize;
+    me->topFrame.bumper = reinterpret_cast<uintptr_t>(top) + sizeof(struct Hunk);
+  }
+
+  { // Try to fit object in the new top hunk.
+    uintptr_t uobj = (me->topFrame.bumper + align-1) & -uintptr_t(align);
+    if (uobj + size <= me->topFrame.end) {
+      me->topFrame.bumper = uobj + size;
+      return reinterpret_cast<void*>(uobj);
+    }
+  }
+
+unhunked:
+  { // We need to allocate the object out-of-band and put an Unhunk proxy in-band
+    // to keep track of it.
+    uintptr_t uproxy = (me->topFrame.bumper + alignof(Unhunk)-1) & -uintptr_t(alignof(Unhunk));
+    Unhunk* proxy = reinterpret_cast<Unhunk*>(uproxy);
+    me->topFrame.bumper = uproxy + sizeof(Unhunk);
+    proxy->next = me->topFrame.unhunks;
+    me->topFrame.unhunks = proxy;
+    mallocSize = size;
+    proxy->obj = malloc(mallocSize);
+    INFO(NCCL_ALLOC, "%s:%d memory stack non-hunk malloc(%llu)", __FILE__, __LINE__, (unsigned long long)mallocSize);
+    if (proxy->obj == nullptr) goto malloc_exhausted;
+    return proxy->obj;
+  }
+
+malloc_exhausted:
+  WARN("%s:%d Unrecoverable error detected: malloc(size=%llu) returned null.", __FILE__, __LINE__, (unsigned long long)mallocSize);
+  abort();
+}
+
+void ncclMemoryStackDestruct(struct ncclMemoryStack* me) {
+  // Free unhunks first because both the frames and unhunk proxies lie within the hunks.
+  struct ncclMemoryStack::Frame* f = &me->topFrame;
+  while (f != nullptr) {
+    struct ncclMemoryStack::Unhunk* u = f->unhunks;
+    while (u != nullptr) {
+      free(u->obj);
+      u = u->next;
+    }
+    f = f->below;
+  }
+  // Free hunks
+  struct ncclMemoryStack::Hunk* h = me->stub.above;
+  while (h != nullptr) {
+    struct ncclMemoryStack::Hunk *h1 = h->above;
+    free(h);
+    h = h1;
+  }
+}