[AMDGPU] Skip redundant per-launch RuntimeContext HtoD - #876
Open
paveltc wants to merge 1 commit into
Open
Conversation
The AMDGPU launcher re-uploads the whole RuntimeContext struct to the per-handle persistent device buffer on every launch. Across repeated launches of the same kernel handle the struct is almost always identical, so most of these ~4us async copies are pure overhead on the launch-bound path. Cache the last-uploaded bytes (plus the device address they were written to) in the per-handle Context and skip the HtoD when both still match. The compare runs after prepare_streaming_checkpoint_state, so any checkpoint_*_ptr mutation forces a re-upload; a moved arg-buffer or reallocated context buffer also forces one. The ephemeral (explicit-stream) path always uploads. To make the struct actually byte-stable, also pin RuntimeContext.result_buffer to the persistent device buffer unconditionally. Previously it was only set for result_buffer_size > 0 kernels, leaving a per-launch-varying host pointer in the field for result-less kernels - which defeated the cache and was a latent hazard on AMDGPU (no UVA fallback for host pointers). Launcher-local, AMDGPU-only; no codegen or ABI change. Measured on CDNA3 (gfx942), 12-kernel x 5000-launch loop: per-launch 17.7us -> 14.4us (~19%), with byte-identical results on both result-less and reduction kernels. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[amdgpu] Skip redundant per-launch RuntimeContext HtoD
Summary
On the AMDGPU launcher's default-stream path,
launch_llvm_kernelre-uploads thewhole
RuntimeContextstruct to the per-handle persistent device buffer onevery 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_*_ptrfields. 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
Contextand skips the HtoD when both stillmatch. Correctness is preserved unconditionally:
prepare_streaming_checkpoint_state, so anycheckpoint_*_ptrchange forces a re-upload.arg_bufferpointer changes) or a reallocated contextbuffer (
context_pointerchanges) also forces a re-upload.buffer per launch and may run concurrently on pool streams.
Making the struct actually byte-stable
The skip-cache is useless unless
RuntimeContextis stable. Profiling showed onefield flipping nearly every launch:
result_buffer(offset 24). The launcheronly set it to the device buffer for
result_buffer_size > 0kernels; forresult-less kernels it left the per-launch-varying host pointer that
LaunchContextBuilderputs there. This PR pinsresult_bufferto the persistentdevice buffer unconditionally:
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}only.Benchmarks (CDNA3, gfx942)
12-kernel × 5000-launch loop (launch-bound):
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.pyharness. 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 = 4096for every scene and a per-steptiming loop (warm 80 steps, best-of-3 over 400 steps) rather than the harness's
warmup/record
runtime_fpsat 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:
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.pynumbers at the suite's env counts are still worthcollecting as a follow-up.
Test plan
mainwith the AMDGPU backend.test_rigid-derived scenes at4096 envs): positive on launch-bound scenes, neutral on compute-bound, no
regressions.
tests/benchmarks/test_rigid.pyA/B at the suite's env counts.qd.checkpoint) kernel — thecompare should force re-upload when
checkpoint_*_ptrchanges.upload).
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.