2.13.4-1

Optimize CUDA graph launch; avoid launching a CPU callback for intra-node operations. Simplify kernel common code to improve the latency of send/recv operations. Strengthen CUDA streams semantics. Change NET API to v6, to add dmabuf support. Add ncclGetLastError() function. Add ncclRemoteError code and use it for remote network errors. Support the use of a different NCCL_NET parameter per communicator. Add support for SHM and P2P transfers using cudaMemcpy. [ROCm/rccl commit: 19ab67d172]
2022-05-24 02:02:31 -07:00
parent 1c5734046d
commit 91154e8df9
62 changed files with 4787 additions and 2496 deletions
@@ -44,12 +44,7 @@ ncclResult_t ArgsCheck(struct ncclInfo* info) {
    return ncclInvalidArgument;
  }
  // Type is OK, compute nbytes. Convert Allgather/Broadcast/P2P calls to chars.
-  info->nBytes = info->count * ncclTypeSize(info->datatype);
-  if (info->coll == ncclFuncAllGather || info->coll == ncclFuncBroadcast) {
-    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)*/
@@ -332,9 +332,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");
@@ -373,7 +374,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;
@@ -409,10 +410,12 @@ ncclResult_t ncclSocketConnect(struct ncclSocket* sock) {

  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");
@@ -424,31 +427,26 @@ ncclResult_t ncclSocketConnect(struct ncclSocket* sock) {
  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) {
@@ -501,7 +499,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;
      }
@@ -521,7 +519,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,272 @@
+/*************************************************************************
+ * 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, cudaStream_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, cudaHostFn_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(cudaStreamCreateWithFlags(&ss->stream, cudaStreamNonBlocking));
+  CUDACHECK(cudaEventCreateWithFlags(&ss->event, cudaEventDisableTiming));
+  #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(cudaStreamDestroy(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, cudaHostFn_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(cudaLaunchHostFunc(ss->stream, fn, arg));
+  #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(cudaLaunchKernel(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(cudaEventRecord(b->event, b->stream));
+    CUDACHECK(cudaStreamWaitEvent(a->stream, b->event, 0));
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamWaitStream(
+    struct ncclCudaGraph graph, struct ncclStrongStream* a, cudaStream_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(cudaEventRecord(a->event, b));
+    CUDACHECK(cudaStreamWaitEvent(a->stream, a->event, 0));
+  #endif
+  return ncclSuccess;
+}
+
+ncclResult_t ncclStrongStreamWaitStream(
+    struct ncclCudaGraph graph, cudaStream_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(cudaEventRecord(b->event, b->stream));
+    CUDACHECK(cudaStreamWaitEvent(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(cudaStreamSynchronize(ss->stream));
+  return ncclSuccess;
+}
@@ -9,6 +9,8 @@

 #include "nvmlwrap.h"

+#include <stdlib.h>
+
 // Get current Compute Capability
 int ncclCudaCompCap() {
  int cudaDev;
@@ -190,3 +192,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;
+  }
+}