Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Nafees01 wants to merge 21 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from 16 commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
6261385
CUDA graph implemented
Nafees01 Jul 25, 2025
b156a00
Implement CUDA graphs with fallback system in libCEED
Nafees01 Aug 6, 2025
83c1d9f
Cleaner and Simplified Implementation of CUDA graph
Nafees01 Aug 6, 2025
bbad4e2
add vector pointer tracking to detect memory changes in Graph
Nafees01 Aug 7, 2025
285b661
Add PETSc vector setup for CUDA Graph compatibility in CUDA-gen backend
Nafees01 Aug 19, 2025
10d6dd2
Clean CUDA-graph implementation
Nafees01 Aug 21, 2025
4abc7ad
Simple and clean implementation of CUDA graph but numerically incorrect
Nafees01 Aug 26, 2025
1fa80ed
per-operator CUDA Graph implementation
Nafees01 Sep 11, 2025
5b8e14f
replace cudaMemset with cudaMemsetAsync
Nafees01 Sep 25, 2025
94bdeca
Auto-detect graph capture and use async memset with cudaStreamPerThread
Nafees01 Oct 21, 2025
a00af53
CUDA Graph support for composite operators in cuda-gen backend
Nafees01 Oct 28, 2025
1ca2b5f
CUDA Graph is working fine for composite operators
Nafees01 Oct 28, 2025
719027b
cuda-gen: CUDA Graph capture and replay working for composite operators
Nafees01 Jul 7, 2026
40f8c3e
Merge remote-tracking branch 'upstream/main' into cuda-graph-dev
Nafees01 Jul 8, 2026
91c8d6d
style: apply clang-format-22
Nafees01 Jul 8, 2026
1249673
cuda-gen: address review feedback for composite CUDA graphs
Nafees01 Jul 22, 2026
c7c5b1a
cuda: add CeedOperatorSetEnableCudaGraph and address review feedback
Nafees01 Jul 24, 2026
8f160e9
cuda: add docs for when graph/CUfunction setters aren't supported
Nafees01 Jul 27, 2026
8f717af
cuda-gen: add output pointer check and address review fixes
Nafees01 Jul 28, 2026
8ef2be7
cuda-gen: address review feedback for ceed handling and async memset
Nafees01 Aug 12, 2026
b170ec4
address review feedback
Nafees01 Aug 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
204 changes: 199 additions & 5 deletions backends/cuda-gen/ceed-cuda-gen-operator.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -31,7 +33,19 @@ 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_instance) {
cudaGraphExecDestroy(impl->graph_instance);
impl->graph_instance = NULL;
}
if (impl->graph) {
cudaGraphDestroy(impl->graph);
impl->graph = NULL;
Comment thread
jeremylt marked this conversation as resolved.
}
impl->graph_created = false;
impl->captured_input_ptr = NULL;

CeedCallBackend(CeedFree(&impl));
CeedCallBackend(CeedDestroy(&ceed));
return CEED_ERROR_SUCCESS;
Expand Down Expand Up @@ -284,7 +298,10 @@ 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));
enum cudaStreamCaptureStatus capture_status;
cudaStreamIsCapturing(cudaStreamPerThread, &capture_status);
Comment thread
jeremylt marked this conversation as resolved.
Outdated
CUstream stream_to_use = (capture_status != cudaStreamCaptureStatusNone) ? cudaStreamPerThread : NULL;
Comment thread
jeremylt marked this conversation as resolved.
Outdated
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));

Expand All @@ -299,7 +316,58 @@ static int CeedOperatorApplyAdd_Cuda_gen(CeedOperator op, CeedVector input_vec,
return CEED_ERROR_SUCCESS;
}

static int CeedOperatorApplyAddComposite_Cuda_gen(CeedOperator op, CeedVector input_vec, CeedVector output_vec, CeedRequest *request) {
// Push passive inputs and QFunction context to device before graph replay.
static int CeedCompositeRefreshForReplay_Cuda_gen(CeedOperator *sub_operators, CeedInt num_suboperators) {
for (CeedInt i = 0; i < num_suboperators; i++) {
bool is_at_points;
CeedInt num_input_fields, num_output_fields;
CeedOperatorField *op_input_fields, *op_output_fields;
CeedQFunction qf = NULL;
CeedQFunctionField *qf_input_fields;
void *d_c = NULL;
Comment thread
jeremylt marked this conversation as resolved.
Outdated

CeedCallBackend(CeedOperatorGetFields(sub_operators[i], &num_input_fields, &op_input_fields, &num_output_fields, &op_output_fields));
CeedCallBackend(CeedOperatorGetQFunction(sub_operators[i], &qf));
CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_input_fields, NULL, NULL));

for (CeedInt j = 0; j < num_input_fields; j++) {
CeedEvalMode eval_mode;

CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[j], &eval_mode));
if (eval_mode == CEED_EVAL_WEIGHT) continue;
{
const CeedScalar *arr;
CeedVector vec;

CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[j], &vec));
if (vec != CEED_VECTOR_ACTIVE && vec != CEED_VECTOR_NONE) {
CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, &arr));
CeedCallBackend(CeedVectorRestoreArrayRead(vec, &arr));
}
CeedCallBackend(CeedVectorDestroy(&vec));
}
}

CeedCallBackend(CeedOperatorIsAtPoints(sub_operators[i], &is_at_points));
if (is_at_points) {
const CeedScalar *arr;
CeedVector vec;

CeedCallBackend(CeedOperatorAtPointsGetPoints(sub_operators[i], NULL, &vec));
CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, &arr));
CeedCallBackend(CeedVectorRestoreArrayRead(vec, &arr));
CeedCallBackend(CeedVectorDestroy(&vec));
}

CeedCallBackend(CeedQFunctionGetInnerContextData(qf, CEED_MEM_DEVICE, &d_c));
CeedCallBackend(CeedQFunctionRestoreInnerContextData(qf, &d_c));
CeedCallBackend(CeedQFunctionDestroy(&qf));
}
return CEED_ERROR_SUCCESS;
}

// Composite apply without CUDA graphs.
static int CeedOperatorApplyAddComposite_NoGraph_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;
Expand Down Expand Up @@ -330,7 +398,7 @@ static int CeedOperatorApplyAddComposite_Cuda_gen(CeedOperator op, CeedVector in
if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArray(output_vec, &output_arr));
CeedCallCuda(ceed, cudaDeviceSynchronize());

// Fallback on unsuccessful run
// Fall back to /gpu/cuda/ref for any sub-operator that couldn't run here.
Comment thread
jeremylt marked this conversation as resolved.
Outdated
for (CeedInt i = 0; i < num_suboperators; i++) {
if (!is_run_good[i]) {
CeedOperator op_fallback;
Expand All @@ -344,6 +412,124 @@ static int CeedOperatorApplyAddComposite_Cuda_gen(CeedOperator op, CeedVector in
return CEED_ERROR_SUCCESS;
}

static int CeedOperatorApplyAddComposite_Cuda_gen(CeedOperator op, CeedVector input_vec, CeedVector output_vec, CeedRequest *request) {
Ceed ceed;
CeedOperator_Cuda_gen *impl;
CeedOperator *sub_operators;
CeedInt num_suboperators;

ceed = CeedOperatorReturnCeed(op);
Comment thread
jeremylt marked this conversation as resolved.
Outdated
CeedCall(CeedOperatorCompositeGetNumSub(op, &num_suboperators));
CeedCall(CeedOperatorCompositeGetSubList(op, &sub_operators));
Comment thread
zatkins-dev marked this conversation as resolved.
Outdated
CeedCallBackend(CeedOperatorGetData(op, &impl));

if (!impl->use_graph || (input_vec == CEED_VECTOR_NONE && output_vec == CEED_VECTOR_NONE)) {
return CeedOperatorApplyAddComposite_NoGraph_Cuda_gen(op, input_vec, output_vec, request);
}

if (!impl->warmup_done) {
CeedCallBackend(CeedOperatorApplyAddComposite_NoGraph_Cuda_gen(op, input_vec, output_vec, request));
impl->warmup_done = true;
return CEED_ERROR_SUCCESS;
}

bool need_build = !impl->graph_created;

if (!need_build && input_vec != CEED_VECTOR_NONE) {
const CeedScalar *in_ptr;

CeedCallBackend(CeedVectorGetArrayRead(input_vec, CEED_MEM_DEVICE, &in_ptr));
need_build = in_ptr != impl->captured_input_ptr;
CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &in_ptr));
}
Comment thread
zatkins-dev marked this conversation as resolved.

if (need_build) {
const CeedScalar *input_arr = NULL;
CeedScalar *output_arr = NULL;
cudaStream_t capture_stream = cudaStreamPerThread;
cudaGraph_t graph = NULL;
bool capture_ok = true;
cudaError_t err;

if (impl->graph_instance) CeedCallCuda(ceed, cudaGraphExecDestroy(impl->graph_instance));
if (impl->graph) CeedCallCuda(ceed, cudaGraphDestroy(impl->graph));
impl->graph = NULL;
impl->graph_instance = NULL;

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));
impl->captured_input_ptr = input_arr;

err = cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeThreadLocal);
if (err != cudaSuccess) capture_ok = false;
if (capture_ok) {
// Still call EndCapture if capture is invalidated mid-way.
for (CeedInt i = 0; i < num_suboperators; i++) {
bool is_run_good = true;

if (CeedOperatorApplyAddCore_Cuda_gen(sub_operators[i], capture_stream, input_arr, output_arr, &is_run_good, request) || !is_run_good) {
capture_ok = false;
break;
}
}
}

err = cudaStreamEndCapture(capture_stream, &graph);
Comment thread
jeremylt marked this conversation as resolved.
if (capture_ok && (err != cudaSuccess || !graph)) capture_ok = false;
if (capture_ok) {
impl->graph = graph;
if (cudaGraphInstantiate(&impl->graph_instance, impl->graph, 0) != cudaSuccess) {
CeedCallCuda(ceed, cudaGraphDestroy(impl->graph));
impl->graph = NULL;
capture_ok = false;
}
} else if (graph) {
cudaGraphDestroy(graph);
Comment thread
zatkins-dev marked this conversation as resolved.
Outdated
}

if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &input_arr));
if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArray(output_vec, &output_arr));

if (!capture_ok) {
cudaGetLastError();
CeedCallCuda(ceed, cudaDeviceSynchronize());
cudaGetLastError();
impl->use_graph = false;
impl->graph_created = false;
impl->captured_input_ptr = NULL;
return CeedOperatorApplyAddComposite_NoGraph_Cuda_gen(op, input_vec, output_vec, request);
}
impl->graph_created = true;
}

if (input_vec != CEED_VECTOR_NONE) {
const CeedScalar *in_arr;

CeedCallBackend(CeedVectorGetArrayRead(input_vec, CEED_MEM_DEVICE, &in_arr));
CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &in_arr));
}
if (output_vec != CEED_VECTOR_NONE) {
CeedScalar *out_arr;

CeedCallBackend(CeedVectorGetArray(output_vec, CEED_MEM_DEVICE, &out_arr));
CeedCallBackend(CeedVectorRestoreArray(output_vec, &out_arr));
}
CeedCallBackend(CeedCompositeRefreshForReplay_Cuda_gen(sub_operators, num_suboperators));

if (cudaGraphLaunch(impl->graph_instance, NULL) != cudaSuccess) {
cudaGetLastError();
if (impl->graph_instance) cudaGraphExecDestroy(impl->graph_instance);
if (impl->graph) cudaGraphDestroy(impl->graph);
impl->graph = NULL;
impl->graph_instance = NULL;
impl->graph_created = false;
impl->captured_input_ptr = NULL;
impl->use_graph = false;
return CeedOperatorApplyAddComposite_NoGraph_Cuda_gen(op, input_vec, output_vec, request);
}
return CEED_ERROR_SUCCESS;
Comment thread
jeremylt marked this conversation as resolved.
}

//------------------------------------------------------------------------------
// QFunction assembly
//------------------------------------------------------------------------------
Expand Down Expand Up @@ -465,7 +651,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));
Expand Down Expand Up @@ -885,6 +1071,13 @@ int CeedOperatorCreate_Cuda_gen(CeedOperator op) {
CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
CeedCallBackend(CeedCalloc(1, &impl));
CeedCallBackend(CeedOperatorSetData(op, impl));

{
const char *env_val = getenv("CEED_DISABLE_GRAPH");

impl->use_graph = !(env_val && !strcmp(env_val, "1"));
}
Comment thread
zatkins-dev marked this conversation as resolved.
Outdated

CeedCall(CeedOperatorIsComposite(op, &is_composite));
if (is_composite) {
CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAddComposite", CeedOperatorApplyAddComposite_Cuda_gen));
Expand All @@ -897,6 +1090,7 @@ int CeedOperatorCreate_Cuda_gen(CeedOperator op) {
CeedOperatorLinearAssembleAddDiagonalAtPoints_Cuda_gen));
CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleSingle", CeedOperatorAssembleSingleAtPoints_Cuda_gen));
}

Comment thread
jeremylt marked this conversation as resolved.
Outdated
if (!is_at_points) {
CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleQFunction", CeedOperatorLinearAssembleQFunction_Cuda_gen));
CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleQFunctionUpdate",
Expand Down
8 changes: 8 additions & 0 deletions backends/cuda-gen/ceed-cuda-gen.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -25,6 +26,13 @@ typedef struct {
Fields_Cuda G;
CeedScalar *W;
Points_Cuda points;

bool use_graph;
Comment thread
jeremylt marked this conversation as resolved.
bool graph_created;
bool warmup_done;
cudaGraph_t graph;
cudaGraphExec_t graph_instance;
const CeedScalar *captured_input_ptr;
} CeedOperator_Cuda_gen;

typedef struct {
Expand Down
11 changes: 10 additions & 1 deletion backends/cuda-ref/ceed-cuda-ref-qfunctioncontext.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}
Comment thread
jeremylt marked this conversation as resolved.

CeedCallBackend(CeedDestroy(&ceed));
return CEED_ERROR_SUCCESS;
}
Expand Down
12 changes: 11 additions & 1 deletion backends/cuda-ref/ceed-cuda-ref-vector.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Comment thread
jeremylt marked this conversation as resolved.
Outdated
// During capture, use async memset with cudaStreamPerThread
CeedCallCuda(CeedVectorReturnCeed(vec), cudaMemsetAsync(impl->d_array, 0, length * sizeof(CeedScalar), cudaStreamPerThread));
Comment thread
jeremylt marked this conversation as resolved.
Outdated
} 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));
}
Expand Down