[BUG] CUDA graph launch when using VBD Integrator with handle_self_contact=True

[shi-eric]

Bug Description

The following code snippet (simplified from the new example_cloth_self_contact.py will fail in the CUDA error verification call if the wp.sim.VBDIntegrator is using handle_self_contact=True:

import os

import numpy as np
from pxr import Usd, UsdGeom

import warp as wp
import warp.examples
import warp.sim

usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "square_cloth.usd"))
usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/cloth/cloth"))

mesh_points = np.array(usd_geom.GetPointsAttr().Get())
mesh_indices = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())

builder = wp.sim.ModelBuilder()
builder.add_cloth_mesh(
    pos=wp.vec3(0.0, 0.0, 0.0),
    rot=wp.quat_identity(),
    scale=1.0,
    vertices=[wp.vec3(v) for v in mesh_points],
    indices=mesh_indices,
    vel=wp.vec3(0.0, 0.0, 0.0),
    density=0.02,
    tri_ke=1.0e5,
    tri_ka=1.0e5,
    tri_kd=3.0e-5,
)
builder.color()
model = builder.finalize()

integrator = wp.sim.VBDIntegrator(
    model,
    iterations=10,
    handle_self_contact=True,
)
state0 = model.state()
state1 = model.state()

dt = 1 / 60.0

with wp.ScopedCapture() as capture:
    integrator.simulate(model, state0, state1, dt)
graph = capture.graph

wp.capture_launch(graph)

wp.synchronize_device()
wp.context.runtime.verify_cuda_device(wp.get_device())

The output is:

Warp 1.5.1 initialized:
   CUDA Toolkit 12.6, Driver 12.6
   Devices:
     "cpu"      : "x86_64"
     "cuda:0"   : "NVIDIA GeForce RTX 3090" (24 GiB, sm_86, mempool enabled)
     "cuda:1"   : "NVIDIA GeForce RTX 3090" (24 GiB, sm_86, mempool enabled)
   CUDA peer access:
     Supported fully (all-directional)
   Kernel cache:
     /home/eshi/.cache/warp/1.5.1
Warp UserWarning: The graph is not optimizable anymore, terminated with a max/min ratio: 2.0 without reaching the target ratio: 1.1
Module warp.sim.graph_coloring f690f7b load on device 'cpu' took 1.28 ms  (cached)
Module warp.sim.integrator_vbd 41ba571 load on device 'cpu' took 0.31 ms  (cached)
Module warp.sim.collide 05171b4 load on device 'cuda:0' took 7.25 ms  (cached)
Module warp.sim.integrator_vbd 41ba571 load on device 'cuda:0' took 2.64 ms  (cached)
Warp CUDA error 700: an illegal memory access was encountered (in function cuda_context_synchronize, /warp/warp/native/warp.cu:1983)
Warp CUDA error 700: an illegal memory access was encountered (in function free_device_async, /warp/warp/native/warp.cu:636)

There will not be an error if handle_self_contact=False.

System Information

Ubuntu 24.04.01 LTS