Enable HIP Graph support

[rtmadduri]

Enable HIP/ROCm Graph-Capture Support

Enables Warp's existing graph-capture infrastructure on HIP/ROCm and validates it on an AMD Instinct MI325X (gfx942, ROCm/HIP 7.2). Phase 1 makes core capture/replay, graph allocations, APIC save, examples, and unit tests work on HIP.

1. Central software gate disabled HIP graph capture

• Issue: Device.supports_graph_capture returned False for HIP, which cascaded through capture_begin() (returned early with graph=None), the test device helpers, and example fallbacks — disabling all graph capture on HIP even though the core native capture path was already device-agnostic.
• Fix: Device.supports_graph_capture now returns True on HIP (warp/_src/context.py); docstring updated. The capture_launch graph is None guard is kept as a safety net.
• Result: In-memory HIP graph capture/replay is enabled and passing. test_graph 12/12 (incl. test_graph_alloc mempool pointer-remapping), and graph-using suites pass (test_map, test_reload, test_sparse, test_array, test_grad, fem/test_fem_integrate, tile/test_tile_shared_memory, geometry/test_mesh).

2. capture_pause/capture_resume miscategorized as conditional-only

• Issue: FEM examples crashed on HIP with "must be CUDA Toolkit 12.4+". wp_cuda_graph_pause_capture/wp_cuda_graph_resume_capture were defined inside the #if !defined(__HIP_PLATFORM_AMD__) && CUDA_VERSION >= 12040 conditional-node block, with no-op HIP stubs that raised. Pause/resume is actually used to suspend capture while JIT-compiling a module mid-capture — unrelated to conditional (if/while) subgraphs.
• Fix: Relocated the real pause/resume implementations out of the conditional-node #if block so they compile unconditionally (they only use hipStreamGetCaptureInfo/cudaStreamEndCapture/cudaStreamBeginCaptureToGraph, all already aliased in hip_util.h); removed the HIP stubs (warp/native/warp.cu).
• Result: The stub error no longer occurs. Mid-capture module JIT (e.g. FEM quadrature caching) works on HIP.

3. Linear solvers required conditional graph nodes

• Issue: The fully-captured convergence loop (check_every=0) in optim/linear.py used wp.capture_while (conditional graph nodes), raising RuntimeError: Conditional graph nodes are not supported on HIP/ROCm for cg/cr/bicgstab/gmres (hit by example_taylor_green).
• Fix: _run_capturable_loop now checks wp.is_conditional_graph_supported() and, when false (HIP), runs a fixed-count loop of maxiter cycles instead of capture_while. This is numerically correct because the CG/CR/BiCGSTAB/GMRES update kernels zero out alpha/beta once the residual falls below tolerance, so iterating past convergence is a no-op.
• Result: All four solvers converge to rel_err ~1e-6 on the MI325X with both use_cuda_graph=False and True.

4. BVH rebuild-in-capture not graph-capturable on HIP

• Issue: test_capture_bvh_rebuild (and grouped) rely on the LBVH builder's rocPRIM/hipcub DeviceRadixSort, which is not graph-capturable during HIP capture (works fine outside capture), leaving BVH with no graph-capture coverage on HIP.
• Fix: Added test_capture_bvh_query (tests/geometry/test_bvh.py) and test_grouped_capture_bvh_query (tests/geometry/test_grouped_bvh.py): build the BVH outside the capture, then capture and replay the query launches against an eager reference. Rebuild-in-capture tests are skipped on HIP.
• Result: BVH keeps graph-capture coverage on HIP; both new tests pass on the MI325X and also run on CUDA.

5. Test suite and example audit for HIP graph capture

• Issue: Several test files and examples had raw not d.is_hip graph guards or stale fallbacks that bypassed the central capability, plus genuine ROCm limitations that would surface as spurious failures.
• Fix: Routed graph guards through supports_graph_capture/get_graph_capture_test_devices (tests/unittest_utils.py), and added targeted, documented skips for genuine ROCm 7.2 limitations:
  • APIC .wrp save/load (capture_save): format encodes a CUDA sm integer arch/device type; HIP gfx strings aren't representable — capture_save raises a clear error on HIP and test_apic/test_apic_mesh exclude HIP.
  • In-graph event timing (test_event_elapsed_time_graph) and external-event cross-graph sync (test_event_external) skipped in tests/cuda/test_streams.py.
  • Async-copy graph-capture variants deferred in tests/cuda/test_async.py.
  • test_fixedarray.py conditional-graph check switched to wp.is_conditional_graph_supported() so it skips correctly on HIP.
• Result: graph/async/streams suites and graph-using tests pass on the MI325X with conditional-capture tests self-skipping. cuda/test_streams 24 OK / 4 skipped, cuda/test_async 1213 OK, cuda/test_conditional_captures 32 OK with 16 self-skipping via is_conditional_graph_supported().

---

Validation

All results validated on AMD Instinct MI325X (gfx942), ROCm/HIP 7.2, native lib rebuilt for gfx942.