Improve and document support for external CUDA graph capture.

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::
 

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);

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 <cuda-graph-capture-ref>`.
 
 CUDECOMP_ENABLE_PERFORMANCE_REPORT
 ------------------------------------

docs/index.rst

@@ -27,6 +27,7 @@ Table of Contents
 
    overview
    basic_usage
+   cuda_graph_capture
    autotuning
    nvshmem
    api

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

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<float>) { 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); }

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<void*, size_t> 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;