[AMDGPU] Skip redundant per-launch RuntimeContext HtoD

[paveltc]

[amdgpu] Skip redundant per-launch RuntimeContext HtoD

Summary

On the AMDGPU launcher's default-stream path, launch_llvm_kernel re-uploads the
whole RuntimeContext struct to the per-handle persistent device buffer on
every launch:

memcpy_host_to_device_async(context_pointer, &ctx.get_context(),
                            sizeof(RuntimeContext), active_stream);

Across repeated launches of the same kernel handle that struct is almost always
byte-identical — the device arg-buffer / runtime / result-buffer pointers are
stable, and only the rare checkpoint kernel mutates the checkpoint_*_ptr
fields. Each of these copies costs ~4 µs of host-side API time (measured), which
is pure overhead on the launch-bound path (Genesis dispatches ~100 kernels per
step).

This PR caches the last-uploaded bytes (plus the device address they were
written to) in the per-handle Context and skips the HtoD when both still
match. Correctness is preserved unconditionally:

• The compare runs after prepare_streaming_checkpoint_state, so any
  checkpoint_*_ptr change forces a re-upload.
• A moved arg-buffer (arg_buffer pointer changes) or a reallocated context
  buffer (context_pointer changes) also forces a re-upload.
• The ephemeral (explicit-stream) path always uploads — it allocates a fresh
  buffer per launch and may run concurrently on pool streams.

Making the struct actually byte-stable

The skip-cache is useless unless RuntimeContext is stable. Profiling showed one
field flipping nearly every launch: result_buffer (offset 24). The launcher
only set it to the device buffer for result_buffer_size > 0 kernels; for
result-less kernels it left the per-launch-varying host pointer that
LaunchContextBuilder puts there. This PR pins result_buffer to the persistent
device buffer unconditionally:

• Result-producing kernels are unchanged (they already got the device pointer).
• Result-less kernels never touch the buffer, so pinning it to a stable, valid
  device address is safe — and is arguably more correct, since a host pointer
  sitting in a device-visible field is a latent hazard on AMDGPU (no UVA
  fallback).

Scope

• quadrants/runtime/amdgpu/kernel_launcher.{cpp,h} (the optimization), a one-line layout-neutral
  int32_t cpu_thread_id{0} in-class initializer in the backend-shared quadrants/program/context.h, and a new
  tests/python/test_amdgpu_context_cache.py.
• AMDGPU-local behavior; no codegen change, no public API. The context.h edit only adds a default initializer
  to an existing field (which value-initialization already zeroed at construction) — it does not change
  RuntimeContext's size or layout, so it is ABI-neutral even though the header is shared across backends.

Benchmarks (CDNA3, gfx942)

12-kernel × 5000-launch loop (launch-bound):

	per-launch	result checksum
baseline	17.7 µs	10223616.000000
this PR	14.4 µs (~19% faster)	10223616.000000 (identical)

Correctness also verified on a workload including a reduction kernel
(result_buffer_size > 0): identical checksum and output sum vs baseline.

End-to-end (Genesis rigid-body scenes, 4096 envs)

The microbench figure is a launch-bound ceiling. On real scenes the async
context copy overlaps with solver compute, so the end-to-end gain scales
inversely with per-step compute.

The numbers below are not from the repository's tests/benchmarks/test_rigid.py
harness. They come from a small custom driver that reuses that file's scene shapes
(anymal / franka / go2 / box-pyramid, same assets and control patterns) but with
its own measurement: a fixed n_envs = 4096 for every scene and a per-step
timing loop (warm 80 steps, best-of-3 over 400 steps) rather than the harness's
warmup/record runtime_fps at the suite's official env counts. It is a paired A/B
(same binary, the optimization gated behind an env flag and toggled back-to-back),
with physics byte-identical between arms:

scene	step time	steps/s Δ
franka (collision-free)	1.8 ms	+4.3%
franka	3.0 ms	+3.1%
anymal (no control)	3.4 ms	+2.8%
go2	4.2 ms	~+1%
anymal (per-env control)	11.9 ms	+0.1%
box pyramid (stacking)	47 ms	−0.4%

So: a consistent ~2–4% throughput gain on launch-bound articulated-body
scenes, tapering to ~0% (never a regression) as scenes become compute/collision
bound. Official test_rigid.py numbers at the suite's env counts are still worth
collecting as a follow-up.

Test plan

• Byte-identical results vs baseline on result-less and reduction kernels.
• Builds and runs on main with the AMDGPU backend.
• Custom end-to-end rigid-body A/B (above, test_rigid-derived scenes at
  4096 envs): positive on launch-bound scenes, neutral on compute-bound, no
  regressions.
• Official tests/benchmarks/test_rigid.py A/B at the suite's env counts.
  (MI300X, gfx942, ROCm 7.2.4; paired A/B — see results comment. Net-neutral
  +0.05% overall at default env counts; surfaced a reproducible,
  significant −1.92% regression on anymal_zero@4096 on the launch-bound
  sweep, flagged for consideration.)
• Confirm parity on a checkpoint-bearing (qd.checkpoint) kernel — the
  compare should force re-upload when checkpoint_*_ptr changes.
  (tests/python/test_checkpoint.py 39/39 pass on MI308X, gfx942; commit
  9e926e0.)
• Confirm parity under an explicit-stream / graph-capture path (must always
  upload). (Verified by construction: the skip-cache exists only on the
  null-stream default path; the ephemeral/explicit-stream path is unchanged
  and always uploads, and HIP graph capture cannot capture the null stream.)

Note on scope

This captures the host-side per-launch win without the higher-risk
kernarg-by-value codegen change (passing the struct by value would also remove
the per-instruction context pointer-loads in the kernel body, but touches shared
codegen and arch guards). It is deliberately launcher-local so it can land
independently and be reverted trivially.

[hughperkins]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 65ffc8cdf4

ℹ️ 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".

[hughperkins]

Checked with Opus:

• Opus says PR broadly looks good
• has some concerns, see below. Notably feels that the tests described in the test plan in this PR description are critical

Please could you check Opus concerns? :

[paveltc]

Thanks - both Opus and Codex caught real things. Pushed 9e926e0 addressing them, validated on MI308X (gfx942, ROCm 7.2.4).

Concern 1 - gate broader than its justification. Fixed. The skip is now gated on default_stream_path (active_stream == nullptr && all_sgid_zero), exactly matching the null-stream-ordering + host-serialization narrative. In the active_stream == nullptr && !all_sgid_zero (parallel-group-stream) case we still use persistent scratch but always re-upload - restoring pre-cache behavior for that path - and we refresh the cache after every upload so a later default-stream launch always compares against accurate bytes. So we never skip under the weaker visibility invariant you flagged.

Concern 2 - memcmp over padding + uninitialized field. Hardened + documented. cpu_thread_id was the one RuntimeContext scalar without an in-class initializer; it now defaults to {0} (layout-neutral). Added a comment noting the raw-byte compare is correctness-safe (a spurious hit is impossible - any meaningful field difference makes memcmp non-zero) and that the key is byte-stable because RuntimeContext is value-initialized at construction (make_unique<RuntimeContext>()), which zeroes padding and every initializer-less scalar. Worst case a perturbed byte costs a redundant re-upload, never a wrong skip.

Concern 3 - test-plan gaps. Addressed:

• checkpoint parity (the merge-blocker): ran tests/python/test_checkpoint.py on MI308X - 39/39 pass. These exercise the streaming-launch path with checkpoint state, i.e. the checkpoint_*_ptr-mutation-forces-reupload branch, which is the subtle one.
• new coverage: added tests/python/test_amdgpu_context_cache.py with three cases - repeated arg-less/result-less launches (cache-hit path + arg_buffer==nullptr), a result-producing reduction launched repeatedly (result_buffer pinning), and changing-ndarray launches that must force a re-upload. All pass on MI308X.
• explicit-stream always-upload: the ephemeral path is unchanged (always uploads, never caches); the changing-ndarray test exercises varied contexts. Official tests/benchmarks/test_rigid.py A/B at the suite env counts is still a good follow-up; my end-to-end numbers were from the derived 4096-env driver as noted in the description.

Concern 4 - thread-safety invariant. Documented: the non-atomic cache (std::vector + raw ptr) assumes same-handle launches are serialized on the host; the comment now states this and notes it holds because the skip path is null-stream-only.

Codex P2 (arg_buffer) is fixed as its own reply on that thread.

CI: the red Linters check was a clang-format reflow in the changed block - now clang-format/black/ruff clean and added comment lines wrap at <=120 (line-wrapping check). The Check test coverage gap is closed by the new test above.

Happy to run the official test_rigid.py A/B as a follow-up if you'd like it in-thread before merge.

[paveltc]

Official test_rigid.py A/B results (MI300X)

Follow-up on the test-plan A/B. Ran on an MI300X (gfx942, full CDNA3), ROCm 7.2.4, with this branch built from source. Paired A/B via a temporary, uncommitted env toggle (QD_DISABLE_CTX_CACHE): base = HtoD skip disabled (pre-optimization behavior), opt = skip enabled (this PR). Each arm runs in its own process (the gate is read once per process). Δ is paired opt/base − 1 on runtime_fps; significance is a paired t-test.

1) Official pytest harness — pytest tests/benchmarks/test_rigid.py::test_speed[<id>], default env counts (n=3)

case (official id)	base fps	opt fps	paired Δ	verdict
franka_free @30000	4,333,624	4,331,117	−0.04%	no sig diff
franka @30000	3,194,920	3,160,685	−1.06%	no sig diff
anymal_zero @30000	2,078,257	2,085,679	+0.36%	no sig diff
anymal_random @30000	1,340,015	1,343,121	+0.23%	no sig diff
go2 @4096	485,706	486,297	+0.12%	no sig diff
box_pyramid_3 @4096	301,579	303,716	+0.71%	no sig diff

Overall paired mean +0.05% (−1.06% … +0.71%). At the suite's default env counts the change is performance-neutral within noise (30000-env steps are compute-bound, so the small RuntimeContext HtoD is negligible there).

2) Launch-bound sweep at n_envs=4096 (n=3)

The official parametrization hard-codes n_envs=30000 for franka/anymal, so to probe the launch-bound regime the PR's custom-driver numbers came from, these reuse test_rigid.py's unmodified make_franka/make_anymal factories and run_benchmark() loop, invoked directly at 4096 (i.e. same scene construction + timing, not the pytest test_speed wrapper).

case @4096	base fps	opt fps	paired Δ	per-trial	verdict
franka	864,784	886,083	+2.46%	+4.8%, +1.5%, +1.1%	no sig diff (n=3)
franka_free	1,000,508	999,189	−0.10%	−0.2%, +3.6%, −3.8%	no sig diff
anymal_zero	917,387	900,473	−1.84%	−2.5%, −1.2%, −1.8%	SLOWER

3) anymal_zero @4096 confirmation (n=5)

The n=3 slowdown reproduced. base mean 917,501 → opt mean 899,838, paired Δ = −1.92% (std 0.90, SE 0.40); per-trial −1.47/−2.68/−2.69/−2.17/−0.59% (all negative); t = −4.79 vs t_crit(0.05,4df)=2.78 → statistically significant (p<0.01).

Bottom line

On the official harness the optimization is net neutral-to-mixed: franka@4096 shows the expected modest win (+2.46%, direction consistent with the PR description's +3.1%), but franka_free/anymal do not reproduce the larger custom-driver gains, and anymal_zero@4096 shows a small but reproducible, significant ~2% regression. That regression appears attributable to the per-launch memcmp of RuntimeContext on the gate path costing more than the saved HtoD in that scene. Net: the skip's benefit is narrow/launch-bound and scene-dependent, and the compare gate has a measurable cost in at least one scene — flagging for consideration.

Note: the toggle used for A/B is local/uncommitted and not part of this PR; the pushed branch is unchanged.

[hughperkins]

Comments from Opus:

Thoughts?

[paveltc]

Followed up on the anymal_zero@4096 regression flag with instrumentation (a local hit/miss counter on the RuntimeContext skip-cache) to settle whether it's real. Re-ran the A/B on a clock-locked MI308X (gfx942, perf_determinism), anymal_zero@4096, full 45s/15s protocol — but with the arm order randomized per pair rather than fixed base-then-opt.

• Cache hit-rate is 99.9% (~78 warmup misses per run). The struct is byte-stable on this scene as intended, so the opt arm does strictly less work per launch — the "compare costs more than the saved HtoD" mechanism can't hold when the compare is essentially always a hit.
• Randomizing the order flips the result: paired Δ = +1.92% (n=5, t=0.59, not significant), versus the earlier fixed-order −1.92%. Same magnitude, opposite sign — the signature of a base-then-opt ordering artifact (opt always ran second and ate a small systematic drift).

So the regression isn't real; it was measurement bias from fixed arm ordering. At ~7% run-to-run sd I'd call this scene net-neutral rather than a win, consistent with the official-harness +0.05%. Correctness was already covered (checkpoint parity 39/39 + the new cache tests). (Instrumentation was a local diagnostic, not part of the PR.)

[paveltc]

Addressed the two remaining minor review points:

1. Test naming/docs (commit 79171fb, test-only). The new tests don't force a RuntimeContext re-upload — the changing ndarray/scalar args ride the separately-uploaded arg_buffer, while RuntimeContext holds only the stable device arg_buffer address and so keeps hitting the cache. Renamed test_repeated_launch_changing_ndarray_forces_reupload → ..._arg_buffer_split and fixed the docstrings to describe them as cache-hit correctness coverage. The genuine forced-re-upload path — a checkpoint_*_ptr mutation inside RuntimeContext — is already covered by test_checkpoint.py.

2. Scope line. Updated to note the PR also touches the backend-shared quadrants/program/context.h (the one-line int32_t cpu_thread_id{0} initializer) and adds tests/python/test_amdgpu_context_cache.py. That initializer only adds a default to an already value-initialized field, so it's ABI-neutral (no change to RuntimeContext size/layout) despite the shared header.

[hughperkins]

Looks like this touches common (non-amd) files => I will run genesis benchmarks, to check for any regression.

[hughperkins]

running genesi benchmarks on bench1

[hughperkins]

Genesis benchmarks:

Seems quesitonale. I'll run bench_interleaved.py

[hughperkins]

rjunning bench_interleaved.py python bench_interleaved.py --quadrants-branch AMD-Ecosystem:feat/amdgpu-skip-redundant-context-h2d --num-runs 5 --partition rtx-mid --filter table_bussing,convexify,g1_fall_accessors,anymal_random --ref 20260831_redundant_1632

[hughperkins]

bench_interleaved.py results look ok-ish:

=== before/after interleaved (5 runs each, single node) ===
before qd=fa3b8a944 gen=6a2a1f0c  after qd=2e55982fc gen=6a2a1f0c

env                     batch  back   gjk    solv   n     before_fps     ±%      after_fps     ±%   delta%  min_sig%
anymal_random           20000  cuda                 4      6949313.2   0.45      6907567.8   1.38    -0.60      1.44
convexify                   0   cpu                 4           51.2   9.61           51.5   3.09     0.49     10.10
g1_fall_accessors        4096  cuda                 4      1269610.8   1.32      1285910.2   1.68     1.28      1.74
table_bussing              50  cuda                 4          342.0   4.85          338.5   4.25    -1.02      5.19
table_bussing              50   cpu                 4          123.5   3.07          122.0   1.51    -1.21      2.97
[14:18:30] Wrote /home/hugh/git/tmp/20260831_redundant_1632/meta.yaml
[14:18:30] Done. NO significant before/after difference

[hughperkins]

Ok, so what we need:

• Geneis benchmarks
• Genesis unit tests
• I check docs
• I check comment/code ratio (from eye-balling)
• I check whether we need agent CI
• this CI or parallel PR CI (depending on above decision) should run clean

[hughperkins]

no docs, so docs ok

[hughperkins]

looks like we already addressed comments/code ratio ✅

[hughperkins]

running genesis unit tests in tests0

[hughperkins]

probably need agent ci, so I'll need to create a parallel pr