diff --git a/backends/cuda-gen/ceed-cuda-gen-operator.c b/backends/cuda-gen/ceed-cuda-gen-operator.c index f8228fbfbd..6e75d7a36f 100644 --- a/backends/cuda-gen/ceed-cuda-gen-operator.c +++ b/backends/cuda-gen/ceed-cuda-gen-operator.c @@ -11,6 +11,8 @@ #include #include #include +#include +#include #include #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", diff --git a/backends/cuda-gen/ceed-cuda-gen.h b/backends/cuda-gen/ceed-cuda-gen.h index 0e04f3c4e4..b9053cbfd2 100644 --- a/backends/cuda-gen/ceed-cuda-gen.h +++ b/backends/cuda-gen/ceed-cuda-gen.h @@ -10,6 +10,7 @@ #include #include #include +#include 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 { diff --git a/backends/cuda-ref/ceed-cuda-ref-qfunctioncontext.c b/backends/cuda-ref/ceed-cuda-ref-qfunctioncontext.c index 491e658338..cadb1025bb 100644 --- a/backends/cuda-ref/ceed-cuda-ref-qfunctioncontext.c +++ b/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; } diff --git a/backends/cuda-ref/ceed-cuda-ref-vector.c b/backends/cuda-ref/ceed-cuda-ref-vector.c index 980d2f0583..ddbc76b81a 100644 --- a/backends/cuda-ref/ceed-cuda-ref-vector.c +++ b/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)); }