[AutoDiff] Speed up reverse-mode kernel launches on GPU backends

[duburcqa]

Speed up reverse-mode kernel launches on GPU backends via an adstack-sizer host-eval fast path

Adds a host-eval fast path on every GPU backend (CUDA / AMDGPU on the LLVM runtime; Metal / Vulkan / MoltenVK on the SPIR-V runtime) that bypasses the on-device sizer kernel + readback whenever every adstack alloca's SizeExpr is host-resolvable. Includes the LLVM-side async pinned-host HtoD optimisation, the SPIR-V predicate's FieldLoad correctness fix, a Codex/Claude-bot review pass on the LLVM path (active-stream routing, program_impl_ null-guard), a parametrized regression test, and several rounds of comment / wrap polish.

TL;DR

Every reverse-mode kernel launch with adstack allocas runs the SizeExpr capacity computation. Pre-PR that always meant a GPU dispatch:

• LLVM (CUDA / AMDGPU): HtoD of the encoded SizeExpr bytecode -> single-thread runtime_eval_adstack_size_expr kernel launch -> synchronous DtoH of the per-thread stride. The DtoH is a stream-sync that stalls the host until the sizer kernel has finished. With ~100 substeps x forward + backward x several reverse-mode tasks per substep, the test launches the sizer thousands of times and pays one host stall per launch.
• SPIR-V (Metal / Vulkan / MoltenVK): explicit flush() + device_->wait_idle() for PSB visibility -> sizer-bytecode upload -> per-task descriptor bind / dispatch -> submit_synced -> blocking metadata readback. Two host-side GPU stalls per kernel launch.

This PR detects the common case where no SizeExpr leaf needs device-resident memory and skips the entire dispatch on both backends.

Reported impact on the motivating workload (Genesis test_differentiable_rigid[gpu]): roughly 4x faster on CUDA after enabling this path.

Why

evaluate_adstack_size_expr already handles every leaf the on-device sizer was designed to interpret, except ExternalTensorRead whose data pointer is GPU-private. Detecting that one leaf at host time and skipping the dispatch is straightforward, and the unrolling_limit baseline is exactly the all-host-resolvable case (no adstacks at all => no SizeExpr at all => no dispatch), so this fast path is the closest the adstack mode can get to the unrolled baseline's launch overhead.

Mechanism

LLVM path (LlvmRuntimeExecutor::publish_adstack_metadata)

The function already had a CPU branch that host-evals each SerializedSizeExpr and writes the metadata arrays directly via copy_h2d. The branch was gated on !is_gpu_llvm. The new code:

1. Scans every alloca's size_expr.nodes for any node whose kind is SizeExpr::Kind::ExternalTensorRead. O(total node count across allocas).
2. Sets use_host_eval = !is_gpu_llvm || all_size_exprs_host_resolvable.
3. Routes to the existing host-eval branch when use_host_eval is true. CUDA / AMDGPU now reach this branch when no ExternalTensorRead is present.
4. Falls back to the existing runtime_eval_adstack_size_expr JIT call for kernels with ExternalTensorRead.

What dropped per launch when the fast path fires: one bytecode HtoD, one device sizer kernel launch, one DtoH stream-sync.

LLVM async pinned-host metadata HtoD

The fast path's three small per-launch HtoD copies (offsets, max_sizes, stride) are issued asynchronously from a pinned-host scratch via cuMemcpyHtoDAsync / hipMemcpyHtoDAsync. The host returns immediately after queueing the three copies on the active CUDA stream (CUDAContext::get_instance().get_stream() so user-set custom streams stream-order against the main-kernel dispatch correctly; AMDGPU keeps nullptr because AMDGPUContext::launch always uses the default stream). Pinned scratch is allocated lazily via cuMemAllocHost / hipHostMalloc and grown amortised-doubling; a per-launch CUDA / HIP event guards scratch reuse against in-flight DMAs. Eliminates the three serial host stalls per launch the synchronous cuMemcpyHtoD_v2 path had.

SPIR-V path (GfxRuntime::publish_adstack_metadata_spirv)

Two helpers in an anonymous namespace:

• all_size_exprs_host_resolvable(adstack_task_indices, task_attribs): scans every adstack-bearing task's allocas for an ExternalTensorRead or FieldLoad leaf. FieldLoad is the correctness gate: the host evaluator's FieldLoad path goes through SNodeRwAccessorsBank::read_int, whose nested accessor-kernel launch from inside the publish corrupts the SPIR-V launcher's per-task metadata-upload state and produces wrong gradients on every kernel that hits it. The on-device sizer was specifically built to handle FieldLoad on-device via PSB loads precisely because of this; the host-eval predicate must therefore reject both ExternalTensorRead (host can't read GPU-private memory) and FieldLoad (nested launch is unsafe).
• eval_per_task_metadata_on_host(adstack_task_indices, task_attribs, prog, host_ctx, per_task_ad_stack): replicates the sizer shader's per-task metadata layout ([stride_float, stride_int, (offset_i, max_size_i)*]) on the host. Float-heap accumulator advances by 2 * max_size (primal + adjoint), Int-heap by max_size, matching the running_off_f / running_off_i arithmetic in quadrants/codegen/spirv/adstack_sizer_shader.cpp.

The fast path runs after the adstack_task_indices early-out and before the sizer-pipeline build / bytecode upload / cmdlist record. When it fires, the function returns the host-computed per_task_ad_stack vector and never touches the sizer pipeline, the bytecode scratch buffer, the per-task metadata-buffer allocation, the flush(), the device_->wait_idle(), the sizer cmdlist record, the submit_synced, or the metadata readback - all of which are skipped entirely.

Per-backend coverage matrix

Backend	Pre-PR per-launch sizer cost	Post-PR (host-resolvable: no ExternalTensorRead, no FieldLoad)	Post-PR (with ExternalTensorRead or FieldLoad on SPIR-V)
CPU	Already host-eval (no GPU dispatch)	Unchanged	Unchanged
CUDA	HtoD bytecode + sizer kernel + DtoH stride sync	3 small async pinned HtoDs on the active stream, no kernel, no sync	Unchanged (sizer kernel + readback)
AMDGPU	HtoD bytecode + sizer kernel + DtoH stride sync	3 small async pinned HtoDs on the default stream, no kernel, no sync	Unchanged
Metal / MoltenVK	flush + wait_idle + bytecode upload + sizer cmdlist + readback	None - host-eval only (when SizeExpr is ExternalTensorRead and FieldLoad free)	Unchanged
Vulkan	Same as Metal (shared SPIR-V path)	Same as Metal	Unchanged

LLVM's host-eval FieldLoad is serviced by SNodeRwAccessorsBank exactly as before - no change for the LLVM CPU / CUDA / AMDGPU paths because the launcher reentrancy issue that gates SPIR-V doesn't apply there.

Tests

• pytest tests/python/test_adstack.py -n 8: 770 passed, 10 xfailed locally on macOS Vulkan / arm64 / Metal.
• test_adstack_sub_of_max_over_range_fusion_does_not_mix_fieldload_and_extread is parametrized on x_unused_val=[0.1, 100.0]. The amplified_unused_x variant pins any future cross-stack push / pop misroute as a 200.0 vs 0.0 mismatch (5+ orders of magnitude) instead of a 0.2 vs 0.0 "looks-like-tolerance" delta - the original 0.1 setup was added by this PR's predecessor and made the SPIR-V FieldLoad-during-publish corruption look like a numerical tolerance issue rather than the structural correctness bug it was.
• test_differentiable_rigid[gpu] end-to-end: ~4x faster on CUDA per the reported repro.
• tests/test_grad.py::test_differentiable_rigid[cpu] end-to-end: passes.

Codex / Claude bot review fixes

Severity	Issue	Fix
P1	Host-eval branch dereferenced program_impl_->program unconditionally; the on-device branch already supports program_impl_ == nullptr (C++-only tests) and falls back to max_size_compile_time	Hoisted the program_impl_ null-check before the evaluate_adstack_size_expr call, mirroring the on-device branch
P2	Async metadata HtoDs were hard-coded to default_stream = nullptr; user calls to CUDAContext::set_stream would leave kernels reading stale metadata	Routed CUDA copies through CUDAContext::get_instance().get_stream() so they stream-order against CUDAContext::launch's dispatch handle. AMDGPU keeps nullptr because AMDGPUContext::launch always passes nullptr to hipLaunchKernel

Side-effect audit

Concern	Where checked	Verdict
Host-eval correctness vs on-device sizer for ExternalTensorRead-free SizeExprs	evaluate_adstack_size_expr is the same function the on-device-bytecode encoder already calls during pre-substitution, so the leaves it can fold (Const / FieldLoad / BoundVariable / ExternalTensorShape / arithmetic / MaxOverRange) produce identical values	Same code path; no divergence between host fold and on-device sizer when the leaf set is identical
SPIR-V FieldLoad reentrancy	Predicate explicitly rejects FieldLoad so the SPIR-V publish never calls read_int from inside publish_adstack_metadata_spirv; the on-device sizer's PSB-load path handles FieldLoad correctly	Documented in the predicate's docstring and pinned by the parametrized regression test
Metadata layout match between SPIR-V host-eval and shader output	Per-stack offset uses 2 * max_size for the Float heap and max_size for the Int heap, matching adstack_sizer_shader.cpp's running_off_f / running_off_i accumulation; final stride values written into metadata[0] / metadata[1]	Bit-identical to shader output for the host-eval'able subset
LLVM byte-offset / 8-byte alignment match	Same align_up_8(sizeof(int64_t) + entry_size_bytes * max_size) formula already used by the existing CPU host-eval branch; no change	Identical to the previously CPU-only branch's behaviour
ExternalTensorRead falls back to on-device sizer	Single linear scan of size_expr.nodes returns false on the first ExternalTensorRead kind; the else arm of the dispatch retains the original LLVM runtime_eval_adstack_size_expr call and the SPIR-V cmdlist record path	Behaviour identical to pre-PR for kernels that hit this case
LLVM async metadata stream-ordering	Active-stream routing via CUDAContext::get_instance().get_stream(); pinned scratch reuse guarded by per-launch CUDA / HIP event so DMAs cannot race the host overwrite	Stream-ordered against the main-kernel dispatch on every CUDA stream the user can configure
Cache-load fallback for empty SizeExpr	When size_expr.nodes.empty() (offline-cache hit, symbolic tree not serialised) the host-eval path uses max_size_compile_time with the same max(_, 1) lower clamp the shader applies	Identical to shader behaviour
ProgramImpl / program back-reference	LLVM host-eval branch now guards on program_impl_ != nullptr before dereferencing program_impl_->program (P1 review fix); SPIR-V path's existing QD_ASSERT_INFO precondition is unchanged	C++-only-tests setups fall back to compile-time bounds rather than crashing

[claude]

I didn't find any bugs, but this PR touches the autodiff runtime on both LLVM and SPIR-V backends and introduces a host-side replica of the sizer shader's metadata layout, so it's worth a human look from someone who owns this code.

Extended reasoning...

Overview

This PR adds a host-eval fast path that bypasses the on-device adstack sizer dispatch when no SizeExpr in any adstack alloca contains an ExternalTensorRead leaf. Two files change: quadrants/runtime/gfx/adstack_sizer_launch.cpp (SPIR-V — Metal/Vulkan/MoltenVK) gets two new anonymous-namespace helpers and an early-return at the top of publish_adstack_metadata_spirv; quadrants/runtime/llvm/llvm_runtime_executor.cpp (CUDA/AMDGPU) generalizes the existing CPU-only host-eval branch to also fire on GPU when the same condition holds.

Security risks

None apparent. No auth, crypto, IPC, or network surface; this is a launch-path optimization in the autodiff runtime. The host-eval path writes the same metadata fields the GPU sizer would have written, via the same copy_h2d helpers already used elsewhere in the file.

Level of scrutiny

Higher than a typical perf tweak, lower than a public-API change. The motivating workload (Genesis test_differentiable_rigid[gpu]) and the reasoning are well-documented, but the SPIR-V path now contains a host replica of the sizer shader's metadata layout ([stride_float, stride_int, (offset_i, max_size_i)*], with 2 * max_size for Float / max_size for Int heap accumulation, and a max(_, 1) floor). Any future change to adstack_sizer_shader.cpp's running_off_f / running_off_i accumulator now has to be mirrored in eval_per_task_metadata_on_host, and divergence will silently mis-size GPU heaps for autodiff kernels — a class of bug that does not surface as a build failure or a clean test signal.

Other factors

The bug hunter found nothing, the description is thorough and includes a side-effect audit, and CPU + test_ad_basics tests pass. The author explicitly notes that the motivating CUDA/GPU benchmark wasn't run locally, which is the main thing I'd want a human owner to verify before merging. The change is also touching the same area as several recent commits (f1dc891, 5d7cf54, 03ab3b3), so a reviewer with that recent context will catch interactions I can't see from the diff alone.

[github-actions]

Coverage Report (4f0688fca)

Metric	Value
Diff coverage (changed lines only)	0%
Overall project coverage	61%

Total: 0 lines, 0 missing, 0% covered

[hughperkins]

Cool. Since I remember the user-facing doc talks a lot about sizing, is there anything that needs updating in them?

[duburcqa]

Cool. Since I remember the user-facing doc talks a lot about sizing, is there anything that needs updating in them?

I don't think so, this is purely internal. It adds a host-eval fast path to bypass the on-device sizer kernel whenever it is possible. Code logics, just different execution path.

[chatgpt-codex-connector]

💡 Codex Review

quadrants/quadrants/runtime/llvm/llvm_runtime_executor.cpp

Line 742 in 4a9d231

if (expr != nullptr && !expr->nodes.empty() && program_impl_->program != nullptr) {

Guard host-eval path when ProgramImpl is unavailable

The new GPU host-eval fast path dereferences program_impl_ unconditionally in program_impl_->program != nullptr. In the same function, the on-device branch explicitly documents/supports program_impl_ == nullptr for C++-only setups; with this change, kernels whose size exprs are host-resolvable can now hit this branch first and crash with a null dereference instead of taking the existing compile-time fallback behavior.

ℹ️ 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]

P1 fixed in e269abe94: gated program_impl_->program behind a null-check on program_impl_, falling back to max_size_compile_time like the on-device branch.

P2 fixed in bf7c12607: routed the CUDA async copies through CUDAContext::get_instance().get_stream() so they share the kernel-launch stream.

[github-actions]

Coverage Report (32b803f0c)

File	Coverage	Missing
🟢 tests/python/test_adstack.py	100%

Diff coverage: 100% · Overall: 73% · 3 lines, 0 missing

Full annotated report

[github-actions]

Coverage Report (037c4530f)

File	Coverage	Missing
🟢 python/quadrants/_kernels.py	100%
🔴 python/quadrants/lang/_fast_caching/args_hasher.py	67%	9,37
🔴 python/quadrants/lang/_func_base.py	75%	21
🟢 python/quadrants/lang/_kernel_impl_dataclass.py	100%
🟢 python/quadrants/lang/_ndarray.py	86%	175,325
🔴 python/quadrants/lang/_ndarray_pickle.py	67%	8
🟢 python/quadrants/lang/_template_mapper_hotpath.py	100%
🔴 python/quadrants/lang/any_array.py	67%	20
🟢 python/quadrants/lang/ast/ast_transformer.py	88%	669
🟢 python/quadrants/lang/ast/ast_transformers/call_transformer.py	100%
🟢 python/quadrants/lang/ast/ast_transformers/function_def_transformer.py	100%
🟢 python/quadrants/lang/field.py	100%
🟢 python/quadrants/lang/impl.py	80%	18
🟢 python/quadrants/lang/kernel.py	100%
🟢 python/quadrants/lang/kernel_arguments.py	100%
🔴 python/quadrants/lang/matrix.py	70%	964,967,973,978,1161,1720,1831
🟢 tests/python/quadrants/lang/test_dlpack.py	100%
🟢 tests/python/test_adstack.py	100%
🟢 tests/python/test_api.py	100%
🟢 tests/python/test_pickle.py	100%

Diff coverage: 84% · Overall: 73% · 99 lines, 16 missing

Full annotated report

[duburcqa]

GS_ENABLE_NDARRAY=1 QD_OFFLINE_CACHE=0 pytest -v -ra --backend gpu --dev --forked ./tests

================== 649 passed, 3 skipped, 2 xfailed in 1108.25s (0:18:28) ===============

[duburcqa]

Comparing Quadrants main vs this PR on Genesis main:

env	batch_size	backend	gjk_collision	constraint_solver	runtime_fps_main	runtime_fps_578	runtime_fps_delta_pct
anymal_random	30000	cuda	-	-	9236721	9279341	+0.46
anymal_uniform	30000	cuda	-	-	12290801	12341098	+0.41
anymal_uniform_kinematic	0	cpu	-	-	2054	2022	-1.56
anymal_uniform_kinematic	30000	cuda	-	-	10466141	10482269	+0.15
anymal_zero	0	cpu	-	-	7257	7148	-1.50
anymal_zero	30000	cuda	-	-	18701483	18688757	-0.07
box_pyramid_3	4096	cuda	-	-	969382	972933	+0.37
box_pyramid_4	4096	cuda	-	-	388357	389805	+0.37
box_pyramid_5	4096	cuda	-	-	141016	139732	-0.91
box_pyramid_6	4096	cuda	False	-	59335	58789	-0.92
box_pyramid_6	4096	cuda	True	-	60333	59872	-0.76
dex_hand	4096	cuda	-	-	17048	17162	+0.67
duck_in_box_easy	30000	cuda	False	-	26662843	26715336	+0.20
duck_in_box_easy	30000	cuda	True	-	9549541	9736801	+1.96
duck_in_box_hard	0	cpu	-	-	5185	5159	-0.50
duck_in_box_hard	30000	cuda	False	-	10239009	10119850	-1.16
duck_in_box_hard	30000	cuda	True	-	3405035	3454206	+1.44
franka	30000	cuda	-	-	22058180	21724297	-1.51
franka_accessors	0	cpu	-	-	1227	1217	-0.81
franka_accessors	30000	cuda	-	-	15590143	15553821	-0.23
franka_free	30000	cuda	-	-	31981995	32202953	+0.69
franka_random	0	cpu	-	-	6457	6324	-2.06
franka_random	30000	cuda	-	CG	16737391	16698827	-0.23
franka_random	30000	cuda	-	Newton	16343418	16525809	+1.12
franka_random	30000	cuda	False	-	16534089	16369501	-1.00
franka_random	30000	cuda	True	-	11429393	11448748	+0.17
g1_fall	4096	cuda	-	Newton	921076	923492	+0.26
go2	4096	cuda	False	CG	3678736	3676077	-0.07
go2	4096	cuda	False	Newton	4438961	4451375	+0.28
go2	4096	cuda	True	-	3273783	3272905	-0.03
shadow_hand_cubes	0	cpu	-	-	40	41	+2.50
shadow_hand_cubes_sparse	0	cpu	-	-	66	65	-1.52

speed_comparison_main_vs_578.csv

[hughperkins]

checklist:

• nothing to update in user facing docs
• Genesis unit tests passing
• Genesis benchmarks appear neutral

=> ok to merge

[github-actions]

Coverage Report (8d0275803)

File	Coverage	Missing
🟢 tests/python/test_adstack.py	100%

Diff coverage: 100% · Overall: 73% · 3 lines, 0 missing

Full annotated report