diff --git a/3rd-party/efa-gda/CUDA/host/efa_cuda_dp.cpp b/3rd-party/efa-gda/CUDA/host/efa_cuda_dp.cpp index 5dd208f1bd..5ea1bce50b 100644 --- a/3rd-party/efa-gda/CUDA/host/efa_cuda_dp.cpp +++ b/3rd-party/efa-gda/CUDA/host/efa_cuda_dp.cpp @@ -4,7 +4,7 @@ #include #include #include -#include +#include "nccl_ofi_cuda.h" #include "efa_cuda_dp.h" #include "efa_cuda_dp_types.h" @@ -26,7 +26,7 @@ static bool is_buf_cleared(void *buf, size_t len) struct efa_cuda_cq *efa_cuda_create_cq(struct efa_cuda_cq_attrs *attrs, uint32_t inlen) { - cudaError_t cuda_err; + int ret; if (inlen > sizeof(*attrs) && !is_ext_cleared(attrs, inlen)) { printf("Incompatible attributes struct\n"); @@ -39,10 +39,9 @@ struct efa_cuda_cq *efa_cuda_create_cq(struct efa_cuda_cq_attrs *attrs, uint32_t } efa_cuda_cq *d_cq; - cuda_err = cudaMalloc(&d_cq, sizeof(efa_cuda_cq)); - if (cuda_err != cudaSuccess) { - printf("Failed to allocate device memory for cq: %s\n", - cudaGetErrorString(cuda_err)); + ret = nccl_net_ofi_gpu_mem_alloc((void **)&d_cq, sizeof(efa_cuda_cq)); + if (ret != 0) { + printf("Failed to allocate device memory for cq: %d\n", ret); return nullptr; } @@ -55,11 +54,10 @@ struct efa_cuda_cq *efa_cuda_create_cq(struct efa_cuda_cq_attrs *attrs, uint32_t h_cq.cc = 0; h_cq.phase = 1; - cuda_err = cudaMemcpy(d_cq, &h_cq, sizeof(efa_cuda_cq), cudaMemcpyHostToDevice); - if (cuda_err != cudaSuccess) { - cudaFree(d_cq); - printf("Failed to copy cq to device: %s\n", - cudaGetErrorString(cuda_err)); + ret = nccl_net_ofi_gpu_mem_copy_host_to_device(d_cq, &h_cq, sizeof(efa_cuda_cq)); + if (ret != 0) { + nccl_net_ofi_gpu_mem_free(d_cq); + printf("Failed to copy cq to device: %d\n", ret); return nullptr; } @@ -68,12 +66,13 @@ struct efa_cuda_cq *efa_cuda_create_cq(struct efa_cuda_cq_attrs *attrs, uint32_t void efa_cuda_destroy_cq(efa_cuda_cq *d_cq) { - cudaFree(d_cq); + if (d_cq) + nccl_net_ofi_gpu_mem_free(d_cq); } struct efa_cuda_qp *efa_cuda_create_qp(struct efa_cuda_qp_attrs *attrs, uint32_t inlen) { - cudaError_t cuda_err; + int ret; if ((inlen > sizeof(*attrs) && !is_ext_cleared(attrs, inlen)) || attrs->reserved) { @@ -88,10 +87,9 @@ struct efa_cuda_qp *efa_cuda_create_qp(struct efa_cuda_qp_attrs *attrs, uint32_t } efa_cuda_qp *d_qp; - cuda_err = cudaMalloc(&d_qp, sizeof(efa_cuda_qp)); - if (cuda_err != cudaSuccess) { - printf("Failed to allocate device memory for qp: %s\n", - cudaGetErrorString(cuda_err)); + ret = nccl_net_ofi_gpu_mem_alloc((void **)&d_qp, sizeof(efa_cuda_qp)); + if (ret != 0) { + printf("Failed to allocate device memory for qp: %d\n", ret); return nullptr; } @@ -124,11 +122,10 @@ struct efa_cuda_qp *efa_cuda_create_qp(struct efa_cuda_qp_attrs *attrs, uint32_t h_qp.rq.wq.pc = 0; h_qp.rq.wq.phase = 1; - cuda_err = cudaMemcpy(d_qp, &h_qp, sizeof(efa_cuda_qp), cudaMemcpyHostToDevice); - if (cuda_err != cudaSuccess) { - cudaFree(d_qp); - printf("Failed to copy qp to device: %s\n", - cudaGetErrorString(cuda_err)); + ret = nccl_net_ofi_gpu_mem_copy_host_to_device(d_qp, &h_qp, sizeof(efa_cuda_qp)); + if (ret != 0) { + nccl_net_ofi_gpu_mem_free(d_qp); + printf("Failed to copy qp to device: %d\n", ret); return nullptr; } @@ -138,7 +135,7 @@ struct efa_cuda_qp *efa_cuda_create_qp(struct efa_cuda_qp_attrs *attrs, uint32_t void efa_cuda_destroy_qp(struct efa_cuda_qp *d_qp) { if (d_qp) - cudaFree(d_qp); + nccl_net_ofi_gpu_mem_free(d_qp); } int efa_cuda_get_version(int *major, int *minor, int *subminor) diff --git a/include/nccl_ofi_cuda.h b/include/nccl_ofi_cuda.h index 1b43c8aeb0..0ed449da5b 100644 --- a/include/nccl_ofi_cuda.h +++ b/include/nccl_ofi_cuda.h @@ -6,6 +6,8 @@ #ifndef NCCL_OFI_CUDA_H_ #define NCCL_OFI_CUDA_H_ +#include + int nccl_net_ofi_gpu_init(void); /* @@ -22,12 +24,19 @@ int nccl_net_ofi_gpu_init(void); int nccl_net_ofi_get_gpu_device_for_addr(void *data, int *dev_id); /* - * @brief Get / set the CUDA device for the calling thread. set_device also - * establishes the thread's primary CUDA context, which bare worker - * threads (e.g. the gdrcopy signal worker) otherwise lack. + * Bind / read the calling thread's current CUDA context. No ownership taken; + * the app (or NCCL) owns the context lifetime. + */ + +/* + * @brief Make ctx the calling thread's current context; NULL detaches. + */ +int nccl_net_ofi_gpu_set_current_context(CUcontext ctx); + +/* + * @brief Return the calling thread's current context; *ctx is NULL if none. */ -int nccl_net_ofi_gpu_get_device(int *dev_id); -int nccl_net_ofi_gpu_set_device(int dev_id); +int nccl_net_ofi_gpu_get_current_context(CUcontext *ctx); /* * @brief Retrieve the base address and size of the VMM segment (cuMemCreate diff --git a/include/rdma/gin/nccl_ofi_gin.h b/include/rdma/gin/nccl_ofi_gin.h index 8caa59408f..84b997ee23 100644 --- a/include/rdma/gin/nccl_ofi_gin.h +++ b/include/rdma/gin/nccl_ofi_gin.h @@ -13,6 +13,7 @@ #include "nccl_ofi_spsc_ring.h" #include "nccl_ofi.h" +#include "nccl_ofi_cuda.h" #include "nccl_ofi_gdrcopy.h" #include "nccl_ofi_tracepoint.h" @@ -652,7 +653,10 @@ class nccl_ofi_rdma_gin_put_comm : public nccl_ofi_gin_put_comm_t { std::unique_ptr metadata_fl; int dev; - int gdrcopy_cuda_dev = -1; /* CUDA device the worker binds to */ + /* CUDA context of the thread that created this comm. Signal-segment + discovery makes it current when it runs on a thread that has none + (NCCL's GIN progress threads are bare std::threads). */ + CUcontext gdrcopy_cuda_ctx = nullptr; /* --- TIER 2: Receiver side — every CQ completion --- */ /* Rail pinned across an aggregated iputSignal sequence: when an op is diff --git a/src/Makefile.am b/src/Makefile.am index b97ae1c439..3795bd3131 100644 --- a/src/Makefile.am +++ b/src/Makefile.am @@ -99,6 +99,7 @@ libefa_cuda_dp_la_SOURCES = \ $(top_srcdir)/3rd-party/efa-gda/CUDA/host/efa_cuda_dp.cpp libefa_cuda_dp_la_CPPFLAGS = \ + -I$(abs_top_srcdir)/include \ $(CUDA_CPPFLAGS) \ -isystem $(abs_top_srcdir)/3rd-party/efa-gda/CUDA/common \ -isystem $(abs_top_srcdir)/3rd-party/efa-gda/CUDA/host @@ -113,12 +114,8 @@ AM_CPPFLAGS += \ -isystem $(abs_top_srcdir)/3rd-party/efa-gda/CUDA/common \ -isystem $(abs_top_srcdir)/3rd-party/efa-gda/CUDA/host -# Make the private host implementation available to the internal plugin for -# later GDAKI integration. It uses CUDA Runtime APIs, so add CUDART only for -# GDAKI builds and preserve the existing non-GDAKI dependency set. libinternal_plugin_la_LIBADD += \ - libefa_cuda_dp.la \ - $(CUDA_RUNTIME_LIBS) + libefa_cuda_dp.la endif lib_LTLIBRARIES = diff --git a/src/nccl_ofi_cuda.cpp b/src/nccl_ofi_cuda.cpp index 07ff69d544..42d5bd5863 100644 --- a/src/nccl_ofi_cuda.cpp +++ b/src/nccl_ofi_cuda.cpp @@ -6,6 +6,7 @@ #include "config.h" #include +#include #include #include #include @@ -73,7 +74,11 @@ static std::unique_ptr cudaruntime_lib; /* Use driver APIs wherever possible - they are version-stable */ DECLARE_CUDA_FUNCTION(cuDriverGetVersion, 2020); +DECLARE_CUDA_FUNCTION(cuGetErrorString, 6000); +DECLARE_CUDA_FUNCTION(cuGetErrorName, 6000); DECLARE_CUDA_FUNCTION(cuCtxGetDevice, 2000); +DECLARE_CUDA_FUNCTION(cuCtxSetCurrent, 4000); +DECLARE_CUDA_FUNCTION(cuCtxGetCurrent, 4000); DECLARE_CUDA_FUNCTION(cuDeviceGetAttribute, 2000); #if HAVE_CUDA_GDRFLUSH_SUPPORT DECLARE_CUDA_FUNCTION(cuFlushGPUDirectRDMAWrites, 11030); @@ -84,6 +89,7 @@ DECLARE_CUDA_FUNCTION(cuMemGetHandleForAddressRange, 11070); DECLARE_CUDA_FUNCTION(cuPointerGetAttributes, 7000); DECLARE_CUDA_FUNCTION(cuMemAlloc, 3020); DECLARE_CUDA_FUNCTION(cuMemFree, 3020); +DECLARE_CUDA_FUNCTION(cuMemsetD8, 3020); DECLARE_CUDA_FUNCTION(cuMemcpyHtoDAsync, 3020); DECLARE_CUDA_FUNCTION(cuStreamCreate, 2000); DECLARE_CUDA_FUNCTION(cuStreamSynchronize, 2000); @@ -104,6 +110,31 @@ DECLARE_CUDA_FUNCTION(cuMemRetainAllocationHandle, 11000); DECLARE_CUDA_FUNCTION(cuMemGetAllocationPropertiesFromHandle, 10020); DECLARE_CUDA_FUNCTION(cuThreadExchangeStreamCaptureMode, 10010); +/* + * Driver-API equivalent of cudaGetErrorString(): renders a CUresult as + * "CUDA_ERROR_OUT_OF_MEMORY (out of memory)". Falls back to placeholders when + * a code cannot be translated, so the result is always safe to log. + */ +static const char *nccl_net_ofi_cuda_error_string(CUresult res) +{ + static thread_local char buf[256]; + const char *name = NULL; + const char *desc = NULL; + + if (pfn_cuGetErrorName == NULL || pfn_cuGetErrorName(res, &name) != CUDA_SUCCESS || + name == NULL) { + name = "unknown error"; + } + + if (pfn_cuGetErrorString == NULL || + pfn_cuGetErrorString(res, &desc) != CUDA_SUCCESS || desc == NULL) { + desc = "no description available"; + } + + (void)snprintf(buf, sizeof(buf), "%s (%s)", name, desc); + return buf; +} + int nccl_net_ofi_gpu_init(void) { int driverVersion = -1; @@ -161,7 +192,11 @@ int nccl_net_ofi_gpu_init(void) #endif RESOLVE_CUDA_FUNCTION(cuDriverGetVersion, 2020); + RESOLVE_CUDA_FUNCTION(cuGetErrorString, 6000); + RESOLVE_CUDA_FUNCTION(cuGetErrorName, 6000); RESOLVE_CUDA_FUNCTION(cuCtxGetDevice, 2000); + RESOLVE_CUDA_FUNCTION(cuCtxSetCurrent, 4000); + RESOLVE_CUDA_FUNCTION(cuCtxGetCurrent, 4000); RESOLVE_CUDA_FUNCTION(cuDeviceGetAttribute, 2000); #if HAVE_CUDA_GDRFLUSH_SUPPORT RESOLVE_CUDA_FUNCTION(cuFlushGPUDirectRDMAWrites, 11030); @@ -172,6 +207,7 @@ int nccl_net_ofi_gpu_init(void) RESOLVE_CUDA_FUNCTION(cuPointerGetAttributes, 7000); RESOLVE_CUDA_FUNCTION(cuMemAlloc, 3020); RESOLVE_CUDA_FUNCTION(cuMemFree, 3020); + RESOLVE_CUDA_FUNCTION(cuMemsetD8, 3020); RESOLVE_CUDA_FUNCTION(cuMemcpyHtoDAsync, 3020); RESOLVE_CUDA_FUNCTION(cuStreamCreate, 2000); RESOLVE_CUDA_FUNCTION(cuStreamSynchronize, 2000); @@ -248,6 +284,7 @@ int nccl_net_ofi_gpu_mem_alloc(void **ptr, size_t size) } if (ret != CUDA_SUCCESS) { + NCCL_OFI_WARN("cuMemAlloc failed: %s", nccl_net_ofi_cuda_error_string(ret)); return -EINVAL; } @@ -265,6 +302,9 @@ int nccl_net_ofi_gpu_mem_free(void *ptr) } ret = pfn_cuMemFree((CUdeviceptr)ptr); + if (ret != CUDA_SUCCESS) { + NCCL_OFI_WARN("cuMemFree failed: %s", nccl_net_ofi_cuda_error_string(ret)); + } CUresult restore_ret = pfn_cuThreadExchangeStreamCaptureMode(&mode); if (restore_ret != CUDA_SUCCESS) { @@ -290,19 +330,19 @@ int nccl_net_ofi_gpu_mem_copy_host_to_device(void *dst, void *src, size_t size) * graph capture on the legacy default stream. */ ret = pfn_cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING); if (ret != CUDA_SUCCESS) { - NCCL_OFI_WARN("cuStreamCreate failed (%d)", ret); + NCCL_OFI_WARN("cuStreamCreate failed: %s", nccl_net_ofi_cuda_error_string(ret)); goto restore; } ret = pfn_cuMemcpyHtoDAsync((CUdeviceptr)dst, src, size, stream); if (ret != CUDA_SUCCESS) { - NCCL_OFI_WARN("cuMemcpyHtoDAsync failed (%d)", ret); + NCCL_OFI_WARN("cuMemcpyHtoDAsync failed: %s", nccl_net_ofi_cuda_error_string(ret)); goto destroy; } ret = pfn_cuStreamSynchronize(stream); if (ret != CUDA_SUCCESS) { - NCCL_OFI_WARN("cuStreamSynchronize failed (%d)", ret); + NCCL_OFI_WARN("cuStreamSynchronize failed: %s", nccl_net_ofi_cuda_error_string(ret)); } destroy: @@ -321,19 +361,29 @@ int nccl_net_ofi_gpu_mem_copy_host_to_device(void *dst, void *src, size_t size) } /* - * Bind the calling thread to a CUDA device (and lazily its primary context). - * The gdrcopy worker thread is a bare std::thread with no CUDA context, so the - * lazy signal-segment discovery it performs (cuMemGetAddressRange) fails there - * unless we set one. + * Thin wrappers over the CUDA driver context primitives. set/get manage which + * context the calling thread is bound to; they take no ownership of it. */ -int nccl_net_ofi_gpu_get_device(int *dev_id) +int nccl_net_ofi_gpu_set_current_context(CUcontext ctx) { - return cudaGetDevice(dev_id) == cudaSuccess ? 0 : -EINVAL; + CUresult res = pfn_cuCtxSetCurrent(ctx); + if (res != CUDA_SUCCESS) { + NCCL_OFI_WARN("cuCtxSetCurrent failed: %s", + nccl_net_ofi_cuda_error_string(res)); + return -EINVAL; + } + return 0; } -int nccl_net_ofi_gpu_set_device(int dev_id) +int nccl_net_ofi_gpu_get_current_context(CUcontext *ctx) { - return cudaSetDevice(dev_id) == cudaSuccess ? 0 : -EINVAL; + CUresult res = pfn_cuCtxGetCurrent(ctx); + if (res != CUDA_SUCCESS) { + NCCL_OFI_WARN("cuCtxGetCurrent failed: %s", + nccl_net_ofi_cuda_error_string(res)); + return -EINVAL; + } + return 0; } int nccl_net_ofi_gpu_get_address_range(void *ptr, void **base_out, size_t *size_out) @@ -497,8 +547,8 @@ int nccl_net_ofi_gpu_vmm_alloc(void **ptr, size_t size, size_t *out_alloc_size) return -1; } - if (cudaMemset((void *)dptr, 0, alloc_size) != cudaSuccess) { - NCCL_OFI_WARN("vmm_alloc: cudaMemset failed"); + if (pfn_cuMemsetD8(dptr, 0, alloc_size) != CUDA_SUCCESS) { + NCCL_OFI_WARN("vmm_alloc: cuMemsetD8 failed"); pfn_cuMemUnmap(dptr, alloc_size); pfn_cuMemAddressFree(dptr, alloc_size); return -EINVAL; diff --git a/src/rdma/gin/nccl_ofi_gin.cpp b/src/rdma/gin/nccl_ofi_gin.cpp index 915ffb2b94..ea9437306e 100644 --- a/src/rdma/gin/nccl_ofi_gin.cpp +++ b/src/rdma/gin/nccl_ofi_gin.cpp @@ -76,17 +76,15 @@ nccl_ofi_rdma_gin_put_comm::nccl_ofi_rdma_gin_put_comm(nccl_ofi_gin_resources &r resources.increment_ref_cnt(); #if HAVE_CUDA - /* Capture the CUDA device on this (context-bearing) thread so the gdrcopy - * worker can bind to it; the worker's lazy signal-segment discovery calls - * cuMemGetAddressRange, which requires a current CUDA context. -1 means we - * could not determine the device, in which case the worker skips binding - * and falls back to whatever context it inherits (see - * run_gdrcopy_worker_loop). cudaGetDevice failing here is unexpected. */ - if (nccl_net_ofi_gpu_get_device(&gdrcopy_cuda_dev) != 0) { - NCCL_OFI_WARN("Could not query CUDA device for gdrcopy worker; " - "it will not bind to a device"); - gdrcopy_cuda_dev = -1; - } + /* Capture the app's context on this context-bearing thread so discovery + * can rebind it on NCCL's context-less GIN progress thread. The app owns + * it and it outlives this comm. */ + if (nccl_net_ofi_gpu_get_current_context(&gdrcopy_cuda_ctx) != 0 || + gdrcopy_cuda_ctx == nullptr) { + NCCL_OFI_WARN("Could not capture a CUDA context; GPU signal " + "delivery may fail"); + gdrcopy_cuda_ctx = nullptr; + } #endif /* Ensure the single process-wide gdrcopy worker exists. It is spawned @@ -1059,15 +1057,12 @@ int nccl_ofi_rdma_gin_put_comm::ensure_signal_seg( } #if HAVE_CUDA - /* Bind this worker to the comm's CUDA device so its lazy segment - discovery (cuMemGetAddressRange) has a valid CUDA context. A failure - here leaves the worker without the intended context, so warn -- lazy - discovery would then fail on the first signal. */ - if (gdrcopy_cuda_dev >= 0 && - nccl_net_ofi_gpu_set_device(gdrcopy_cuda_dev) != 0) { - NCCL_OFI_WARN("gdrcopy worker failed to bind CUDA device %d", - gdrcopy_cuda_dev); - } + /* cuMemGetAddressRange needs a current context; the GIN progress thread + * has none. Bind the comm's captured context -- per comm, not per thread, + * since one thread serves comms across GPUs. No restore needed. */ + if (gdrcopy_cuda_ctx != nullptr) { + nccl_net_ofi_gpu_set_current_context(gdrcopy_cuda_ctx); + } #endif void *seg_base_ptr = nullptr; @@ -1161,10 +1156,8 @@ int nccl_ofi_rdma_gin_put_comm::build_signal_work( uintptr_t signal_va = metadata.signal_base_address + metadata.signal_offset; - /* Lazily discover and pin the segment this signal lands in. - This runs on the proxy thread (under ep_lock) where the CUDA - context is current, so cuMemGetAddressRange succeeds. The worker - only ever touches the pre-resolved segment handle. */ + /* Lazily discover and pin the segment this signal lands in; + ensure_signal_seg binds the comm's context for the lookup. */ signal_seg_t *seg = nullptr; int ret = ensure_signal_seg(mr_handle, signal_va, &seg); if (OFI_UNLIKELY(ret != 0)) { diff --git a/tests/unit/Makefile.am b/tests/unit/Makefile.am index e28fe33ba8..c7ad5ab6fe 100644 --- a/tests/unit/Makefile.am +++ b/tests/unit/Makefile.am @@ -45,6 +45,7 @@ if !ENABLE_NEURON LDADD += $(CUDA_LIBS) noinst_PROGRAMS += gdrcopy gdrcopy_SOURCES = $(base_sources) gdrcopy.cpp + gdrcopy_LDADD = $(LDADD) $(CUDA_RUNTIME_LIBS) noinst_PROGRAMS += ctrl_msg ctrl_msg_SOURCES = $(base_sources) ctrl_msg.cpp if WANT_PLATFORM_AWS