[Perf] Per-launch host sync: drop wait_idle on SPIR-V, pin stream and drop stream_synchronize on CUDA/AMDGPU

[duburcqa]

Launcher redundant-sync drop: fold per-launch stream barriers into existing sync semantics on the default-stream fast path

Two commits, no behavior change on the explicit-stream path. On the default-stream fast path (entry stream_ == nullptr, every task on stream_parallel_group_id == 0) the launcher chains HtoD / kernel / DtoH on a single stream and the explicit stream_synchronize barriers between phases collapse onto sync DtoH semantics. Outside the fast path - user-supplied stream OR any task on stream_parallel_group_id != 0 - main's barrier set stays untouched. Measured on the default path: Metal rigid-step bench 200 -> 220 FPS (~ +10%), CUDA same bench 531 -> 576 FPS (~ +8.5%). AMDGPU same code shape, unverified.

Why

Each backend's per-launch path issued a synchronous stream drain (device_->wait_idle() on SPIR-V, stream_synchronize(stream) on CUDA, AMDGPU) immediately before a sequence of synchronous host-blocking transfers (Device::readback_data, cuMemcpyDtoH, hipMemcpyDtoH) that themselves drain the stream they target. The explicit drain was paying a full round-trip per launch on top of the round-trip the readback already pays. With one launch per substep across thousands of substeps in differentiable-sim workloads, that compound stall dominates host-side launch overhead.

Mechanism

SPIR-V (runtime/gfx/runtime.cpp)

Drop device_->wait_idle() between submit(cmdlist) and readback_data(...) in HostDeviceContextBlitter::device_to_host. readback_data allocates a staging buffer, records a buffer_copy cmdlist, and calls submit_synced internally; on Metal submit_synced -> command_sync() drains all pending cmdbufs, on Vulkan submit_synced -> vkQueueWaitIdle does the same. The kernel cmdlist completes as part of that wait without a separate wait_idle(). No interaction with the streams plumbing - the SPIR-V path predates and bypasses CUDAContext::stream_ / AMDGPUContext::stream_.

LLVM-GPU (runtime/cuda/kernel_launcher.cpp, runtime/amdgpu/kernel_launcher.cpp, rhi/cuda/cuda_stream_pin.h, rhi/amdgpu/amdgpu_stream_pin.h)

Both LLVM-GPU launchers now compute a default_stream_path flag at entry from (active_stream == nullptr) && all_sgid_zero, where all_sgid_zero walks offloaded_tasks for stream_parallel_group_id == 0. On the fast path:

• Engage a CudaDefaultStreamPinGuard / AmdgpuDefaultStreamPinGuard (RAII swap-and-restore on CUDAContext::stream_ / AMDGPUContext::stream_) so the kernel dispatch through cuda_module->launch / amdgpu_module->launch lands on the null stream regardless of any inner helper that swaps the context stream. Both guards are no-ops when entry stream is already nullptr but they cover transient swaps inside dispatch_max_reducers_for_tasks, publish_adstack_metadata, publish_per_task_bound_count_device.
• Drop the post-HtoD stream_synchronize(active_stream) (HtoD on null + kernel on null serialises by null-stream ordering).
• Drop the pre-DtoH-loop stream_synchronize(active_stream) (kernel on null + DtoH on null).
• Convert the async memcpy_device_to_host_async transfer loop to sync memcpy_device_to_host (host-blocking, drains the null stream); drop the post-loop stream_synchronize.
• Drop the result-buffer-only stream_synchronize (async DtoH to pageable host memory on null stream is implicit-sync per CUDA / HIP spec).

Outside the fast path - user-supplied active_stream != nullptr, OR any task on stream_parallel_group_id != 0 (per-group acquired streams differ from active_stream, cross-stream HtoD -> kernel and kernel -> DtoH need explicit barriers) - the pin guard stays disengaged and main's barrier set runs unchanged.

The CUDA pin guard closes the cross-stream-visibility break that cuMemAllocAsync-pool buffers show on pre-Ampere hardware (Turing T4 faults at cuLaunchKernel with illegal-address when consumer + producer streams differ) on the default path; this is the original motivation. The AMDGPU guard is symmetric defensive coverage - the equivalent fault has not been observed on AMDGPU hardware (CI + amdgpu box clean on the default path), but AMDGPUContext::launch now forwards stream_ to hipModuleLaunchKernel (post-streams plumbing in #653) so the same same-stream-invariant rationale applies and keeping the two backends symmetric simplifies future hardware shifts.

The guards' destructors do NOT call stream_synchronize(nullptr). The trailing sync would pay a forced host-GPU round-trip on every launcher exit; the launcher's sync DtoH calls inside the scope already drain the null stream where host visibility is needed, and the kernel-only fast path (no host-side sync DtoH) deliberately stays asynchronous.

Note: header layout - rhi/cuda/cuda_stream_pin.h (moved from runtime/llvm/adstack_lazy_claim/cuda_stream_pin.h) plus the new rhi/amdgpu/amdgpu_stream_pin.h. Both are generic GPU-context utilities, not adstack-specific; the existing dispatch_max_reducers_for_tasks / publish_adstack_metadata / publish_per_task_bound_count_device callsites of the CUDA guard are unchanged.

Per-backend matrix

Backend	Default-stream change	Explicit-stream / sgid != 0 change	Verified
Metal (SPIR-V via MoltenVK)	drop wait_idle() in blitter	n/a (no stream plumbing)	yes - repro_metal_rigid_step.py 200 -> 220 FPS (~+10%)
Vulkan (SPIR-V)	same path as Metal	n/a (no stream plumbing)	not benched, same code path
CUDA	pin stream + drop post-HtoD / pre-DtoH-loop syncs + convert DtoH transfer loop to sync DtoH + drop result-buffer-only sync	unchanged from main	yes - 531 -> 576 FPS (~+8.5%) on default path
AMDGPU	symmetric to CUDA, with AmdgpuDefaultStreamPinGuard	unchanged from main	builds + test_adstack.py 1143 passed, 11 xfailed
CPU	n/a	n/a	n/a

Side-effect audit

Concern	Verdict
SPIR-V kernel completes before readback_data writes its host buffer	yes - submit_synced inside readback_data drains the entire compute stream including the kernel cmdbuf
CUDA pre-Ampere cuMemAllocAsync pool buffers visible across HtoD / kernel / DtoH on the default path	yes - the launcher-scope pin guard ensures all three land on the null stream, matching LlvmRuntimeExecutor::allocate_memory_on_device's pool-allocation stream
Streams PR (#653) explicit-stream behavior	unchanged - default_stream_path == false keeps every barrier from main
Stream-parallel-group launches (stream_parallel_group_id != 0)	unchanged - all_sgid_zero == false disengages the pin guard and keeps every barrier from main; the per-group stream_synchronize(stream_by_id[...]) inside launch_offloaded_tasks continues to handle cross-group ordering
CUDA Graphs path	unaffected - the pin guard is instantiated AFTER the ctx.use_graph early-return and graph_manager_.try_launch
Existing adstack helper callsites of the CUDA pin guard	unchanged - they all issue a sync DtoH inside the scope which drains the null stream before the destructor runs; dropping the destructor's redundant sync is a no-op for them
AMDGPU kernel sees HtoD'd data on the default path	yes - sync hipMemcpyHtoD is fully synchronous w.r.t. host; the subsequent kernel runs on the null stream the HtoD targeted

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 72af99d253

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[duburcqa]

No impact on Genesis perf benchmarks:

env	batch_size	backend	gjk_collision	constraint_solver	runtime_fps_635	runtime_fps_654	runtime_fps_delta_pct
anymal_random	30000	cuda	-	-	9254908	9329350	+0.80
anymal_uniform	30000	cuda	-	-	12281694	12324555	+0.35
anymal_uniform_kinematic	0	cpu	-	-	1927	1978	+2.65
anymal_uniform_kinematic	30000	cuda	-	-	10390082	10457477	+0.65
anymal_zero	0	cpu	-	-	7209	7203	-0.08
anymal_zero	30000	cuda	-	-	19327082	18996395	-1.71
box_pyramid_3	4096	cuda	-	-	998985	989306	-0.97
box_pyramid_4	4096	cuda	-	-	390705	389188	-0.39
box_pyramid_5	4096	cuda	-	-	138762	138571	-0.14
box_pyramid_6	4096	cuda	False	-	58264	58856	+1.02
box_pyramid_6	4096	cuda	True	-	60600	60382	-0.36
dex_hand	4096	cuda	-	-	17567	17350	-1.24
duck_in_box_easy	30000	cuda	False	-	26829174	26982636	+0.57
duck_in_box_easy	30000	cuda	True	-	9644717	9561036	-0.87
duck_in_box_hard	0	cpu	-	-	5117	5196	+1.54
duck_in_box_hard	30000	cuda	False	-	10082166	10287954	+2.04
duck_in_box_hard	30000	cuda	True	-	3424350	3404519	-0.58
franka	30000	cuda	-	-	21538775	21520877	-0.08
franka_accessors	0	cpu	-	-	1144	1137	-0.61
franka_accessors	30000	cuda	-	-	15575265	15391071	-1.18
franka_free	30000	cuda	-	-	32637695	32453488	-0.56
franka_random	0	cpu	-	-	6259	6291	+0.51
franka_random	30000	cuda	-	CG	16824033	16816713	-0.04
franka_random	30000	cuda	-	Newton	16450657	16569780	+0.72
franka_random	30000	cuda	False	-	16699440	16623269	-0.46
franka_random	30000	cuda	True	-	11450874	11618053	+1.46
g1_fall	4096	cuda	-	Newton	921650	931884	+1.11
go2	4096	cuda	False	CG	3729044	3739424	+0.28
go2	4096	cuda	False	Newton	4558736	4551336	-0.16
go2	4096	cuda	True	-	3290047	3366242	+2.32
shadow_hand_cubes	0	cpu	-	-	41	41	+0.00
shadow_hand_cubes_sparse	0	cpu	-	-	65	66	+1.54

speed_test_654.txt

[duburcqa]

========= 642 passed, 3 skipped, 2 xfailed in 1140.09s (0:19:00) =========

[github-actions]

Total: 7 file(s) changed, +29 -36 code lines.

[github-actions]

Total: 7 file(s) changed, +29 -36 code lines.

[github-actions]

Total: 7 file(s) changed, +29 -36 code lines.

[github-actions]

Diff coverage: 100% · 5 lines, 0 missing

[github-actions]

Total: 7 file(s) changed, +29 -36 code lines.

[github-actions]

Diff coverage: 100% · 5 lines, 0 missing

[github-actions]

Total: 7 file(s) changed, +29 -36 code lines.

[github-actions]

Diff coverage: 0% · 0 lines, 0 missing

[github-actions]

Total: 7 file(s) changed, +29 -36 code lines.

[github-actions]

Diff coverage: 0% · 0 lines, 0 missing

[hughperkins]

Checklist:
[X] No user facing doc changes needed
[X] file line count changes look resaonable to me
[X] Genesis unit tests passing
[X] Genesis benchmarks at least ~neutral

[hughperkins]

Hmmm, but:

• why is there cuda stream pin.h , and not amd gpu stream pin.h?

[duburcqa]

why is there cuda stream pin.h , and not amd gpu stream pin.h?

First, because the issue that is being solved here was not observed on AMDGPU. Second, ADMGPU stream is not exposed in Quadrants before your PR for stream 1.

[hughperkins]

why is there cuda stream pin.h , and not amd gpu stream pin.h?

First, because the issue that is being solved here was not observed on AMDGPU. Second, ADMGPU stream is not exposed in Quadrants before your PR for stream 1.

ideally, I'd like Quadrants to be cross-gpu portable, on the whole.

since it seems streams pr is ready-ish to merge 🤞 , perhaps I get streams into a mergeable state first, we merge that, and then we can make this change symmetric across CUDA and AMDGPU?

[duburcqa]

ideally, I'd like Quadrants to be cross-gpu portable, on the whole.

it is cross-gpu portable. Here I'm using stream to get around some arguably faulty GPU driver behaviour on pre-Ampere architecture. It works perfectly fine on all AMDGPU machines I tested (CI + amdgpu box). As it does on Nvidia Blackwell.

[duburcqa]

since it seems streams pr is ready-ish to merge 🤞 , perhaps I get streams into a mergeable state first, we merge that, and then we can make this change symmetric across CUDA and AMDGPU?

We can, but I'm not in favour of speculatively fixing a potential bug that was never observed on AMDGPU.

[hughperkins]

since it seems streams pr is ready-ish to merge 🤞 , perhaps I get streams into a mergeable state first, we merge that, and then we can make this change symmetric across CUDA and AMDGPU?

We can, but I'm not in favour of speculatively fixing a bug that was never observed on AMDGPU.

Ideally, the CUDA and AMDGPU implementations should match as much as possible. This vastly facilitates maintenance and debugging I feel. I'd like to avoid forking the code to be different between them as much as possible, unless there is a strong reason, I feel.

[hughperkins]

(Ideally it'd be good to fuse the code somehow at some point, ideally, I feel. some kind of template class or something)

[hughperkins]

=> ok to merge

[github-actions]

Total: 7 file(s) changed, +29 -36 code lines.

[github-actions]

Diff coverage: 0% · 0 lines, 0 missing

[github-actions]

Total: 8 file(s) changed, +92 -41 code lines.

[github-actions]

Diff coverage: 0% · 0 lines, 0 missing