[Perf] CUDA graph 2: graph_do_while

[hughperkins]

Issue: #

Brief Summary

In this PR, we add graph_do_while for CUDA, and do NOT add fallbacks on other platforms. A later PR will add fallbacks on other platforms.

• in addition, must always use the same physical ndarray condition object when calling into a do_while graph graph

The do-while is implemented by usng a cuda graph conditional node

• we need to use PTX for this node, and compile that, using libcudadevrta.a

copilot:summary

Walkthrough

copilot:walkthrough

[hughperkins]

Opus 4.6 review:

What it does

Adds graph_while — a GPU-side iteration primitive built on the CUDA graph infrastructure from MVP-1. Annotating a kernel with @qd.kernel(graph_while="counter") repeats the kernel body while a device-side i32 flag is non-zero:

• SM 9.0+ (Hopper): Uses CUDA conditional while graph nodes. A JIT-compiled PTX condition kernel calls cudaGraphSetConditional to control the loop entirely on-GPU.
• Older CUDA / CPU / AMDGPU: Falls back to a host-side do-while loop with stream_synchronize + D2H memcpy to poll the flag each iteration.

The implementation threads through cleanly: Python decorator → graph_while_arg_id on LaunchContextBuilder → backend-specific loop logic.

What's good

1. Clean layering. The Python side (kernel_impl.py, kernel.py) is minimal — just plumbs the arg name through and resolves the C++ arg index. All the heavy lifting is in the C++ launchers.

2. Cross-platform fallback. CPU and AMDGPU get do-while fallbacks with identical semantics. cuda_graph=True is implicitly set, so the user writes one annotation and it works everywhere.

3. Solid documentation. The new cuda_graphs.md page explains both cuda_graph=True and graph_while, provides counter-based and boolean-flag patterns, and clearly warns about do-while semantics.

4. Good test coverage. Both CUDA-specific tests and cross-backend tests covering counter, boolean flag, multi-loop, replay, and single-iteration edge cases.

5. ABI guard. The static_assert(sizeof(CudaGraphNodeParams) == 256, ...) is a smart safety net for the hand-rolled struct.

Concerns

1. Graph replay bug with changed ndarray pointers (correctness)

When the CUDA graph is replayed (cache hit), the work kernels are fine because the arg buffer is re-uploaded. But the condition kernel's flag_ptr is baked into the graph at creation time:

// kernel_launcher.cpp, inside launch_llvm_kernel_graph
void *flag_ptr = ctx.graph_while_flag_dev_ptr;
void *cond_args[2] = {&cond_handle, &flag_ptr};
// ... added as a kernel node — pointer captured permanently

If a user creates a new ndarray for the counter between calls, the condition kernel still reads from the old device address. The tests pass because they reuse the same ndarray object — but that's a latent bug.

Fix: Store the condition kernel node in CachedCudaGraph and update its params on replay (via cuGraphExecKernelNodeSetParams or equivalent), or invalidate the cache when the flag pointer changes.

2. Fragile ABI-coupled struct (maintainability)

CudaGraphNodeParams is a hand-rolled mirror of CUgraphNodeParams with manual padding:

// kernel_launcher.h
struct CudaGraphNodeParams {
  unsigned int type;
  int reserved0[3];
  unsigned long long handle;
  unsigned int condType;
  unsigned int size;
  void *phGraph_out;
  void *ctx;
  char _pad[232 - 8 - 4 - 4 - 8 - 8];
  long long reserved2;
};

The static_assert on total size catches gross drift, but internal layout drift (e.g., NVIDIA changing field offsets within the union) would silently corrupt data and still pass the size check. A comment referencing the exact CUDA driver API version this was modeled against would help future maintainers know when to re-verify.

3. No validation on the graph_while parameter type (usability)

There's no check that the named parameter is actually a scalar i32 ndarray. Passing a f32 array, a 1D array, or a misspelled name would produce undefined behavior or a silent wrong result rather than a clear error.

Fix: Validate at kernel definition or first-call time that the named arg exists and has the expected dtype/shape.

4. Do-while footgun with counter=0 (usability)

The docs warn about it, which is good. But there's no runtime guard — counter=0 will infinite-loop on CPU/AMDGPU fallbacks (work, sync, read 0, decrement to −1, −1 ≠ 0, loop forever). A simple QD_ASSERT(initial_value >= 1) or at least a QD_WARN on the fallback paths would be cheap insurance.

5. Host-side fallback performance (docs)

The AMDGPU and CUDA non-graph fallbacks call stream_synchronize + D2H memcpy every iteration. For fine-grained loops this could be slower than just running the iterations without the graph_while machinery. Worth a note in the docs that the fallback is semantically correct but not a performance win.

6. Minor: stream_synchronize(nullptr) in AMDGPU (correctness)

The AMDGPU fallback passes nullptr for the stream, which may synchronize all streams rather than just the relevant one. The CUDA fallback correctly uses CUDAContext::get_instance().get_stream().

Summary

Well-structured feature with clean layering, good docs, and solid tests. #1 (replay pointer staleness) is a correctness bug that should be fixed before merge. #3 (input validation) is a close second for user-facing ergonomics. The rest are hardening items that could be tracked as follow-ups.

[hughperkins]

Update from Opus on concern 4:

Thinking this through more carefully: counter=0 causes an infinite loop on all paths, not just the fallbacks. On the CUDA graph path, defaultLaunchValue=1 forces the body to
execute once, the body decrements 0 to -1, the condition kernel sees -1 != 0, and it loops forever too.
The question is whether a runtime guard is practical. The counter lives on device — to check its initial value at launch time you'd need a D2H memcpy + sync, which adds
latency on every launch and undermines the whole point of keeping things on the GPU. On CPU you could peek cheaply, but adding a check only on one backend is inconsistent.
You could do a post-first-iteration sanity check (if the counter went negative, something is wrong), but that's complex, still requires a sync on the fallback paths, and
doesn't help the CUDA graph path where the loop is entirely on-device.
I think this is adequately handled by the documentation warning. The do-while semantics are inherent to CUDA conditional while nodes — it's not a quirk of this implementation,
it's how the hardware primitive works. Users of this API are writing GPU iteration loops and should be expected to read the docs. The name validation we just added catches
the most common real-world mistake (typos); counter=0 is a logic error that's hard to guard against cheaply.
I wouldn't add anything here.

[hughperkins]

Update from Opus on concern 5:

Concern 5 (host-side fallback performance) is just a documentation suggestion — noting that the fallback is correct but not a performance win. Looking at it again, I think the
existing docs already cover this adequately:

docs/source/user_guide/user_guide/cuda_graphs.md lines 68-70

• On SM 9.0+ (Hopper), this uses CUDA conditional while nodes — the entire iteration runs on the GPU with no host involvement.
• On older CUDA GPUs and non-CUDA backends, it falls back to a host-side do-while loop.
• graph_while implicitly enables cuda_graph=True.

The fact that it's a "host-side do-while loop" already implies sync-per-iteration. Anyone who understands enough to care about this overhead will read "host-side" and
understand the implication. Adding a paragraph spelling out "this means a stream synchronize and D2H memcpy every iteration" would be accurate but is probably over-explaining
for the target audience.
I wouldn't change anything.

[hughperkins]

update from Opus for concern 6:

Looking at the code, the existing AMDGPU backend already uses stream_synchronize(nullptr) everywhere — line 95 has the same pattern, predating graph_while. The new code at
line 129 is just following the established convention for this backend.
This isn't something the branch introduced; it's how the AMDGPU backend works. Changing it would be a broader AMDGPU backend concern, not something to fix in the graph_while
PR. I'd drop this from the review.

[hughperkins]

Addressed all 6 of Opus's concerns

[hughperkins]

this should be removed

[hughperkins]

raise error if not in cuda graph arlreayd

[hughperkins]

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.

[hughperkins]

env.sh shouldnt be here