cuda-gen: Add CUDA Graph capture and replay for composite operators

CHANGELOG.md

@@ -23,6 +23,7 @@ These functions will be removed when the SYCL backends are updated to reflect th
 
 ### New features
 
+- Add `CeedOperatorSetEnableCudaGraph` for CUDA Graph capture/replay on `/gpu/cuda/gen` composite operators. Enabled by default; use `CEED_ENABLE_CUDA_GRAPH=0` to turn off.
 - Add `CeedOperatorCreateAtPoints` which evaluates the `CeedQFunction` at arbitrary locations in each element, for use in Particle in Cell, Material Point Method, and similar methods.
 - Add `CeedElemRestrictionGetLLayout` to provide L-vector layout for strided `CeedElemRestriction` created with `CEED_BACKEND_STRIDES`.
 - Add `CeedVectorReturnCeed` and similar when parent `Ceed` context for a libCEED object is only needed once in a calling scope.

backends/cuda-gen/ceed-cuda-gen.h

@@ -10,6 +10,7 @@
 #include <ceed/backend.h>
 #include <ceed/jit-source/cuda/cuda-types.h>
 #include <cuda.h>
+#include <cuda_runtime.h>
 
 typedef struct {
   bool           use_fallback, use_assembly_fallback;
@@ -25,6 +26,15 @@ typedef struct {
   Fields_Cuda    G;
   CeedScalar    *W;
   Points_Cuda    points;
+
+  // Graph capture data
+  bool              use_graph;
+  bool              graph_created;
+  bool              warmup_done;
+  cudaGraph_t       graph;
+  cudaGraphExec_t   graph_instance;
+  const CeedScalar *captured_input_ptr;
+  CeedScalar       *captured_output_ptr;
 } CeedOperator_Cuda_gen;
 
 typedef struct {

backends/cuda-ref/ceed-cuda-ref-qfunctioncontext.c

@@ -36,7 +36,17 @@ static inline int CeedQFunctionContextSyncH2D_Cuda(const CeedQFunctionContext ct
     CeedCallCuda(ceed, cudaMalloc((void **)&impl->d_data_owned, ctx_size));
     impl->d_data = impl->d_data_owned;
   }
-  CeedCallCuda(ceed, cudaMemcpy(impl->d_data, impl->h_data, ctx_size, cudaMemcpyHostToDevice));
+
+  // Use async memcpy during CUDA Graph capture for compatibility
+  enum cudaStreamCaptureStatus capture_status;
+
+  cudaStreamIsCapturing(cudaStreamPerThread, &capture_status);
+  if (capture_status != cudaStreamCaptureStatusNone) {
+    CeedCallCuda(ceed, cudaMemcpyAsync(impl->d_data, impl->h_data, ctx_size, cudaMemcpyHostToDevice, cudaStreamPerThread));
+  } else {
+    CeedCallCuda(ceed, cudaMemcpy(impl->d_data, impl->h_data, ctx_size, cudaMemcpyHostToDevice));
+  }
+
   CeedCallBackend(CeedDestroy(&ceed));
   return CEED_ERROR_SUCCESS;
 }

backends/cuda-ref/ceed-cuda-ref-vector.c

@@ -326,7 +326,7 @@ static int CeedVectorSetValue_Cuda(CeedVector vec, CeedScalar val) {
   }
   if (impl->d_array) {
     if (val == 0) {
-      CeedCallCuda(CeedVectorReturnCeed(vec), cudaMemset(impl->d_array, 0, length * sizeof(CeedScalar)));
+      CeedCallCuda(CeedVectorReturnCeed(vec), cudaMemsetAsync(impl->d_array, 0, length * sizeof(CeedScalar), cudaStreamPerThread));
     } else {
       CeedCallBackend(CeedDeviceSetValue_Cuda(impl->d_array, length, val));
     }

include/ceed-impl.h

@@ -368,6 +368,7 @@ struct CeedOperator_private {
   int (*ApplyAdd)(CeedOperator, CeedVector, CeedVector, CeedRequest *);
   int (*ApplyAddComposite)(CeedOperator, CeedVector, CeedVector, CeedRequest *);
   int (*ApplyJacobian)(CeedOperator, CeedVector, CeedVector, CeedVector, CeedVector, CeedRequest *);
+  int (*SetEnableCudaGraph)(CeedOperator, bool);
   int (*Destroy)(CeedOperator);
   CeedOperatorField        *input_fields;
   CeedOperatorField        *output_fields;

include/ceed/cuda.h

@@ -13,3 +13,4 @@
 #include <cuda.h>
 
 CEED_EXTERN int CeedQFunctionSetCUDAUserFunction(CeedQFunction qf, CUfunction f);
+CEED_EXTERN int CeedOperatorSetEnableCudaGraph(CeedOperator op, bool enable_graph);

interface/ceed-cuda.c

@@ -12,7 +12,10 @@
 #include <cuda.h>
 
 /**
-  @brief Set CUDA function pointer to evaluate action at quadrature points
+  @brief Set CUDA function pointer to evaluate action at quadrature points.
+
+  If the backend does not support `CUfunction` pointers for QFunctions, then the call succeeds without effect.
+  When unsupported, a message is emitted via `CeedDebug`.
 
   @param[in,out] qf `CeedQFunction` to set device pointer
   @param[in]     f  Device function pointer to evaluate action at quadrature points
@@ -29,3 +32,25 @@ int CeedQFunctionSetCUDAUserFunction(CeedQFunction qf, CUfunction f) {
   }
   return CEED_ERROR_SUCCESS;
 }
+
+/**
+  @brief Enable or disable CUDA Graph capture/replay for a `CeedOperator`.
+
+  If the backend does not support CUDA Graphs for operators, then the call succeeds without effect.
+  When unsupported, a message is emitted via `CeedDebug`.
+
+  @param[in,out] op           `CeedOperator`
+  @param[in]     enable_graph Boolean flag to enable CUDA Graph use
+
+  @return An error code: 0 - success, otherwise - failure
+
+  @ref User
+**/
+int CeedOperatorSetEnableCudaGraph(CeedOperator op, bool enable_graph) {
+  if (!op->SetEnableCudaGraph) {
+    CeedDebug(CeedOperatorReturnCeed(op), "Backend does not support CUDA Graphs for operators.");
+  } else {
+    CeedCall(op->SetEnableCudaGraph(op, enable_graph));
+  }
+  return CEED_ERROR_SUCCESS;
+}

interface/ceed.c

@@ -1374,6 +1374,7 @@ int CeedInit(const char *resource, Ceed *ceed) {
       CEED_FTABLE_ENTRY(CeedOperator, ApplyAdd),
       CEED_FTABLE_ENTRY(CeedOperator, ApplyAddComposite),
       CEED_FTABLE_ENTRY(CeedOperator, ApplyJacobian),
+      CEED_FTABLE_ENTRY(CeedOperator, SetEnableCudaGraph),
       CEED_FTABLE_ENTRY(CeedOperator, Destroy),
       {NULL, 0}  // End of lookup table - used in SetBackendFunction loop
   };