ROCm capability gaps for downstream mujoco_warp/mjlab: hipArray textures and conditional-graph gating via is_cuda

[shailensobhee]

Summary

While bringing up an mjlab / mujoco_warp RL training stack on MI300X (gfx942) with this ROCm Warp build (built from amd-integration, reports Warp 1.13.0+rocm.0), I hit two ROCm capability gaps that block otherwise CUDA-only downstream code. Both are understood and worked around on my side; filing here so they are tracked and, ideally, surfaced more clearly to downstream users.

The training itself works end to end on MI300X once these are handled (physics steps and the PPO update run on the GPU, sustained 100% utilization), so this is about closing capability gaps, not a fundamental blocker.

1. wp.Texture2D (hipArray textures) unsupported on ROCm

wp.Texture2D(...) fails on ROCm with:

RuntimeError: Failed to create CUDA texture: Warp CUDA error 801: operation not supported (in wp_texture_create_device, texture.cpp:171)

This surfaces in mujoco_warp's create_render_context, which materializes model textures even for raycast-only contexts. I have opened a mujoco_warp PR to guard that path with its existing use_textures flag (google-deepmind/mujoco_warp#1619), which avoids the call when no rendering is requested. But any downstream code that legitimately needs textures on ROCm (camera sensors) will still hit the underlying limitation.

Questions:

• Is hipArray texture-object support (wp_texture_create_device) on the roadmap for this fork, or is there a documented reason it cannot be supported on gfx942?
• If it is a hard limitation, would it be possible to raise a clearer, HIP-specific error message (as is already done for conditional graphs, see below) so downstream users immediately understand it is a backend capability gap rather than a usage error?

2. Conditional graph nodes: is_cuda is True on ROCm, so downstream gating does not trip

The fork correctly and clearly refuses conditional graph nodes:

RuntimeError: Conditional graph nodes are not supported on HIP/ROCm

(assert_conditional_graph_support in warp/_src/context.py). The clear message is appreciated.

The interop trap is that a ROCm device reports as CUDA. mujoco_warp's solver gates conditional-graph use on wp.get_device().is_cuda (solver.py in main), assuming is_cuda == True implies conditional graphs are available. On ROCm, is_cuda is True (the device is exposed as cuda:0), so the gate stays open and the solver calls wp.capture_while, which then raises the error above. Downstream I work around it by forcing opt.graph_conditional = False.

Question:

• Would it make sense to expose a capability query (for example wp.is_conditional_graph_supported(device) returning False on ROCm) so downstream libraries can gate on the actual capability instead of inferring it from is_cuda? That would let mujoco_warp/mjlab select the non-conditional solver path automatically on ROCm.

Environment

• GPU: AMD Instinct MI300X (gfx942), ROCm 10.0
• Warp: built from amd-integration, banner Warp 1.13.0+rocm.0
• Downstream: mujoco_warp 3.8.1, mjlab 1.3.0, PyTorch 2.10.0+rocm7.0

Happy to provide full repros or test any changes on MI300X hardware.

[shailensobhee]

Follow-up repro on newer hardware/software stack, confirming gap #1 (hipArray textures) still reproduces and that the mujoco_warp side is now merged.

Environment (all measured today, 2026-08-31):

• GPU: AMD Instinct MI300X VF (gfx942), ROCm 7.2.0
• Warp: built from this fork amd-integration-halo branch, banner Warp 1.17.0.dev4 (HIP). This is the 1.17-based branch (not the 1.13 amd-integration used in the original report), and it satisfies mujoco_warp mains warp-lang>=1.15requirement, soimport mujoco_warp` works directly with no shims.
• Downstream: google-deepmind/mujoco_warp main @ 7ec1b57, which now includes the merged PR Honor use_textures in create_render_context google-deepmind/mujoco_warp#1619 (merged 2026-08-31).

What works on MI300X with this stack:

• wp.get_devices() shows cuda:0 = "AMD Instinct MI300X VF" (192 GiB, gfx942).
• Full physics: 4096 envs, freejoint sphere dropped from z=1.0, timestep 0.004, settles to z=0.09963 by step 200 and holds through step 500, zero NaNs. All solver/constraint/cholesky/linesearch kernels compile and run on cuda:0.
• create_render_context(use_textures=False) (the raycast-only path enabled by the merged PR) returns cleanly with an empty texture registry.

What still hits gap #1:

• create_render_context(use_textures=True) fails at create_warp_texture -> wp.Texture2D:

RuntimeError: Failed to create CUDA texture: Warp CUDA error 801: operation not supported (in function wp_texture_create_device, /tmp/warp-halo/warp/native/texture.cpp:205)

• In the mujoco_warp render test suite on the GPU device: 9 pass, 5 fail. Four of the five fail purely on this texture-object 801 (test_create_warp_texture, test_get_segmentation, test_get_segmentation_preserves_flex_ids, and the use_textures=True half of test_use_textures_flag; its use_textures=False half passes). The fifth (test_render_context_with_textures) fails earlier on a separate mesh points/indices device-mismatch in build_mesh_bvh, unrelated to textures.

So the merged use_textures guard lets all texture-free paths (physics, raycast, depth, segmentation-from-camera-output) run on MI300X, but any path that genuinely samples a texture (RGB rasterization, camera sensors) is still blocked by the underlying hipArray limitation. Still happy to test any wp_texture_create_device changes on this MI300X.