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

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

@@ -11,6 +11,8 @@
 #include <cuda.h>
 #include <cuda_runtime.h>
 #include <stddef.h>
+#include <stdint.h>
+#include <stdlib.h>
 #include <string.h>
 
 #include "../cuda/ceed-cuda-common.h"
@@ -31,7 +33,25 @@ static int CeedOperatorDestroy_Cuda_gen(CeedOperator op) {
   if (impl->module_assemble_full) CeedCallCuda(ceed, cuModuleUnload(impl->module_assemble_full));
   if (impl->module_assemble_diagonal) CeedCallCuda(ceed, cuModuleUnload(impl->module_assemble_diagonal));
   if (impl->module_assemble_qfunction) CeedCallCuda(ceed, cuModuleUnload(impl->module_assemble_qfunction));
-  if (impl->points.num_per_elem) CeedCallCuda(ceed, cudaFree((void **)impl->points.num_per_elem));
+  if (impl->points.num_per_elem) CeedCallCuda(ceed, cudaFree((void *)impl->points.num_per_elem));
+
+  if (impl->graph_created && impl->graph_launches > 0) {
+    char *op_name = NULL;
+    CeedOperatorGetName(op, (const char **)&op_name);
+    printf("[CUDA Graph] Summary for operator '%s': %d graph launches, %d fallbacks\n", op_name ? op_name : "unnamed", impl->graph_launches,
+           impl->fallbacks);
+  }
+  if (impl->graph_instance) {
+    cudaGraphExecDestroy(impl->graph_instance);
+    impl->graph_instance = NULL;
+  }
+  if (impl->graph) {
+    cudaGraphDestroy(impl->graph);
+    impl->graph = NULL;
+  }
+  impl->graph_created      = false;
+  impl->captured_input_ptr = NULL;
+
   CeedCallBackend(CeedFree(&impl));
   CeedCallBackend(CeedDestroy(&ceed));
   return CEED_ERROR_SUCCESS;
@@ -284,7 +304,11 @@ static int CeedOperatorApplyAdd_Cuda_gen(CeedOperator op, CeedVector input_vec,
   // Try to run kernel
   if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(input_vec, CEED_MEM_DEVICE, &input_arr));
   if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArray(output_vec, CEED_MEM_DEVICE, &output_arr));
-  CeedCallBackend(CeedOperatorApplyAddCore_Cuda_gen(op, NULL, input_arr, output_arr, &is_run_good, request));
+  // During graph capture use the capturing stream, otherwise the default stream.
+  enum cudaStreamCaptureStatus capture_status;
+  cudaStreamIsCapturing(cudaStreamPerThread, &capture_status);
+  CUstream stream_to_use = (capture_status != cudaStreamCaptureStatusNone) ? cudaStreamPerThread : NULL;
+  CeedCallBackend(CeedOperatorApplyAddCore_Cuda_gen(op, stream_to_use, input_arr, output_arr, &is_run_good, request));
   if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &input_arr));
   if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArray(output_vec, &output_arr));
 
@@ -299,48 +323,215 @@ static int CeedOperatorApplyAdd_Cuda_gen(CeedOperator op, CeedVector input_vec,
   return CEED_ERROR_SUCCESS;
 }
 
+// Push each suboperator's QFunction context to device. Replay skips the normal
+// apply path, so we do this by hand to keep time/load parameters current.
+static int CeedCompositeRefreshContexts_Cuda_gen(CeedOperator *sub_operators, CeedInt num_suboperators) {
+  for (CeedInt i = 0; i < num_suboperators; i++) {
+    CeedQFunction qf  = NULL;
+    void         *d_c = NULL;
+
+    CeedCallBackend(CeedOperatorGetQFunction(sub_operators[i], &qf));
+    CeedCallBackend(CeedQFunctionGetInnerContextData(qf, CEED_MEM_DEVICE, &d_c));
+    CeedCallBackend(CeedQFunctionRestoreInnerContextData(qf, &d_c));
+    CeedCallBackend(CeedQFunctionDestroy(&qf));
+  }
+  return CEED_ERROR_SUCCESS;
+}
+
 static int CeedOperatorApplyAddComposite_Cuda_gen(CeedOperator op, CeedVector input_vec, CeedVector output_vec, CeedRequest *request) {
-  bool              is_run_good[CEED_COMPOSITE_MAX] = {false}, is_sequential;
-  CeedInt           num_suboperators;
-  const CeedScalar *input_arr  = NULL;
-  CeedScalar       *output_arr = NULL;
-  Ceed              ceed;
-  CeedOperator     *sub_operators;
-  cudaStream_t      stream = NULL;
+  Ceed                   ceed;
+  CeedOperator_Cuda_gen *impl;
+  CeedOperator          *sub_operators;
+  CeedInt                num_suboperators;
+  char                  *op_name = NULL;
 
-  CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
+  ceed = CeedOperatorReturnCeed(op);
   CeedCall(CeedOperatorCompositeGetNumSub(op, &num_suboperators));
   CeedCall(CeedOperatorCompositeGetSubList(op, &sub_operators));
-  CeedCall(CeedOperatorCompositeIsSequential(op, &is_sequential));
-  if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(input_vec, CEED_MEM_DEVICE, &input_arr));
-  if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArray(output_vec, CEED_MEM_DEVICE, &output_arr));
-  if (is_sequential) CeedCallCuda(ceed, cudaStreamCreate(&stream));
-  for (CeedInt i = 0; i < num_suboperators; i++) {
-    CeedInt num_elem = 0;
+  CeedCallBackend(CeedOperatorGetData(op, &impl));
+  CeedCallBackend(CeedOperatorGetName(op, (const char **)&op_name));
+
+  // CEED_FORCE_BASELINE=1: skip graphs, run via /gpu/cuda/ref
+  static bool force_baseline         = false;
+  static bool force_baseline_checked = false;
+  if (!force_baseline_checked) {
+    char *env_val          = getenv("CEED_FORCE_BASELINE");
+    force_baseline         = (env_val != NULL && strcmp(env_val, "1") == 0);
+    force_baseline_checked = true;
+  }
+  if (force_baseline) {
+    CeedOperator op_fallback;
+    CeedCallBackend(CeedOperatorGetFallback(op, &op_fallback));
+    CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
+    return CEED_ERROR_SUCCESS;
+  }
 
-    CeedCall(CeedOperatorGetNumElements(sub_operators[i], &num_elem));
-    if (num_elem > 0) {
-      if (!is_sequential) CeedCallCuda(ceed, cudaStreamCreate(&stream));
-      CeedCallBackend(CeedOperatorApplyAddCore_Cuda_gen(sub_operators[i], stream, input_arr, output_arr, &is_run_good[i], request));
-      if (!is_sequential) CeedCallCuda(ceed, cudaStreamDestroy(stream));
+  // CEED_DISABLE_GRAPH=1: run suboperators directly, no capture/replay (for benchmarking).
+  static bool disable_graph         = false;
+  static bool disable_graph_checked = false;
+  if (!disable_graph_checked) {
+    char *env_val         = getenv("CEED_DISABLE_GRAPH");
+    disable_graph         = (env_val != NULL && strcmp(env_val, "1") == 0);
+    disable_graph_checked = true;
+  }
+  if (disable_graph) {
+    for (CeedInt i = 0; i < num_suboperators; i++) {
+      CeedCallBackend(CeedOperatorApplyAdd(sub_operators[i], input_vec, output_vec, request));
     }
+    return CEED_ERROR_SUCCESS;
   }
-  if (is_sequential) CeedCallCuda(ceed, cudaStreamDestroy(stream));
-  if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &input_arr));
-  if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArray(output_vec, &output_arr));
-  CeedCallCuda(ceed, cudaDeviceSynchronize());
 
-  // Fallback on unsuccessful run
-  for (CeedInt i = 0; i < num_suboperators; i++) {
-    if (!is_run_good[i]) {
-      CeedOperator op_fallback;
+  // No real I/O buffers to capture; just run directly.
+  if (input_vec == CEED_VECTOR_NONE || output_vec == CEED_VECTOR_NONE) {
+    for (CeedInt i = 0; i < num_suboperators; i++) {
+      CeedCallBackend(CeedOperatorApplyAdd(sub_operators[i], input_vec, output_vec, request));
+    }
+    return CEED_ERROR_SUCCESS;
+  }
+
+  // Phase 1: first call runs directly so lazy allocations happen before capture.
+  if (!impl->warmup_done) {
+    for (CeedInt i = 0; i < num_suboperators; i++) {
+      CeedCallBackend(CeedOperatorApplyAdd(sub_operators[i], input_vec, output_vec, request));
+    }
+    CeedCallCuda(ceed, cudaDeviceSynchronize());
+    impl->warmup_done = true;
+    return CEED_ERROR_SUCCESS;
+  }
+
+  // Phase 2: capture.
+  if (!impl->graph_created) {
+    {
+      const CeedScalar *in_ptr;
+      CeedCallBackend(CeedVectorGetArrayRead(input_vec, CEED_MEM_DEVICE, &in_ptr));
+      impl->captured_input_ptr = in_ptr;
+      CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &in_ptr));
+    }
+
+    printf("[CUDA Graph] Phase 2: Recording graph for operator '%s'...\n", op_name ? op_name : "unnamed");
+    cudaStream_t capture_stream = cudaStreamPerThread;
+    bool         capture_ok     = true;
 
-      CeedDebug(ceed, "\nFalling back to /gpu/cuda/ref CeedOperator for ApplyAdd\n");
-      CeedCallBackend(CeedOperatorGetFallback(sub_operators[i], &op_fallback));
+    cudaError_t err = cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeThreadLocal);
+    if (err != cudaSuccess) {
+      capture_ok = false;
+    }
+
+    if (capture_ok) {
+      // Check errors by hand (not CeedCallBackend) so we always reach the
+      // cudaStreamEndCapture below. Some sub-operators (e.g. contact, which calls
+      // cudaMalloc) invalidate the capture, and we still need to close it out.
+      for (CeedInt i = 0; i < num_suboperators; i++) {
+        if (CeedOperatorApplyAdd(sub_operators[i], input_vec, output_vec, CEED_REQUEST_IMMEDIATE)) {
+          capture_ok = false;
+          break;
+        }
+      }
+    }
+
+    // Always end capture so the stream is usable again, even if it failed.
+    {
+      cudaGraph_t partial = NULL;
+      err                 = cudaStreamEndCapture(capture_stream, &partial);
+      if (capture_ok && (err != cudaSuccess || !partial)) capture_ok = false;
+      if (capture_ok)
+        impl->graph = partial;
+      else if (partial)
+        cudaGraphDestroy(partial);
+    }
+
+    if (capture_ok) {
+      err = cudaGraphInstantiate(&impl->graph_instance, impl->graph, 0);
+      if (err != cudaSuccess) {
+        cudaGraphDestroy(impl->graph);
+        impl->graph = NULL;
+        capture_ok  = false;
+      }
+    }
+
+    // Clear the leftover error from a failed capture so later CUDA calls don't inherit it.
+    if (!capture_ok) {
+      cudaGetLastError();
+      cudaDeviceSynchronize();
+      cudaGetLastError();
+    }
+
+    impl->graph_created  = true;
+    impl->graph_launches = 0;
+
+    if (capture_ok)
+      printf("[CUDA Graph] Graph created for operator '%s'\n", op_name ? op_name : "unnamed");
+    else
+      printf("[CUDA Graph] Capture disabled for operator '%s'; falling back to direct apply\n", op_name ? op_name : "unnamed");
+
+    // Capture doesn't run the kernels, so apply directly to get this call's output.
+    for (CeedInt i = 0; i < num_suboperators; i++) {
+      CeedCallBackend(CeedOperatorApplyAdd(sub_operators[i], input_vec, output_vec, request));
+    }
+    return CEED_ERROR_SUCCESS;
+  }
+
+  // Phase 3: replay.
+  {
+    // graph_instance == NULL means capture failed; direct-apply every call.
+    if (!impl->graph_instance) {
+      for (CeedInt i = 0; i < num_suboperators; i++) {
+        CeedCallBackend(CeedOperatorApplyAdd(sub_operators[i], input_vec, output_vec, request));
+      }
+      return CEED_ERROR_SUCCESS;
+    }
+
+    cudaStream_t stream = NULL;
+
+    // Sync contexts so kernels see updated time/load parameters.
+    CeedCallBackend(CeedCompositeRefreshContexts_Cuda_gen(sub_operators, num_suboperators));
+
+    // If the input buffer moved since capture, recapture.
+    {
+      const CeedScalar *in_ptr;
+      bool              ptr_ok;
+
+      CeedCallBackend(CeedVectorGetArrayRead(input_vec, CEED_MEM_DEVICE, &in_ptr));
+      ptr_ok = (in_ptr == impl->captured_input_ptr);
+      CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &in_ptr));
+      if (!ptr_ok) goto use_fallback;
+    }
+
+    cudaError_t err = cudaGraphLaunch(impl->graph_instance, stream);
+    if (err != cudaSuccess) {
+      printf("CUDA Graph launch failed: %s - using fallback\n", cudaGetErrorString(err));
+      CeedOperator op_fallback;
+      CeedCallBackend(CeedOperatorGetFallback(op, &op_fallback));
       CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
+      return CEED_ERROR_SUCCESS;
     }
+
+    if (impl->graph_launches == 0) {
+      printf("[CUDA Graph] ✓ Replaying graph for operator '%s'\n", op_name ? op_name : "unnamed");
+    }
+    impl->graph_launches++;
+
+    return CEED_ERROR_SUCCESS;
   }
-  CeedCallBackend(CeedDestroy(&ceed));
+
+// Drop the graph and run via /gpu/cuda/ref; next call will recapture.
+use_fallback:
+  if (impl->graph_instance) {
+    cudaGraphExecDestroy(impl->graph_instance);
+    impl->graph_instance = NULL;
+  }
+  if (impl->graph) {
+    cudaGraphDestroy(impl->graph);
+    impl->graph = NULL;
+  }
+  impl->graph_created      = false;
+  impl->captured_input_ptr = NULL;
+  impl->fallbacks++;
+
+  CeedOperator op_fallback;
+  CeedCallBackend(CeedOperatorGetFallback(op, &op_fallback));
+  CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
+
   return CEED_ERROR_SUCCESS;
 }
 
@@ -465,7 +656,7 @@ static int CeedOperatorLinearAssembleQFunctionCore_Cuda_gen(CeedOperator op, boo
 
     // Assemble QFunction
     void *opargs[] = {(void *)&num_elem, &qf_data->d_c, &data->indices, &data->fields, &data->B, &data->G, &data->W, &data->points, &assembled_array};
-    bool  is_tensor = false;
+    bool  is_tensor;
     int   max_threads_per_block, min_grid_size, grid;
 
     CeedCallBackend(CeedOperatorHasTensorBases(op, &is_tensor));
@@ -885,6 +1076,15 @@ int CeedOperatorCreate_Cuda_gen(CeedOperator op) {
   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
   CeedCallBackend(CeedCalloc(1, &impl));
   CeedCallBackend(CeedOperatorSetData(op, impl));
+
+  impl->graph_created      = false;
+  impl->warmup_done        = false;
+  impl->graph              = NULL;
+  impl->graph_instance     = NULL;
+  impl->graph_launches     = 0;
+  impl->fallbacks          = 0;
+  impl->captured_input_ptr = NULL;
+
   CeedCall(CeedOperatorIsComposite(op, &is_composite));
   if (is_composite) {
     CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAddComposite", CeedOperatorApplyAddComposite_Cuda_gen));
@@ -897,6 +1097,7 @@ int CeedOperatorCreate_Cuda_gen(CeedOperator op) {
                                            CeedOperatorLinearAssembleAddDiagonalAtPoints_Cuda_gen));
     CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleSingle", CeedOperatorAssembleSingleAtPoints_Cuda_gen));
   }
+
   if (!is_at_points) {
     CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleQFunction", CeedOperatorLinearAssembleQFunction_Cuda_gen));
     CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleQFunctionUpdate",

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;
+
+  // CUDA graph state
+  bool              graph_created;
+  bool              warmup_done;
+  cudaGraph_t       graph;
+  cudaGraphExec_t   graph_instance;
+  int               graph_launches;
+  int               fallbacks;
+  const CeedScalar *captured_input_ptr;  // device address at capture; checked before each replay
 } CeedOperator_Cuda_gen;
 
 typedef struct {

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

@@ -36,7 +36,16 @@ 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,17 @@ 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)));
+      // Check if we're in CUDA Graph capture mode
+      enum cudaStreamCaptureStatus capture_status;
+      cudaStreamIsCapturing(cudaStreamPerThread, &capture_status);
+
+      if (capture_status != cudaStreamCaptureStatusNone) {
+        // During capture, use async memset with cudaStreamPerThread
+        CeedCallCuda(CeedVectorReturnCeed(vec), cudaMemsetAsync(impl->d_array, 0, length * sizeof(CeedScalar), cudaStreamPerThread));
+      } else {
+        // Normal execution, use blocking memset
+        CeedCallCuda(CeedVectorReturnCeed(vec), cudaMemset(impl->d_array, 0, length * sizeof(CeedScalar)));
+      }
     } else {
       CeedCallBackend(CeedDeviceSetValue_Cuda(impl->d_array, length, val));
     }