Amd/hip graph support - #6
Open
zhihuidu-amd wants to merge 4 commits into
Open
Conversation
Makes wp.ScopedCapture() + mjw.step() + wp.capture_launch() work correctly
on AMD MI300X/MI325X (ROCm 7.x) without any environment variables.
Three changes enable zero-overhead graph replay:
1. solver.py: cache solver context on Data._solver_ctx on first call.
On HIP/ROCm, temporarily disable memory pool during allocation so buffers
use hipMalloc (stable addresses) not hipMallocAsync. hipMallocAsync
pointers appear as memAlloc nodes inside a hipGraph and re-execute on every
replay, adding ~5ms overhead per step. On CUDA this code path has no effect.
2. smooth.py: same fix for tendon scratch buffers (ten_Jdot, ten_bias_coef).
Both are wp.zeros() calls that fire inside graph capture on HIP, adding
two more memAlloc/memFree node pairs. Cached on Data with stable hipMalloc.
3. forward.py + io.py: optional multi-stream parallelism in fwd_position.
put_data() pre-creates two dedicated streams (stream_collision,
stream_secondary). fwd_position() uses them to run collision detection
and mass-matrix kinematics concurrently via fork-join.
The join uses stream.record_event / stream.wait_event (GPU-side events,
capturable as graph dependency edges) rather than synchronize_stream
(CPU-blocking, raises RuntimeError inside graph capture).
Falls back to sequential execution when sleep is enabled or streams
are unavailable.
Measured on AMD MI325X (gfx942, ROCm 7.2), humanoid model, 256 worlds:
Before: graph replay 7.1 ms/step (0.33x slower than eager 2.4 ms)
After: graph replay 1.4 ms/step (1.7x faster than eager 2.4 ms)
Usage (unchanged from upstream convention):
with wp.ScopedCapture() as cap:
mjw.step(model, data)
wp.capture_launch(cap.graph) # 1.7x faster on AMD ROCm
Provides a mempool-aware wrapper around wp.ScopedCapture() that works
correctly with all Warp versions on AMD ROCm:
with mjw.hip_graph_capture() as cap:
mjw.step(model, data)
wp.capture_launch(cap.graph) # ~1.7x faster on AMD ROCm
On HIP/ROCm: auto-enables mempool before ScopedCapture (required for
hipGraph capture) and restores pool state after. Some Warp versions
disable mempool globally in put_data() for ROCm 7.2 stability; this
helper ensures capture always has a compatible memory state.
On CUDA: behaves identically to wp.ScopedCapture() (no-op wrapper).
Note: wp.ScopedCapture() also works directly if the caller manages
mempool state manually.
Takes model+data, runs warmup_steps (default 3) eager steps before capture to trigger all lazy wp.zeros/wp.empty calls (solver context, tendon scratch, RK4 buffers, collision structures, etc.). After warmup these are cached on Data and do not fire during capture, ensuring the graph has zero memAlloc nodes and replays at full speed. Also enables mempool on HIP before capture and restores state after. Updated signature: mjw.hip_graph_capture(model, data, device=None, warmup_steps=3)
…m graph Multi-stream fork-join benefits eager execution but adds event-record/wait nodes to the captured graph. On ROCm 7.2, these event nodes add overhead during replay. During ScopedCapture (wp.get_device().is_capturing=True), use the sequential single-stream path to produce a minimal graph with only kernel nodes. Multi-stream continues to benefit eager execution.
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
1 participant
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.
PR Description: Opt A + HG (hipGraph capture + multi-stream)
Use for: zhihuidu-amd PR#1, AMD-Ecosystem new PR, google-deepmind PR#1556
Title:
feat(hip): HIP graph capture + multi-stream parallelism for AMD ROCm (Opt A+HG)Summary
Enables correct and efficient HIP graph capture for AMD ROCm in mujoco_warp, and adds multi-stream parallelism in
fwd_position. Makes the standard upstream graph capture pattern work on AMD MI300X/MI325X (ROCm 7.x) with zero user API changes.Measured on AMD MI325X (gfx942, ROCm 7.2), Unitree G1, 16,384 environments:
AMD MI325X with this PR surpasses H200 by 44% on Unitree G1 locomotion.
Root Cause of the Previous 3× Regression
On HIP/ROCm,
hipMallocAsynccalls insideScopedCapturebecomememAllocgraph nodes that re-execute every replay. Before this fix: 24 memAlloc + 24 memFree nodes → replay 7.1ms (0.33× slower than eager 2.4ms). After: 0 memAlloc nodes → 1.4ms (1.7× faster).Changes
io.py— Pre-allocate scratch buffers + dedicated streamsPre-allocates solver context, tendon scratch, and step_size_cost buffers using
hipMalloc(stable addresses, no graph nodes) before capture. Also pre-creates dedicated streams:forward.py— Multi-stream fork-join infwd_positionCollision detection and mass-matrix kinematics run concurrently after
fwd_kinematics, joined via GPU events (capturable as graph dependency edges — unlikesynchronize_streamwhich raises RuntimeError during HIP capture):Falls back to sequential when sleep is enabled or streams unavailable. Single-stream path used during graph capture for minimal graph overhead.
solver.py+smooth.py— Eliminate 24 memAlloc nodesCache
_solver_ctx,_ten_Jdot,_ten_bias_coefusinghipMallocon first call. Graph: 282 nodes → 234 nodes, 0 memAlloc.__init__.py—mjw.hip_graph_capture()context managerBenchmark Results
Step() throughput (physics only)
Graph node composition
End-to-end training throughput (16,384 envs, Unitree G1)
Validation (11/11 checks pass, AMD MI325X ROCm 7.2)
put_datapre-creates_stream_collision+_stream_secondary✅Dependencies
Companion PR to NVIDIA/warp #1702 for three HIP platform bug fixes required by this PR:
hipMalloc(0)returns NULL on ROCm → 1-byte minimum for zero-size arrayszero_()on zero-size array raiseshipErrorInvalidValue→ size guardWARP_DISABLE_MEMPOOLS_FOR_TORCHopt-outRelated