Merge pull request #2006 from CEED/test/cuda-graph

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

Author: jeremylt

CHANGELOG.md

@@ -35,6 +35,7 @@ Specifically, directories set with `CeedAddJitSourceRoot(ceed, "foo/bar")` will
 - Add `CeedGetBuildConfiguration()` to access compilers, flags, and related information about the build environment.
 - Add support for full `CeedOperator` assembly for operators with multiple active fields with different bases for CPU backends and `/gpu/cuda/ref` and `/gpu/hip/gen` backends.
 - Add `CeedVectorFilter` to zero out components of a `CeedVector` that have absolute value below a specified threshold value.
+- 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.
 
 ### Examples
 

backends/cuda-gen/ceed-cuda-gen-operator.c

@@ -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-gen/ceed-cuda-gen.h

@@ -36,7 +36,9 @@ 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));
+
+  CeedCallCuda(ceed, cudaMemcpyAsync(impl->d_data, impl->h_data, ctx_size, cudaMemcpyHostToDevice, cudaStreamPerThread));
+
   CeedCallBackend(CeedDestroy(&ceed));
   return CEED_ERROR_SUCCESS;
 }

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

@@ -57,7 +57,7 @@ static inline int CeedVectorSyncH2D_Cuda(const CeedVector vec) {
     CeedCallCuda(CeedVectorReturnCeed(vec), cudaMalloc((void **)&impl->d_array_owned, bytes));
     impl->d_array = impl->d_array_owned;
   }
-  CeedCallCuda(CeedVectorReturnCeed(vec), cudaMemcpy(impl->d_array, impl->h_array, bytes, cudaMemcpyHostToDevice));
+  CeedCallCuda(CeedVectorReturnCeed(vec), cudaMemcpyAsync(impl->d_array, impl->h_array, bytes, cudaMemcpyHostToDevice, cudaStreamPerThread));
   return CEED_ERROR_SUCCESS;
 }
 
@@ -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));
     }

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

@@ -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-impl.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);

include/ceed/cuda.h

@@ -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-cuda.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
   };