feat: add support for cuda-oxide and cuda-graphs

Cargo.toml

@@ -88,11 +88,20 @@ nexus_viewer3d = { version = "0.5.0", path = "crates/nexus_viewer3d" }
 
 [workspace.lints]
 rust.unexpected_cfgs = { level = "warn", check-cfg = [
-    'cfg(feature, values("dim2", "dim3", "cpu", "cuda", "metal"))',
+    'cfg(feature, values("dim2", "dim3", "cpu", "cuda", "cuda-oxide", "metal"))',
     'cfg(target_arch_is_gpu)',
 ] }
 
 [patch.crates-io]
+# khal and vortx from GitHub: the `cuda-oxide` feature, the WebGPU SPIR-V
+# passthrough entry-point fix, and compute-graph capture are not on crates.io
+# yet. Drop once khal 0.3.1 / vortx 0.4.1 are published.
+khal = { git = "https://github.com/dimforge/khal", branch = "main" }
+khal-derive = { git = "https://github.com/dimforge/khal", branch = "main" }
+khal-std = { git = "https://github.com/dimforge/khal", branch = "main" }
+khal-builder = { git = "https://github.com/dimforge/khal", branch = "main" }
+vortx = { git = "https://github.com/dimforge/vortx", branch = "main" }
+vortx-shaders = { git = "https://github.com/dimforge/vortx", branch = "main" }
 #rapier2d = { path = "../rapier/crates/rapier2d" }
 #rapier3d = { path = "../rapier/crates/rapier3d" }
 #rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" }

crates/examples2d/Cargo.toml

@@ -10,6 +10,7 @@ default = []
 cpu = ["nexus_viewer2d/cpu"]
 cpu-parallel = ["cpu", "nexus_viewer2d/cpu-parallel"]
 cuda = ["nexus_viewer2d/cuda"]
+cuda-oxide = ["cuda", "nexus_viewer2d/cuda-oxide"]
 metal = ["nexus_viewer2d/metal"]
 
 [dependencies]

crates/examples2d/all_examples2.rs

@@ -80,6 +80,8 @@ struct CliOptions {
     example: Option<String>,
     list: bool,
     cpu: bool,
+    cuda: bool,
+    compute_graphs: bool,
     metal: bool,
     run: bool,
 }
@@ -90,6 +92,8 @@ fn parse_command_line() -> CliOptions {
         example: None,
         list: false,
         cpu: false,
+        cuda: false,
+        compute_graphs: false,
         metal: false,
         run: false,
     };
@@ -99,6 +103,8 @@ fn parse_command_line() -> CliOptions {
             "--example" => opts.example = args.next(),
             "--list" => opts.list = true,
             "--cpu" => opts.cpu = true,
+            "--cuda" => opts.cuda = true,
+            "--compute-graphs" => opts.compute_graphs = true,
             "--metal" => opts.metal = true,
             "--run" => opts.run = true,
             _ => {}
@@ -139,6 +145,13 @@ pub async fn main() {
     if opts.cpu {
         viewer = viewer.with_cpu();
     }
+    #[cfg(feature = "cuda")]
+    if opts.cuda {
+        viewer = viewer.with_backend(nexus_viewer2d::BackendType::Cuda);
+    }
+    if opts.compute_graphs {
+        viewer = viewer.with_compute_graphs(true);
+    }
     #[cfg(feature = "metal")]
     if opts.metal {
         viewer = viewer.with_backend(nexus_viewer2d::BackendType::Metal);

crates/examples3d/Cargo.toml

@@ -10,6 +10,7 @@ default = []
 cpu = ["nexus_viewer3d/cpu"]
 cpu-parallel = ["cpu", "nexus_viewer3d/cpu-parallel"]
 cuda = ["nexus_viewer3d/cuda"]
+cuda-oxide = ["cuda", "nexus_viewer3d/cuda-oxide"]
 metal = ["nexus_viewer3d/metal"]
 
 [dependencies]

crates/examples3d/all_examples3.rs

@@ -115,6 +115,7 @@ struct CliOptions {
     list: bool,
     cpu: bool,
     cuda: bool,
+    compute_graphs: bool,
     metal: bool,
     run: bool,
 }
@@ -126,6 +127,7 @@ fn parse_command_line() -> CliOptions {
         list: false,
         cpu: false,
         cuda: false,
+        compute_graphs: false,
         metal: false,
         run: false,
     };
@@ -136,6 +138,7 @@ fn parse_command_line() -> CliOptions {
             "--list" => opts.list = true,
             "--cpu" => opts.cpu = true,
             "--cuda" => opts.cuda = true,
+            "--compute-graphs" => opts.compute_graphs = true,
             "--metal" => opts.metal = true,
             "--run" => opts.run = true,
             _ => {}
@@ -182,6 +185,9 @@ pub async fn main() {
     if opts.cuda {
         viewer = viewer.with_backend(nexus_viewer3d::BackendType::Cuda);
     }
+    if opts.compute_graphs {
+        viewer = viewer.with_compute_graphs(true);
+    }
     #[cfg(feature = "metal")]
     if opts.metal {
         viewer = viewer.with_backend(nexus_viewer3d::BackendType::Metal);

crates/nexus2d/Cargo.toml

@@ -28,6 +28,7 @@ metal = ["nexus_rbd2d?/metal", "nexus_mpm2d?/metal"]
 cpu = ["nexus_rbd2d?/cpu", "nexus_mpm2d?/cpu"]
 cpu-parallel = ["cpu", "nexus_rbd2d?/cpu-parallel", "nexus_mpm2d?/cpu-parallel"]
 cuda = ["nexus_rbd2d?/cuda", "nexus_mpm2d?/cuda"]
+cuda-oxide = ["cuda", "nexus_rbd2d?/cuda-oxide", "nexus_mpm2d?/cuda-oxide"]
 
 rbd = ["dep:nexus_rbd2d"]
 mpm = ["dep:nexus_mpm2d"]

crates/nexus3d/Cargo.toml

@@ -28,6 +28,7 @@ metal = ["nexus_rbd3d?/metal", "nexus_mpm3d?/metal"]
 cpu = ["nexus_rbd3d?/cpu", "nexus_mpm3d?/cpu"]
 cpu-parallel = ["cpu", "nexus_rbd3d?/cpu-parallel", "nexus_mpm3d?/cpu-parallel"]
 cuda = ["nexus_rbd3d?/cuda", "nexus_mpm3d?/cuda"]
+cuda-oxide = ["cuda", "nexus_rbd3d?/cuda-oxide", "nexus_mpm3d?/cuda-oxide"]
 
 rbd = ["dep:nexus_rbd3d"]
 mpm = ["dep:nexus_mpm3d"]

crates/nexus_mpm2d/Cargo.toml

@@ -29,6 +29,10 @@ metal = ["khal/metal", "nexus_rbd2d/metal"]
 cpu = ["nexus_mpm_shaders2d/cpu", "nexus_rbd2d/cpu", "vortx/cpu"]
 cpu-parallel = ["cpu", "nexus_mpm_shaders2d/cpu-parallel", "nexus_rbd2d/cpu-parallel", "vortx/cpu-parallel"]
 cuda = ["khal/cuda", "khal-builder/cuda", "nexus_mpm_shaders2d/cuda"]
+# Native CUDA with kernels compiled by cuda-oxide instead of rust-cuda (implies `cuda`).
+# Not forwarded to the shader crate: khal-builder enables its `cuda-oxide` in
+# the separate `cargo oxide` build; on the host it would break the CPU backend.
+cuda-oxide = ["cuda", "khal-builder/cuda-oxide"]
 
 [dependencies]
 nexus_mpm_shaders2d = { workspace = true }

crates/nexus_mpm3d/Cargo.toml

@@ -29,6 +29,10 @@ metal = ["khal/metal", "nexus_rbd3d/metal"]
 cpu = ["nexus_mpm_shaders3d/cpu", "nexus_rbd3d/cpu", "vortx/cpu"]
 cpu-parallel = ["cpu", "nexus_mpm_shaders3d/cpu-parallel", "nexus_rbd3d/cpu-parallel", "vortx/cpu-parallel"]
 cuda = ["khal/cuda", "khal-builder/cuda", "nexus_mpm_shaders3d/cuda"]
+# Native CUDA with kernels compiled by cuda-oxide instead of rust-cuda (implies `cuda`).
+# Not forwarded to the shader crate: khal-builder enables its `cuda-oxide` in
+# the separate `cargo oxide` build; on the host it would break the CPU backend.
+cuda-oxide = ["cuda", "khal-builder/cuda-oxide"]
 
 [dependencies]
 nexus_mpm_shaders3d = { workspace = true }

crates/nexus_mpm_shaders2d/Cargo.toml

@@ -32,6 +32,8 @@ web-compat = ["nexus_rbd_shaders2d/web-compat"]
 cpu = []
 cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"]
 cuda = []
+# cuda-oxide compiler for the CUDA kernels (implies `cuda`).
+cuda-oxide = ["cuda", "khal-std/cuda-oxide", "vortx-shaders/cuda-oxide", "nexus_rbd_shaders2d/cuda-oxide"]
 
 [dependencies]
 nexus_rbd_shaders2d = { workspace = true }

crates/nexus_mpm_shaders3d/Cargo.toml

@@ -31,6 +31,8 @@ web-compat = ["nexus_rbd_shaders3d/web-compat"]
 cpu = []
 cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"]
 cuda = []
+# cuda-oxide compiler for the CUDA kernels (implies `cuda`).
+cuda-oxide = ["cuda", "khal-std/cuda-oxide", "vortx-shaders/cuda-oxide", "nexus_rbd_shaders3d/cuda-oxide"]
 
 [dependencies]
 nexus_rbd_shaders3d = { workspace = true }

crates/nexus_python3d/Cargo.toml

@@ -29,6 +29,7 @@ default = ["webgpu", "extension-module"]
 webgpu = ["nexus_viewer3d/webgpu"]
 metal = ["nexus_viewer3d/metal"]
 cuda = ["nexus_viewer3d/cuda"]
+cuda-oxide = ["cuda", "nexus_viewer3d/cuda-oxide"]
 cpu = ["nexus_viewer3d/cpu"]
 cpu-parallel = ["nexus_viewer3d/cpu-parallel"]
 # Build as a real Python extension module (no libpython link).

crates/nexus_python3d/src/nexus.rs

@@ -472,6 +472,12 @@ impl NexusState {
         self.0.set_rbd_steps_per_frame(steps);
     }
 
+    /// Replays each frame's physics through a compute graph on backends that
+    /// support graph capture (ignored elsewhere).
+    fn set_compute_graphs_enabled(&mut self, enabled: bool) {
+        self.0.set_compute_graphs_enabled(enabled);
+    }
+
     fn set_rbd_gravity(&mut self, viewer: PyRef<NexusViewer>, gravity: Vec3) {
         self.0
             .set_rbd_gravity(viewer.backend(), [gravity.0.x, gravity.0.y, gravity.0.z]);

crates/nexus_rbd2d/Cargo.toml

@@ -29,6 +29,10 @@ metal = ["khal/metal"]
 cpu = ["nexus_rbd_shaders2d/cpu", "vortx/cpu"]
 cpu-parallel = ["cpu", "vortx/cpu-parallel", "nexus_rbd_shaders2d/cpu-parallel"]
 cuda = ["khal/cuda", "khal-builder/cuda", "nexus_rbd_shaders2d/cuda"]
+# Native CUDA with kernels compiled by cuda-oxide instead of rust-cuda (implies `cuda`).
+# Not forwarded to the shader crate: khal-builder enables its `cuda-oxide` in
+# the separate `cargo oxide` build; on the host it would break the CPU backend.
+cuda-oxide = ["cuda", "khal-builder/cuda-oxide"]
 
 [dependencies]
 nexus_rbd_shaders2d = { workspace = true }

crates/nexus_rbd3d/Cargo.toml

@@ -29,6 +29,10 @@ metal = ["khal/metal"]
 cpu = ["nexus_rbd_shaders3d/cpu", "vortx/cpu"]
 cpu-parallel = ["cpu", "vortx/cpu-parallel", "nexus_rbd_shaders3d/cpu-parallel"]
 cuda = ["khal/cuda", "khal-builder/cuda", "nexus_rbd_shaders3d/cuda"]
+# Native CUDA with kernels compiled by cuda-oxide instead of rust-cuda (implies `cuda`).
+# Not forwarded to the shader crate: khal-builder enables its `cuda-oxide` in
+# the separate `cargo oxide` build; on the host it would break the CPU backend.
+cuda-oxide = ["cuda", "khal-builder/cuda-oxide"]
 
 [dependencies]
 nexus_rbd_shaders3d = { workspace = true }

crates/nexus_rbd_shaders2d/Cargo.toml

@@ -34,6 +34,8 @@ web-compat = []
 cpu = []
 cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"]
 cuda = []
+# cuda-oxide compiler for the CUDA kernels (implies `cuda`).
+cuda-oxide = ["cuda", "khal-std/cuda-oxide", "vortx-shaders/cuda-oxide"]
 
 [dependencies]
 vortx-shaders = { workspace = true }

crates/nexus_rbd_shaders3d/Cargo.toml

@@ -34,6 +34,8 @@ web-compat = []
 cpu = []
 cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"]
 cuda = []
+# cuda-oxide compiler for the CUDA kernels (implies `cuda`).
+cuda-oxide = ["cuda", "khal-std/cuda-oxide", "vortx-shaders/cuda-oxide"]
 
 [dependencies]
 vortx-shaders = { workspace = true }

crates/nexus_viewer2d/Cargo.toml

@@ -27,6 +27,7 @@ metal = ["khal/metal", "nexus2d/metal"]
 cpu = ["nexus2d/cpu", "vortx/cpu"]
 cpu-parallel = ["cpu", "nexus2d/cpu-parallel", "vortx/cpu-parallel"]
 cuda = ["nexus2d/cuda"]
+cuda-oxide = ["cuda", "nexus2d/cuda-oxide"]
 
 [dependencies]
 nexus2d = { workspace = true, features = ["rbd", "mpm"]}

crates/nexus_viewer3d/Cargo.toml

@@ -27,6 +27,7 @@ metal = ["khal/metal", "nexus3d/metal"]
 cpu = ["nexus3d/cpu", "vortx/cpu"]
 cpu-parallel = ["cpu", "nexus3d/cpu-parallel", "vortx/cpu-parallel"]
 cuda = ["nexus3d/cuda"]
+cuda-oxide = ["cuda", "nexus3d/cuda-oxide"]
 
 [dependencies]
 nexus3d = { workspace = true, features = ["rbd", "mpm"]}