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

[Nafees01]

Summary

Adds CUDA Graph capture and replay to the 'cuda-gen' backend for composite operators. On the second apply, the kernel sequence is captured into a cudaGraph; all subsequent applies replay the instantiated graph. QFunction contexts are re-synced to device before each replay so time- and load-dependent parameters stay correct.

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply, which CUDA doesn't allow inside a capture. When this happens, the capture is properly closed to leave the stream in a clean state, any leftover CUDA error is cleared, and that operator falls back to running normally on every call.

Validated against both /gpu/cuda/ref and cuda-gen without graphs across static, quasistatic, contact, multi-material, and dynamic elasticity examples. Strain energy and displacements match to ~1e-13 in every case, with zero graph fallbacks on capturable operators.

Performance-wise, graphs are within ±1% of cuda-gen without graphs across all sizes, so no regression, but no real speedup either. The reason is that cuda-gen already fuses everything into a few large kernels, so there's very little launch overhead left for graphs to remove.

Validated correctness against /gpu/cuda/ref and cuda-gen (CEED_DISABLE_GRAPH=1) on the following Ratel examples:

• ex01-static-elasticity-linear-mms
• ex02-quasistatic-elasticity-linear-platen (contact)
• ex02-quasistatic-elasticity-neo-hookean-current-sinker (multi-material)
• ex03-dynamic-elasticity-mooney-rivlin-current
  Strain energy matches to ~1e-13 in all cases. Zero graph fallbacks on capturable operators.

[zatkins-dev]

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply

Are there specific operators you were unable to use in Ratel? None of the QFunctions there should alloc during operator application.

[zatkins-dev]

There's some style errors here, be sure to make format each commit

[zatkins-dev]

I'm sure you're still working on this, but here's an initial review to correct some high-level design issues and removed functionality that I think should be fixed.

[zatkins-dev]

This removed the is_sequential checks, which are there to ensure that all suboperators run in a single stream.

[Nafees01]

Good catch, restored. The is_sequential handling now lives in CeedOperatorApplyAddComposite_NoGraph_Cuda_gen, which is the shared non-graph path for composite operators.

[zatkins-dev]

This should be stored on a per-operator (or at least per-Ceed) basis, not as static variables. It's okay to use the environment variable to set a default value, but there should be a setter (e.g. CeedSetUseCudaGraph(Ceed ceed, bool use_graph))

[Nafees01]

Addressed. Removed the static disable_graph variable. Graph enable/disable is now stored per-operator as impl->use_graph in CeedOperator_Cuda_gen.

[zatkins-dev]

Should this be?

Suggested change

	if (input_vec == CEED_VECTOR_NONE || output_vec == CEED_VECTOR_NONE) {
	if (input_vec == CEED_VECTOR_NONE && output_vec == CEED_VECTOR_NONE) {

It's possible to have a passive input and active output

[Nafees01]

Yes, good catch; updated to &&. We only skip graph capture when both input_vec and output_vec are CEED_VECTOR_NONE.

[zatkins-dev]

This call pattern prevents each suboperator from running on separate streams. You won't see a perf difference until you have large, multioperator applications, but the perf difference can be significant.

[zatkins-dev]

The disabled graph case should restore the previous implementation, with multi-stream operators, sequential checks, etc

[zatkins-dev]

Given you do this a bunch of places, I'd recommend a helper function which looks like:

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;
  CeedScalar       *output_arr = NULL;
  Ceed              ceed;
  CeedOperator     *sub_operators;
  cudaStream_t      stream = NULL;

  CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
  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;

    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));
    }
  }
  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;

      CeedDebug(ceed, "\nFalling back to /gpu/cuda/ref CeedOperator for ApplyAdd\n");
      CeedCallBackend(CeedOperatorGetFallback(sub_operators[i], &op_fallback));
      CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
    }
  }
  CeedCallBackend(CeedDestroy(&ceed));
  return CEED_ERROR_SUCCESS;
}

[Nafees01]

The disabled graph case should restore the previous implementation, with multi-stream operators, sequential checks, etc

Restored. The graph-disabled path no longer loops over suboperators with the public CeedOperatorApplyAdd interface. It now goes through CeedOperatorApplyAddComposite_NoGraph_Cuda_gen, which brings back the original multi-stream behavior: is_sequential composites run on one shared stream, non-sequential composites get per-suboperator streams for overlap, and we call CeedOperatorApplyAddCore_Cuda_gen directly with per-suboperator /ref fallback only as a last resort. All non-graph routes (disable graph, passive-only composite, warmup, capture/launch failure) use this helper.

[Nafees01]

Given you do this a bunch of places, I'd recommend a helper function which looks like:

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;
  CeedScalar       *output_arr = NULL;
  Ceed              ceed;
  CeedOperator     *sub_operators;
  cudaStream_t      stream = NULL;

  CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
  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;

    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));
    }
  }
  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;

      CeedDebug(ceed, "\nFalling back to /gpu/cuda/ref CeedOperator for ApplyAdd\n");
      CeedCallBackend(CeedOperatorGetFallback(sub_operators[i], &op_fallback));
      CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
    }
  }
  CeedCallBackend(CeedDestroy(&ceed));
  return CEED_ERROR_SUCCESS;
}

Done; added CeedOperatorApplyAddComposite_NoGraph_Cuda_gen as suggested. It centralizes the non-graph composite path (is_sequential stream handling, core apply, sync, per-suboperator ref fallback). All routes that shouldn't use a graph go through it: graph disabled, passive-only composite, warmup, and capture/launch failure fallbacks.

[zatkins-dev]

Need to use the same helper here

[Nafees01]

Updated. This path now calls CeedOperatorApplyAddComposite_NoGraph_Cuda_gen instead of looping over suboperators with the public ApplyAdd interface. Same helper is used for the other non-graph routes.

[zatkins-dev]

Same here -- use the helper, not fall back.

[Nafees01]

Fixed. Routed through the NoGraph helper instead of whole-operator fallback.

[zatkins-dev]

Use helper here

[Nafees01]

Fixed. Capture failure and other non-graph paths now use the NoGraph helper instead of the old Phase 3 inline direct-apply.

[zatkins-dev]

This check seems like it should happen before the graph is built, that way you can rebuild the graph if needed instead of waiting until the next apply call

[Nafees01]

Done. Moved the input pointer check before graph build. If the device address changed, we set need_build and recapture immediately in the same apply instead of falling back and waiting until the next call.

[zatkins-dev]

This same issue is going to happen for any passive inputs which have been updated on the CPU -- there should be a more robust way to handle this

[Nafees01]

Good point; extended pre-replay sync beyond QFunction contexts. CeedCompositeRefreshForReplay_Cuda_gen now touches each sub-operator's passive input vectors (and at-points coordinates) with GetArrayRead(CEED_MEM_DEVICE) / Restore, along with the existing context refresh.

[zatkins-dev]

This isn't the right way to benchmark against ref -- use the /gpu/cuda/ref backend directly.

Suggested change

	// 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;
	}

[Nafees01]

Done. Removed CEED_FORCE_BASELINE. Ref benchmarks use -ceed /gpu/cuda/ref directly.

[zatkins-dev]

I would be interested to see if this speeds up e.g. /gpu/cuda/ref, when recording the graph over the entire CeedOperatorApplyCore_Cuda_Ref function

[jeremylt]

Note - merging main into your branch instead of rebasing means I'd recommend we squash-merge this PR

[Nafees01]

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply

Are there specific operators you were unable to use in Ratel? None of the QFunctions there should alloc during operator application.

The only example where the contact (Nitsche) sub-operator couldn't be captured was the quasistatic contact case (ex02-quasistatic-elasticity-linear-platen). Everything else I tried static, dynamic, multi-material, neo-Hookean, Mooney-Rivlin captured and replayed fine.

And you're right that the QFunction doesn't allocate; the cudaMalloc comes from libCEED's own device-side handling of the contact data, not the QFunction itself. It runs either way correctly; the contact operators just fall back to a normal apply rather than being replayed.

[zatkins-dev]

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply

Are there specific operators you were unable to use in Ratel? None of the QFunctions there should alloc during operator application.

The only example where the contact (Nitsche) sub-operator couldn't be captured was the quasistatic contact case (ex02-quasistatic-elasticity-linear-platen). Everything else I tried static, dynamic, multi-material, neo-Hookean, Mooney-Rivlin captured and replayed fine.

And you're right that the QFunction doesn't allocate; the cudaMalloc comes from libCEED's own device-side handling of the contact data, not the QFunction itself. It runs either way correctly; the contact operators just fall back to a normal apply rather than being replayed.

This is a pretty big red flag to me. The Nitsche contact operators don't actually do anything out of the ordinary on the libCEED side. Them failing to work means that something about this setup (very possibly something with passive inputs and outputs) is violating the libCEED interface assumptions.

[Nafees01]

Sub-operators that cannot be captured, for example, contact operators that allocate memory (cudaMalloc) during apply

Are there specific operators you were unable to use in Ratel? None of the QFunctions there should alloc during operator application.

The only example where the contact (Nitsche) sub-operator couldn't be captured was the quasistatic contact case (ex02-quasistatic-elasticity-linear-platen). Everything else I tried static, dynamic, multi-material, neo-Hookean, Mooney-Rivlin captured and replayed fine.
And you're right that the QFunction doesn't allocate; the cudaMalloc comes from libCEED's own device-side handling of the contact data, not the QFunction itself. It runs either way correctly; the contact operators just fall back to a normal apply rather than being replayed.

This is a pretty big red flag to me. The Nitsche contact operators don't actually do anything out of the ordinary on the libCEED side. Them failing to work means that something about this setup (very possibly something with passive inputs and outputs) is violating the libCEED interface assumptions.

Thanks for flagging this. On re-checking linear-platen with CEED_DEBUG=1, the contact operators appear to fail earlier than graph capture: cuda-gen hits an NVRTC compile error (GradTransposeTensor3dFlattened defined twice in cuda-shared-basis-tensor-flattened-templates.h) and then falls back to /gpu/cuda/ref. So this looks like a pre-existing cuda-gen JIT issue triggered by the contact operator's basis configuration, rather than something specific to Nitsche or CUDA graphs.

[zatkins-dev]

On re-checking linear-platen with CEED_DEBUG=1, the contact operators appear to fail earlier than graph capture: cuda-gen hits an NVRTC compile error (GradTransposeTensor3dFlattened defined twice in cuda-shared-basis-tensor-flattened-templates.h) and then falls back to /gpu/cuda/ref

Sweet, good catch. I'll fix that separately.

[jeremylt]

Suggested change

	// Fall back to /gpu/cuda/ref for any sub-operator that couldn't run here.
	// Fall back to /gpu/cuda/ref for any sub-operator that couldn't run here

[Nafees01]

Fixed.

[jeremylt]

same here

[Nafees01]

Fixed, wrapped in CeedCallCuda here as well.

[jeremylt]

Generally not great if the only way to control something is via an env variable. I'd also add a way to set this per operator from user facing code in ceed-operator.c with CeedOperatorSetCudaEnableCudaGraph(op, enable_graph) or similar.

Q: Is this graph also valid to use with ref composite operators, or no?

Either way, I'd call CeedOperatorSetCudaEnableCudaGraph here in this block instead of directly modifying impl->use_graph here, and I'd also call it when capture fails.

Also, maybe it's a me thing, but I'd go with CEED_ENABLE_CUDA_GRAPH 0 and CEED_ENABLE_CUDA_GRAPH 1 and assume CEED_ENABLE_CUDA_GRAPH 1 is the default?

[Nafees01]

Programmatic control: Added CeedOperatorSetEnableCudaGraph(op, enable_graph) in include/ceed/cuda.h / interface/ceed-cuda.c (same pattern as CeedQFunctionSetCUDAUserFunction). The cuda-gen backend registers a SetEnableCudaGraph hook on operator creation.

Env var: Switched from CEED_DISABLE_GRAPH to CEED_ENABLE_CUDA_GRAPH. Default is on (unset or any value other than 0/false). CEED_ENABLE_CUDA_GRAPH=0 disables graphs.

Create / failure paths: Operator create reads the env var and calls CeedOperatorSetEnableCudaGraph instead of setting impl->use_graph directly. Capture and launch failure also call CeedOperatorSetEnableCudaGraph(op, false) rather than touching the impl field.

Ref composite operators: Not in this PR. Graph capture/replay is only implemented for /gpu/cuda/gen composite ApplyAddComposite. /gpu/cuda/ref still uses the standard multi-kernel apply path with no graph support; adding graphs there would be a separate follow-up

Verification: On noether, multi-density-force with 50 TS steps (-ts_time_step 0.02): gen+graph 4.26 s, gen with CEED_ENABLE_CUDA_GRAPH=0 6.02 s, ref 6.15 s. Strain energy and displacement match across all three. CUDA graphs are ~29% faster than no-graph gen (1.41×) and ~31% faster than ref (1.44×).

[jeremylt]

Woof, that's a lengthy comment that looks LLM generated. If it was, don't forget to disclose LLM usage!

I say woof because.... that's a lot to parse all at once, which makes it less good at communication. If you are and want to use LLMs to write comments, I would definitely try to make a human summary that is easier to understand to go with it. If its human written - same thing really, I'd make a small summary to preceed such a lengthy reply

[Nafees01]

Fair point, that was too much to read at once. I did try to keep it concise, but I kept adding summary bits and asked the LLM to clean up the phrasing afterward. From next time, I will use a short summary so that it's easy to understand. My bad on the length.

Here is the short summary: added a real setter CeedOperatorSetEnableCudaGraph instead of hardcoding the env var check, so users/backends can turn graphs on/off directly. Renamed the env var to CEED_ENABLE_CUDA_GRAPH (on by default, 0 turns it off). Capture/launch failures now go through that same setter. Ref backend still isn't touched; that's a separate follow-up. Tested on noether with 50 steps of multi-density-force, graph is ~30% faster than both no-graph and ref, and results match across all three.

[jeremylt]

sounds good, and much easier for me to understand than what the LLM did with your words ❤️

[jeremylt]

We should probably note that this function (and the one above) are silently successful if the backend doesn't support this feature

[Nafees01]

Makes sense. Planning to add something like: if the backend does not support CUfunction pointers for QFunctions, the call succeeds without effect, and if the backend does not support CUDA Graphs for operators, the call succeeds without effect. Does that sound good?

[jeremylt]

Yea, that works for me. Maybe also something like 'When this function has no effect, it is logged in debugging output."