From 73c824c5fdf1886b0e18f2eb77f487b130333846 Mon Sep 17 00:00:00 2001 From: romerojosh Date: Tue, 8 Sep 2026 22:13:03 +0000 Subject: [PATCH] Improve and document support for external CUDA graph capture. Signed-off-by: romerojosh --- docs/basic_usage.rst | 5 +- docs/cuda_graph_capture.rst | 84 ++++++++++ docs/env_vars.rst | 1 + docs/index.rst | 1 + include/cudecomp.h | 35 ++--- include/internal/comm_routines.h | 15 +- include/internal/common.h | 20 +++ include/internal/transpose.h | 29 ++-- src/cudecomp.cc | 112 +++++++++++-- tests/ctest/CMakeLists.txt | 35 ++++- tests/ctest/automatic_workspace_tests.cc | 1 + tests/ctest/cuda_graph_tests.cc | 192 +++++++++++++++++++++++ tests/ctest/halo_tests.cc | 83 +++++++++- tests/ctest/transpose_tests.cc | 88 +++++++++-- 14 files changed, 619 insertions(+), 82 deletions(-) create mode 100644 docs/cuda_graph_capture.rst create mode 100644 tests/ctest/cuda_graph_tests.cc diff --git a/docs/basic_usage.rst b/docs/basic_usage.rst index 2631e138..8b9c35d1 100644 --- a/docs/basic_usage.rst +++ b/docs/basic_usage.rst @@ -587,9 +587,8 @@ is required for NVSHMEM operations (see NVSHMEM documentation for more details). Alternatively, applications may let cuDecomp manage workspace memory by skipping the size queries and allocations above and passing :code:`CUDECOMP_WORKSPACE_AUTO` directly to each operation. cuDecomp then allocates, grows, and reuses handle-owned workspace memory and preserves ordering with the stream passed to each operation. All participating ranks -must consistently choose automatic or explicit workspace management for a given operation. Automatic workspace -management cannot be used while the caller's stream is being captured by a CUDA Graph; provide an explicit workspace -in that case. +must consistently choose automatic or explicit workspace management for a given operation. +See :ref:`cuda-graph-capture-ref` before capturing cuDecomp operations in a caller-owned CUDA Graph. .. tabs:: diff --git a/docs/cuda_graph_capture.rst b/docs/cuda_graph_capture.rst new file mode 100644 index 00000000..4ca7799f --- /dev/null +++ b/docs/cuda_graph_capture.rst @@ -0,0 +1,84 @@ +.. _cuda-graph-capture-ref: + +CUDA Graph Capture +================== + +cuDecomp transpose and halo operations can be captured as part of a caller-owned CUDA Graph. This allows an +application to capture a larger workflow containing cuDecomp operations and replay it with the standard CUDA Graph +APIs. + +Compatibility +------------- + +.. list-table:: External capture compatibility + :header-rows: 1 + + * - Communication family + - Supported + - Notes + * - MPI + - No + - Performs host-side MPI communication and CUDA stream synchronization. + * - NCCL + - Yes + - Includes the pipelined transpose backend. + * - NVSHMEM + - Yes + - Includes the pipelined and device-initiated transpose backends; requires NVSHMEM 2.9 or newer. NVSHMEM 2.6 + through 2.8 remain supported for normal, uncaptured operation. + +External capture is separate from cuDecomp's internal graph optimization. The +:code:`CUDECOMP_ENABLE_CUDA_GRAPHS` environment variable must be disabled. If it is enabled, a cuDecomp operation +called during external capture returns :code:`CUDECOMP_RESULT_NOT_SUPPORTED`. Performance reporting must also be +disabled during external capture. + +Workspace preparation +--------------------- + +Caller-owned workspace can be used directly during capture. The workspace must remain valid until all graph +executables and graphs that reference it have been destroyed. + +Automatic workspace requires a completed warmup before capture. Run each operation once, using the same datatype and +configuration that will be captured, and synchronize its stream before beginning capture. The warmup allocates enough +workspace and establishes any required communication-buffer registrations. An operation captured before its automatic +workspace is ready returns :code:`CUDECOMP_RESULT_NOT_SUPPORTED`. + +The automatic workspace allocation domain is frozen when it is first used during capture. A later operation that +would require the allocation to grow returns :code:`CUDECOMP_RESULT_NOT_SUPPORTED`. Warm up the largest operation that +the handle will execute before capturing any operation, or use a separate handle for operations with different +workspace requirements. + +Execution and lifetime requirements +----------------------------------- + +All participating ranks must capture the same collective operation and must consistently use caller-owned or +automatic workspace. Operations using the same cuDecomp handle must not overlap in execution, including eager +operations and graph replays. Order them on the same stream or with explicit CUDA dependencies. Applications that +require concurrent cuDecomp operations must use separate handles. + +Complete all graph launches and destroy graph executables and graphs before destroying the associated grid descriptor +or finalizing its cuDecomp handle. + +Example +------- + +The following example warms up an automatically managed workspace, captures one transpose, and instantiates the +resulting graph. Error checking is omitted for brevity. + +.. code-block:: c++ + + cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, stream); + cudaStreamSynchronize(stream); + + cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal); + cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, stream); + cudaStreamEndCapture(stream, &graph); + + cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0); + cudaGraphLaunch(graph_exec, stream); + cudaStreamSynchronize(stream); + + cudaGraphExecDestroy(graph_exec); + cudaGraphDestroy(graph); diff --git a/docs/env_vars.rst b/docs/env_vars.rst index 15426e83..a83af56b 100644 --- a/docs/env_vars.rst +++ b/docs/env_vars.rst @@ -35,6 +35,7 @@ CUDECOMP_ENABLE_CUDA_GRAPHS and communication overlap of packing kernels in large scale cases. Default setting is off (:code:`0`). Setting this variable to :code:`1` will enable this feature. +This option must remain disabled when using :ref:`external CUDA Graph capture `. CUDECOMP_ENABLE_PERFORMANCE_REPORT ------------------------------------ diff --git a/docs/index.rst b/docs/index.rst index 394e01d1..12a1fbf6 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -27,6 +27,7 @@ Table of Contents overview basic_usage + cuda_graph_capture autotuning nvshmem api diff --git a/include/cudecomp.h b/include/cudecomp.h index 790e7c46..00788b3a 100644 --- a/include/cudecomp.h +++ b/include/cudecomp.h @@ -530,9 +530,8 @@ cudecompResult_t cudecompGetShiftedRank(cudecompHandle_t handle, cudecompGridDes * the same, operation is performed in-place * @param[in] work A pointer to transpose workspace memory. For handle-owned workspace management, a NULL pointer (or * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and - * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same - * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while - * stream is being captured. + * runs the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. * @param[in] dtype The cuDecomp datatype to use for the transpose operation * @param[in] input_halo_extents An array of three integers to define halo region extents of the input data, in global * order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along @@ -563,9 +562,8 @@ cudecompResult_t cudecompTransposeXToY(cudecompHandle_t handle, cudecompGridDesc * the same, operation is performed in-place * @param[in] work A pointer to transpose workspace memory. For handle-owned workspace management, a NULL pointer (or * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and - * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same - * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while - * stream is being captured. + * runs the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. * @param[in] dtype The cuDecomp datatype to use for the transpose operation * @param[in] input_halo_extents An array of three integers to define halo region extents of the input data, in global * order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along @@ -596,9 +594,8 @@ cudecompResult_t cudecompTransposeYToZ(cudecompHandle_t handle, cudecompGridDesc * the same, operation is performed in-place * @param[in] work A pointer to transpose workspace memory. For handle-owned workspace management, a NULL pointer (or * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and - * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same - * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while - * stream is being captured. + * runs the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. * @param[in] dtype The cuDecomp datatype to use for the transpose operation * @param[in] input_halo_extents An array of three integers to define halo region extents of the input data, in global * order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along @@ -629,9 +626,8 @@ cudecompResult_t cudecompTransposeZToY(cudecompHandle_t handle, cudecompGridDesc * the same, operation is performed in-place * @param[in] work A pointer to transpose workspace memory. For handle-owned workspace management, a NULL pointer (or * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and - * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same - * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while - * stream is being captured. + * runs the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. * @param[in] dtype The cuDecomp datatype to use for the transpose operation * @param[in] input_halo_extents An array of three integers to define halo region extents of the input data, in global * order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along @@ -662,9 +658,8 @@ cudecompResult_t cudecompTransposeYToX(cudecompHandle_t handle, cudecompGridDesc * completion, this buffer will contain the input X-axis aligned pencil data with the specified halo regions updated. * @param[in] work A pointer to halo workspace memory. For handle-owned workspace management, a NULL pointer (or * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and - * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same - * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while - * stream is being captured. + * runs the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. * @param[in] dtype The cuDecomp datatype to use for the halo operation * @param[in] halo_extents An array of three integers to define halo region extents of the input data, in global order. * The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the @@ -693,9 +688,8 @@ cudecompResult_t cudecompUpdateHalosX(cudecompHandle_t handle, cudecompGridDesc_ * completion, this buffer will contain the input Y-axis aligned pencil data with the specified halo regions updated. * @param[in] work A pointer to halo workspace memory. For handle-owned workspace management, a NULL pointer (or * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and - * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same - * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while - * stream is being captured. + * runs the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. * @param[in] dtype The cuDecomp datatype to use for the halo operation * @param[in] halo_extents An array of three integers to define halo region extents of the input data, in global order. * The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the @@ -724,9 +718,8 @@ cudecompResult_t cudecompUpdateHalosY(cudecompHandle_t handle, cudecompGridDesc_ * completion, this buffer will contain the input Z-axis aligned pencil data with the specified halo regions updated. * @param[in] work A pointer to halo workspace memory. For handle-owned workspace management, a NULL pointer (or * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and - * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same - * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while - * stream is being captured. + * runs the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. * @param[in] dtype The cuDecomp datatype to use for the halo operation * @param[in] halo_extents An array of three integers to define halo region extents of the input data, in global order. * The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the diff --git a/include/internal/comm_routines.h b/include/internal/comm_routines.h index 1a3fe03c..a347238b 100644 --- a/include/internal/comm_routines.h +++ b/include/internal/comm_routines.h @@ -42,6 +42,10 @@ namespace cudecomp { +#ifdef ENABLE_NVSHMEM +inline constexpr uint64_t nvshmem_signal_set_value = 1; // Graph-stable source for mapped-peer signal writes +#endif + static inline MPI_Datatype getMpiDataType(float) { return MPI_FLOAT; } static inline MPI_Datatype getMpiDataType(double) { return MPI_DOUBLE; } static inline MPI_Datatype getMpiDataType(cuda::std::complex) { return MPI_C_FLOAT_COMPLEX; } @@ -503,13 +507,16 @@ cudecompAlltoallPipelined(const cudecompHandle_t& handle, const cudecompGridDesc int dst_rank = dst_ranks[i]; int dst_rank_global = getGlobalRank(handle, grid_desc, comm_axis, dst_rank); - if (!nvshmem_ptr(recv_buff, dst_rank_global) || src_rank == self_rank) { continue; } + auto peer_recv_buff = nvshmem_ptr(recv_buff, dst_rank_global); + if (!peer_recv_buff || src_rank == self_rank) { continue; } CHECK_CUDA(cudaStreamWaitEvent(pl_stream, grid_desc->events[dst_rank], 0)); - nvshmemx_putmem_signal_on_stream(recv_buff + recv_offsets_nvshmem[dst_rank], send_buff + send_offsets[dst_rank], - send_counts[dst_rank] * sizeof(T), &comm_info.nvshmem_signals[comm_info.rank], - 1, NVSHMEM_SIGNAL_SET, dst_rank_global, pl_stream); + nvshmemx_putmem_on_stream(recv_buff + recv_offsets_nvshmem[dst_rank], send_buff + send_offsets[dst_rank], + send_counts[dst_rank] * sizeof(T), dst_rank_global, pl_stream); + auto peer_signal = nvshmem_ptr(&comm_info.nvshmem_signals[comm_info.rank], dst_rank_global); + CHECK_CUDA(cudaMemcpyAsync(peer_signal, &nvshmem_signal_set_value, sizeof(nvshmem_signal_set_value), + cudaMemcpyHostToDevice, pl_stream)); } if (need_quiet) { nvshmemx_quiet_on_stream(pl_stream); } diff --git a/include/internal/common.h b/include/internal/common.h index 62311740..b8863524 100644 --- a/include/internal/common.h +++ b/include/internal/common.h @@ -71,6 +71,7 @@ struct nvshmemRuntimeState { } bool initialized = false; // Flag to track NVSHMEM initialization + bool external_graph_capture_supported = false; // NVSHMEM runtime supports CUDA Graph capture size_t nvshmem_symmetric_size = 0; // NVSHMEM symmetric size bool nvshmem_vmm = true; // Flag to track if NVSHMEM is using VMM allocations std::unordered_map nvshmem_allocations; // Table to record NVSHMEM allocations @@ -127,6 +128,7 @@ struct cudecompHandle { struct ManagedWorkspace { void* ptr = nullptr; size_t size = 0; + bool capture_frozen = false; #ifdef ENABLE_NVSHMEM cudecomp::nvshmemRuntime nvshmem_runtime; #endif @@ -433,6 +435,24 @@ static inline bool haloBackendRequiresNvshmem(cudecompHaloCommBackend_t comm_bac return (comm_backend == CUDECOMP_HALO_COMM_NVSHMEM || comm_backend == CUDECOMP_HALO_COMM_NVSHMEM_BLOCKING); } +static inline bool transposeBackendSupportsExternalGraphCapture(cudecompTransposeCommBackend_t comm_backend) { + return transposeBackendRequiresNccl(comm_backend) || comm_backend == CUDECOMP_TRANSPOSE_COMM_NVSHMEM || + comm_backend == CUDECOMP_TRANSPOSE_COMM_NVSHMEM_PL || comm_backend == CUDECOMP_TRANSPOSE_COMM_NVSHMEM_SM; +} + +static inline bool haloBackendSupportsExternalGraphCapture(cudecompHaloCommBackend_t comm_backend) { + return haloBackendRequiresNccl(comm_backend) || haloBackendRequiresNvshmem(comm_backend); +} + +static inline bool nvshmemSupportsExternalGraphCapture(const cudecompGridDesc_t grid_desc) { +#ifdef ENABLE_NVSHMEM + return grid_desc->nvshmem_runtime && grid_desc->nvshmem_runtime->external_graph_capture_supported; +#else + (void)grid_desc; + return false; +#endif +} + static inline bool isManagedPointer(void* ptr) { // Check if input pointer is managed cudaPointerAttributes attr; diff --git a/include/internal/transpose.h b/include/internal/transpose.h index ff4b66c2..d5a39dd3 100644 --- a/include/internal/transpose.h +++ b/include/internal/transpose.h @@ -218,6 +218,9 @@ static void cudecompTranspose_(int ax, int dir, const cudecompHandle_t handle, c bool output_has_halos_padding = anyNonzeros(output_halo_extents) || anyNonzeros(output_padding); bool pipelined = isTransposeCommPipelined(grid_desc->config.transpose_comm_backend); int memcpy_limit = pipelined ? 1 : CUDECOMP_BATCHED_D2D_3D_PARAM_CAPACITY; + cudaStreamCaptureStatus capture_status; + CHECK_CUDA(cudaStreamIsCapturing(stream, &capture_status)); + bool use_internal_graphs = handle->cuda_graphs_enable && capture_status == cudaStreamCaptureStatusNone; // Set axis values int ax_a = ax; @@ -459,11 +462,11 @@ static void cudecompTranspose_(int ax, int dir, const cudecompHandle_t handle, c auto dtype = getCudecompDataType(); auto key = std::tie(i1, o1, ax, dir, pinfo_a_h, pinfo_b_h, dtype); - if (handle->cuda_graphs_enable && grid_desc->graph_cache.cached(key)) { + if (use_internal_graphs && grid_desc->graph_cache.cached(key)) { grid_desc->graph_cache.replay(key, stream); } else { cudaStream_t graph_stream = stream; - if (handle->cuda_graphs_enable && splits_a.size() > 1) { + if (use_internal_graphs && splits_a.size() > 1) { graph_stream = grid_desc->graph_cache.startCapture(key, stream); } @@ -501,18 +504,15 @@ static void cudecompTranspose_(int ax, int dir, const cudecompHandle_t handle, c localPermute(handle, extents, order, strides_in, strides_out, src, dst, graph_stream); #if CUDART_VERSION >= 11010 - cudaStreamCaptureStatus capture_status; - CHECK_CUDA(cudaStreamIsCapturing(graph_stream, &capture_status)); CHECK_CUDA(cudaEventRecordWithFlags(grid_desc->events[dst_rank], graph_stream, - capture_status == cudaStreamCaptureStatusActive - ? cudaEventRecordExternal - : cudaEventRecordDefault)); + use_internal_graphs && splits_a.size() > 1 ? cudaEventRecordExternal + : cudaEventRecordDefault)); #else CHECK_CUDA(cudaEventRecord((grid_desc->events[dst_rank], graph_stream)); #endif } - if (handle->cuda_graphs_enable && splits_a.size() > 1) { + if (use_internal_graphs && splits_a.size() > 1) { grid_desc->graph_cache.endCapture(key); grid_desc->graph_cache.replay(key, stream); } @@ -538,11 +538,11 @@ static void cudecompTranspose_(int ax, int dir, const cudecompHandle_t handle, c auto dtype = getCudecompDataType(); auto key = std::tie(i1, o1, ax, dir, pinfo_a_h, pinfo_b_h, dtype); - if (handle->cuda_graphs_enable && grid_desc->graph_cache.cached(key)) { + if (use_internal_graphs && grid_desc->graph_cache.cached(key)) { grid_desc->graph_cache.replay(key, stream); } else { cudaStream_t graph_stream = stream; - if (handle->cuda_graphs_enable && pipelined && splits_a.size() > 1) { + if (use_internal_graphs && pipelined && splits_a.size() > 1) { graph_stream = grid_desc->graph_cache.startCapture(key, stream); } @@ -598,18 +598,15 @@ static void cudecompTranspose_(int ax, int dir, const cudecompHandle_t handle, c } #if CUDART_VERSION >= 11010 if (pipelined) { - cudaStreamCaptureStatus capture_status; - CHECK_CUDA(cudaStreamIsCapturing(graph_stream, &capture_status)); CHECK_CUDA(cudaEventRecordWithFlags(grid_desc->events[dst_rank], graph_stream, - capture_status == cudaStreamCaptureStatusActive - ? cudaEventRecordExternal - : cudaEventRecordDefault)); + use_internal_graphs && splits_a.size() > 1 ? cudaEventRecordExternal + : cudaEventRecordDefault)); } #else if (pipelined) CHECK_CUDA(cudaEventRecord((grid_desc->events[dst_rank], graph_stream)); #endif } - if (handle->cuda_graphs_enable && pipelined && splits_a.size() > 1) { + if (use_internal_graphs && pipelined && splits_a.size() > 1) { grid_desc->graph_cache.endCapture(key); grid_desc->graph_cache.replay(key, stream); } diff --git a/src/cudecomp.cc b/src/cudecomp.cc index dd9bd950..d7eca15b 100644 --- a/src/cudecomp.cc +++ b/src/cudecomp.cc @@ -758,7 +758,7 @@ static void inspectNvshmemEnvVars(nvshmemRuntimeState& runtime) { } } -static void checkNvshmemVersion(cudecompHandle_t& handle) { +static void checkNvshmemVersion(cudecompHandle_t& handle, bool& external_graph_capture_supported) { int major, minor, patch; char name[NVSHMEM_MAX_NAME_LEN]; nvshmem_info_get_name(name); @@ -786,6 +786,8 @@ static void checkNvshmemVersion(cudecompHandle_t& handle) { CHECK_MPI(MPI_Barrier(handle->mpi_comm)); setenv("NVSHMEM_CUMEM_GRANULARITY", "2147483648", 1); } + + external_graph_capture_supported = major > 2 || (major == 2 && minor >= 9); } static nvshmemRuntime acquireNvshmemRuntime(cudecompHandle_t& handle) { @@ -797,7 +799,7 @@ static nvshmemRuntime acquireNvshmemRuntime(cudecompHandle_t& handle) { if (process_nvshmem_state.init_world_ranks.empty()) { process_nvshmem_state.init_world_ranks = std::move(world_ranks); } - checkNvshmemVersion(handle); + checkNvshmemVersion(handle, runtime->external_graph_capture_supported); initNvshmemFromMPIComm(handle->mpi_comm); runtime->initialized = true; process_nvshmem_state.active_runtime = runtime; @@ -807,7 +809,7 @@ static nvshmemRuntime acquireNvshmemRuntime(cudecompHandle_t& handle) { if (!runtime->initialized) { runtime->finalize(); inspectNvshmemEnvVars(*runtime); - checkNvshmemVersion(handle); + checkNvshmemVersion(handle, runtime->external_graph_capture_supported); initNvshmemFromMPIComm(handle->mpi_comm); runtime->initialized = true; } @@ -935,6 +937,7 @@ static void releaseManagedWorkspaces(cudecompHandle_t handle) { freeOrdinaryWorkspace(handle, handle->ordinary_workspace.ptr); handle->ordinary_workspace.ptr = nullptr; handle->ordinary_workspace.size = 0; + handle->ordinary_workspace.capture_frozen = false; #ifdef ENABLE_NVSHMEM freeNvshmemWorkspace(handle->nvshmem_workspace.nvshmem_runtime, handle->nvshmem_workspace.ptr); @@ -942,6 +945,7 @@ static void releaseManagedWorkspaces(cudecompHandle_t handle) { #endif handle->nvshmem_workspace.ptr = nullptr; handle->nvshmem_workspace.size = 0; + handle->nvshmem_workspace.capture_frozen = false; handle->workspace_ingress_event.reset(); handle->workspace_egress_event.reset(); @@ -978,12 +982,41 @@ static void registerWorkspaceWithNccl(cudecompHandle_t handle, cudecompGridDesc_ #endif } +static bool workspaceRegisteredWithNccl(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, void* buffer) { +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0) + if (!handle->nccl_enable_ubr) return true; + + auto entry = handle->nccl_ubr_handles.find(buffer); + if (entry == handle->nccl_ubr_handles.end()) return false; + auto registered_with = [&](const ncclComm& comm) { + return !comm || std::any_of(entry->second.begin(), entry->second.end(), + [&](const auto& registration) { return registration.first.get() == comm.get(); }); + }; + return registered_with(grid_desc->nccl_comm) && registered_with(grid_desc->nccl_local_comm); +#else + (void)handle; + (void)grid_desc; + (void)buffer; + return true; +#endif +} + +static bool allRanksTrue(cudecompHandle_t handle, bool local_value) { + int value = local_value; + CHECK_MPI(MPI_Allreduce(MPI_IN_PLACE, &value, 1, MPI_INT, MPI_MIN, handle->mpi_comm)); + return value != 0; +} + static void* prepareManagedWorkspace(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, size_t required_size, bool use_nvshmem) { if (required_size == 0) return nullptr; required_size = getWorkspaceAllocationSize(handle, required_size, use_nvshmem); auto& workspace = use_nvshmem ? handle->nvshmem_workspace : handle->ordinary_workspace; + if (workspace.capture_frozen && !allRanksTrue(handle, workspace.ptr != nullptr && workspace.size >= required_size)) { + THROW_NOT_SUPPORTED( + "automatic workspace cannot grow after being used by an externally captured CUDA Graph; use a new handle"); + } if (workspace.size < required_size) { CHECK_CUDA(cudaStreamSynchronize(*handle->workspace_stream)); if (workspace.ptr) { @@ -1014,20 +1047,54 @@ static void* prepareManagedWorkspace(cudecompHandle_t handle, cudecompGridDesc_t template static void runWithWorkspace(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, void* work, size_t required_size, - bool use_nvshmem, cudaStream_t caller_stream, Function&& function) { + bool use_nvshmem, bool graph_capture_supported, cudaStream_t caller_stream, + Function&& function) { + cudaStreamCaptureStatus capture_status; + CHECK_CUDA(cudaStreamIsCapturing(caller_stream, &capture_status)); + if (capture_status == cudaStreamCaptureStatusInvalidated) { + THROW_NOT_SUPPORTED("cuDecomp cannot be used with an invalidated CUDA stream capture"); + } + if (capture_status == cudaStreamCaptureStatusActive && handle->cuda_graphs_enable) { + THROW_NOT_SUPPORTED("external CUDA Graph capture cannot be combined with CUDECOMP_ENABLE_CUDA_GRAPHS"); + } + if (capture_status == cudaStreamCaptureStatusActive && !graph_capture_supported) { + THROW_NOT_SUPPORTED("the selected communication backend does not support external CUDA Graph capture"); + } +#ifdef ENABLE_NVSHMEM + if (capture_status == cudaStreamCaptureStatusActive && use_nvshmem && + !nvshmemSupportsExternalGraphCapture(grid_desc)) { + THROW_NOT_SUPPORTED("NVSHMEM 2.9.0 or newer is required for external CUDA Graph capture"); + } +#endif + if (capture_status == cudaStreamCaptureStatusActive && handle->performance_report_enable) { + THROW_NOT_SUPPORTED("performance reporting is not supported during external CUDA Graph capture"); + } + if (work) { function(work, caller_stream); return; } - cudaStreamCaptureStatus capture_status; - CHECK_CUDA(cudaStreamIsCapturing(caller_stream, &capture_status)); - if (capture_status != cudaStreamCaptureStatusNone) { - THROW_NOT_SUPPORTED( - "automatic workspace management is not supported during CUDA stream capture; use an explicit workspace"); + std::lock_guard lock(handle->workspace_mutex); + auto& workspace = use_nvshmem ? handle->nvshmem_workspace : handle->ordinary_workspace; + if (capture_status == cudaStreamCaptureStatusActive) { + bool capture_ready = true; + if (required_size != 0) { + capture_ready = capture_ready && workspace.ptr != nullptr && workspace.size >= required_size; + if (!use_nvshmem) { + capture_ready = capture_ready && workspaceRegisteredWithNccl(handle, grid_desc, workspace.ptr); + } + } + if (!allRanksTrue(handle, capture_ready)) { + THROW_NOT_SUPPORTED( + "automatic workspace is not prepared for CUDA Graph capture; warm up the operation outside capture"); + } + + function(workspace.ptr, caller_stream); + if (required_size != 0) { workspace.capture_frozen = true; } + return; } - std::lock_guard lock(handle->workspace_mutex); if (!handle->workspace_stream) { handle->workspace_stream = std::make_unique(); } if (!handle->workspace_ingress_event) { handle->workspace_ingress_event = std::make_unique(); } if (!handle->workspace_egress_event) { handle->workspace_egress_event = std::make_unique(); } @@ -1932,7 +1999,9 @@ cudecompResult_t cudecompTransposeXToY(cudecompHandle_t handle, cudecompGridDesc } }; bool use_nvshmem = transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend); - runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); + bool graph_capture_supported = + transposeBackendSupportsExternalGraphCapture(grid_desc->config.transpose_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, graph_capture_supported, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -1982,7 +2051,9 @@ cudecompResult_t cudecompTransposeYToZ(cudecompHandle_t handle, cudecompGridDesc } }; bool use_nvshmem = transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend); - runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); + bool graph_capture_supported = + transposeBackendSupportsExternalGraphCapture(grid_desc->config.transpose_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, graph_capture_supported, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -2032,7 +2103,9 @@ cudecompResult_t cudecompTransposeZToY(cudecompHandle_t handle, cudecompGridDesc } }; bool use_nvshmem = transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend); - runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); + bool graph_capture_supported = + transposeBackendSupportsExternalGraphCapture(grid_desc->config.transpose_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, graph_capture_supported, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -2082,7 +2155,9 @@ cudecompResult_t cudecompTransposeYToX(cudecompHandle_t handle, cudecompGridDesc } }; bool use_nvshmem = transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend); - runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); + bool graph_capture_supported = + transposeBackendSupportsExternalGraphCapture(grid_desc->config.transpose_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, graph_capture_supported, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -2133,7 +2208,8 @@ cudecompResult_t cudecompUpdateHalosX(cudecompHandle_t handle, cudecompGridDesc_ } }; bool use_nvshmem = haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend); - runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); + bool graph_capture_supported = haloBackendSupportsExternalGraphCapture(grid_desc->config.halo_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, graph_capture_supported, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -2184,7 +2260,8 @@ cudecompResult_t cudecompUpdateHalosY(cudecompHandle_t handle, cudecompGridDesc_ } }; bool use_nvshmem = haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend); - runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); + bool graph_capture_supported = haloBackendSupportsExternalGraphCapture(grid_desc->config.halo_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, graph_capture_supported, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -2235,7 +2312,8 @@ cudecompResult_t cudecompUpdateHalosZ(cudecompHandle_t handle, cudecompGridDesc_ } }; bool use_nvshmem = haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend); - runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); + bool graph_capture_supported = haloBackendSupportsExternalGraphCapture(grid_desc->config.halo_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, graph_capture_supported, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; diff --git a/tests/ctest/CMakeLists.txt b/tests/ctest/CMakeLists.txt index 150fa846..eb15a34c 100644 --- a/tests/ctest/CMakeLists.txt +++ b/tests/ctest/CMakeLists.txt @@ -155,6 +155,7 @@ endif() add_executable(cudecomp_test_transpose) target_sources(cudecomp_test_transpose PRIVATE + cuda_graph_tests.cc mpi_test_main.cc transpose_tests.cc ) @@ -243,10 +244,40 @@ if (CUDECOMP_ENABLE_NVSHMEM) ) endif() -add_cudecomp_transpose_test(cudecomp_cuda_graphs "CudaGraphMpiBackends/*" "cuda_graphs;mpi") -set_tests_properties(cudecomp_cuda_graphs PROPERTIES +add_cudecomp_transpose_test(cudecomp_cuda_graphs_internal "CudaGraphMpiBackends/*" "cuda_graphs;mpi") +set_tests_properties(cudecomp_cuda_graphs_internal PROPERTIES ENVIRONMENT "CUDECOMP_ENABLE_CUDA_GRAPHS=1" ) +add_cudecomp_transpose_test(cudecomp_cuda_graphs_external_transpose_nccl + "ExternalCudaGraphNcclBackends/*:CudaGraphContractTest.*" + "cuda_graphs;workspace;nccl" +) +set_tests_properties(cudecomp_cuda_graphs_external_transpose_nccl PROPERTIES + ENVIRONMENT "CUDECOMP_TEST_KEEPALIVE_BACKEND=nccl" +) +add_cudecomp_halo_test(cudecomp_cuda_graphs_external_halo_nccl + "ExternalCudaGraphNcclBackends/*" + "cuda_graphs;workspace;nccl" +) +set_tests_properties(cudecomp_cuda_graphs_external_halo_nccl PROPERTIES + ENVIRONMENT "CUDECOMP_TEST_KEEPALIVE_BACKEND=nccl" +) +if (CUDECOMP_ENABLE_NVSHMEM) + add_cudecomp_transpose_test(cudecomp_cuda_graphs_external_transpose_nvshmem + "ExternalCudaGraphNvshmemBackends/*" + "cuda_graphs;workspace;nvshmem" + ) + set_tests_properties(cudecomp_cuda_graphs_external_transpose_nvshmem PROPERTIES + ENVIRONMENT "CUDECOMP_TEST_KEEPALIVE_BACKEND=nvshmem;NVSHMEM_DISABLE_NCCL=1" + ) + add_cudecomp_halo_test(cudecomp_cuda_graphs_external_halo_nvshmem + "ExternalCudaGraphNvshmemBackends/*" + "cuda_graphs;workspace;nvshmem" + ) + set_tests_properties(cudecomp_cuda_graphs_external_halo_nvshmem PROPERTIES + ENVIRONMENT "CUDECOMP_TEST_KEEPALIVE_BACKEND=nvshmem;NVSHMEM_DISABLE_NCCL=1" + ) +endif() add_cudecomp_transpose_test(cudecomp_nccl_user_buffer_registration "NcclUserBufferRegistration/*" "nccl;nccl_ubr") set_tests_properties(cudecomp_nccl_user_buffer_registration PROPERTIES diff --git a/tests/ctest/automatic_workspace_tests.cc b/tests/ctest/automatic_workspace_tests.cc index f281bae9..7c33a0ac 100644 --- a/tests/ctest/automatic_workspace_tests.cc +++ b/tests/ctest/automatic_workspace_tests.cc @@ -237,4 +237,5 @@ TEST(AutomaticWorkspaceTest, KeepsAllocationDomainsSeparate) { EXPECT_EQ(handle->ordinary_workspace.ptr, ordinary_workspace); EXPECT_EQ(handle->nvshmem_workspace.ptr, nvshmem_workspace); } + #endif diff --git a/tests/ctest/cuda_graph_tests.cc b/tests/ctest/cuda_graph_tests.cc new file mode 100644 index 00000000..a0d650e6 --- /dev/null +++ b/tests/ctest/cuda_graph_tests.cc @@ -0,0 +1,192 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include + +#include +#include + +#include "cudecomp.h" +#include "internal/common.h" + +#include "gpu_test_utils.h" +#include "mpi_test_utils.h" +#include "test_utils.h" + +namespace { + +constexpr std::array kGdims{9, 10, 11}; +constexpr std::array kPdims{2, 2}; + +cudecompGridDescConfig_t makeConfig() { + cudecompGridDescConfig_t config; + EXPECT_EQ(CUDECOMP_RESULT_SUCCESS, cudecompGridDescConfigSetDefaults(&config)); + std::copy(kGdims.begin(), kGdims.end(), config.gdims); + std::copy(kPdims.begin(), kPdims.end(), config.pdims); + config.transpose_comm_backend = CUDECOMP_TRANSPOSE_COMM_NCCL; + config.halo_comm_backend = CUDECOMP_HALO_COMM_NCCL; + return config; +} + +} // namespace + +TEST(CudaGraphContractTest, AutomaticWorkspaceCaptureRequiresWarmup) { + const auto world_comm = cudecomp_test::MpiTestComm::world(); + if (world_comm.size() != 4) { GTEST_SKIP() << "automatic workspace capture test requires exactly four ranks"; } + + const auto setup_decision = cudecomp_test::initializeGpuForTest(world_comm, true); + ASSERT_FALSE(setup_decision.fail) << setup_decision.reason; + if (setup_decision.skip) { GTEST_SKIP() << setup_decision.reason; } + + cudecompHandle_t handle = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompInit(&handle, world_comm.mpiComm())); + cudecomp_test::cudecompHandleGuard handle_guard(handle); + + auto config = makeConfig(); + cudecompGridDesc_t grid_desc = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGridDescCreate(handle, &grid_desc, &config, nullptr)); + cudecomp_test::gridDescGuard grid_desc_guard(handle, grid_desc); + + cudecompPencilInfo_t x_pinfo; + cudecompPencilInfo_t y_pinfo; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &x_pinfo, 0, nullptr, nullptr)); + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &y_pinfo, 1, nullptr, nullptr)); + const int64_t data_elements = std::max(x_pinfo.size, y_pinfo.size); + float* input = nullptr; + float* output = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&input, data_elements * sizeof(*input))); + cudecomp_test::cudaBufferGuard input_guard(input); + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&output, data_elements * sizeof(*output))); + cudecomp_test::cudaBufferGuard output_guard(output); + + cudaStream_t stream = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + CHECK_CUDA_GLOBAL(world_comm, cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + EXPECT_EQ(CUDECOMP_RESULT_NOT_SUPPORTED, + cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, CUDECOMP_FLOAT, nullptr, + nullptr, nullptr, nullptr, stream)); + cudaGraph_t graph = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaStreamEndCapture(stream, &graph)); + if (graph) { CHECK_CUDA_GLOBAL(world_comm, cudaGraphDestroy(graph)); } + CHECK_CUDA_GLOBAL(world_comm, cudaStreamDestroy(stream)); + EXPECT_FALSE(handle->ordinary_workspace.capture_frozen); +} + +TEST(CudaGraphContractTest, InternalGraphsRejectExternalCapture) { + const auto world_comm = cudecomp_test::MpiTestComm::world(); + if (world_comm.size() != 4) { GTEST_SKIP() << "CUDA Graph mode conflict test requires exactly four ranks"; } + + const auto setup_decision = cudecomp_test::initializeGpuForTest(world_comm, true); + ASSERT_FALSE(setup_decision.fail) << setup_decision.reason; + if (setup_decision.skip) { GTEST_SKIP() << setup_decision.reason; } + + cudecompHandle_t handle = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompInit(&handle, world_comm.mpiComm())); + cudecomp_test::cudecompHandleGuard handle_guard(handle); + + auto config = makeConfig(); + cudecompGridDesc_t grid_desc = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGridDescCreate(handle, &grid_desc, &config, nullptr)); + cudecomp_test::gridDescGuard grid_desc_guard(handle, grid_desc); + + cudecompPencilInfo_t x_pinfo; + cudecompPencilInfo_t y_pinfo; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &x_pinfo, 0, nullptr, nullptr)); + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &y_pinfo, 1, nullptr, nullptr)); + const int64_t data_elements = std::max(x_pinfo.size, y_pinfo.size); + float* input = nullptr; + float* output = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&input, data_elements * sizeof(*input))); + cudecomp_test::cudaBufferGuard input_guard(input); + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&output, data_elements * sizeof(*output))); + cudecomp_test::cudaBufferGuard output_guard(output); + + int64_t workspace_elements = 0; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetTransposeWorkspaceSize(handle, grid_desc, &workspace_elements)); + void* explicit_workspace = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, + cudecompMalloc(handle, grid_desc, &explicit_workspace, workspace_elements * sizeof(float))); + cudecomp_test::cudecompBufferGuard workspace_guard(handle, grid_desc, explicit_workspace); + + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + handle->cuda_graphs_enable = true; + + for (void* workspace : {explicit_workspace, CUDECOMP_WORKSPACE_AUTO}) { + cudaStream_t stream = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + CHECK_CUDA_GLOBAL(world_comm, cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + EXPECT_EQ(CUDECOMP_RESULT_NOT_SUPPORTED, + cudecompTransposeXToY(handle, grid_desc, input, output, workspace, CUDECOMP_FLOAT, nullptr, nullptr, + nullptr, nullptr, stream)); + cudaGraph_t graph = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaStreamEndCapture(stream, &graph)); + if (graph) { CHECK_CUDA_GLOBAL(world_comm, cudaGraphDestroy(graph)); } + CHECK_CUDA_GLOBAL(world_comm, cudaStreamDestroy(stream)); + } +} + +TEST(CudaGraphContractTest, CapturedAutomaticWorkspaceCannotGrow) { + const auto world_comm = cudecomp_test::MpiTestComm::world(); + if (world_comm.size() != 4) { GTEST_SKIP() << "automatic workspace freeze test requires exactly four ranks"; } + + const auto setup_decision = cudecomp_test::initializeGpuForTest(world_comm, true); + ASSERT_FALSE(setup_decision.fail) << setup_decision.reason; + if (setup_decision.skip) { GTEST_SKIP() << setup_decision.reason; } + + cudecompHandle_t handle = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompInit(&handle, world_comm.mpiComm())); + cudecomp_test::cudecompHandleGuard handle_guard(handle); + + auto config = makeConfig(); + cudecompGridDesc_t grid_desc = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGridDescCreate(handle, &grid_desc, &config, nullptr)); + cudecomp_test::gridDescGuard grid_desc_guard(handle, grid_desc); + + cudecompPencilInfo_t x_pinfo; + cudecompPencilInfo_t y_pinfo; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &x_pinfo, 0, nullptr, nullptr)); + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &y_pinfo, 1, nullptr, nullptr)); + const int64_t data_elements = std::max(x_pinfo.size, y_pinfo.size); + std::complex* input = nullptr; + std::complex* output = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&input, data_elements * sizeof(*input))); + cudecomp_test::cudaBufferGuard input_guard(input); + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&output, data_elements * sizeof(*output))); + cudecomp_test::cudaBufferGuard output_guard(output); + + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + void* captured_workspace = handle->ordinary_workspace.ptr; + const size_t captured_size = handle->ordinary_workspace.size; + + cudaStream_t stream = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + CHECK_CUDA_GLOBAL(world_comm, cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, stream)); + cudaGraph_t graph = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaStreamEndCapture(stream, &graph)); + ASSERT_NE(graph, nullptr); + CHECK_CUDA_GLOBAL(world_comm, cudaGraphDestroy(graph)); + CHECK_CUDA_GLOBAL(world_comm, cudaStreamDestroy(stream)); + EXPECT_TRUE(handle->ordinary_workspace.capture_frozen); + + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + EXPECT_EQ(CUDECOMP_RESULT_NOT_SUPPORTED, + cudecompTransposeXToY(handle, grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, CUDECOMP_DOUBLE_COMPLEX, + nullptr, nullptr, nullptr, nullptr, nullptr)); + EXPECT_EQ(handle->ordinary_workspace.ptr, captured_workspace); + EXPECT_EQ(handle->ordinary_workspace.size, captured_size); +} diff --git a/tests/ctest/halo_tests.cc b/tests/ctest/halo_tests.cc index ddace8b7..1799a6bf 100644 --- a/tests/ctest/halo_tests.cc +++ b/tests/ctest/halo_tests.cc @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -165,6 +166,18 @@ std::vector haloCasesForLabel(const char* label) { return cases; } +std::vector externalCudaGraphHaloCases(const char* label) { + std::vector cases; + for (const auto& backend : cudecomp_test::haloBackends()) { + if (std::string(backend.label) != label) continue; + + auto test_case = makeCase(backend, "ExternalCudaGraph", 0); + cases.push_back(test_case); + cases.push_back(withAutomaticWorkspace(std::move(test_case))); + } + return cases; +} + bool isInternal(const cudecompPencilInfo_t& pinfo, const std::array& local) { return local[0] >= pinfo.halo_extents[pinfo.order[0]] && local[0] < pinfo.shape[0] - pinfo.halo_extents[pinfo.order[0]] - pinfo.padding[pinfo.order[0]] && @@ -284,17 +297,17 @@ testing::AssertionResult pencilMatches(const std::vector& expected, const std template cudecompResult_t runHalo(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, int axis, T* input, void* work, cudecompDataType_t dtype, const cudecompPencilInfo_t& pinfo, - const std::array& halo_periods, int dim) { + const std::array& halo_periods, int dim, cudaStream_t stream = nullptr) { switch (axis) { case 0: return cudecompUpdateHalosX(handle, grid_desc, input, work, dtype, pinfo.halo_extents, halo_periods.data(), dim, - pinfo.padding, 0); + pinfo.padding, stream); case 1: return cudecompUpdateHalosY(handle, grid_desc, input, work, dtype, pinfo.halo_extents, halo_periods.data(), dim, - pinfo.padding, 0); + pinfo.padding, stream); case 2: return cudecompUpdateHalosZ(handle, grid_desc, input, work, dtype, pinfo.halo_extents, halo_periods.data(), dim, - pinfo.padding, 0); + pinfo.padding, stream); } return CUDECOMP_RESULT_INVALID_USAGE; } @@ -302,8 +315,9 @@ cudecompResult_t runHalo(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, } // namespace class HaloCorrectnessTest : public ::testing::TestWithParam {}; +class ExternalCudaGraphHaloCorrectnessTest : public ::testing::TestWithParam {}; -template void runHaloCase(const HaloCase& test_case) { +template void runHaloCase(const HaloCase& test_case, bool capture_external_graph = false) { const int active_ranks = test_case.pdims[0] * test_case.pdims[1]; const auto world_comm = cudecomp_test::MpiTestComm::world(); @@ -402,6 +416,50 @@ template void runHaloCase(const HaloCase& test_case) { ? handle->nvshmem_workspace.ptr : handle->ordinary_workspace.ptr); } + + if (capture_external_graph) { + cudaStream_t capture_stream = nullptr; + CHECK_CUDA_GLOBAL(active_comm, cudaStreamCreateWithFlags(&capture_stream, cudaStreamNonBlocking)); + CHECK_CUDA_GLOBAL(active_comm, cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeThreadLocal)); + cudecompResult_t capture_result = CUDECOMP_RESULT_SUCCESS; + for (int dim = 0; dim < 3 && capture_result == CUDECOMP_RESULT_SUCCESS; ++dim) { + capture_result = runHalo(handle, grid_desc, test_case.axis, data_d, work_d, test_case.dtype, pinfo, + test_case.halo_periods, dim, capture_stream); + } + cudaGraph_t graph = nullptr; + CHECK_CUDA_GLOBAL(active_comm, cudaStreamEndCapture(capture_stream, &graph)); + + bool nvshmem_capture_unsupported = cudecomp::haloBackendRequiresNvshmem(test_case.backend.backend) && + !cudecomp::nvshmemSupportsExternalGraphCapture(grid_desc); + if (handle->cuda_graphs_enable || !cudecomp::haloBackendSupportsExternalGraphCapture(test_case.backend.backend) || + nvshmem_capture_unsupported) { + EXPECT_EQ(CUDECOMP_RESULT_NOT_SUPPORTED, capture_result); + } else { + EXPECT_EQ(CUDECOMP_RESULT_SUCCESS, capture_result); + ASSERT_NE(graph, nullptr); + if (test_case.automatic_workspace) { + EXPECT_TRUE(cudecomp::haloBackendRequiresNvshmem(test_case.backend.backend) + ? handle->nvshmem_workspace.capture_frozen + : handle->ordinary_workspace.capture_frozen); + } + + CHECK_CUDA_GLOBAL(active_comm, cudaMemset(data_d, 0, pinfo.size * sizeof(*data_d))); + CHECK_CUDA_GLOBAL(active_comm, + cudaMemcpy(data_d, initial.data(), initial.size() * sizeof(*data_d), cudaMemcpyHostToDevice)); + cudaGraphExec_t graph_exec = nullptr; + CHECK_CUDA_GLOBAL(active_comm, cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0)); + CHECK_CUDA_GLOBAL(active_comm, cudaGraphLaunch(graph_exec, capture_stream)); + CHECK_CUDA_GLOBAL(active_comm, cudaStreamSynchronize(capture_stream)); + + std::vector actual(expected.size(), unsetValue()); + CHECK_CUDA_GLOBAL(active_comm, + cudaMemcpy(actual.data(), data_d, actual.size() * sizeof(*data_d), cudaMemcpyDeviceToHost)); + EXPECT_TRUE(pencilMatches(expected, actual, pinfo)); + CHECK_CUDA_GLOBAL(active_comm, cudaGraphExecDestroy(graph_exec)); + } + if (graph) { CHECK_CUDA_GLOBAL(active_comm, cudaGraphDestroy(graph)); } + CHECK_CUDA_GLOBAL(active_comm, cudaStreamDestroy(capture_stream)); + } } TEST_P(HaloCorrectnessTest, UpdateHalos) { @@ -415,7 +473,22 @@ TEST_P(HaloCorrectnessTest, UpdateHalos) { } } +TEST_P(ExternalCudaGraphHaloCorrectnessTest, CapturesAndReplaysOperation) { + const auto test_case = GetParam(); + switch (test_case.dtype) { + case CUDECOMP_FLOAT: runHaloCase(test_case, true); break; + case CUDECOMP_FLOAT_COMPLEX: runHaloCase>(test_case, true); break; + case CUDECOMP_DOUBLE: runHaloCase(test_case, true); break; + case CUDECOMP_DOUBLE_COMPLEX: runHaloCase>(test_case, true); break; + default: FAIL() << "unsupported test dtype " << test_case.dtype; + } +} + INSTANTIATE_TEST_SUITE_P(MpiBackends, HaloCorrectnessTest, ::testing::ValuesIn(haloCasesForLabel("mpi")), paramName); INSTANTIATE_TEST_SUITE_P(NcclBackends, HaloCorrectnessTest, ::testing::ValuesIn(haloCasesForLabel("nccl")), paramName); INSTANTIATE_TEST_SUITE_P(NvshmemBackends, HaloCorrectnessTest, ::testing::ValuesIn(haloCasesForLabel("nvshmem")), paramName); +INSTANTIATE_TEST_SUITE_P(ExternalCudaGraphNcclBackends, ExternalCudaGraphHaloCorrectnessTest, + ::testing::ValuesIn(externalCudaGraphHaloCases("nccl")), paramName); +INSTANTIATE_TEST_SUITE_P(ExternalCudaGraphNvshmemBackends, ExternalCudaGraphHaloCorrectnessTest, + ::testing::ValuesIn(externalCudaGraphHaloCases("nvshmem")), paramName); diff --git a/tests/ctest/transpose_tests.cc b/tests/ctest/transpose_tests.cc index bd89d7a2..a2ae3ecc 100644 --- a/tests/ctest/transpose_tests.cc +++ b/tests/ctest/transpose_tests.cc @@ -29,6 +29,7 @@ namespace { enum class TransposeOperation { XToY, YToX, YToZ, ZToY }; +enum class GraphExecution { Eager, Internal, External }; constexpr std::array kTransposeOperations{TransposeOperation::XToY, TransposeOperation::YToX, TransposeOperation::YToZ, TransposeOperation::ZToY}; @@ -321,6 +322,19 @@ std::vector cudaGraphTransposeCases() { return cases; } +std::vector externalCudaGraphTransposeCases(const char* label) { + std::vector cases; + for (const auto& backend : cudecomp_test::transposeBackends()) { + if (std::string(backend.label) != label) continue; + + auto test_case = + makeCase(backend, "ExternalCudaGraph", TransposeOperation::XToY, kBaselineGdims, {2, 2}, CUDECOMP_FLOAT, true); + cases.push_back(test_case); + cases.push_back(withAutomaticWorkspace(std::move(test_case))); + } + return cases; +} + std::vector ncclUserBufferRegistrationCases() { std::vector cases; for (const auto& backend : cudecomp_test::transposeBackends()) { @@ -523,6 +537,7 @@ testing::AssertionResult nvshmemTeamsMatchProcessGrid(cudecompHandle_t handle, c class TransposeCorrectnessTest : public ::testing::TestWithParam {}; class CudaGraphTransposeCorrectnessTest : public ::testing::TestWithParam {}; +class ExternalCudaGraphTransposeCorrectnessTest : public ::testing::TestWithParam {}; class NcclUserBufferRegistrationTest : public ::testing::TestWithParam {}; TEST(NcclCommunicatorLifecycleTest, RetainsLocalCommunicatorUntilAllDescriptorsReleaseIt) { @@ -585,7 +600,7 @@ testing::AssertionResult ncclUserBufferRegistrationIsActive(cudecompHandle_t han } template -void runTransposeCase(const TransposeCase& test_case, bool check_cuda_graph_replay = false, +void runTransposeCase(const TransposeCase& test_case, GraphExecution graph_execution = GraphExecution::Eager, bool check_nccl_user_buffer_registration = false) { const int active_ranks = test_case.pdims[0] * test_case.pdims[1]; const auto world_comm = cudecomp_test::MpiTestComm::world(); @@ -624,7 +639,7 @@ void runTransposeCase(const TransposeCase& test_case, bool check_cuda_graph_repl cudecompHandle_t handle = test_context.handle(); ASSERT_NE(handle, nullptr); ASSERT_TRUE(applySyntheticHostnames(handle, test_case.synthetic_host_groups)); - if (check_cuda_graph_replay) { ASSERT_TRUE(handle->cuda_graphs_enable); } + if (graph_execution == GraphExecution::Internal) { ASSERT_TRUE(handle->cuda_graphs_enable); } cudecompGridDesc_t grid_desc = nullptr; const cudecompResult_t grid_desc_create_result = cudecompGridDescCreate(handle, &grid_desc, &config, nullptr); @@ -663,7 +678,7 @@ void runTransposeCase(const TransposeCase& test_case, bool check_cuda_graph_repl cudecomp_test::cudaBufferGuard input_buffer(input_d); CHECK_CUDA_GLOBAL(active_comm, input_alloc_result); - if (check_cuda_graph_replay) { + if (graph_execution == GraphExecution::Internal) { ASSERT_TRUE(test_case.out_of_place); T* output_a_d = nullptr; @@ -733,15 +748,45 @@ void runTransposeCase(const TransposeCase& test_case, bool check_cuda_graph_repl EXPECT_EQ(cached_workspace, cudecomp::transposeBackendRequiresNvshmem(test_case.backend.backend) ? handle->nvshmem_workspace.ptr : handle->ordinary_workspace.ptr); + } + if (graph_execution == GraphExecution::External) { cudaStream_t capture_stream = nullptr; CHECK_CUDA_GLOBAL(active_comm, cudaStreamCreateWithFlags(&capture_stream, cudaStreamNonBlocking)); CHECK_CUDA_GLOBAL(active_comm, cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeThreadLocal)); - EXPECT_EQ(CUDECOMP_RESULT_NOT_SUPPORTED, - runTranspose(handle, grid_desc, test_case.operation, input_d, output_d, CUDECOMP_WORKSPACE_AUTO, - test_case.dtype, input_info, output_info, capture_stream)); + const cudecompResult_t capture_result = + runTranspose(handle, grid_desc, test_case.operation, input_d, output_d, work_d, test_case.dtype, input_info, + output_info, capture_stream); cudaGraph_t graph = nullptr; CHECK_CUDA_GLOBAL(active_comm, cudaStreamEndCapture(capture_stream, &graph)); + + bool nvshmem_capture_unsupported = cudecomp::transposeBackendRequiresNvshmem(test_case.backend.backend) && + !cudecomp::nvshmemSupportsExternalGraphCapture(grid_desc); + if (handle->cuda_graphs_enable || + !cudecomp::transposeBackendSupportsExternalGraphCapture(test_case.backend.backend) || + nvshmem_capture_unsupported) { + EXPECT_EQ(CUDECOMP_RESULT_NOT_SUPPORTED, capture_result); + } else { + EXPECT_EQ(CUDECOMP_RESULT_SUCCESS, capture_result); + ASSERT_NE(graph, nullptr); + if (test_case.automatic_workspace) { + EXPECT_TRUE(cudecomp::transposeBackendRequiresNvshmem(test_case.backend.backend) + ? handle->nvshmem_workspace.capture_frozen + : handle->ordinary_workspace.capture_frozen); + } + + cudaGraphExec_t graph_exec = nullptr; + CHECK_CUDA_GLOBAL(active_comm, cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0)); + CHECK_CUDA_GLOBAL(active_comm, cudaGraphLaunch(graph_exec, capture_stream)); + CHECK_CUDA_GLOBAL(active_comm, cudaGraphLaunch(graph_exec, capture_stream)); + CHECK_CUDA_GLOBAL(active_comm, cudaStreamSynchronize(capture_stream)); + + std::vector actual(output_ref.size()); + CHECK_CUDA_GLOBAL(active_comm, + cudaMemcpy(actual.data(), output_d, actual.size() * sizeof(*output_d), cudaMemcpyDeviceToHost)); + EXPECT_TRUE(pencilMatches(output_ref, actual, output_info)); + CHECK_CUDA_GLOBAL(active_comm, cudaGraphExecDestroy(graph_exec)); + } if (graph) { CHECK_CUDA_GLOBAL(active_comm, cudaGraphDestroy(graph)); } CHECK_CUDA_GLOBAL(active_comm, cudaStreamDestroy(capture_stream)); } @@ -770,10 +815,21 @@ TEST_P(TransposeCorrectnessTest, DirectOperation) { TEST_P(CudaGraphTransposeCorrectnessTest, CapturesAndReplaysPackingGraph) { const auto test_case = GetParam(); switch (test_case.dtype) { - case CUDECOMP_FLOAT: runTransposeCase(test_case, true); break; - case CUDECOMP_FLOAT_COMPLEX: runTransposeCase>(test_case, true); break; - case CUDECOMP_DOUBLE: runTransposeCase(test_case, true); break; - case CUDECOMP_DOUBLE_COMPLEX: runTransposeCase>(test_case, true); break; + case CUDECOMP_FLOAT: runTransposeCase(test_case, GraphExecution::Internal); break; + case CUDECOMP_FLOAT_COMPLEX: runTransposeCase>(test_case, GraphExecution::Internal); break; + case CUDECOMP_DOUBLE: runTransposeCase(test_case, GraphExecution::Internal); break; + case CUDECOMP_DOUBLE_COMPLEX: runTransposeCase>(test_case, GraphExecution::Internal); break; + default: FAIL() << "unsupported test dtype " << test_case.dtype; + } +} + +TEST_P(ExternalCudaGraphTransposeCorrectnessTest, CapturesAndReplaysOperation) { + const auto test_case = GetParam(); + switch (test_case.dtype) { + case CUDECOMP_FLOAT: runTransposeCase(test_case, GraphExecution::External); break; + case CUDECOMP_FLOAT_COMPLEX: runTransposeCase>(test_case, GraphExecution::External); break; + case CUDECOMP_DOUBLE: runTransposeCase(test_case, GraphExecution::External); break; + case CUDECOMP_DOUBLE_COMPLEX: runTransposeCase>(test_case, GraphExecution::External); break; default: FAIL() << "unsupported test dtype " << test_case.dtype; } } @@ -784,10 +840,10 @@ TEST_P(NcclUserBufferRegistrationTest, DirectOperation) { #else const auto test_case = GetParam(); switch (test_case.dtype) { - case CUDECOMP_FLOAT: runTransposeCase(test_case, false, true); break; - case CUDECOMP_FLOAT_COMPLEX: runTransposeCase>(test_case, false, true); break; - case CUDECOMP_DOUBLE: runTransposeCase(test_case, false, true); break; - case CUDECOMP_DOUBLE_COMPLEX: runTransposeCase>(test_case, false, true); break; + case CUDECOMP_FLOAT: runTransposeCase(test_case, GraphExecution::Eager, true); break; + case CUDECOMP_FLOAT_COMPLEX: runTransposeCase>(test_case, GraphExecution::Eager, true); break; + case CUDECOMP_DOUBLE: runTransposeCase(test_case, GraphExecution::Eager, true); break; + case CUDECOMP_DOUBLE_COMPLEX: runTransposeCase>(test_case, GraphExecution::Eager, true); break; default: FAIL() << "unsupported test dtype " << test_case.dtype; } #endif @@ -801,5 +857,9 @@ INSTANTIATE_TEST_SUITE_P(NvshmemBackends, TransposeCorrectnessTest, ::testing::ValuesIn(transposeCasesForLabel("nvshmem")), paramName); INSTANTIATE_TEST_SUITE_P(CudaGraphMpiBackends, CudaGraphTransposeCorrectnessTest, ::testing::ValuesIn(cudaGraphTransposeCases()), paramName); +INSTANTIATE_TEST_SUITE_P(ExternalCudaGraphNcclBackends, ExternalCudaGraphTransposeCorrectnessTest, + ::testing::ValuesIn(externalCudaGraphTransposeCases("nccl")), paramName); +INSTANTIATE_TEST_SUITE_P(ExternalCudaGraphNvshmemBackends, ExternalCudaGraphTransposeCorrectnessTest, + ::testing::ValuesIn(externalCudaGraphTransposeCases("nvshmem")), paramName); INSTANTIATE_TEST_SUITE_P(NcclUserBufferRegistration, NcclUserBufferRegistrationTest, ::testing::ValuesIn(ncclUserBufferRegistrationCases()), paramName);