[Perf] Add CUDA Graph support

[hughperkins]

Issue: #

Brief Summary

Example usage:

    @qd.kernel(graph_while='counter')
    def inc(x: qd.types.ndarray(qd.i32, ndim=1),
            counter: qd.types.ndarray(qd.i32, ndim=0)):
        for i in range(x.shape[0]):
            x[i] = x[i] + 1
        for i in range(1):
            counter[None] = counter[None] - 1

    @qd.kernel(graph_while='keep_going')
    def increment_until_all_done(
            x: qd.types.ndarray(qd.i32, ndim=1),
            thresholds: qd.types.ndarray(qd.i32, ndim=1),
            keep_going: qd.types.ndarray(qd.i32, ndim=0)):
        # Work: increment elements that haven't reached their threshold
        for i in range(x.shape[0]):
            if x[i] < thresholds[i]:
                x[i] = x[i] + 1

        # Reduction: reset flag, then OR-reduce per-element conditions
        for i in range(1):
            keep_going[None] = 0
        for i in range(x.shape[0]):
            if x[i] < thresholds[i]:
                keep_going[None] = 1

copilot:summary

Walkthrough

copilot:walkthrough

[hughperkins]

opus-4.6 review:

1. Per-kernel CUDA graph capture/replay (cuda_graph=True)

The Python API is @qd.kernel(cuda_graph=True). When a kernel has 2+ top-level for loops (offloaded tasks), it captures them into a CUDA graph on first launch and replays
on subsequent launches, eliminating per-kernel launch overhead.
Key design decisions:
• Uses the explicit graph node API (cuGraphAddKernelNode) rather than stream capture -- good choice, gives more control and avoids the pitfalls of stream capture.
• Persistent device buffers for arg/result survive across replays. On each replay, the arg buffer is re-uploaded so changed ndarray pointers are picked up correctly.
• Falls back silently to the normal launch path for single-loop kernels or non-CUDA backends (harmless no-op).

2. GPU-side conditional while loops (graph_while='flag_arg')

The API is @qd.kernel(graph_while='counter') where 'counter' names a scalar i32 ndarray argument. This wraps the kernel tasks in a CUDA conditional while node (SM 9.0+ /
Hopper required), enabling GPU-side iteration without CPU round-trips -- very relevant for convergence loops in physics solvers.
A small PTX condition kernel is JIT-linked with libcudadevrt.a at runtime to call cudaGraphSetConditional. Cross-platform fallback (CPU, AMDGPU, CUDA non-graph) uses a
host-side do-while loop with device readback.

Things that look good

• Clean API layering: flags flow from Python decorator through LaunchContextBuilder to the C++ launcher without polluting intermediate layers.
• Cross-platform parity: all three backends (CPU, CUDA, AMDGPU) get the graph_while do-while fallback, so semantics are identical.
• Arg-change handling: re-resolving ndarray device pointers and re-uploading the arg buffer on replay is the right approach.
• Test coverage: comprehensive -- two-loop, three-loop, single-loop fallback, no-annotation, changed args, graph_while with counter/boolean/multi-loop/replay, and
cross-backend variants.

Concerns / things I'd flag

1. CudaGraphNodeParams struct layout is fragile -- it manually mirrors the CUDA driver API struct with hardcoded offsets and padding. If NVIDIA changes the layout
   between driver versions, this breaks silently. Consider using the actual CUDA headers or at least a static_assert on sizeof.
2. No compute capability check -- the PTX targets sm_90 and conditional graph nodes require Hopper. If you run graph_while on an Ampere or older GPU, it will presumably
   fail at JIT link time with an unhelpful error. A runtime check with a clear error message would be better.
3. libcudadevrt.a discovery is brittle -- only two hardcoded paths are checked. Custom toolkit installs (e.g. via CUDA_HOME or conda) won't be found. Should check
   environment variables like CUDA_HOME/CUDA_PATH.
4. No graph cache eviction -- cuda_graph_cache_ grows monotonically. Each entry holds persistent device arg/result buffers. In a long-running process with many distinct
   kernels, this could leak significant GPU memory.
5. Removed documentation comments -- the diff strips several explanatory comments from the existing CUDA kernel launcher (e.g., the purpose of transfers and
   device_ptrs). In a diff that adds substantial complexity, removing existing documentation is undesirable.
6. Do-while semantics -- the body always executes at least once. This is documented in tests but could surprise users who set counter=0. Worth documenting explicitly in
   the API or raising an error for counter <= 0.
7. Thread safety -- cuda_graph_cache_, cond_kernel_func_, and cond_kernel_module_ have no synchronization. Fine if Quadrants is single-threaded, but worth a comment
   noting this assumption.
8. No test for autodiff interaction -- adjoint kernels also get use_cuda_graph and graph_while_arg set. There are no tests verifying this works (or is properly ignored)
   for reverse-mode AD kernels.

────────────────────────────────────────

Overall, it's a solid MVP. The architecture is sound, the explicit graph node API is the right choice, and the graph_while feature is a compelling addition for iterative
GPU workloads. The main risks are around portability (hardcoded paths, struct layouts, no SM version check) and missing edge-case handling.

[hughperkins]

additional information from opus on point 4:

thoughts on 'No graph cache eviction'?

Honestly, on closer look I think this is a non-issue and I was wrong to flag it.
The cache is keyed by launch_id, which corresponds to a distinct compiled kernel registration — not to individual calls. So the cache has at most one entry per
@qd.kernel(cuda_graph=True) definition. It doesn't grow with the number of invocations.
In practice:
• Users explicitly opt in per-kernel, so the number of entries is small (realistically 1-10 in a typical app)
• Each entry holds two small device buffers (arg buffer is typically a few hundred bytes, result buffer is ~8 bytes minimum) plus one graph exec object — negligible GPU
memory

The only way this could blow up is if someone programmatically generates many distinct kernels with cuda_graph=True in a loop, which would be a misuse with much bigger
problems (JIT compilation cost per kernel).
I'd leave it as-is. Adding LRU eviction would be complexity for no real benefit.

[hughperkins]

another opus update:

thoughts on 'Thread safety -- cuda_graph_cache_, cond_kernel_func_, and cond_kernel_module_ have no synchronization.'?

Not a real concern. Quadrants kernel launches go through a single KernelLauncher instance per device, and the launch path is inherently sequential -- you can't
meaningfully launch two CUDA kernels on the same stream from different threads simultaneously. The CUDA driver context itself serializes stream operations.
If Quadrants ever adds multi-threaded launch support, it would require much broader synchronization changes (the existing contexts_ vector has the same non-thread-safe
pattern), so cuda_graph_cache_ wouldn't be the first thing to break. Not worth adding a mutex for a scenario that doesn't exist.
I'd leave it as-is.

[hughperkins]

repalced by smaller prs