Extend CPU graph capture coverage (#3501)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: c0d1f1ed

CHANGELOG.md

@@ -111,6 +111,7 @@
 - Preserve muscles and rigid-body color groups when copying or replicating a `ModelBuilder`.
 - Fix `ModelBuilder.add_usd()` to honor `PhysicsScene.gravityDirection`, including stage-to-builder rotation and per-world imports.
 - Fix stale overlay layers remaining visible after switching examples in the OpenGL viewer.
+- Fix `SolverKamino` CG/CR solves silently under-iterating on CPU graph capture; the capture-safe loop path now runs on any capturing device, not only CUDA, so CPU captures no longer record a stale host-readback convergence decision at record time.
 - Reject incompatible custom attribute and frequency definitions before composing `ModelBuilder` instances.
 - Fix `cloth_franka` example rendering particles at simulation scale (cm) instead of viewer scale (m)
 - Fix `ModelBuilder` merges to accept array-valued transform fields and plain-list particle color groups.

docs/concepts/simulation_tuning_solvers.rst

@@ -157,15 +157,21 @@ repository examples spend tuning effort, not a shared solver API.
        ``particle_topological_contact_filter_threshold``,
        ``particle_rest_shape_contact_exclusion_radius``.
      - Contact history requires matched contacts, for example
-       ``CollisionPipeline(contact_matching="latest")``. When recording VBD
-       steps in a CUDA graph, construct :class:`~newton.CollisionPipeline`
-       before :class:`~newton.solvers.SolverVBD` so contact history is
-       pre-allocated, or run one uncaptured solver step before capture. Buffer
-       sizes that are too small can drop contacts; sizes that are too large cost
-       memory and performance. Examples commonly tune ``iterations``, particle
-       self-contact radius and margin, particle contact buffers and filters,
-       ``particle_collision_detection_interval``, ``particle_enable_tile_solve``,
-       ``rigid_body_contact_buffer_size``,
+       ``CollisionPipeline(contact_matching="latest")``. Contact history is
+       cross-replay-persistent state, so it must always be pre-allocated
+       before graph capture on any device: allocating it inside a graph
+       records a ``wp.zeros`` fill that wipes the warm-start buffers on every
+       replay. Construct :class:`~newton.CollisionPipeline` before
+       :class:`~newton.solvers.SolverVBD` so contact history is pre-allocated,
+       or run one uncaptured solver step before capture. Ordinary contact
+       buffers can still grow on demand during graph capture on CPU and on
+       CUDA with the memory pool enabled; only CUDA capture without a memory
+       pool requires that they also be pre-allocated. Buffer sizes that are
+       too small can drop contacts; sizes that are too large cost memory and
+       performance. Examples commonly tune
+       ``iterations``, particle self-contact radius and margin, particle
+       contact buffers and filters, ``particle_collision_detection_interval``,
+       ``particle_enable_tile_solve``, ``rigid_body_contact_buffer_size``,
        ``rigid_body_particle_contact_buffer_size``, ``rigid_contact_hard``,
        ``rigid_contact_history``, and ``rigid_avbd_contact_alpha``.
    * - :class:`~newton.solvers.SolverFeatherstone`

newton/_src/solvers/kamino/_src/linalg/conjugate.py

@@ -338,7 +338,7 @@ def _run_capturable_loop(
     maxiter: wp.array[int],
     atol_sq: wp.array[Any],
     callback: Callable | None,
-    use_cuda_graph: bool,
+    use_graph: bool,
     use_graph_conditionals: bool = True,
     maxiter_host: int | None = None,
     loop_granularity: int = 1,
@@ -384,7 +384,7 @@ def do_cycle_with_condition():
         if callback_launch is not None:
             callback_launch.launch()
 
-    if use_cuda_graph and device.is_cuda and device.is_capturing:
+    if use_graph and device.is_capturing:
         if use_graph_conditionals:
             wp.capture_while(global_condition, do_cycle_with_condition)
         else:
@@ -541,7 +541,7 @@ class ConjugateSolver(Generic[ScalarType, IndexType]):
         maxiter: Maximum iterations per world. If None, defaults to 1.5 * maxdims.
         Mi: Operator applying the inverse preconditioner M^-1, such that Mi @ A has a smaller condition number than A.
         callback: Optional callback kernel invoked each iteration.
-        use_cuda_graph: Whether to use CUDA graph capture for the solve loop.
+        use_graph: Whether to use graph capture for the solve loop.
         loop_granularity: Number of iterations before termination criteria are checked.
     """
 
@@ -555,7 +555,7 @@ def __init__(
         maxiter: wp.array[wp.int32] | None = None,
         Mi: BatchedLinearOperator[ScalarType, IndexType] | None = None,
         callback: Callable | None = None,
-        use_cuda_graph: bool = True,
+        use_graph: bool = True,
         use_graph_conditionals: bool = True,
         loop_granularity: int = 1,
     ):
@@ -593,7 +593,7 @@ def __init__(
         self.loop_granularity = loop_granularity
 
         self.callback = callback
-        self.use_cuda_graph = use_cuda_graph
+        self.use_graph = use_graph
 
         self.dot_tile_size = min(2048, 2 ** math.ceil(math.log(self.maxdims, 2)))
         self.tiled_dot_kernel = make_dot_kernel(self.dot_tile_size, self.maxdims)
@@ -756,7 +756,7 @@ def solve(
             self.maxiter,
             self.atol_sq,
             self.callback,
-            self.use_cuda_graph,
+            self.use_graph,
             use_graph_conditionals=self.use_graph_conditionals,
             maxiter_host=self.maxiter_host,
             loop_granularity=min(self.loop_granularity, self.maxiter_host),
@@ -890,7 +890,7 @@ def solve(
             self.maxiter,
             self.atol_sq,
             self.callback,
-            self.use_cuda_graph,
+            self.use_graph,
             use_graph_conditionals=self.use_graph_conditionals,
             maxiter_host=self.maxiter_host,
             loop_granularity=min(self.loop_granularity, self.maxiter_host),

newton/_src/solvers/kamino/_src/linalg/linear.py

@@ -776,7 +776,7 @@ def _allocate_impl(self, operator, **kwargs: dict[str, Any]) -> None:
             maxiter=self._maxiter,
             Mi=self._Mi,
             callback=None,
-            use_cuda_graph=True,
+            use_graph=True,
             use_graph_conditionals=self._use_graph_conditionals,
             loop_granularity=self.loop_granularity,
         )
@@ -790,7 +790,7 @@ def _allocate_impl(self, operator, **kwargs: dict[str, Any]) -> None:
                 maxiter=self._maxiter,
                 Mi=self._Mi,
                 callback=None,
-                use_cuda_graph=True,
+                use_graph=True,
                 use_graph_conditionals=self._use_graph_conditionals,
                 loop_granularity=self.loop_granularity,
             )
@@ -888,7 +888,7 @@ def _allocate_impl(self, operator, **kwargs: dict[str, Any]) -> None:
             maxiter=self._maxiter,
             Mi=self._Mi,
             callback=None,
-            use_cuda_graph=True,
+            use_graph=True,
             use_graph_conditionals=self._use_graph_conditionals,
             loop_granularity=self.loop_granularity,
         )
@@ -902,7 +902,7 @@ def _allocate_impl(self, operator, **kwargs: dict[str, Any]) -> None:
                 maxiter=self._maxiter,
                 Mi=self._Mi,
                 callback=None,
-                use_cuda_graph=True,
+                use_graph=True,
                 use_graph_conditionals=self._use_graph_conditionals,
                 loop_granularity=self.loop_granularity,
             )

newton/_src/solvers/kamino/tests/test_linalg_solve_cg.py

@@ -75,7 +75,7 @@ def _test_solve(self, solver_cls, problem_params, device):
             maxiter=maxiter,
             Mi=None,
             callback=None,
-            use_cuda_graph=False,
+            use_graph=False,
         )
         cur_iter, r_norm_sq, atol_sq = solver.solve(b_wp, x_wp)
 
@@ -139,6 +139,98 @@ def test_solve_cr_cuda(self):
             with self.subTest(problem=problem_name, solver=solver_cls.__name__):
                 self._test_solve(solver_cls, problem_params, device)
 
+    def _test_capture_replay_matches_eager(self, solver_cls, device):
+        """Regression: capture-and-replay must match an eager solve.
+
+        Guards against the class of bug where the capture-safe path in
+        ``_run_capturable_loop`` silently under-iterates. Concretely, if the
+        break decision is frozen at record time via a host readback of stale
+        pre-capture memory, the recorded graph bakes in a fixed cycle count
+        and ``cur_iter`` under capture will be far below the eager count.
+        """
+        device = wp.get_device(device)
+        problem = RandomProblemLLT(
+            maxdims=8,
+            dims=[5, 8],
+            seed=self.seed,
+            np_dtype=np.float32,
+            wp_dtype=wp.float32,
+            device=device,
+        )
+        n_worlds = problem.num_blocks
+
+        info = DenseSquareMultiLinearInfo()
+        info.finalize(dimensions=problem.maxdims, dtype=wp.float32, device=device)
+        info.dim = problem.dim_wp
+        operator = DenseLinearOperatorData(info=info, mat=problem.A_wp)
+        A = BatchedLinearOperator.from_dense(operator)
+
+        world_active = wp.full(n_worlds, True, dtype=wp.bool, device=device)
+        maxdim = max(problem.maxdims)
+        atol = wp.full(n_worlds, 1.0e-4, dtype=problem.wp_dtype, device=device)
+        rtol = wp.full(n_worlds, 1.0e-5, dtype=problem.wp_dtype, device=device)
+        maxiter = wp.full(n_worlds, max(3 * maxdim, 50), dtype=int, device=device)
+
+        def new_solver(*, use_graph):
+            return solver_cls(
+                A=A,
+                world_active=world_active,
+                atol=atol,
+                rtol=rtol,
+                maxiter=maxiter,
+                Mi=None,
+                callback=None,
+                use_graph=use_graph,
+            )
+
+        # Eager reference.
+        x_eager = wp.zeros(info.total_vec_size, dtype=wp.float32, device=device)
+        eager_cur, _, _ = new_solver(use_graph=False).solve(problem.b_wp, x_eager)
+
+        # Captured replay. Warm up outside capture so kernels compile before
+        # recording; zero the output between the warmup and the capture body
+        # so both solves see the same initial state.
+        solver = new_solver(use_graph=True)
+        x_capture = wp.zeros(info.total_vec_size, dtype=wp.float32, device=device)
+        solver.solve(problem.b_wp, x_capture)
+        x_capture.zero_()
+        # Warp CPU graph capture currently rejects wp.copy() on non-contiguous
+        # arrays. Both the buggy pre-fix eager path (via strided
+        # ``dot_partial_sums[:, :, 0]`` in ``dot_product``) and the fixed
+        # capture path (via ``rz_old.assign(rz_new)`` in ``do_iteration``)
+        # hit that limitation on CPU, so the parity assertions below only
+        # activate once Warp lifts it. The CUDA counterparts exercise the
+        # same gate on a path where the dot buffer is contiguous and no
+        # skip is needed.
+        try:
+            with wp.ScopedCapture(device) as cap:
+                cap_cur, _, _ = solver.solve(problem.b_wp, x_capture)
+        except NotImplementedError as e:
+            if "non-contiguous" not in str(e):
+                raise
+            self.skipTest(f"Warp graph capture limitation on {device}: {e}")
+        assert cap.graph is not None
+        wp.capture_launch(cap.graph)
+
+        np.testing.assert_array_equal(cap_cur.numpy(), eager_cur.numpy())
+        np.testing.assert_array_equal(x_capture.numpy(), x_eager.numpy())
+
+    def test_capture_replay_cg_cpu(self):
+        self._test_capture_replay_matches_eager(CGSolver, "cpu")
+
+    def test_capture_replay_cr_cpu(self):
+        self._test_capture_replay_matches_eager(CRSolver, "cpu")
+
+    def test_capture_replay_cg_cuda(self):
+        if not wp.get_cuda_devices():
+            self.skipTest("No CUDA devices found")
+        self._test_capture_replay_matches_eager(CGSolver, wp.get_cuda_device())
+
+    def test_capture_replay_cr_cuda(self):
+        if not wp.get_cuda_devices():
+            self.skipTest("No CUDA devices found")
+        self._test_capture_replay_matches_eager(CRSolver, wp.get_cuda_device())
+
     def _test_sparse_solve(self, solver_cls, dims, block_size, device):
         """Test CG/CR with sparse matrices built from random SPD matrices.
 
@@ -223,7 +315,7 @@ def _test_sparse_solve(self, solver_cls, dims, block_size, device):
             maxiter=None,
             Mi=None,
             callback=None,
-            use_cuda_graph=False,
+            use_graph=False,
         )
         solver_dense.solve(b_wp, x_dense)
 
@@ -237,7 +329,7 @@ def _test_sparse_solve(self, solver_cls, dims, block_size, device):
             maxiter=None,
             Mi=None,
             callback=None,
-            use_cuda_graph=False,
+            use_graph=False,
         )
         solver_sparse.solve(b_wp, x_sparse)
 
@@ -335,7 +427,7 @@ def test_sparse_cg_solve_simple(self):
             rtol=rtol,
             maxiter=None,
             Mi=None,
-            use_cuda_graph=False,
+            use_graph=False,
         )
         solver.solve(b_wp, x_wp)
 
@@ -477,7 +569,7 @@ def _test_solve_heterogeneous(self, solver_cls, problem_params, device):
             maxiter=maxiter,
             Mi=None,
             callback=None,
-            use_cuda_graph=False,
+            use_graph=False,
         )
         solver.solve(b, x_wp)
 
@@ -570,7 +662,7 @@ def _test_solve_heterogeneous_jacobi(self, solver_cls, problem_params, device):
             maxiter=maxiter,
             Mi=Mi,
             callback=None,
-            use_cuda_graph=False,
+            use_graph=False,
         )
         solver.solve(b, x_wp)
 

newton/_src/solvers/vbd/solver_vbd.py

@@ -22,6 +22,7 @@
     State,
     StateFlags,
 )
+from ...utils import is_graph_capture_allocation_enabled
 from ...utils.deprecation import deprecate_nonkeyword_arguments
 from ..coupled.interface import CouplingInterface
 from ..solver import SolverBase
@@ -134,13 +135,16 @@ class SolverVBD(SolverBase, CouplingInterface):
     Buffer sizing:
         SolverVBD pre-allocates contact state from capacities populated by
         :class:`~newton.CollisionPipeline` when available; otherwise, the first
-        :meth:`step` lazily sizes buffers from ``Contacts``. During CUDA graph
-        recording, ordinary lazy resizing is supported only when Warp's memory pool
-        is enabled; otherwise, the solver raises with guidance to pre-size before
-        capture. Rigid contact history must be allocated before capture regardless
-        of memory-pool support. With ``rigid_contact_history=True``, construct
-        :class:`~newton.CollisionPipeline` before ``SolverVBD``, or run one
-        uncaptured solver step before capture.
+        :meth:`step` lazily sizes buffers from ``Contacts``. During graph capture,
+        ordinary lazy resizing is supported on CPU and on CUDA with Warp's
+        stream-ordered memory pool enabled; otherwise the solver raises with
+        guidance to pre-size before capture. Rigid contact history is
+        cross-replay-persistent state, so it must always be allocated before
+        capture regardless of the device's allocation-during-capture support --
+        allocating it inside a graph records a `wp.zeros` fill that wipes the
+        warm-start buffers on every replay. With ``rigid_contact_history=True``,
+        construct :class:`~newton.CollisionPipeline` before ``SolverVBD``, or run
+        one uncaptured solver step before capture.
 
     References:
         - Anka He Chen, Ziheng Liu, Yin Yang, and Cem Yuksel. 2024. Vertex Block Descent. ACM Trans. Graph. 43, 4, Article 116 (July 2024), 16 pages.
@@ -341,7 +345,7 @@ def __init__(
                 Requires contacts with ``rigid_contact_match_index`` populated; use
                 ``CollisionPipeline(contact_matching="latest")`` for VBD warm-starting. Ignored
                 when ``integrate_with_external_rigid_solver=True`` or ``model.body_count == 0``.
-                For CUDA graph capture, construct :class:`~newton.CollisionPipeline` before
+                During graph capture, construct :class:`~newton.CollisionPipeline` before
                 ``SolverVBD`` so history is pre-allocated, or run one uncaptured solver step
                 before capture.
             rigid_contact_stick_motion_eps: Tangential contact residual threshold for marking hard
@@ -1180,8 +1184,6 @@ def _init_rigid_contact_warmstart(self, rigid_contact_max: int) -> None:
         self._prev_contact_normal = wp.zeros(cap, dtype=wp.vec3, device=self.device)
 
     def _raise_if_capturing_resize(self, name: str, current: int, required: int) -> None:
-        from ...utils import is_graph_capture_allocation_enabled  # noqa: PLC0415
-
         if self.device.is_capturing and not is_graph_capture_allocation_enabled(self.device):
             raise RuntimeError(
                 f"SolverVBD {name} buffer needs to grow from {current} to {required} "
@@ -1809,7 +1811,7 @@ def step(
 
         Raises:
             RuntimeError: If required rigid contact-matching data is unavailable, or contact-history storage would
-                need to be allocated or grown during CUDA graph capture.
+                need to be allocated or grown during graph capture.
         """
         self._apply_module_options()
         update_rigid = self._update_rigid_history
@@ -2141,11 +2143,17 @@ def _initialize_rigid_bodies(
         internal_rigid = model.body_count > 0 and not self.integrate_with_external_rigid_solver
         rigid_capacity = contacts.rigid_contact_max if contacts is not None else 0
 
+        # Rigid contact history is cross-replay-persistent state: allocating it
+        # during capture records a `wp.zeros` fill into the graph, which then
+        # re-zeros the warm-start buffers on every replay -- silently
+        # equivalent to `rigid_contact_history=False`. So this guard fires
+        # unconditionally when capturing, regardless of the device's
+        # allocation-during-capture support.
         if self.device.is_capturing and internal_rigid and self.rigid_contact_history:
             history_capacity = 0 if self._prev_contact_lambda is None else self._prev_contact_lambda.shape[0]
             if history_capacity < rigid_capacity:
                 raise RuntimeError(
-                    "SolverVBD contact history must be allocated before CUDA graph capture. "
+                    "SolverVBD contact history must be allocated before graph capture. "
                     "Construct CollisionPipeline before SolverVBD, or run one uncaptured solver step before capture."
                 )
 

newton/examples/cable/example_cable_cross_slide_table.py

@@ -775,7 +775,7 @@ def __init__(self, viewer, args):
         self.control = self.model.control()
         self.contacts = self.collision_pipeline.contacts()
 
-        # Device arrays used by kernels during simulation and CUDA graph replay.
+        # Device arrays used by kernels during simulation and captured replay.
         self.kinematic_body_indices = wp.array(
             kinematic_body_indices,
             dtype=wp.int32,
@@ -834,13 +834,10 @@ def __init__(self, viewer, args):
         self.capture()
 
     def capture(self):
-        """Capture the simulation update when running on CUDA."""
-        if self.solver.device.is_cuda:
-            with wp.ScopedCapture() as capture:
-                self.simulate()
-            self.graph = capture.graph
-        else:
-            self.graph = None
+        """Capture the simulation update into a graph for replay."""
+        with wp.ScopedCapture() as capture:
+            self.simulate()
+        self.graph = capture.graph
 
     def simulate(self):
         """Advance the XY table simulation by one rendered frame."""