Enable HIP Graph support

warp/_src/context.py

@@ -3858,6 +3858,12 @@ def __init__(self, runtime, alias, ordinal=-1, is_primary=False, context=None):
             if warp.config.enable_mempools_at_init:
                 # enable if supported
                 self.is_mempool_enabled = self.is_mempool_supported
+            elif self.is_hip and self.is_mempool_supported:
+                # HIP/ROCm: enable mempool by default when graph capture is supported
+                # (hipGraph requires mempool for in-capture allocations).
+                # Previously disabled due to hipMemsetAsync unreliability, but
+                # graph capture is now validated on ROCm 7.2 (Warp PR #15).
+                self.is_mempool_enabled = True
             else:
                 # disable by default
                 self.is_mempool_enabled = False
@@ -3970,15 +3976,18 @@ def is_capturing(self) -> bool:
     def supports_graph_capture(self) -> bool:
         """A boolean indicating whether native graph capture is supported on this device.
 
-        Returns ``True`` for CPU (always recordable via APIC) and CUDA devices.
-        Returns ``False`` for HIP/ROCm devices: hipGraph is not yet mature in
-        Warp's runtime (mempool pointer remapping bugs, event timing inside
-        graphs, etc.), so :func:`capture_begin` / :class:`ScopedCapture` are
-        no-ops on HIP and :func:`capture_save` is unavailable for HIP-targeted
-        graphs.
+        Returns ``True`` for CPU (always recordable via APIC), CUDA, and HIP/ROCm
+        devices. HIP graph capture uses the same native capture/replay path as
+        CUDA (all ``hipGraph*`` entry points are aliased in ``hip_util.h``).
+
+        Note: two graph features remain gated separately on HIP:
+
+        - Conditional graph nodes (if/else/while subgraphs), tracked by
+          :func:`is_conditional_graph_supported` (still ``False`` on HIP).
+        - APIC ``.wrp`` serialization via :func:`capture_save`, which encodes a
+          CUDA ``sm`` architecture integer and does not yet support HIP ``gfx``
+          string architectures.
         """
-        if self.is_hip:
-            return False
         return True
 
     @property
@@ -8743,13 +8752,21 @@ def synchronize():
         # save the original context to avoid side effects
         saved_context = runtime.core.wp_cuda_context_get_current()
 
+        # Refuse to synchronize while a graph capture is active on ANY device
+        # before performing any partial synchronization. On HIP/ROCm graph
+        # capture is process-global, so synchronizing a non-capturing device's
+        # context while another device is capturing fails and invalidates the
+        # in-progress capture. Checking all devices up front keeps the error
+        # clean and leaves the capture intact on both CUDA and HIP.
+        for device in runtime.cuda_devices:
+            # avoid creating primary context if the device has not been used yet
+            if device.has_context and device.is_capturing:
+                raise RuntimeError(f"Cannot synchronize device {device} while graph capture is active")
+
         # TODO: only synchronize devices that have outstanding work
         for device in runtime.cuda_devices:
             # avoid creating primary context if the device has not been used yet
             if device.has_context:
-                if device.is_capturing:
-                    raise RuntimeError(f"Cannot synchronize device {device} while graph capture is active")
-
                 runtime.core.wp_cuda_context_synchronize(device.context)
 
         # restore the original context to avoid side effects
@@ -9453,9 +9470,7 @@ def capture_begin(
     if stream is None:
         stream = device.stream
 
-    # HIP/ROCm graph capture is not yet mature (mempool pointer remapping bugs,
-    # event timing inside graphs, etc.). Disable until hipGraph stabilizes.
-    # See ``Device.supports_graph_capture`` -- callers (notably
+    # Devices that cannot record a native graph return early. Callers (notably
     # :class:`ScopedCapture`) treat a ``False`` return as "capture disabled" and
     # skip the matching :func:`capture_end`.
     if not device.supports_graph_capture:
@@ -9946,11 +9961,9 @@ def capture_launch(graph: Graph, stream: Stream | None = None):
 
     if graph is None:
         # ScopedCapture leaves ``graph=None`` when ``capture_begin`` returned
-        # False (currently only on HIP, where graph capture is unsupported;
-        # see ``Device.supports_graph_capture``).
+        # False on a device where ``Device.supports_graph_capture`` is False.
         raise RuntimeError(
-            "capture_launch() received graph=None: capture was not started. "
-            "On HIP/ROCm, native graph capture is unsupported "
+            "capture_launch() received graph=None: capture was not started "
             "(Device.supports_graph_capture is False); guard call sites or "
             "filter test devices with get_graph_capture_test_devices()."
         )
@@ -10026,6 +10039,15 @@ def capture_save(graph: Graph, path: str, inputs: dict | None = None, outputs: d
     import os  # noqa: PLC0415
     import shutil  # noqa: PLC0415
 
+    if graph.device.is_hip:
+        # The APIC .wrp format encodes a CUDA sm architecture integer and derives
+        # the device type from it; HIP gfx string architectures are not yet
+        # representable. In-memory HIP graph capture/replay is unaffected.
+        raise RuntimeError(
+            "capture_save() is not supported on HIP/ROCm: the APIC .wrp format does not "
+            "yet support HIP gfx architectures. In-memory graph capture/replay is supported."
+        )
+
     if not graph.apic:
         raise RuntimeError(
             "Graph was not captured with apic=True. Pass apic=True to capture_begin() or ScopedCapture()."

warp/_src/optim/linear.py

@@ -1112,21 +1112,33 @@ def do_cycle_with_condition():
         if callback_launch is not None:
             callback_launch.launch()
 
+    def fixed_count_loop():
+        # Fallback that avoids conditional graph nodes: run a fixed number of
+        # cycles. The CG/CR update kernels zero out alpha/beta once the residual
+        # drops below tolerance (see _cg_kernel_1/_cg_kernel_2), so iterating past
+        # convergence is a no-op and the result is unchanged. When invoked inside an
+        # outer capture, these launches are recorded into that graph.
+        for _ in range(0, maxiter, cycle_size):
+            do_cycle_with_condition()
+
     if use_cuda_graph and device.is_cuda:
-        if device.is_capturing:
+        if not wp.is_conditional_graph_supported():
+            # Conditional graph nodes are unavailable (e.g. HIP/ROCm). Use a
+            # fixed-count loop instead of wp.capture_while so the solver still runs
+            # (and stays capturable) without requiring if/while subgraph support.
+            fixed_count_loop()
+        elif device.is_capturing:
             wp.capture_while(condition, do_cycle_with_condition)
         else:
             with wp.ScopedCapture() as capture:
                 wp.capture_while(condition, do_cycle_with_condition)
             if capture.graph is not None:
                 wp.capture_launch(capture.graph)
             else:
-                # Graph capture not supported (e.g. HIP), fall back to eager loop
-                for _ in range(0, maxiter, cycle_size):
-                    do_cycle_with_condition()
+                # Graph capture not supported, fall back to eager loop
+                fixed_count_loop()
     else:
-        for _ in range(0, maxiter, cycle_size):
-            do_cycle_with_condition()
+        fixed_count_loop()
 
     return cur_iter, r_norm_sq, atol_sq
 

warp/examples/fem/example_kelvin_helmholtz.py

@@ -510,14 +510,19 @@ def __init__(
         self.renderer = fem_example_utils.Plot()
         self.renderer.add_field("rho", self.rho_field)
 
-        # Capture CUDA graph for the simulation step
-        self.use_cuda_graph = wp.get_device().is_cuda
+        # Capture CUDA graph for the simulation step.
+        # NOTE: on HIP/ROCm 7.2, instantiating the large, mempool-allocation-heavy
+        # graph this DG step produces crashes inside hipGraphInstantiate (heap
+        # corruption in libamdhip64). Capture itself succeeds, so the graph-is-None
+        # check below does not catch it; skip capture up front and run eagerly.
+        device = wp.get_device()
+        self.use_cuda_graph = device.is_cuda and not device.is_hip
         if self.use_cuda_graph:
             with wp.ScopedCapture() as capture:
                 self.simulate()
             self.graph = capture.graph
             if self.graph is None:
-                self.use_cuda_graph = False  # Graph capture is disabled on HIP.
+                self.use_cuda_graph = False
 
     def _state_delta(self, trial_state):
         """Evaluate the DG spatial operator: M^{-1} [volume_rhs + side_rhs]."""

warp/examples/fem/example_shallow_water.py

@@ -433,14 +433,19 @@ def __init__(self, quiet=False, resolution=50, degree=1, domain_size=1.0):
         self.renderer = fem_example_utils.Plot()
         self.renderer.add_field("height", self.height_field)
 
-        # Capture CUDA graph for the simulation step
-        self.use_cuda_graph = wp.get_device().is_cuda
+        # Capture CUDA graph for the simulation step.
+        # NOTE: on HIP/ROCm 7.2, instantiating the large, mempool-allocation-heavy
+        # graph this DG step produces crashes inside hipGraphInstantiate (heap
+        # corruption in libamdhip64). Capture itself succeeds, so the graph-is-None
+        # check below does not catch it; skip capture up front and run eagerly.
+        device = wp.get_device()
+        self.use_cuda_graph = device.is_cuda and not device.is_hip
         if self.use_cuda_graph:
             with wp.ScopedCapture() as capture:
                 self.simulate()
             self.graph = capture.graph
             if self.graph is None:
-                self.use_cuda_graph = False  # Graph capture is disabled on HIP.
+                self.use_cuda_graph = False
 
     def _state_delta(self, trial_state):
         """Evaluate the DG spatial operator: M^{-1} [volume_rhs + side_rhs].

warp/examples/fem/example_taylor_green.py

@@ -557,7 +557,12 @@ def __init__(self, quiet=False, degree=2, resolution=25, Re=100.0, mesh: str = "
         # reduced launch overhead.  Before each replay, step() computes
         # the time-dependent BC contributions into the working buffers
         # (_u_bd_rhs_buf, _p_rhs) that the graph references.
-        self.use_cuda_graph = wp.get_device().is_cuda
+        # NOTE: on HIP/ROCm 7.2, instantiating the large, mempool-allocation-heavy
+        # graph this step produces crashes inside hipGraphInstantiate (heap
+        # corruption in libamdhip64). Capture itself succeeds, so the graph-is-None
+        # check below does not catch it; skip capture up front and run eagerly.
+        device = wp.get_device()
+        self.use_cuda_graph = device.is_cuda and not device.is_hip
         if self.use_cuda_graph:
             import gc  # noqa: PLC0415
 
@@ -569,7 +574,7 @@ def __init__(self, quiet=False, degree=2, resolution=25, Re=100.0, mesh: str = "
             finally:
                 gc.enable()
             if self.graph is None:
-                self.use_cuda_graph = False  # Graph capture is disabled on HIP.
+                self.use_cuda_graph = False
 
     def _compute_bc(self, t):
         """Compute time-dependent Dirichlet BC contributions at time ``t``.

warp/native/cuda_util.cpp

@@ -798,7 +798,12 @@ CUresult cuCtxEnablePeerAccess_f(CUcontext peer_ctx, unsigned int flags)
 #if defined(__HIP_PLATFORM_AMD__)
     if (!peer_ctx)
         return CUDA_ERROR_NOT_INITIALIZED;
-    return hipDeviceEnablePeerAccess(peer_ctx->device, flags);
+    hipError_t status = hipDeviceEnablePeerAccess(peer_ctx->device, flags);
+    // hipDeviceEnablePeerAccess is a runtime API that sets the sticky last-error state (e.g.
+    // hipErrorPeerAccessAlreadyEnabled/704). Clear it so a benign status does not leak into an
+    // unrelated runtime call (such as a subsequent hipcub call) and get reported there.
+    ignore_cuda_error(hipGetLastError());
+    return status;
 #else
     return pfn_cuCtxEnablePeerAccess ? pfn_cuCtxEnablePeerAccess(peer_ctx, flags) : DRIVER_ENTRY_POINT_ERROR;
 #endif  // defined(__HIP_PLATFORM_AMD__)