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]

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.

[zatkins-dev]

Note to me and @jeremylt -- we should probably make these async/non-blocking during normal execution; it shouldn't affect determinism and will likely improve perf. If blocking is needed, that's likely a bug anyway.

[jeremylt]

I think honestly than would be cleaner and clearer

[jeremylt]

yea, the more I think on it, the better off I think we are just defaulting to async and skipping this special case logic

[zatkins-dev]

I think this looks pretty good. I'll hold off for @jeremylt on final approval, but this is a really cool feature add! Thanks @Nafees01!

[Nafees01]

I think this looks pretty good. I'll hold off for @jeremylt on final approval, but this is a really cool feature add! Thanks @Nafees01!

Thanks @zatkins-dev and thanks @jeremylt for all the detailed feedback - really helped get this in good shape. Appreciate you both taking the time!

[jeremylt]

call should be wrapped in error catching macro

[jeremylt]

Suggested change

	CUstream stream_to_use = (capture_status != cudaStreamCaptureStatusNone) ? cudaStreamPerThread : NULL;
	CUstream stream_to_use = (capture_status != cudaStreamCaptureStatusNone) ? cudaStreamPerThread : NULL;

Always blank line after variable declaration

[jeremylt]

variable name should use words

[jeremylt]

This function should not be used to get a persistent ceed

Suggested change

	ceed = CeedOperatorReturnCeed(op);
	CeedCallBackend(CeedOperatorGetCeed(op, &ceed));

[jeremylt]

And destroy that Ceed down here

[jeremylt]

really, before each of the return statements

[jeremylt]

You could, alternatively, use CeedOperatorReturnCeed() exactly in the sites where you use the ceed. Just don't use CeedOperatorReturnCeed() to get around having to destroy a persistent ceed variable

[jeremylt]

Suggested change

	CeedCallBackend(CeedOperatorIsComposite(op, &is_composite));
	CeedCallBackend(CeedOperatorIsComposite(op, &is_composite));

[jeremylt]

ok, either add blank lines before those if/else blocks or don't please - the mix is confusing

[jeremylt]

Suggested change

	bool use_graph;
	// Graph capture data
	bool use_graph;

Either a comment here or no blank line

[jeremylt]

Suggested change

	// 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)));
	}
	CeedCallCuda(CeedVectorReturnCeed(vec), cudaMemsetAsync(impl->d_array, 0, length * sizeof(CeedScalar), cudaStreamPerThread));