[Draft/RFC] Native-CUDA (cuda-oxide) backend: device build, fixed-grid dispatch, CUDA-graph capture

Cargo.toml

@@ -26,7 +26,7 @@ exclude = ["**/.DS_Store"]
 [workspace.dependencies]
 glamx = { version = "0.3", default-features = false, features = ["nostd-libm", "bytemuck"] }
 include_dir = "0.7"
-khal-std = "0.2"
+khal-std = { version = "0.2", default-features = false }
 khal = { version = "0.2", features = ["derive"]}
 vortx = { version = "0.3", features = ["unsafe_remove_boundchecks"] }
 bytemuck = { version = "1", features = ["derive"] }
@@ -110,9 +110,9 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = [
 #rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" }
 #rapier3d-urdf = { path = "../rapier/crates/rapier3d-urdf" }
 #rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" }
-#khal = { path = "../khal/crates/khal" }
-#khal-derive = { path = "../khal/crates/khal-derive" }
-#khal-std = { path = "../khal/crates/khal-std" }
-#khal-builder = { path = "../khal/crates/khal-builder" }
-#vortx = { path = "../vortx" }
-#vortx-shaders = { path = "../vortx/vortx-shaders" }
+khal = { path = "../khal/crates/khal" }
+khal-derive = { path = "../khal/crates/khal-derive" }
+khal-std = { path = "../khal/crates/khal-std" }
+khal-builder = { path = "../khal/crates/khal-builder" }
+vortx = { path = "../vortx" }
+vortx-shaders = { path = "../vortx/vortx-shaders" }

crates/nexus_python3d/src/nexus.rs

@@ -345,6 +345,31 @@ impl NexusPipeline {
             .map_err(gpu_err)
     }
 
+    /// Captures one frame's rigid-body step sequence
+    /// (`rbd_steps_per_frame` solver steps) into a CUDA graph, executing it
+    /// once. Subsequent `replay_cuda_graph` calls replay the whole sequence
+    /// with a single `cuGraphLaunch` — the fast path for capture/eval loops.
+    ///
+    /// Returns `False` when the backend is not CUDA. Call after the scene is
+    /// finalized and a few warmup `simulate` calls (the graph records raw
+    /// buffer addresses; buffer growth after capture invalidates it).
+    #[cfg(feature = "cuda")]
+    fn capture_cuda_graph(
+        &mut self,
+        viewer: PyRef<NexusViewer>,
+        mut state: PyRefMut<NexusState>,
+    ) -> PyResult<bool> {
+        pollster::block_on(self.0.capture_rbd_graph(viewer.backend(), &mut state.0))
+            .map_err(gpu_err)
+    }
+
+    /// Replays the captured rigid-body CUDA graph (see `capture_cuda_graph`).
+    /// Returns `False` when no graph has been captured.
+    #[cfg(feature = "cuda")]
+    fn replay_cuda_graph(&mut self) -> PyResult<bool> {
+        self.0.replay_rbd_graph().map_err(gpu_err)
+    }
+
     /// Advances the simulation by one frame. Blocks on the async GPU work.
     #[pyo3(signature = (viewer, state, timestamps=None))]
     fn simulate(

crates/nexus_rbd_shaders3d/Cargo.toml

@@ -29,6 +29,7 @@ push_constants = []
 cpu = []
 cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"]
 cuda = []
+cuda-oxide = ["khal-std/cuda-oxide", "vortx-shaders/cuda-oxide"]
 
 [dependencies]
 vortx-shaders = { workspace = true }

src/pipeline.rs

@@ -13,6 +13,12 @@ bitflags::bitflags! {
 #[derive(Default)]
 pub struct NexusPipeline {
     pub rbd_pipeline: Option<RbdPipeline>,
+    /// CUDA-graph capture of one frame's rigid-body step sequence
+    /// (`rbd_steps_per_frame × step`). Replaying it costs a single
+    /// `cuGraphLaunch` instead of re-encoding every kernel dispatch from the
+    /// host. See [`Self::capture_rbd_graph`].
+    #[cfg(feature = "cuda")]
+    rbd_graph: Option<khal::backend::cuda::CapturedGraph>,
 }
 
 impl NexusPipeline {
@@ -36,6 +42,71 @@ impl NexusPipeline {
     /// In addition, resources are loaded lazily on the GPU, so the first step
     /// after inserting/removing entities can be slower too. Call `Self::finalize`
     /// to pay that cost upfront.
+    /// Captures one frame's rigid-body dispatch sequence
+    /// (`rbd_steps_per_frame × step`) into a CUDA graph and stores it on the
+    /// pipeline. Returns `false` (and captures nothing) when the backend is
+    /// not CUDA or there is no rigid-body state.
+    ///
+    /// Requirements: call after the scene is final (`finalize` runs here) and
+    /// after a few warmup `simulate` calls so the buffer sizes are stable —
+    /// the graph records raw buffer addresses, so any later reallocation
+    /// (e.g. `auto_resize_buffers` growth) invalidates it. Fixed-grid dispatch
+    /// must be active (it is the CUDA default): indirect dispatches read
+    /// counts on the host and cannot be captured.
+    ///
+    /// Timestamps and buffer auto-resize are skipped during capture and
+    /// replay; `run_stats` stops updating while replaying.
+    #[cfg(feature = "cuda")]
+    pub async fn capture_rbd_graph(
+        &mut self,
+        backend: &GpuBackend,
+        state: &mut NexusState,
+    ) -> Result<bool, GpuBackendError> {
+        state.finalize(backend).await?;
+        use khal::backend::Backend as _;
+        let Some(cuda) = backend.as_cuda() else {
+            return Ok(false);
+        };
+        let Some(rbd) = state.rbd.as_mut() else {
+            return Ok(false);
+        };
+        self.preload_pipelines(backend, NexusPipelineMask::RBD)?;
+        let pipeline = self.rbd_pipeline.as_mut().unwrap_or_else(|| unreachable!());
+        let steps = state.rbd_steps_per_frame.max(1);
+        cuda.begin_capture().map_err(khal::backend::GpuBackendError::Cuda)?;
+        let mut step_result = Ok(state.run_stats.clone());
+        for _ in 0..steps {
+            step_result = pipeline.step(backend, rbd, None);
+            if step_result.is_err() {
+                break;
+            }
+        }
+        // Always end the capture, even on error, so the stream isn't left in
+        // capture mode.
+        let graph = cuda.end_capture().map_err(khal::backend::GpuBackendError::Cuda)?;
+        state.run_stats = step_result?;
+        graph.upload().map_err(khal::backend::GpuBackendError::Cuda)?;
+        // Capture only records; execute the captured sequence once so this
+        // call has the same effect as a `simulate`.
+        graph.launch().map_err(khal::backend::GpuBackendError::Cuda)?;
+        self.rbd_graph = Some(graph);
+        Ok(true)
+    }
+
+    /// Replays the captured rigid-body graph (one `cuGraphLaunch` for the whole
+    /// `rbd_steps_per_frame × step` sequence). Returns `false` when no graph
+    /// has been captured.
+    #[cfg(feature = "cuda")]
+    pub fn replay_rbd_graph(&self) -> Result<bool, GpuBackendError> {
+        match &self.rbd_graph {
+            Some(g) => {
+                g.launch().map_err(khal::backend::GpuBackendError::Cuda)?;
+                Ok(true)
+            }
+            None => Ok(false),
+        }
+    }
+
     pub async fn simulate(
         &mut self,
         backend: &GpuBackend,

src_rbd/broad_phase/narrow_phase.rs

@@ -50,12 +50,20 @@ impl GpuNarrowPhase {
         collider_materials: &Tensor<crate::shaders::queries::ColliderMaterial>,
     ) -> Result<(), GpuBackendError> {
         let num_batches = contacts_len.len() as u32;
+        // Capacity-based grids for fixed-grid dispatch (see `crate::dispatch_grid`).
+        let nb = num_batches.max(1);
+        let pairs_grid = [
+            (collision_pairs.len() as u32 / nb).max(1).div_ceil(64),
+            num_batches,
+            1,
+        ];
+        let pfm_grid = [(pfm_pairs.len() as u32 / nb).max(1).div_ceil(64), num_batches, 1];
         self.reset_narrow_phase
             .call(pass, [1u32, num_batches, 1], contacts_len, pfm_pairs_len)?;
 
         self.narrow_phase.call(
             pass,
-            collision_pairs_indirect,
+            crate::dispatch_grid(collision_pairs_indirect, pairs_grid),
             collision_pairs,
             collision_pairs_len,
             poses,
@@ -71,7 +79,7 @@ impl GpuNarrowPhase {
         // separate dispatch so each pass fits 8 storage buffers).
         self.narrow_phase_deferred.call(
             pass,
-            collision_pairs_indirect,
+            crate::dispatch_grid(collision_pairs_indirect, pairs_grid),
             collision_pairs,
             collision_pairs_len,
             poses,
@@ -87,7 +95,7 @@ impl GpuNarrowPhase {
             .call(pass, 1u32, pfm_pairs_len, pfm_pairs_indirect)?;
         self.narrow_phase_pfm_pfm.call(
             pass,
-            &*pfm_pairs_indirect,
+            crate::dispatch_grid(&*pfm_pairs_indirect, pfm_grid),
             contacts,
             contacts_len,
             pfm_pairs,

src_rbd/dynamics/coloring.rs

@@ -69,6 +69,18 @@ pub struct ColoringArgs<'a> {
     pub body_group: &'a Tensor<u32>,
 }
 
+
+/// Capacity-based dispatch grid for the coloring kernels (used when
+/// fixed-grid dispatch replaces the indirect buffer; see `crate::dispatch_grid`).
+fn coloring_grid(args: &ColoringArgs) -> [u32; 3] {
+    let nb = (args.contacts_len.len() as u32).max(1);
+    [
+        (args.constraints.len() as u32 / nb).max(1).div_ceil(64),
+        nb,
+        1,
+    ]
+}
+
 impl GpuColoring {
     /// Dispatches the reset_luby kernel.
     fn dispatch_reset_luby(
@@ -78,7 +90,7 @@ impl GpuColoring {
     ) -> Result<(), GpuBackendError> {
         self.reset_luby_kernel.call(
             pass,
-            args.contacts_len_indirect,
+            crate::dispatch_grid(args.contacts_len_indirect, coloring_grid(args)),
             args.constraints_colors,
             args.constraints_rands,
             args.contacts_len,
@@ -95,7 +107,7 @@ impl GpuColoring {
     ) -> Result<(), GpuBackendError> {
         self.step_graph_coloring_luby_kernel.call(
             pass,
-            args.contacts_len_indirect,
+            crate::dispatch_grid(args.contacts_len_indirect, coloring_grid(args)),
             args.body_constraint_counts,
             args.body_constraint_ids,
             args.constraints,
@@ -118,7 +130,7 @@ impl GpuColoring {
     ) -> Result<(), GpuBackendError> {
         self.reset_topo_gc_kernel.call(
             pass,
-            args.contacts_len_indirect,
+            crate::dispatch_grid(args.contacts_len_indirect, coloring_grid(args)),
             args.constraints_colors,
             args.colored,
             args.contacts_len,
@@ -135,7 +147,7 @@ impl GpuColoring {
     ) -> Result<(), GpuBackendError> {
         self.step_graph_coloring_topo_gc_kernel.call(
             pass,
-            args.contacts_len_indirect,
+            crate::dispatch_grid(args.contacts_len_indirect, coloring_grid(args)),
             args.body_constraint_counts,
             args.body_constraint_ids,
             args.constraints,
@@ -157,7 +169,7 @@ impl GpuColoring {
     ) -> Result<(), GpuBackendError> {
         self.fix_conflicts_topo_gc_kernel.call(
             pass,
-            args.contacts_len_indirect,
+            crate::dispatch_grid(args.contacts_len_indirect, coloring_grid(args)),
             args.body_constraint_counts,
             args.body_constraint_ids,
             args.constraints,

src_rbd/dynamics/solver.rs

@@ -58,6 +58,19 @@ pub struct GpuSolver {
 
 /// Arguments for constraint solver dispatch, used by [`GpuSolver::prepare`] and
 /// [`GpuSolver::solve_tgs`].
+
+/// Capacity-based dispatch grid for the contact-count-driven solver kernels
+/// (used when fixed-grid dispatch replaces the indirect buffer; see
+/// `crate::dispatch_grid`).
+fn contacts_grid(args: &SolverArgs) -> [u32; 3] {
+    let nb = args.num_batches.max(1);
+    [
+        (args.contacts.len() as u32 / nb).max(1).div_ceil(64),
+        args.num_batches.max(1),
+        1,
+    ]
+}
+
 pub struct SolverArgs<'a> {
     /// Total number of colors from graph coloring.
     pub num_colors: u32,
@@ -166,7 +179,7 @@ impl GpuSolver {
 
         self.init_constraints.call(
             pass,
-            args.contacts_len_indirect,
+            crate::dispatch_grid(args.contacts_len_indirect, contacts_grid(&args)),
             args.contacts,
             args.constraints,
             args.constraint_builders,
@@ -182,7 +195,7 @@ impl GpuSolver {
         // build pass above stays within 8 storage buffers.
         self.count_constraints.call(
             pass,
-            args.contacts_len_indirect,
+            crate::dispatch_grid(args.contacts_len_indirect, contacts_grid(&args)),
             args.contacts,
             args.body_constraint_counts,
             args.body_group,
@@ -200,7 +213,7 @@ impl GpuSolver {
 
         self.sort_constraints.call(
             pass,
-            args.contacts_len_indirect,
+            crate::dispatch_grid(args.contacts_len_indirect, contacts_grid(&args)),
             args.body_constraint_counts,
             args.mprops,
             args.contacts,
@@ -307,7 +320,7 @@ impl GpuSolver {
             mb_phase!(substep_build_constraints);
             self.update_constraints.call(
                 pass,
-                args.contacts_len_indirect,
+                crate::dispatch_grid(args.contacts_len_indirect, contacts_grid(&args)),
                 args.constraints,
                 args.constraint_builders,
                 args.contacts_len,
@@ -320,7 +333,7 @@ impl GpuSolver {
             for _ in 0..args.num_colors {
                 self.warmstart.call(
                     pass,
-                    args.contacts_len_indirect,
+                    crate::dispatch_grid(args.contacts_len_indirect, contacts_grid(&args)),
                     args.constraints,
                     args.solver_vels,
                     args.constraints_colors,
@@ -340,7 +353,7 @@ impl GpuSolver {
             for _ in 0..args.num_colors {
                 self.step_gauss_seidel.call(
                     pass,
-                    args.contacts_len_indirect,
+                    crate::dispatch_grid(args.contacts_len_indirect, contacts_grid(&args)),
                     args.constraints,
                     args.solver_vels,
                     args.constraints_colors,
@@ -371,7 +384,7 @@ impl GpuSolver {
             joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?;
             self.remove_cfm_and_bias_kernel.call(
                 pass,
-                args.contacts_len_indirect,
+                crate::dispatch_grid(args.contacts_len_indirect, contacts_grid(&args)),
                 args.constraints,
                 args.contacts_len,
                 args.batch_indices,
@@ -380,7 +393,7 @@ impl GpuSolver {
             for _ in 0..args.num_colors {
                 self.step_gauss_seidel.call(
                     pass,
-                    args.contacts_len_indirect,
+                    crate::dispatch_grid(args.contacts_len_indirect, contacts_grid(&args)),
                     args.constraints,
                     args.solver_vels,
                     args.constraints_colors,

src_rbd/dynamics/warmstart.rs

@@ -49,9 +49,15 @@ impl GpuWarmstart {
         pass: &mut GpuPass,
         args: WarmstartArgs<'a>,
     ) -> Result<(), GpuBackendError> {
+        let nb = (args.contacts_len.len() as u32).max(1);
+        let ws_grid = [
+            (args.new_constraints.len() as u32 / nb).max(1).div_ceil(64),
+            nb,
+            1,
+        ];
         self.transfer_warmstart_impulses_kernel.call(
             pass,
-            args.contacts_len_indirect,
+            crate::dispatch_grid(args.contacts_len_indirect, ws_grid),
             args.old_body_constraint_counts,
             args.old_body_constraint_ids,
             args.old_constraints,