[Perf] CUDA graph 4: call from multiple locations

[hughperkins]

Issue: #

Brief Summary

When multiple locations call the same graph, very likely the counter ndarray will be different physically between each calling site.

Prior to this PR, this causes an exception.

In this PR, we handle making it possible to pass in a different ndarray object, without either triggering a recompile, throwing an exception, or being incorrect (which were the three options before this fix).

copilot:summary

Walkthrough

copilot:walkthrough

[hughperkins]

reviewer hat: note this constructor was created as a result of my earlier comment to do all alocation in the constructor, in one place. After squashed commit, we lose this history unforutnatley. But the previous commit history was a mess, s owe needed to squash.

[hughperkins]

Opus 4.6 review:

What it does

The branch solves a real usability problem: previously, passing a different counter ndarray to a graph_do_while kernel on a subsequent call would raise a RuntimeError because
the counter's device pointer was baked directly into the captured CUDA graph. Users had to reuse the exact same ndarray object every time.
The fix introduces a device-side indirection slot — a persistent void* allocation whose address is baked into the graph, but whose contents (the actual counter pointer) can be
updated cheaply via memcpy_host_to_device before each launch. The PTX condition kernel now does one extra pointer dereference to read through the slot.

What I like

• Clean design. The indirection slot is the textbook solution here — avoids graph rebuilds while keeping the graph structure immutable.
• Constructor consolidation. Moving the scattered allocation logic from try_launch into a proper CachedCudaGraph constructor is a nice cleanup. Fewer places to forget an init
step.
• Move-assignment swap idiom (commit 2). The old manual approach was a leak waiting to happen — and in fact was leaking counter_ptr_slot. The swap-through-RAII-guard pattern
handles cleanup uniformly via the destructor. Good catch.
• total_builds_ counter. Exposing this all the way to Python lets the tests assert "the graph was reused, not rebuilt," which is exactly the right thing to verify.
• Test coverage. Replacing the "raises on change" test with two positive tests (one-time swap + alternating back-and-forth across 6 calls) covers the important cases well.

Minor observations

1. Async memcpy TODO — the two memcpy_host_to_device calls before launch (counter_ptr_slot + arg_buffer) are synchronous. The TODO to make them async is valid, but note that
   flag_ptr is a stack local, so going async would require pinned memory or keeping it alive until stream sync. Worth a comment when that's tackled.
2. Test robustness — the tests assume c1 and c2 get different device pointers (otherwise the indirection isn't really being exercised). That's almost certainly true in
   practice, but an explicit assertion like assert c1.data_ptr() != c2.data_ptr() would make the test self-documenting about why two ndarrays are needed.
3. PTX is hand-maintained — the inline PTX with the extra ld.global.u64 for the indirection looks correct (%rd<5> register count updated, load chain is slot -> pointer ->
   value). Just noting that hand-edited PTX is inherently fragile; a comment about re-generating from CUDA C (as the existing header comment describes) would help future
   maintainers know the source of truth.

Overall, this is solid work — the right abstraction, proper resource management, and good test coverage.

[hughperkins]

opus minor obsevation 2 addressed in f755aa2
opus minor observation 1 aligns with not attempting async for now, sounds complciated

[hughperkins]

Asked opus to explain comment 3 more:

The existing header comment on the PTX already says how it was originally produced:

quadrants/runtime/cuda/cuda_graph_manager.cpp lines 83-83

// Compiled from CUDA C with: nvcc -ptx -arch=sm_90 -rdc=true

The idea is: instead of hand-editing the PTX directly (bumping register counts, adding load instructions), you'd maintain a CUDA C source file as the source of truth, compile
it to PTX, and paste the output into the string literal.
Here's what that would look like concretely. The current PTX corresponds to something like this CUDA C kernel (pre-indirection version):

#include <cuda.h>
#include <cudaGraphSetConditional.h>
extern "C" global void _qd_graph_do_while_cond(
cudaGraphConditionalHandle handle,
int32_t *flag) {
unsigned int val = (flag != 0) ? 1u : 0u;
// actually: val = (*flag != 0) ? 1 : 0
cudaGraphSetConditional(handle, *flag != 0 ? 1 : 0);
}

With the indirection change, it becomes:

extern "C" global void _qd_graph_do_while_cond(
cudaGraphConditionalHandle handle,
int32_t **flag_slot) {
int32_t *flag = *flag_slot; // dereference the slot
unsigned int val = (*flag != 0) ? 1u : 0u;
cudaGraphSetConditional(handle, val);
}

The workflow would be:

1. Keep a .cu file (e.g. quadrants/runtime/cuda/condition_kernel.cu) with the above source
2. Compile to PTX: nvcc -ptx -arch=sm_90 -rdc=true condition_kernel.cu -o condition_kernel.ptx
3. Copy the .entry block from the generated PTX into the kConditionKernelPTX string literal

That's basically it. The reason it's not already set up this way is that it's a tiny kernel unlikely to change often, and adding a build step for one PTX snippet is overhead.
The tradeoff is:
• Current approach (inline PTX, hand-edited): Zero build complexity, but you need to understand PTX register conventions, address space conversions, and the
cudaGraphSetConditional device-side API at the PTX level. The %rd<5> register declaration, cvta.to.global, ld.global.u64 vs ld.global.u32 distinctions are easy to get
wrong.
• Keeping a .cu source file: nvcc handles register allocation, address space conversion, and instruction selection. You edit C code, re-run the compiler, paste the result.
But it requires having nvcc with SM 9.0 support available, and you'd want to version-control the .cu alongside the .cpp.

A lightweight middle ground — which is what my observation was really suggesting — is just to keep the CUDA C source in a comment next to the PTX string, so the next person who
needs to modify it can regenerate rather than hand-edit. The existing comment mentions the nvcc command but doesn't include the actual C source it was compiled from.

[hughperkins]

opus observation 3 addressed in 06f1d9f.

[hughperkins]

With the exception of the generated PTX sourcecode, I have read every line added in this PR, and reviewed the lines. I take responsibilty for the lines added and removed in this PR, and won't blame any issues on Opus.