diff --git a/README.md b/README.md index e83c6b16..089a071b 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,31 @@ Nexus is organized into independent physics modules, each available in 2D and 3D materials, fluids, etc. Supports one-way coupling with rigid-bodies (= rigid-bodies can push particles but particles cannot push rigid-bodies). +## Rigid-body friction (3D) + +Choose the contact model globally with `RbdSimParams::friction_model`: + +```rust ignore +use nexus_rbd3d::dynamics::{FrictionModel, RbdSimParams}; + +let params = RbdSimParams { + friction_model: FrictionModel::Simplified, + ..RbdSimParams::default() +}; +``` + +`Coulomb` (the default) solves sliding friction at each contact point. +`Simplified` follows Rapier's twist-friction model: per-point normal constraints, +one sliding-friction constraint at the manifold center, and one angular constraint +resisting twist about its normal. A single-point contact has no twist constraint. +The models approximate friction differently and can produce different motion. +Multibody contacts always use Coulomb, as in Rapier; 2D is unchanged. + +Pass these parameters to `RbdState::from_rapier` (identically for every environment), +or call `NexusState::set_rbd_friction_model` before `finalize`. A live `RbdState` +can switch all environments with `set_friction_model(&backend, model)`. +Switching clears incompatible friction warmstarts and preserves normal warmstarts. + ## Prerequisites ### Install `cargo gpu` diff --git a/crates/examples2d/all_examples2.rs b/crates/examples2d/all_examples2.rs index 0eae8b57..1de538e6 100644 --- a/crates/examples2d/all_examples2.rs +++ b/crates/examples2d/all_examples2.rs @@ -25,6 +25,9 @@ mod mpm_hourglass2; mod mpm_sand2; mod mpm_snowball2; +// Large stress scenes, listed in their own picker group. +mod stress_scenes; + /// Declares the demo registry: a `(name, kind)` list for the picker UI and a /// name -> `run()` dispatcher. Keeping both in one macro keeps them in sync. macro_rules! demos { @@ -32,6 +35,7 @@ macro_rules! demos { fn demo_list() -> Vec<(String, DemoKind)> { let mut demos: Vec<(String, DemoKind)> = vec![ $( ($name.to_string(), DemoKind::$kind) ),* ]; + demos.extend(stress_scenes::demo_list()); demos.sort_by(|a, b| match (a.0.starts_with('('), b.0.starts_with('(')) { (true, true) | (false, false) => a.0.cmp(&b.0), (true, false) => std::cmp::Ordering::Greater, @@ -46,7 +50,11 @@ macro_rules! demos { // migrated to the `NexusState` API); discard whatever it yields // so every arm has the same `()` type. $( $name => { let _ = $module::run(viewer, pipeline).await; }, )* - _ => eprintln!("Unknown demo: '{name}'"), + _ => { + if stress_scenes::run(name, viewer, pipeline).await.is_none() { + eprintln!("Unknown demo: '{name}'"); + } + } } } }; @@ -85,6 +93,8 @@ struct CliOptions { metal: bool, run: bool, deterministic: bool, + snapshot: Option<(u32, String)>, + headless: bool, } fn parse_command_line() -> CliOptions { @@ -98,6 +108,8 @@ fn parse_command_line() -> CliOptions { metal: false, run: false, deterministic: false, + snapshot: None, + headless: false, }; while let Some(arg) = args.next() { @@ -110,6 +122,11 @@ fn parse_command_line() -> CliOptions { "--metal" => opts.metal = true, "--run" => opts.run = true, "--deterministic" => opts.deterministic = true, + "--headless" => opts.headless = true, + "--snapshot" => { + let frame = args.next().and_then(|f| f.parse().ok()); + opts.snapshot = frame.zip(args.next()); + } _ => {} } } @@ -143,8 +160,14 @@ pub async fn main() { } } - let mut viewer = NexusViewer::new(demos.clone()).await; - viewer = viewer.with_selected_demo(selected); + let mut viewer = if opts.headless { + NexusViewer::new_headless_with_size(demos.clone(), 1600, 1200).await + } else { + NexusViewer::new(demos.clone()).await + }; + viewer = viewer + .with_selected_demo(selected) + .with_wait_for_render(true); if opts.cpu { viewer = viewer.with_cpu(); } @@ -165,6 +188,10 @@ pub async fn main() { if opts.deterministic { viewer = viewer.with_deterministic(true); } + #[cfg(not(target_arch = "wasm32"))] + if let Some((frame, path)) = opts.snapshot.clone() { + viewer = viewer.with_snapshot(frame, path); + } // The GPU pipelines are owned here (not by `NexusState`) so they can be // compiled once up-front and reused across demos. A backend switch drops and diff --git a/crates/examples2d/stress_scenes/box_rain.rs b/crates/examples2d/stress_scenes/box_rain.rs new file mode 100644 index 00000000..41ff78d0 --- /dev/null +++ b/crates/examples2d/stress_scenes/box_rain.rs @@ -0,0 +1,49 @@ +//! A grid of boxes dropped into a walled container. + +use super::builder::{Camera, Random, Scene}; + +/// A `cols` × `rows` grid of boxes (slightly jittered sizes) dropped into a walled container. +fn box_rain(camera: Camera, cols: usize, rows: usize) -> Scene { + let mut s = Scene::new(camera); + let width = cols as f32 * 1.2 + 4.0; + let wall_height = rows as f32 * 1.4 + 10.0; + s.rigid([width + 2.0, 1.0], 0.0, 0.5, [0.0, -0.5, 0.0]); + s.rigid( + [1.0, wall_height], + 0.0, + 0.5, + [-width / 2.0, wall_height / 2.0, 0.0], + ); + s.rigid( + [1.0, wall_height], + 0.0, + 0.5, + [width / 2.0, wall_height / 2.0, 0.0], + ); + let mut rand = Random::new(1); + for y in 0..rows { + for x in 0..cols { + let w = 0.6 + rand.next() * 0.5; + let h = 0.6 + rand.next() * 0.5; + let pos = [ + x as f32 * 1.2 - (cols - 1) as f32 * 0.6, + 2.0 + y as f32 * 1.4, + rand.next() * 0.5, + ]; + s.rigid([w, h], 1.0, 0.5, pos); + } + } + s +} + +pub fn box_rain_40x25() -> Scene { + box_rain(Camera::new(0.0, 15.0, 12.0), 40, 25) +} + +pub fn box_rain_100x50() -> Scene { + box_rain(Camera::new(0.0, 30.0, 5.5), 100, 50) +} + +pub fn box_rain_900x100() -> Scene { + box_rain(Camera::new(0.0, 60.0, 0.7), 900, 100) +} diff --git a/crates/examples2d/stress_scenes/builder.rs b/crates/examples2d/stress_scenes/builder.rs new file mode 100644 index 00000000..a793cf3c --- /dev/null +++ b/crates/examples2d/stress_scenes/builder.rs @@ -0,0 +1,355 @@ +//! A small scene description (boxes and welds), translated into a nexus scene when run. +//! +//! Conventions: box sizes are full widths, positions are `[x, y, angle]`, a zero density makes +//! a body static, and gravity is -10. + +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; +use rapier2d::glamx::IVec2; +use rapier2d::prelude::*; + +/// Collision-group bit shared by every body that isn't excluded from colliding with another one. +const FREE_GROUP: u32 = 1 << 31; +/// Number of collision-group bits available to the other bodies. +const NUM_CLASSES: u32 = 31; + +pub struct Rigid { + /// Full widths along each local axis. + pub size: Vector, + pub density: f32, + pub friction: f32, + pub position: Vector, + pub angle: f32, + pub linvel: Vector, + pub angvel: f32, +} + +/// A weld between two body-local anchors. +pub struct FixedJoint { + pub a: usize, + pub b: usize, + pub ra: Vector, + pub rb: Vector, +} + +/// Camera center and zoom (pixels per meter). +pub struct Camera { + pub center: Vector, + pub zoom: f32, +} + +impl Camera { + pub fn new(x: f32, y: f32, zoom: f32) -> Self { + Self { + center: Vector::new(x, y), + zoom, + } + } +} + +pub struct Scene { + pub bodies: Vec, + pub joints: Vec, + /// Body pairs that must not collide (on top of the jointed ones). + pub ignored: Vec<(usize, usize)>, + pub camera: Camera, +} + +impl Scene { + pub fn new(camera: Camera) -> Self { + Self { + bodies: vec![], + joints: vec![], + ignored: vec![], + camera, + } + } + + /// A box of full widths `size` at `[x, y, angle]`; static if `density` is zero. + pub fn rigid(&mut self, size: [f32; 2], density: f32, friction: f32, pos: [f32; 3]) -> usize { + self.bodies.push(Rigid { + size: Vector::from(size), + density, + friction, + position: Vector::new(pos[0], pos[1]), + angle: pos[2], + linvel: Vector::ZERO, + angvel: 0.0, + }); + self.bodies.len() - 1 + } + + /// Sets a body's `[vx, vy, angular velocity]`. + pub fn set_velocity(&mut self, id: usize, velocity: [f32; 3]) { + self.bodies[id].linvel = Vector::new(velocity[0], velocity[1]); + self.bodies[id].angvel = velocity[2]; + } + + pub fn fixed_joint(&mut self, a: usize, b: usize, ra: [f32; 2], rb: [f32; 2]) { + self.joints.push(FixedJoint { + a, + b, + ra: Vector::from(ra), + rb: Vector::from(rb), + }); + } + + pub fn ignore_collision(&mut self, a: usize, b: usize) { + self.ignored.push((a, b)); + } + + /// Every pair of bodies that must not collide: jointed pairs and explicitly ignored ones. + fn excluded_pairs(&self) -> impl Iterator + '_ { + self.joints + .iter() + .map(|j| (j.a, j.b)) + .chain(self.ignored.iter().copied()) + } +} + +/// Deterministic pseudo-random numbers in [0, 1) (a linear congruential generator). +pub struct Random(u32); + +impl Random { + pub fn new(seed: u32) -> Self { + Self(seed) + } + + pub fn next(&mut self) -> f32 { + self.0 = self.0.wrapping_mul(1103515245).wrapping_add(12345); + (self.0 as f64 / 4294967296.0) as f32 + } +} + +/// Assigns collision groups so that no excluded pair collides. +/// +/// Every body excluded from colliding with another gets one of 31 group bits, chosen to differ +/// from its excluded neighbors and their own neighbors, and filters out its neighbors' bits. +/// The other bodies share the last bit and collide with everything. The price is that an +/// excluded body also ignores far-away bodies that happen to share a neighbor's bit. +fn collision_groups(scene: &Scene) -> Vec { + let n = scene.bodies.len(); + let mut adjacency = vec![vec![]; n]; + for (a, b) in scene.excluded_pairs() { + adjacency[a].push(b); + adjacency[b].push(a); + } + + const UNASSIGNED: u32 = u32::MAX; + let mut classes = vec![UNASSIGNED; n]; + let mut next = 0; + for i in 0..n { + if adjacency[i].is_empty() { + continue; + } + let mut near = 0u32; + let mut far = 0u32; + for &j in &adjacency[i] { + if classes[j] != UNASSIGNED { + near |= 1 << classes[j]; + } + for &k in &adjacency[j] { + if k != i && classes[k] != UNASSIGNED { + far |= 1 << classes[k]; + } + } + } + // Round-robin from `next` spreads the classes evenly. + let pick = |taken: u32| { + (0..NUM_CLASSES) + .map(|k| (next + k) % NUM_CLASSES) + .find(|c| taken & (1 << c) == 0) + }; + let class = pick(near | far).or_else(|| pick(near)).unwrap_or(next); + classes[i] = class; + next = (class + 1) % NUM_CLASSES; + } + + (0..n) + .map(|i| { + if classes[i] == UNASSIGNED { + InteractionGroups::new( + Group::from_bits_retain(FREE_GROUP), + Group::ALL, + InteractionTestMode::And, + ) + } else { + let excluded = adjacency[i] + .iter() + .fold(0u32, |acc, &j| acc | (1 << classes[j])); + InteractionGroups::new( + Group::from_bits_retain(1 << classes[i]), + Group::from_bits_retain(!excluded), + InteractionTestMode::And, + ) + } + }) + .collect() +} + +/// The number of body pairs whose bounding boxes overlap (within the prediction distance) at +/// the start of `scene`, and the most such pairs of a single body. +fn initial_pairs(scene: &Scene) -> (usize, usize) { + // The default prediction distance, plus some slack. + const MARGIN: f32 = 0.025; + let aabbs: Vec<(Vector, Vector)> = scene + .bodies + .iter() + .map(|b| { + let (sin, cos) = b.angle.sin_cos(); + let h = b.size / 2.0; + let half = Vector::new( + cos.abs() * h.x + sin.abs() * h.y, + sin.abs() * h.x + cos.abs() * h.y, + ) + Vector::splat(MARGIN); + (b.position - half, b.position + half) + }) + .collect(); + + // Grid cells of twice the median body size; larger bodies are tested against every body. + let mut sizes: Vec = aabbs + .iter() + .map(|(lo, hi)| (*hi - *lo).max_element()) + .collect(); + sizes.sort_by(f32::total_cmp); + let cell = 2.0 * sizes.get(sizes.len() / 2).copied().unwrap_or(1.0); + let cell_of = |p: Vector| (p / cell).floor().as_ivec2(); + let overlap = |a: usize, b: usize| { + let ((lo1, hi1), (lo2, hi2)) = (aabbs[a], aabbs[b]); + lo1.cmple(hi2).all() && lo2.cmple(hi1).all() + }; + let solved = |i: usize| scene.bodies[i].density > 0.0; + + let mut grid: std::collections::HashMap> = Default::default(); + let mut large = vec![]; + for (i, (lo, hi)) in aabbs.iter().enumerate() { + if (*hi - *lo).max_element() > cell { + large.push(i); + continue; + } + let (c0, c1) = (cell_of(*lo), cell_of(*hi)); + for x in c0.x..=c1.x { + for y in c0.y..=c1.y { + grid.entry(IVec2::new(x, y)).or_default().push(i); + } + } + } + + let mut degree = vec![0usize; aabbs.len()]; + let mut count = 0; + let mut add = |a: usize, b: usize, degree: &mut Vec| { + if (solved(a) || solved(b)) && overlap(a, b) { + count += 1; + degree[a] += 1; + degree[b] += 1; + } + }; + for (key, bodies) in &grid { + for (k, &a) in bodies.iter().enumerate() { + for &b in &bodies[k + 1..] { + // Count a pair once: in the cell holding the low corner of the boxes' overlap. + let low = aabbs[a].0.max(aabbs[b].0); + if cell_of(low) == *key { + add(a, b, &mut degree); + } + } + } + } + for (k, &a) in large.iter().enumerate() { + for b in 0..aabbs.len() { + // Large-large pairs once, large-small pairs from the large side. + if b != a && (!large.contains(&b) || large[..k].contains(&b)) { + add(a, b, &mut degree); + } + } + } + + // Static bodies aren't colored. + let max_degree = (0..aabbs.len()) + .filter(|&i| solved(i)) + .map(|i| degree[i]) + .max() + .unwrap_or(0); + (count, max_degree) +} + +/// The gravity of the stress scenes. +pub const GRAVITY: [f32; 3] = [0.0, -10.0, 0.0]; + +/// Builds `scene` into a fresh [`NexusState`], with each body's handle and shape for rendering. +pub fn build_state(scene: &Scene) -> (NexusState, Vec<(RigidBodyHandle, SharedShape)>) { + let n = scene.bodies.len(); + // The buffers grow from lagging readbacks: start them with room for the scene's initial + // contacts (twice, for the bodies coming into contact later) instead of letting early steps + // overflow. + let (pairs, max_degree) = initial_pairs(scene); + let mut capacities = NexusCapacities::default() + .rbd_bodies((n as u32 + 16).max(65536)) + .rbd_collisions((2 * pairs as u32).max(4096)); + // A body's constraints all need distinct colors. + capacities.rbd.solver_colors = (max_degree as u32 + 8).max(capacities.rbd.solver_colors); + let mut state = NexusState::new(capacities); + let groups = collision_groups(scene); + + let mut shapes = Vec::with_capacity(n); + for (body, groups) in scene.bodies.iter().zip(groups) { + let builder = if body.density > 0.0 { + RigidBodyBuilder::dynamic() + } else { + RigidBodyBuilder::fixed() + }; + let rb = builder + .translation(body.position) + .rotation(body.angle) + .linvel(body.linvel) + .angvel(body.angvel) + .build(); + // Frictions mix as sqrt(a * b): multiplying square roots gives exactly that. + let collider = ColliderBuilder::cuboid(body.size.x / 2.0, body.size.y / 2.0) + .density(body.density.max(0.0)) + .friction(body.friction.sqrt()) + .friction_combine_rule(CoefficientCombineRule::Multiply) + .collision_groups(groups) + .build(); + let shape = collider.shared_shape().clone(); + let handle = state.insert_rigid_body(rb, collider, RbdCoupling::None); + shapes.push((handle, shape)); + } + + for joint in &scene.joints { + let builder = FixedJointBuilder::new() + .local_anchor1(joint.ra) + .local_anchor2(joint.rb); + state.insert_impulse_joint(shapes[joint.a].0, shapes[joint.b].0, builder); + } + (state, shapes) +} + +/// Builds `scene` into a fresh [`NexusState`] and runs it until the viewer switches demos. +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, + scene: Scene, +) -> anyhow::Result { + let (mut state, shapes) = build_state(&scene); + for (handle, shape) in &shapes { + viewer.insert_shape(*handle, shape, Pose::IDENTITY); + } + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer.set_camera_2d(scene.camera.center, scene.camera.zoom); + state.finalize(viewer.backend()).await?; + state.set_rbd_gravity(viewer.backend(), GRAVITY); + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples2d/stress_scenes/joint_lattice.rs b/crates/examples2d/stress_scenes/joint_lattice.rs new file mode 100644 index 00000000..3465ea9a --- /dev/null +++ b/crates/examples2d/stress_scenes/joint_lattice.rs @@ -0,0 +1,47 @@ +//! A lattice of welded boxes hanging from its two top corners. + +use super::builder::{Camera, Scene}; + +/// A `w` × `h` lattice of unit boxes welded to their neighbors. +fn joint_lattice(camera: Camera, w: usize, h: usize) -> Scene { + let mut s = Scene::new(camera); + let mut grid = vec![]; + for x in 0..w { + let mut column = vec![]; + for y in 0..h { + let pinned = y == h - 1 && (x == 0 || x == w - 1); + let density = if pinned { 0.0 } else { 1.0 }; + column.push(s.rigid([1.0, 1.0], density, 0.5, [x as f32, y as f32, 0.0])); + } + grid.push(column); + } + for (left, right) in grid.iter().zip(&grid[1..]) { + for (&a, &b) in left.iter().zip(right) { + s.fixed_joint(a, b, [0.5, 0.0], [-0.5, 0.0]); + } + } + for column in &grid { + for (&below, &above) in column.iter().zip(&column[1..]) { + s.fixed_joint(below, above, [0.0, 0.5], [0.0, -0.5]); + } + } + for x in 1..w { + for y in 1..h { + s.ignore_collision(grid[x - 1][y - 1], grid[x][y]); + s.ignore_collision(grid[x][y - 1], grid[x - 1][y]); + } + } + s +} + +pub fn joint_lattice_64x64() -> Scene { + joint_lattice(Camera::new(32.0, 20.0, 7.0), 64, 64) +} + +pub fn joint_lattice_320x320() -> Scene { + joint_lattice(Camera::new(160.0, 120.0, 1.6), 320, 320) +} + +pub fn joint_lattice_512x512() -> Scene { + joint_lattice(Camera::new(256.0, 180.0, 1.0), 512, 512) +} diff --git a/crates/examples2d/stress_scenes/mod.rs b/crates/examples2d/stress_scenes/mod.rs new file mode 100644 index 00000000..5a81821e --- /dev/null +++ b/crates/examples2d/stress_scenes/mod.rs @@ -0,0 +1,11 @@ +//! Large rigid-body scenes for performance work, grouped under their own entry in the demo +//! picker. + +mod box_rain; +mod builder; +mod joint_lattice; +mod pyramid; +mod registry; +mod wrecking_ball; + +pub use registry::{demo_list, run}; diff --git a/crates/examples2d/stress_scenes/pyramid.rs b/crates/examples2d/stress_scenes/pyramid.rs new file mode 100644 index 00000000..8e3c3697 --- /dev/null +++ b/crates/examples2d/stress_scenes/pyramid.rs @@ -0,0 +1,33 @@ +//! A large pyramid of boxes. + +use super::builder::{Camera, Scene}; + +/// A pyramid of `size` rows of 1 × 0.5 boxes. +fn pyramid(camera: Camera, size: usize) -> Scene { + let mut s = Scene::new(camera); + let ground_width = (size as f32 * 1.5).max(100.0); + s.rigid([ground_width, 0.5], 0.0, 0.5, [0.0, -2.0, 0.0]); + for y in 0..size { + for x in 0..size - y { + let pos = [ + x as f32 * 1.1 + y as f32 * 0.5 - size as f32 / 2.0, + y as f32 * 0.85, + 0.0, + ]; + s.rigid([1.0, 0.5], 1.0, 0.5, pos); + } + } + s +} + +pub fn pyramid_50() -> Scene { + pyramid(Camera::new(0.0, 10.0, 12.0), 50) +} + +pub fn pyramid_100() -> Scene { + pyramid(Camera::new(0.0, 22.0, 6.0), 100) +} + +pub fn pyramid_200() -> Scene { + pyramid(Camera::new(0.0, 45.0, 2.8), 200) +} diff --git a/crates/examples2d/stress_scenes/registry.rs b/crates/examples2d/stress_scenes/registry.rs new file mode 100644 index 00000000..d5915d5b --- /dev/null +++ b/crates/examples2d/stress_scenes/registry.rs @@ -0,0 +1,54 @@ +use super::builder::{self, Scene}; +use super::{box_rain, joint_lattice, pyramid, wrecking_ball}; +use nexus_viewer2d::{DemoKind, NexusViewer}; +use nexus2d::prelude::NexusPipeline; + +/// A stress scene: its display name and its builder. +type SceneEntry = (&'static str, fn() -> Scene); + +const SCENES: &[SceneEntry] = &[ + ("Pyramid 50 (1.3k)", pyramid::pyramid_50), + ("Pyramid 100 (5k)", pyramid::pyramid_100), + ("Pyramid 200 (20k)", pyramid::pyramid_200), + ("Box rain 40x25 (1k)", box_rain::box_rain_40x25), + ("Box rain 100x50 (5k)", box_rain::box_rain_100x50), + ("Box rain 900x100 (90k)", box_rain::box_rain_900x100), + ( + "Joint lattice 64x64 (4k)", + joint_lattice::joint_lattice_64x64, + ), + ( + "Joint lattice 320x320 (100k)", + joint_lattice::joint_lattice_320x320, + ), + ( + "Joint lattice 512x512 (262k)", + joint_lattice::joint_lattice_512x512, + ), + ( + "Wrecking ball 100x40 (4k)", + wrecking_ball::wrecking_ball_100x40, + ), + ( + "Wrecking ball 400x100 (40k)", + wrecking_ball::wrecking_ball_400x100, + ), +]; + +/// The stress scenes for the demo picker. +pub fn demo_list() -> Vec<(String, DemoKind)> { + SCENES + .iter() + .map(|(name, _)| (name.to_string(), DemoKind::Stress)) + .collect() +} + +/// Runs the stress scene called `name`, or returns `None` if there is none. +pub async fn run( + name: &str, + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> Option> { + let (_, build) = SCENES.iter().find(|(scene, _)| *scene == name)?; + Some(builder::run(viewer, pipeline, build()).await.map(|_| ())) +} diff --git a/crates/examples2d/stress_scenes/wrecking_ball.rs b/crates/examples2d/stress_scenes/wrecking_ball.rs new file mode 100644 index 00000000..d54773a0 --- /dev/null +++ b/crates/examples2d/stress_scenes/wrecking_ball.rs @@ -0,0 +1,37 @@ +//! A brick wall hit by a heavy, fast block. + +use super::builder::{Camera, Scene}; + +/// A wall of `cols` × `rows` bricks hit by a heavy, fast block. +fn wrecking_ball(camera: Camera, cols: usize, rows: usize) -> Scene { + let mut s = Scene::new(camera); + s.rigid([cols as f32 * 3.0 + 200.0, 1.0], 0.0, 0.6, [0.0, -0.5, 0.0]); + for y in 0..rows { + let offset = if y % 2 == 0 { 0.0 } else { 0.5 }; + for x in 0..cols { + let pos = [ + x as f32 + offset - cols as f32 / 2.0, + 0.25 + y as f32 * 0.5, + 0.0, + ]; + s.rigid([1.0, 0.5], 1.0, 0.6, pos); + } + } + let size = (rows as f32 * 0.15).max(3.0); + let pos = [ + -(cols as f32) / 2.0 - size * 2.0 - 5.0, + rows as f32 * 0.25, + 0.0, + ]; + let block = s.rigid([size, size], 20.0, 0.5, pos); + s.set_velocity(block, [40.0, 0.0, 0.0]); + s +} + +pub fn wrecking_ball_100x40() -> Scene { + wrecking_ball(Camera::new(-10.0, 10.0, 8.0), 100, 40) +} + +pub fn wrecking_ball_400x100() -> Scene { + wrecking_ball(Camera::new(-80.0, 25.0, 2.2), 400, 100) +} diff --git a/crates/examples3d/all_examples3.rs b/crates/examples3d/all_examples3.rs index 80a4d29e..4b97d58a 100644 --- a/crates/examples3d/all_examples3.rs +++ b/crates/examples3d/all_examples3.rs @@ -38,6 +38,9 @@ mod mpm_jelly_drop3; mod mpm_sand3; mod mpm_snow3; +// Large stress scenes, listed in their own picker group. +mod stress_scenes; + /// Declares the demo registry: a `(name, kind)` list for the picker UI and a /// name -> `run()` dispatcher. Keeping both in one macro keeps them in sync. /// Entries can carry attributes (e.g. `#[cfg(...)]`) to exclude a demo from @@ -53,6 +56,7 @@ macro_rules! demos { $(#[$attr])* demos.push(($name.to_string(), DemoKind::$kind)); )* + demos.extend(stress_scenes::demo_list()); // Lexicographic sort, with stress tests (names starting with '(') // moved to the end of the list. demos.sort_by(|a, b| match (a.0.starts_with('('), b.0.starts_with('(')) { @@ -69,7 +73,11 @@ macro_rules! demos { // migrated to the `NexusState` API); discard whatever it yields // so every arm has the same `()` type. $( $(#[$attr])* $name => { let _ = $module::run(viewer, pipeline).await; }, )* - _ => eprintln!("Unknown demo: '{name}'"), + _ => { + if stress_scenes::run(name, viewer, pipeline).await.is_none() { + eprintln!("Unknown demo: '{name}'"); + } + } } } }; @@ -119,6 +127,8 @@ struct CliOptions { metal: bool, run: bool, deterministic: bool, + snapshot: Option<(u32, String)>, + headless: bool, } fn parse_command_line() -> CliOptions { @@ -132,6 +142,8 @@ fn parse_command_line() -> CliOptions { metal: false, run: false, deterministic: false, + snapshot: None, + headless: false, }; while let Some(arg) = args.next() { @@ -144,6 +156,11 @@ fn parse_command_line() -> CliOptions { "--metal" => opts.metal = true, "--run" => opts.run = true, "--deterministic" => opts.deterministic = true, + "--snapshot" => { + let frame = args.next().and_then(|f| f.parse().ok()); + opts.snapshot = frame.zip(args.next()); + } + "--headless" => opts.headless = true, _ => {} } } @@ -179,8 +196,14 @@ pub async fn main() { } } - let mut viewer = NexusViewer::new(demos.clone()).await; - viewer = viewer.with_selected_demo(selected); + let mut viewer = if opts.headless { + NexusViewer::new_headless_with_size(demos.clone(), 1600, 1200).await + } else { + NexusViewer::new(demos.clone()).await + }; + viewer = viewer + .with_selected_demo(selected) + .with_wait_for_render(true); if opts.cpu { viewer = viewer.with_cpu(); } @@ -201,6 +224,10 @@ pub async fn main() { if opts.deterministic { viewer = viewer.with_deterministic(true); } + #[cfg(not(target_arch = "wasm32"))] + if let Some((frame, path)) = opts.snapshot.clone() { + viewer = viewer.with_snapshot(frame, path); + } // The GPU pipelines are owned here (not by `NexusState`) so they can be // compiled once up-front and reused across demos. A backend switch drops and diff --git a/crates/examples3d/stress_scenes/box_columns.rs b/crates/examples3d/stress_scenes/box_columns.rs new file mode 100644 index 00000000..4b497156 --- /dev/null +++ b/crates/examples3d/stress_scenes/box_columns.rs @@ -0,0 +1,24 @@ +//! A forest of box columns. + +use super::builder::{Camera, Scene}; +use super::common::ground; + +/// `n` × `n` columns of `h` unit boxes, each resting on the one below. +pub fn box_columns_100k() -> Scene { + let (n, h) = (100, 10); + let mut s = Scene::new(Camera::new(190.0, [0.0, 0.0, 5.0]).elevation(0.5)); + ground(&mut s, n as f32 * 1.5); + for x in 0..n { + for y in 0..n { + for z in 0..h { + let pos = [ + (x as f32 - n as f32 / 2.0) * 1.5, + (y as f32 - n as f32 / 2.0) * 1.5, + 1.0 + z as f32, + ]; + s.rigid([1.0, 1.0, 1.0], 1.0, 0.5, pos); + } + } + } + s +} diff --git a/crates/examples3d/stress_scenes/box_pile.rs b/crates/examples3d/stress_scenes/box_pile.rs new file mode 100644 index 00000000..6569dc13 --- /dev/null +++ b/crates/examples3d/stress_scenes/box_pile.rs @@ -0,0 +1,47 @@ +//! A block of randomly sized and turned boxes dropped into a pile. + +use super::builder::{Camera, Random, Scene}; +use super::common::ground; +use rapier3d::prelude::*; + +/// An `n` × `n` × `h` block of boxes. +fn box_pile(camera: Camera, n: usize, h: usize) -> Scene { + let mut s = Scene::new(camera); + ground(&mut s, n as f32 * 1.6); + let mut rand = Random::new(12345); + for z in 0..h { + for x in 0..n { + for y in 0..n { + let size = [ + 0.5 + rand.next() * 0.7, + 0.5 + rand.next() * 0.7, + 0.5 + rand.next() * 0.7, + ]; + let pos = [ + (x as f32 - n as f32 / 2.0) * 1.6, + (y as f32 - n as f32 / 2.0) * 1.6, + 2.0 + z as f32 * 1.6, + ]; + let id = s.rigid(size, 1.0, 0.5, pos); + let q = Rotation::from_xyzw( + rand.next() - 0.5, + rand.next() - 0.5, + rand.next() - 0.5, + rand.next() - 0.5, + ); + s.set_rotation(id, q); + } + } + } + s +} + +pub fn box_pile_4k() -> Scene { + let camera = Camera::new(55.0, [0.0, 0.0, 4.0]).elevation(0.45); + box_pile(camera, 20, 10) +} + +pub fn box_pile_32k() -> Scene { + let camera = Camera::new(110.0, [0.0, 0.0, 6.0]).elevation(0.45); + box_pile(camera, 40, 20) +} diff --git a/crates/examples3d/stress_scenes/brick_ring.rs b/crates/examples3d/stress_scenes/brick_ring.rs new file mode 100644 index 00000000..0a3b1503 --- /dev/null +++ b/crates/examples3d/stress_scenes/brick_ring.rs @@ -0,0 +1,54 @@ +//! A ring wall of bricks smashed by a sphere rolling in from outside. + +use super::builder::{Camera, Scene}; +use super::common::{brick, ground, rolling_ball}; +use std::f32::consts::PI; + +/// `courses` courses in running bond on rings from `radius` outwards, `rows` bricks deep at the +/// bottom and one fewer every `tier` courses. +fn brick_ring(camera: Camera, radius: f32, courses: usize, tier: usize, rows: usize) -> Scene { + let mut s = Scene::new(camera); + ground(&mut s, radius + rows as f32); + for c in 0..courses { + let depth = rows.saturating_sub(c / tier).max(1); + for j in 0..depth { + let r = radius + 0.25 + 0.52 * j as f32; + let n = (2.0 * PI * (r - 0.25) / 1.02).floor() as usize; + for i in 0..n { + let a = (i as f32 + (c % 2) as f32 * 0.5) / n as f32 * 2.0 * PI; + brick( + &mut s, + r * a.cos(), + r * a.sin(), + 0.75 + 0.5 * c as f32, + a + PI / 2.0, + ); + } + } + } + let r = courses as f32 / 5.0; + rolling_ball( + &mut s, + r, + 10.0, + 0.0, + -(radius + 0.52 * rows as f32 + r + 2.0), + 30.0, + ); + s +} + +/// The 110k-brick ring at a quarter of the bricks (same proportions). +pub fn brick_ring_28k() -> Scene { + let camera = Camera::new(110.0, [0.0, -8.0, 0.0]) + .azimuth(-100.0) + .elevation(0.3); + brick_ring(camera, 40.0, 20, 2, 10) +} + +pub fn brick_ring_110k() -> Scene { + let camera = Camera::new(215.0, [0.0, -15.0, 0.0]) + .azimuth(-100.0) + .elevation(0.3); + brick_ring(camera, 80.0, 40, 4, 10) +} diff --git a/crates/examples3d/stress_scenes/brick_walls.rs b/crates/examples3d/stress_scenes/brick_walls.rs new file mode 100644 index 00000000..d9bb5e03 --- /dev/null +++ b/crates/examples3d/stress_scenes/brick_walls.rs @@ -0,0 +1,68 @@ +//! Rows of triangular brick walls smashed by two heavy spheres. + +use super::builder::{Camera, Scene, hsl}; +use super::common::{brick, ground, rolling_ball}; + +/// `columns` × `rows` triangular brick walls, `base` bricks wide, smashed by two heavy spheres +/// rolling down the columns `ball_columns`. +fn brick_walls( + camera: Camera, + columns: usize, + rows: usize, + base: usize, + ball_columns: [usize; 2], +) -> Scene { + let mut s = Scene::new(camera); + let (pitch_x, pitch_y) = ((base + 2) as f32, 5.0); + ground( + &mut s, + (columns as f32 * pitch_x).max(rows as f32 * pitch_y) / 2.0, + ); + for cx in 0..columns { + for cy in 0..rows { + let x0 = (cx as f32 - (columns - 1) as f32 / 2.0) * pitch_x; + let y = (cy as f32 - (rows - 1) as f32 / 2.0) * pitch_y; + for k in 0..base { + let m = base - k; + for i in 0..m { + let x = x0 + (i as f32 - (m - 1) as f32 / 2.0) * 1.01; + brick(&mut s, x, y, 0.75 + 0.5 * k as f32, 0.0); + } + } + } + } + let r = base as f32 / 7.0; + for cx in ball_columns { + let x = (cx as f32 - (columns - 1) as f32 / 2.0) * pitch_x; + let y = -((rows - 1) as f32 / 2.0) * pitch_y - r - 12.0; + rolling_ball(&mut s, r, 50.0, x, y, 30.0); + } + s +} + +/// 8 × 16 walls of 210 bricks. +pub fn brick_walls_27k() -> Scene { + let (columns, rows, base) = (8, 16, 20); + let ball_columns = [2, 5]; + let camera = Camera::new(86.0, [15.0, 0.0, 2.0]) + .azimuth(-69.0) + .elevation(0.12); + brick_walls(camera, columns, rows, base, ball_columns) +} + +/// 16 × 68 walls of 465 bricks (505,920 bricks) and two rolling spheres. +pub fn pyramids_500k() -> Scene { + let camera = Camera::new(650.0, [0.0, -25.0, 5.0]) + .azimuth(-69.0) + .elevation(0.55); + let mut scene = brick_walls(camera, 16, 68, 30, [6, 9]); + scene.set_color(0, 0x808080); + for i in 1..scene.bodies.len() - 2 { + // A repeatable palette makes individual bricks visible in the dense field. + scene.set_color(i, hsl((i * 37 % 101) as f32 / 101.0, 0.45, 0.65)); + } + let n = scene.bodies.len(); + scene.set_color(n - 2, 0xad70d6); + scene.set_color(n - 1, 0xd677a7); + scene +} diff --git a/crates/examples3d/stress_scenes/builder.rs b/crates/examples3d/stress_scenes/builder.rs new file mode 100644 index 00000000..b63c6921 --- /dev/null +++ b/crates/examples3d/stress_scenes/builder.rs @@ -0,0 +1,454 @@ +//! A small scene description (boxes, spheres and ball joints), translated into a nexus scene +//! when run. +//! +//! Conventions: Z is up, box sizes are full widths, a zero density makes a body static, and +//! gravity is -10. + +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; +use rapier3d::glamx::IVec3; +use rapier3d::prelude::*; + +/// Collision-group bit shared by every body that isn't jointed to another one. +const FREE_GROUP: u32 = 1 << 31; +/// Number of collision-group bits available to jointed bodies. +const NUM_CLASSES: u32 = 31; + +pub struct Rigid { + /// Full widths along each local axis (a sphere's is `2r` on each side). + pub size: Vector, + pub density: f32, + pub friction: f32, + pub position: Vector, + pub rotation: Rotation, + pub linvel: Vector, + pub angvel: Vector, + /// Set for spheres, `None` for boxes. + pub radius: Option, + pub color: Option, +} + +/// A ball joint between two body-local anchors. +pub struct BallJoint { + pub a: usize, + pub b: usize, + pub ra: Vector, + pub rb: Vector, +} + +/// Orbit camera framing (azimuth in degrees, elevation in radians). +pub struct Camera { + pub distance: f32, + pub target: Vector, + pub azimuth: f32, + pub elevation: f32, +} + +impl Camera { + /// A camera with the default azimuth (90°) and elevation (0.35 rad). + pub fn new(distance: f32, target: [f32; 3]) -> Self { + Self { + distance, + target: Vector::from(target), + azimuth: 90.0, + elevation: 0.35, + } + } + + pub fn azimuth(mut self, azimuth: f32) -> Self { + self.azimuth = azimuth; + self + } + + pub fn elevation(mut self, elevation: f32) -> Self { + self.elevation = elevation; + self + } + + fn eye(&self) -> Vector { + let az = self.azimuth.to_radians(); + let el = self.elevation; + self.target + + self.distance * Vector::new(el.cos() * az.cos(), el.cos() * az.sin(), el.sin()) + } +} + +pub struct Scene { + pub bodies: Vec, + pub joints: Vec, + pub camera: Camera, +} + +impl Scene { + pub fn new(camera: Camera) -> Self { + Self { + bodies: vec![], + joints: vec![], + camera, + } + } + + /// A box of full widths `size`; static if `density` is zero. + pub fn rigid( + &mut self, + size: [f32; 3], + density: f32, + friction: f32, + position: [f32; 3], + ) -> usize { + self.bodies.push(Rigid { + size: Vector::from(size), + density, + friction, + position: Vector::from(position), + rotation: Rotation::IDENTITY, + linvel: Vector::ZERO, + angvel: Vector::ZERO, + radius: None, + color: None, + }); + self.bodies.len() - 1 + } + + /// A solid sphere of radius `r`; static if `density` is zero. + pub fn sphere(&mut self, r: f32, density: f32, friction: f32, position: [f32; 3]) -> usize { + let id = self.rigid([2.0 * r; 3], density, friction, position); + self.bodies[id].radius = Some(r); + id + } + + pub fn set_rotation(&mut self, id: usize, rotation: Rotation) { + self.bodies[id].rotation = rotation.normalize(); + } + + pub fn set_linvel(&mut self, id: usize, linvel: [f32; 3]) { + self.bodies[id].linvel = Vector::from(linvel); + } + + pub fn set_angvel(&mut self, id: usize, angvel: [f32; 3]) { + self.bodies[id].angvel = Vector::from(angvel); + } + + /// Sets a body's color from a `0xrrggbb` value. + pub fn set_color(&mut self, id: usize, rgb: u32) { + self.bodies[id].color = Some(rgb_color(rgb)); + } + + /// A ball joint between `a` and `b`, which then don't collide with each other. + pub fn ball_joint(&mut self, a: usize, b: usize, ra: [f32; 3], rb: [f32; 3]) { + self.joints.push(BallJoint { + a, + b, + ra: Vector::from(ra), + rb: Vector::from(rb), + }); + } +} + +/// Converts a `0xrrggbb` color to an opaque RGBA color. +pub fn rgb_color(rgb: u32) -> Vec4 { + let c = |shift: u32| ((rgb >> shift) & 0xff) as f32 / 255.0; + Vec4::new(c(16), c(8), c(0), 1.0) +} + +/// HSL (all in 0..1) to a `0xrrggbb` color, as three.js' `Color.setHSL`. +pub fn hsl(h: f32, s: f32, l: f32) -> u32 { + let a = s * l.min(1.0 - l); + let f = |n: f32| { + let k = (n + h * 12.0) % 12.0; + let v = l - a * (k - 3.0).min(9.0 - k).clamp(-1.0, 1.0); + (v.clamp(0.0, 1.0) * 255.0).round() as u32 + }; + (f(0.0) << 16) | (f(8.0) << 8) | f(4.0) +} + +/// Deterministic pseudo-random numbers in [0, 1) (a linear congruential generator). +pub struct Random(u32); + +impl Random { + pub fn new(seed: u32) -> Self { + Self(seed) + } + + pub fn next(&mut self) -> f32 { + self.0 = self.0.wrapping_mul(1103515245).wrapping_add(12345); + (self.0 as f64 / 4294967296.0) as f32 + } +} + +/// Assigns collision groups so that no jointed pair collides. +/// +/// Every jointed body gets one of 31 group bits, chosen to differ from its jointed neighbors +/// and their own neighbors, and filters out its neighbors' bits. Unjointed bodies share the +/// last bit and collide with everything. The price is that a jointed body also ignores +/// far-away bodies that happen to share a neighbor's bit. +fn collision_groups(scene: &Scene) -> Vec { + let n = scene.bodies.len(); + let mut adjacency = vec![vec![]; n]; + for joint in &scene.joints { + adjacency[joint.a].push(joint.b); + adjacency[joint.b].push(joint.a); + } + + const UNASSIGNED: u32 = u32::MAX; + let mut classes = vec![UNASSIGNED; n]; + let mut next = 0; + for i in 0..n { + if adjacency[i].is_empty() { + continue; + } + let mut near = 0u32; + let mut far = 0u32; + for &j in &adjacency[i] { + if classes[j] != UNASSIGNED { + near |= 1 << classes[j]; + } + for &k in &adjacency[j] { + if k != i && classes[k] != UNASSIGNED { + far |= 1 << classes[k]; + } + } + } + // Round-robin from `next` spreads the classes evenly. + let pick = |taken: u32| { + (0..NUM_CLASSES) + .map(|k| (next + k) % NUM_CLASSES) + .find(|c| taken & (1 << c) == 0) + }; + let class = pick(near | far).or_else(|| pick(near)).unwrap_or(next); + classes[i] = class; + next = (class + 1) % NUM_CLASSES; + } + + (0..n) + .map(|i| { + if classes[i] == UNASSIGNED { + InteractionGroups::new( + Group::from_bits_retain(FREE_GROUP), + Group::ALL, + InteractionTestMode::And, + ) + } else { + let excluded = adjacency[i] + .iter() + .fold(0u32, |acc, &j| acc | (1 << classes[j])); + InteractionGroups::new( + Group::from_bits_retain(1 << classes[i]), + Group::from_bits_retain(!excluded), + InteractionTestMode::And, + ) + } + }) + .collect() +} + +/// The number of body pairs whose bounding boxes overlap (within the prediction distance) at +/// the start of `scene`, and the most such pairs of a single body. +fn initial_pairs(scene: &Scene) -> (usize, usize) { + // The default prediction distance, plus some slack. + const MARGIN: f32 = 0.025; + let aabbs: Vec<(Vector, Vector)> = scene + .bodies + .iter() + .map(|b| { + let half = match b.radius { + Some(r) => Vector::splat(r), + None => { + let m = Mat3::from_quat(b.rotation); + let h = b.size / 2.0; + m.x_axis.abs() * h.x + m.y_axis.abs() * h.y + m.z_axis.abs() * h.z + } + }; + let half = half + Vector::splat(MARGIN); + (b.position - half, b.position + half) + }) + .collect(); + + // Grid cells of twice the median body size; larger bodies are tested against every body. + let mut sizes: Vec = aabbs + .iter() + .map(|(lo, hi)| (*hi - *lo).max_element()) + .collect(); + sizes.sort_by(f32::total_cmp); + let cell = 2.0 * sizes.get(sizes.len() / 2).copied().unwrap_or(1.0); + let cell_of = |p: Vector| (p / cell).floor().as_ivec3(); + let overlap = |a: usize, b: usize| { + let ((lo1, hi1), (lo2, hi2)) = (aabbs[a], aabbs[b]); + lo1.cmple(hi2).all() && lo2.cmple(hi1).all() + }; + let solved = |i: usize| scene.bodies[i].density > 0.0; + + let mut grid: std::collections::HashMap> = Default::default(); + let mut large = vec![]; + for (i, (lo, hi)) in aabbs.iter().enumerate() { + if (*hi - *lo).max_element() > cell { + large.push(i); + continue; + } + let (c0, c1) = (cell_of(*lo), cell_of(*hi)); + for x in c0.x..=c1.x { + for y in c0.y..=c1.y { + for z in c0.z..=c1.z { + grid.entry(IVec3::new(x, y, z)).or_default().push(i); + } + } + } + } + + let mut degree = vec![0usize; aabbs.len()]; + let mut count = 0; + let mut add = |a: usize, b: usize, degree: &mut Vec| { + if (solved(a) || solved(b)) && overlap(a, b) { + count += 1; + degree[a] += 1; + degree[b] += 1; + } + }; + for (key, bodies) in &grid { + for (k, &a) in bodies.iter().enumerate() { + for &b in &bodies[k + 1..] { + // Count a pair once: in the cell holding the low corner of the boxes' overlap. + let low = aabbs[a].0.max(aabbs[b].0); + if cell_of(low) == *key { + add(a, b, &mut degree); + } + } + } + } + for (k, &a) in large.iter().enumerate() { + for b in 0..aabbs.len() { + // Large-large pairs once, large-small pairs from the large side. + if b != a && (!large.contains(&b) || large[..k].contains(&b)) { + add(a, b, &mut degree); + } + } + } + + // Static bodies aren't colored. + let max_degree = (0..aabbs.len()) + .filter(|&i| solved(i)) + .map(|i| degree[i]) + .max() + .unwrap_or(0); + (count, max_degree) +} + +/// A body's render shape and optional color, as built by [`build_state`]. +pub struct RenderShape { + pub handle: RigidBodyHandle, + pub shape: SharedShape, + pub color: Option, +} + +/// Builds `scene` into a fresh [`NexusState`], returning the render shape of each body. +pub fn build_state(scene: &Scene) -> (NexusState, Vec) { + let n = scene.bodies.len(); + // The buffers grow from lagging readbacks: start them with room for the scene's initial + // contacts (twice, for the debris of an impact) instead of letting early steps overflow. + let (pairs, max_degree) = initial_pairs(scene); + let mut capacities = NexusCapacities::default() + .rbd_bodies((n as u32 + 16).max(65536)) + .rbd_collisions((2 * pairs as u32).max(4096)); + // A body's constraints all need distinct colors. + capacities.rbd.solver_colors = (max_degree as u32 + 8).max(capacities.rbd.solver_colors); + let mut state = NexusState::new(capacities); + let groups = collision_groups(scene); + + let mut shapes = Vec::with_capacity(n); + for (body, groups) in scene.bodies.iter().zip(groups) { + let builder = if body.density > 0.0 { + RigidBodyBuilder::dynamic() + } else { + RigidBodyBuilder::fixed() + }; + let rb = builder + .can_sleep(false) + .pose(Pose::from_parts(body.position, body.rotation)) + .linvel(body.linvel) + .angvel(body.angvel) + .build(); + let collider = match body.radius { + Some(r) => ColliderBuilder::ball(r), + None => { + ColliderBuilder::cuboid(body.size.x / 2.0, body.size.y / 2.0, body.size.z / 2.0) + } + }; + // Frictions mix as sqrt(a * b): multiplying square roots gives exactly that. + let collider = collider + .density(body.density.max(0.0)) + .friction(body.friction.sqrt()) + .friction_combine_rule(CoefficientCombineRule::Multiply) + .collision_groups(groups) + .build(); + let shape = collider.shared_shape().clone(); + let handle = state.insert_rigid_body(rb, collider, RbdCoupling::None); + shapes.push(RenderShape { + handle, + shape, + color: body.color, + }); + } + + for joint in &scene.joints { + let builder = SphericalJointBuilder::new() + .local_anchor1(joint.ra) + .local_anchor2(joint.rb); + state.insert_impulse_joint(shapes[joint.a].handle, shapes[joint.b].handle, builder); + } + + (state, shapes) +} + +/// Gravity of the stress scenes (Z-up). +pub const GRAVITY: [f32; 3] = [0.0, 0.0, -10.0]; + +/// Builds `scene` into a fresh [`NexusState`] and runs it until the viewer switches demos. +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, + scene: Scene, +) -> anyhow::Result { + let (mut state, shapes) = build_state(&scene); + for s in &shapes { + match s.color { + Some(color) => { + viewer.insert_shape_with_color(s.handle, &s.shape, Pose::IDENTITY, color) + } + None => viewer.insert_shape(s.handle, &s.shape, Pose::IDENTITY), + } + } + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, 2.0, -3.0)); + viewer.set_up_axis(Vector::Z); + viewer.set_camera(scene.camera.eye(), scene.camera.target); + + let result = simulate(viewer, pipeline, &mut state, &mut timestamps).await; + // The other demos are Y-up. + viewer.set_up_axis(Vector::Y); + result.map(|_| state) +} + +async fn simulate( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, + state: &mut NexusState, + timestamps: &mut GpuTimestamps, +) -> anyhow::Result<()> { + state.finalize(viewer.backend()).await?; + state.set_rbd_gravity(viewer.backend(), GRAVITY); + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), state, Some(timestamps)) + .await?; + } + viewer.sync(state, Some(timestamps)).await?; + } + + Ok(()) +} diff --git a/crates/examples3d/stress_scenes/common.rs b/crates/examples3d/stress_scenes/common.rs new file mode 100644 index 00000000..89308621 --- /dev/null +++ b/crates/examples3d/stress_scenes/common.rs @@ -0,0 +1,72 @@ +//! Pieces shared by several stress scenes. + +use super::builder::Scene; +use rapier3d::prelude::*; + +/// A ground slab sized to hold `extent` meters of content, its top face at z = 0.5. +pub fn ground(s: &mut Scene, extent: f32) -> usize { + let size = (4.0 * extent).max(200.0); + s.rigid([size, size, 1.0], 0.0, 0.5, [0.0; 3]) +} + +/// A 1 × 0.5 × 0.5 brick turned `angle` about z. +pub fn brick(s: &mut Scene, x: f32, y: f32, z: f32, angle: f32) -> usize { + let id = s.rigid([1.0, 0.5, 0.5], 1.0, 0.5, [x, y, z]); + s.set_rotation(id, Rotation::from_rotation_z(angle)); + id +} + +/// A sphere resting on the ground at (x, y), rolling along +y at `speed`. +pub fn rolling_ball(s: &mut Scene, r: f32, density: f32, x: f32, y: f32, speed: f32) -> usize { + let ball = s.sphere(r, density, 0.5, [x, y, 0.5 + r]); + s.set_linvel(ball, [0.0, speed, 0.0]); + s.set_angvel(ball, [-speed / r, 0.0, 0.0]); + ball +} + +/// An `n` × `n` grid of thin plates ball-jointed at their edges, static where `pinned`. +pub fn plate_grid( + s: &mut Scene, + n: usize, + pitch: f32, + z: f32, + friction: f32, + pinned: impl Fn(usize, usize) -> bool, +) -> Vec> { + let mut plates = vec![]; + for x in 0..n { + let mut column = vec![]; + for y in 0..n { + let density = if pinned(x, y) { 0.0 } else { 1.0 }; + let pos = [ + (x as f32 - (n - 1) as f32 / 2.0) * pitch, + (y as f32 - (n - 1) as f32 / 2.0) * pitch, + z, + ]; + column.push(s.rigid([pitch * 0.9, pitch * 0.9, 0.1], density, friction, pos)); + } + plates.push(column); + } + let h = pitch / 2.0; + for x in 0..n { + for y in 0..n { + if x > 0 { + s.ball_joint( + plates[x - 1][y], + plates[x][y], + [h, 0.0, 0.0], + [-h, 0.0, 0.0], + ); + } + if y > 0 { + s.ball_joint( + plates[x][y - 1], + plates[x][y], + [0.0, h, 0.0], + [0.0, -h, 0.0], + ); + } + } + } + plates +} diff --git a/crates/examples3d/stress_scenes/jointed_drop.rs b/crates/examples3d/stress_scenes/jointed_drop.rs new file mode 100644 index 00000000..807f64ce --- /dev/null +++ b/crates/examples3d/stress_scenes/jointed_drop.rs @@ -0,0 +1,83 @@ +//! Jointed blocks of cubes dropped onto a chain-mail net. + +use super::builder::{Camera, Scene}; +use super::common::{ground, plate_grid}; + +/// Plates of 5 × 5 × 2 small cubes ball-jointed at their shared face centers, `grid` × `grid` +/// per layer, drop in `layers` staggered layers onto a chain-mail net pinned along its border. +pub fn jointed_drop() -> Scene { + let (grid, layers, net) = (10, 6, 64); + let camera = Camera::new(55.0, [0.0, 0.0, 4.0]) + .azimuth(-120.0) + .elevation(0.5); + let mut s = Scene::new(camera); + ground(&mut s, net as f32 * 0.5); + let z_net = 4.0; + let edge = |x: usize, y: usize| x == 0 || y == 0 || x == net - 1 || y == net - 1; + let links = plate_grid(&mut s, net, 0.5, z_net, 0.5, edge); + for id in links.into_iter().flatten() { + s.set_color(id, 0xb8bcc2); + } + + let c = 0.4; + let (nx, ny, nz) = (5, 5, 2); + let pitch = 3.0; + for layer in 0..layers { + let offset = if layer % 2 == 0 { 0.0 } else { pitch / 2.0 }; + for gx in 0..grid { + for gy in 0..grid { + let origin = [ + (gx as f32 - (grid - 1) as f32 / 2.0) * pitch + offset - 0.75, + (gy as f32 - (grid - 1) as f32 / 2.0) * pitch + offset - 0.75, + z_net + 3.0 + layer as f32 * 2.5, + ]; + let mut cells = vec![]; + for x in 0..nx { + for y in 0..ny { + for z in 0..nz { + let pos = [ + origin[0] + (x as f32 - (nx - 1) as f32 / 2.0) * c, + origin[1] + (y as f32 - (ny - 1) as f32 / 2.0) * c, + origin[2] + z as f32 * c, + ]; + cells.push(s.rigid([c; 3], 1.0, 0.5, pos)); + } + } + } + let at = |x: usize, y: usize, z: usize| cells[(x * ny + y) * nz + z]; + let h = c / 2.0; + for x in 0..nx { + for y in 0..ny { + for z in 0..nz { + if x > 0 { + s.ball_joint( + at(x - 1, y, z), + at(x, y, z), + [h, 0.0, 0.0], + [-h, 0.0, 0.0], + ); + } + if y > 0 { + s.ball_joint( + at(x, y - 1, z), + at(x, y, z), + [0.0, h, 0.0], + [0.0, -h, 0.0], + ); + } + if z > 0 { + s.ball_joint( + at(x, y, z - 1), + at(x, y, z), + [0.0, 0.0, h], + [0.0, 0.0, -h], + ); + } + } + } + } + } + } + } + s +} diff --git a/crates/examples3d/stress_scenes/mod.rs b/crates/examples3d/stress_scenes/mod.rs new file mode 100644 index 00000000..6f00ea72 --- /dev/null +++ b/crates/examples3d/stress_scenes/mod.rs @@ -0,0 +1,14 @@ +//! Large rigid-body scenes for performance work, grouped under their own entry in the demo +//! picker. + +mod box_columns; +mod box_pile; +mod brick_ring; +mod brick_walls; +mod builder; +mod common; +mod jointed_drop; +mod ragdolls_on_cloth; +mod registry; + +pub use registry::{demo_list, run}; diff --git a/crates/examples3d/stress_scenes/ragdolls_on_cloth.rs b/crates/examples3d/stress_scenes/ragdolls_on_cloth.rs new file mode 100644 index 00000000..46c2ed37 --- /dev/null +++ b/crates/examples3d/stress_scenes/ragdolls_on_cloth.rs @@ -0,0 +1,119 @@ +//! A block of ragdolls dropped onto a cloth of jointed plates. + +use super::builder::{Camera, Random, Scene, hsl}; +use super::common::{ground, plate_grid}; +use rapier3d::prelude::*; +use std::f32::consts::PI; + +/// A ragdoll standing on z = 0 in its own frame: rods (from, to, width) and a head. +const RAGDOLL_RODS: [([f32; 3], [f32; 3], f32); 9] = [ + ([0.0, 0.0, 0.5], [0.0, 0.0, 0.93], 0.24), // torso + ([0.0, 0.19, 0.9], [0.0, 0.19, 0.66], 0.09), // upper arms + ([0.0, -0.19, 0.9], [0.0, -0.19, 0.66], 0.09), // + ([0.0, 0.19, 0.66], [0.0, 0.19, 0.42], 0.085), // forearms + ([0.0, -0.19, 0.66], [0.0, -0.19, 0.42], 0.085), + ([0.0, 0.08, 0.5], [0.0, 0.08, 0.27], 0.11), // thighs + ([0.0, -0.08, 0.5], [0.0, -0.08, 0.27], 0.11), + ([0.0, 0.08, 0.27], [0.0, 0.08, 0.03], 0.1), // shins + ([0.0, -0.08, 0.27], [0.0, -0.08, 0.03], 0.1), +]; +const RAGDOLL_HEAD: ([f32; 3], f32) = ([0.0, 0.0, 1.1], 0.12); +/// Ball joints: (part, part, point), parts indexing the rods then the head (9). +const RAGDOLL_JOINTS: [(usize, usize, [f32; 3]); 9] = [ + (0, 9, [0.0, 0.0, 1.0]), + (0, 1, [0.0, 0.19, 0.9]), + (0, 2, [0.0, -0.19, 0.9]), + (1, 3, [0.0, 0.19, 0.66]), + (2, 4, [0.0, -0.19, 0.66]), + (0, 5, [0.0, 0.08, 0.5]), + (0, 6, [0.0, -0.08, 0.5]), + (5, 7, [0.0, 0.08, 0.27]), + (6, 8, [0.0, -0.08, 0.27]), +]; +/// The ragdoll's middle (its frame is turned about this). +const RAGDOLL_CENTER: [f32; 3] = [0.0, 0.0, 0.6]; + +/// The rotation taking the x axis onto the unit vector `d`. +fn align_x(d: Vector) -> Rotation { + if 1.0 + d.x < 1.0e-9 { + Rotation::from_xyzw(0.0, 0.0, 1.0, 0.0) + } else { + Rotation::from_rotation_arc(Vector::X, d) + } +} + +/// A floppy ragdoll (ball joints) at `origin`, turned by `q`, in one color. +fn ragdoll(s: &mut Scene, origin: Vector, q: Rotation, color: u32) { + let center = Vector::from(RAGDOLL_CENTER); + let place = |c: Vector| (origin + q * (c - center)).into(); + // Each part's body, center and rotation in the ragdoll's frame. + let mut parts = vec![]; + for (from, to, width) in RAGDOLL_RODS { + let (from, to) = (Vector::from(from), Vector::from(to)); + let len = (to - from).length(); + let c = (from + to) / 2.0; + let rot = align_x((to - from) / len); + let body = s.rigid([len + width, width, width], 1.0, 0.6, place(c)); + s.set_rotation(body, q * rot); + s.set_color(body, color); + parts.push((body, c, rot)); + } + let head_center = Vector::from(RAGDOLL_HEAD.0); + let head = s.sphere(RAGDOLL_HEAD.1, 1.0, 0.6, place(head_center)); + s.set_rotation(head, q); + s.set_color(head, color); + parts.push((head, head_center, Rotation::IDENTITY)); + + // Anchors in each part's own frame (the same in the ragdoll's frame as in the world's). + let local = |k: usize, point: Vector| { + let (_, c, rot) = parts[k]; + (rot.inverse() * (point - c)).into() + }; + for (a, b, point) in RAGDOLL_JOINTS { + let point = Vector::from(point); + s.ball_joint(parts[a].0, parts[b].0, local(a, point), local(b, point)); + } +} + +/// A block of ragdolls (`side` × `side` × `layers`, colored in a rainbow) dropped onto a cloth +/// of plates pinned at its four corners. +pub fn ragdolls_on_cloth() -> Scene { + let (side, layers, cloth) = (12, 10, 96); + let camera = Camera::new(52.0, [0.0, 0.0, 6.0]) + .azimuth(-115.0) + .elevation(0.3); + let mut s = Scene::new(camera); + let pitch = 0.3; + ground(&mut s, cloth as f32 * pitch * 0.5); + let z_cloth = 8.0; + let corner = |x: usize, y: usize| (x == 0 || x == cloth - 1) && (y == 0 || y == cloth - 1); + let plates = plate_grid(&mut s, cloth, pitch, z_cloth, 0.6, corner); + for id in plates.into_iter().flatten() { + s.set_color(id, 0xe8e2d4); + } + + let mut rand = Random::new(1414); + let spacing = 1.4; + for layer in 0..layers { + for gx in 0..side { + for gy in 0..side { + // A random orientation, uniform over rotations. + let (u1, u2, u3) = (rand.next(), rand.next(), rand.next()); + let q = Rotation::from_xyzw( + (1.0 - u1).sqrt() * (2.0 * PI * u2).sin(), + (1.0 - u1).sqrt() * (2.0 * PI * u2).cos(), + u1.sqrt() * (2.0 * PI * u3).sin(), + u1.sqrt() * (2.0 * PI * u3).cos(), + ); + let origin = Vector::new( + (gx as f32 - (side - 1) as f32 / 2.0) * spacing, + (gy as f32 - (side - 1) as f32 / 2.0) * spacing, + z_cloth + 1.5 + layer as f32 * spacing, + ); + let color = hsl(0.82 * gx as f32 / (side - 1).max(1) as f32, 0.72, 0.6); + ragdoll(&mut s, origin, q, color); + } + } + } + s +} diff --git a/crates/examples3d/stress_scenes/registry.rs b/crates/examples3d/stress_scenes/registry.rs new file mode 100644 index 00000000..d4209289 --- /dev/null +++ b/crates/examples3d/stress_scenes/registry.rs @@ -0,0 +1,40 @@ +use super::builder::{self, Scene}; +use super::{box_columns, box_pile, brick_ring, brick_walls, jointed_drop, ragdolls_on_cloth}; +use nexus_viewer3d::{DemoKind, NexusViewer}; +use nexus3d::prelude::NexusPipeline; + +/// A stress scene: its display name and its builder. +type SceneEntry = (&'static str, fn() -> Scene); + +const SCENES: &[SceneEntry] = &[ + ("Brick ring (28k)", brick_ring::brick_ring_28k), + ("Brick ring (110k)", brick_ring::brick_ring_110k), + ("Brick walls (27k)", brick_walls::brick_walls_27k), + ("Pyramids (500k)", brick_walls::pyramids_500k), + ( + "Ragdolls on cloth (24k)", + ragdolls_on_cloth::ragdolls_on_cloth, + ), + ("Box pile (4k)", box_pile::box_pile_4k), + ("Box pile (32k)", box_pile::box_pile_32k), + ("Jointed drop (34k)", jointed_drop::jointed_drop), + ("Box columns (100k)", box_columns::box_columns_100k), +]; + +/// The stress scenes for the demo picker. +pub fn demo_list() -> Vec<(String, DemoKind)> { + SCENES + .iter() + .map(|(name, _)| (name.to_string(), DemoKind::Stress)) + .collect() +} + +/// Runs the stress scene called `name`, or returns `None` if there is none. +pub async fn run( + name: &str, + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> Option> { + let (_, build) = SCENES.iter().find(|(scene, _)| *scene == name)?; + Some(builder::run(viewer, pipeline, build()).await.map(|_| ())) +} diff --git a/crates/nexus3d/tests/friction_model.rs b/crates/nexus3d/tests/friction_model.rs new file mode 100644 index 00000000..f7e7fba8 --- /dev/null +++ b/crates/nexus3d/tests/friction_model.rs @@ -0,0 +1,95 @@ +//! Live testbed friction changes must not rebuild bodies from their initial poses. +#![cfg(feature = "rbd")] + +use khal::backend::{Backend, GpuBackend, WebGpu}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; +use nexus3d::rbd::dynamics::FrictionModel; +use rapier3d::prelude::{ColliderBuilder, RigidBodyBuilder, Vec3}; + +#[test] +#[ignore = "requires a WebGPU adapter"] +fn switching_friction_preserves_running_bodies() { + pollster::block_on(async { + let backend = GpuBackend::WebGpu( + WebGpu::new(Default::default(), Default::default()) + .await + .unwrap(), + ); + let mut pipeline = NexusPipeline::default(); + for (batches, reserve) in [(2, 0), (1, 8)] { + let mut state = NexusState::new( + NexusCapacities::default() + .rbd_batches(batches) + .rbd_collisions(64), + ); + state.reserve_rigid_bodies(reserve); + state.set_rbd_friction_model(FrictionModel::Simplified); + for env in 0..batches as usize { + if env > 0 { + state.add_environment(); + } + state.insert_rigid_body_in( + env, + RigidBodyBuilder::dynamic() + .translation(Vec3::new(env as f32, 3.0, 0.0)) + .linvel(Vec3::X) + .can_sleep(false) + .build(), + ColliderBuilder::cuboid(0.2, 0.2, 0.2).build(), + RbdCoupling::None, + ); + } + state.finalize(&backend).await.unwrap(); + assert_eq!( + state.rbd.as_ref().unwrap().friction_model(), + FrictionModel::Simplified + ); + for _ in 0..6 { + pipeline.simulate(&backend, &mut state, None).await.unwrap(); + } + for model in [FrictionModel::Coulomb, FrictionModel::Simplified] { + let rbd = state.rbd.as_ref().unwrap(); + let poses = backend + .slow_read_vec(rbd.body_poses().buffer()) + .await + .unwrap(); + let velocities = backend.slow_read_vec(rbd.vels().buffer()).await.unwrap(); + assert!(poses.iter().any(|pose| pose.translation.y < 2.99)); + let key = rbd.graph_key(1); + state.set_rbd_friction_model(model); + state.finalize(&backend).await.unwrap(); + let rbd = state.rbd.as_ref().unwrap(); + assert_eq!(rbd.friction_model(), model); + assert_ne!( + rbd.graph_key(1), + key, + "model switch must invalidate captured work" + ); + for env in 0..batches as usize { + assert_eq!(state.rbd_sim_params(env).unwrap().friction_model, model); + } + let after_poses = backend + .slow_read_vec(rbd.body_poses().buffer()) + .await + .unwrap(); + let after_velocities = backend.slow_read_vec(rbd.vels().buffer()).await.unwrap(); + assert_eq!(poses.len(), after_poses.len()); + assert_eq!(velocities.len(), after_velocities.len()); + for (before, after) in poses.iter().zip(after_poses) { + assert_eq!(before.translation, after.translation); + assert_eq!(before.rotation, after.rotation); + } + for (before, after) in velocities.iter().zip(after_velocities) { + assert_eq!(before.linear, after.linear); + assert_eq!(before.angular, after.angular); + } + let key = state.rbd.as_ref().unwrap().graph_key(1); + // The UI pushes its selection each frame; unchanged selections are cheap. + state.set_rbd_friction_model(model); + state.finalize(&backend).await.unwrap(); + assert_eq!(state.rbd.as_ref().unwrap().graph_key(1), key); + pipeline.simulate(&backend, &mut state, None).await.unwrap(); + } + } + }); +} diff --git a/crates/nexus3d/tests/substeps.rs b/crates/nexus3d/tests/substeps.rs new file mode 100644 index 00000000..f8da381e --- /dev/null +++ b/crates/nexus3d/tests/substeps.rs @@ -0,0 +1,104 @@ +//! Live TGS substep changes preserve body state and the full physics-step duration. +#![cfg(feature = "rbd")] + +use khal::backend::{Backend, GpuBackend, WebGpu}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; +use rapier3d::prelude::{ColliderBuilder, RigidBodyBuilder, Vec3}; + +#[test] +#[ignore = "requires a WebGPU adapter"] +fn changing_substeps_preserves_state_and_elapsed_time() { + pollster::block_on(async { + let backend = GpuBackend::WebGpu( + WebGpu::new(Default::default(), Default::default()) + .await + .unwrap(), + ); + let mut pipeline = NexusPipeline::default(); + const DT: f32 = 1.0 / 50.0; + for (batches, reserve) in [(2, 0), (1, 8)] { + let mut state = NexusState::new( + NexusCapacities::default() + .rbd_batches(batches) + .rbd_collisions(64), + ); + state.reserve_rigid_bodies(reserve); + for env in 0..batches as usize { + if env > 0 { + state.add_environment(); + } + state.insert_rigid_body_in( + env, + RigidBodyBuilder::dynamic() + .translation(Vec3::new(env as f32, 3.0, 0.0)) + .linvel(Vec3::X) + .can_sleep(false) + .build(), + ColliderBuilder::cuboid(0.2, 0.2, 0.2).build(), + RbdCoupling::None, + ); + } + state.set_rbd_timestep(DT, 4); + state.set_rbd_substeps(&backend, 3); + state.finalize(&backend).await.unwrap(); + state.set_rbd_gravity(&backend, [0.0; 3]); + assert_eq!(state.rbd.as_ref().unwrap().num_solver_iterations(), 3); + for _ in 0..3 { + pipeline.simulate(&backend, &mut state, None).await.unwrap(); + } + + for substeps in [1, 8, 2, 20, 0] { + let rbd = state.rbd.as_ref().unwrap(); + let poses = backend + .slow_read_vec(rbd.body_poses().buffer()) + .await + .unwrap(); + let velocities = backend.slow_read_vec(rbd.vels().buffer()).await.unwrap(); + let key = rbd.graph_key(1); + state.set_rbd_substeps(&backend, substeps); + state.finalize(&backend).await.unwrap(); + let rbd = state.rbd.as_ref().unwrap(); + assert_eq!(rbd.num_solver_iterations(), substeps.max(1)); + assert_eq!(rbd.multibodies().num_solver_iterations(), substeps.max(1)); + assert_eq!(state.rbd_substeps(), substeps.max(1)); + assert_ne!( + rbd.graph_key(1), + key, + "invalidate captured substep dispatches" + ); + for env in 0..batches as usize { + let params = state.rbd_sim_params(env).unwrap(); + assert_eq!(params.num_solver_iterations, substeps.max(1)); + assert_eq!(params.dt, DT); + } + let after_poses = backend + .slow_read_vec(rbd.body_poses().buffer()) + .await + .unwrap(); + let after_velocities = backend.slow_read_vec(rbd.vels().buffer()).await.unwrap(); + for (before, after) in poses.iter().zip(after_poses) { + assert_eq!(before.translation, after.translation); + assert_eq!(before.rotation, after.rotation); + } + for (before, after) in velocities.iter().zip(after_velocities) { + assert_eq!(before.linear, after.linear); + assert_eq!(before.angular, after.angular); + } + let key = rbd.graph_key(1); + // The viewer pushes the selection every frame; unchanged is a no-op. + state.set_rbd_substeps(&backend, substeps); + assert_eq!(state.rbd.as_ref().unwrap().graph_key(1), key); + pipeline.simulate(&backend, &mut state, None).await.unwrap(); + let rbd = state.rbd.as_ref().unwrap(); + let after = backend + .slow_read_vec(rbd.body_poses().buffer()) + .await + .unwrap(); + for env in 0..batches as usize { + let dx = after[env].translation.x - poses[env].translation.x; + assert!((dx - DT).abs() < 1.0e-5, "{substeps} substeps: dx={dx}"); + } + } + } + }); +} diff --git a/crates/nexus_python3d/README.md b/crates/nexus_python3d/README.md index d30030d6..b38d83e8 100644 --- a/crates/nexus_python3d/README.md +++ b/crates/nexus_python3d/README.md @@ -39,6 +39,21 @@ pip install target/wheels/dimforge_nexus3d-*.whl > so a single wheel works across all supported Python versions — you don't need > to match the build interpreter to the run interpreter. +## Contact friction + +Before `finalize`, select one model for all environments: + +```python +state.set_rbd_solver_params(friction_model="simplified") +assert state.rbd_solver_params()["friction_model"] == "simplified" +``` + +`"coulomb"` is the existing default, with friction at each contact point. +`"simplified"` uses Rapier's central sliding-friction and angular twist-friction +constraints. Normal constraints stay per point; single-point contacts have no twist +constraint. Multibody contacts always use Coulomb. Invalid model names raise +`ValueError` before any solver parameters change. + ## Robots, state access and sensor cameras Beyond the demo-oriented API, the module exposes what a robotics environment diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 1361d8cb..ec5d42d1 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -1179,12 +1179,15 @@ impl NexusState { /// tolerated penetration in length units, `max_corrective_velocity` caps /// penetration recovery, `prediction_distance` is the contact detection /// margin, `internal_pgs_iterations` the biased-pass PGS iterations per - /// substep (rigid-body and multibody sweeps alike), and `friction_in_bias_pass` whether friction + /// substep (rigid-body and multibody alike), and `friction_in_bias_pass` whether friction /// rows are solved in every biased PGS iteration instead of only in the - /// per-substep stabilization sweep (rapier's default, `False`); `True` + /// per-substep stabilization iteration (rapier's default, `False`); `True` /// gives friction as many iterations as the normal rows, which holds - /// grasps and resting contacts far more firmly. - #[pyo3(signature = (contact_natural_frequency=None, contact_damping_ratio=None, static_contact_natural_frequency=None, static_contact_damping_ratio=None, allowed_linear_error=None, max_corrective_velocity=None, prediction_distance=None, internal_pgs_iterations=None, friction_in_bias_pass=None))] + /// grasps and resting contacts far more firmly. `friction_model` is + /// "coulomb" (default, friction at each contact) or "simplified" (Rapier's + /// central friction plus twist). It applies to every environment; multibody + /// contacts always use Coulomb. + #[pyo3(signature = (contact_natural_frequency=None, contact_damping_ratio=None, static_contact_natural_frequency=None, static_contact_damping_ratio=None, allowed_linear_error=None, max_corrective_velocity=None, prediction_distance=None, internal_pgs_iterations=None, friction_in_bias_pass=None, friction_model=None))] #[allow(clippy::too_many_arguments)] fn set_rbd_solver_params( &mut self, @@ -1197,7 +1200,19 @@ impl NexusState { prediction_distance: Option, internal_pgs_iterations: Option, friction_in_bias_pass: Option, - ) { + friction_model: Option<&str>, + ) -> PyResult<()> { + use nexus3d::rbd::dynamics::FrictionModel; + let friction_model = match friction_model { + None => None, + Some("coulomb") => Some(FrictionModel::Coulomb), + Some("simplified") => Some(FrictionModel::Simplified), + Some(_) => { + return Err(pyo3::exceptions::PyValueError::new_err( + "friction_model must be 'coulomb' or 'simplified'", + )); + } + }; for env in 0..self.0.num_environments() { let Some(mut params) = self.0.rbd_sim_params(env) else { continue; @@ -1229,8 +1244,12 @@ impl NexusState { if let Some(v) = friction_in_bias_pass { params.friction_in_bias_pass = v as u32; } + if let Some(v) = friction_model { + params.friction_model = v; + } self.0.set_rbd_sim_params(env, params); } + Ok(()) } /// Implicit (default) or explicit treatment of the robots' Coriolis and @@ -1257,6 +1276,14 @@ impl NexusState { use pyo3::types::PyDict; let dict = PyDict::new(py); if let Some(p) = self.0.rbd_sim_params(0) { + dict.set_item( + "friction_model", + if p.friction_model == nexus3d::rbd::dynamics::FrictionModel::Simplified { + "simplified" + } else { + "coulomb" + }, + )?; dict.set_item("dt", p.dt)?; dict.set_item("substeps", p.num_solver_iterations)?; dict.set_item("contact_natural_frequency", p.contact_natural_frequency)?; @@ -1390,37 +1417,24 @@ impl NexusState { py: Python<'py>, viewer: PyRef, ) -> PyResult>> { - use nexus3d::rbd::shaders::dynamics::TwoBodyConstraint; use pyo3::types::PyDict; let Some(rbd) = self.0.rbd.as_ref() else { return Ok(Vec::new()); }; - let cons: Vec = pollster::block_on( - viewer - .backend() - .slow_read_vec(rbd.rigid_contact_constraints().buffer()), - ) - .map_err(gpu_err)?; + let cons = pollster::block_on(rbd.read_rigid_contact_impulses(viewer.backend())) + .map_err(gpu_err)?; let mut out = Vec::new(); for c in cons.iter().filter(|c| c.len > 0) { let dict = PyDict::new(py); - dict.set_item("body_a", c.solver_body_a)?; - dict.set_item("body_b", c.solver_body_b)?; + dict.set_item("body_a", c.body_a)?; + dict.set_item("body_b", c.body_b)?; dict.set_item("dir_a", [c.dir_a.x, c.dir_a.y, c.dir_a.z])?; - dict.set_item("friction", c.limit)?; - dict.set_item("inv_mass_a", c.im_a.x)?; - dict.set_item("inv_mass_b", c.im_b.x)?; - let normal: Vec = (0..c.len as usize) - .map(|k| c.elements[k].normal_part.impulse) - .collect(); - let tangent: Vec<[f32; 2]> = (0..c.len as usize) - .map(|k| { - let t = c.elements[k].tangent_part.impulse; - [t.x, t.y] - }) - .collect(); - dict.set_item("normal_impulse", normal)?; - dict.set_item("tangent_impulse", tangent)?; + dict.set_item("friction", c.friction)?; + dict.set_item("inv_mass_a", c.inv_mass_a.x)?; + dict.set_item("inv_mass_b", c.inv_mass_b.x)?; + dict.set_item("normal_impulse", &c.normal_impulse)?; + dict.set_item("tangent_impulse", &c.tangent_impulse)?; + dict.set_item("twist_impulse", c.twist_impulse)?; out.push(dict); } Ok(out) diff --git a/crates/nexus_viewer2d/Cargo.toml b/crates/nexus_viewer2d/Cargo.toml index 17bb6bd2..133aba2c 100644 --- a/crates/nexus_viewer2d/Cargo.toml +++ b/crates/nexus_viewer2d/Cargo.toml @@ -37,5 +37,6 @@ vortx = { workspace = true } rapier2d = { workspace = true, features = ["default", "debug-render"] } bitflags = { workspace = true } bytemuck = { workspace = true } +futures = { workspace = true } kiss3d = { workspace = true, features = ["egui"] } web-time = { workspace = true } diff --git a/crates/nexus_viewer3d/Cargo.toml b/crates/nexus_viewer3d/Cargo.toml index f34241c1..417d05f4 100644 --- a/crates/nexus_viewer3d/Cargo.toml +++ b/crates/nexus_viewer3d/Cargo.toml @@ -37,5 +37,6 @@ vortx = { workspace = true } rapier3d = { workspace = true, features = ["default", "debug-render"] } bitflags = { workspace = true } bytemuck = { workspace = true } +futures = { workspace = true } kiss3d = { workspace = true, features = ["egui", "rt_switcher"] } web-time = { workspace = true } diff --git a/src/state.rs b/src/state.rs index 57bdfcbc..5c9e5078 100644 --- a/src/state.rs +++ b/src/state.rs @@ -434,7 +434,7 @@ impl NexusState { } /// Sets how many PGS iterations the biased pass runs per substep (rigid-body - /// and multibody sweeps alike). + /// and multibody alike). #[cfg(all(feature = "rbd", feature = "dim3"))] pub fn set_rbd_num_internal_pgs_iterations(&mut self, backend: &GpuBackend, n: u32) { for params in &mut self.rbd_sim_params { @@ -582,7 +582,13 @@ impl NexusState { /// poses. pub fn add_environment(&mut self) -> usize { self.rbd_envs.push(PhysicsWorld::default()); - self.rbd_sim_params.push(RbdSimParams::tgs_soft()); + #[allow(unused_mut)] + let mut params = RbdSimParams::tgs_soft(); + #[cfg(feature = "dim3")] + { + params.friction_model = self.rbd_sim_params[0].friction_model; + } + self.rbd_sim_params.push(params); self.rbd2gpu.push(Coarena::new()); self.rbd_dirty = true; self.rbd_envs.len() - 1 @@ -603,6 +609,17 @@ impl NexusState { self.rbd_sim_params.get(env).copied() } + /// Selects the friction model across all environments, applied by `finalize`. + /// A running scene keeps its body poses and velocities; only contact storage + /// and warmstarts are reset when the model changes. Multibody contacts always + /// use Coulomb. + #[cfg(feature = "dim3")] + pub fn set_rbd_friction_model(&mut self, model: crate::rbd::dynamics::FrictionModel) { + for params in &mut self.rbd_sim_params { + params.friction_model = model; + } + } + pub fn set_rbd_sim_params(&mut self, env: usize, params: RbdSimParams) { self.rbd_sim_params[env] = params; self.rbd_dirty = true; @@ -852,6 +869,24 @@ impl NexusState { } } + /// Sets the TGS substep count for every environment, clamped to at least one. + /// Updates a running simulation without resetting bodies, preserving each + /// physics step's duration. + pub fn set_rbd_substeps(&mut self, backend: &GpuBackend, substeps: u32) { + let substeps = substeps.max(1); + for params in &mut self.rbd_sim_params { + params.num_solver_iterations = substeps; + } + if let Some(rbd) = self.rbd.as_mut() { + rbd.set_timestep(backend, self.rbd_sim_params[0].dt, substeps); + } + } + + /// Number of TGS substeps per rigid-body physics step. + pub fn rbd_substeps(&self) -> u32 { + self.rbd_sim_params[0].num_solver_iterations + } + /// The multibody rooted at (or containing) `body` in environment `env`, as /// its index in the environment's multibody iteration order, the GPU /// descriptor of that multibody, and every link's rigid body paired with its @@ -1608,10 +1643,20 @@ impl NexusState { if self.deterministic { rbd_state.set_deterministic(backend, true); } + // Reserved-slot builds start with default parameters too: apply the scene's timestep. + let params = self.rbd_sim_params[0]; + rbd_state.set_timestep(backend, params.dt, params.num_solver_iterations); self.rbd = Some(rbd_state); self.rbd_dirty = false; } + // Apply a pending friction-model edit without rebuilding the rigid bodies + // from their original CPU poses. Also covers reserved-slot initial builds. + #[cfg(feature = "dim3")] + if let Some(rbd) = self.rbd.as_mut() { + rbd.set_friction_model(backend, self.rbd_sim_params[0].friction_model); + } + // MPM/rapier coupling. Boundary colliders are inserted into environment 0 // as rigid bodies tagged `RbdCoupling::Mpm*`; rebuild the coupling // (sampled rigid particles, uploaded body set) whenever those bodies or diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index 97eed52f..cc2a5655 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/src_rbd/broad_phase/lbvh.rs @@ -7,9 +7,10 @@ use crate::math::Pose; use crate::shaders::PaddedVector; use crate::shaders::bounding_volumes::Aabb; use crate::shaders::broad_phase::{ - CollisionPair, GpuBfComputeAabbs, GpuBfFindPairs, GpuFlatListDispatch, GpuLbvhBuild, - GpuLbvhComputeDomain, GpuLbvhComputeMorton, GpuLbvhFindCollisionPairs, GpuLbvhRefitInternal, - GpuLbvhRefitLeaves, GpuLbvhResetCollisionPairs, LbvhNode, + CollisionPair, DOMAIN_WORKGROUPS, GpuBfComputeAabbs, GpuBfFindPairs, GpuFlatListDispatch, + GpuLbvhBuild, GpuLbvhComputeDomain, GpuLbvhComputeMorton, GpuLbvhDomainMerge, GpuLbvhEscapes, + GpuLbvhFindCollisionPairs, GpuLbvhRefitChunks, GpuLbvhRefitFrontier, GpuLbvhRefitLeaves, + GpuLbvhResetCollisionPairs, LbvhNode, }; use crate::shaders::shapes::Shape; use crate::utils::{RadixSort, RadixSortWorkspace}; @@ -26,10 +27,15 @@ use vortx::tensor::Tensor; #[derive(Shader)] pub struct GpuLbvh { compute_domain: GpuLbvhComputeDomain, + domain_merge: GpuLbvhDomainMerge, compute_morton: GpuLbvhComputeMorton, build: GpuLbvhBuild, + escapes: GpuLbvhEscapes, refit_leaves: GpuLbvhRefitLeaves, - refit_internal: GpuLbvhRefitInternal, + refit_chunks: GpuLbvhRefitChunks, + refit_frontier: GpuLbvhRefitFrontier, + #[cfg(target_arch = "wasm32")] + refit_plan: crate::shaders::broad_phase::GpuLbvhRefitPlan, reset_collision_pairs: GpuLbvhResetCollisionPairs, find_collision_pairs: GpuLbvhFindCollisionPairs, /// Writes the `[total/64, 1, 1]` indirect grid from the single global pair @@ -47,6 +53,8 @@ pub struct GpuLbvh { pub struct LbvhState { buffer_usages: BufferUsages, domain_aabb: Tensor, + /// Per-workgroup domains of each batch, merged into `domain_aabb`. + domain_partials: Tensor, n_sort: Tensor, /// Per-batch active key count currently uploaded to `n_sort`, as /// `(active_per_batch, num_batches)`. `None` forces a re-upload (e.g. after @@ -60,6 +68,13 @@ pub struct LbvhState { unsorted_colliders: Tensor, sorted_colliders: Tensor, tree: Tensor, + /// Per batch (strided like the colliders): the nodes left by the chunked refit for its + /// second phase. + refit_frontier: Tensor, + /// Per batch: the length of its refit frontier. + refit_frontier_len: Tensor, + /// GPU-written uniform bound for all batches of the web frontier refit. + refit_rounds: Tensor, sort_workspace: RadixSortWorkspace, /// Per-collider world AABBs, only used by the brute-force tiny-batch path /// (strided by the per-batch collider capacity, like the other @@ -92,11 +107,15 @@ impl LbvhState { n_sort_active: None, generation: 0, domain_aabb: Tensor::scalar_uninit(backend, usages).unwrap(), + domain_partials: Tensor::vector_uninit(backend, DOMAIN_WORKGROUPS, usages).unwrap(), unsorted_morton_keys: Tensor::vector_uninit(backend, 0, usages).unwrap(), sorted_morton_keys: Tensor::vector_uninit(backend, 0, usages).unwrap(), unsorted_colliders: Tensor::vector_uninit(backend, 0, usages).unwrap(), sorted_colliders: Tensor::vector_uninit(backend, 0, usages).unwrap(), tree: Tensor::vector_uninit(backend, 0, usages).unwrap(), + refit_frontier: Tensor::vector_uninit(backend, 0, usages).unwrap(), + refit_frontier_len: Tensor::vector_uninit(backend, 0, usages).unwrap(), + refit_rounds: Tensor::scalar(backend, 0, usages | BufferUsages::UNIFORM).unwrap(), sort_workspace: RadixSortWorkspace::new(backend), aabbs: Tensor::vector_uninit(backend, 0, usages).unwrap(), buffer_usages: usages, @@ -116,6 +135,14 @@ impl LbvhState { self.generation += 1; self.domain_aabb = Tensor::vector_uninit(backend, num_batches, self.buffer_usages).unwrap(); + self.domain_partials = + Tensor::vector_uninit(backend, num_batches * DOMAIN_WORKGROUPS, self.buffer_usages) + .unwrap(); + } + // Its own check: it starts empty, while `domain_aabb` starts with one batch. + if (self.refit_frontier_len.len() as u32) < num_batches { + self.refit_frontier_len = + Tensor::vector_uninit(backend, num_batches, self.buffer_usages).unwrap(); } // NOTE: colliders_len is the total colliders count, already taking all batches into account. @@ -138,6 +165,8 @@ impl LbvhState { Tensor::vector_uninit(backend, colliders_len, self.buffer_usages).unwrap(); self.tree = Tensor::vector_uninit(backend, 2 * colliders_len, self.buffer_usages).unwrap(); + self.refit_frontier = + Tensor::vector_uninit(backend, colliders_len, self.buffer_usages).unwrap(); // FIXME: this doesn’t account for batches having mismatched numbers of colliders. // n_sort is a per-batch vector: each element is the per-batch *active* @@ -215,13 +244,23 @@ impl Lbvh { let colliders_per_batch = active_per_batch; let mut pass = encoder.begin_pass("[RBD] lbvh-compute-domain", timestamps.as_deref_mut()); + // About 4k colliders per workgroup. + let domain_workgroups = colliders_per_batch + .div_ceil(4096) + .clamp(1, DOMAIN_WORKGROUPS); self.shaders.compute_domain.call( &mut pass, - [1u32, num_batches, 1], + [domain_workgroups * 128, num_batches, 1], poses, - &mut state.domain_aabb, + &mut state.domain_partials, batch_indices, )?; + self.shaders.domain_merge.call( + &mut pass, + [DOMAIN_WORKGROUPS, num_batches, 1], + &state.domain_partials, + &mut state.domain_aabb, + )?; drop(pass); let mut pass = encoder.begin_pass("[RBD] lbvh-compute-morton", timestamps.as_deref_mut()); @@ -258,6 +297,12 @@ impl Lbvh { &mut state.tree, batch_indices, )?; + self.shaders.escapes.call( + &mut pass, + [(2 * colliders_per_batch).saturating_sub(1), num_batches, 1], + &mut state.tree, + batch_indices, + )?; drop(pass); let mut pass = encoder.begin_pass("[RBD] lbvh-refit_leaves", timestamps.as_deref_mut()); @@ -269,16 +314,37 @@ impl Lbvh { &state.sorted_colliders, &mut state.tree, batch_indices, + &mut state.refit_frontier_len, vertex_buffers, )?; drop(pass); let mut pass = encoder.begin_pass("[RBD] lbvh-refit-internal", timestamps); - self.shaders.refit_internal.call( + // One workgroup per chunk of sorted leaves, then one per batch for the top of the tree. + self.shaders.refit_chunks.call( + &mut pass, + [colliders_per_batch, num_batches, 1], + &mut state.tree, + &mut state.refit_frontier, + &mut state.refit_frontier_len, + batch_indices, + )?; + #[cfg(target_arch = "wasm32")] + self.shaders.refit_plan.call( + &mut pass, + 1u32, + &state.refit_frontier_len, + &mut state.refit_rounds, + batch_indices, + )?; + self.shaders.refit_frontier.call( &mut pass, [1u32, num_batches, 1], &mut state.tree, + &state.refit_frontier, + &state.refit_frontier_len, batch_indices, + &state.refit_rounds, )?; drop(pass); diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 6a53faa0..dfe131f9 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -123,7 +123,7 @@ impl GpuNarrowPhase { contacts: &mut Tensor, contacts_indirect: &mut Tensor<[u32; 3]>, contact_plan: &mut Tensor, - mb_sweep_indirect: &mut Tensor<[u32; 3]>, + mb_dispatch_indirect: &mut Tensor<[u32; 3]>, pfm_pairs: &mut Tensor, pfm_pairs_len: &mut Tensor, pfm_pairs_indirect: &mut Tensor<[u32; 3]>, @@ -170,8 +170,8 @@ impl GpuNarrowPhase { indices, )?; - // Clamped totals + every derived grid (contacts sweep, PFM sweep, - // multibody contact sweep) in one serial thread. + // Clamped totals + every derived grid (contacts dispatch, PFM dispatch, + // multibody contact dispatch) in one serial thread. self.shaders.contact_plan.call( pass, 1u32, @@ -181,7 +181,7 @@ impl GpuNarrowPhase { &mut pfm_sort.sort_len, contacts_indirect, pfm_pairs_indirect, - mb_sweep_indirect, + mb_dispatch_indirect, batch_indices, )?; diff --git a/src_rbd/dynamics/coloring.rs b/src_rbd/dynamics/coloring.rs index 211a88d5..d42fc1bf 100644 --- a/src_rbd/dynamics/coloring.rs +++ b/src_rbd/dynamics/coloring.rs @@ -9,15 +9,18 @@ use crate::pipeline::RunStats; use crate::shaders::broad_phase::ContactPlan; -use crate::shaders::dynamics::TwoBodyConstraint; use crate::shaders::dynamics::{ - GpuClearCompletionFlagTopoGc, GpuColorBucketsCount, GpuColorBucketsReset, - GpuColorBucketsScatter, GpuFixConflictsTopoGc, GpuResetCompletionFlagTopoGc, GpuResetLuby, - GpuResetTopoGc, GpuStepGraphColoringLuby, GpuStepGraphColoringTopoGc, + ColorStats, ContactLink, GpuClearCompletionFlagTopoGc, GpuColorBucketsCount, + GpuColorBucketsReset, GpuColorBucketsScatter, GpuColorDispatchGrid, GpuFixConflictsTopoGc, + GpuResetCompletionFlagTopoGc, GpuResetLuby, GpuResetTopoGc, GpuStepGraphColoringLuby, + GpuStepGraphColoringTopoGc, }; use crate::utils::{GpuPrefixSum, PrefixSumWorkspace}; use khal::Shader; -use khal::backend::{Backend, Encoder, GpuBackend, GpuBackendError, GpuPass, GpuTimestamps}; +use khal::backend::{ + AsGpuSlice, AsGpuSliceMut, Backend, Encoder, GpuBackend, GpuBackendError, GpuBufferSlice, + GpuBufferSliceMut, GpuPass, +}; use vortx::tensor::Tensor; /// GPU shaders for constraint graph coloring. @@ -37,6 +40,8 @@ pub struct GpuColoring { fix_conflicts_topo_gc_kernel: GpuFixConflictsTopoGc, reset_completion_flag_topo_gc: GpuResetCompletionFlagTopoGc, clear_completion_flag_topo_gc: GpuClearCompletionFlagTopoGc, + /// Sizes the colored dispatches from the color buckets. + color_dispatch_grid: GpuColorDispatchGrid, // Workspace for bucket-sorting constraint ids by color so each color iteration // only touches their own constraint. color_buckets_reset: GpuColorBucketsReset, @@ -44,24 +49,52 @@ pub struct GpuColoring { color_buckets_scatter: GpuColorBucketsScatter, } +/// The [`ColorStats`] of a step. Stored as words so the resize readback can gather them with the +/// other `u32` counters, and seen by the kernels as a [`ColorStats`]. +pub struct ColorStatsBuffer(pub(crate) Tensor); + +impl ColorStatsBuffer { + /// Number of `u32` words in a [`ColorStats`]. + pub(crate) const WORDS: usize = size_of::() / size_of::(); +} + +impl AsGpuSlice for ColorStatsBuffer { + fn as_gpu_slice(&self) -> GpuBufferSlice<'_, ColorStats> { + self.0.buffer_slice().reinterpret() + } +} + +impl AsGpuSliceMut for ColorStatsBuffer { + fn as_gpu_slice_mut(&mut self) -> GpuBufferSliceMut<'_, ColorStats> { + self.0.buffer_slice_mut().reinterpret() + } +} + /// Buffers for the per-color constraint bucket sort. pub struct ColorBucketsArgs<'a> { /// Flat dispatch grid over the whole contacts range. pub contacts_len_indirect: &'a Tensor<[u32; 3]>, /// Color assigned to each constraint by graph coloring. pub constraints_colors: &'a Tensor, - /// The colored constraints (batch recovered from their global body ids). - pub constraints: &'a Tensor, - /// Clamped per-frame list totals (the flat contact sweep bound). + /// The links of the colored constraints (batch recovered from their global body ids). + pub links: &'a Tensor, + /// Clamped per-frame list totals (the flat contact dispatch bound). pub contact_plan: &'a Tensor, /// The single `(color, batch)` bucket buffer, color-major /// (`solver_color_buckets_stride * num_batches` entries): counts, then /// scanned exclusive starts, then post-scatter exclusive ends. pub color_buckets: &'a mut Tensor, - /// Constraint ids bucket-sorted by `(color, batch)`. - pub color_sorted_ids: &'a mut Tensor, + /// Output: the constraint links bucket-sorted by `(color, batch)`, inactive past the last + /// constraint. + pub sorted_links: &'a mut Tensor, + /// Output: the index of each constraint in the color order. + pub constraint_indices: &'a mut Tensor, /// Shared per-batch capacity / section-offset uniform. pub batch_indices: &'a Tensor, + /// Output: `[largest bucket, highest color, size of colors 0..64]`. + pub color_stats: &'a mut ColorStatsBuffer, + /// Output: the grid of the colored dispatches. + pub dispatch_indirect: &'a mut Tensor<[u32; 3]>, } /// Arguments for graph coloring dispatch. @@ -74,8 +107,8 @@ pub struct ColoringArgs<'a> { pub body_constraint_counts: &'a Tensor, /// Constraint IDs associated with each body. pub body_constraint_ids: &'a Tensor, - /// The constraints to be colored. - pub constraints: &'a Tensor, + /// The links of the constraints to be colored. + pub links: &'a Tensor, /// Output: color assigned to each constraint. pub constraints_colors: &'a mut Tensor, /// Color picked by each constraint in the current round, before the conflict pass. @@ -89,7 +122,7 @@ pub struct ColoringArgs<'a> { pub uncolored: &'a mut Tensor, /// Staging buffer for reading uncolored count on CPU. pub uncolored_staging: &'a Tensor, - /// Clamped per-frame list totals (the flat contact sweep bound). + /// Clamped per-frame list totals (the flat contact dispatch bound). pub contact_plan: &'a Tensor, /// Buffer tracking which constraints are colored. pub colored: &'a mut Tensor, @@ -99,6 +132,8 @@ pub struct ColoringArgs<'a> { /// different bodies of the same multibody share a group and never share /// a color. For free bodies, `body_group[i] = i`. pub body_group: &'a Tensor, + /// Dispatch grid of the current topo-gc iteration (empty once converged). + pub coloring_indirect: &'a mut Tensor<[u32; 3]>, } impl GpuColoring { @@ -113,7 +148,7 @@ impl GpuColoring { args.contacts_len_indirect, args.constraints_colors, args.constraints_rands, - args.constraints, + args.links, args.contact_plan, )?; Ok(()) @@ -130,7 +165,7 @@ impl GpuColoring { args.contacts_len_indirect, args.body_constraint_counts, args.body_constraint_ids, - args.constraints, + args.links, args.constraints_colors, args.constraints_rands, args.uncolored, @@ -152,8 +187,9 @@ impl GpuColoring { args.contacts_len_indirect, args.constraints_colors, args.colored, - args.constraints, + args.links, args.contact_plan, + args.uncolored, args.constraints_pending_colors, )?; Ok(()) @@ -167,10 +203,10 @@ impl GpuColoring { ) -> Result<(), GpuBackendError> { self.step_graph_coloring_topo_gc_kernel.call( pass, - args.contacts_len_indirect, + &*args.coloring_indirect, args.body_constraint_counts, args.body_constraint_ids, - args.constraints, + args.links, args.constraints_colors, args.colored, args.uncolored, @@ -190,10 +226,10 @@ impl GpuColoring { ) -> Result<(), GpuBackendError> { self.fix_conflicts_topo_gc_kernel.call( pass, - args.contacts_len_indirect, + &*args.coloring_indirect, args.body_constraint_counts, args.body_constraint_ids, - args.constraints, + args.links, args.constraints_colors, args.colored, args.uncolored, @@ -257,7 +293,7 @@ impl GpuColoring { /// Bucket-sorts the constraint ids by `(color, batch)`: zero the buckets, /// count, exclusive-prefix-scan them in place, then scatter (which turns - /// the starts into exclusive ends, the form the sweeps read). + /// the starts into exclusive ends, the form the solve kernels read). pub fn dispatch_build_color_buckets( &self, backend: &GpuBackend, @@ -273,20 +309,31 @@ impl GpuColoring { pass, args.contacts_len_indirect, args.constraints_colors, - args.constraints, + args.links, args.contact_plan, args.color_buckets, args.batch_indices, + args.sorted_links, )?; prefix_sum.launch(backend, pass, prefix_workspace, args.color_buckets, 1)?; self.color_buckets_scatter.call( pass, args.contacts_len_indirect, args.constraints_colors, - args.constraints, + args.links, args.contact_plan, args.color_buckets, - args.color_sorted_ids, + args.sorted_links, + args.batch_indices, + args.constraint_indices, + )?; + // A single workgroup reduces over the colors. + self.color_dispatch_grid.call( + pass, + 64u32, + &*args.color_buckets, + args.color_stats, + args.dispatch_indirect, args.batch_indices, )?; Ok(()) @@ -324,14 +371,24 @@ impl GpuColoring { max_colors: u32, ) -> Result<(), GpuBackendError> { // The first fix-conflicts pass must validate the seeded colors even when the first step - // colors nothing, so it starts from a cleared ("not converged") flag. At least two rounds - // run so the last pass can still record the color count. - self.clear_completion_flag_topo_gc - .call(pass, 1u32, args.uncolored)?; + // colors nothing, so it starts from a cleared ("not converged") flag and a full grid. At + // least two rounds run so the last pass can still record the color count. + self.clear_completion_flag_topo_gc.call( + pass, + 1u32, + &mut *args.uncolored, + args.contacts_len_indirect, + &mut *args.coloring_indirect, + )?; for i in 0..max_colors.max(2) { if i > 0 { - self.reset_completion_flag_topo_gc - .call(pass, 1u32, args.uncolored)?; + self.reset_completion_flag_topo_gc.call( + pass, + 1u32, + &mut *args.uncolored, + args.contacts_len_indirect, + &mut *args.coloring_indirect, + )?; } self.dispatch_step_topo_gc(pass, &mut args)?; self.dispatch_fix_conflicts_topo_gc(pass, &mut args)?; @@ -339,61 +396,4 @@ impl GpuColoring { Ok(()) } - - /// Executes the TOPO-GC (Topological Graph Coloring) algorithm. - /// - /// Returns `Some(num_colors)` (1-indexed) on success, or `None` if convergence - /// fails (caller should fall back to [`dispatch_luby`](Self::dispatch_luby)). - pub async fn dispatch_topo_gc<'a>( - &self, - backend: &GpuBackend, - mut args: ColoringArgs<'a>, - stats: &mut RunStats, - mut timestamps: Option<&mut GpuTimestamps>, - ) -> Option { - // Initialize TOPO-GC state - { - let mut encoder = backend.begin_encoding(); - let mut pass = encoder.begin_pass("topo-gc-coloring-reset", timestamps.as_deref_mut()); - self.dispatch_reset_topo_gc(&mut pass, &mut args).unwrap(); - drop(pass); - backend.submit(encoder).unwrap(); - } - - let mut num_loops = 0; - loop { - num_loops += 1; - if num_loops > 64 { - return None; - } - - // Batch multiple iterations to reduce CPU-GPU sync overhead - { - let mut encoder = backend.begin_encoding(); - let mut pass = - encoder.begin_pass("topo-gc-coloring-step", timestamps.as_deref_mut()); - for _ in 0..10 { - self.reset_completion_flag_topo_gc - .call(&mut pass, 1u32, args.uncolored) - .unwrap(); - self.dispatch_step_topo_gc(&mut pass, &mut args).unwrap(); - self.dispatch_fix_conflicts_topo_gc(&mut pass, &mut args) - .unwrap(); - } - drop(pass); - backend.submit(encoder).unwrap(); - } - - let max_color = backend - .slow_read_vec(args.uncolored.buffer()) - .await - .unwrap()[0]; - - if max_color != 0 { - stats.num_colors = max_color; - stats.coloring_iterations = num_loops; - return Some(max_color + 1); // NOTE: color indices are 1-based. - } - } - } } diff --git a/src_rbd/dynamics/contact_kernels.rs b/src_rbd/dynamics/contact_kernels.rs new file mode 100644 index 00000000..ee737ed8 --- /dev/null +++ b/src_rbd/dynamics/contact_kernels.rs @@ -0,0 +1,336 @@ +//! Host dispatch over the contact tiles of the simulation's friction model. Only the selected +//! model's tiles are allocated; kernel selection happens once per dispatch on the CPU. +use crate::math::Pose; +use crate::shaders::broad_phase::ContactPlan; +use crate::shaders::dynamics::contact_tiles::{TILE_LEN, tile_lane}; +use crate::shaders::dynamics::coulomb_tiles::CoulombTile; +#[cfg(feature = "dim3")] +use crate::shaders::dynamics::twist_tiles::TwistTile; +use crate::shaders::dynamics::*; +use crate::shaders::queries::IndexedManifold; +use crate::shaders::utils::BatchIndices; +use khal::backend::{Backend, DispatchGrid, GpuBackend, GpuBackendError, GpuPass}; +use vortx::tensor::Tensor; + +/// The color-ordered tiles of the contact constraints, in the layout of the simulation's +/// friction model. +pub enum ContactTiles { + /// Friction at each contact point. + Coulomb(Tensor), + #[cfg(feature = "dim3")] + /// Central sliding friction and twist friction per manifold. + Simplified(Tensor), +} + +impl ContactTiles { + /// Uninitialized tiles for `capacity` constraints of `params`' friction model. + pub(crate) fn new( + backend: &GpuBackend, + capacity: u32, + params: &RbdSimParams, + ) -> Result { + let _ = params; + let tiles = capacity.div_ceil(TILE_LEN as u32).max(1); + let usage = khal::BufferUsages::STORAGE | khal::BufferUsages::COPY_SRC; + #[cfg(feature = "dim3")] + if params.friction_model == FrictionModel::Simplified { + return Ok(Self::Simplified(Tensor::vector_uninit( + backend, tiles, usage, + )?)); + } + Ok(Self::Coulomb(Tensor::vector_uninit(backend, tiles, usage)?)) + } + + /// The Coulomb tiles, if this simulation uses Coulomb friction. + pub fn coulomb(&self) -> Option<&Tensor> { + match self { + Self::Coulomb(x) => Some(x), + #[cfg(feature = "dim3")] + _ => None, + } + } + + #[cfg(feature = "dim3")] + /// The twist tiles, if this simulation uses simplified friction. + pub fn simplified(&self) -> Option<&Tensor> { + match self { + Self::Simplified(x) => Some(x), + _ => None, + } + } +} + +/// The contact constraints of one step: their links in contact order, and their solver data +/// sorted by color. +pub struct ContactConstraints { + /// The topology and identity of each contact's constraint. + pub links: Tensor, + /// The index of each active constraint in `tiles`. + pub constraint_indices: Tensor, + /// The constraint tiles. + pub tiles: ContactTiles, +} + +impl ContactConstraints { + /// Uninitialized constraints for `capacity` contact slots of `params`' friction model. + pub(crate) fn new( + backend: &GpuBackend, + capacity: u32, + usage: khal::BufferUsages, + params: &RbdSimParams, + ) -> Result { + Ok(Self { + links: Tensor::vector_uninit(backend, capacity, usage)?, + constraint_indices: Tensor::vector_uninit(backend, capacity, usage)?, + tiles: ContactTiles::new(backend, capacity, params)?, + }) + } + + /// Reads the constraints of the first `bound` contact slots, in contact order. This is a + /// blocking GPU synchronization when driven to completion; keep it outside timed runs. + pub async fn read_impulses( + &self, + backend: &GpuBackend, + bound: usize, + ) -> Result, GpuBackendError> { + self.tiles + .read_impulses(backend, &self.links, &self.constraint_indices, bound) + .await + } +} + +// Every tile argument is checked before launch; the old and new tiles always change model +// together when the simulation-wide model is changed. +macro_rules! tile_kernel { + ($name:ident, $primary:ident, [$($other:ident),*], ($($arg:ident : $ty:ty),* $(,)?)) => { + pub(crate) struct $name { + coulomb: crate::shaders::dynamics::coulomb_tiles::$name, + #[cfg(feature="dim3")] + simplified: crate::shaders::dynamics::twist_tiles::$name, + } + impl $name { + #[cfg(feature="cpu")] + pub const __ERROR__SHADER_CRATE_IS_MISSING_FEATURE_NAMED____CPU: () = crate::shaders::dynamics::coulomb_tiles::$name::__ERROR__SHADER_CRATE_IS_MISSING_FEATURE_NAMED____CPU; + #[cfg(feature="cpu-parallel")] + pub const __ERROR__SHADER_CRATE_IS_MISSING_FEATURE_NAMED____CPU_PARALLEL: () = crate::shaders::dynamics::coulomb_tiles::$name::__ERROR__SHADER_CRATE_IS_MISSING_FEATURE_NAMED____CPU_PARALLEL; + #[cfg(feature="cuda")] + pub const __ERROR__SHADER_CRATE_IS_MISSING_FEATURE_NAMED____CUDA: () = crate::shaders::dynamics::coulomb_tiles::$name::__ERROR__SHADER_CRATE_IS_MISSING_FEATURE_NAMED____CUDA; + pub fn from_dir(backend: &GpuBackend, dir: &include_dir::Dir<'static>) -> Result { + Ok(Self { + coulomb: crate::shaders::dynamics::coulomb_tiles::$name::from_dir(backend,dir)?, + #[cfg(feature="dim3")] + simplified: crate::shaders::dynamics::twist_tiles::$name::from_dir(backend,dir)?, + }) + } + #[allow(clippy::too_many_arguments)] + pub fn call<'a>(&self, pass: &mut GpuPass, grid: impl Into>, $($arg:$ty),*) -> Result<(),GpuBackendError> { + let grid = grid.into(); + match $primary { + ContactTiles::Coulomb($primary) => { + $(let $other = $other.coulomb().expect("tile models agree");)* + self.coulomb.call(pass,grid,$($arg),*) + } + #[cfg(feature="dim3")] + ContactTiles::Simplified($primary) => { + $(let $other = $other.simplified().expect("tile models agree");)* + self.simplified.call(pass,grid,$($arg),*) + } + } + } + } + }; +} + +tile_kernel!(GpuPrepareConstraints, tiles, [old_tiles], ( + sorted_links: &Tensor, + contacts: &Tensor, + mprops: &Tensor, + solver_poses: &Tensor, + vels: &Tensor, + states: &Tensor, + old_tiles: &ContactTiles, + tiles: &mut ContactTiles, + plan: &Tensor, + recycle_offsets: &Tensor, +)); +tile_kernel!(GpuScaleConstraintImpulses, tiles, [], ( + buckets: &Tensor, + tiles: &mut ContactTiles, + ids: &Tensor, + params: &Tensor, +)); + +/// Declares the wrappers of the solve kernels, which share their arguments. +macro_rules! solve_constraints_kernels { + ($($name:ident),*) => {$( + tile_kernel!($name, tiles, [], ( + tiles: &mut ContactTiles, + solver_vels: &mut Tensor, + buckets: &Tensor, + solver_body_poses: &Tensor, + color: &Tensor, + ids: &Tensor, + mode: &Tensor, + params: &Tensor, + )); + )*}; +} +solve_constraints_kernels!( + GpuSolveConstraints, + GpuSolveConstraintsBiased, + GpuSolveConstraintsUnbiased, + GpuSolveConstraintsUnbiasedCached, + GpuSolveConstraintsTail, + GpuSolveConstraintsTailBiased, + GpuSolveConstraintsTailUnbiased, + GpuSolveConstraintsTailUnbiasedCached, + GpuSolveConstraintsFused, + GpuSolveConstraintsFusedCached +); +tile_kernel!(GpuWarmstartConstraints, tiles, [], ( + tiles: &ContactTiles, + solver_vels: &mut Tensor, + buckets: &Tensor, + curr_color: &Tensor, + ids: &Tensor, +)); +tile_kernel!(GpuWarmstartConstraintsFused, tiles, [], ( + tiles: &ContactTiles, + solver_vels: &mut Tensor, + buckets: &Tensor, + num_colors: &Tensor, + ids: &Tensor, +)); +tile_kernel!(GpuGatherWarmstartVelocities, tiles, [], ( + counts: &Tensor, + constraint_ids: &Tensor, + constraint_indices: &Tensor, + tiles: &ContactTiles, + solver_vels: &mut Tensor, + hub_first_slot: &Tensor, + ids: &Tensor, +)); +tile_kernel!(GpuHubWarmstartConstraints, tiles, [], ( + solver_vels: &mut Tensor, + tiles: &ContactTiles, + hub_slot_constraint: &Tensor, + hub_counts: &Tensor, + constraint_indices: &Tensor, +)); + +/// CPU diagnostic view; tangent impulses are repeated per point for the +/// simplified model, matching Rapier's contact-data convention. +pub struct ContactImpulseSnapshot { + /// Global solver index of the first body. + pub body_a: u32, + /// Global solver index of the second body. + pub body_b: u32, + /// Number of active contact points. + pub len: u32, + /// World normal force direction on the first body. + pub dir_a: crate::math::Vector, + /// Combined friction coefficient. + pub friction: f32, + /// Axis-dependent inverse mass of the first body. + pub inv_mass_a: crate::math::Vector, + /// Axis-dependent inverse mass of the second body. + pub inv_mass_b: crate::math::Vector, + /// Accumulated normal impulse per contact. + pub normal_impulse: Vec, + /// Per-contact sliding impulses (central impulse repeated for simplified friction). + pub tangent_impulse: Vec<[f32; SUB_LEN]>, + /// Accumulated manifold twist impulse; zero for Coulomb. + pub twist_impulse: f32, +} + +impl ContactImpulseSnapshot { + /// The snapshot of a gap, inert or multibody-owned slot. + fn inactive(link: &ContactLink) -> Self { + Self { + body_a: link.solver_body_a, + body_b: link.solver_body_b, + len: 0, + dir_a: Default::default(), + friction: 0.0, + inv_mass_a: Default::default(), + inv_mass_b: Default::default(), + normal_impulse: vec![], + tangent_impulse: vec![], + twist_impulse: 0.0, + } + } +} + +impl ContactTiles { + async fn read_impulses( + &self, + backend: &GpuBackend, + links: &Tensor, + constraint_indices: &Tensor, + bound: usize, + ) -> Result, GpuBackendError> { + let links: Vec = backend.slow_read_vec(links.buffer()).await?; + let constraint_indices: Vec = + backend.slow_read_vec(constraint_indices.buffer()).await?; + let mut out = Vec::new(); + match self { + Self::Coulomb(buffer) => { + let tiles: Vec = backend.slow_read_vec(buffer.buffer()).await?; + for (link, index) in links.iter().zip(constraint_indices).take(bound) { + if link.len == 0 { + out.push(ContactImpulseSnapshot::inactive(link)); + continue; + } + let (tile, lane) = tile_lane(index as usize); + let c = tiles[tile].constraint(lane); + let n = c.len as usize; + #[cfg(feature = "dim3")] + let tangent = c.points[..n] + .iter() + .map(|p| p.tangent_impulse.to_array()) + .collect(); + #[cfg(feature = "dim2")] + let tangent = c.points[..n].iter().map(|p| p.tangent_impulse).collect(); + out.push(ContactImpulseSnapshot { + body_a: c.solver_body_a, + body_b: c.solver_body_b, + len: c.len, + dir_a: c.dir_a, + friction: c.limit, + inv_mass_a: c.im_a, + inv_mass_b: c.im_b, + normal_impulse: c.points[..n].iter().map(|p| p.normal_impulse).collect(), + tangent_impulse: tangent, + twist_impulse: 0.0, + }); + } + } + #[cfg(feature = "dim3")] + Self::Simplified(buffer) => { + let tiles: Vec = backend.slow_read_vec(buffer.buffer()).await?; + for (link, index) in links.iter().zip(constraint_indices).take(bound) { + if link.len == 0 { + out.push(ContactImpulseSnapshot::inactive(link)); + continue; + } + let (tile, lane) = tile_lane(index as usize); + let c = tiles[tile].constraint(lane); + let n = c.len as usize; + out.push(ContactImpulseSnapshot { + body_a: c.solver_body_a, + body_b: c.solver_body_b, + len: c.len, + dir_a: c.dir_a, + friction: c.limit, + inv_mass_a: c.im_a, + inv_mass_b: c.im_b, + normal_impulse: c.points[..n].iter().map(|p| p.normal_impulse).collect(), + tangent_impulse: vec![c.friction.tangent_impulse.to_array(); n], + twist_impulse: c.friction.twist_impulse, + }); + } + } + } + Ok(out) + } +} diff --git a/src_rbd/dynamics/contact_recycling.rs b/src_rbd/dynamics/contact_recycling.rs new file mode 100644 index 00000000..dfb87d5e --- /dev/null +++ b/src_rbd/dynamics/contact_recycling.rs @@ -0,0 +1,68 @@ +//! Double-buffered recycling metadata with one storage binding. +use crate::shaders::dynamics::{ContactRecycleOffsets, ContactRecycleState}; +use khal::BufferUsages; +use khal::backend::{GpuBackend, GpuBackendError}; +use vortx::tensor::Tensor; + +/// Previous and current contact metadata in disjoint halves of one allocation. +/// Constraint layouts are unchanged. Swapping frames only selects an immutable +/// offset uniform, so captured GPU work never sees a rewritten offset buffer. +pub struct ContactRecycleStates { + storage: Tensor, + offsets: [Tensor; 2], + phase: usize, +} + +impl ContactRecycleStates { + /// Allocate equally sized previous and current metadata regions. + pub fn new( + backend: &GpuBackend, + capacity: u32, + usage: BufferUsages, + ) -> Result { + Ok(Self { + storage: Tensor::vector_uninit( + backend, + capacity + .checked_mul(2) + .expect("recycling capacity overflow"), + usage, + )?, + offsets: [ + Tensor::scalar( + backend, + ContactRecycleOffsets { + old_base: 0, + new_base: capacity, + ..Default::default() + }, + BufferUsages::UNIFORM, + )?, + Tensor::scalar( + backend, + ContactRecycleOffsets { + old_base: capacity, + new_base: 0, + ..Default::default() + }, + BufferUsages::UNIFORM, + )?, + ], + phase: 0, + }) + } + + pub(crate) fn bindings( + &mut self, + ) -> ( + &mut Tensor, + &Tensor, + ) { + (&mut self.storage, &self.offsets[self.phase]) + } + + /// Advance alongside the old/new constraint-buffer swap after each step. + pub fn swap_frames(&mut self) { + self.phase ^= 1; + } +} diff --git a/src_rbd/dynamics/joint.rs b/src_rbd/dynamics/joint.rs index 39fe7111..777b20fb 100644 --- a/src_rbd/dynamics/joint.rs +++ b/src_rbd/dynamics/joint.rs @@ -309,7 +309,7 @@ impl GpuImpulseJointSet { // Build the flat joint buffer [num_batches * max_joints], padded with // zeroed joints and batch-interleaved (`joint * num_batches + batch`) - // so the solver's flat sweeps put the same joint of consecutive + // so the solver's flat dispatches put the same joint of consecutive // batches on adjacent lanes. let dummy_joint = ImpulseJoint::zeroed(); let mut all_joints = vec![dummy_joint; num_batches as usize * max_joints as usize]; diff --git a/src_rbd/dynamics/mass_splitting.rs b/src_rbd/dynamics/mass_splitting.rs new file mode 100644 index 00000000..520716bb --- /dev/null +++ b/src_rbd/dynamics/mass_splitting.rs @@ -0,0 +1,204 @@ +//! Mass splitting of bodies with many contacts ("hubs"). +//! +//! See the `mass_splitting` shader module: a hub body is split into one sub-body per contact, +//! each with its own solver velocity slot, so its contacts don't all need different colors. + +use super::contact_kernels::{ContactTiles, GpuHubWarmstartConstraints}; +use crate::shaders::dynamics::{ + ContactLink, GpuHubAssign, GpuHubAverage, GpuHubDispatch, GpuHubScatter, + GpuHubSplitConstraints, HUB_MIN_CONSTRAINTS, HubCounts, NOT_A_HUB, Velocity, + WorldMassProperties, +}; +use crate::shaders::utils::BatchIndices; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::{BufferUsages, Shader}; +use vortx::tensor::Tensor; + +/// Buffers splitting bodies with many contacts into one sub-body per contact. +pub struct HubState { + /// Per body: its first sub-body slot, or `NOT_A_HUB`. + pub(crate) first_slot: Tensor, + /// The split bodies of the current step. + list: Tensor, + /// Per sub-body slot: the split body. + slot_body: Tensor, + /// Per sub-body slot: the constraint acting on it. + slot_constraint: Tensor, + /// The hub counters of the current step. + pub(crate) counts: Tensor, + /// Dispatch grid over the allocated sub-body slots. + slots_indirect: Tensor<[u32; 3]>, + /// Dispatch grid of the averaging kernel (one workgroup per hub). + average_indirect: Tensor<[u32; 3]>, +} + +impl HubState { + /// Number of sub-body slots for a scene with `num_body_slots` body slots. + fn pool(num_body_slots: usize) -> usize { + num_body_slots.max(1024) + } + + pub(crate) fn new(backend: &GpuBackend, num_body_slots: usize) -> Self { + let storage = BufferUsages::STORAGE | BufferUsages::COPY_SRC; + let indirect = BufferUsages::STORAGE | BufferUsages::INDIRECT; + let pool = Self::pool(num_body_slots); + let max_hubs = pool / HUB_MIN_CONSTRAINTS as usize + 1; + Self { + first_slot: Tensor::vector(backend, vec![NOT_A_HUB; num_body_slots], storage).unwrap(), + list: Tensor::vector(backend, vec![0u32; max_hubs], storage).unwrap(), + slot_body: Tensor::vector(backend, vec![0u32; pool], storage).unwrap(), + slot_constraint: Tensor::vector(backend, vec![0u32; pool], storage).unwrap(), + counts: Tensor::scalar( + backend, + HubCounts { + hubs: 0, + slots: 0, + base: num_body_slots as u32, + pool: pool as u32, + }, + storage, + ) + .unwrap(), + slots_indirect: Tensor::scalar(backend, [0u32, 1, 1], indirect).unwrap(), + average_indirect: Tensor::scalar(backend, [0u32, 1, 1], indirect).unwrap(), + } + } + + /// The solver velocity buffer: `vels` followed by the sub-body slots. + pub(crate) fn solver_vels( + backend: &GpuBackend, + vels: &[Velocity], + usage: BufferUsages, + ) -> Tensor { + let mut all = vels.to_vec(); + all.resize(vels.len() + Self::pool(vels.len()), Velocity::default()); + Tensor::vector(backend, &all, usage).unwrap() + } +} + +/// Kernels of the mass splitting. +#[derive(Shader)] +pub struct GpuMassSplitting { + assign: GpuHubAssign, + dispatch: GpuHubDispatch, + split_constraints: GpuHubSplitConstraints, + scatter: GpuHubScatter, + warmstart: GpuHubWarmstartConstraints, + average: GpuHubAverage, +} + +/// Inputs of [`GpuMassSplitting::split`]. +pub struct SplitArgs<'a> { + /// Number of body slots (all batches). + pub num_body_slots: u32, + /// Cumulative per-body constraint counts (each body's list end). + pub body_constraint_counts: &'a Tensor, + /// Per-body constraint lists. + pub body_constraint_ids: &'a Tensor, + /// The contact links of the step. + pub links: &'a mut Tensor, + /// World-space mass properties. + pub mprops: &'a Tensor, + /// Per-body graph-coloring group (multibody links are never split). + pub body_group: &'a Tensor, + /// Shared per-batch indices. + pub batch_indices: &'a Tensor, +} + +impl GpuMassSplitting { + /// Picks the hubs and points their constraint links at the sub-bodies. Runs once per step, + /// after the per-body constraint lists are built and before the coloring. + pub fn split( + &self, + pass: &mut GpuPass, + hubs: &mut HubState, + args: SplitArgs<'_>, + ) -> Result<(), GpuBackendError> { + self.assign.call( + pass, + args.num_body_slots, + args.body_constraint_counts, + args.mprops, + args.body_group, + &mut hubs.first_slot, + &mut hubs.list, + &mut hubs.counts, + args.batch_indices, + )?; + self.dispatch.call( + pass, + 1u32, + &hubs.counts, + &mut hubs.slots_indirect, + &mut hubs.average_indirect, + )?; + self.split_constraints.call( + pass, + &hubs.average_indirect, + args.body_constraint_counts, + args.body_constraint_ids, + &hubs.first_slot, + &mut *args.links, + &mut hubs.slot_body, + &mut hubs.slot_constraint, + &hubs.counts, + &hubs.list, + )?; + Ok(()) + } + + /// Copies each hub's velocity into its sub-bodies, before a contact iteration. + pub fn scatter( + &self, + pass: &mut GpuPass, + hubs: &HubState, + solver_vels: &mut Tensor, + ) -> Result<(), GpuBackendError> { + self.scatter.call( + pass, + &hubs.slots_indirect, + solver_vels, + &hubs.slot_body, + &hubs.counts, + ) + } + + /// Applies each hub constraint's warmstart impulse to its sub-body. + pub fn warmstart( + &self, + pass: &mut GpuPass, + hubs: &HubState, + solver_vels: &mut Tensor, + tiles: &ContactTiles, + constraint_indices: &Tensor, + ) -> Result<(), GpuBackendError> { + self.warmstart.call( + pass, + &hubs.slots_indirect, + solver_vels, + tiles, + &hubs.slot_constraint, + &hubs.counts, + constraint_indices, + ) + } + + /// Averages the sub-body velocities back into their hub, after a contact iteration. + pub fn average( + &self, + pass: &mut GpuPass, + hubs: &HubState, + solver_vels: &mut Tensor, + body_constraint_counts: &Tensor, + ) -> Result<(), GpuBackendError> { + self.average.call( + pass, + &hubs.average_indirect, + solver_vels, + &hubs.list, + &hubs.first_slot, + body_constraint_counts, + &hubs.counts, + ) + } +} diff --git a/src_rbd/dynamics/mod.rs b/src_rbd/dynamics/mod.rs index 16146a5f..c22ced8d 100644 --- a/src_rbd/dynamics/mod.rs +++ b/src_rbd/dynamics/mod.rs @@ -1,10 +1,15 @@ //! Rigid-body dynamics: forces, velocities, constraints, and solvers. +#[cfg(feature = "dim3")] +pub use crate::shaders::dynamics::FrictionModel; pub use crate::shaders::dynamics::RbdSimParams; pub use body::{BodyCoupling, BodyCouplingEntry, BodyDesc, GpuBodySet}; pub use canonical_order::{CanonicalContactsArgs, GpuCanonicalOrder, contact_sort_collider_shift}; -pub use coloring::{ColorBucketsArgs, ColoringArgs, GpuColoring}; +pub use coloring::{ColorBucketsArgs, ColorStatsBuffer, ColoringArgs, GpuColoring}; +pub use contact_kernels::{ContactConstraints, ContactImpulseSnapshot, ContactTiles}; +pub use contact_recycling::ContactRecycleStates; pub use joint::{GpuImpulseJointSet, GpuJointSolver, JointSolverArgs, convert_joint_motor}; +pub use mass_splitting::{GpuMassSplitting, HubState, SplitArgs}; pub use mprops_update::{GpuMpropsUpdate, GpuSyncColliderPosesShader}; #[cfg(feature = "dim3")] pub use multibody::{ @@ -12,12 +17,15 @@ pub use multibody::{ }; pub use prep_render::{RbdInstanceDesc, WgRbdPrepRender}; pub use solver::{GpuSolver, SolverArgs}; -pub use warmstart::{GpuWarmstart, WarmstartArgs}; +pub use warmstart::{GpuWarmstart, SeedColorsArgs}; pub mod body; mod canonical_order; mod coloring; +pub(crate) mod contact_kernels; +mod contact_recycling; mod joint; +mod mass_splitting; mod mprops_update; #[cfg(feature = "dim3")] pub(crate) mod multibody; diff --git a/src_rbd/dynamics/multibody/loop_closing_joints.rs b/src_rbd/dynamics/multibody/loop_closing_joints.rs index a301767b..9f7d3382 100644 --- a/src_rbd/dynamics/multibody/loop_closing_joints.rs +++ b/src_rbd/dynamics/multibody/loop_closing_joints.rs @@ -175,7 +175,7 @@ impl GpuMultibodySet { // the two sides of every joint. Two joints get the same color // only if they share no node, so within a color every joint // writes disjoint `dof_state` / `solver_vels`, making the - // per-color sweep an exact (race-free) Gauss–Seidel step. + // per-color iteration an exact (race-free) Gauss–Seidel step. let num_mb = mb_set.multibodies().count() as u32; // Unified node id: MB side → mb_idx; free body → num_mb + // local_body_id; FIXED → none. diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index cd4ef8a4..3703ad32 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -717,7 +717,7 @@ impl GpuMultibodySet { ) .unwrap(), // Per-multibody Delassus blocks for the constraint-space contact - // sweep: MAX_MB_CONTACT_CONSTRAINTS_PER_MB² floats each (147 KB + // iteration: MAX_MB_CONTACT_CONSTRAINTS_PER_MB² floats each (147 KB // in 3D), so only small total multibody counts get them; larger // batched scenes keep the dof-space solve. contact_delassus: { diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 1c7a2373..25226e0d 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -210,7 +210,7 @@ pub struct GpuMultibodySet { /// Per-batch slot counts for the interleaved impulse-joint buffers above /// (`mb_imp_joints_per_batch` is mirrored into `BatchIndices` as the flat - /// sweeps' loop bound; the others only size the buffers). + /// iterations' loop bound; the others only size the buffers). pub(super) mb_imp_joints_per_batch: u32, pub(super) mb_imp_joint_constraints_per_batch: u32, pub(super) mb_imp_joint_jacobians_per_batch: u32, @@ -312,7 +312,7 @@ impl GpuMultibodySet { self.num_active_multibodies == 0 || self.links_per_batch == 0 } - /// Number of colors used by the colored multibody impulse-joint sweeps. + /// Number of colors used by the colored multibody impulse-joint iterations. pub fn mb_imp_joint_num_colors(&self) -> u32 { self.mb_imp_joint_num_colors } diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index c03218a3..c12b5812 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -47,13 +47,13 @@ pub struct GpuMultibodySolver { /// Advances the actuator-delay step counter, once per physics step. delay_tick: GpuMbDelayTick, /// Contact force-sensor readout, dispatched once per step after the last - /// substep's stabilization sweep and only when sensors are configured. + /// substep's stabilization iteration and only when sensors are configured. sense_contact_impulses: GpuMbSenseContactImpulses, finalize_contact_constraints: GpuMbFinalizeContactConstraints, /// Fused joint+contact PGS iteration (one workgroup per multibody, shared- /// memory dof velocities). solve_constraints: GpuMbSolveConstraints, - /// Joint-only half of the sweep, used with the Delassus contact path + /// Joint-only half of the iteration, used with the Delassus contact path /// (one kernel binding both joint and Delassus buffers would exceed the /// 8-storage-buffer budget). solve_joints: GpuMbSolveJoints, @@ -61,9 +61,9 @@ pub struct GpuMultibodySolver { /// coupling) right after the contact columns are finalized. build_contact_delassus: GpuMbBuildContactDelassus, /// Reduces the per-multibody contact-constraint counts to their maximum, - /// the trip count of the `web-compat` contact sweeps. + /// the trip count of the `web-compat` contact iterations. compute_solve_bounds: GpuMbComputeSolveBounds, - /// Constraint-space contact sweep: `a = J·u` tracked incrementally in + /// Constraint-space contact iteration: `a = J·u` tracked incrementally in /// shared memory via the Delassus rows, breaking the per-iteration /// dof-space latency chain. solve_contacts_delassus: GpuMbSolveContactsDelassus, @@ -73,7 +73,7 @@ pub struct GpuMultibodySolver { /// layout is computed (the warmstart transfer matches against them). save_prev_cons_bounds: GpuMbSavePrevConsBounds, /// Predicts the per-(multibody, batch) contact-constraint slot and index - /// counts (one flat sweep over the contacts). + /// counts (one flat dispatch over the contacts). count_contact_constraints: GpuMbCountContactConstraints, /// Turns the predictions into dynamic segments of the flat constraint /// buffer and of the contact→multibody index (and publishes the total @@ -116,7 +116,7 @@ pub struct MultibodySolverArgs<'a> { /// Flat contact manifold list (filled by narrow-phase; positional slots). pub contacts: &'a Tensor, /// Clamped per-frame list totals (see `gpu_contact_plan`); `[PLAN_BOUND]` - /// bounds the flat contact sweeps. + /// bounds the flat contact dispatches. pub contact_plan: &'a Tensor, /// Flat dispatch grid over the whole contacts range (written by /// `gpu_contact_plan`). @@ -140,7 +140,7 @@ pub struct MultibodySolverArgs<'a> { pub friction_in_bias_pass: bool, /// GPU-written workgroup grid for the per-multibody contact-constraint /// dispatches: `[multibodies_batch_capacity, num_batches, 1]`. - pub mb_sweep_indirect: &'a Tensor<[u32; 3]>, + pub mb_dispatch_indirect: &'a Tensor<[u32; 3]>, } impl GpuMultibodySolver { @@ -200,7 +200,7 @@ impl GpuMultibodySolver { /// Lay out this frame's dynamic contact-constraint segments and build the /// contact→multibody index, once per step after the narrow phase: - /// predict the per-(multibody, batch) slot / index counts (one flat sweep + /// predict the per-(multibody, batch) slot / index counts (one flat dispatch /// over the contacts), prefix-scan them into segment starts (also /// publishing the total demand for the auto-resize readback), then /// scatter each contact's slot into its owner's index segment. @@ -339,7 +339,7 @@ impl GpuMultibodySolver { encoder.begin_pass("[RBD] mbb/transfer-warmstart", timestamps.as_deref_mut()); self.transfer_contact_warmstart.call( &mut pass, - args.mb_sweep_indirect, + args.mb_dispatch_indirect, &mb.multibody_info, &mut mb.contact_constraints, &mb.old_contact_constraints, @@ -354,7 +354,7 @@ impl GpuMultibodySolver { encoder.begin_pass("[RBD] mbb/seed-restitution", timestamps.as_deref_mut()); self.seed_contact_restitution.call( &mut pass, - args.mb_sweep_indirect, + args.mb_dispatch_indirect, &mb.multibody_info, &mut mb.contact_constraints, &mb.contact_jac_cols, @@ -374,7 +374,7 @@ impl GpuMultibodySolver { // when no batch has any contact this step. self.warmstart_contact_constraints.call( &mut pass, - args.mb_sweep_indirect, + args.mb_dispatch_indirect, &mb.multibody_info, &mb.contact_constraints, &mb.contact_jac_cols, @@ -478,7 +478,7 @@ impl GpuMultibodySolver { encoder.begin_pass("[RBD] mbb/finalize-contact", timestamps.as_deref_mut()); self.finalize_contact_constraints.call( &mut pass, - args.mb_sweep_indirect, + args.mb_dispatch_indirect, &mb.multibody_info, &mb.mass_matrices, &mb.lu_pivots, @@ -501,13 +501,13 @@ impl GpuMultibodySolver { )?; } - // Delassus blocks for the constraint-space contact sweep (consumes + // Delassus blocks for the constraint-space contact iteration (consumes // the columns finalized just above). if let Some(delassus) = &mut mb.contact_delassus { let mut pass = encoder.begin_pass("[RBD] mbb/build-delassus", timestamps); self.build_contact_delassus.call( &mut pass, - args.mb_sweep_indirect, + args.mb_dispatch_indirect, &mb.multibody_info, &mb.contact_constraints, &mb.contact_jac_cols, @@ -549,7 +549,7 @@ impl GpuMultibodySolver { // on contact-free steps. self.solve_contacts_delassus.call( pass, - args.mb_sweep_indirect, + args.mb_dispatch_indirect, &mb.multibody_info, &mut mb.contact_constraints, &mb.contact_jac_cols, @@ -576,12 +576,12 @@ impl GpuMultibodySolver { args.solver_vels, )?; } else { - // No joint limits/motors anywhere: the fused sweep is contact-only + // No joint limits/motors anywhere: the fused iteration is contact-only // work, so the indirect grid (zero workgroups on contact-free // steps) replaces the full per-(multibody, batch) launch. self.solve_constraints.call( pass, - args.mb_sweep_indirect, + args.mb_dispatch_indirect, &mb.multibody_info, &mut mb.joint_constraints, &mb.joint_constraint_columns, @@ -599,7 +599,7 @@ impl GpuMultibodySolver { /// P3: one PGS iteration with bias over the joint, contact, and multibody- /// touching impulse-joint constraints. The rigid-body solver interleaves - /// its own sweep after each call, `num_internal_pgs_iterations` times per + /// its own iteration after each call, `num_internal_pgs_iterations` times per /// substep; `first_iteration` (re)builds the impulse-joint constraints. pub fn substep_solve_with_bias( &self, @@ -628,7 +628,7 @@ impl GpuMultibodySolver { // Multibody-touching impulse joints — generic (rb-mb / mb-mb) // constraints, built on the first iteration of the substep. if mb.mb_imp_joints_per_batch > 0 && first_iteration { - // Flat 1-D sweep over the interleaved joint slots. + // Flat 1-D dispatch over the interleaved joint slots. let imp_dispatch = [mb.mb_imp_joints_per_batch * mb.num_batches, 1, 1]; self.update_impulse_joint_constraints.call( pass, @@ -742,11 +742,11 @@ impl GpuMultibodySolver { return Ok(()); } - // Stabilization sweep: `use_bias = 0` (`color_uniforms[0] == 0`). + // Stabilization iteration: `use_bias = 0` (`color_uniforms[0] == 0`). let solve_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.dispatch_solve(pass, mb, args, solve_dispatch, 0)?; if mb.mb_imp_joints_per_batch > 0 { - // Flat 1-D sweep over the interleaved joint slots. + // Flat 1-D dispatch over the interleaved joint slots. let imp_dispatch = [mb.mb_imp_joints_per_batch * mb.num_batches, 1, 1]; self.remove_impulse_joint_constraint_bias.call( pass, @@ -755,8 +755,8 @@ impl GpuMultibodySolver { &mut mb.mb_imp_joint_constraints, args.batch_indices, )?; - // Final stabilization sweep WITHOUT bias — colored, one - // dispatch per color (see the with-bias sweep above). + // Final stabilization iteration WITHOUT bias — colored, one + // dispatch per color (see the with-bias iteration above). for c in 0..mb.mb_imp_joint_num_colors as usize { self.solve_impulse_joint_constraints.call( pass, @@ -786,7 +786,7 @@ impl GpuMultibodySolver { /// Contact force-sensor readout: folds each sensed link's normal-contact /// impulses into `contact_sensor_out`. Run it once per step, after the last - /// substep's stabilization sweep and before [`Self::apply_restitution`], so + /// substep's stabilization iteration and before [`Self::apply_restitution`], so /// the value is the accumulated contact impulse rather than a /// restitution-adjusted one. A no-op when no sensors are configured. pub fn sense_contact_impulses( @@ -821,7 +821,7 @@ impl GpuMultibodySolver { } self.apply_contact_restitution.call( pass, - args.mb_sweep_indirect, + args.mb_dispatch_indirect, &mb.multibody_info, &mut mb.contact_constraints, &mb.contact_jac_cols, diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 5d76c7ad..e2b443ad 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -3,8 +3,21 @@ //! Constraint-based physics solver running entirely on the GPU, using graph //! coloring to solve constraints in parallel without data races. Uses the //! `Soft-TGS` algorithm (as in Rapier). +//! +//! The contact constraints are stored in color-ordered tiles (see the `contact_tiles` shader +//! module): every iteration walks a contiguous range of them. Their topology and identity stay +//! in contact order, in the contact links used for the matching, mass splitting and coloring. +use super::contact_kernels::{ + ContactConstraints, GpuGatherWarmstartVelocities, GpuPrepareConstraints, + GpuScaleConstraintImpulses, GpuSolveConstraints, GpuSolveConstraintsBiased, + GpuSolveConstraintsFused, GpuSolveConstraintsFusedCached, GpuSolveConstraintsTail, + GpuSolveConstraintsTailBiased, GpuSolveConstraintsTailUnbiased, + GpuSolveConstraintsTailUnbiasedCached, GpuSolveConstraintsUnbiased, + GpuSolveConstraintsUnbiasedCached, GpuWarmstartConstraints, GpuWarmstartConstraintsFused, +}; use crate::dynamics::joint::{GpuJointSolver, JointSolverArgs}; +use crate::dynamics::mass_splitting::{GpuMassSplitting, HubState, SplitArgs}; #[cfg(feature = "dim3")] use crate::dynamics::multibody::{GpuMultibodySet, GpuMultibodySolver, MultibodySolverArgs}; use crate::math::Pose; @@ -12,49 +25,62 @@ use crate::queries::GpuIndexedContact; use crate::shaders::broad_phase::ContactPlan; #[cfg(feature = "dim3")] use crate::shaders::dynamics::MbContactIndexEntry; +use crate::shaders::dynamics::contact_tiles::TAIL_COLOR; use crate::shaders::dynamics::{BIAS_MODE_BIAS, BIAS_MODE_BIAS_FRICTION}; use crate::shaders::dynamics::{ - GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized, - GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, GpuSolverInitConstraints, - GpuSolverRefreshRhsWoBias, GpuSolverSortConstraints, GpuSolverUpdateConstraints, - GpuStepGaussSeidel, GpuStepGaussSeidelFused, GpuWarmstart, GpuWarmstartFused, - GpuWarmstartWithoutColors, LocalMassProperties, RbdSimParams, TwoBodyConstraint, - TwoBodyConstraintBuilder, Velocity, WorldMassProperties, + ContactLink, GpuApplySolverVelsInc, GpuInitContactLinks, GpuInitSolverBodies, + GpuInitSolverVelsInc, GpuIntegrateLinearized, GpuMatchContactLinks, GpuSolverCleanup, + GpuSolverCountConstraints, GpuSolverFinalize, GpuSolverSortConstraints, LocalMassProperties, + RbdSimParams, Velocity, WorldMassProperties, }; use crate::utils::{GpuPrefixSum, PrefixSumWorkspace}; use khal::Shader; -use khal::backend::{Encoder, GpuBackend, GpuBackendError, GpuEncoder, GpuPass, GpuTimestamps}; +use khal::backend::{ + DispatchGrid, Encoder, GpuBackend, GpuBackendError, GpuEncoder, GpuPass, GpuTimestamps, +}; use vortx::tensor::Tensor; /// GPU shader bundle for the constraint solver. #[derive(Shader)] pub struct GpuSolver { - sort_constraints: GpuSolverSortConstraints, - /// Initializes constraints from contact manifolds. - init_constraints: GpuSolverInitConstraints, - /// Companion counting pass to `init_constraints` (split out to keep each - /// pass within the 8-storage-buffer WebGPU limit). + /// Initializes the contact links from the contact manifolds. + init_links: GpuInitContactLinks, + /// Matches the contact links with the previous step's (warmstart and contact recycling). + match_links: GpuMatchContactLinks, + /// Counts the constraints of each body. count_constraints: GpuSolverCountConstraints, - /// Updates nonlinear constraint terms during substeps. - update_constraints: GpuSolverUpdateConstraints, - /// Refreshes the unbiased normal rhs from the post-integration poses, - /// between the position integration and the no-bias sweeps of each substep. - refresh_rhs_wo_bias: GpuSolverRefreshRhsWoBias, + /// Lists the constraints of each body. + sort_constraints: GpuSolverSortConstraints, + /// Builds the contact constraints and gives them their previous impulses, or recycled + /// contacts. + prepare_constraints: GpuPrepareConstraints, + /// Scales the warmstart impulses (warmstart coefficients other than 1). + scale_constraint_impulses: GpuScaleConstraintImpulses, + /// Sums the cached warmstart velocity changes of each body's constraints. + gather_warmstart_velocities: GpuGatherWarmstartVelocities, + /// Applies the warmstart of one color's constraints (scatter-style). + warmstart_constraints: GpuWarmstartConstraints, + /// Applies the warmstart of every color, with one workgroup per batch. + warmstart_constraints_fused: GpuWarmstartConstraintsFused, + /// Gauss-Seidel iteration over one color, in the mode given by a uniform. + solve_constraints: GpuSolveConstraints, + /// Fixed bias-only mode for the usual biased iteration. + solve_constraints_biased: GpuSolveConstraintsBiased, + /// Fixed friction/stabilization mode for the unbiased iteration. + solve_constraints_unbiased: GpuSolveConstraintsUnbiased, + /// Unbiased iteration that also caches the warmstart of the next substep. + solve_constraints_unbiased_cached: GpuSolveConstraintsUnbiasedCached, + /// The iterations over the sparse colors from [`TAIL_COLOR`], in one workgroup. + solve_constraints_tail: GpuSolveConstraintsTail, + solve_constraints_tail_biased: GpuSolveConstraintsTailBiased, + solve_constraints_tail_unbiased: GpuSolveConstraintsTailUnbiased, + solve_constraints_tail_unbiased_cached: GpuSolveConstraintsTailUnbiasedCached, + /// The iterations over every color with one workgroup per batch, used when per-batch + /// constraint counts are small. + solve_constraints_fused: GpuSolveConstraintsFused, + solve_constraints_fused_cached: GpuSolveConstraintsFusedCached, /// Clears solver velocities and constraint counts. cleanup: GpuSolverCleanup, - /// Applies warmstart impulses from previous frame. - #[allow(dead_code)] - warmstart: GpuWarmstart, - /// Applies warmstart impulses from previous frame, without relying on graph coloring. - warmstart_without_colors: GpuWarmstartWithoutColors, - /// Gauss-Seidel iteration step (sequential per color). - step_gauss_seidel: GpuStepGaussSeidel, - /// Fused variant of the colored warmstart sweep: one workgroup per batch - /// loops every color internally (barrier between colors). Used when - /// per-batch constraint counts are small. - warmstart_fused: GpuWarmstartFused, - /// Fused variant of the colored Gauss-Seidel sweep (same rationale). - step_gauss_seidel_fused: GpuStepGaussSeidelFused, /// Initializes solver velocity increments. init_solver_vels_inc: GpuInitSolverVelsInc, /// Seeds the COM-centered solver poses from the body world poses @@ -69,8 +95,7 @@ pub struct GpuSolver { finalize: GpuSolverFinalize, } -/// Arguments for constraint solver dispatch, used by [`GpuSolver::prepare`] and -/// [`GpuSolver::solve_tgs`]. +/// Arguments for constraint preparation and TGS dispatch. pub struct SolverArgs<'a> { /// Total number of colors from graph coloring. pub num_colors: u32, @@ -81,7 +106,7 @@ pub struct SolverArgs<'a> { /// Contact manifolds generated by narrow-phase. pub contacts: &'a Tensor, /// Clamped per-frame list totals (see `gpu_contact_plan`); `[PLAN_BOUND]` - /// bounds the flat contact sweeps. + /// bounds the flat contact dispatches. pub contact_plan: &'a Tensor, /// Contact→multibody index storage, rebuilt each step by /// `GpuMultibodySolver::layout_contact_constraints`. @@ -92,10 +117,16 @@ pub struct SolverArgs<'a> { pub stable_mb_contact_index: Option<&'a mut Tensor>, /// Flat dispatch grid over the whole contacts range. pub contacts_len_indirect: &'a Tensor<[u32; 3]>, - /// Solver constraints (output from constraint initialization). - pub constraints: &'a mut Tensor, - /// Builder data for initializing constraints. - pub constraint_builders: &'a mut Tensor, + /// This step's contact constraints. + pub constraints: &'a mut ContactConstraints, + /// The previous step's contact constraints, with their final impulses. + pub old_constraints: &'a ContactConstraints, + /// Previous-step per-body adjacency ends. + pub old_body_constraint_counts: &'a Tensor, + /// Previous-step adjacency entries. + pub old_body_constraint_ids: &'a Tensor, + /// Previous and current contact validity metadata (updated during preparation). + pub recycle_states: &'a mut super::ContactRecycleStates, /// Global simulation parameters. pub sim_params: &'a Tensor, /// Rigid body world-origin poses. Mirrors rapier's `RigidBody::position`. @@ -133,10 +164,17 @@ pub struct SolverArgs<'a> { /// to the multibody’s root. pub body_constraint_ids: &'a mut Tensor, /// The `(color, batch)` bucket buffer (color-major, post-scatter exclusive - /// ends): bucket `k` of `color_sorted_ids` spans `[buckets[k-1], buckets[k])`. + /// ends): bucket `k` of `sorted_links` spans `[buckets[k-1], buckets[k])`. pub color_buckets: &'a Tensor, - /// Constraint ids bucket-sorted by `(color, batch)`. - pub color_sorted_ids: &'a Tensor, + /// The constraint links bucket-sorted by `(color, batch)`. + pub sorted_links: &'a Tensor, + /// Previous-frame per-color thread counts. Zero means feedback is not available. + pub color_dispatch_threads: &'a [u32; 64], + /// Whether previous-frame statistics suggest a short late-color suffix, solved by one + /// workgroup. + pub fuse_tail_colors: bool, + /// Current GPU grid, used before feedback arrives or when readback is disabled. + pub color_dispatch_indirect: &'a Tensor<[u32; 3]>, /// Per-color-index uniform tensors: `color_uniforms[c] == c`. pub color_uniforms: &'a [Tensor], /// Prefix sum shader for building constraint ranges. @@ -144,7 +182,7 @@ pub struct SolverArgs<'a> { /// Number of solver iterations (max across all environments). pub num_solver_iterations: u32, /// PGS iterations of the biased pass per substep, shared by the - /// rigid-body and multibody sweeps (`RbdSimParams::num_internal_pgs_iterations`). + /// rigid-body and multibody iterations (`RbdSimParams::num_internal_pgs_iterations`). pub num_internal_pgs_iterations: u32, /// Per-body graph-coloring group id (multibody-aware). pub body_group: &'a Tensor, @@ -161,7 +199,7 @@ pub struct SolverArgs<'a> { /// /// This is generally used when the number of constraints is small wrt. /// the number of environments. - pub fused_color_sweeps: bool, + pub fused_color_dispatches: bool, /// `true` when every rigid-body contact constraint is provably a no-op. pub rb_contacts_inert: bool, /// Solve friction rows during the biased pass too @@ -171,13 +209,29 @@ pub struct SolverArgs<'a> { pub batch_indices: &'a Tensor, /// The one gravity uniform every rigid-body and multibody kernel reads. pub gravity: &'a Tensor, + /// Mass-splitting state of the bodies with many contacts. + pub hubs: &'a mut HubState, + /// Mass-splitting kernels. + pub mass_splitting: &'a GpuMassSplitting, + /// Whether the warmstart impulses are scaled (a warmstart coefficient other than 1). + pub scale_warmstart_impulses: bool, /// GPU-written workgroup grid for the per-multibody contact-constraint /// dispatches (zero workgroups on contact-free steps). - pub mb_sweep_indirect: &'a Tensor<[u32; 3]>, + pub mb_dispatch_indirect: &'a Tensor<[u32; 3]>, +} + +impl<'a> SolverArgs<'a> { + fn color_dispatch_grid(&self, color: u32) -> DispatchGrid<'a, GpuBackend> { + match self.color_dispatch_threads.get(color as usize).copied() { + Some(threads) if threads > 0 => threads.into(), + _ => self.color_dispatch_indirect.into(), + } + } } impl GpuSolver { - /// Prepares constraints for solving (init, prefix sum, sort). + /// Prepares the constraints before their coloring: their links, per-body lists, match + /// with the previous step and mass splitting. pub fn prepare<'a>( &self, backend: &GpuBackend, @@ -194,13 +248,13 @@ impl GpuSolver { args.vels, args.mprops, args.batch_indices, + &mut args.hubs.first_slot, + &mut args.hubs.counts, )?; // Seed `solver_body_poses` from `body_poses`: rapier's // `SolverBodies::copy_from`. After this, only the COM-centered solver // poses are touched until the final `finalize` writeback. - // Independent of `cleanup` (different buffers), but `init_constraints` - // below reads `solver_body_poses` so we barrier before it. self.init_solver_bodies.call( pass, args.num_colliders * args.num_batches, @@ -214,23 +268,15 @@ impl GpuSolver { return Ok(()); } - self.init_constraints.call( + self.init_links.call( pass, args.contacts_len_indirect, args.contacts, - args.constraints, - args.constraint_builders, - args.contact_plan, + &mut args.constraints.links, args.body_is_multibody, - args.collider_world_poses, - args.solver_body_poses, - args.vels, - args.mprops, - args.sim_params, + args.contact_plan, )?; - // Counting runs as a separate dispatch (same indirect grid) so the - // build pass above stays within 8 storage buffers. self.count_constraints.call( pass, args.contacts_len_indirect, @@ -262,9 +308,67 @@ impl GpuSolver { args.body_is_multibody, )?; + let (recycle_states, recycle_offsets) = args.recycle_states.bindings(); + self.match_links.call( + pass, + args.contacts_len_indirect, + args.contacts, + &mut args.constraints.links, + args.old_body_constraint_counts, + args.old_body_constraint_ids, + &args.old_constraints.links, + &args.old_constraints.constraint_indices, + recycle_states, + args.collider_world_poses, + args.contact_plan, + args.sim_params, + recycle_offsets, + )?; + + args.mass_splitting.split( + pass, + args.hubs, + SplitArgs { + num_body_slots: args.num_colliders * args.num_batches, + body_constraint_counts: args.body_constraint_counts, + body_constraint_ids: args.body_constraint_ids, + links: &mut args.constraints.links, + mprops: args.mprops, + body_group: args.body_group, + batch_indices: args.batch_indices, + }, + )?; + Ok(()) } + /// Builds the contact constraints once their order is known (after the color bucket sort), + /// and caches their warmstart for the first substep. + pub fn build_constraints( + &self, + pass: &mut GpuPass, + args: &mut SolverArgs<'_>, + ) -> Result<(), GpuBackendError> { + if args.rb_contacts_inert { + return Ok(()); + } + let (recycle_states, recycle_offsets) = args.recycle_states.bindings(); + self.prepare_constraints.call( + pass, + args.contacts_len_indirect, + args.sorted_links, + args.contacts, + args.mprops, + &*args.solver_body_poses, + &*args.vels, + &*recycle_states, + &args.old_constraints.tiles, + &mut args.constraints.tiles, + args.contact_plan, + recycle_offsets, + ) + } + /// Solves constraints using the TGS (Total Gauss-Seidel) algorithm. /// /// When `multibody` is `Some`, multibody substep work is interleaved with @@ -288,7 +392,9 @@ impl GpuSolver { let skip_rb = args.rb_contacts_inert; let joints_empty = joint_args.joints.is_empty(); #[cfg(feature = "dim3")] - let stable_mb_contact_index = args.stable_mb_contact_index; + let mut args = args; + #[cfg(feature = "dim3")] + let stable_mb_contact_index = args.stable_mb_contact_index.take(); /* * Init solver vel increments. @@ -326,7 +432,7 @@ impl GpuSolver { batch_indices: args.batch_indices, gravity: args.gravity, color_uniforms: args.color_uniforms, - mb_sweep_indirect: args.mb_sweep_indirect, + mb_dispatch_indirect: args.mb_dispatch_indirect, friction_in_bias_pass: args.friction_in_bias_pass, }; solver.layout_contact_constraints( @@ -359,7 +465,7 @@ impl GpuSolver { batch_indices: args.batch_indices, gravity: args.gravity, color_uniforms: args.color_uniforms, - mb_sweep_indirect: args.mb_sweep_indirect, + mb_dispatch_indirect: args.mb_dispatch_indirect, friction_in_bias_pass: args.friction_in_bias_pass, }; solver.$method(&mut pass, state, &mut mb_args $(, $extra)*)?; @@ -414,7 +520,7 @@ impl GpuSolver { batch_indices: args.batch_indices, gravity: args.gravity, color_uniforms: args.color_uniforms, - mb_sweep_indirect: args.mb_sweep_indirect, + mb_dispatch_indirect: args.mb_dispatch_indirect, friction_in_bias_pass: args.friction_in_bias_pass, }; solver.substep_build_constraints( @@ -430,14 +536,15 @@ impl GpuSolver { let mut pass = encoder.begin_pass("[RBD] slv/rb-build-warmstart", timestamps.as_deref_mut()); let pass = &mut pass; - if !skip_rb { - self.update_constraints.call( + // The solve kernels compute the normal right-hand sides themselves: this pass only + // rescales the warmstart impulses, for warmstart coefficients other than 1. + if !skip_rb && args.scale_warmstart_impulses { + self.scale_constraint_impulses.call( pass, args.contacts_len_indirect, - args.constraints, - args.constraint_builders, - args.contact_plan, - args.solver_body_poses, + args.color_buckets, + &mut args.constraints.tiles, + args.batch_indices, args.sim_params, )?; } @@ -446,49 +553,77 @@ impl GpuSolver { // Contact warmstart skipped: no rigid-body contact // constraint can carry an impulse here. } else if args.colorless_warmstart { - self.warmstart_without_colors.call( + self.gather_warmstart_velocities.call( pass, args.num_colliders * args.num_batches, args.body_constraint_counts, args.body_constraint_ids, - args.constraints, + &args.constraints.constraint_indices, + &args.constraints.tiles, args.solver_vels, + &args.hubs.first_slot, args.batch_indices, )?; - } else if args.fused_color_sweeps { + // Split bodies are warmstarted through their sub-bodies. + args.mass_splitting + .scatter(pass, args.hubs, args.solver_vels)?; + args.mass_splitting.warmstart( + pass, + args.hubs, + args.solver_vels, + &args.constraints.tiles, + &args.constraints.constraint_indices, + )?; + args.mass_splitting.average( + pass, + args.hubs, + args.solver_vels, + args.body_constraint_counts, + )?; + } else if args.fused_color_dispatches { + args.mass_splitting + .scatter(pass, args.hubs, args.solver_vels)?; // One dispatch, one workgroup per batch, colors looped // internally. `color_uniforms[num_colors]` holds the // constant `num_colors`. - self.warmstart_fused.call( + self.warmstart_constraints_fused.call( pass, [64, args.num_batches, 1], - args.constraints, + &args.constraints.tiles, args.solver_vels, args.color_buckets, - args.color_sorted_ids, &args.color_uniforms[args.num_colors as usize], args.batch_indices, )?; } else { + args.mass_splitting + .scatter(pass, args.hubs, args.solver_vels)?; // NOTE: contact colors start at 1 (0 = unassigned). for c in 1..=args.num_colors { - self.warmstart.call( + self.warmstart_constraints.call( pass, - args.contacts_len_indirect, - args.constraints, + args.color_dispatch_grid(c), + &args.constraints.tiles, args.solver_vels, args.color_buckets, - args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, )?; } } + if !skip_rb && !args.colorless_warmstart { + args.mass_splitting.average( + pass, + args.hubs, + args.solver_vels, + args.body_constraint_counts, + )?; + } } /* * Solve all joints + contacts with bias. The multibody and - * rigid-body sweeps interleave, one iteration each, so a body + * rigid-body iterations interleave, one at a time, so a body * squeezed between a multibody link and a rigid body sees both * sides converge at the same rate. */ @@ -507,37 +642,71 @@ impl GpuSolver { encoder.begin_pass("[RBD] slv/rb-solve-bias", timestamps.as_deref_mut()); let pass = &mut pass; joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?; + if !skip_rb { + args.mass_splitting + .scatter(pass, args.hubs, args.solver_vels)?; + } if skip_rb { - // Contact sweeps skipped (inert constraints). - } else if args.fused_color_sweeps { - self.step_gauss_seidel_fused.call( + // Contact iterations skipped (inert constraints). + } else if args.fused_color_dispatches { + self.solve_constraints_fused.call( pass, [64, args.num_batches, 1], - args.constraints, + &mut args.constraints.tiles, args.solver_vels, args.color_buckets, - args.color_sorted_ids, + args.solver_body_poses, &args.color_uniforms[args.num_colors as usize], args.batch_indices, // Biased pass: `color_uniforms[bias_mode]` holds `bias_mode`. &args.color_uniforms[bias_mode], + args.sim_params, )?; } else { for c in 1..=args.num_colors { - self.step_gauss_seidel.call( - pass, - args.contacts_len_indirect, - args.constraints, - args.solver_vels, - args.color_buckets, - args.color_sorted_ids, - &args.color_uniforms[c as usize], - args.batch_indices, - // Biased pass: `color_uniforms[bias_mode]` holds `bias_mode`. - &args.color_uniforms[bias_mode], - )?; + let tail = args.fuse_tail_colors && c >= TAIL_COLOR; + if tail && c > TAIL_COLOR { + continue; + } + // The tail kernel solves every color from `TAIL_COLOR` to its + // color uniform, in one workgroup. + let (grid, color): (DispatchGrid, u32) = if tail { + (64u32.into(), args.num_colors) + } else { + (args.color_dispatch_grid(c), c) + }; + macro_rules! solve_color { + ($kernel:expr) => { + $kernel.call( + pass, + grid, + &mut args.constraints.tiles, + args.solver_vels, + args.color_buckets, + args.solver_body_poses, + &args.color_uniforms[color as usize], + args.batch_indices, + &args.color_uniforms[bias_mode], + args.sim_params, + )? + }; + } + match (tail, args.friction_in_bias_pass) { + (true, true) => solve_color!(self.solve_constraints_tail), + (true, false) => solve_color!(self.solve_constraints_tail_biased), + (false, true) => solve_color!(self.solve_constraints), + (false, false) => solve_color!(self.solve_constraints_biased), + } } } + if !skip_rb { + args.mass_splitting.average( + pass, + args.hubs, + args.solver_vels, + args.body_constraint_counts, + )?; + } } } @@ -566,53 +735,89 @@ impl GpuSolver { * Solve all joints + contacts without bias (stabilization). */ mb_phase!("[RBD] slv/mb-solve-nobias", substep_solve_no_bias); + // The last iteration of a substep caches the warmstart of the next one, gathered per + // body. Rescaled impulses recompute it anyway. + let cache_warmstart = + !is_last_substep && args.colorless_warmstart && !args.scale_warmstart_impulses; if !skip_rb || !joints_empty { let mut pass = encoder.begin_pass("[RBD] slv/rb-solve-nobias", timestamps.as_deref_mut()); let pass = &mut pass; + joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?; if !skip_rb { - self.refresh_rhs_wo_bias.call( - pass, - args.contacts_len_indirect, - args.constraints, - args.constraint_builders, - args.contact_plan, - args.solver_body_poses, - args.sim_params, - )?; + args.mass_splitting + .scatter(pass, args.hubs, args.solver_vels)?; } - joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?; if skip_rb { - // Contact sweeps skipped (inert constraints). - } else if args.fused_color_sweeps { - self.step_gauss_seidel_fused.call( - pass, - [64, args.num_batches, 1], - args.constraints, - args.solver_vels, - args.color_buckets, - args.color_sorted_ids, - &args.color_uniforms[args.num_colors as usize], - args.batch_indices, - // use_bias = 0 (the `color_uniform[0]` contains the value 0) - &args.color_uniforms[0], - )?; + // Contact iterations skipped (inert constraints). + } else if args.fused_color_dispatches { + // `color_uniforms[0]` holds 0: no bias, with friction. + macro_rules! solve_fused { + ($kernel:expr) => { + $kernel.call( + pass, + [64, args.num_batches, 1], + &mut args.constraints.tiles, + args.solver_vels, + args.color_buckets, + args.solver_body_poses, + &args.color_uniforms[args.num_colors as usize], + args.batch_indices, + &args.color_uniforms[0], + args.sim_params, + )? + }; + } + if cache_warmstart { + solve_fused!(self.solve_constraints_fused_cached); + } else { + solve_fused!(self.solve_constraints_fused); + } } else { for c in 1..=args.num_colors { - self.step_gauss_seidel.call( - pass, - args.contacts_len_indirect, - args.constraints, - args.solver_vels, - args.color_buckets, - args.color_sorted_ids, - &args.color_uniforms[c as usize], - args.batch_indices, - // use_bias = 0 (the `color_uniform[0]` contains the value 0) - &args.color_uniforms[0], - )?; + let tail = args.fuse_tail_colors && c >= TAIL_COLOR; + if tail && c > TAIL_COLOR { + continue; + } + let (grid, color): (DispatchGrid, u32) = if tail { + (64u32.into(), args.num_colors) + } else { + (args.color_dispatch_grid(c), c) + }; + macro_rules! solve_color { + ($kernel:expr) => { + $kernel.call( + pass, + grid, + &mut args.constraints.tiles, + args.solver_vels, + args.color_buckets, + args.solver_body_poses, + &args.color_uniforms[color as usize], + args.batch_indices, + &args.color_uniforms[0], + args.sim_params, + )? + }; + } + match (tail, cache_warmstart) { + (true, true) => { + solve_color!(self.solve_constraints_tail_unbiased_cached) + } + (true, false) => solve_color!(self.solve_constraints_tail_unbiased), + (false, true) => solve_color!(self.solve_constraints_unbiased_cached), + (false, false) => solve_color!(self.solve_constraints_unbiased), + } } } + if !skip_rb { + args.mass_splitting.average( + pass, + args.hubs, + args.solver_vels, + args.body_constraint_counts, + )?; + } } } diff --git a/src_rbd/dynamics/warmstart.rs b/src_rbd/dynamics/warmstart.rs index 4d17016a..69b9206f 100644 --- a/src_rbd/dynamics/warmstart.rs +++ b/src_rbd/dynamics/warmstart.rs @@ -1,61 +1,28 @@ //! Warmstarting: reuses previous-frame impulses for faster solver convergence. +//! +//! The contacts are matched with the previous step's by `GpuSolver::prepare`, and the +//! constraints receive their previous impulses in `GpuSolver::build_constraints`. The coloring +//! also starts from the previous colors of the matched constraints. use crate::shaders::broad_phase::ContactPlan; -use crate::shaders::dynamics::{ - GpuSeedColorsFromWarmstart, GpuTransferWarmstartImpulses, TwoBodyConstraint, - TwoBodyConstraintBuilder, -}; +use crate::shaders::dynamics::{ContactLink, GpuSeedColorsFromWarmstart}; use khal::Shader; use khal::backend::{GpuBackendError, GpuPass}; use vortx::tensor::Tensor; -/// GPU shader for transferring warmstart impulses between frames. -/// -/// This shader matches new contacts against old contacts and transfers impulse -/// accumulators when a match is found. +/// GPU shader seeding the coloring with the previous step's colors. #[derive(Shader)] pub struct GpuWarmstart { - /// Compute pipeline that matches contacts and transfers impulses. - transfer_warmstart_impulses_kernel: GpuTransferWarmstartImpulses, - /// Seeds the topo-gc coloring from the previous frame's colors (same - /// old/new body-pair matching as the impulse transfer). + /// Seeds the topo-gc coloring from the previous frame's colors of the matched constraints. seed_colors_kernel: GpuSeedColorsFromWarmstart, } -/// Arguments for warmstart dispatch. -/// -/// Contains buffers for both old (previous frame) and new (current frame) constraint data. -pub struct WarmstartArgs<'a> { - /// Clamped per-frame list totals (the flat contact sweep bound). - pub contact_plan: &'a Tensor, - /// Constraint counts per body from previous frame. - pub old_body_constraint_counts: &'a Tensor, - /// Constraint IDs per body from previous frame. - pub old_body_constraint_ids: &'a Tensor, - /// Solver constraints from previous frame. - pub old_constraints: &'a Tensor, - /// Constraint builders from previous frame. - pub old_constraint_builders: &'a Tensor, - /// Solver constraints for current frame (to be warmstarted). - pub new_constraints: &'a mut Tensor, - /// Constraint builders for current frame. - pub new_constraint_builders: &'a Tensor, - /// Indirect dispatch arguments based on contact count. - pub contacts_len_indirect: &'a Tensor<[u32; 3]>, -} - /// Arguments for the coloring seed dispatch. pub struct SeedColorsArgs<'a> { - /// Clamped per-frame list totals (the flat contact sweep bound). + /// Clamped per-frame list totals (the flat contact dispatch bound). pub contact_plan: &'a Tensor, - /// Constraint counts per body from previous frame. - pub old_body_constraint_counts: &'a Tensor, - /// Constraint IDs per body from previous frame. - pub old_body_constraint_ids: &'a Tensor, - /// Solver constraints from previous frame. - pub old_constraints: &'a Tensor, - /// Solver constraints for current frame. - pub new_constraints: &'a Tensor, + /// The contact links of the current frame, matched with the previous frame's. + pub links: &'a Tensor, /// Colors assigned to the previous frame's constraints. pub old_constraints_colors: &'a Tensor, /// Output: colors for the current frame's constraints (seeded slots only). @@ -67,25 +34,6 @@ pub struct SeedColorsArgs<'a> { } impl GpuWarmstart { - /// Transfers warmstart impulses from old constraints to new constraints. - pub fn transfer_warmstart_impulses<'a>( - &self, - pass: &mut GpuPass, - args: WarmstartArgs<'a>, - ) -> Result<(), GpuBackendError> { - self.transfer_warmstart_impulses_kernel.call( - pass, - args.contacts_len_indirect, - args.old_body_constraint_counts, - args.old_body_constraint_ids, - args.old_constraints, - args.old_constraint_builders, - args.new_constraints, - args.new_constraint_builders, - args.contact_plan, - ) - } - /// Seeds the topo-gc coloring from the previous frame's colors. Must run /// after the topo-gc reset and before its iterations. pub fn seed_colors_from_warmstart( @@ -96,10 +44,7 @@ impl GpuWarmstart { self.seed_colors_kernel.call( pass, args.contacts_len_indirect, - args.old_body_constraint_counts, - args.old_body_constraint_ids, - args.old_constraints, - args.new_constraints, + args.links, args.old_constraints_colors, args.constraints_colors, args.colored, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index cdc5d6b0..546c6916 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -17,6 +17,7 @@ use crate::utils::PrefixSumWorkspace; use std::ops::Range; use super::rbd_state::*; +use crate::dynamics::HubState; #[cfg(feature = "dim3")] use crate::rapier::dynamics::{MultibodyJointSet, RigidBodySet}; use khal::BufferUsages; @@ -145,10 +146,7 @@ impl RbdState { .unwrap(); // Readback: pair count, PFM count, uncolored count (+ the multibody // contact-constraint demand on dim3). - #[cfg(feature = "dim3")] - let resize_readback = GpuReadback::new(backend, 4).unwrap(); - #[cfg(not(feature = "dim3"))] - let resize_readback = GpuReadback::new(backend, 3).unwrap(); + let resize_readback = GpuReadback::new(backend, ResizeFeedback::WORDS).unwrap(); let collision_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let contacts = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); @@ -156,7 +154,7 @@ impl RbdState { let mb_contact_index = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let contacts_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); - let mb_sweep_indirect = + let mb_dispatch_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let pfm_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); @@ -169,12 +167,23 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); - let old_constraints = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); - let old_constraint_builders = - Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); - let new_constraints = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); - let new_constraint_builders = - Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let old_constraints = crate::dynamics::ContactConstraints::new( + backend, + contacts_capacity, + storage, + &base_sim_params, + ) + .unwrap(); + let new_constraints = crate::dynamics::ContactConstraints::new( + backend, + contacts_capacity, + storage, + &base_sim_params, + ) + .unwrap(); + let recycle_states = + crate::dynamics::ContactRecycleStates::new(backend, contacts_capacity, storage) + .unwrap(); let constraints_colors = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let constraints_pending_colors = @@ -194,7 +203,7 @@ impl RbdState { storage | BufferUsages::UNIFORM, ) .unwrap(); - let color_sorted_ids = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let sorted_links = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); // Zeroed: see `RbdState::from_rapier`. let old_constraints_counts = Tensor::vector( backend, @@ -251,8 +260,9 @@ impl RbdState { vels: Tensor::vector(backend, &all_vels, rw).unwrap(), #[cfg(feature = "dim3")] reset_templates_bodies: None, - solver_vels: Tensor::vector(backend, &all_vels, storage).unwrap(), + solver_vels: HubState::solver_vels(backend, &all_vels, storage), solver_vels_inc: Tensor::vector(backend, &all_vels, storage).unwrap(), + hubs: HubState::new(backend, all_vels.len()), joints, #[cfg(feature = "dim3")] multibodies, @@ -278,7 +288,8 @@ impl RbdState { collision_pairs_indirect, contacts_capacity_cpu, collision_pairs_capacity_cpu, - collision_pairs_len_cpu: 0, + collision_pairs_len_cpu: None, + stepped_since_readback: false, graph_generation: 0, compute_graph: Default::default(), #[cfg(feature = "dim3")] @@ -290,22 +301,21 @@ impl RbdState { pfm_sort, #[cfg(feature = "dim3")] mb_contact_index, - mb_sweep_indirect, + mb_dispatch_indirect, pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, old_constraints, - old_constraint_builders, old_constraints_counts, new_constraints, - new_constraint_builders, + recycle_states, new_constraints_counts, constraints_colors, old_constraints_colors, colored, constraints_rands, color_buckets, - color_sorted_ids, + sorted_links, curr_color: Tensor::scalar( backend, 0u32, @@ -321,6 +331,7 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, ) .unwrap(), + coloring_dispatch: ColoringDispatch::new(backend), uncolored_staging: Tensor::scalar( backend, 0, diff --git a/src_rbd/pipeline/mod.rs b/src_rbd/pipeline/mod.rs index dc0aeb9d..1f79841e 100644 --- a/src_rbd/pipeline/mod.rs +++ b/src_rbd/pipeline/mod.rs @@ -22,3 +22,18 @@ pub use rbd_state::{ DebugContact, DebugLbvhNode, RbdCapacities, RbdGraphKey, RbdResizePolicy, RbdState, RunStats, }; pub use rbd_step::RbdPipeline; + +#[cfg(test)] +mod test_large_warmstart; + +#[cfg(all(test, feature = "dim3"))] +mod test_tail_colors; + +#[cfg(all(test, feature = "dim3"))] +mod test_twist_friction; + +#[cfg(test)] +mod test_color_buckets; + +#[cfg(test)] +mod test_web_preparation; diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index f7977cff..0a14bbaf 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -1,18 +1,18 @@ //! GPU-resident rigid-body state ([`RbdState`]): buffer definitions, accessors, //! run statistics and capacity/resize policies. use crate::broad_phase::{BRUTE_FORCE_MAX_COLLIDERS, LbvhState, PfmSortState}; -use crate::dynamics::GpuImpulseJointSet; #[cfg(feature = "dim3")] use crate::dynamics::GpuMultibodySet; +use crate::dynamics::{ColorStatsBuffer, GpuImpulseJointSet, HubState}; use crate::math::{Pose, Vector}; use crate::queries::{GpuColliderMaterial, GpuIndexedContact}; use crate::shaders::PaddedVector; use crate::shaders::broad_phase::{CollisionPair, ContactPlan, LbvhNode, NarrowPhasePfmPair}; +use crate::shaders::dynamics::ColorStats; #[cfg(feature = "dim3")] use crate::shaders::dynamics::MbContactIndexEntry; use crate::shaders::dynamics::{ - LocalMassProperties as GpuLocalMassProperties, RbdSimParams, TwoBodyConstraint, - TwoBodyConstraintBuilder, Velocity as GpuVelocity, + LocalMassProperties as GpuLocalMassProperties, RbdSimParams, Velocity as GpuVelocity, WorldMassProperties as GpuWorldMassProperties, }; use crate::shaders::queries::MAX_MANIFOLD_POINTS; @@ -118,15 +118,15 @@ pub struct RbdCapacities { /// GPU->CPU buffer readback, resulting in a larger performance gain than just setting /// only one of them to `Fixed`. pub collisions_resize_policy: RbdResizePolicy, - /// Maximum number of colors used by the solver for constraints coloring. + /// Initial/minimum coloring budget. Large adaptive scenes start with additional + /// convergence margin; sparse contact graphs can shrink back to this minimum. pub solver_colors: u32, /// How internal constraints coloring gets automatically adjusted (or not). /// /// While this doesn’t change any buffer allocation, this affects the number of /// iterations the constraints coloring step applies, which has a computational cost. /// - /// Note that `RbdResizePolicy::Fit` for solver colors will currently act like `::Grow` - /// (i.e. the color won’t go back down yet). + /// Adaptive policies grow on failed convergence and shrink above the minimum budget. /// /// Note that setting both [`Self::collisions_resize_policy`] and /// [`Self::solver_colors_resize_policy`] to [`RbdResizePolicy::Fixed`] eliminates a @@ -149,6 +149,28 @@ impl Default for RbdCapacities { } } +impl RbdCapacities { + pub(super) fn minimum_solver_colors( + &self, + active_colliders: u32, + highest_color: Option, + ) -> u32 { + // Before feedback, and for dense contact graphs, leave room to converge before + // CPU feedback arrives. Sparse graphs need no such floor: extra empty color + // dispatches can cost more than their actual solver work. + let dense = highest_color.is_none_or(|highest| highest >= 16); + if active_colliders >= 4096 + && dense + && self.solver_colors_resize_policy != RbdResizePolicy::Fixed + { + self.solver_colors + .max(if active_colliders >= 65536 { 64 } else { 32 }) + } else { + self.solver_colors + } + } +} + /// Governs the way the rigid-body dynamics pipeline automatically resizes internal buffers storing /// data with unpredictable size (like collisions). #[derive(Copy, Clone, Debug, PartialEq, Eq, Default)] @@ -164,6 +186,64 @@ pub enum RbdResizePolicy { Fit, } +/// Dispatch grids and stats of the graph coloring and the colored dispatches. +pub(crate) struct ColoringDispatch { + /// Grid of the current topo-gc iteration: empty once the coloring converged. + pub(crate) coloring_indirect: Tensor<[u32; 3]>, + /// Grid of the colored dispatches: sized to the largest color bucket. + pub(crate) dispatch_indirect: Tensor<[u32; 3]>, + /// CPU scheduling hints, in threads, for colors 0..64. Never used as loop bounds. + pub(crate) dispatch_threads: [u32; 64], + /// Whether the sparse colors from `TAIL_COLOR` are solved by a single workgroup. Hint + /// only: that workgroup still visits the full current GPU bucket range. + pub(crate) fuse_tail_colors: bool, + /// The color statistics of the current step. + pub(crate) color_stats: ColorStatsBuffer, +} + +/// The counters read back after each step, to resize the buffers and size the next dispatches. +/// [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers) +/// gathers them in field order. +#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)] +#[repr(C)] +pub(super) struct ResizeFeedback { + /// Total collision-pair count, over all batches. + pub collision_pairs: u32, + /// Total PFM work-list count, over all batches. + pub pfm_pairs: u32, + /// Non-zero when the topo-gc coloring converged. + pub coloring_converged: u32, + /// Total multibody contact-constraint demand. + #[cfg(feature = "dim3")] + pub mb_cons_demand: u32, + pub color_stats: ColorStats, +} + +impl ResizeFeedback { + /// Number of `u32` words in a [`ResizeFeedback`]. + pub(super) const WORDS: usize = size_of::() / size_of::(); +} + +impl ColoringDispatch { + pub(crate) fn new(backend: &GpuBackend) -> Self { + let indirect = BufferUsages::STORAGE | BufferUsages::INDIRECT; + Self { + coloring_indirect: Tensor::scalar(backend, [0u32, 1, 1], indirect).unwrap(), + dispatch_indirect: Tensor::scalar(backend, [0u32, 1, 1], indirect).unwrap(), + dispatch_threads: [0; 64], + fuse_tail_colors: false, + color_stats: ColorStatsBuffer( + Tensor::vector( + backend, + vec![0u32; ColorStatsBuffer::WORDS], + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + ) + .unwrap(), + ), + } + } +} + /// GPU-resident physics simulation state containing all rigid bodies, shapes, and solver data. /// /// Holds all the buffers needed for a complete physics simulation on the GPU @@ -191,8 +271,12 @@ pub struct RbdState { /// `publish_reset_templates` and consumed by `reset_envs_from_templates`. #[cfg(feature = "dim3")] pub(super) reset_templates_bodies: Option, + /// Solver velocities: one slot per body, followed by the sub-body slots of split bodies + /// (see [`HubState`]). pub(super) solver_vels: Tensor, pub(super) solver_vels_inc: Tensor, + /// Mass-splitting state of bodies with many contacts. + pub(super) hubs: HubState, pub(super) vertex_buffers: Tensor, pub(super) index_buffers: Tensor, pub(super) shapes: Tensor, @@ -218,8 +302,7 @@ pub struct RbdState { /// Single global live collision-pair count (length 1): every batch appends /// to the same flat pair buffer. pub(super) collision_pairs_len: Tensor, - /// Non-blocking readback of `[collision_pairs_len, pfm_pairs_len, - /// uncolored]` used by + /// Non-blocking readback of the [`ResizeFeedback`] words, used by /// [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers) /// to grow buffers without stalling. pub(super) resize_readback: GpuReadback, @@ -234,8 +317,12 @@ pub struct RbdState { pub(super) collision_pairs_capacity_cpu: u32, /// Most recently read live collision-pair count — the total across all /// batches, harvested by the non-blocking readback in [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers). - /// Surfaced in the viewer UI; lags the GPU by a frame or two like the resize. - pub(super) collision_pairs_len_cpu: u32, + /// Surfaced in the viewer UI; lags the GPU by a frame or two like the resize. `None` until + /// the first readback completes. + pub(super) collision_pairs_len_cpu: Option, + /// Whether a step ran since the last resize readback request: the counters are only read + /// back once a step wrote them. + pub(super) stepped_since_readback: bool, /// Bumped on every GPU buffer (re)allocation or capacity change; part of /// [`Self::graph_key`]. pub(super) graph_generation: u64, @@ -259,7 +346,7 @@ pub struct RbdState { /// Flat dispatch grid over the whole contacts range, written by /// `gpu_contact_plan`. pub(super) contacts_indirect: Tensor<[u32; 3]>, - /// Clamped per-frame list totals (see `gpu_contact_plan`): the sweep + /// Clamped per-frame list totals (see `gpu_contact_plan`): the dispatch /// bound and the positional-slot bases of the flat contacts buffer. pub(super) contact_plan: Tensor, /// Buffers backing the per-pair PFM sort of the contact-reduction path. @@ -271,13 +358,13 @@ pub struct RbdState { pub(super) mb_contact_index: Tensor, /// Workgroup grid for the per-multibody contact-constraint dispatches: /// `[multibodies_batch_capacity, num_batches, 1]`. - pub(super) mb_sweep_indirect: Tensor<[u32; 3]>, - pub(super) new_constraints: Tensor, - pub(super) new_constraint_builders: Tensor, + pub(super) mb_dispatch_indirect: Tensor<[u32; 3]>, + pub(super) new_constraints: crate::dynamics::ContactConstraints, pub(super) new_constraints_counts: Tensor, pub(super) new_body_constraint_ids: Tensor, - pub(super) old_constraints: Tensor, - pub(super) old_constraint_builders: Tensor, + pub(super) old_constraints: crate::dynamics::ContactConstraints, + /// Previous and current recycling metadata in one storage allocation. + pub(super) recycle_states: crate::dynamics::ContactRecycleStates, pub(super) old_constraints_counts: Tensor, pub(super) old_body_constraint_ids: Tensor, pub(super) constraints_colors: Tensor, @@ -286,15 +373,17 @@ pub struct RbdState { pub(super) constraints_rands: Tensor, /// The single `(color, batch)` bucket buffer, color-major, of length /// `(max_colors + 3) * num_batches`: counts, then scanned exclusive - /// starts, then post-scatter exclusive ends (what the sweeps read). + /// starts, then post-scatter exclusive ends (what the solve kernels read). pub(super) color_buckets: Tensor, - /// Constraint indices bucket-sorted by `(color, batch)`. - pub(super) color_sorted_ids: Tensor, + /// The constraint links bucket-sorted by `(color, batch)`. + pub(super) sorted_links: Tensor, pub(super) curr_color: Tensor, /// Pre-built per-color-index uniforms: `color_uniforms[c] == c`. /// [`Self::ensure_color_uniforms`]. pub(super) color_uniforms: Vec>, pub(super) uncolored: Tensor, + /// Dispatch grids and stats of the coloring and colored dispatches. + pub(super) coloring_dispatch: ColoringDispatch, pub(super) uncolored_staging: Tensor, pub(super) lbvh: LbvhState, pub(super) joints: GpuImpulseJointSet, @@ -540,7 +629,6 @@ impl RbdState { zero(backend, &mut self.determinism.pending_colors); zero(backend, &mut self.old_constraints_colors); zero(backend, &mut self.colored); - zero(backend, &mut self.color_sorted_ids); zero(backend, &mut self.determinism.contact_sort_keys); zero(backend, &mut self.determinism.contact_sort_ids); zero(backend, &mut self.determinism.contact_sort_keys_out); @@ -593,7 +681,7 @@ impl RbdState { /// harvested by the non-blocking readback in [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers). Lags the GPU by a /// frame or two; `0` until the first readback completes. pub fn collision_pairs_len(&self) -> u32 { - self.collision_pairs_len_cpu + self.collision_pairs_len_cpu.unwrap_or(0) } /// GPU buffer of the broad-phase collision pairs found this step. @@ -635,8 +723,32 @@ impl RbdState { self.sim_params_cpu = params; } + /// Sets the duration of a full physics step and its TGS substep count, + /// without rebuilding bodies. + pub fn set_timestep(&mut self, backend: &GpuBackend, dt: f32, substeps: u32) { + let substeps = substeps.max(1); + let substep_dt = dt / substeps as f32; + if self.num_solver_iterations == substeps && self.sim_params_cpu.dt == substep_dt { + return; + } + let mut params = self.sim_params_cpu; + params.dt = substep_dt; + params.num_solver_iterations = substeps; + let _ = backend.write_buffer(self.sim_params.buffer_mut(), 0, &[params]); + self.sim_params_cpu = params; + self.num_solver_iterations = substeps; + #[cfg(feature = "dim3")] + { + self.multibodies.set_num_solver_iterations(substeps); + self.multibodies.set_visible_dt(backend, dt); + self.multibodies.set_constraint_softness(backend, ¶ms); + } + // Multibody uniforms were replaced and captured dispatch counts changed. + self.graph_generation += 1; + } + /// Sets how many PGS iterations the biased pass runs per substep (rigid-body - /// and multibody sweeps alike), without rebuilding the GPU state. + /// and multibody alike), without rebuilding the GPU state. #[cfg(feature = "dim3")] pub fn set_num_internal_pgs_iterations(&mut self, backend: &GpuBackend, n: u32) { let n = n.max(1); @@ -649,6 +761,62 @@ impl RbdState { self.sim_params_cpu = params; } + /// Sets the drift below which a contact pair keeps its contacts instead of computing them + /// again, before scaling by `length_unit` (see + /// [`RbdSimParams::normalized_contact_recycle_distance`]). `0` disables contact recycling. + pub fn set_contact_recycle_distance(&mut self, backend: &GpuBackend, normalized_distance: f32) { + let mut params = self.sim_params_cpu; + params.normalized_contact_recycle_distance = normalized_distance; + let _ = backend.write_buffer(self.sim_params.buffer_mut(), 0, &[params]); + self.sim_params_cpu = params; + } + + /// Selects the friction model for every rigid-body contact in every batch. + /// Switching models reallocates contact storage and resets contact warmstarts + /// and color seeds. Multibody contacts always use Coulomb. + #[cfg(feature = "dim3")] + pub fn set_friction_model( + &mut self, + backend: &GpuBackend, + model: crate::dynamics::FrictionModel, + ) { + if self.sim_params_cpu.friction_model != model { + let mut params = self.sim_params_cpu; + params.friction_model = model; + let usage = khal::BufferUsages::STORAGE | khal::BufferUsages::COPY_SRC; + let old = crate::dynamics::ContactConstraints::new( + backend, + self.contacts_capacity_cpu, + usage, + ¶ms, + ) + .expect("allocate contact model storage"); + let new = crate::dynamics::ContactConstraints::new( + backend, + self.contacts_capacity_cpu, + usage, + ¶ms, + ) + .expect("allocate contact model storage"); + // Adjacency counts gate every previous-contact read; clear them before + // exposing the newly allocated buffers to the next step. + let empty_counts = vec![0u32; self.old_constraints_counts.len() as usize]; + self.old_constraints_counts = + Tensor::vector(backend, &empty_counts, usage).expect("reset contact warmstarts"); + self.old_constraints = old; + self.new_constraints = new; + self.sim_params_cpu.friction_model = model; + let _ = backend.write_buffer(self.sim_params.buffer_mut(), 0, &[self.sim_params_cpu]); + self.graph_generation += 1; + } + } + + /// The simulation-wide rigid-body friction model. + #[cfg(feature = "dim3")] + pub fn friction_model(&self) -> crate::dynamics::FrictionModel { + self.sim_params_cpu.friction_model + } + /// PGS iterations per substep in the biased pass. #[cfg(feature = "dim3")] pub fn num_internal_pgs_iterations(&self) -> u32 { @@ -754,8 +922,8 @@ impl RbdState { /// GPU buffer holding the rigid-body contact constraints of the current /// step (impulses included). For debugging. - pub fn rigid_contact_constraints(&self) -> &Tensor { - &self.new_constraints + pub fn rigid_contact_constraints(&self) -> &crate::dynamics::ContactConstraints { + &self.old_constraints } /// Debug: reads the contacts back from the GPU, as world-space points of all batches. @@ -882,9 +1050,6 @@ impl RbdState { /// /// Reads the `old_*` buffers, i.e. the constraints and colors of the last step. pub fn debug_constraint_colors(&self, backend: &GpuBackend) -> Vec<(u32, u32, u32, u32, u32)> { - let cons: Vec = - futures::executor::block_on(backend.slow_read_vec(self.old_constraints.buffer())) - .unwrap_or_default(); let colors: Vec = futures::executor::block_on( backend.slow_read_vec(self.old_constraints_colors.buffer()), ) @@ -897,17 +1062,30 @@ impl RbdState { .and_then(|plan| plan.first().map(|p| p.bound as usize)) .unwrap_or(0); + let cons = futures::executor::block_on(self.old_constraints.read_impulses(backend, bound)) + .unwrap_or_default(); let mut out = Vec::new(); for (i, c) in cons.iter().enumerate().take(bound) { if c.len == 0 { continue; } let color = colors.get(i).copied().unwrap_or(u32::MAX); - out.push((i as u32, c.solver_body_a, c.solver_body_b, color, c.len)); + out.push((i as u32, c.body_a, c.body_b, color, c.len)); } out } + /// Diagnostic readback of the last step's active contact range. + pub async fn read_rigid_contact_impulses( + &self, + backend: &GpuBackend, + ) -> Result, khal::backend::GpuBackendError> { + let plan: Vec = backend.slow_read_vec(self.contact_plan.buffer()).await?; + self.old_constraints + .read_impulses(backend, plan[0].bound as usize) + .await + } + /// The number of colliders per batch. pub fn num_colliders_per_batch(&self) -> u32 { self.num_colliders_per_batch @@ -1286,7 +1464,7 @@ pub struct RbdGraphKey { pub num_active_colliders: u32, /// Collider slots per environment. pub num_colliders_per_batch: u32, - /// Maximum number of graph colors the solver sweeps. + /// Maximum number of graph colors the solver iterates over. pub max_colors: u32, /// Number of per-color uniform buffers. pub num_color_uniforms: usize, @@ -1311,8 +1489,8 @@ pub struct RbdGraphKey { /// Number of graph colors of the multibody impulse joints. #[cfg(feature = "dim3")] pub mb_imp_joint_num_colors: u32, - /// Whether the fused colored-sweep kernels are used. - pub fused_color_sweeps: bool, + /// Whether the fused colored kernels are used. + pub fused_color_dispatches: bool, } impl RbdState { @@ -1343,7 +1521,7 @@ impl RbdState { multibodies_empty: self.multibodies.is_empty(), #[cfg(feature = "dim3")] mb_imp_joint_num_colors: self.multibodies.mb_imp_joint_num_colors(), - fused_color_sweeps: crate::pipeline::RbdPipeline::fused_color_sweeps(self), + fused_color_dispatches: crate::pipeline::RbdPipeline::fused_color_dispatches(self), } } } diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index b60df585..e4f9c6cd 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -15,6 +15,7 @@ use crate::shaders::utils::BatchIndices; use crate::utils::PrefixSumWorkspace; use super::rbd_state::*; +use crate::dynamics::HubState; use khal::BufferUsages; use khal::backend::{GpuBackend, GpuReadback}; use vortx::tensor::Tensor; @@ -687,10 +688,7 @@ impl RbdState { .unwrap(); // Readback: pair count, PFM count, uncolored count (+ the multibody // contact-constraint demand on dim3). - #[cfg(feature = "dim3")] - let resize_readback = GpuReadback::new(backend, 4).unwrap(); - #[cfg(not(feature = "dim3"))] - let resize_readback = GpuReadback::new(backend, 3).unwrap(); + let resize_readback = GpuReadback::new(backend, ResizeFeedback::WORDS).unwrap(); let collision_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); @@ -704,7 +702,7 @@ impl RbdState { let mb_contact_index = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let contacts_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); - let mb_sweep_indirect = + let mb_dispatch_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let pfm_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); @@ -721,12 +719,23 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); - let old_constraints = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); - let old_constraint_builders = - Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); - let new_constraints = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); - let new_constraint_builders = - Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let old_constraints = crate::dynamics::ContactConstraints::new( + backend, + contacts_capacity, + storage, + &sim_params, + ) + .unwrap(); + let new_constraints = crate::dynamics::ContactConstraints::new( + backend, + contacts_capacity, + storage, + &sim_params, + ) + .unwrap(); + let recycle_states = + crate::dynamics::ContactRecycleStates::new(backend, contacts_capacity, storage) + .unwrap(); let constraints_colors = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let constraints_pending_colors = @@ -735,7 +744,8 @@ impl RbdState { Tensor::vector(backend, vec![0u32; contacts_capacity as usize], storage).unwrap(); let colored = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let constraints_rands = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); - let color_buckets_stride = capacities.solver_colors + 3; + let max_colors = capacities.minimum_solver_colors(num_colliders as u32, None); + let color_buckets_stride = max_colors + 3; let color_buckets = Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); // Clamped per-frame list totals (see `gpu_contact_plan`). @@ -749,7 +759,7 @@ impl RbdState { ) .unwrap(); let pfm_sort = PfmSortState::new(backend, pairs_capacity); - let color_sorted_ids = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let sorted_links = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); // Zeroed (not uninit): the first frame's warmstart transfer walks the // "old" counts before any step has written them; zero counts = empty // ranges. @@ -811,8 +821,9 @@ impl RbdState { vels: Tensor::vector(backend, &all_vels, storage | BufferUsages::COPY_DST).unwrap(), #[cfg(feature = "dim3")] reset_templates_bodies: None, - solver_vels: Tensor::vector(backend, &all_vels, storage).unwrap(), + solver_vels: HubState::solver_vels(backend, &all_vels, storage), solver_vels_inc: Tensor::vector(backend, &all_vels, storage).unwrap(), + hubs: HubState::new(backend, all_vels.len()), joints, #[cfg(feature = "dim3")] multibodies, @@ -859,7 +870,8 @@ impl RbdState { collision_pairs_indirect, contacts_capacity_cpu, collision_pairs_capacity_cpu, - collision_pairs_len_cpu: 0, + collision_pairs_len_cpu: None, + stepped_since_readback: false, graph_generation: 0, compute_graph: Default::default(), #[cfg(feature = "dim3")] @@ -871,22 +883,21 @@ impl RbdState { pfm_sort, #[cfg(feature = "dim3")] mb_contact_index, - mb_sweep_indirect, + mb_dispatch_indirect, pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, old_constraints, - old_constraint_builders, old_constraints_counts, new_constraints, - new_constraint_builders, + recycle_states, new_constraints_counts, constraints_colors, old_constraints_colors, colored, constraints_rands, color_buckets, - color_sorted_ids, + sorted_links, curr_color: Tensor::scalar( backend, 0u32, @@ -903,6 +914,7 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, ) .unwrap(), + coloring_dispatch: ColoringDispatch::new(backend), uncolored_staging: Tensor::scalar( backend, 0, @@ -914,7 +926,7 @@ impl RbdState { prefix_sum_workspace: PrefixSumWorkspace::default(), bucket_prefix_workspace: PrefixSumWorkspace::default(), lbvh: LbvhState::with_usages(backend, lbvh_usages), - max_colors: capacities.solver_colors, + max_colors, rb_contacts_inert, num_active_colliders: num_colliders as u32, num_active_bodies: num_bodies as u32, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index c79f38b1..8b470893 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -4,9 +4,9 @@ use crate::broad_phase::{GpuNarrowPhase, Lbvh}; #[cfg(feature = "dim3")] use crate::dynamics::GpuMultibodySolver; use crate::dynamics::{ - CanonicalContactsArgs, ColoringArgs, GpuCanonicalOrder, GpuColoring, GpuJointSolver, - GpuMpropsUpdate, GpuSolver, GpuWarmstart, JointSolverArgs, SolverArgs, - warmstart::WarmstartArgs, + CanonicalContactsArgs, ColorStatsBuffer, ColoringArgs, GpuCanonicalOrder, GpuColoring, + GpuJointSolver, GpuMassSplitting, GpuMpropsUpdate, GpuSolver, GpuWarmstart, JointSolverArgs, + SolverArgs, }; use crate::shaders::broad_phase::LbvhNode; use crate::utils::GpuPrefixSum; @@ -14,8 +14,10 @@ use khal::Shader; use super::lbvh_validation::validate_lbvh_topology; use super::rbd_state::*; +use bytemuck::Zeroable; use khal::BufferUsages; use khal::backend::{Backend, Encoder, GpuBackend, GpuBackendError, GpuTimestamps}; +use std::slice; use vortx::tensor::Tensor; /// The main GPU physics pipeline coordinating all simulation stages. @@ -31,6 +33,7 @@ pub struct RbdPipeline { lbvh: Lbvh, coloring: GpuColoring, warmstart: GpuWarmstart, + mass_splitting: GpuMassSplitting, /// Optional (default `false`): merge each collider pair's manifolds /// (e.g. per-triangle trimesh contacts) into one before the solvers. pub contact_reduction: bool, @@ -57,6 +60,7 @@ impl RbdPipeline { lbvh: Lbvh::from_backend(backend), coloring: GpuColoring::from_backend(backend)?, warmstart: GpuWarmstart::from_backend(backend)?, + mass_splitting: GpuMassSplitting::from_backend(backend)?, contact_reduction: false, canonical_order: GpuCanonicalOrder::from_backend(backend)?, }) @@ -94,7 +98,7 @@ impl RbdPipeline { self.step_impl(backend, state, timestamps, encoder, false) } - /// Whether the step uses the fused colored-sweep kernels (one workgroup per + /// Whether the step uses the fused colored kernels (one workgroup per /// batch walking every color) instead of one dispatch per color. /// /// Chosen from the expected pair count: small pair counts with many @@ -102,18 +106,17 @@ impl RbdPipeline { /// at any size, just serialized past ~64 lanes. /// /// The estimate is per batch: the read-back counter (or the capacity when - /// the readback is disabled) is a total over the whole flat pair buffer. - pub fn fused_color_sweeps(state: &RbdState) -> bool { + /// the readback is disabled or hasn't completed yet) is a total over the whole flat pair + /// buffer. + pub fn fused_color_dispatches(state: &RbdState) -> bool { let readback_enabled = state.capacities.solver_colors_resize_policy != RbdResizePolicy::Fixed || state.capacities.collisions_resize_policy != RbdResizePolicy::Fixed; - let est_pairs = if readback_enabled { - state.collision_pairs_len_cpu.div_ceil(state.num_batches) - } else { - state - .collision_pairs_capacity_cpu - .div_ceil(state.num_batches) + let pairs = match state.collision_pairs_len_cpu { + Some(len) if readback_enabled => len, + _ => state.collision_pairs_capacity_cpu, }; + let est_pairs = pairs.div_ceil(state.num_batches); est_pairs <= 128 } @@ -136,6 +139,7 @@ impl RbdPipeline { Ok(()) }; let mut stats = RunStats::default(); + state.stepped_since_readback = true; // A multibody capacity edit (e.g. reserving the dry-friction // constraint slots on the first `set_dof_frictionloss`) leaves the @@ -176,7 +180,7 @@ impl RbdPipeline { solver_vels: &mut state.solver_vels, batch_indices: &state.batch_indices, color_uniforms: &state.color_uniforms, - mb_sweep_indirect: &state.mb_sweep_indirect, + mb_dispatch_indirect: &state.mb_dispatch_indirect, gravity: &state.gravity, friction_in_bias_pass: state.sim_params_cpu.friction_in_bias_pass != 0, }; @@ -295,11 +299,11 @@ impl RbdPipeline { } } - let fused_color_sweeps = Self::fused_color_sweeps(state); + let fused_color_dispatches = Self::fused_color_dispatches(state); // In small scenes, submit less frequently. In big scenes submit more // to overlap compute and encoding. - let merge_submits = fused_color_sweeps && state.num_batches <= 64; + let merge_submits = fused_color_dispatches && state.num_batches <= 64; // Phase 2a: Narrow phase. Split out from solver-prep + coloring // so its CPU encoding overlaps with Phase 1's GPU work and its @@ -319,7 +323,7 @@ impl RbdPipeline { &mut state.contacts, &mut state.contacts_indirect, &mut state.contact_plan, - &mut state.mb_sweep_indirect, + &mut state.mb_dispatch_indirect, &mut state.pfm_pairs, &mut state.pfm_pairs_len, &mut state.pfm_pairs_indirect, @@ -382,7 +386,10 @@ impl RbdPipeline { stable_mb_contact_index: None, contacts_len_indirect: &state.contacts_indirect, constraints: &mut state.new_constraints, - constraint_builders: &mut state.new_constraint_builders, + old_constraints: &state.old_constraints, + old_body_constraint_counts: &state.old_constraints_counts, + old_body_constraint_ids: &state.old_body_constraint_ids, + recycle_states: &mut state.recycle_states, sim_params: &state.sim_params, body_poses: &mut state.body_poses, solver_body_poses: &mut state.solver_body_poses, @@ -396,7 +403,10 @@ impl RbdPipeline { body_constraint_counts: &mut state.new_constraints_counts, body_constraint_ids: &mut state.new_body_constraint_ids, color_buckets: &state.color_buckets, - color_sorted_ids: &state.color_sorted_ids, + sorted_links: &state.sorted_links, + color_dispatch_indirect: &state.coloring_dispatch.dispatch_indirect, + color_dispatch_threads: &state.coloring_dispatch.dispatch_threads, + fuse_tail_colors: state.coloring_dispatch.fuse_tail_colors, color_uniforms: &state.color_uniforms, prefix_sum: &self.prefix_sum, num_colors: 0, @@ -407,12 +417,15 @@ impl RbdPipeline { body_group: &state.body_group, body_is_multibody: &state.body_is_multibody, batch_indices: &state.batch_indices, - mb_sweep_indirect: &state.mb_sweep_indirect, + mb_dispatch_indirect: &state.mb_dispatch_indirect, colorless_warmstart: false, - fused_color_sweeps, + fused_color_dispatches, rb_contacts_inert: state.rb_contacts_inert, friction_in_bias_pass: state.sim_params_cpu.friction_in_bias_pass != 0, gravity: &state.gravity, + hubs: &mut state.hubs, + mass_splitting: &self.mass_splitting, + scale_warmstart_impulses: state.sim_params_cpu.warmstart_coefficient != 1.0, }; self.solver.prepare( backend, @@ -442,26 +455,14 @@ impl RbdPipeline { stats.num_colors = state.max_colors + 1; drop(pass); } else { - // Warmstart - let warmstart_args = WarmstartArgs { - contact_plan: &state.contact_plan, - old_body_constraint_counts: &state.old_constraints_counts, - old_constraint_builders: &state.old_constraint_builders, - old_body_constraint_ids: &state.old_body_constraint_ids, - old_constraints: &state.old_constraints, - new_constraints: &mut state.new_constraints, - new_constraint_builders: &state.new_constraint_builders, - contacts_len_indirect: &state.contacts_indirect, - }; - - self.warmstart - .transfer_warmstart_impulses(&mut pass, warmstart_args)?; + drop(pass); + let mut pass = encoder.begin_pass("[RBD] prep/coloring", timestamps.as_deref_mut()); let coloring_args = ColoringArgs { contacts_len_indirect: &state.contacts_indirect, body_constraint_counts: &state.new_constraints_counts, body_constraint_ids: &state.new_body_constraint_ids, - constraints: &state.new_constraints, + links: &state.new_constraints.links, constraints_colors: &mut state.constraints_colors, constraints_pending_colors: &mut state.determinism.pending_colors, constraints_rands: &mut state.constraints_rands, @@ -472,6 +473,7 @@ impl RbdPipeline { colored: &mut state.colored, batch_indices: &state.batch_indices, body_group: &state.body_group, + coloring_indirect: &mut state.coloring_dispatch.coloring_indirect, }; self.coloring .dispatch_topo_gc_reset(&mut pass, coloring_args)?; @@ -481,10 +483,7 @@ impl RbdPipeline { // topo-gc iterations converge in 1-2 rounds instead of ~num_colors). let seed_args = crate::dynamics::warmstart::SeedColorsArgs { contact_plan: &state.contact_plan, - old_body_constraint_counts: &state.old_constraints_counts, - old_body_constraint_ids: &state.old_body_constraint_ids, - old_constraints: &state.old_constraints, - new_constraints: &state.new_constraints, + links: &state.new_constraints.links, old_constraints_colors: &state.old_constraints_colors, constraints_colors: &mut state.constraints_colors, colored: &mut state.colored, @@ -497,7 +496,7 @@ impl RbdPipeline { contacts_len_indirect: &state.contacts_indirect, body_constraint_counts: &state.new_constraints_counts, body_constraint_ids: &state.new_body_constraint_ids, - constraints: &state.new_constraints, + links: &state.new_constraints.links, constraints_colors: &mut state.constraints_colors, constraints_pending_colors: &mut state.determinism.pending_colors, constraints_rands: &mut state.constraints_rands, @@ -508,6 +507,7 @@ impl RbdPipeline { colored: &mut state.colored, batch_indices: &state.batch_indices, body_group: &state.body_group, + coloring_indirect: &mut state.coloring_dispatch.coloring_indirect, }; self.coloring.dispatch_topo_gc_iterations( &mut pass, @@ -516,15 +516,20 @@ impl RbdPipeline { )?; // Bucket-sort the constraint ids by color so each colored solver - // sweep only touches its own constraints. + // iteration only touches its own constraints. + drop(pass); + let mut pass = encoder.begin_pass("[RBD] prep/buckets", timestamps.as_deref_mut()); let bucket_args = crate::dynamics::ColorBucketsArgs { contacts_len_indirect: &state.contacts_indirect, constraints_colors: &state.constraints_colors, - constraints: &state.new_constraints, + links: &state.new_constraints.links, contact_plan: &state.contact_plan, color_buckets: &mut state.color_buckets, - color_sorted_ids: &mut state.color_sorted_ids, + sorted_links: &mut state.sorted_links, + constraint_indices: &mut state.new_constraints.constraint_indices, batch_indices: &state.batch_indices, + color_stats: &mut state.coloring_dispatch.color_stats, + dispatch_indirect: &mut state.coloring_dispatch.dispatch_indirect, }; self.coloring.dispatch_build_color_buckets( backend, @@ -548,7 +553,7 @@ impl RbdPipeline { let num_colors = stats.num_colors; // Create solver_args for solve phase (after coloring is complete) - let solver_args = SolverArgs { + let mut solver_args = SolverArgs { contacts: &state.contacts, contact_plan: &state.contact_plan, #[cfg(feature = "dim3")] @@ -560,7 +565,10 @@ impl RbdPipeline { .then_some(&mut state.determinism.stable_mb_contact_index), contacts_len_indirect: &state.contacts_indirect, constraints: &mut state.new_constraints, - constraint_builders: &mut state.new_constraint_builders, + old_constraints: &state.old_constraints, + old_body_constraint_counts: &state.old_constraints_counts, + old_body_constraint_ids: &state.old_body_constraint_ids, + recycle_states: &mut state.recycle_states, sim_params: &state.sim_params, body_poses: &mut state.body_poses, solver_body_poses: &mut state.solver_body_poses, @@ -574,7 +582,10 @@ impl RbdPipeline { body_constraint_counts: &mut state.new_constraints_counts, body_constraint_ids: &mut state.new_body_constraint_ids, color_buckets: &state.color_buckets, - color_sorted_ids: &state.color_sorted_ids, + sorted_links: &state.sorted_links, + color_dispatch_indirect: &state.coloring_dispatch.dispatch_indirect, + color_dispatch_threads: &state.coloring_dispatch.dispatch_threads, + fuse_tail_colors: state.coloring_dispatch.fuse_tail_colors, color_uniforms: &state.color_uniforms, prefix_sum: &self.prefix_sum, num_colors, @@ -585,19 +596,28 @@ impl RbdPipeline { body_group: &state.body_group, body_is_multibody: &state.body_is_multibody, batch_indices: &state.batch_indices, - mb_sweep_indirect: &state.mb_sweep_indirect, + mb_dispatch_indirect: &state.mb_dispatch_indirect, // The gather warmstart is only valid without multibody grouping; // see `SolverArgs::colorless_warmstart`. #[cfg(feature = "dim3")] colorless_warmstart: state.multibodies.is_empty(), #[cfg(not(feature = "dim3"))] colorless_warmstart: true, - fused_color_sweeps, + fused_color_dispatches, rb_contacts_inert: state.rb_contacts_inert, friction_in_bias_pass: state.sim_params_cpu.friction_in_bias_pass != 0, gravity: &state.gravity, + hubs: &mut state.hubs, + mass_splitting: &self.mass_splitting, + scale_warmstart_impulses: state.sim_params_cpu.warmstart_coefficient != 1.0, }; + // The constraints follow the color order, known only now. + { + let mut pass = encoder.begin_pass("[RBD] prep/constraints", timestamps.as_deref_mut()); + self.solver.build_constraints(&mut pass, &mut solver_args)?; + } + // Phase 3: Solve constraints let joint_solver_args = JointSolverArgs { num_batches: state.num_batches, @@ -635,10 +655,7 @@ impl RbdPipeline { // Swap buffers for warm-starting next frame std::mem::swap(&mut state.old_constraints, &mut state.new_constraints); - std::mem::swap( - &mut state.old_constraint_builders, - &mut state.new_constraint_builders, - ); + state.recycle_states.swap_frames(); std::mem::swap( &mut state.old_body_constraint_ids, &mut state.new_body_constraint_ids, @@ -670,29 +687,62 @@ impl RbdPipeline { != RbdResizePolicy::Fixed || state.capacities.collisions_resize_policy != RbdResizePolicy::Fixed; - // The readback holds `[total collision-pair count, total PFM-pair - // count, uncolored count]` (+ the multibody contact-constraint demand - // on dim3). - #[cfg(feature = "dim3")] - let mut counts = [0u32; 4]; - #[cfg(not(feature = "dim3"))] - let mut counts = [0u32; 3]; - if state.resize_readback.try_take(backend, &mut counts) { + let mut feedback = ResizeFeedback::zeroed(); + if state.resize_readback.try_take( + backend, + bytemuck::cast_slice_mut(slice::from_mut(&mut feedback)), + ) { + let stats = feedback.color_stats; + for (threads, size) in state + .coloring_dispatch + .dispatch_threads + .iter_mut() + .zip(&stats.sizes) + { + // Always launch some work, even if the previous frame's bucket was empty. + // A stale/small hint only changes how much the shader's for loop strides. + *threads = size.div_ceil(64).clamp(8, 256) * 64; + } + // Empty/short late colors share one workgroup. An old hint never skips + // contacts: if the graph changes, that workgroup simply does more work. + // Color 64 has no individual readback slot, so require it to be unused. + state.coloring_dispatch.fuse_tail_colors = stats.highest_color < 64 + && stats.sizes[crate::shaders::dynamics::contact_tiles::TAIL_COLOR as usize..] + .iter() + .map(|&n| n as u64) + .sum::() + <= 512; // TODO: make the coloring update optional (and pre-configurable) too? // The flat pair and PFM work-lists share one buffer capacity, so // whichever is larger drives that resize. - let pairs_len = counts[0].max(counts[1]); - let coloring_converged = counts[2]; - state.collision_pairs_len_cpu = counts[0]; + let pairs_len = feedback.collision_pairs.max(feedback.pfm_pairs); + let coloring_converged = feedback.coloring_converged; + state.collision_pairs_len_cpu = Some(feedback.collision_pairs); let nb = state.num_batches; // TODO: Fit will act like Grow. To be able to auto-shrink the max color count, we need // to readback the actual color count. This would also allow us to grow the color // count earlier, before it gets a chance to fail. + let min_colors = state + .capacities + .minimum_solver_colors( + state.num_active_colliders, + (coloring_converged != 0).then_some(stats.highest_color), + ) + .max(1); let grow_colors = state.capacities.solver_colors_resize_policy != RbdResizePolicy::Fixed - && coloring_converged == 0 + && (coloring_converged == 0 || state.max_colors < min_colors) && !state.rb_contacts_inert; + // Every color costs a coloring iteration and a dispatch per iteration: once the + // coloring converges, shrink the budget back to a few colors above those in use. + let highest_color = stats.highest_color; + let shrink_colors = state.capacities.solver_colors_resize_policy + != RbdResizePolicy::Fixed + && coloring_converged != 0 + && !state.rb_contacts_inert + && highest_color + 8 < state.max_colors + && state.max_colors > min_colors; // Decide every resize up front, then drain the GPU once before // applying any of them: `rebuild_batch_indices` rewrites the shared @@ -713,7 +763,7 @@ impl RbdPipeline { #[cfg(feature = "dim3")] let (resize_mb, new_mb) = { - let mb_demand = counts[3]; + let mb_demand = feedback.mb_cons_demand; state.mb_cons_demand_cpu = mb_demand; let mb_capacity = state.multibodies.contact_constraints_capacity(); let safe_mb = mb_demand.saturating_add(mb_demand / 4); @@ -736,13 +786,17 @@ impl RbdPipeline { #[cfg(not(feature = "dim3"))] let resize_mb = false; - if grow_colors || resize_pairs || resize_mb { + if grow_colors || shrink_colors || resize_pairs || resize_mb { backend.synchronize()?; state.graph_generation += 1; } - if grow_colors { - state.max_colors += 5; + if grow_colors || shrink_colors { + if grow_colors { + state.max_colors = (state.max_colors + 5).max(min_colors); + } else { + state.max_colors = (highest_color + 4).max(min_colors); + } // The color-bucket buffer is strided by `max_colors + 3`: // regrow it and update the stride in `BatchIndices`. @@ -778,14 +832,22 @@ impl RbdPipeline { { state.mb_contact_index = Tensor::vector_uninit(backend, new_contacts, storage)?; } - state.old_constraints = Tensor::vector_uninit(backend, new_contacts, storage)?; - state.old_constraint_builders = - Tensor::vector_uninit(backend, new_contacts, storage)?; + state.old_constraints = crate::dynamics::ContactConstraints::new( + backend, + new_contacts, + storage, + &state.sim_params_cpu, + )?; state.old_body_constraint_ids = Tensor::vector_uninit(backend, new_contacts * 2, storage)?; - state.new_constraints = Tensor::vector_uninit(backend, new_contacts, storage)?; - state.new_constraint_builders = - Tensor::vector_uninit(backend, new_contacts, storage)?; + state.new_constraints = crate::dynamics::ContactConstraints::new( + backend, + new_contacts, + storage, + &state.sim_params_cpu, + )?; + state.recycle_states = + crate::dynamics::ContactRecycleStates::new(backend, new_contacts, storage)?; state.new_body_constraint_ids = Tensor::vector_uninit(backend, new_contacts * 2, storage)?; state.constraints_colors = Tensor::vector_uninit(backend, new_contacts, storage)?; @@ -797,7 +859,7 @@ impl RbdPipeline { Tensor::vector(backend, vec![0u32; new_contacts as usize], storage)?; state.colored = Tensor::vector_uninit(backend, new_contacts, storage)?; state.constraints_rands = Tensor::vector_uninit(backend, new_contacts, storage)?; - state.color_sorted_ids = Tensor::vector_uninit(backend, new_contacts, storage)?; + state.sorted_links = Tensor::vector_uninit(backend, new_contacts, storage)?; // The old counts index the old (now discarded) constraint list: // zero them so the next warmstart transfer sees empty ranges // instead of stale offsets into the fresh buffers. @@ -817,26 +879,26 @@ impl RbdPipeline { } } - if readback_enabled && state.resize_readback.is_idle() { - #[cfg(feature = "dim3")] - state.resize_readback.request( - backend, - &[ - (state.collision_pairs_len.buffer(), 0, 1), - (state.pfm_pairs_len.buffer(), 0, 1), - (state.uncolored.buffer(), 0, 1), - (state.multibodies.mb_cons_demand().buffer(), 0, 1), - ], - )?; - #[cfg(not(feature = "dim3"))] - state.resize_readback.request( - backend, - &[ - (state.collision_pairs_len.buffer(), 0, 1), - (state.pfm_pairs_len.buffer(), 0, 1), - (state.uncolored.buffer(), 0, 1), - ], - )?; + if readback_enabled && state.stepped_since_readback && state.resize_readback.is_idle() { + state.stepped_since_readback = false; + // Gathered in the field order of `ResizeFeedback`. + let sources = [ + (state.collision_pairs_len.buffer(), 0, 1), + (state.pfm_pairs_len.buffer(), 0, 1), + (state.uncolored.buffer(), 0, 1), + #[cfg(feature = "dim3")] + (state.multibodies.mb_cons_demand().buffer(), 0, 1), + ( + state.coloring_dispatch.color_stats.0.buffer(), + 0, + ColorStatsBuffer::WORDS, + ), + ]; + debug_assert_eq!( + sources.iter().map(|s| s.2).sum::(), + ResizeFeedback::WORDS + ); + state.resize_readback.request(backend, &sources)?; } Ok(()) diff --git a/src_rbd/pipeline/test_batched_stacks.rs b/src_rbd/pipeline/test_batched_stacks.rs index 26e995a2..39d89847 100644 --- a/src_rbd/pipeline/test_batched_stacks.rs +++ b/src_rbd/pipeline/test_batched_stacks.rs @@ -420,12 +420,26 @@ fn build_mb_env(ball_colliders: bool) -> (RigidBodySet, ColliderSet, MultibodyJo #[serial_test::serial] #[ignore] async fn test_stacks_6_multibody() { + run_multibody_case(crate::dynamics::FrictionModel::Coulomb).await; +} + +#[futures_test::test] +#[serial_test::serial] +#[ignore = "requires a WebGPU adapter"] +async fn twist_friction_keeps_multibody_contacts_on_coulomb() { + run_multibody_case(crate::dynamics::FrictionModel::Simplified).await; +} + +async fn run_multibody_case(friction_model: crate::dynamics::FrictionModel) { let backend = test_backend().await; let num_envs = 4u32; let envs: Vec<_> = (0..num_envs).map(|_| build_mb_env(false)).collect(); let joints = ImpulseJointSet::new(); - let params = RbdSimParams::tgs_soft(); + let params = RbdSimParams { + friction_model, + ..RbdSimParams::tgs_soft() + }; let refs: Vec<_> = envs .iter() .map(|(b, c, mb)| (b, c, &joints, mb, ¶ms)) diff --git a/src_rbd/pipeline/test_color_buckets.html b/src_rbd/pipeline/test_color_buckets.html new file mode 100644 index 00000000..fbeedd36 --- /dev/null +++ b/src_rbd/pipeline/test_color_buckets.html @@ -0,0 +1,139 @@ + + +Contact bucket WebGPU regression +
RUNNING
+ diff --git a/src_rbd/pipeline/test_color_buckets.rs b/src_rbd/pipeline/test_color_buckets.rs new file mode 100644 index 00000000..5f08dad2 --- /dev/null +++ b/src_rbd/pipeline/test_color_buckets.rs @@ -0,0 +1,231 @@ +//! GPU checks for bucket counts, membership, and the generic scatter fallback. +use crate::shaders::broad_phase::ContactPlan; +use crate::shaders::dynamics::*; +use crate::shaders::utils::BatchIndices; +use khal::BufferUsages; +use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; +use vortx::tensor::Tensor; + +#[futures_test::test] +#[serial_test::serial] +#[ignore = "requires a WebGPU adapter"] +async fn color_bucket_histograms_preserve_membership_and_bounds() { + let backend = GpuBackend::WebGpu( + WebGpu::new(Default::default(), Default::default()) + .await + .unwrap(), + ); + let count = GpuColorBucketsCount::from_dir(&backend, &crate::SPIRV_DIR).unwrap(); + let scatter = GpuColorBucketsScatter::from_dir(&backend, &crate::SPIRV_DIR).unwrap(); + let storage = BufferUsages::STORAGE | BufferUsages::COPY_SRC; + for n in [0u32, 1, 63, 64, 65, 193] { + for nb in [1u32, 3] { + for stride in [5u32, 67, 129, 133] { + for mixed in [false, true] { + let mut colors = vec![0u32; n.max(1) as usize]; + let mut input = vec![ContactLink::default(); n.max(1) as usize]; + let mut expected = vec![Vec::::new(); (stride * nb) as usize]; + for i in 0..n as usize { + let color = if mixed { + [0, 1, 2, 63, 64, 127, stride - 1, u32::MAX][i % 8] + } else { + 1 + }; + colors[i] = color; + input[i].solver_body_a = (i as u32 % 17) * nb + (i as u32 / 7) % nb; + input[i].contact = i as u32; + // Some inactive slots, left out of the tiles. + input[i].len = u32::from(!mixed || i % 5 != 4); + if input[i].len == 0 { + continue; + } + // Active constraints left uncolored go to the last bucket. + let bucket = if color != 0 && color < stride - 1 { + color + } else { + stride - 1 + }; + expected[(bucket * nb + input[i].solver_body_a % nb) as usize] + .push(i as u32); + } + let links = Tensor::vector(&backend, input, storage).unwrap(); + let colors = Tensor::vector(&backend, colors, storage).unwrap(); + let plan = Tensor::scalar( + &backend, + ContactPlan { + bound: n, + ..Default::default() + }, + BufferUsages::UNIFORM, + ) + .unwrap(); + let ids = Tensor::scalar( + &backend, + BatchIndices { + num_batches: nb, + solver_color_buckets_stride: stride, + ..Default::default() + }, + BufferUsages::UNIFORM, + ) + .unwrap(); + let sizes = expected.iter().map(|x| x.len() as u32).collect::>(); + let mut starts = Vec::new(); + let mut ends = Vec::new(); + let mut total = 0u32; + for &size in &sizes { + starts.push(total); + total += size; + ends.push(total); + } + let mut grids = vec![64u32, n.max(1).div_ceil(64) * 64]; + grids.dedup(); + for grid in grids { + let mut counts = + Tensor::vector(&backend, vec![0u32; sizes.len()], storage).unwrap(); + let mut cursors = + Tensor::vector(&backend, starts.clone(), storage).unwrap(); + // Stale links, which the count deactivates past the constraints. + let stale = ContactLink { + len: 7, + ..Default::default() + }; + let mut sorted_links = + Tensor::vector(&backend, vec![stale; n.max(1) as usize], storage) + .unwrap(); + let mut constraint_indices = + Tensor::vector(&backend, vec![0u32; n.max(1) as usize], storage) + .unwrap(); + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("color-bucket-test", None); + count + .call( + &mut pass, + grid, + &colors, + &links, + &plan, + &mut counts, + &ids, + &mut sorted_links, + ) + .unwrap(); + scatter + .call( + &mut pass, + grid, + &colors, + &links, + &plan, + &mut cursors, + &mut sorted_links, + &ids, + &mut constraint_indices, + ) + .unwrap(); + drop(pass); + backend.submit(encoder).unwrap(); + assert_eq!(backend.slow_read_vec(counts.buffer()).await.unwrap(), sizes); + assert_eq!(backend.slow_read_vec(cursors.buffer()).await.unwrap(), ends); + let actual: Vec = + backend.slow_read_vec(sorted_links.buffer()).await.unwrap(); + assert!( + actual[total as usize..n as usize] + .iter() + .all(|l| l.len == 0) + ); + let constraint_indices = backend + .slow_read_vec(constraint_indices.buffer()) + .await + .unwrap(); + for i in (0..n).filter(|&i| !mixed || i % 5 != 4) { + assert_eq!(actual[constraint_indices[i as usize] as usize].contact, i); + } + for (bucket, expected) in expected.iter().enumerate() { + let mut members: Vec = actual + [starts[bucket] as usize..ends[bucket] as usize] + .iter() + .map(|l| l.contact) + .collect(); + members.sort_unstable(); + assert_eq!( + &members, expected, + "n={n}, nb={nb}, stride={stride}, grid={grid}, bucket={bucket}" + ); + } + } + } + } + } + } +} + +/// Exercise the exact SPIR-V -> WGSL path used by wgpu's browser backend. +/// Set NEXUS_WEB_TEST_DIR to export the browser regression page and shaders. +/// Native shader validation does not catch Dawn's stricter uniformity errors. +#[test] +fn color_bucket_shaders_translate_for_web() { + use khal::re_exports::wgpu::naga; + use std::mem::{offset_of, size_of}; + + let out = std::env::var_os("NEXUS_WEB_TEST_DIR").map(std::path::PathBuf::from); + if let Some(out) = &out { + std::fs::create_dir_all(out).unwrap(); + std::fs::write( + out.join("index.html"), + include_str!("test_color_buckets.html"), + ) + .unwrap(); + } + let dim = if cfg!(feature = "dim3") { "3d" } else { "2d" }; + let mut descriptions = Vec::new(); + for model in ["links"] { + for phase in ["count", "scatter"] { + let file = format!("dynamics-color_buckets-gpu_color_buckets_{phase}.spv"); + let spirv = crate::SPIRV_DIR.get_file(&file).unwrap().contents(); + let module = naga::front::spv::parse_u8_slice( + spirv, + &naga::front::spv::Options { + adjust_coordinate_space: false, + strict_capabilities: true, + block_ctx_dump_prefix: None, + }, + ) + .unwrap(); + let info = naga::valid::Validator::new( + naga::valid::ValidationFlags::all(), + naga::valid::Capabilities::all(), + ) + .validate(&module) + .unwrap(); + let wgsl = naga::back::wgsl::write_string( + &module, + &info, + naga::back::wgsl::WriterFlags::empty(), + ) + .unwrap(); + if let Some(out) = &out { + std::fs::write(out.join(format!("{dim}-{model}-{phase}.wgsl")), wgsl).unwrap(); + } + } + descriptions.push(format!( + r#"{{"model":"{model}","linkSize":{},"bodyOffset":{},"lenOffset":{},"contactOffset":{}}}"#, + size_of::(), + offset_of!(ContactLink, solver_body_a), + offset_of!(ContactLink, len), + offset_of!(ContactLink, contact), + )); + } + if let Some(out) = &out { + let metadata = format!( + r#"{{"planSize":{},"boundOffset":{},"batchSize":{},"numBatchesOffset":{},"strideOffset":{},"models":[{}]}}"#, + size_of::(), + offset_of!(ContactPlan, bound), + size_of::(), + offset_of!(BatchIndices, num_batches), + offset_of!(BatchIndices, solver_color_buckets_stride), + descriptions.join(","), + ); + std::fs::write(out.join(format!("{dim}.json")), metadata).unwrap(); + } +} diff --git a/src_rbd/pipeline/test_large_warmstart.rs b/src_rbd/pipeline/test_large_warmstart.rs new file mode 100644 index 00000000..63e9f3fb --- /dev/null +++ b/src_rbd/pipeline/test_large_warmstart.rs @@ -0,0 +1,127 @@ +//! Exercises the large-scene warmstart scratch on WebGPU, in both dimensions. +use super::{RbdCapacities, RbdPipeline, RbdState}; +use crate::math::Pose; +use crate::rapier::prelude::*; +use crate::shaders::dynamics::RbdSimParams; +use khal::backend::{Backend, GpuBackend, WebGpu}; + +#[futures_test::test] +#[serial_test::serial] +#[ignore = "requires a WebGPU adapter"] +async fn large_warmstart_keeps_all_boxes_on_ground() { + check_large_warmstart(16, RbdSimParams::tgs_soft(), false).await; +} + +#[futures_test::test] +#[serial_test::serial] +#[ignore = "requires a WebGPU adapter"] +async fn large_warmstart_supports_default_binding_limit() { + check_large_warmstart(8, RbdSimParams::tgs_soft(), false).await; + #[cfg(feature = "dim3")] + { + let mut params = RbdSimParams::tgs_soft(); + params.friction_model = crate::shaders::dynamics::FrictionModel::Simplified; + check_large_warmstart(8, params, false).await; + } +} + +#[futures_test::test] +#[serial_test::serial] +#[ignore = "requires a WebGPU adapter"] +async fn large_warmstart_resizes_within_default_binding_limit() { + check_large_warmstart(8, RbdSimParams::tgs_soft(), true).await; + #[cfg(feature = "dim3")] + { + let mut params = RbdSimParams::tgs_soft(); + params.friction_model = crate::shaders::dynamics::FrictionModel::Simplified; + check_large_warmstart(8, params, true).await; + } +} + +async fn check_large_warmstart(storage_buffers: u32, params: RbdSimParams, resize: bool) { + let backend = GpuBackend::WebGpu( + WebGpu::new( + Default::default(), + khal::re_exports::wgpu::Limits { + max_storage_buffers_per_shader_stage: storage_buffers, + ..Default::default() + }, + ) + .await + .unwrap(), + ); + let mut bodies = RigidBodySet::new(); + let mut colliders = ColliderSet::new(); + #[cfg(feature = "dim3")] + let (ground_pos, ground_shape) = ( + Vector::new(0.0, -0.5, 0.0), + ColliderBuilder::cuboid(50.0, 0.5, 50.0), + ); + #[cfg(feature = "dim2")] + let (ground_pos, ground_shape) = (Vector::new(0.0, -0.5), ColliderBuilder::cuboid(2000.0, 0.5)); + let ground = bodies.insert(RigidBodyBuilder::fixed().translation(ground_pos)); + colliders.insert_with_parent(ground_shape, ground, &mut bodies); + for i in 0..4096 { + #[cfg(feature = "dim3")] + let (pos, shape) = ( + Vector::new((i % 64) as f32 * 0.7, 0.25, (i / 64) as f32 * 0.7), + ColliderBuilder::cuboid(0.2, 0.2, 0.2), + ); + #[cfg(feature = "dim2")] + let (pos, shape) = ( + Vector::new((i as f32 - 2048.0) * 0.7, 0.25), + ColliderBuilder::cuboid(0.2, 0.2), + ); + let handle = bodies.insert( + RigidBodyBuilder::dynamic() + .can_sleep(false) + .translation(pos), + ); + colliders.insert_with_parent(shape, handle, &mut bodies); + } + let mut state = RbdState::from_rapier( + &backend, + &[( + &bodies, + &colliders, + &ImpulseJointSet::new(), + &MultibodyJointSet::new(), + ¶ms, + )], + RbdCapacities { + collisions_capacity: if resize { 64 } else { 8192 }, + ..Default::default() + }, + ); + let pipeline = RbdPipeline::new(&backend).unwrap(); + let initial_contact_capacity = state.contacts_capacity_cpu; + for step in 0..120 { + if resize { + // Readbacks must finish so the deliberately undersized buffers grow + // before falling boxes cross the ground. Toggle the ordinary + // deterministic path after both arena phases have been exercised. + backend.synchronize().unwrap(); + if step == 61 || step == 80 { + state.set_deterministic(&backend, step == 61); + } + } + pipeline.auto_resize_buffers(&backend, &mut state).unwrap(); + pipeline.step(&backend, &mut state, None).unwrap(); + } + backend.synchronize().unwrap(); + if resize { + assert!(state.contacts_capacity_cpu > initial_contact_capacity); + } + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + for (i, pose) in poses.iter().enumerate().skip(1).take(4096) { + assert!(pose.translation.is_finite(), "box {i}: non-finite pose"); + assert!( + (pose.translation.y - 0.2).abs() < 0.01, + "box {i}: height {}", + pose.translation.y + ); + } +} diff --git a/src_rbd/pipeline/test_polyline_warmstart.rs b/src_rbd/pipeline/test_polyline_warmstart.rs index 00c7b29b..620c24ac 100644 --- a/src_rbd/pipeline/test_polyline_warmstart.rs +++ b/src_rbd/pipeline/test_polyline_warmstart.rs @@ -102,6 +102,10 @@ async fn rest_report(num_envs: usize, settle: u32, steps: u32) -> RestReport { .map(|(b, c)| (b, c, &impulse_joints, &multibody_joints, ¶ms)) .collect(); let capacities = RbdCapacities { + // Size each nine-body fixture explicitly: defaults multiplied by 256 + // environments exceed WebGPU's default buffer-size limit. + body_capacity: 32, + collisions_capacity: 512, batches: num_envs as u32, ..Default::default() }; diff --git a/src_rbd/pipeline/test_tail_colors.rs b/src_rbd/pipeline/test_tail_colors.rs new file mode 100644 index 00000000..84d38f28 --- /dev/null +++ b/src_rbd/pipeline/test_tail_colors.rs @@ -0,0 +1,248 @@ +//! Compares fused color tails with dispatch-by-dispatch GPU execution. +use crate::dynamics::contact_kernels::{ + ContactTiles, GpuSolveConstraints, GpuSolveConstraintsBiased, GpuSolveConstraintsTail, + GpuSolveConstraintsTailBiased, GpuSolveConstraintsTailUnbiased, + GpuSolveConstraintsTailUnbiasedCached, GpuSolveConstraintsUnbiased, + GpuSolveConstraintsUnbiasedCached, +}; +use crate::math::Pose; +use crate::shaders::dynamics::contact_tiles::{TAIL_COLOR, TILE_LEN, tile_lane}; +use crate::shaders::dynamics::coulomb_tiles::CoulombTile; +use crate::shaders::dynamics::twist_tiles::TwistTile; +use crate::shaders::dynamics::*; +use crate::shaders::utils::BatchIndices; +use glamx::{Vec2, Vec3}; +use khal::BufferUsages; +use khal::backend::{Backend, Encoder, GpuBackend, WebGpu}; +use vortx::tensor::Tensor; + +#[futures_test::test] +#[serial_test::serial] +#[ignore = "requires a WebGPU adapter"] +async fn tail_colors_match_separate_color_dispatches() { + for model in [FrictionModel::Coulomb, FrictionModel::Simplified] { + check_tail_colors(model).await; + } +} + +async fn check_tail_colors(model: FrictionModel) { + let backend = GpuBackend::WebGpu( + WebGpu::new(Default::default(), Default::default()) + .await + .unwrap(), + ); + let dir = &crate::SPIRV_DIR; + let biased = GpuSolveConstraintsBiased::from_dir(&backend, dir).unwrap(); + let generic = GpuSolveConstraints::from_dir(&backend, dir).unwrap(); + let warm = GpuSolveConstraintsUnbiasedCached::from_dir(&backend, dir).unwrap(); + let final_iteration = GpuSolveConstraintsUnbiased::from_dir(&backend, dir).unwrap(); + let tail_biased = GpuSolveConstraintsTailBiased::from_dir(&backend, dir).unwrap(); + let tail_generic = GpuSolveConstraintsTail::from_dir(&backend, dir).unwrap(); + let tail_warm = GpuSolveConstraintsTailUnbiasedCached::from_dir(&backend, dir).unwrap(); + let tail_final = GpuSolveConstraintsTailUnbiased::from_dir(&backend, dir).unwrap(); + let storage = BufferUsages::STORAGE | BufferUsages::COPY_SRC; + let params = Tensor::scalar(&backend, RbdSimParams::tgs_soft(), BufferUsages::UNIFORM).unwrap(); + let colors: Vec<_> = (0..=64) + .map(|c| Tensor::scalar(&backend, c, BufferUsages::UNIFORM).unwrap()) + .collect(); + // Alternate disjoint matchings of a body ring. Consecutive colors share bodies, + // so missing inter-color visibility changes the result. Each occupied color has + // more than 64 constraints; empty colors and an incomplete final tile are also covered. + let mut input = Vec::new(); + let mut boundaries = vec![0u32]; + for color in 1..=64u32 { + if [TAIL_COLOR, TAIL_COLOR + 1, 63, 64].contains(&color) { + for index in 0..95u32 { + let a = (index * 2 + color % 2) % 192; + let b = (a + 1) % 192; + let mut c = TwoBodyConstraint { + len: 1 + index % 4, + dir_a: Vec3::Y, + tangent_a: Vec3::Z, + limit: 0.7, + im_a: Vec3::splat(0.4), + im_b: Vec3::splat(0.6), + solver_body_a: a, + solver_body_b: b, + vel_slot_a: a + if a % 7 == 0 { 192 } else { 0 }, + vel_slot_b: b + if b % 7 == 0 { 192 } else { 0 }, + ..Default::default() + }; + c.ii_a.diag = Vec3::splat(0.2); + c.ii_b.diag = Vec3::splat(0.3); + for k in 0..c.len as usize { + let p = &mut c.points[k]; + p.r_a = Vec3::new(k as f32 * 0.07, -0.2, 0.1); + p.r_b = Vec3::new(k as f32 * 0.07, 0.2, 0.1); + p.local_pt_a = p.r_a; + p.local_pt_b = p.r_b; + p.dist = -0.01; + p.normal_impulse = 0.04 * (k + 1) as f32; + p.tangent_impulse = Vec2::new(0.01, -0.02); + } + c.compute_effective_masses(); + input.push(c); + } + } + boundaries.push(input.len() as u32); + } + boundaries.push(input.len() as u32); + let poses = Tensor::vector(&backend, vec![Pose::IDENTITY; 192], storage).unwrap(); + let initial_vels: Vec<_> = (0..384) + .map(|i| Velocity { + linear: Vec3::new((i % 7) as f32 * 0.03, -0.1, 0.07), + angular: Vec3::new(0.1, 0.03, (i % 3) as f32 * 0.02), + ..Default::default() + }) + .collect(); + for nb in [1u32, 3] { + let buckets = Tensor::vector( + &backend, + boundaries + .iter() + .flat_map(|&end| vec![end; nb as usize]) + .collect::>(), + storage, + ) + .unwrap(); + let ids = Tensor::scalar( + &backend, + BatchIndices { + num_batches: nb, + solver_color_buckets_stride: 67, + ..Default::default() + }, + BufferUsages::UNIFORM, + ) + .unwrap(); + for variant in 0..4 { + let mut results = Vec::new(); + for fused in [false, true] { + // The constraints in tiles, zeroed so that their unused points compare equal. + let num_tiles = input.len().div_ceil(TILE_LEN); + let mut tiles = if model == FrictionModel::Simplified { + let compact: Vec<_> = input + .iter() + .map(|c| { + let mut t = TwistConstraint { + dir_a: c.dir_a, + tangent_a: c.tangent_a, + len: c.len, + limit: c.limit, + im_a: c.im_a, + im_b: c.im_b, + ii_a: c.ii_a, + ii_b: c.ii_b, + solver_body_a: c.solver_body_a, + solver_body_b: c.solver_body_b, + vel_slot_a: c.vel_slot_a, + vel_slot_b: c.vel_slot_b, + offset_b: c.points[0].r_b - c.points[0].r_a, + frame_a: glamx::Quat::from_rotation_x(0.13), + frame_b: glamx::Quat::from_rotation_y(-0.08), + ..Default::default() + }; + for k in 0..t.len as usize { + let p = c.points[k]; + t.points[k] = TwistContactPoint { + r_a: p.r_a, + dist: p.dist, + normal_impulse: p.normal_impulse, + normal_mass: 0.0, + normal_vel: p.normal_vel, + radius: 0.0, + }; + } + t.friction.tangent_impulse = Vec2::new(0.01, -0.02); + t.friction.twist_impulse = if t.len > 1 { 0.01 } else { 0.0 }; + t.compute_effective_masses(); + t + }) + .collect(); + let mut tiles = vec![TwistTile::default(); num_tiles]; + for (index, c) in compact.iter().enumerate() { + let (tile, lane) = tile_lane(index); + tiles[tile].write_constraint(lane, c); + } + ContactTiles::Simplified(Tensor::vector(&backend, tiles, storage).unwrap()) + } else { + let mut tiles = vec![CoulombTile::default(); num_tiles]; + for (index, c) in input.iter().enumerate() { + let (tile, lane) = tile_lane(index); + tiles[tile].write_constraint(lane, c); + } + ContactTiles::Coulomb(Tensor::vector(&backend, tiles, storage).unwrap()) + }; + let mut vels = Tensor::vector(&backend, initial_vels.clone(), storage).unwrap(); + let mut encoder = backend.begin_encoding(); + { + let mut pass = encoder.begin_pass("tail-colors-regression", None); + let first = if fused { 64 } else { TAIL_COLOR }; + for color in first..=64 { + macro_rules! call { + ($kernel:expr) => { + $kernel + .call( + &mut pass, + if fused { 64u32 } else { 128u32 }, + &mut tiles, + &mut vels, + &buckets, + &poses, + &colors[color as usize], + &ids, + &colors[if variant == 1 { + 2 + } else if variant == 0 { + 1 + } else { + 0 + }], + ¶ms, + ) + .unwrap() + }; + } + match (fused, variant) { + (false, 0) => call!(biased), + (true, 0) => call!(tail_biased), + (false, 1) => call!(generic), + (true, 1) => call!(tail_generic), + (false, 2) => call!(warm), + (true, 2) => call!(tail_warm), + (false, _) => call!(final_iteration), + (true, _) => call!(tail_final), + } + } + } + backend.submit(encoder).unwrap(); + backend.synchronize().unwrap(); + let mut bytes: Vec = Vec::new(); + macro_rules! read { + ($tensor:expr, $ty:ty) => { + bytes.extend_from_slice(bytemuck::cast_slice( + &backend + .slow_read_vec::<$ty>($tensor.buffer()) + .await + .unwrap(), + )) + }; + } + match &tiles { + ContactTiles::Coulomb(x) => { + read!(x, CoulombTile); + } + ContactTiles::Simplified(x) => { + read!(x, TwistTile); + } + } + read!(vels, Velocity); + results.push(bytes); + } + assert!( + results[0] == results[1], + "fused tail differs for variant {variant}, batches {nb}" + ); + } + } +} diff --git a/src_rbd/pipeline/test_twist_friction.rs b/src_rbd/pipeline/test_twist_friction.rs new file mode 100644 index 00000000..97cecc9c --- /dev/null +++ b/src_rbd/pipeline/test_twist_friction.rs @@ -0,0 +1,158 @@ +use super::{RbdCapacities, RbdPipeline, RbdState}; +use crate::dynamics::FrictionModel; +use crate::math::Pose; +use crate::rapier::prelude::*; +use crate::shaders::dynamics::{RbdSimParams, Velocity}; +use khal::backend::{Backend, GpuBackend, WebGpu}; + +#[futures_test::test] +#[serial_test::serial] +#[ignore = "requires a WebGPU adapter"] +async fn twist_friction_stops_loaded_spin_on_small_and_large_paths() { + for (count, bindings) in [(1, 8), (4096, 8), (4096, 10), (4096, 16)] { + let backend = GpuBackend::WebGpu( + WebGpu::new( + Default::default(), + khal::re_exports::wgpu::Limits { + max_storage_buffers_per_shader_stage: bindings, + ..Default::default() + }, + ) + .await + .unwrap(), + ); + let mut bodies = RigidBodySet::new(); + let mut colliders = ColliderSet::new(); + let ground = + bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, -0.5, 0.0))); + colliders.insert_with_parent( + ColliderBuilder::cuboid(100.0, 0.5, 100.0).friction(0.8), + ground, + &mut bodies, + ); + for i in 0..count { + let body = bodies.insert( + RigidBodyBuilder::dynamic() + .can_sleep(false) + .translation(Vec3::new((i % 64) as f32, 0.21, (i / 64) as f32)) + .angvel(Vec3::Y * 6.0), + ); + colliders.insert_with_parent( + ColliderBuilder::cuboid(0.2, 0.2, 0.2) + .friction(if i % 2 == 0 { 0.8 } else { 0.0 }) + .friction_combine_rule(CoefficientCombineRule::Min), + body, + &mut bodies, + ); + } + let mut params = RbdSimParams::tgs_soft(); + params.friction_model = FrictionModel::Simplified; + params.warmstart_coefficient = if bindings == 16 { 1.0 } else { 0.75 }; + let mut state = RbdState::from_rapier( + &backend, + &[( + &bodies, + &colliders, + &ImpulseJointSet::new(), + &MultibodyJointSet::new(), + ¶ms, + )], + RbdCapacities { + collisions_capacity: (count * 2).max(256), + ..Default::default() + }, + ); + let pipeline = RbdPipeline::new(&backend).unwrap(); + for step in 0..180 { + // Both directions exercise stale-contact invalidation and graph changes. + if step == 60 { + state.set_friction_model(&backend, FrictionModel::Coulomb); + } + if step == 80 { + state.set_friction_model(&backend, FrictionModel::Simplified); + } + pipeline.auto_resize_buffers(&backend, &mut state).unwrap(); + pipeline.step(&backend, &mut state, None).unwrap(); + if [20, 70, 90].contains(&step) { + let plan: Vec = backend + .slow_read_vec(state.contact_plan.buffer()) + .await + .unwrap(); + // step swaps buffers: old_constraints now holds the just-solved frame. + let old = &state.old_constraints; + if state.friction_model() == FrictionModel::Simplified { + assert!(old.tiles.coulomb().is_none()); + let links: Vec = + backend.slow_read_vec(old.links.buffer()).await.unwrap(); + let constraint_indices: Vec = backend + .slow_read_vec(old.constraint_indices.buffer()) + .await + .unwrap(); + let tiles: Vec = backend + .slow_read_vec(old.tiles.simplified().unwrap().buffer()) + .await + .unwrap(); + let active: Vec<_> = (0..plan[0].bound as usize) + .filter(|&i| links[i].len > 0) + .map(|i| { + let (tile, lane) = crate::shaders::dynamics::contact_tiles::tile_lane( + constraint_indices[i] as usize, + ); + tiles[tile].constraint(lane) + }) + .collect(); + assert!(!active.is_empty()); + for c in &active { + let points = &c.points[..c.len as usize]; + let twist_limit = c.limit + * points + .iter() + .map(|p| p.normal_impulse * p.radius) + .sum::(); + let tangent_limit = + c.limit * points.iter().map(|p| p.normal_impulse).sum::(); + assert!(c.friction.twist_impulse.abs() <= twist_limit + 1e-4); + assert!(c.friction.tangent_impulse.length() <= tangent_limit + 1e-4); + } + } else { + assert!(old.tiles.simplified().is_none()); + let snapshots = old + .read_impulses(&backend, plan[0].bound as usize) + .await + .unwrap(); + assert!(snapshots.iter().any(|c| c.len > 0)); + } + } + } + backend.synchronize().unwrap(); + assert_eq!(state.friction_model(), FrictionModel::Simplified); + let velocities: Vec = backend.slow_read_vec(state.vels().buffer()).await.unwrap(); + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + for i in 0..count as usize { + let v = velocities[i + 1]; + let p = poses[i + 1]; + assert!(p.translation.is_finite() && v.angular.is_finite()); + assert!( + (p.translation.y - 0.2).abs() < 0.015, + "bindings={bindings} box={i} height={}", + p.translation.y + ); + if i % 2 == 0 { + assert!( + v.angular.length() < 0.1, + "bindings={bindings} box={i} angular={:?}", + v.angular + ); + } else { + assert!( + v.angular.y > 5.0, + "frictionless box {i} lost spin: {:?}", + v.angular + ); + } + } + } +} diff --git a/src_rbd/pipeline/test_web_preparation.rs b/src_rbd/pipeline/test_web_preparation.rs new file mode 100644 index 00000000..4d39226c --- /dev/null +++ b/src_rbd/pipeline/test_web_preparation.rs @@ -0,0 +1,186 @@ +//! The same checks run natively and in the browser regression harness. +use crate::shaders::Vector; +use crate::shaders::bounding_volumes::Aabb; +use crate::shaders::broad_phase::{ + GpuLbvhRefitChunks, GpuLbvhRefitFrontier, GpuLbvhRefitPlan, LbvhNode, +}; +use crate::shaders::utils::BatchIndices; +use khal::BufferUsages; +use khal::backend::{Backend, Encoder, GpuBackend}; +use vortx::tensor::Tensor; + +fn tree_fixture(n: usize, capacity: usize, batch: usize) -> Vec { + fn build( + nodes: &mut [LbvhNode], + next: &mut usize, + n: usize, + first: usize, + end: usize, + parent: u32, + batch: usize, + depth: usize, + ) -> usize { + if end - first == 1 { + let id = n - 1 + first; + let center = Vector::X * ((first * 13 % 97) as f32 + batch as f32 * 128.0) + + Vector::Y * (first % 31) as f32; + let half = Vector::splat((first % 5 + 1) as f32 * 0.125); + nodes[id] = LbvhNode { + aabb: Aabb::new(center - half, center + half), + first_leaf: first as u32, + last_leaf: first as u32, + parent, + max_extent: half.x * 2.0, + ..Default::default() + }; + return id; + } + let id = *next; + *next += 1; + // Distinct frontier sizes per batch, including a non-balanced tree. + let split = if batch % 2 == 1 && depth < 12 { + 1 + } else { + (end - first) / 2 + }; + let a = build( + nodes, + next, + n, + first, + first + split, + id as u32, + batch, + depth + 1, + ); + let b = build( + nodes, + next, + n, + first + split, + end, + id as u32, + batch, + depth + 1, + ); + nodes[id] = LbvhNode { + aabb: nodes[a].aabb.merged(&nodes[b].aabb), + left: a as u32, + right: b as u32, + parent, + first_leaf: first as u32, + last_leaf: (end - 1) as u32, + max_extent: nodes[a].max_extent.max(nodes[b].max_extent), + ..Default::default() + }; + id + } + let mut nodes = vec![LbvhNode::default(); 2 * capacity]; + if n > 0 { + build(&mut nodes, &mut 0, n, 0, n, 0, batch, 0); + } + nodes +} + +/// Validate every node against CPU subtree bounds, including empty and unequal frontiers. +pub async fn check_refit(backend: &GpuBackend) -> usize { + let chunks = GpuLbvhRefitChunks::from_dir(backend, &crate::SPIRV_DIR).unwrap(); + let plan = GpuLbvhRefitPlan::from_dir(backend, &crate::SPIRV_DIR).unwrap(); + let frontier = GpuLbvhRefitFrontier::from_dir(backend, &crate::SPIRV_DIR).unwrap(); + let usage = BufferUsages::STORAGE | BufferUsages::COPY_SRC; + let mut cases = 0; + for n in [0usize, 1, 2, 255, 256, 257, 513, 1025, 4097] { + for nb in [1usize, 3] { + let capacity = n + 7; + let expected: Vec<_> = (0..nb).flat_map(|b| tree_fixture(n, capacity, b)).collect(); + let mut initial = expected.clone(); + let mut expected_counts = vec![0u32; nb]; + for (batch, expected_count) in expected_counts.iter_mut().enumerate() { + let start = batch * capacity * 2; + for id in 1..n.saturating_mul(2).saturating_sub(1) { + let node = &expected[start + id]; + let parent = &expected[start + node.parent as usize]; + if node.first_leaf / 256 == node.last_leaf / 256 + && parent.first_leaf / 256 != parent.last_leaf / 256 + { + *expected_count += 1; + } + } + for node in &mut initial[start..start + n.saturating_sub(1)] { + node.aabb = Aabb::default(); + node.max_extent = 0.0; + } + } + let mut tree = Tensor::vector(backend, initial, usage).unwrap(); + let mut entries = + Tensor::vector(backend, vec![u32::MAX; capacity * nb], usage).unwrap(); + let mut counts = Tensor::vector(backend, vec![0u32; nb], usage).unwrap(); + let mut rounds = + Tensor::scalar(backend, 777u32, usage | BufferUsages::UNIFORM).unwrap(); + let ids = Tensor::scalar( + backend, + BatchIndices { + num_batches: nb as u32, + colliders_len: n as u32, + colliders_batch_capacity: capacity as u32, + ..Default::default() + }, + BufferUsages::UNIFORM, + ) + .unwrap(); + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("web-refit-regression", None); + // An extra chunk exercises inactive workgroups and partial final chunks. + chunks + .call( + &mut pass, + [n as u32 + 256, nb as u32, 1], + &mut tree, + &mut entries, + &mut counts, + &ids, + ) + .unwrap(); + plan.call(&mut pass, 1u32, &counts, &mut rounds, &ids) + .unwrap(); + frontier + .call( + &mut pass, + [1u32, nb as u32, 1], + &mut tree, + &entries, + &counts, + &ids, + &rounds, + ) + .unwrap(); + drop(pass); + backend.submit(encoder).unwrap(); + let actual = backend.slow_read_vec(tree.buffer()).await.unwrap(); + for (id, (a, b)) in actual.iter().zip(&expected).enumerate() { + assert_eq!(a.aabb.mins, b.aabb.mins, "n={n}, batches={nb}, node={id}"); + assert_eq!(a.aabb.maxs, b.aabb.maxs, "n={n}, batches={nb}, node={id}"); + assert_eq!(a.max_extent, b.max_extent, "n={n}, batches={nb}, node={id}"); + } + assert_eq!( + backend.slow_read_vec(counts.buffer()).await.unwrap(), + expected_counts + ); + assert_eq!( + backend.slow_read_vec(rounds.buffer()).await.unwrap(), + [expected_counts.iter().copied().max().unwrap().div_ceil(256)] + ); + cases += 1; + } + } + cases +} + +#[cfg(all(test, not(target_arch = "wasm32")))] +#[futures_test::test] +#[serial_test::serial] +#[ignore = "requires a WebGPU adapter"] +async fn web_preparation_refit_preserves_every_subtree() { + let backend = GpuBackend::WebGpu(khal::backend::WebGpu::default().await.unwrap()); + assert_eq!(check_refit(&backend).await, 18); +} diff --git a/src_rbd_shaders/bounding_volumes/aabb.rs b/src_rbd_shaders/bounding_volumes/aabb.rs index aa7ec709..f3538fb0 100644 --- a/src_rbd_shaders/bounding_volumes/aabb.rs +++ b/src_rbd_shaders/bounding_volumes/aabb.rs @@ -159,26 +159,26 @@ impl Aabb { #[cfg(feature = "dim2")] #[inline] pub fn intersects(&self, other: &Aabb) -> bool { - // TODO PERF: is could we use some sort of `self.mins <= other.maxs` - // directly instead of detailing each component. - self.mins.x <= other.maxs.x - && self.maxs.x >= other.mins.x - && self.mins.y <= other.maxs.y - && self.maxs.y >= other.mins.y + // Non-short-circuit `&`: `&&` compiles to nested branches, each loading its operands + // only after the previous test, which serializes the loads of a storage AABB. + (self.mins.x <= other.maxs.x) + & (self.maxs.x >= other.mins.x) + & (self.mins.y <= other.maxs.y) + & (self.maxs.y >= other.mins.y) } /// Tests if this AABB intersects another AABB. #[cfg(feature = "dim3")] #[inline] pub fn intersects(&self, other: &Aabb) -> bool { - // TODO PERF: is could we use some sort of `self.mins <= other.maxs` - // directly instead of detailing each component. - self.mins.x <= other.maxs.x - && self.maxs.x >= other.mins.x - && self.mins.y <= other.maxs.y - && self.maxs.y >= other.mins.y - && self.mins.z <= other.maxs.z - && self.maxs.z >= other.mins.z + // Non-short-circuit `&`: `&&` compiles to nested branches, each loading its operands + // only after the previous test, which serializes the loads of a storage AABB. + (self.mins.x <= other.maxs.x) + & (self.maxs.x >= other.mins.x) + & (self.mins.y <= other.maxs.y) + & (self.maxs.y >= other.mins.y) + & (self.mins.z <= other.maxs.z) + & (self.maxs.z >= other.mins.z) } /// Tests if this AABB contains a point. diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 8633b11c..24fb86c3 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -41,14 +41,34 @@ pub struct LbvhNode { pub right: u32, /// Parent node index. pub parent: u32, - /// During refit: bottom-up arrival counter; each thread atomically - /// increments it and only the second one continues upward (both children - /// ready). After refit, `refit_leaves` sets it to the maximum sorted leaf - /// index in this node’s subtree, which the pair traversal uses to prune + /// Bottom-up refit arrival counter (internal nodes): each child's thread atomically + /// increments it and only the second one continues upward (both children ready). + pub refit_count: u32, + /// The first sorted leaf index in this node's subtree. + pub first_leaf: u32, + /// The last sorted leaf index in this node's subtree. The pair traversal uses it to prune /// subtrees that can only produce duplicate pairs. - pub refit_count_or_max_subtree_index: u32, + pub last_leaf: u32, + /// The node following this one's subtree in depth-first order ([`NO_ESCAPE`] for the + /// last one), for stackless traversals. + pub escape: u32, + /// The largest AABB extent of the leaves in this node's subtree. + pub max_extent: f32, } +/// The largest extent of an AABB. +#[inline(always)] +fn max_extent(aabb: &Aabb) -> f32 { + let e = aabb.maxs - aabb.mins; + #[cfg(feature = "dim3")] + return e.x.max(e.y).max(e.z); + #[cfg(feature = "dim2")] + return e.x.max(e.y); +} + +/// The escape index of the nodes ending a depth-first traversal. +pub const NO_ESCAPE: u32 = u32::MAX; + /// Resets the (single, global) collision pairs counter. #[spirv_bindgen] #[spirv(compute(threads(64)))] @@ -84,24 +104,33 @@ pub fn gpu_flat_list_dispatch( } } -/// Runs a reduction to compute the AABB of the collider positions. -/// Needs to be called with a single workgroup. +/// Most workgroups per batch of [`gpu_lbvh_compute_domain`]. +pub const DOMAIN_WORKGROUPS: u32 = 64; + +/// Min/max reduction of the collider positions, first pass: each of the batch's workgroups +/// reduces a share of them into `partials`, merged by [`gpu_lbvh_domain_merge`]. #[spirv_bindgen] #[spirv(compute(threads(128)))] pub fn gpu_lbvh_compute_domain( - #[spirv(global_invocation_id)] global_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] poses: &[Pose], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] domain_aabb: &mut [Aabb], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] partials: &mut [Aabb], #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, #[spirv(workgroup)] workspace_mins: &mut [Vector; 128], #[spirv(workgroup)] workspace_maxs: &mut [Vector; 128], ) { - let batch_id = global_id.y; - let thread_id = global_id.x; + let batch_id = workgroup_id.y; + let thread_id = local_id.x; + let group = workgroup_id.x; + let num_groups = num_workgroups.x.min(DOMAIN_WORKGROUPS); + let first = group * REDUCTION_WORKGROUP_SIZE + thread_id; + let stride = num_groups * REDUCTION_WORKGROUP_SIZE; *workspace_mins.at_mut(thread_id as usize) = Vector::splat(MAX_FLT); *workspace_maxs.at_mut(thread_id as usize) = Vector::splat(-MAX_FLT); - for i in StepRng::new(thread_id..batch_ids.colliders_len, REDUCTION_WORKGROUP_SIZE) { + for i in StepRng::new(first..batch_ids.colliders_len, stride) { let val_i = poses.at(batch_ids.body_global(batch_id, i)).translation; *workspace_mins.at_mut(thread_id as usize) = workspace_mins.at(thread_id as usize).min(val_i); @@ -131,6 +160,55 @@ pub fn gpu_lbvh_compute_domain( step_reduce!(2); step_reduce!(1); + let base = (batch_id * DOMAIN_WORKGROUPS) as usize; + if thread_id == 0 && group < num_groups { + partials.at_mut(base + group as usize).mins = *workspace_mins.at(0); + partials.at_mut(base + group as usize).maxs = *workspace_maxs.at(0); + } + // The first workgroup empties the slots no workgroup fills. + if group == 0 && thread_id >= num_groups && thread_id < DOMAIN_WORKGROUPS { + partials.at_mut(base + thread_id as usize).mins = Vector::splat(MAX_FLT); + partials.at_mut(base + thread_id as usize).maxs = Vector::splat(-MAX_FLT); + } +} + +/// Min/max reduction of the collider positions, second pass: one workgroup per batch merges +/// the partial domains of [`gpu_lbvh_compute_domain`]. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_lbvh_domain_merge( + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] partials: &[Aabb], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] domain_aabb: &mut [Aabb], + #[spirv(workgroup)] workspace_mins: &mut [Vector; 64], + #[spirv(workgroup)] workspace_maxs: &mut [Vector; 64], +) { + let batch_id = workgroup_id.y; + let thread_id = local_id.x; + let partial = partials.at((batch_id * DOMAIN_WORKGROUPS + thread_id) as usize); + *workspace_mins.at_mut(thread_id as usize) = partial.mins; + *workspace_maxs.at_mut(thread_id as usize) = partial.maxs; + workgroup_memory_barrier_with_group_sync(); + + macro_rules! step_reduce( + ($stride: expr) => { + if thread_id < $stride { + *workspace_mins.at_mut(thread_id as usize) = workspace_mins.at(thread_id as usize) + .min(*workspace_mins.at((thread_id + $stride) as usize)); + *workspace_maxs.at_mut(thread_id as usize) = workspace_maxs.at(thread_id as usize) + .max(*workspace_maxs.at((thread_id + $stride) as usize)); + } + workgroup_memory_barrier_with_group_sync(); + } + ); + step_reduce!(32); + step_reduce!(16); + step_reduce!(8); + step_reduce!(4); + step_reduce!(2); + step_reduce!(1); + if thread_id == 0 { domain_aabb.at_mut(batch_id as usize).mins = *workspace_mins.at(0); domain_aabb.at_mut(batch_id as usize).maxs = *workspace_maxs.at(0); @@ -263,13 +341,48 @@ pub fn gpu_lbvh_build( tree.at_mut(node_id as usize).left = left as u32; tree.at_mut(node_id as usize).right = right as u32; - tree.at_mut(node_id as usize) - .refit_count_or_max_subtree_index = 0; // Might as well reset the refit count here. + tree.at_mut(node_id as usize).refit_count = 0; + tree.at_mut(node_id as usize).first_leaf = ii.min(j) as u32; + tree.at_mut(node_id as usize).last_leaf = ii.max(j) as u32; tree.at_mut(left as usize).parent = node_id; tree.at_mut(right as usize).parent = node_id; } } +/// Computes the escape index of every node: the right sibling of its first ancestor (itself +/// included) that is a left child. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_lbvh_escapes( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] tree: &mut [LbvhNode], + #[spirv(uniform, descriptor_set = 0, binding = 1)] batch_ids: &BatchIndices, +) { + let num_threads = num_workgroups.x * WORKGROUP_SIZE; + let batch_id = invocation_id.y; + let colliders_start = scratch_start(batch_ids, batch_id); + let num_nodes = (2 * batch_ids.colliders_len).max(2) - 1; + let mut tree = SliceMut(tree, root_id(colliders_start) as usize); + + for i in StepRng::new(invocation_id.x..num_nodes, num_threads) { + let mut escape = NO_ESCAPE; + let mut node = i; + for _ in 0..REFIT_MAX_DEPTH { + if node == 0 { + break; + } + let parent = tree.at(node as usize).parent; + if tree.at(parent as usize).left == node { + escape = tree.at(parent as usize).right; + break; + } + node = parent; + } + tree.at_mut(i as usize).escape = escape; + } +} + /// Computes leaf AABBs from shapes. #[spirv_bindgen] #[spirv(compute(threads(64)))] @@ -281,11 +394,9 @@ pub fn gpu_lbvh_refit_leaves( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] sorted_colliders: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] tree: &mut [LbvhNode], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] refit_frontier_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] vertices: &[PaddedVector], ) { - // TODO PERF: we could use shared memory atomics between threads belonging to the same - // workgroup. - // Bottom-up refit. Leaf index starts at `num_colliders`. let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; let colliders_start = scratch_start(batch_ids, batch_id); @@ -303,191 +414,265 @@ pub fn gpu_lbvh_refit_leaves( let leaf_pose = poses[leaf_collider as usize]; let leaf_shape = &shapes[leaf_collider as usize]; - tree.at_mut(curr_leaf_id as usize).aabb = leaf_shape.compute_aabb(leaf_pose, vertices); + let aabb = leaf_shape.compute_aabb(leaf_pose, vertices); + tree.at_mut(curr_leaf_id as usize).aabb = aabb; + tree.at_mut(curr_leaf_id as usize).max_extent = max_extent(&aabb); tree.at_mut(curr_leaf_id as usize).left = leaf_collider; - // For leaves, we can set their index here. They don’t use the `refit_count` - // mechanism which is for internal nodes only. - tree.at_mut(curr_leaf_id as usize) - .refit_count_or_max_subtree_index = i; + tree.at_mut(curr_leaf_id as usize).first_leaf = i; + tree.at_mut(curr_leaf_id as usize).last_leaf = i; + } + + // The bottom-up refit appends to this batch's frontier list. + if invocation_id.x == 0 { + refit_frontier_len.write(batch_id as usize, 0); } } -/// Bottom-up AABB propagation using atomic synchronization. -/// This version uses uniform control flow for web compatibility. +/// Number of sorted leaves refit by one workgroup of [`gpu_lbvh_refit_chunks`]. +pub const REFIT_CHUNK: u32 = 256; +/// Bound on the depth of the tree walked by the bottom-up refits. +const REFIT_MAX_DEPTH: u32 = 64; + +/// Workgroup barrier making storage writes of the workgroup visible to its other lanes. +#[inline(always)] +fn refit_barrier() { + control_barrier::< + { khal_std::memory::Scope::Workgroup as u32 }, + { khal_std::memory::Scope::QueueFamily as u32 }, + { + khal_std::memory::Semantics::UNIFORM_MEMORY.bits() + | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() + }, + >(); +} + +/// Whether any lane of the workgroup is still `active`. Must be reached by every lane. +#[cfg(not(feature = "web-compat"))] +#[inline(always)] +fn any_lane_active(lane: u32, active: bool, flag: &mut [u32; 1]) -> bool { + if lane == 0 { + flag.write(0, 0); + } + workgroup_memory_barrier_with_group_sync(); + if active { + // Every active lane stores the same value. + flag.write(0, 1); + } + workgroup_memory_barrier_with_group_sync(); + let any = flag.read(0) != 0; + // Lane 0 resets the flag at the next call: keep it from racing with these reads. + workgroup_memory_barrier_with_group_sync(); + any +} + +/// Merges the AABBs of `node`'s children into its own. +#[inline(always)] +fn refit_node(tree: &mut SliceMut, node: u32) { + let left = tree.at(node as usize).left; + let right = tree.at(node as usize).right; + let aabb = tree + .at(left as usize) + .aabb + .merged(&tree.at(right as usize).aabb); + let extent = tree + .at(left as usize) + .max_extent + .max(tree.at(right as usize).max_extent); + tree.at_mut(node as usize).aabb = aabb; + tree.at_mut(node as usize).max_extent = extent; +} + +/// First phase of the bottom-up refit: one workgroup per chunk of [`REFIT_CHUNK`] sorted +/// leaves refits every node whose leaves all lie in its chunk. +/// +/// Both children of such a node are refit by lanes of the same workgroup, so the usual +/// arrival counter only needs workgroup barriers. A lane stops at a node whose parent's leaves +/// cross the chunk boundary, and appends it to the batch's frontier list, which +/// [`gpu_lbvh_refit_frontier`] finishes. #[spirv_bindgen] #[spirv(compute(threads(256)))] -pub fn gpu_lbvh_refit_internal( +pub fn gpu_lbvh_refit_chunks( #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(workgroup_id)] workgroup_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] tree: &mut [LbvhNode], - #[spirv(uniform, descriptor_set = 0, binding = 1)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] refit_frontier: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] refit_frontier_len: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(workgroup)] any_active_flag: &mut [u32; 1], ) { - // TODO PERF: we could use shared memory atomics between threads belonging to the same - // workgroup. - // Bottom-up refit. Leaf index starts at `num_colliders`. - let num_threads = 256u32; + let lane = local_id.x; + #[cfg(feature = "web-compat")] + let _ = any_active_flag; + let chunk = workgroup_id.x; let batch_id = workgroup_id.y; let colliders_start = scratch_start(batch_ids, batch_id); - let num_bodies = batch_ids.colliders_len; - let first_leaf_id = num_bodies - 1; - + let num_leaves = batch_ids.colliders_len; + if num_leaves < 2 { + return; + } + // Uniform across the workgroup, but derived from `workgroup_id`, which the WGSL uniformity + // analysis doesn't trust: on the web, out-of-range chunks run the loop inactive instead. + #[cfg(not(feature = "web-compat"))] + if chunk * REFIT_CHUNK >= num_leaves { + return; + } + let first_leaf_id = num_leaves - 1; let mut tree = SliceMut(tree, root_id(colliders_start) as usize); - let num_iterations = num_bodies.div_ceil(num_threads); - // NOTE: using unchecked indexing (via MaybeIndexUnchecked) because otherwise the bounds - // checking inserted by rustgpu breaks the shader when targeting some NVidia graphics - // (works fine on AMD). - for iter in 0..num_iterations { - let i = local_id.x + iter * num_threads; - let mut thread_is_active = i < num_bodies; + let leaf = chunk * REFIT_CHUNK + lane; + let mut active = leaf < num_leaves; + let mut child = first_leaf_id + leaf; + let mut node = 0u32; + if active { + node = tree.at(child as usize).parent; + } - let mut curr_id = 0u32; - if thread_is_active { - let curr_leaf_id = first_leaf_id + i; - curr_id = tree.at(curr_leaf_id as usize).parent; + for _ in 0..REFIT_MAX_DEPTH { + // On the web, every level runs: the early exit reads workgroup memory, which the WGSL + // uniformity analysis can't prove uniform (no `workgroupUniformLoad` in rust-gpu). + #[cfg(not(feature = "web-compat"))] + if !any_lane_active(lane, active, any_active_flag) { + break; } - // Process the tree level by level with uniform barriers. - // Maximum tree depth is log2(num_colliders), but we use 32 as a safe upper bound. - for _level in 0..32u32 { - if thread_is_active { - let refit_count = atomic_add_u32( - &mut tree - .at_mut(curr_id as usize) - .refit_count_or_max_subtree_index, - 1, - ); - - if refit_count == 0 { - // If `refit_count` was 0 then the other thread hasn't reached this node - // yet and the sibling aabb might not be available yet. - // Stop the propagation to the parents here, the other thread will do it. - thread_is_active = false; + if active { + let first = tree.at(node as usize).first_leaf; + let last = tree.at(node as usize).last_leaf; + if first / REFIT_CHUNK != chunk || last / REFIT_CHUNK != chunk { + // `child` is complete but its parent spans several chunks. + let k = atomic_add_u32(refit_frontier_len.at_mut(batch_id as usize), 1); + refit_frontier.write((colliders_start + k) as usize, child); + active = false; + } else if atomic_add_u32(&mut tree.at_mut(node as usize).refit_count, 1) == 0 { + // The sibling isn't ready: its lane continues when it arrives. + active = false; + } else { + refit_node(&mut tree, node); + if node == 0 { + // The root: the whole tree fits in this chunk. + active = false; } else { - // If `refit_count` was 1 then the other thread has already reached this node - // and we know the sibblings aabb is available. So we can continue the propagation. - - // TODO PERF: instead of re-reading both aabbs, we could keep the aabb from the - // previous loop so we don't have to re-fetch one of the two aabbs. - let left_idx = tree.at(curr_id as usize).left; - let right_idx = tree.at(curr_id as usize).right; - let left = tree.at(left_idx as usize).aabb; - let right = tree.at(right_idx as usize).aabb; - tree.at_mut(curr_id as usize).aabb = left.merged(&right); - - // Set `refit_count_or_max_subtree_index` to the max subtree leaf index. - let max_l = tree.at(left_idx as usize).refit_count_or_max_subtree_index; - let max_r = tree.at(right_idx as usize).refit_count_or_max_subtree_index; - tree.at_mut(curr_id as usize) - .refit_count_or_max_subtree_index = max_l.max(max_r); - - if curr_id == 0 { - // We reached the root, can't go higher. - thread_is_active = false; - } else { - curr_id = tree.at(curr_id as usize).parent; - } + child = node; + node = tree.at(node as usize).parent; } } - - // workgroup_memory_barrier_with_group_sync(); - // Barrier ensures all AABB writes (to device/storage buffer memory) are complete - // before the next iteration's atomics. Uses QueueFamily scope (device-equivalent - // under the Vulkan memory model) with UNIFORM_MEMORY to cover storage buffers. - control_barrier::< - { khal_std::memory::Scope::Workgroup as u32 }, - { khal_std::memory::Scope::QueueFamily as u32 }, - { - khal_std::memory::Semantics::UNIFORM_MEMORY.bits() - | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() - }, - >(); } + + refit_barrier(); } } -/// Full refit: computes leaf AABBs and propagates to ancestors. +/// Reduce the actual frontier sizes to a uniform loop bound for the web refit. +/// A storage load indexed by `workgroup_id` does not pass WGSL uniformity analysis; +/// a separate uniform buffer does, without iterating over every possible leaf. #[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_lbvh_refit( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(num_workgroups)] num_workgroups: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] poses: &[Pose], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] shapes: &[Shape], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] sorted_colliders: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] tree: &mut [LbvhNode], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, - #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] vertices: &[PaddedVector], +#[spirv(compute(threads(1)))] +pub fn gpu_lbvh_refit_plan( + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] frontier_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] rounds: &mut u32, + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, ) { - // TODO PERF: we could use shared memory atomics between threads belonging to the same - // workgroup. - // Bottom-up refit. Leaf index starts at `num_colliders`. - let batch_id = invocation_id.y; - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let colliders_start = scratch_start(batch_ids, batch_id); - let num_bodies = batch_ids.colliders_len; - let first_leaf_id = num_bodies - 1; + let mut largest = 0; + for batch in 0..batch_ids.num_batches { + largest = largest.max( + frontier_len + .read(batch as usize) + .min(batch_ids.colliders_len), + ); + } + *rounds = largest.div_ceil(REFIT_CHUNK); +} - let poses = batch_ids.ib(batch_id, poses); - let shapes = batch_ids.ib(batch_id, shapes); - let sorted_colliders = Slice(sorted_colliders, colliders_start as usize); +/// Second phase of the bottom-up refit: one workgroup per batch refits the nodes spanning +/// several chunks, starting from the frontier left by [`gpu_lbvh_refit_chunks`]. +#[spirv_bindgen] +#[spirv(compute(threads(256)))] +pub fn gpu_lbvh_refit_frontier( + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] tree: &mut [LbvhNode], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] refit_frontier: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] refit_frontier_len: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 4)] refit_rounds: &u32, + #[spirv(workgroup)] any_active_flag: &mut [u32; 1], +) { + let lane = local_id.x; + #[cfg(feature = "web-compat")] + let _ = any_active_flag; + let batch_id = workgroup_id.y; + let colliders_start = scratch_start(batch_ids, batch_id); + let num_leaves = batch_ids.colliders_len; let mut tree = SliceMut(tree, root_id(colliders_start) as usize); + // Uniform across the workgroup: every lane reads the same count. + let frontier_len = refit_frontier_len.read(batch_id as usize).min(num_leaves); + // The web bound is the largest actual frontier over the batches, uploaded by + // the preceding GPU planning pass. Every lane reaches the same barriers. + #[cfg(not(feature = "web-compat"))] + let num_rounds = frontier_len.div_ceil(REFIT_CHUNK); + #[cfg(feature = "web-compat")] + let num_rounds = *refit_rounds; + #[cfg(not(feature = "web-compat"))] + let _ = refit_rounds; + + for round in 0..num_rounds { + let k = round * REFIT_CHUNK + lane; + let mut active = k < frontier_len; + let mut node = 0u32; + if active { + let child = refit_frontier.read((colliders_start + k) as usize); + node = tree.at(child as usize).parent; + } - for i in StepRng::new(invocation_id.x..num_bodies, num_threads) { - let curr_leaf_id = first_leaf_id + i; - let leaf_collider = sorted_colliders[i as usize]; - let leaf_pose = poses[leaf_collider as usize]; - let leaf_shape = &shapes[leaf_collider as usize]; - - tree.at_mut(curr_leaf_id as usize).aabb = leaf_shape.compute_aabb(leaf_pose, vertices); - tree.at_mut(curr_leaf_id as usize).left = leaf_collider; - - // Propagate to ancestors. - let mut curr_id = tree.at(curr_leaf_id as usize).parent; - - // NOTE: bounded `for` (tree depth <= 32 in practice) instead of `loop` - // to avoid the MacOS miscompilation bug. - for _ in 0..32u32 { - let refit_count = atomic_add_u32( - &mut tree - .at_mut(curr_id as usize) - .refit_count_or_max_subtree_index, - 1, - ); - - if refit_count == 0 { - // If `refit_count` was 0 then the other thread hasn't reached this node - // yet and the sibling aabb might not be available yet. - // Stop the propagation to the parents here, the other thread will do it. + for _ in 0..REFIT_MAX_DEPTH { + // See `gpu_lbvh_refit_chunks`: every level runs on the web. + #[cfg(not(feature = "web-compat"))] + if !any_lane_active(lane, active, any_active_flag) { break; } - // If `refit_count` was 1 then the other thread has already reached this node, - // and we know the siblings aabb is available. So we can continue the propagation. - - // TODO PERF: instead of re-reading both aabbs, we could keep the aabb from the - // previous loop so we don't have to re-fetch one of the two aabbs. - let left_idx = tree.at(curr_id as usize).left; - let right_idx = tree.at(curr_id as usize).right; - let left = tree.at(left_idx as usize).aabb; - let right = tree.at(right_idx as usize).aabb; - tree.at_mut(curr_id as usize).aabb = left.merged(&right); - - if curr_id == 0 { - // We reached the root, can't go higher. - break; + if active { + if atomic_add_u32(&mut tree.at_mut(node as usize).refit_count, 1) == 0 { + active = false; + } else { + refit_node(&mut tree, node); + if node == 0 { + // The root. + active = false; + } else { + node = tree.at(node as usize).parent; + } + } } - curr_id = tree.at(curr_id as usize).parent; + refit_barrier(); } } } +/// `if c { a } else { b }` as a select instead of a branch. +#[inline(always)] +fn select_u32(c: bool, a: u32, b: u32) -> u32 { + let mask = 0u32.wrapping_sub(c as u32); + (a & mask) | (b & !mask) +} + +/// Pairs each thread of [`gpu_lbvh_find_collision_pairs`] buffers in workgroup memory. +const PAIRS_PER_THREAD: usize = 8; + /// Finds collision pairs by traversing the LBVH tree. +/// +/// A pair belongs to its smaller leaf when one leaf is more than twice the other's size, and to +/// its first sorted leaf otherwise: a large leaf (e.g. the ground) then prunes the tree right +/// away instead of visiting every small leaf it overlaps on a single thread. +/// +/// Each workgroup gathers its pairs in workgroup memory and reserves their range of the pair +/// buffer with one atomic, so a workgroup's pairs (of nearby leaves) stay contiguous. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_lbvh_find_collision_pairs( #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] tree: &[LbvhNode], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] @@ -497,8 +682,13 @@ pub fn gpu_lbvh_find_collision_pairs( collision_groups: &[InteractionGroups], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] pair_filter: &[[u32; 2]], + #[spirv(workgroup)] wg_pairs: &mut [UVec2; PAIRS_PER_THREAD * WORKGROUP_SIZE as usize], + #[spirv(workgroup)] wg_scan: &mut [u32; WORKGROUP_SIZE as usize], + #[spirv(workgroup)] wg_base: &mut [u32; 1], ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; + let tid = local_id.x as usize; + let mut num_buffered = 0usize; let batch_id = invocation_id.y; let colliders_start = scratch_start(batch_ids, batch_id); let num_bodies = batch_ids.colliders_len; @@ -513,95 +703,121 @@ pub fn gpu_lbvh_find_collision_pairs( let groups_i = collision_groups[i as usize]; let filter_i = pair_filter[i as usize]; let mut aabb1 = tree.at((first_leaf_id + leaf_i) as usize).aabb; + let extent1 = max_extent(&aabb1); let prediction = 2.0e-3; // TODO: should be configurable. let dilation = Vector::splat(prediction); aabb1.mins -= dilation; aabb1.maxs += dilation; - // Traverse the tree using a stack. - let mut stack = [0u32; 64]; - let mut stack_len = 1u32; - stack.write(0, 0); - + // Stackless depth-first traversal: a node's escape index skips its subtree. + let mut curr_id = 0u32; // NOTE: we use a fixed-size for loop to avoid miscompilation issues of - // while loops on MacOs. Each tree node is pushed at most once per + // while loops on MacOs. Each node is visited at most once per // traversal, so `2 * num_bodies` (≥ node count) bounds the loop. for _ in 0..2 * num_bodies { - if stack_len == 0 { + if curr_id == NO_ESCAPE { break; } - stack_len -= 1; - let curr_id = stack.read(stack_len as usize); let node = tree.at(curr_id as usize); + let is_leaf = curr_id >= first_leaf_id; + + // Only subtrees with a leaf this one owns a pair with: a much larger leaf, or a + // similar one after `leaf_i`. For a leaf, that's whether it owns the pair. The next + // node is selected without branching, so the node's fields load at once. + let owned = (node.max_extent > 2.0 * extent1) + | ((leaf_i < node.last_leaf) & (node.max_extent >= 0.5 * extent1)); + let hit = owned & aabb1.intersects(&node.aabb); + let descend = hit & !is_leaf; + curr_id = select_u32(descend, node.left, node.escape); + if !(hit & is_leaf) { + continue; + } - if curr_id >= first_leaf_id { - // We reached a leaf, register a collision pair. - let j = node.left; - let groups_j = collision_groups[j as usize]; + // We reached a leaf, register a collision pair. + let j = node.left; + let groups_j = collision_groups[j as usize]; - // Skip pairs whose collision groups don't authorize an interaction. - if !groups_i.test(groups_j) { - continue; - } + // Skip pairs whose collision groups don't authorize an interaction. + if !groups_i.test(groups_j) { + continue; + } - // Apply computed filters (same-body and self-contact). - let filter_j = pair_filter[j as usize]; - if filter_i[0] == filter_j[0] || (filter_i[1] != 0 && filter_i[1] == filter_j[1]) { - continue; - } + // Apply computed filters (same-body and self-contact). + let filter_j = pair_filter[j as usize]; + if filter_i[0] == filter_j[0] || (filter_i[1] != 0 && filter_i[1] == filter_j[1]) { + continue; + } - // Single global counter: every batch appends to the same flat - // buffer (pairs from different batches interleave freely). - let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(0), 1); - - // NOTE: if the index is out-of-bounds (meaning the `collision_pairs` isn't - // big enough), don't write. But keep traversing so we get the exact count we need - // for reallocating the buffers. - if target_pair_index < batch_ids.collision_pairs_capacity { - // NOTE: we only store the collider pair here, with global - // collider ids (the owning batch is recovered from the - // id downstream). The parent body ids are resolved - // lazily, at the very last moment, when the - // narrow-phase writes the `IndexedManifold` consumed - // by the solver — keeping this hot buffer (and the - // intermediate pfm-pair buffer) narrow, and keeping - // `collider_parent` out of the broad phase entirely. - let (ci, cj) = if i < j { (i, j) } else { (j, i) }; - collision_pairs.write( - target_pair_index as usize, - CollisionPair { - colliders: UVec2::new( - batch_ids.body_global(batch_id, ci) as u32, - batch_ids.body_global(batch_id, cj) as u32, - ), - }, - ); - } - } else { - let left = node.left; - let right = node.right; - - // Descend only if the subtree contains leaf id smaller than - // `leaf_i`. That way we prune the part of the tree that’s - // on the "left" of `leaf_i`, avoiding duplicate pairs/traversals. - if leaf_i < tree.at(left as usize).refit_count_or_max_subtree_index - && aabb1.intersects(&tree.at(left as usize).aabb) - && stack_len < 64 - { - stack.write(stack_len as usize, node.left); - stack_len += 1; - } + let (ci, cj) = if i < j { (i, j) } else { (j, i) }; + let pair = UVec2::new( + batch_ids.body_global(batch_id, ci) as u32, + batch_ids.body_global(batch_id, cj) as u32, + ); + // Slot-major, so that consecutive threads hit consecutive words. + if num_buffered < PAIRS_PER_THREAD { + wg_pairs.write(num_buffered * WORKGROUP_SIZE as usize + tid, pair); + num_buffered += 1; + continue; + } - if leaf_i < tree.at(right as usize).refit_count_or_max_subtree_index - && aabb1.intersects(&tree.at(right as usize).aabb) - && stack_len < 64 - { - stack.write(stack_len as usize, node.right); - stack_len += 1; - } + // Single global counter: every batch appends to the same flat + // buffer (pairs from different batches interleave freely). + let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(0), 1); + + // NOTE: if the index is out-of-bounds (meaning the `collision_pairs` isn't + // big enough), don't write. But keep traversing so we get the exact count we need + // for reallocating the buffers. + if target_pair_index < batch_ids.collision_pairs_capacity { + // NOTE: we only store the collider pair here, with global + // collider ids (the owning batch is recovered from the + // id downstream). The parent body ids are resolved + // lazily, at the very last moment, when the + // narrow-phase writes the `IndexedManifold` consumed + // by the solver — keeping this hot buffer (and the + // intermediate pfm-pair buffer) narrow, and keeping + // `collider_parent` out of the broad phase entirely. + collision_pairs.write( + target_pair_index as usize, + CollisionPair { colliders: pair }, + ); } } } + + // Offsets of the buffered pairs in the workgroup's range (inclusive Hillis-Steele scan). + wg_scan.write(tid, num_buffered as u32); + let mut offset = 1usize; + for _ in 0..6u32 { + workgroup_memory_barrier_with_group_sync(); + let v = if tid >= offset { + wg_scan.read(tid - offset) + } else { + 0 + }; + workgroup_memory_barrier_with_group_sync(); + let s = wg_scan.read(tid) + v; + wg_scan.write(tid, s); + offset *= 2; + } + workgroup_memory_barrier_with_group_sync(); + if tid == WORKGROUP_SIZE as usize - 1 { + let total = wg_scan.read(tid); + wg_base.write(0, atomic_add_u32(collision_pairs_len.at_mut(0), total)); + } + workgroup_memory_barrier_with_group_sync(); + + let first = wg_base.read(0) + wg_scan.read(tid) - num_buffered as u32; + for k in 0..PAIRS_PER_THREAD { + let target_pair_index = first + k as u32; + if k < num_buffered && target_pair_index < batch_ids.collision_pairs_capacity { + collision_pairs.write( + target_pair_index as usize, + CollisionPair { + colliders: wg_pairs.read(k * WORKGROUP_SIZE as usize + tid), + }, + ); + } + } } /// Expands a 10-bit integer into 30 bits by inserting 2 zeros after each bit (3D). @@ -708,3 +924,32 @@ fn root_id(collider_start_id: u32) -> u32 { // `- b`. collider_start_id * 2 } + +#[cfg(all(test, not(target_arch_is_gpu)))] +mod tests { + use super::*; + + #[test] + fn web_refit_plan_covers_all_batches_and_clamps_stale_counts() { + for (counts, leaves, expected) in [ + (vec![], 513, 0), + (vec![0, 0, 0], 513, 0), + (vec![1, 256, 0], 513, 1), + (vec![0, 257, 512], 513, 2), + (vec![513, 1, 0], 513, 3), + (vec![u32::MAX, 1], 256, 1), + ] { + let mut rounds = 999; + gpu_lbvh_refit_plan( + &counts, + &mut rounds, + &BatchIndices { + num_batches: counts.len() as u32, + colliders_len: leaves, + ..Default::default() + }, + ); + assert_eq!(rounds, expected); + } + } +} diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 8857e3de..ad4ce8d9 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -35,12 +35,12 @@ const WORKGROUP_SIZE: u32 = 64; /// `i`-th PFM entry (in sorted order when the sort runs) owns slot /// `pfm_base + i`. Every slot in `[0, bound)` is (re)written each frame by /// exactly one producer (`len = 0` when the pair yields no manifold), so no -/// zeroing pass is needed and the flat consumers can sweep the whole bound. +/// zeroing pass is needed and the flat consumers can cover the whole bound. #[derive(Copy, Clone, Default)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] pub struct ContactPlan { - /// Total contact-slot bound (`pfm_base + pfm_len`): the sweep range of + /// Total contact-slot bound (`pfm_base + pfm_len`): the dispatch range of /// every flat contacts-keyed kernel. pub bound: u32, /// Clamped flat collision-pair total; also the base contact slot of the @@ -64,8 +64,8 @@ pub fn gpu_reset_narrow_phase( } /// Publishes this frame's clamped list totals (the `contact_plan`) and every -/// grid derived from them: the flat contacts sweep grid, the PFM sweep grid, -/// the (clamped) PFM sort count, and the per-multibody contact sweep grid +/// grid derived from them: the flat contacts dispatch grid, the PFM dispatch grid, +/// the (clamped) PFM sort count, and the per-multibody contact dispatch grid /// (`[multibodies_batch_capacity, num_batches, 1]`, zero workgroups when the /// frame cannot produce any contact). /// @@ -84,7 +84,7 @@ pub fn gpu_contact_plan( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_sort_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contacts_indirect: &mut [u32; 3], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] pfm_indirect: &mut [u32; 3], - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] mb_sweep_indirect: &mut [u32; 3], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] mb_dispatch_indirect: &mut [u32; 3], #[spirv(uniform, descriptor_set = 0, binding = 7)] batch_ids: &BatchIndices, ) { let pairs = @@ -106,13 +106,13 @@ pub fn gpu_contact_plan( *pfm_indirect.at_mut(1) = 1; *pfm_indirect.at_mut(2) = 1; - *mb_sweep_indirect.at_mut(0) = if bound > 0 { + *mb_dispatch_indirect.at_mut(0) = if bound > 0 { batch_ids.multibodies_batch_capacity } else { 0 }; - *mb_sweep_indirect.at_mut(1) = batch_ids.num_batches; - *mb_sweep_indirect.at_mut(2) = 1; + *mb_dispatch_indirect.at_mut(1) = batch_ids.num_batches; + *mb_dispatch_indirect.at_mut(2) = 1; } /// Copies each PFM entry's originating pair index into the flat sort-key @@ -357,7 +357,7 @@ pub fn gpu_reduce_contacts( /// Contact slots are positional: pair `t` owns contact slot `t` and this pass /// writes every slot in `[0, pairs_total)` exactly once (`len = 0` when the /// pair yields no manifold here: separated, same-body, or deferred), so the -/// flat consumers can sweep the whole bound without a zeroing pass. +/// flat consumers can cover the whole bound without a zeroing pass. /// /// The complex cases (generic convex via PFM, trimesh, polyline) are deferred /// to `gpu_narrow_phase_shape_shape_deferred`. @@ -396,6 +396,7 @@ pub fn gpu_narrow_phase_shape_shape( let body1 = collider_parent.read(pair.colliders.x as usize); let body2 = collider_parent.read(pair.colliders.y as usize); let mut manifold = ContactManifold::default(); + let mut recycle_extent = 0.0; if body1 != body2 { let pose1 = poses.read(pair.colliders.x as usize); let pose2 = poses.read(pair.colliders.y as usize); @@ -404,6 +405,9 @@ pub fn gpu_narrow_phase_shape_shape( let shape_ty1 = shape1.shape_type(); let shape_ty2 = shape2.shape_type(); let pose12 = pose1.inverse() * pose2; + recycle_extent = shape1 + .primitive_local_extent() + .max(shape2.primitive_local_extent()); // Ball - Convex if shape_ty1 == SHAPE_TYPE_BALL { @@ -455,7 +459,7 @@ pub fn gpu_narrow_phase_shape_shape( friction: mat1.combined_friction(&mat2), restitution: mat1.combined_restitution(&mat2), subshape: 0, - _padding: 0.0, + recycle_extent, }, ); } else { @@ -564,7 +568,10 @@ pub fn gpu_narrow_phase_shape_shape_deferred( colliders: pair.colliders, pair_index: t, subshape: 0, - _padding: [0; 2], + recycle_extent: shape1 + .local_extent(vertices) + .max(shape2.local_extent(vertices)), + _padding: [0; 1], }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // NOTE: if we exceed capacity, just skip the pair. @@ -715,7 +722,9 @@ fn trimesh_convex( colliders, pair_index, subshape: idx.shape_index + 1, - _padding: [0; 2], + // One manifold per triangle or segment: not recyclable. + recycle_extent: -1.0, + _padding: [0; 1], }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // Skip (don’t write) on overflow; the caller resizes and re-runs. @@ -793,7 +802,9 @@ fn polyline_convex( colliders, pair_index, subshape: idx.shape_index + 1, - _padding: [0; 2], + // One manifold per triangle or segment: not recyclable. + recycle_extent: -1.0, + _padding: [0; 1], }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // Skip (don’t write) on overflow; the caller resizes and re-runs. @@ -829,7 +840,9 @@ pub struct NarrowPhasePfmPair { pair_index: u32, /// See [`IndexedManifold::subshape`]. subshape: u32, - _padding: [u32; 2], + /// See [`IndexedManifold::recycle_extent`]. + recycle_extent: f32, + _padding: [u32; 1], } /// PFM (GJK/EPA) manifold computation for the deferred work-list entries. @@ -906,7 +919,7 @@ pub fn gpu_narrow_phase_pfm_pfm( friction: mat1.combined_friction(&mat2), restitution: mat1.combined_restitution(&mat2), subshape: pair.subshape, - _padding: 0.0, + recycle_extent: pair.recycle_extent, }, ); } else { diff --git a/src_rbd_shaders/dynamics/body.rs b/src_rbd_shaders/dynamics/body.rs index 10bae6d3..c0e57371 100644 --- a/src_rbd_shaders/dynamics/body.rs +++ b/src_rbd_shaders/dynamics/body.rs @@ -210,7 +210,7 @@ impl Force { } /// Linear and angular velocity of a rigid body. -#[derive(Clone, Copy, Default)] +#[derive(Clone, Copy, Default, PartialEq, Debug)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] pub struct Velocity { diff --git a/src_rbd_shaders/dynamics/color_buckets.rs b/src_rbd_shaders/dynamics/color_buckets.rs index b8cff93f..3cd4d411 100644 --- a/src_rbd_shaders/dynamics/color_buckets.rs +++ b/src_rbd_shaders/dynamics/color_buckets.rs @@ -1,22 +1,43 @@ //! Bucket-sort of contact constraints by graph-coloring color. //! -//! After the per-step (global) coloring converges, the constraint indices are -//! bucket-sorted by `(color, batch)` into `color_sorted_ids`. Buckets are laid +//! After the per-step (global) coloring converges, the constraint links are +//! bucket-sorted by `(color, batch)` into `sorted_links`. Buckets are laid //! out color-major (`bucket = color * num_batches + batch`, buffer length //! `solver_color_buckets_stride * num_batches`), so one color's constraints -//! are contiguous across every batch (per-color solver sweeps) while each -//! `(color, batch)` cell stays contiguous too (fused per-batch sweeps). +//! are contiguous across every batch (per-color solver dispatches) while each +//! `(color, batch)` cell stays contiguous too (fused per-batch dispatches). +//! +//! The position of a constraint in `sorted_links` is its index in the contact tiles. +//! Active constraints the bounded coloring left uncolored (or past the solved colors) are never +//! solved, but still need a place in the tiles for their warmstart: they go to the last color +//! bucket. use crate::broad_phase::ContactPlan; +use crunchy::unroll; use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::atomic_add_u32_workgroup; +use khal_std::sync::workgroup_memory_barrier_with_group_sync; use khal_std::{index::MaybeIndexUnchecked, iter::StepRng, sync::atomic_add_u32}; -use super::constraint::TwoBodyConstraint; +use super::ContactLink; use crate::utils::{BatchIndices, Slice}; const WORKGROUP_SIZE: u32 = 64; +/// The bucket color of a constraint: its color if it is solved, the last color bucket if it is +/// active but not solved, or 0 for gap, inert and multibody-owned slots (no bucket). +#[inline(always)] +fn bucket_color(color: u32, len: u32, stride: u32) -> u32 { + if len == 0 { + 0 + } else if color != 0 && color < stride - 1 { + color + } else { + stride - 1 + } +} + /// Zeroes the `(color, batch)` bucket counts (flat 1-D grid over the whole /// bucket buffer). #[spirv_bindgen] @@ -31,17 +52,21 @@ pub fn gpu_color_buckets_reset( } } -/// Counts how many constraints fall in each `(color, batch)` bucket. +/// Counts how many constraints fall in each `(color, batch)` bucket, and marks every sorted +/// link as inactive before the scatter fills them. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_color_buckets_count( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &[ContactLink], #[spirv(uniform, descriptor_set = 0, binding = 2)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_buckets: &mut [u32], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] sorted_links: &mut [ContactLink], + #[spirv(workgroup)] local_counts: &mut [u32; 128], ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let nb = batch_ids.num_batches; @@ -49,46 +74,199 @@ pub fn gpu_color_buckets_count( let total = contact_plan.bound; let constraints = Slice(constraints, 0); - for i in StepRng::new(invocation_id.x..total, num_threads) { - let color = constraints_colors.read(i as usize); - // Color 0 (uncolored / gap slots) is never swept; colors past the - // swept range (bounded coloring didn't converge) are dropped. They - // were never solved before either. Skipping color 0 also keeps the - // stale body ids of gap slots from being dereferenced. - if color != 0 && color < stride - 1 { - let batch = batch_ids.collider_batch(constraints[i as usize].solver_body_a); - atomic_add_u32(color_buckets.at_mut((color * nb + batch) as usize), 1); + // Naga lowers num_workgroups through a private global, so browser WGSL + // validation cannot prove this condition uniform. Keep every barrier outside + // the branch. Run the generic path last: its checked batch division can lower + // to an early exit, which would also make subsequent barriers non-uniform. + // Fallback lanes initialize an empty histogram and reach every barrier. The last bucket, + // `stride - 1`, must fit the 128-entry histogram. + let use_histogram = nb == 1 && stride <= 128 && num_threads >= total; + let lane = local_id.x as usize; + local_counts.write(lane, 0); + local_counts.write(lane + 64, 0); + workgroup_memory_barrier_with_group_sync(); + if use_histogram && invocation_id.x < total { + let i = invocation_id.x as usize; + let color = bucket_color(constraints_colors.read(i), constraints[i].len, stride); + if color != 0 { + atomic_add_u32_workgroup(local_counts.at_mut(color as usize), 1); + } + } + workgroup_memory_barrier_with_group_sync(); + crunchy::unroll! { for k in 0..2 { + let bucket = lane + k * 64; + let size = local_counts.read(bucket); + if size != 0 { + atomic_add_u32(color_buckets.at_mut(bucket), size); + } + } } + + if !use_histogram { + for i in StepRng::new(invocation_id.x..total, num_threads) { + let link = &constraints[i as usize]; + // Skipping the slots without bucket also keeps the stale body ids of gap slots + // from being dereferenced. + let color = bucket_color(constraints_colors.read(i as usize), link.len, stride); + if color != 0 { + let batch = batch_ids.collider_batch(link.solver_body_a); + atomic_add_u32(color_buckets.at_mut((color * nb + batch) as usize), 1); + } } } + for i in StepRng::new(invocation_id.x..total, num_threads) { + sorted_links.at_mut(i as usize).len = 0; + } } -/// Scatters each constraint index into its `(color, batch)` bucket. The bucket +/// Scatters each constraint link into its `(color, batch)` bucket. The bucket /// buffer holds the scanned exclusive starts, used as cursors; after this pass /// every entry is its bucket's exclusive end. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_color_buckets_scatter( - #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(global_invocation_id)] gid: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &[ContactLink], #[spirv(uniform, descriptor_set = 0, binding = 2)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_buckets: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] color_sorted_ids: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] sorted_links: &mut [ContactLink], #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] constraint_indices: &mut [u32], + #[spirv(workgroup)] local_counts: &mut [u32; 128], + #[spirv(workgroup)] local_starts: &mut [u32; 128], ) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; let nb = batch_ids.num_batches; let stride = batch_ids.solver_color_buckets_stride; let total = contact_plan.bound; - let constraints = Slice(constraints, 0); + let lane = local_id.x as usize; + let num_threads = num_workgroups.x * WORKGROUP_SIZE; + // Match the count fast path's bounds. Reserving one contiguous span per local + // color reduces global contention and keeps nearby input constraints together. + // As in count, fallback lanes must reach all barriers for WGSL uniformity. + let use_histogram = nb == 1 && stride <= 128 && num_threads >= total; + local_counts.write(lane, 0); + local_counts.write(lane + 64, 0); + workgroup_memory_barrier_with_group_sync(); + let color = if use_histogram && gid.x < total { + let i = gid.x as usize; + bucket_color(constraints_colors.read(i), constraints.at(i).len, stride) + } else { + 0 + }; + let valid = color != 0; + let rank = if valid { + atomic_add_u32_workgroup(local_counts.at_mut(color as usize), 1) + } else { + 0 + }; + workgroup_memory_barrier_with_group_sync(); + crunchy::unroll! { for k in 0..2 { + let bucket = lane + k * 64; + let size = local_counts.read(bucket); + if size != 0 { + let start = atomic_add_u32(color_buckets.at_mut(bucket), size); + local_starts.write(bucket, start); + } + } } + workgroup_memory_barrier_with_group_sync(); + if valid { + let index = local_starts.read(color as usize) + rank; + sorted_links.write(index as usize, constraints.read(gid.x as usize)); + constraint_indices.write(gid.x as usize, index); + } - for i in StepRng::new(invocation_id.x..total, num_threads) { - let color = constraints_colors.read(i as usize); - if color != 0 && color < stride - 1 { - let batch = batch_ids.collider_batch(constraints[i as usize].solver_body_a); - let dst = atomic_add_u32(color_buckets.at_mut((color * nb + batch) as usize), 1); - color_sorted_ids.write(dst as usize, i); + if !use_histogram { + for i in StepRng::new(gid.x..total, num_threads) { + let link = constraints.at(i as usize); + let color = bucket_color(constraints_colors.read(i as usize), link.len, stride); + if color != 0 { + let batch = batch_ids.collider_batch(link.solver_body_a); + let dst = atomic_add_u32(color_buckets.at_mut((color * nb + batch) as usize), 1); + sorted_links.write(dst as usize, *link); + constraint_indices.write(i as usize, dst); + } } } } + +/// Statistics of the colored constraints, read back to size the next step's dispatches. +#[derive(Copy, Clone, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct ColorStats { + /// Size of the largest color bucket, over all batches. + pub max_bucket: u32, + /// Highest non-empty color. + pub highest_color: u32, + /// Size of each of the colors `0..64`, over all batches. + pub sizes: [u32; 64], +} + +/// Sizes the colored dispatches from the color buckets: their grid only needs to cover the +/// largest color (the solve kernels stride over their bucket), not every contact. Also records the +/// highest color in use, read back to adapt the color budget. One workgroup. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_color_dispatch_grid( + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] color_buckets: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] color_stats: &mut ColorStats, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dispatch_indirect: &mut [[u32; 3]], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(workgroup)] max_sizes: &mut [u32; 64], + #[spirv(workgroup)] max_colors: &mut [u32; 64], +) { + let lane = local_id.x as usize; + let nb = batch_ids.num_batches; + // Swept colors are `1..=stride - 2`; color `c` spans `[ends[c*nb - 1], ends[(c+1)*nb - 1])`. + let num_colors = batch_ids.solver_color_buckets_stride - 2; + + // Per-color sizes are CPU scheduling hints for subsequent steps. They never bound + // the actual dispatch: every dispatch traverses its current GPU bucket with a strided for. + let color = lane as u32; + let size = if color > 0 && color <= num_colors { + color_buckets.read(((color + 1) * nb - 1) as usize) + - color_buckets.read((color * nb - 1) as usize) + } else { + 0 + }; + color_stats.sizes.write(lane, size); + + let mut max_size = 0u32; + let mut max_color = 0u32; + for color in StepRng::new(1 + lane as u32..num_colors + 1, WORKGROUP_SIZE) { + let start = color_buckets.read((color * nb - 1) as usize); + let end = color_buckets.read(((color + 1) * nb - 1) as usize); + let size = end - start; + max_size = max_size.max(size); + if size > 0 { + max_color = max_color.max(color); + } + } + max_sizes.write(lane, max_size); + max_colors.write(lane, max_color); + + let mut stride = 32usize; + for _ in 0..6u32 { + workgroup_memory_barrier_with_group_sync(); + if lane < stride { + let s = max_sizes.read(lane).max(max_sizes.read(lane + stride)); + let c = max_colors.read(lane).max(max_colors.read(lane + stride)); + max_sizes.write(lane, s); + max_colors.write(lane, c); + } + stride /= 2; + } + workgroup_memory_barrier_with_group_sync(); + + if lane == 0 { + let max_size = max_sizes.read(0); + color_stats.max_bucket = max_size; + color_stats.highest_color = max_colors.read(0); + // Every color shares this grid. Limit the empty workgroups in the smaller buckets; + // the solve kernels stride over large buckets, so no constraints are skipped. + dispatch_indirect.write(0, [max_size.div_ceil(WORKGROUP_SIZE).min(256), 1, 1]); + } +} diff --git a/src_rbd_shaders/dynamics/coloring.rs b/src_rbd_shaders/dynamics/coloring.rs index b47760e7..b6c8579a 100644 --- a/src_rbd_shaders/dynamics/coloring.rs +++ b/src_rbd_shaders/dynamics/coloring.rs @@ -6,15 +6,12 @@ use crate::broad_phase::ContactPlan; use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; -use khal_std::{ - iter::StepRng, - sync::{atomic_add_u32, atomic_max_u32}, -}; +use khal_std::{iter::StepRng, sync::atomic_add_u32}; use crate::utils::{BatchIndices, Slice, SliceMut}; use khal_std::index::MaybeIndexUnchecked; -use super::constraint::TwoBodyConstraint; +use super::ContactLink; const WORKGROUP_SIZE: u32 = 64; @@ -38,6 +35,19 @@ fn hash(packed_key: u32) -> u32 { key } +/// Lowest free color. The upper word only contributes when the lower word is full. +/// Adding both words' trailing-zero counts can select an occupied color when a previous +/// frame seeded colors above 31, repeatedly recreating the same conflict. +#[inline(always)] +pub(super) fn first_free_color(mask: (u32, u32)) -> u32 { + let low = (!mask.0).trailing_zeros(); + if low < 32 { + low + } else { + 32 + (!mask.1).trailing_zeros() + } +} + /// Returns the `n`-th free color of a 64-bit mask (`n = 0` is the lowest), or 63 if none. /// Spreads the picks of the same round over different colors to avoid conflicts. #[inline] @@ -76,7 +86,7 @@ pub fn gpu_reset_luby( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints_rands: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[ContactLink], #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, ) { let total = contact_plan.bound; @@ -104,7 +114,7 @@ pub fn gpu_step_graph_coloring_luby( #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] body_constraint_counts: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_constraint_ids: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[ContactLink], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] constraints_rands: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] uncolored: &mut u32, @@ -212,13 +222,19 @@ pub fn gpu_reset_topo_gc( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] colored: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[ContactLink], #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] pending_colors: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] num_colors: &mut u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] pending_colors: &mut [u32], ) { let total = contact_plan.bound; let i = invocation_id.x; + if i == 0 { + // Not converged: the first iteration runs. + *num_colors = 0; + } + if i < total { let idx = i as usize; // Color 0 is reserved for "uncolored" state @@ -232,27 +248,30 @@ pub fn gpu_reset_topo_gc( } } -/// Resets the convergence flag for Topo-GC. +/// Resets the convergence flag for Topo-GC, and sizes this iteration's dispatches: none once +/// the coloring has converged (the flag stayed non-zero through the previous iteration). #[spirv_bindgen] #[spirv(compute(threads(1)))] pub fn gpu_reset_completion_flag_topo_gc( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] num_colors: &mut u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_indirect: &[[u32; 3]], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] coloring_indirect: &mut [[u32; 3]], ) { if invocation_id.x == 0 { - // NOTE: this `for` loop is silly. It doesn’t do anything - // more than a `*num_colors = 1` in a convoluted - // way because otherwise rustgpu apparently does not generate - // the spirv for this kernel (seems to happen if the kernel is - // too trivial. - for k in 0..1 { - // Non-zero value indicates algorithm should continue - *num_colors = k + 1; - } + let converged = *num_colors != 0; + let grid = if converged { + [0, 1, 1] + } else { + contacts_indirect.read(0) + }; + coloring_indirect.write(0, grid); + *num_colors = 1; } } -/// Clears the convergence flag so the first fix-conflicts pass validates every color. +/// Clears the convergence flag and opens the full coloring grid so the first fix-conflicts pass +/// validates every color. /// /// Colors seeded from the previous frame are marked colored, so the first step may color /// nothing and leave the flag set, which would skip the validation of the seeds. @@ -261,12 +280,12 @@ pub fn gpu_reset_completion_flag_topo_gc( pub fn gpu_clear_completion_flag_topo_gc( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] num_colors: &mut u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_indirect: &[[u32; 3]], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] coloring_indirect: &mut [[u32; 3]], ) { if invocation_id.x == 0 { - // NOTE: same trivial-kernel workaround as `gpu_reset_completion_flag_topo_gc`. - for k in 0..1 { - *num_colors = k; - } + coloring_indirect.write(0, contacts_indirect.read(0)); + *num_colors = 0; } } @@ -280,7 +299,7 @@ pub fn gpu_step_graph_coloring_topo_gc( #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] body_constraint_counts: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_constraint_ids: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[ContactLink], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] colored: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] num_colors: &mut u32, @@ -322,14 +341,23 @@ pub fn gpu_step_graph_coloring_topo_gc( } else { 0 }; - let last_constraint_id_a = body_constraint_counts[body_a as usize] as usize; + let mut last_constraint_id_a = body_constraint_counts[body_a as usize] as usize; let first_constraint_id_b = if body_b != 0 { body_constraint_counts[body_b as usize - 1] as usize } else { 0 }; - let last_constraint_id_b = body_constraint_counts[body_b as usize] as usize; + let mut last_constraint_id_b = body_constraint_counts[body_b as usize] as usize; + + // A split body's constraints each act on their own sub-body: they don't conflict. + // The split points them at a sub-body velocity slot instead of the body's own. + if constraints[i].vel_slot_a != constraints[i].solver_body_a { + last_constraint_id_a = first_constraint_id_a; + } + if constraints[i].vel_slot_b != constraints[i].solver_body_b { + last_constraint_id_b = first_constraint_id_b; + } // Traverse all constraints from body A. for j in first_constraint_id_a..last_constraint_id_a { @@ -382,8 +410,7 @@ pub fn gpu_step_graph_coloring_topo_gc( // `colored` is not written here since the rank loops read it. // The fix pass sets it, based on `pending_colors[i]`. } else { - constraints_colors[i] = - (!color_mask.0).trailing_zeros() + (!color_mask.1).trailing_zeros(); + constraints_colors[i] = first_free_color(color_mask); colored[i] = 1; } // We are not finished coloring. 0 indicates the algorithm must continue. @@ -403,7 +430,7 @@ pub fn gpu_fix_conflicts_topo_gc( #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] body_constraint_counts: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_constraint_ids: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[ContactLink], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] colored: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] num_colors: &mut u32, @@ -451,17 +478,10 @@ pub fn gpu_fix_conflicts_topo_gc( colored[i] = PENDING_COLORED; } - // NOTE: this `num_colors` read doesn't need to be atomic. Any non-zero value is indicative of a finished - // algorithm. - // If `num_colors > 0u` then we know that the coloring algorithm has converged. So we use this dispatch - // as an opportunity to compute the colors count that will be ready back to the CPU side. - if *num_colors > 0 { - // TODO PERF: not sure if that would have a significant impact but we could keep track of - // whether the last iteration of the TOPO-GC algorithm already finished, in which - // case we can skip the atomic max entirely and just early-exit. - - atomic_max_u32(num_colors, color_i); - } else { + // A non-zero `num_colors` means the previous step iteration colored nothing new: the + // coloring has converged and there is nothing left to fix. (The CPU only reads it as a + // convergence flag.) + if *num_colors == 0 { // Map raw body ids to graph-coloring GROUP ids (multibody-aware). let body_a = body_group[constraints[i].solver_body_a as usize]; let body_b = body_group[constraints[i].solver_body_b as usize]; @@ -471,21 +491,30 @@ pub fn gpu_fix_conflicts_topo_gc( } else { 0 }; - let last_constraint_id_a = body_constraint_counts[body_a as usize] as usize; + let mut last_constraint_id_a = body_constraint_counts[body_a as usize] as usize; let first_constraint_id_b = if body_b != 0 { body_constraint_counts[body_b as usize - 1] as usize } else { 0 }; - let last_constraint_id_b = body_constraint_counts[body_b as usize] as usize; + let mut last_constraint_id_b = body_constraint_counts[body_b as usize] as usize; + + // A split body's constraints each act on their own sub-body: they don't conflict. + // The split points them at a sub-body velocity slot instead of the body's own. + if constraints[i].vel_slot_a != constraints[i].solver_body_a { + last_constraint_id_a = first_constraint_id_a; + } + if constraints[i].vel_slot_b != constraints[i].solver_body_b { + last_constraint_id_b = first_constraint_id_b; + } // Traverse all constraints from body A. On a conflict, the largest index keeps // its color, the others pick again in the next round. for j in first_constraint_id_a..last_constraint_id_a { let constraint_j = body_constraint_ids[j]; - if constraint_j != constraint_i { + if constraint_i < constraint_j { let cj = constraint_j as usize; let pending_j = if deterministic { pending_colors[cj] @@ -497,32 +526,36 @@ pub fn gpu_fix_conflicts_topo_gc( } else { constraints_colors[cj] }; - if color_i == color_j && constraint_i < constraint_j { + if color_i == color_j { // Found a conflict, uncolor this node. colored[i] = 0; + break; } } } - // Traverse all constraints from body B. - for j in first_constraint_id_b..last_constraint_id_b { - let constraint_j = body_constraint_ids[j]; + if colored[i] != 0 { + // Traverse all constraints from body B. + for j in first_constraint_id_b..last_constraint_id_b { + let constraint_j = body_constraint_ids[j]; - if constraint_j != constraint_i { - let cj = constraint_j as usize; - let pending_j = if deterministic { - pending_colors[cj] - } else { - MAX_U32 - }; - let color_j = if pending_j != MAX_U32 { - pending_j - } else { - constraints_colors[cj] - }; - if color_i == color_j && constraint_i < constraint_j { - // Found a conflict, uncolor this node. - colored[i] = 0; + if constraint_i < constraint_j { + let cj = constraint_j as usize; + let pending_j = if deterministic { + pending_colors[cj] + } else { + MAX_U32 + }; + let color_j = if pending_j != MAX_U32 { + pending_j + } else { + constraints_colors[cj] + }; + if color_i == color_j { + // Found a conflict, uncolor this node. + colored[i] = 0; + break; + } } } } @@ -536,3 +569,7 @@ pub fn gpu_fix_conflicts_topo_gc( } } } + +#[cfg(all(test, not(target_arch_is_gpu)))] +#[path = "../tests/coloring.rs"] +mod tests; diff --git a/src_rbd_shaders/dynamics/constraint.rs b/src_rbd_shaders/dynamics/constraint.rs index 5ff6d312..0e0c3dfc 100644 --- a/src_rbd_shaders/dynamics/constraint.rs +++ b/src_rbd_shaders/dynamics/constraint.rs @@ -1,9 +1,17 @@ //! Contact constraint data structures for the iterative solver. +//! +//! A [`TwoBodyConstraint`] is one contact manifold between two bodies. It is kept compact +//! since every solver iteration streams it: each contact point only stores its lever arms, +//! anchors, impulses and effective masses, and the constraint stores both bodies' inverse +//! mass and (symmetric) world inverse inertia, from which the angular Jacobians are +//! recomputed on the fly. -use crate::{AngVector, Pad, Vector}; +use crate::{AngVector, Pose, Vector, gcross, gdot}; +use glamx::UVec2; +use khal_std::index::MaybeIndexUnchecked; #[cfg(feature = "dim3")] -use glamx::Vec2; +use glamx::{Vec2, Vec3}; #[cfg(feature = "dim2")] /// Number of tangent constraint directions (2D: one tangent perpendicular to normal). @@ -21,240 +29,357 @@ pub const MAX_CONSTRAINTS_PER_MANIFOLD: usize = 2; /// Maximum number of contact points per contact manifold (3D: up to 4). pub const MAX_CONSTRAINTS_PER_MANIFOLD: usize = 4; -/// Metadata for building a two-body constraint from a contact point. -/// -/// This data is used to initialize and update constraints at each solver substep. -#[derive(Clone, Copy, Default)] -#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] -#[repr(C)] -#[cfg(feature = "dim2")] // Same as dim3, but with a different fields ordering to limit padding. -pub struct TwoBodyConstraintInfos { - /// Tangent velocity (for conveyor belt effects, currently unused). - pub tangent_vel: Vector, // TODO PERF: could be one float less, be shared by both contact point infos? - - /// Contact point in body A's local coordinates (for warmstarting). - /// Stored in local space to detect matching contacts across frames. - pub local_pt_a: Vector, - - /// Contact point in body B's local coordinates (for warmstarting). - pub local_pt_b: Vector, - - /// Normal relative velocity at the contact point (for restitution). - pub normal_vel: f32, - - /// Penetration depth (negative = penetration, positive = separation). - pub dist: f32, -} - +/// A world-space inverse angular inertia tensor (symmetric). #[derive(Clone, Copy, Default)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] #[cfg(feature = "dim3")] -pub struct TwoBodyConstraintInfos { - /// Tangent velocity (for conveyor belt effects, currently unused). - pub tangent_vel: Vector, // TODO PERF: could be one float less, be shared by both contact point infos? - - /// Normal relative velocity at the contact point (for restitution). - pub normal_vel: f32, - - /// Contact point in body A's local coordinates (for warmstarting). - /// Stored in local space to detect matching contacts across frames. - pub local_pt_a: Vector, +pub struct SymInertia { + /// The diagonal `(xx, yy, zz)`. + pub diag: Vec3, pub _padding0: u32, - - /// Contact point in body B's local coordinates (for warmstarting). - pub local_pt_b: Vector, - - /// Penetration depth (negative = penetration, positive = separation). - pub dist: f32, + /// The off-diagonal terms `(xy, xz, yz)`. + pub off: Vec3, + pub _padding1: u32, } -/// Builder data for constructing constraints from contact manifolds. -/// -/// Stores auxiliary information needed to update constraints at each solver substep. -#[derive(Clone, Copy, Default)] -#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] -#[repr(C, align(16))] -pub struct TwoBodyConstraintBuilder { - /// Information for each contact point in the manifold. - pub infos: [TwoBodyConstraintInfos; MAX_CONSTRAINTS_PER_MANIFOLD], +#[cfg(feature = "dim3")] +impl SymInertia { + /// The symmetric part of `m`'s upper-left 3x3 block. + #[inline(always)] + pub fn from_mat4(m: glamx::Mat4) -> Self { + Self { + diag: Vec3::new(m.x_axis.x, m.y_axis.y, m.z_axis.z), + _padding0: 0, + off: Vec3::new(m.y_axis.x, m.z_axis.x, m.z_axis.y), + _padding1: 0, + } + } + + /// Multiplies `v` by this tensor. + #[inline(always)] + pub fn mul(&self, v: Vec3) -> Vec3 { + let d = self.diag; + let o = self.off; + Vec3::new( + d.x * v.x + o.x * v.y + o.y * v.z, + o.x * v.x + d.y * v.y + o.z * v.z, + o.y * v.x + o.z * v.y + d.z * v.z, + ) + } + + /// Scales every term of this tensor. + #[inline(always)] + pub fn scale(&mut self, s: f32) { + self.diag *= s; + self.off *= s; + } } -/// A contact constraint between two rigid bodies. -/// -/// Encodes all the data needed to solve a contact constraint, including: -/// - Constraint directions (normal and tangent). -/// - Effective masses and inverse masses. -/// - Solver coefficients (CFM factor, friction limit). -/// - Per-contact-point constraint elements. -// PERF: differentiate two-bodies and one-body constraints? +/// A contact manifold between two rigid bodies, as the solver sees it. #[derive(Clone, Copy, Default)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] +#[cfg(feature = "dim3")] pub struct TwoBodyConstraint { - /// Contact normal direction from body A's perspective (points away from A). - /// Normal impulses are applied along this direction to prevent penetration. - pub dir_a: Vector, // Non-penetration force direction for the first body. - /// Collider A of the source manifold (3D; stored in the `dir_a` padding lane). - #[cfg(feature = "dim3")] - pub warmstart_collider_a: u32, - - #[cfg(feature = "dim3")] - /// First tangent direction (3D only, orthogonal to normal). - /// Used for friction in the contact plane. - pub tangent_a: Vector, // One of the friction force directions. - /// Collider B of the source manifold (3D; stored in the `tangent_a` padding lane). - #[cfg(feature = "dim3")] - pub warmstart_collider_b: u32, - + /// Contact normal from body A's perspective (points away from A). + pub dir_a: Vector, + /// Number of active contact points in this manifold. + pub len: u32, + /// First friction direction (orthogonal to the normal); the second is `dir_a × tangent_a`. + pub tangent_a: Vector, + /// Friction coefficient. + pub limit: f32, /// Inverse mass of body A along each axis. - /// Used to compute linear velocity changes from impulses. pub im_a: Vector, - /// [`IndexedManifold::subshape`] of the source manifold (3D; `im_a` padding lane). + /// Index of body A (poses, mass properties). + pub solver_body_a: u32, + /// Inverse mass of body B along each axis. + pub im_b: Vector, + /// Index of body B. + pub solver_body_b: u32, + /// World-space inverse inertia of body A. + pub ii_a: SymInertia, + /// World-space inverse inertia of body B. + pub ii_b: SymInertia, + /// Index of body A's velocity in the solver velocity buffer: `solver_body_a`, or one of + /// the sub-body slots when body A is split (see `mass_splitting`). + pub vel_slot_a: u32, + /// Index of body B's velocity in the solver velocity buffer. + pub vel_slot_b: u32, + /// Collider A of the source manifold, which identifies it across steps with + /// `warmstart_collider_b` and `warmstart_subshape` (for the warmstart and color seeding). + pub warmstart_collider_a: u32, + /// Collider B of the source manifold. + pub warmstart_collider_b: u32, + /// [`IndexedManifold::subshape`] of the source manifold. /// /// [`IndexedManifold::subshape`]: crate::queries::IndexedManifold::subshape - #[cfg(feature = "dim3")] pub warmstart_subshape: u32, + /// Combined restitution coefficient. + pub restitution: f32, + pub _padding: [u32; 2], + /// The contact points (the first `len` are active). + pub points: [ContactPoint; MAX_CONSTRAINTS_PER_MANIFOLD], +} +/// A contact manifold between two rigid bodies, as the solver sees it. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +#[cfg(feature = "dim2")] +pub struct TwoBodyConstraint { + /// Contact normal from body A's perspective (points away from A). + pub dir_a: Vector, + /// Number of active contact points in this manifold. + pub len: u32, + /// Friction coefficient. + pub limit: f32, + /// Inverse mass of body A along each axis. + pub im_a: Vector, /// Inverse mass of body B along each axis. pub im_b: Vector, - #[cfg(feature = "dim3")] - pub _padding3: f32, - - /// Constraint Force Mixing (CFM) factor for regularization. - /// Softens the constraint: new_impulse = cfm_factor * (old_impulse - ...) - pub cfm_factor: f32, - - /// Friction coefficient (μ in Coulomb friction model). - /// Friction impulse magnitude limited to: |f_tangent| <= limit * f_normal - pub limit: f32, - - /// Index of body A in the solver arrays. + /// World-space inverse inertia of body A. + pub ii_a: f32, + /// World-space inverse inertia of body B. + pub ii_b: f32, + /// Index of body A (poses, mass properties). pub solver_body_a: u32, - - /// Index of body B in the solver arrays. + /// Index of body B. pub solver_body_b: u32, + /// Index of body A's velocity in the solver velocity buffer: `solver_body_a`, or one of + /// the sub-body slots when body A is split (see `mass_splitting`). + pub vel_slot_a: u32, + /// Index of body B's velocity in the solver velocity buffer. + pub vel_slot_b: u32, + /// Collider A of the source manifold, which identifies it across steps with + /// `warmstart_collider_b` and `warmstart_subshape` (for the warmstart and color seeding). + pub warmstart_collider_a: u32, + /// Collider B of the source manifold. + pub warmstart_collider_b: u32, + /// [`IndexedManifold::subshape`] of the source manifold. + /// + /// [`IndexedManifold::subshape`]: crate::queries::IndexedManifold::subshape + pub warmstart_subshape: u32, + /// Combined restitution coefficient. + pub restitution: f32, + pub _padding: [u32; 2], + /// The contact points (the first `len` are active). + pub points: [ContactPoint; MAX_CONSTRAINTS_PER_MANIFOLD], +} - /// Per-contact-point constraint data (up to MAX_CONSTRAINTS_PER_MANIFOLD). - pub elements: [TwoBodyConstraintElement; MAX_CONSTRAINTS_PER_MANIFOLD], - - /// Number of active contact points in this manifold (1-4 in 3D, 1-2 in 2D). +/// The topology and identity of a contact constraint, in contact order. +/// +/// Matching, mass splitting and graph coloring only read these; the solver data of the +/// constraints lives in color-ordered tiles (see `contact_tiles`). +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct ContactLink { + /// Number of active contact points: 0 for gap, inert and multibody-owned slots. pub len: u32, - /// Collider A of the source manifold (2D; appended, no spare padding lane). - #[cfg(feature = "dim2")] + /// Index of body A (poses, mass properties). + pub solver_body_a: u32, + /// Index of body B. + pub solver_body_b: u32, + /// Index of body A's velocity in the solver velocity buffer: `solver_body_a`, or one of + /// the sub-body slots when body A is split (see `mass_splitting`). + pub vel_slot_a: u32, + /// Index of body B's velocity in the solver velocity buffer. + pub vel_slot_b: u32, + /// Collider A of the source manifold, which identifies it across steps with + /// `warmstart_collider_b` and `warmstart_subshape` (for the warmstart and color seeding). pub warmstart_collider_a: u32, - /// Collider B of the source manifold (2D). - #[cfg(feature = "dim2")] + /// Collider B of the source manifold. pub warmstart_collider_b: u32, - /// [`IndexedManifold::subshape`] of the source manifold (2D). + /// [`IndexedManifold::subshape`] of the source manifold. /// /// [`IndexedManifold::subshape`]: crate::queries::IndexedManifold::subshape - #[cfg(feature = "dim2")] pub warmstart_subshape: u32, - #[cfg(feature = "dim3")] - pub _padding: [u32; 3], + /// Matching previous-frame constraint, or `u32::MAX`. + pub previous_constraint: u32, + /// Index of the matching previous-frame constraint in the previous frame's tiles. + pub previous_constraint_index: u32, + /// Whether the previous constraint's contacts are reused instead of recomputed (contact + /// recycling). + pub recycled: u32, + /// Factor of body A's inverse mass and inertia: its sub-body count when split, else 1. + pub mass_scale_a: f32, + /// Factor of body B's inverse mass and inertia. + pub mass_scale_b: f32, + /// Combined restitution coefficient. + pub restitution: f32, + /// The contact slot of this link, for its color-ordered copy (see `gpu_color_buckets_scatter`). + pub contact: u32, + pub _padding: u32, } -/// Constraint data for a single contact point. +/// Element offsets into the two disjoint halves of the recycling-state buffer. +/// The host swaps immutable uniforms each frame; no state data is copied. #[derive(Clone, Copy, Default)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] -pub struct TwoBodyConstraintElement { - /// Normal constraint: prevents penetration. - pub normal_part: TwoBodyConstraintNormalPart, - - /// Tangent constraint(s): models friction. - pub tangent_part: TwoBodyConstraintTangentPart, +pub struct ContactRecycleOffsets { + pub old_base: u32, + pub new_base: u32, + pub _padding0: u32, + pub _padding1: u32, } -/// Normal constraint data (non-penetration). -/// -/// Implements the constraint: C >= 0 (bodies cannot interpenetrate) -/// Solved as a unilateral constraint (impulse >= 0). +/// When the contacts of a [`TwoBodyConstraint`] were last computed, for contact recycling (see +/// `RbdSimParams::normalized_contact_recycle_distance`). #[derive(Clone, Copy, Default)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] -pub struct TwoBodyConstraintNormalPart { - /// Angular contribution for body A: r_a × normal - /// (In 2D: scalar cross product, in 3D: vector cross product) - pub torque_dir_a: AngVector, - #[cfg(feature = "dim3")] - pub _padding0: f32, - /// `torque_dir_a` multiplied by the inverse angular inertia tensor. - pub ii_torque_dir_a: AngVector, - #[cfg(feature = "dim3")] - pub _padding1: f32, - - /// Angular contribution for body B: r_b × normal - pub torque_dir_b: AngVector, - #[cfg(feature = "dim3")] - pub _padding2: f32, - /// `torque_dir_b` multiplied by the inverse angular inertia tensor. - pub ii_torque_dir_b: AngVector, - - /// Right-hand side: target relative velocity (includes bias for correction). - /// rhs = desired_velocity + bias_velocity - pub rhs: f32, - - /// Right-hand side without bias term (used in substep iterations). - /// rhs_wo_bias = desired_velocity (without position correction) - pub rhs_wo_bias: f32, - - /// Current impulse magnitude for this iteration. - /// Updated during solving, used for warmstarting next frame. - pub impulse: f32, - - /// Inverse effective mass: 1 / (m_eff) - /// where m_eff = projected mass along constraint direction - pub r: f32, - - /// Per-point softness: separated (speculative, `dist > 0`) points are solved - /// rigidly (`cfm_factor = 1`): perfectly inelastic touchdowns damp stack - /// rocking; a constraint-wide cfm under-stops each touchdown and pumps tall - /// stacks. Refreshed by `update_constraint` each substep. - pub cfm_factor: f32, +pub struct ContactRecycleState { + /// World pose of the first collider. + pub pose_a: Pose, + /// World pose of the second collider. + pub pose_b: Pose, + /// The collider pair. + pub colliders: UVec2, + /// See `IndexedManifold::recycle_extent`. + pub max_extent: f32, + /// The drift allowed before the contacts are computed again (0 if they can't be recycled). + pub max_drift: f32, } -/// Tangent constraint data (friction). -/// -/// Implements Coulomb friction: |f_tangent| <= μ * f_normal -/// Solved as a bilateral constraint with limits. +/// One contact point of a [`TwoBodyConstraint`]. #[derive(Clone, Copy, Default)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] -pub struct TwoBodyConstraintTangentPart { - /// Angular contributions for body A (one per tangent direction). - pub torque_dir_a: [Pad; SUB_LEN], - /// `torque_dir_a` multiplied by the inverse angular inertia tensor. - pub ii_torque_dir_a: [Pad; SUB_LEN], - - /// Angular contributions for body B (one per tangent direction). - pub torque_dir_b: [Pad; SUB_LEN], - /// `torque_dir_b` multiplied by the inverse angular inertia tensor. - pub ii_torque_dir_b: [Pad; SUB_LEN], - - /// Right-hand sides (one per tangent direction). - pub rhs: [f32; SUB_LEN], - - /// Right-hand sides without bias (one per tangent direction). - pub rhs_wo_bias: [f32; SUB_LEN], - - #[cfg(feature = "dim2")] - /// Current tangent impulse (2D: single scalar for one tangent direction). - pub impulse: [f32; 1], +#[cfg(feature = "dim3")] +pub struct ContactPoint { + /// The contact point relative to body A's center of mass, in world space (frozen at the + /// start of the step, like the Jacobians). + pub r_a: Vector, + /// Accumulated normal impulse. + pub normal_impulse: f32, + /// The contact point relative to body B's center of mass, in world space. + pub r_b: Vector, + /// Inverse of the effective mass along the normal. + pub normal_mass: f32, + /// The contact point in body A's (center-of-mass) frame, to track its current position. + pub local_pt_a: Vector, + /// Signed distance at the start of the step (negative = penetration). + pub dist: f32, + /// The contact point in body B's frame. + pub local_pt_b: Vector, + /// Target normal relative velocity (restitution). + pub normal_vel: f32, + /// Accumulated friction impulses (along `tangent_a` and `dir_a × tangent_a`). + pub tangent_impulse: Vec2, + /// Effective mass of the friction directions: `[k00, k11, 2 k01]` (not inverted). + pub tangent_k: [f32; 3], + pub _padding: [u32; 3], +} - #[cfg(feature = "dim2")] - /// Inverse effective mass (2D: single value). - pub r: [f32; 1], +/// One contact point of a [`TwoBodyConstraint`]. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +#[cfg(feature = "dim2")] +pub struct ContactPoint { + /// The contact point relative to body A's center of mass, in world space. + pub r_a: Vector, + /// The contact point relative to body B's center of mass, in world space. + pub r_b: Vector, + /// The contact point in body A's (center-of-mass) frame. + pub local_pt_a: Vector, + /// The contact point in body B's frame. + pub local_pt_b: Vector, + /// Accumulated normal impulse. + pub normal_impulse: f32, + /// Inverse of the effective mass along the normal. + pub normal_mass: f32, + /// Signed distance at the start of the step (negative = penetration). + pub dist: f32, + /// Target normal relative velocity (restitution). + pub normal_vel: f32, + /// Accumulated friction impulse. + pub tangent_impulse: [f32; 1], + /// Inverse of the effective mass along the tangent. + pub tangent_mass: f32, + pub _padding: [u32; 2], +} - #[cfg(feature = "dim3")] - /// Current tangent impulses (3D: vec2 for two tangent directions). - pub impulse: Vec2, +#[inline(always)] +fn inv(x: f32) -> f32 { + if x == 0.0 { 0.0 } else { 1.0 / x } +} - #[cfg(feature = "dim3")] - /// Inverse effective mass components (3D: 3 values for 2x2 mass matrix). - /// `r[0]` = r_00, `r[1]` = r_11, `r[2]` = r_01 (symmetric, so r_10 = r_01) - pub r: [f32; 3], - #[cfg(feature = "dim3")] - pub _padding0: [u32; 3], +impl TwoBodyConstraint { + /// Body A's inverse inertia times `v`. + #[inline(always)] + pub fn ii_a_mul(&self, v: AngVector) -> AngVector { + #[cfg(feature = "dim2")] + return self.ii_a * v; + #[cfg(feature = "dim3")] + return self.ii_a.mul(v); + } + + /// Body B's inverse inertia times `v`. + #[inline(always)] + pub fn ii_b_mul(&self, v: AngVector) -> AngVector { + #[cfg(feature = "dim2")] + return self.ii_b * v; + #[cfg(feature = "dim3")] + return self.ii_b.mul(v); + } + + /// The friction directions. + #[inline(always)] + pub fn tangents(&self) -> [Vector; SUB_LEN] { + #[cfg(feature = "dim2")] + return [Vector::new(-self.dir_a.y, self.dir_a.x)]; + #[cfg(feature = "dim3")] + return [self.tangent_a, self.dir_a.cross(self.tangent_a)]; + } + + /// Computes the effective masses of every contact point from the lever arms, inverse + /// masses and inertias. + #[inline(always)] + pub fn compute_effective_masses(&mut self) { + let dir = self.dir_a; + let imsum = self.im_a + self.im_b; + let tangents = self.tangents(); + + for k in 0..self.len as usize { + let r_a = self.points.at(k).r_a; + let r_b = self.points.at(k).r_b; + + let ta = gcross(r_a, dir); + let tb = gcross(r_b, -dir); + let k_n = + dir.dot(imsum * dir) + gdot(self.ii_a_mul(ta), ta) + gdot(self.ii_b_mul(tb), tb); + self.points.at_mut(k).normal_mass = inv(k_n); + + #[cfg(feature = "dim2")] + { + let t = tangents.read(0); + let ta = gcross(r_a, t); + let tb = gcross(r_b, -t); + let k_t = + t.dot(imsum * t) + gdot(self.ii_a_mul(ta), ta) + gdot(self.ii_b_mul(tb), tb); + self.points.at_mut(k).tangent_mass = inv(k_t); + } + + #[cfg(feature = "dim3")] + { + let t0 = tangents.read(0); + let t1 = tangents.read(1); + let ta0 = gcross(r_a, t0); + let tb0 = gcross(r_b, -t0); + let ta1 = gcross(r_a, t1); + let tb1 = gcross(r_b, -t1); + let iia1 = self.ii_a_mul(ta1); + let iib1 = self.ii_b_mul(tb1); + let k00 = + t0.dot(imsum * t0) + self.ii_a_mul(ta0).dot(ta0) + self.ii_b_mul(tb0).dot(tb0); + let k11 = t1.dot(imsum * t1) + iia1.dot(ta1) + iib1.dot(tb1); + let k01 = 2.0 * (ta0.dot(iia1) + tb0.dot(iib1)); + self.points.at_mut(k).tangent_k = [k00, k11, k01]; + } + } + } } diff --git a/src_rbd_shaders/dynamics/contact_links.rs b/src_rbd_shaders/dynamics/contact_links.rs new file mode 100644 index 00000000..6e0bdb6d --- /dev/null +++ b/src_rbd_shaders/dynamics/contact_links.rs @@ -0,0 +1,230 @@ +//! The contact-order links of the contact constraints: their topology for mass splitting and +//! coloring, and their match with the previous frame's constraints (warmstarting and contact +//! recycling). +//! +//! A contact is matched with the previous constraint of the same body pair, collider pair and +//! sub-shape. A pair that barely moved since its contacts were computed keeps them instead of +//! recomputing them (rapier's contact recycling). + +use crate::broad_phase::ContactPlan; +use crate::queries::IndexedManifold; +use crate::{Pose, Rotation}; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +#[allow(unused_imports)] // Needed on the GPU for `sqrt`. +use khal_std::num_traits::Float; + +use super::RbdSimParams; +use super::constraint::{ContactLink, ContactRecycleOffsets, ContactRecycleState}; + +/// Initializes the link of every contact. Manifolds touching a multibody link belong to the +/// multibody contact solver: their link stays inactive, like gap and inert slots. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_init_contact_links( + #[spirv(global_invocation_id)] gid: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &[IndexedManifold], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links: &mut [ContactLink], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_is_multibody: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] plan: &ContactPlan, +) { + let i = gid.x as usize; + if gid.x >= plan.bound { + return; + } + let im = contacts.at(i); + let multibody = + body_is_multibody.read(im.bodies.x as usize) | body_is_multibody.read(im.bodies.y as usize); + let len = if multibody != 0 { 0 } else { im.contact.len }; + links.write( + i, + ContactLink { + len, + solver_body_a: im.bodies.x, + solver_body_b: im.bodies.y, + vel_slot_a: im.bodies.x, + vel_slot_b: im.bodies.y, + warmstart_collider_a: im.colliders.x, + warmstart_collider_b: im.colliders.y, + warmstart_subshape: im.subshape, + previous_constraint: u32::MAX, + previous_constraint_index: u32::MAX, + recycled: 0, + mass_scale_a: 1.0, + mass_scale_b: 1.0, + restitution: im.restitution, + contact: i as u32, + _padding: 0, + }, + ); +} + +/// Matches every active link with the previous frame's, decides whether its contacts are +/// recycled, and records the recycling state of its contacts. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_match_contact_links( + #[spirv(global_invocation_id)] gid: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &[IndexedManifold], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links: &mut [ContactLink], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] old_counts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] old_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] old_links: &[ContactLink], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] old_constraint_indices: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] states: &mut [ContactRecycleState], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] collider_poses: &[Pose], + #[spirv(uniform, descriptor_set = 0, binding = 8)] plan: &ContactPlan, + #[spirv(uniform, descriptor_set = 0, binding = 9)] params: &RbdSimParams, + #[spirv(uniform, descriptor_set = 0, binding = 10)] recycle_offsets: &ContactRecycleOffsets, +) { + let i = gid.x as usize; + if gid.x >= plan.bound { + return; + } + let link = links.read(i); + // Gap slots may hold stale body ids: they must not be dereferenced. + if link.len == 0 { + return; + } + let im = contacts.at(i); + let previous = find_previous_constraint(old_counts, old_ids, old_links, &link); + let pose_a = collider_poses.read(im.colliders.x as usize); + let pose_b = collider_poses.read(im.colliders.y as usize); + let new_state = recycle_offsets.new_base as usize + i; + let mut recycled = 0; + let mut previous_constraint_index = u32::MAX; + if previous != u32::MAX { + previous_constraint_index = old_constraint_indices.read(previous as usize); + let old_state = states.read(recycle_offsets.old_base as usize + previous as usize); + if can_recycle(&old_state, pose_a, pose_b) { + recycled = 1; + states.write(new_state, old_state); + } + } + if recycled == 0 { + states.write( + new_state, + ContactRecycleState { + pose_a, + pose_b, + colliders: im.colliders, + max_extent: im.recycle_extent, + max_drift: if im.recycle_extent >= 0.0 { + params.contact_recycle_distance() + } else { + 0.0 + }, + }, + ); + } + let dst = links.at_mut(i); + dst.previous_constraint = previous; + dst.previous_constraint_index = previous_constraint_index; + dst.recycled = recycled; +} + +/// The previous constraint of the same body pair, collider pair and sub-shape, searched in +/// the shorter of the two bodies' previous adjacency lists, or `u32::MAX`. +#[inline(always)] +fn find_previous_constraint( + counts: &[u32], + ids: &[u32], + old_links: &[ContactLink], + new: &ContactLink, +) -> u32 { + let a = new.solver_body_a as usize; + let b = new.solver_body_b as usize; + let first_a = if a == 0 { 0 } else { counts.read(a - 1) }; + let last_a = counts.read(a); + let first_b = if b == 0 { 0 } else { counts.read(b - 1) }; + let last_b = counts.read(b); + let len_a = last_a - first_a; + let len_b = last_b - first_b; + // Static bodies have empty lists: search the other one. + let (first, last) = if len_b == 0 || (len_a != 0 && len_a < len_b) { + (first_a, last_a) + } else { + (first_b, last_b) + }; + for entry in first..last { + let cid = ids.read(entry as usize); + let old = old_links.at(cid as usize); + if (old.solver_body_a == new.solver_body_a) + & (old.solver_body_b == new.solver_body_b) + & (old.warmstart_collider_a == new.warmstart_collider_a) + & (old.warmstart_collider_b == new.warmstart_collider_b) + & (old.warmstart_subshape == new.warmstart_subshape) + { + return cid; + } + } + u32::MAX +} + +/// Seeds the coloring with the colors of the matched previous constraints. Runs after the +/// topo-gc reset, before the coloring iterations. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_seed_colors_from_warmstart( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] links: &[ContactLink], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] old_constraints_colors: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints_colors: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] colored: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] contact_plan: &ContactPlan, +) { + let i = invocation_id.x as usize; + if invocation_id.x < contact_plan.bound { + let link = links.at(i); + if link.len != 0 && link.previous_constraint != u32::MAX { + let old_color = old_constraints_colors.read(link.previous_constraint as usize); + // The conflict pass validates every seed, including stale colors. + if old_color > 0 && old_color < 64 { + constraints_colors.write(i, old_color); + colored.write(i, 1); + } + } + } +} + +/// Whether a pair whose contacts were computed at `state` can keep them at the collider poses +/// `pose_a` and `pose_b` (rapier's `relative_pose_drift` and `relative_rot_cos` tests). +#[inline(always)] +fn can_recycle(state: &ContactRecycleState, pose_a: Pose, pose_b: Pose) -> bool { + let pos12_old = state.pose_a.inverse() * state.pose_b; + let pos12 = pose_a.inverse() * pose_b; + let trans = (pos12.translation - pos12_old.translation).length(); + let drot = pos12.rotation * pos12_old.rotation.inverse(); + // Bound on the rotation's displacement of a point within `max_extent` of the origin. + #[cfg(feature = "dim3")] + let chord = 2.0 * glamx::Vec3::new(drot.x, drot.y, drot.z).length() * state.max_extent; + #[cfg(feature = "dim2")] + let chord = { + let denom = 2.0 * (1.0 + drot.re); + let half_sin = if denom > 1.0e-6 { + drot.im.abs() / denom.sqrt() + } else { + 1.0 + }; + 2.0 * half_sin * state.max_extent + }; + // The frozen normal and lever arms also bound each collider's own rotation. + let rot_cos = rot_cos(state.pose_a.rotation, pose_a.rotation) + .min(rot_cos(state.pose_b.rotation, pose_b.rotation)); + state.max_drift > 0.0 && trans + chord <= state.max_drift && rot_cos > 0.98 +} + +/// The cosine of the angle between two rotations. +#[inline(always)] +fn rot_cos(a: Rotation, b: Rotation) -> f32 { + #[cfg(feature = "dim3")] + { + let c = a.dot(b); + 2.0 * c * c - 1.0 + } + #[cfg(feature = "dim2")] + { + a.re * b.re + a.im * b.im + } +} diff --git a/src_rbd_shaders/dynamics/contact_tiles.rs b/src_rbd_shaders/dynamics/contact_tiles.rs new file mode 100644 index 00000000..e29895ab --- /dev/null +++ b/src_rbd_shaders/dynamics/contact_tiles.rs @@ -0,0 +1,309 @@ +//! Typed tiles of contact constraints, the solver's storage for them. +//! +//! Constraints are sorted by color, so every solver iteration walks a contiguous range of them. +//! A tile holds [`TILE_LEN`] consecutive constraints field by field: each field is an array +//! with one entry per constraint, so neighboring lanes read neighboring memory. + +use super::Velocity; +use super::solver_utils::ContactSolverData; +use khal_std::index::MaybeIndexUnchecked; + +#[cfg(feature = "dim2")] +use khal_std::glamx::Vec2; +#[cfg(feature = "dim3")] +use {super::SymInertia, khal_std::glamx::Vec3}; + +/// Number of neighboring constraints stored together in a tile. +pub const TILE_LEN: usize = 32; +/// First sparse color solved by a single workgroup, with barriers between colors. +pub const TAIL_COLOR: u32 = 24; + +/// The tile holding the constraint at `index`, and its lane in it. +#[inline(always)] +pub fn tile_lane(index: usize) -> (usize, usize) { + (index / TILE_LEN, index % TILE_LEN) +} + +/// The velocity changes of the constraints' two bodies applying their accumulated impulses +/// (their warmstart), each tagged with its body index in `padding1`. A constraint's pair is +/// contiguous, for the per-body warmstart gather. +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct WarmstartTile { + pub bodies: [[Velocity; 2]; TILE_LEN], +} + +/// A vector and a scalar of a contact point. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct VecScalar { + pub vec: Vec3, + pub scalar: f32, +} + +/// Contact normal and number of active points. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct NormalLen { + pub normal: Vec3, + pub len: u32, +} + +/// First friction direction and friction coefficient. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct TangentLimit { + pub tangent: Vec3, + pub limit: f32, +} + +/// Inverse mass of a body along each axis, and the body's index. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct MassBody { + pub inv_mass: Vec3, + pub body: u32, +} + +/// Inverse inertia terms, and a velocity slot. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct InertiaSlot { + pub inertia: Vec3, + pub vel_slot: u32, +} + +/// The constraints' common solver header: normal, friction direction, masses and velocity +/// slots. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct HeaderTile { + pub normal: [NormalLen; TILE_LEN], + pub tangent: [TangentLimit; TILE_LEN], + pub body_a: [MassBody; TILE_LEN], + pub body_b: [MassBody; TILE_LEN], + /// Body A's inverse inertia diagonal, and its velocity slot. + pub inertia_diag_a: [InertiaSlot; TILE_LEN], + /// Body A's inverse inertia off-diagonal terms, and body B's velocity slot. + pub inertia_off_a: [InertiaSlot; TILE_LEN], + /// Body B's inverse inertia diagonal; the twist model stores its twist mass alongside. + pub inertia_diag_b: [VecScalar; TILE_LEN], + /// Body B's inverse inertia off-diagonal terms; the twist model stores its twist impulse + /// alongside. + pub inertia_off_b: [VecScalar; TILE_LEN], +} + +#[cfg(feature = "dim3")] +impl HeaderTile { + #[inline(always)] + pub(crate) fn read(&self, lane: usize) -> ContactSolverData { + let normal = self.normal.read(lane); + let tangent = self.tangent.read(lane); + let body_a = self.body_a.read(lane); + let body_b = self.body_b.read(lane); + let diag_a = self.inertia_diag_a.read(lane); + let off_a = self.inertia_off_a.read(lane); + ContactSolverData { + dir_a: normal.normal, + len: normal.len, + tangent_a: tangent.tangent, + limit: tangent.limit, + im_a: body_a.inv_mass, + solver_body_a: body_a.body, + im_b: body_b.inv_mass, + solver_body_b: body_b.body, + ii_a: SymInertia { + diag: diag_a.inertia, + off: off_a.inertia, + _padding0: 0, + _padding1: 0, + }, + ii_b: SymInertia { + diag: self.inertia_diag_b.read(lane).vec, + off: self.inertia_off_b.read(lane).vec, + _padding0: 0, + _padding1: 0, + }, + vel_slot_a: diag_a.vel_slot, + vel_slot_b: off_a.vel_slot, + } + } + + /// Writes the header of `lane`, with the model-specific scalars stored next to body B's + /// inertia. + #[inline(always)] + pub(crate) fn write(&mut self, lane: usize, h: &ContactSolverData, scalars_b: [f32; 2]) { + self.normal.write( + lane, + NormalLen { + normal: h.dir_a, + len: h.len, + }, + ); + self.tangent.write( + lane, + TangentLimit { + tangent: h.tangent_a, + limit: h.limit, + }, + ); + self.body_a.write( + lane, + MassBody { + inv_mass: h.im_a, + body: h.solver_body_a, + }, + ); + self.body_b.write( + lane, + MassBody { + inv_mass: h.im_b, + body: h.solver_body_b, + }, + ); + self.inertia_diag_a.write( + lane, + InertiaSlot { + inertia: h.ii_a.diag, + vel_slot: h.vel_slot_a, + }, + ); + self.inertia_off_a.write( + lane, + InertiaSlot { + inertia: h.ii_a.off, + vel_slot: h.vel_slot_b, + }, + ); + self.inertia_diag_b.write( + lane, + VecScalar { + vec: h.ii_b.diag, + scalar: scalars_b[0], + }, + ); + self.inertia_off_b.write( + lane, + VecScalar { + vec: h.ii_b.off, + scalar: scalars_b[1], + }, + ); + } +} + +/// Contact normal, number of active points and friction coefficient. +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct NormalLimit { + pub normal: Vec2, + pub len: u32, + pub limit: f32, +} + +/// Inverse mass of a body along each axis, its inverse inertia, and its index. +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct MassBody { + pub inv_mass: Vec2, + pub inv_inertia: f32, + pub body: u32, +} + +/// The velocity slots of the two bodies. +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct VelSlots { + pub a: u32, + pub b: u32, +} + +/// The constraints' common solver header: normal, masses and velocity slots. +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct HeaderTile { + pub normal: [NormalLimit; TILE_LEN], + pub body_a: [MassBody; TILE_LEN], + pub body_b: [MassBody; TILE_LEN], + pub vel_slots: [VelSlots; TILE_LEN], +} + +#[cfg(feature = "dim2")] +impl HeaderTile { + #[inline(always)] + pub(crate) fn read(&self, lane: usize) -> ContactSolverData { + let normal = self.normal.read(lane); + let body_a = self.body_a.read(lane); + let body_b = self.body_b.read(lane); + let slots = self.vel_slots.read(lane); + ContactSolverData { + dir_a: normal.normal, + len: normal.len, + limit: normal.limit, + im_a: body_a.inv_mass, + solver_body_a: body_a.body, + im_b: body_b.inv_mass, + solver_body_b: body_b.body, + ii_a: body_a.inv_inertia, + ii_b: body_b.inv_inertia, + vel_slot_a: slots.a, + vel_slot_b: slots.b, + } + } + + #[inline(always)] + pub(crate) fn write(&mut self, lane: usize, h: &ContactSolverData) { + self.normal.write( + lane, + NormalLimit { + normal: h.dir_a, + len: h.len, + limit: h.limit, + }, + ); + self.body_a.write( + lane, + MassBody { + inv_mass: h.im_a, + inv_inertia: h.ii_a, + body: h.solver_body_a, + }, + ); + self.body_b.write( + lane, + MassBody { + inv_mass: h.im_b, + inv_inertia: h.ii_b, + body: h.solver_body_b, + }, + ); + self.vel_slots.write( + lane, + VelSlots { + a: h.vel_slot_a, + b: h.vel_slot_b, + }, + ); + } +} diff --git a/src_rbd_shaders/dynamics/coulomb_tiles.rs b/src_rbd_shaders/dynamics/coulomb_tiles.rs new file mode 100644 index 00000000..b169ef38 --- /dev/null +++ b/src_rbd_shaders/dynamics/coulomb_tiles.rs @@ -0,0 +1,514 @@ +//! Coulomb contact constraints (friction at each contact point) in typed tiles, and their +//! kernels. +use super::constraint::{ContactPoint, TwoBodyConstraint}; +use super::contact_tiles::{HeaderTile, TILE_LEN, WarmstartTile, tile_lane}; +use super::solver_utils::{ContactPointAccess, ContactSolverData}; +use super::{RbdSimParams, Velocity}; +use crate::utils::Slice; +use crate::{Pose, Vector, gcross}; +use khal_std::index::MaybeIndexUnchecked; + +#[cfg(feature = "dim2")] +use khal_std::glamx::Vec2; +#[cfg(feature = "dim3")] +use {super::contact_tiles::VecScalar, khal_std::glamx::Vec2}; + +/// Accumulated friction impulses, and the first two friction effective-mass terms. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct FrictionRows { + pub impulse: Vec2, + pub k: Vec2, +} + +/// One contact point of the constraints of a [`CoulombTile`]. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct CoulombPointTile { + /// Lever arm from body A's center of mass, and the accumulated normal impulse. + pub arm_a: [VecScalar; TILE_LEN], + /// Lever arm from body B's center of mass, and the inverse effective normal mass. + pub arm_b: [VecScalar; TILE_LEN], + /// The point in body A's frame, and the signed distance at the start of the step. + pub local_a: [VecScalar; TILE_LEN], + /// The point in body B's frame, and the target normal velocity. + pub local_b: [VecScalar; TILE_LEN], + pub friction: [FrictionRows; TILE_LEN], +} + +/// [`TILE_LEN`] consecutive Coulomb constraints. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct CoulombTile { + pub header: HeaderTile, + /// Each point's friction coupling term (`ContactPoint::tangent_k[2]`). + pub coupling: [[f32; 4]; TILE_LEN], + pub points: [CoulombPointTile; 4], + pub warmstart: WarmstartTile, +} + +/// Two vectors of a contact point, one per body. +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct VectorPair { + pub a: Vec2, + pub b: Vec2, +} + +/// The normal row of a contact point. +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct NormalRow { + /// Accumulated impulse. + pub impulse: f32, + /// Inverse of the effective mass. + pub mass: f32, + /// Signed distance at the start of the step. + pub dist: f32, + /// Target velocity (restitution). + pub vel: f32, +} + +/// The friction row of a contact point. +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct TangentRow { + /// Accumulated impulse. + pub impulse: f32, + /// Inverse of the effective mass. + pub mass: f32, +} + +/// One contact point of the constraints of a [`CoulombTile`]. +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct CoulombPointTile { + /// Lever arms from the centers of mass of bodies A and B. + pub arms: [VectorPair; TILE_LEN], + /// The point in the frames of bodies A and B. + pub anchors: [VectorPair; TILE_LEN], + pub normal: [NormalRow; TILE_LEN], + pub tangent: [TangentRow; TILE_LEN], +} + +/// [`TILE_LEN`] consecutive Coulomb constraints. +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct CoulombTile { + pub header: HeaderTile, + pub points: [CoulombPointTile; 2], + pub warmstart: WarmstartTile, +} + +impl CoulombTile { + #[cfg(feature = "dim3")] + #[inline(always)] + fn point(&self, lane: usize, k: usize) -> ContactPoint { + let p = self.points.at(k); + let arm_a = p.arm_a.read(lane); + let arm_b = p.arm_b.read(lane); + let local_a = p.local_a.read(lane); + let local_b = p.local_b.read(lane); + let friction = p.friction.read(lane); + ContactPoint { + r_a: arm_a.vec, + normal_impulse: arm_a.scalar, + r_b: arm_b.vec, + normal_mass: arm_b.scalar, + local_pt_a: local_a.vec, + dist: local_a.scalar, + local_pt_b: local_b.vec, + normal_vel: local_b.scalar, + tangent_impulse: friction.impulse, + tangent_k: [friction.k.x, friction.k.y, self.coupling.at(lane).read(k)], + _padding: [0; 3], + } + } + + #[cfg(feature = "dim2")] + #[inline(always)] + fn point(&self, lane: usize, k: usize) -> ContactPoint { + let p = self.points.at(k); + let arms = p.arms.read(lane); + let anchors = p.anchors.read(lane); + let normal = p.normal.read(lane); + let tangent = p.tangent.read(lane); + ContactPoint { + r_a: arms.a, + r_b: arms.b, + local_pt_a: anchors.a, + local_pt_b: anchors.b, + normal_impulse: normal.impulse, + normal_mass: normal.mass, + dist: normal.dist, + normal_vel: normal.vel, + tangent_impulse: [tangent.impulse], + tangent_mass: tangent.mass, + _padding: [0; 2], + } + } + + #[cfg(feature = "dim3")] + #[inline(always)] + fn write_point(&mut self, lane: usize, k: usize, p: &ContactPoint) -> f32 { + let dst = self.points.at_mut(k); + dst.arm_a.write( + lane, + VecScalar { + vec: p.r_a, + scalar: p.normal_impulse, + }, + ); + dst.arm_b.write( + lane, + VecScalar { + vec: p.r_b, + scalar: p.normal_mass, + }, + ); + dst.local_a.write( + lane, + VecScalar { + vec: p.local_pt_a, + scalar: p.dist, + }, + ); + dst.local_b.write( + lane, + VecScalar { + vec: p.local_pt_b, + scalar: p.normal_vel, + }, + ); + dst.friction.write( + lane, + FrictionRows { + impulse: p.tangent_impulse, + k: Vec2::new(p.tangent_k[0], p.tangent_k[1]), + }, + ); + p.tangent_k[2] + } + + #[cfg(feature = "dim2")] + #[inline(always)] + fn write_point(&mut self, lane: usize, k: usize, p: &ContactPoint) { + let dst = self.points.at_mut(k); + dst.arms.write(lane, VectorPair { a: p.r_a, b: p.r_b }); + dst.anchors.write( + lane, + VectorPair { + a: p.local_pt_a, + b: p.local_pt_b, + }, + ); + dst.normal.write( + lane, + NormalRow { + impulse: p.normal_impulse, + mass: p.normal_mass, + dist: p.dist, + vel: p.normal_vel, + }, + ); + dst.tangent.write( + lane, + TangentRow { + impulse: p.tangent_impulse[0], + mass: p.tangent_mass, + }, + ); + } + + /// Writes the constraint of `lane`. + #[inline(always)] + pub fn write_constraint(&mut self, lane: usize, c: &TwoBodyConstraint) { + let h = ContactSolverData::from_constraint(c); + #[cfg(feature = "dim2")] + { + self.header.write(lane, &h); + crate::dynamics::solver_utils::for_contact_point!(k, c.len, { + self.write_point(lane, k, c.points.at(k)); + }); + } + #[cfg(feature = "dim3")] + { + self.header.write(lane, &h, [0.0; 2]); + let mut coupling = [0.0; 4]; + crate::dynamics::solver_utils::for_contact_point!(k, c.len, { + coupling.write(k, self.write_point(lane, k, c.points.at(k))); + }); + self.coupling.write(lane, coupling); + } + } + + /// The constraint of `lane`. Its identity fields are left to their defaults: they live in the + /// contact links. + #[inline(always)] + pub fn constraint(&self, lane: usize) -> TwoBodyConstraint { + let h = self.header.read(lane); + let mut c = TwoBodyConstraint { + dir_a: h.dir_a, + len: h.len, + limit: h.limit, + im_a: h.im_a, + im_b: h.im_b, + ii_a: h.ii_a, + ii_b: h.ii_b, + solver_body_a: h.solver_body_a, + solver_body_b: h.solver_body_b, + vel_slot_a: h.vel_slot_a, + vel_slot_b: h.vel_slot_b, + ..Default::default() + }; + #[cfg(feature = "dim3")] + { + c.tangent_a = h.tangent_a; + } + crate::dynamics::solver_utils::for_contact_point!(k, h.len, { + c.points.write(k, self.point(lane, k)); + }); + c + } +} + +/// The tile type of this model's kernels (see `tile_kernels.rs`). +pub type ContactTile = CoulombTile; + +#[inline(always)] +fn header(tiles: &[CoulombTile], index: usize) -> ContactSolverData { + let (tile, lane) = tile_lane(index); + tiles.at(tile).header.read(lane) +} + +#[inline(always)] +fn read_constraint(tiles: &[CoulombTile], index: usize) -> TwoBodyConstraint { + let (tile, lane) = tile_lane(index); + tiles.at(tile).constraint(lane) +} + +#[inline(always)] +fn write_constraint(tiles: &mut [CoulombTile], index: usize, c: &TwoBodyConstraint) { + let (tile, lane) = tile_lane(index); + tiles.at_mut(tile).write_constraint(lane, c); +} + +/// The velocity changes of the bodies of `c` applying its accumulated impulses, tagged with +/// their body index (for the warmstart gather). +#[inline(always)] +fn constraint_warmstart(c: &TwoBodyConstraint) -> (Velocity, Velocity) { + let (mut a, mut b) = (Velocity::default(), Velocity::default()); + c.warmstart_constraint(&mut a, &mut b); + a.padding1 = c.solver_body_a; + b.padding1 = c.solver_body_b; + (a, b) +} + +struct TilePoints<'a> { + tiles: &'a mut [CoulombTile], + index: usize, +} +impl ContactPointAccess for TilePoints<'_> { + #[inline(always)] + fn read_point(&self, k: usize) -> ContactPoint { + let (tile, lane) = tile_lane(self.index); + self.tiles.at(tile).point(lane, k) + } + #[cfg(feature = "dim3")] + #[inline(always)] + fn write_normal_impulse(&mut self, k: usize, value: f32) { + let (tile, lane) = tile_lane(self.index); + self.tiles + .at_mut(tile) + .points + .at_mut(k) + .arm_a + .at_mut(lane) + .scalar = value; + } + #[cfg(feature = "dim3")] + #[inline(always)] + fn write_tangent_impulse(&mut self, k: usize, value: Vec2) { + let (tile, lane) = tile_lane(self.index); + let point = self.tiles.at_mut(tile).points.at_mut(k); + point.friction.at_mut(lane).impulse = value; + } + #[cfg(feature = "dim2")] + #[inline(always)] + fn write_normal_impulse(&mut self, k: usize, value: f32) { + let (tile, lane) = tile_lane(self.index); + self.tiles + .at_mut(tile) + .points + .at_mut(k) + .normal + .at_mut(lane) + .impulse = value; + } + #[cfg(feature = "dim2")] + #[inline(always)] + fn write_tangent_impulse(&mut self, k: usize, value: [f32; 1]) { + let (tile, lane) = tile_lane(self.index); + self.tiles + .at_mut(tile) + .points + .at_mut(k) + .tangent + .at_mut(lane) + .impulse = value[0]; + } +} + +#[inline(always)] +#[allow(clippy::too_many_arguments)] +fn solve_constraint( + h: &ContactSolverData, + tiles: &mut [CoulombTile], + index: usize, + poses: &Slice, + params: &RbdSimParams, + a: &mut Velocity, + b: &mut Velocity, + bias: bool, + friction: bool, +) { + h.solve_points( + &mut TilePoints { tiles, index }, + poses, + params, + a, + b, + bias, + friction, + ); +} + +/// The velocity changes of the bodies of the constraint at `index` applying its accumulated +/// impulses, tagged with their body index (for the warmstart gather). +#[inline(always)] +fn tile_warmstart( + tiles: &[CoulombTile], + index: usize, + h: &ContactSolverData, +) -> (Velocity, Velocity) { + let (tile, lane) = tile_lane(index); + let tile = tiles.at(tile); + let tangents = h.tangents(); + let mut force = Vector::ZERO; + let mut torque_a = crate::AngVector::default(); + let mut torque_b = crate::AngVector::default(); + crate::dynamics::solver_utils::for_contact_point!(k, h.len as usize, { + let p = tile.point(lane, k); + #[cfg(feature = "dim2")] + let f = h.dir_a * p.normal_impulse + tangents.read(0) * p.tangent_impulse[0]; + #[cfg(feature = "dim3")] + let f = h.dir_a * p.normal_impulse + + tangents.read(0) * p.tangent_impulse.x + + tangents.read(1) * p.tangent_impulse.y; + force += f; + torque_a += gcross(p.r_a, f); + torque_b += gcross(p.r_b, f); + }); + let mut a = Velocity::default(); + let mut b = Velocity::default(); + a.linear += h.im_a * force; + a.angular += h.ii_a_mul(torque_a); + b.linear -= h.im_b * force; + b.angular -= h.ii_b_mul(torque_b); + a.padding1 = h.solver_body_a; + b.padding1 = h.solver_body_b; + (a, b) +} + +// The kernels are shared by both friction models, each instantiating them with its own tiles. +#[allow(clippy::duplicate_mod)] +#[path = "tile_kernels.rs"] +mod kernels; +pub use kernels::*; + +/// Asserts that solving `initial` from a tile gives `expected` and its body velocities, for +/// constraints at the start and the end of tiles. +#[cfg(all(test, not(target_arch_is_gpu)))] +#[allow(clippy::too_many_arguments)] +pub(crate) fn assert_tile_equivalent( + initial: TwoBodyConstraint, + poses: &Slice, + params: &RbdSimParams, + initial_a: Velocity, + initial_b: Velocity, + use_bias: bool, + solve_friction: bool, + expected: &TwoBodyConstraint, + expected_a: Velocity, + expected_b: Velocity, +) { + for index in [0usize, 31, 32, 63] { + let mut tiles = vec![CoulombTile::default(); 2]; + write_constraint(&mut tiles, index, &initial); + let h = header(&tiles, index); + let (mut a, mut b) = (initial_a, initial_b); + solve_constraint( + &h, + &mut tiles, + index, + poses, + params, + &mut a, + &mut b, + use_bias, + solve_friction, + ); + let solved = read_constraint(&tiles, index); + crate::dynamics::solver_utils::for_contact_point!(k, expected.len, { + assert_eq!( + solved.points[k].normal_impulse, + expected.points[k].normal_impulse + ); + assert_eq!( + solved.points[k].tangent_impulse, + expected.points[k].tangent_impulse + ); + }); + assert_eq!( + (a.linear, a.angular, b.linear, b.angular), + ( + expected_a.linear, + expected_a.angular, + expected_b.linear, + expected_b.angular + ) + ); + let (cached_a, cached_b) = tile_warmstart(&tiles, index, &h); + let (mut warm_a, mut warm_b) = (Velocity::default(), Velocity::default()); + expected.warmstart_constraint(&mut warm_a, &mut warm_b); + assert_eq!( + ( + cached_a.linear, + cached_a.angular, + cached_b.linear, + cached_b.angular + ), + (warm_a.linear, warm_a.angular, warm_b.linear, warm_b.angular) + ); + } +} + +#[cfg(all(test, not(target_arch_is_gpu)))] +#[path = "../tests/contact_constraints.rs"] +mod tests; diff --git a/src_rbd_shaders/dynamics/mass_splitting.rs b/src_rbd_shaders/dynamics/mass_splitting.rs new file mode 100644 index 00000000..0d1c8288 --- /dev/null +++ b/src_rbd_shaders/dynamics/mass_splitting.rs @@ -0,0 +1,289 @@ +//! Mass splitting for bodies with many contacts ("hubs"). +//! +//! Graph coloring assigns different colors to constraints sharing a body, so a body touching +//! `n` others needs at least `n` colors. A heavy ball ploughing through a brick wall touches +//! hundreds of bricks at once, which blows up the color count (and every colored dispatch). +//! +//! Instead, a hub body is split into one sub-body per contact (Tonge et al. 2012, "Mass +//! splitting for jitter-free parallel rigid body simulation"): each sub-body has `1/n`-th of +//! the mass and inertia, owns a private velocity slot past the regular bodies in the solver +//! velocity buffer, and its constraints are colored ignoring the hub. Before each iteration the hub +//! velocity is scattered to its sub-bodies, and after the iteration their velocities are averaged +//! back into the hub, which applies the sum of every contact impulse to the real body. + +use crate::utils::{BatchIndices, Slice, SliceMut}; +use crate::{AngVector, Vector}; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::iter::StepRng; +use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::{atomic_add_u32, workgroup_memory_barrier_with_group_sync}; + +use super::ContactLink; +use super::body::{Velocity, WorldMassProperties}; + +const WORKGROUP_SIZE: u32 = 64; + +/// A body is split when more constraints than this touch it. +pub const HUB_MIN_CONSTRAINTS: u32 = 32; +/// Marks a body that isn't split in `hub_first_slot`. +pub const NOT_A_HUB: u32 = u32::MAX; + +/// The counters shared by the mass splitting kernels. +#[derive(Copy, Clone, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct HubCounts { + /// Number of hubs picked this step, possibly past the hub list capacity. + pub hubs: u32, + /// Number of sub-body slots requested this step, possibly past `pool`. + pub slots: u32, + /// First sub-body velocity slot (the number of regular body slots). + pub base: u32, + /// Sub-body slot capacity. + pub pool: u32, +} + +impl HubCounts { + /// Capacity of the hub list. + #[inline(always)] + pub fn max_hubs(&self) -> u32 { + self.pool / HUB_MIN_CONSTRAINTS + 1 + } + + /// Number of hubs recorded in the hub list. + #[inline(always)] + pub fn live_hubs(&self) -> u32 { + self.hubs.min(self.max_hubs()) + } + + /// Number of sub-body slots in use. + #[inline(always)] + pub fn live_slots(&self) -> u32 { + self.slots.min(self.pool) + } +} + +/// Number of workgroups averaging the sub-bodies back into their hubs (at most). +pub const HUB_AVERAGE_WORKGROUPS: u32 = 256; + +/// The constraint range of `body` in the cumulative per-body constraint counts. +#[inline(always)] +fn constraint_range(counts: &Slice, body: u32) -> (u32, u32) { + let first = if body != 0 { + counts[body as usize - 1] + } else { + 0 + }; + (first, counts[body as usize]) +} + +/// Picks the hubs and reserves their sub-body slots. +/// +/// Runs after the per-body constraint lists are built (`body_constraint_counts` holds each +/// body's cumulative end). `hub_first_slot` and the hub counters must have been reset. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_hub_assign( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] body_constraint_counts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mprops: &[WorldMassProperties], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_group: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] hub_first_slot: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] hub_list: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] hub_counts: &mut HubCounts, + #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, +) { + let num_threads = num_workgroups.x * WORKGROUP_SIZE; + let num_slots = batch_ids.colliders_batch_capacity * batch_ids.num_batches; + let counts = Slice(body_constraint_counts, 0); + let pool = hub_counts.pool; + let max_hubs = hub_counts.max_hubs(); + + for body in StepRng::new(invocation_id.x..num_slots, num_threads) { + let (first, last) = constraint_range(&counts, body); + let n = last - first; + // Multibody links are colored as their multibody's group and never split. + let splittable = n > HUB_MIN_CONSTRAINTS + && body_group.read(body as usize) == body + && mprops.at(body as usize).inv_mass != Vector::ZERO; + + if splittable { + let s0 = atomic_add_u32(&mut hub_counts.slots, n); + if s0 + n <= pool { + let k = atomic_add_u32(&mut hub_counts.hubs, 1); + if k < max_hubs { + hub_first_slot.write(body as usize, s0); + hub_list.write(k as usize, body); + } + } + } + } +} + +/// Points each hub constraint at its sub-body velocity slot, and scales the hub side's inverse +/// mass and inertia by the number of sub-bodies (applied when the constraints are built). +/// +/// One workgroup per hub, striding over its adjacency entries. Ordinary bodies +/// never enter this pass; no per-contact binary search of body ranges is needed. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_hub_split_constraints( + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] body_constraint_counts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_constraint_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] hub_first_slot: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] constraints: &mut [ContactLink], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] hub_slot_body: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] hub_slot_constraint: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] hub_counts: &HubCounts, + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] hub_list: &[u32], +) { + let counts = Slice(body_constraint_counts, 0); + let base = hub_counts.base; + let hubs = hub_counts.live_hubs(); + for hub in StepRng::new(workgroup_id.x..hubs, num_workgroups.x) { + let body = hub_list.read(hub as usize); + let s0 = hub_first_slot.read(body as usize); + let (first, last) = constraint_range(&counts, body); + for entry in StepRng::new(first + local_id.x..last, WORKGROUP_SIZE) { + let n = (last - first) as f32; + let slot = s0 + (entry - first); + let cid = body_constraint_ids.read(entry as usize); + let constraint = constraints.at_mut(cid as usize); + + if constraint.solver_body_a == body { + constraint.vel_slot_a = base + slot; + constraint.mass_scale_a = n; + } else { + constraint.vel_slot_b = base + slot; + constraint.mass_scale_b = n; + } + + hub_slot_body.write(slot as usize, body); + hub_slot_constraint.write(slot as usize, cid); + } + } +} + +/// Sizes hub-only dispatches after assignment, including the zero-hub case. +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_hub_dispatch( + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] hub_counts: &HubCounts, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] slots_indirect: &mut [[u32; 3]], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] average_indirect: &mut [[u32; 3]], +) { + let slots = hub_counts.live_slots(); + let hubs = hub_counts.live_hubs(); + slots_indirect.write(0, [slots.div_ceil(WORKGROUP_SIZE), 1, 1]); + average_indirect.write(0, [hubs.min(HUB_AVERAGE_WORKGROUPS), 1, 1]); +} + +/// Copies each hub's velocity into its sub-body slots. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_hub_scatter( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] hub_slot_body: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] hub_counts: &HubCounts, +) { + let slot = invocation_id.x; + if slot < hub_counts.live_slots() { + let base = hub_counts.base; + let body = hub_slot_body.read(slot as usize); + let vel = solver_vels.read(body as usize); + solver_vels.write((base + slot) as usize, vel); + } +} + +/// Averages the sub-body velocities back into their hub. One workgroup per hub. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_hub_average( + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] hub_list: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] hub_first_slot: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] body_constraint_counts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] hub_counts: &HubCounts, + #[spirv(workgroup)] linear_sums: &mut [Vector; 64], + #[spirv(workgroup)] angular_sums: &mut [AngVector; 64], +) { + let lane = local_id.x as usize; + let counts = Slice(body_constraint_counts, 0); + let base = hub_counts.base; + let num_hubs = hub_counts.live_hubs(); + let mut solver_vels = SliceMut(solver_vels, 0); + + // Uniform across the workgroup, so the barriers below are reached by every lane. The WGSL + // uniformity analysis doesn't trust `workgroup_id`, so on the web the rounds are bounded by + // `max_hubs` (read-only storage) and the out-of-range hubs run them inactive. + #[cfg(not(feature = "web-compat"))] + #[allow(clippy::implicit_saturating_sub)] + let num_rounds = if num_hubs > workgroup_id.x { + (num_hubs - workgroup_id.x).div_ceil(num_workgroups.x) + } else { + 0 + }; + #[cfg(feature = "web-compat")] + let num_rounds = hub_counts.max_hubs().div_ceil(HUB_AVERAGE_WORKGROUPS); + + for round in 0..num_rounds { + let k = workgroup_id.x + round * num_workgroups.x; + let active = k < num_hubs; + let body = if active { hub_list.read(k as usize) } else { 0 }; + let s0 = if active { + hub_first_slot.read(body as usize) + } else { + 0 + }; + let (first, last) = if active { + constraint_range(&counts, body) + } else { + (0, 0) + }; + let n = last - first; + + let mut linear = Vector::ZERO; + let mut angular = AngVector::default(); + for s in StepRng::new(lane as u32..n, WORKGROUP_SIZE) { + let vel = solver_vels[(base + s0 + s) as usize]; + linear += vel.linear; + angular += vel.angular; + } + linear_sums.write(lane, linear); + angular_sums.write(lane, angular); + + let mut stride = 32usize; + for _ in 0..6u32 { + workgroup_memory_barrier_with_group_sync(); + if lane < stride { + let l = linear_sums.read(lane) + linear_sums.read(lane + stride); + let a = angular_sums.read(lane) + angular_sums.read(lane + stride); + linear_sums.write(lane, l); + angular_sums.write(lane, a); + } + stride /= 2; + } + workgroup_memory_barrier_with_group_sync(); + + if lane == 0 && n > 0 { + let inv_n = 1.0 / n as f32; + solver_vels[body as usize].linear = linear_sums.read(0) * inv_n; + solver_vels[body as usize].angular = angular_sums.read(0) * inv_n; + } + // The shared sums are rewritten by the next hub. + workgroup_memory_barrier_with_group_sync(); + } +} + +#[cfg(all(test, not(target_arch_is_gpu)))] +#[path = "../tests/hub_split.rs"] +mod tests; diff --git a/src_rbd_shaders/dynamics/mod.rs b/src_rbd_shaders/dynamics/mod.rs index fb56a9d5..2a4be617 100644 --- a/src_rbd_shaders/dynamics/mod.rs +++ b/src_rbd_shaders/dynamics/mod.rs @@ -16,15 +16,22 @@ mod joint_constraint_builder; mod multibody; mod sim_params; mod solver_utils; -mod warmstart; +#[cfg(feature = "dim3")] +mod twist_friction; // GPU compute shader kernels mod canonical_order; mod color_buckets; mod coloring; +mod contact_links; +pub mod contact_tiles; +pub mod coulomb_tiles; +mod mass_splitting; mod mprops_update; mod prep_render; mod solver; +#[cfg(feature = "dim3")] +pub mod twist_tiles; pub use body::*; pub use constraint::*; @@ -41,8 +48,10 @@ pub use sim_params::*; pub use canonical_order::*; pub use color_buckets::*; pub use coloring::*; +pub use contact_links::*; +pub use mass_splitting::*; pub use mprops_update::*; pub use prep_render::*; pub use solver::*; -pub use solver_utils::warmstart_body; -pub use warmstart::*; +#[cfg(feature = "dim3")] +pub use twist_friction::*; diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index faac3aa0..68e899cf 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -37,7 +37,7 @@ use super::types::{ MultibodyLinkStatic, }; -/// Workgroup width of the flat contact sweeps below (must match the +/// Workgroup width of the flat contact dispatches below (must match the /// `contacts_indirect` grid written by `gpu_contact_plan`). const MB_SWEEP_WG: u32 = 64; use super::utils::zero_kinematic_dofs; @@ -208,7 +208,7 @@ fn mb_contact_owner( /// Predicts, per (multibody, batch), how many contact-constraint slots the /// emission pass needs, and how many contacts touch each multibody. One flat -/// sweep over the contacts: each contact resolves its owner and atomically +/// dispatch over the contacts: each contact resolves its owner and atomically /// bumps that owner's counters, so the pass costs one visit per contact /// instead of one full-list scan per (multibody, batch). /// @@ -310,7 +310,7 @@ pub fn gpu_mb_cons_offsets_scan( mb_cons_demand.write(0, demand); } -/// Builds the contact→multibody index: one flat sweep over the contacts, +/// Builds the contact→multibody index: one flat dispatch over the contacts, /// each contact appending its entry to its owner's segment (laid out by the /// offsets scan; `mb_index_counts` still holds the per-multibody counts and /// serves as write cursors counting down to zero). Entry order within a segment follows the @@ -1281,13 +1281,13 @@ pub fn gpu_mb_apply_contact_restitution( workgroup_memory_barrier_with_group_sync(); #[cfg(feature = "web-compat")] - let contact_sweep_len = *max_contact_constraints; + let contact_iteration_len = *max_contact_constraints; #[cfg(not(feature = "web-compat"))] - let contact_sweep_len = count; + let contact_iteration_len = count; #[cfg(not(feature = "web-compat"))] let _ = max_contact_constraints; - for s in 0..contact_sweep_len { + for s in 0..contact_iteration_len { let slot_active = active && s < count; let cons_idx = if slot_active { cons_base + s as usize diff --git a/src_rbd_shaders/dynamics/multibody/contact_sensor.rs b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs index 89589861..66c1b151 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_sensor.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs @@ -10,7 +10,7 @@ use khal_std::macros::{spirv, spirv_bindgen}; pub const MAX_CONTACT_SENSORS: u32 = 4; /// Per sensed link, sums the accumulated normal-constraint impulses. Dispatch -/// it once per step, after the last substep's stabilization sweep: the value is +/// it once per step, after the last substep's stabilization iteration: the value is /// then the step's total accumulated normal impulse (divide by the step `dt` /// for an average force) when the constraints are built once per step, or the /// last substep's impulse when they are rebuilt per substep. diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs index dd74d16c..b65b7ff2 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs @@ -52,7 +52,7 @@ pub fn gpu_mb_update_impulse_joint_constraints( let lock_cfm_coeff = softness.joint_cfm_coeff; let max_corr_velocity = softness.max_corr_velocity; - // Flat sweep over the interleaved builder slots (slot `t` = joint `t / nb` + // Flat dispatch over the interleaved builder slots (slot `t` = joint `t / nb` // of batch `t % nb`). Iterating to `cap * nb` instead of the per-batch // lengths lets us drop the `num_joints` storage binding — the host pads // unused slots with `side_a_kind == SIDE_KIND_FIXED && side_b_kind == @@ -170,9 +170,8 @@ pub fn gpu_mb_finalize_impulse_joint_constraints( /// /// Joints are graph-colored at init time (see `set_impulse_joints`): within /// one color no two joints share a multibody or a free body, so the color's -/// joints run race-free in parallel. The host dispatches one color per -/// iteration, giving an exact sequential Gauss–Seidel sweep in color-sorted -/// order. +/// joints run race-free in parallel. The host dispatches one color at a time, +/// giving an exact sequential Gauss–Seidel iteration in color-sorted order. #[spirv_bindgen] #[spirv(compute(threads(64, 1, 1)))] pub fn gpu_mb_solve_impulse_joint_constraints( @@ -326,7 +325,7 @@ pub fn gpu_mb_solve_impulse_joint_constraints( } /// Strip the positional bias from each active constraint's `rhs` for the -/// stabilization sweep — mirrors `gpu_mb_remove_joint_constraint_bias`. +/// stabilization iteration — mirrors `gpu_mb_remove_joint_constraint_bias`. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_remove_impulse_joint_constraint_bias( diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/mod.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/mod.rs index b92043f9..b9c36213 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/mod.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/mod.rs @@ -11,7 +11,7 @@ //! 3. `gpu_mb_solve_impulse_joint_constraints` — one PGS iteration, updates both //! sides' velocities. //! 4. `gpu_mb_remove_impulse_joint_constraint_bias` — strips the positional -//! bias from `rhs` before the stabilization sweep. +//! bias from `rhs` before the stabilization iteration. mod helper; mod jacobians; diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index fe361342..aa32fffe 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -1,7 +1,7 @@ //! Multibody joint limit / motor constraints. //! //! Each constraint targets a single generalized DOF and is solved with PGS -//! sweeps. Per-multibody, all constraint slots are scanned (`kind == 0` ones +//! iterations. Per-multibody, all constraint slots are scanned (`kind == 0` ones //! are skipped). use glamx::Vec4; diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index d1140b97..c513140d 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -154,12 +154,12 @@ pub fn gpu_mb_solve_constraints( workgroup_memory_barrier_with_group_sync(); #[cfg(feature = "web-compat")] - let joint_sweep_len = batch_ids.mb_max_joint_constraints; + let joint_iteration_len = batch_ids.mb_max_joint_constraints; #[cfg(not(feature = "web-compat"))] - let joint_sweep_len = mb.max_constraints; + let joint_iteration_len = mb.max_constraints; // Joint limits/motors - for s in 0..joint_sweep_len { + for s in 0..joint_iteration_len { let slot_active = active && s < mb.max_constraints; let cons_idx = if slot_active { jcons_base + s as usize @@ -218,13 +218,13 @@ pub fn gpu_mb_solve_constraints( // together so their impulse can be capped to the friction cone; the second // row is handled by its sibling and skipped here. #[cfg(feature = "web-compat")] - let contact_sweep_len = *max_contact_constraints; + let contact_iteration_len = *max_contact_constraints; #[cfg(not(feature = "web-compat"))] - let contact_sweep_len = contact_count; + let contact_iteration_len = contact_count; #[cfg(not(feature = "web-compat"))] let _ = max_contact_constraints; - for s in 0..contact_sweep_len { + for s in 0..contact_iteration_len { let slot_active = active && s < contact_count; let cons_idx = if slot_active { ccons_base + s as usize @@ -337,7 +337,7 @@ pub fn gpu_mb_solve_constraints( // Normal: clamp to ≥ 0. Friction: cap the tangent pair to the // circular cone `μ · normal_impulse`. let (new0, new1) = if is_tangent { - // The paired normal was updated earlier in this sweep by this + // The paired normal was updated earlier in this iteration by this // same lane, so the storage read observes the fresh value. let limit = cons.friction_coeff * contact_constraints @@ -451,11 +451,11 @@ pub fn gpu_mb_solve_joints( workgroup_memory_barrier_with_group_sync(); #[cfg(feature = "web-compat")] - let joint_sweep_len = batch_ids.mb_max_joint_constraints; + let joint_iteration_len = batch_ids.mb_max_joint_constraints; #[cfg(not(feature = "web-compat"))] - let joint_sweep_len = mb.max_constraints; + let joint_iteration_len = mb.max_constraints; - for s in 0..joint_sweep_len { + for s in 0..joint_iteration_len { let slot_active = active && s < mb.max_constraints; let cons_idx = if slot_active { jcons_base + s as usize @@ -593,7 +593,7 @@ pub fn gpu_mb_build_contact_delassus( } } -/// Constraint-space contact sweep: tracks `a[s] = J_s · u` incrementally in +/// Constraint-space contact iteration: tracks `a[s] = J_s · u` incrementally in /// workgroup memory using the precomputed Delassus rows, so each PGS /// iteration is a couple of shared-memory scalars plus one lane-parallel row /// update. @@ -691,7 +691,7 @@ pub fn gpu_mb_solve_contacts_delassus( } workgroup_memory_barrier_with_group_sync(); - // Fresh `a[s] = J_s · u` under the current (post-joint-sweep, post- + // Fresh `a[s] = J_s · u` under the current (post-joint-iteration, post- // warmstart) velocities. if active { for s in StepRng::new(lane..count, LANES) { @@ -714,13 +714,13 @@ pub fn gpu_mb_solve_contacts_delassus( // their impulse can be capped to the friction cone; the second row is // handled by its sibling and skipped here. #[cfg(feature = "web-compat")] - let contact_sweep_len = *max_contact_constraints; + let contact_iteration_len = *max_contact_constraints; #[cfg(not(feature = "web-compat"))] - let contact_sweep_len = count; + let contact_iteration_len = count; #[cfg(not(feature = "web-compat"))] let _ = max_contact_constraints; - for s in 0..contact_sweep_len { + for s in 0..contact_iteration_len { let slot_active = active && s < count; let meta = if slot_active { meta_shared.read(s as usize) @@ -732,7 +732,7 @@ pub fn gpu_mb_solve_contacts_delassus( let free_active = (meta >> 24) != 0; let is_tangent = kind == MB_CONTACT_KIND_TANGENT; - // Friction is solved during the stabilization sweep (and during the + // Friction is solved during the stabilization iteration (and during the // biased pass when `friction_in_bias_pass` is set); in 3D a tangent // pair is solved by its first row only. #[cfg(feature = "dim3")] diff --git a/src_rbd_shaders/dynamics/multibody/types.rs b/src_rbd_shaders/dynamics/multibody/types.rs index 64884878..a9069a51 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -314,7 +314,7 @@ pub struct MultibodyContactConstraint { /// `J·v_target + bias` — bias from penetration (`erp_inv_dt · depth`) /// for normals, surface velocity for tangents. pub rhs: f32, - /// `rhs` without the positional bias (used by the stabilization sweep). + /// `rhs` without the positional bias (used by the stabilization iteration). pub rhs_wo_bias: f32, /// Accumulated impulse (warmstart-able). pub impulse: f32, diff --git a/src_rbd_shaders/dynamics/sim_params.rs b/src_rbd_shaders/dynamics/sim_params.rs index f0825eed..6ca77e49 100644 --- a/src_rbd_shaders/dynamics/sim_params.rs +++ b/src_rbd_shaders/dynamics/sim_params.rs @@ -2,6 +2,26 @@ use crate::MAX_FLT; +/// Simulation-wide friction choice for 3D rigid-body contacts. +/// +/// Multibody contacts always use Coulomb, as in Rapier. This transparent value +/// has a stable four-byte GPU representation; unknown values behave as Coulomb. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(transparent)] +pub struct FrictionModel(u32); + +#[cfg(feature = "dim3")] +#[allow(non_upper_case_globals)] +impl FrictionModel { + /// One Coulomb friction constraint at every contact point (Nexus's default). + pub const Coulomb: Self = Self(0); + /// Rapier's simplified model: one central friction constraint and one + /// angular twist constraint per manifold of up to four contact points. + pub const Simplified: Self = Self(1); +} + /// Two times pi (2π), used for converting natural frequency to angular frequency. pub const TWO_PI: f32 = core::f32::consts::TAU; @@ -18,7 +38,7 @@ pub struct ConstraintSoftness { /// overshoots and jitters. pub erp_inv_dt: f32, /// Contact `1 / (1 + cfm_coeff)` — multiplies the contact impulse each PGS - /// sweep for constraint-force-mixing compliance. + /// iteration for constraint-force-mixing compliance. pub cfm_factor: f32, /// Geometric slop distance. pub allowed_lin_err: f32, @@ -70,7 +90,7 @@ impl ConstraintSoftness { } /// Bias-mode uniform handed to the constraint solve kernels: the unbiased -/// stabilization sweep, the biased pass with friction rows skipped (rapier's +/// stabilization iteration, the biased pass with friction rows skipped (rapier's /// default scheduling), or the biased pass with friction rows solved /// (`RbdSimParams::friction_in_bias_pass`). The values double as indices into /// the solver's constant uniforms (`color_uniforms[c] == c`). @@ -184,7 +204,7 @@ pub struct RbdSimParams { pub contact_merge_cos: f32, /// PGS iterations over the joint + contact constraints run per substep in - /// the biased pass (default: `1`). The rigid-body and multibody sweeps + /// the biased pass (default: `1`). The rigid-body and multibody iterations /// interleave, one iteration each, so both sides of a rigid-body/multibody /// contact converge at the same rate. /// @@ -192,7 +212,7 @@ pub struct RbdSimParams { pub num_internal_pgs_iterations: u32, /// Nonzero: friction rows are also solved during the biased pass instead - /// of only during the unbiased stabilization sweep (default: `0`, matching + /// of only during the unbiased stabilization iteration (default: `0`, matching /// rapier's `friction_in_bias_pass`). Turning it on gives friction as many /// PGS iterations as the normal rows, which stiffens grasps and resting /// contacts at a small cost per iteration. @@ -200,11 +220,19 @@ pub struct RbdSimParams { /// Host-side only: it selects the bias-mode uniform passed to the solve /// kernels, never read by a shader. pub friction_in_bias_pass: u32, - // Uniform-layout padding to a 16-byte multiple (scalars: an array member - // here would itself need 16-byte alignment). + + /// A contact pair whose relative pose drifted less than this since its contacts were computed + /// keeps them, frozen on the bodies (default: `0.05` as in rapier, `0` disables recycling). + /// + /// This value is implicitly scaled by `length_unit`. + pub normalized_contact_recycle_distance: f32, + // Reserved scalar, matching the 3D friction-model slot. + #[cfg(feature = "dim2")] pub _padding0: u32, - pub _padding1: u32, - pub _padding2: u32, + /// Friction model for all rigid-body contact manifolds. Multibody contacts + /// remain Coulomb. Defaults to [`FrictionModel::Coulomb`]. + #[cfg(feature = "dim3")] + pub friction_model: FrictionModel, } impl RbdSimParams { @@ -231,9 +259,11 @@ impl RbdSimParams { length_unit: 1.0, num_internal_pgs_iterations: 1, friction_in_bias_pass: 0, + normalized_contact_recycle_distance: 0.05, + #[cfg(feature = "dim2")] _padding0: 0, - _padding1: 0, - _padding2: 0, + #[cfg(feature = "dim3")] + friction_model: FrictionModel::Coulomb, } } } @@ -379,6 +409,12 @@ impl RbdSimParams { self.normalized_prediction_distance * self.length_unit } + /// The distance below which a contact pair's drift lets it keep its contacts (default: + /// `0.05` multiplied by `length_unit`, `0` when recycling is disabled). + pub fn contact_recycle_distance(&self) -> f32 { + self.normalized_contact_recycle_distance * self.length_unit + } + /// Maximum linear velocity a body may have after each solver substep. /// /// This is equal to `normalized_max_linear_velocity` multiplied by diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 01b22c05..d00e4cf7 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -7,81 +7,19 @@ use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; use crate::{AngVector, Pose, Vector}; -use khal_std::{ - index::MaybeIndexUnchecked, - iter::StepRng, - sync::{atomic_add_u32, control_barrier}, -}; +use khal_std::{index::MaybeIndexUnchecked, iter::StepRng, sync::atomic_add_u32}; use super::body::{LocalMassProperties, Velocity, WorldMassProperties}; -use super::constraint::{TwoBodyConstraint, TwoBodyConstraintBuilder}; -use super::sim_params::{RbdSimParams, decode_bias_mode}; -use super::solver_utils::warmstart_body; +use super::mass_splitting::{HubCounts, NOT_A_HUB}; +use super::sim_params::RbdSimParams; use crate::queries::IndexedManifold; use crate::utils::{BatchIndices, Slice, SliceMut}; const WORKGROUP_SIZE: u32 = 64; -/// Initializes constraints from contact manifolds. -/// -/// Split into two passes to stay within WebGPU's 8-storage-buffer per-stage -/// limit: this pass builds the per-contact constraint/builder; -/// `gpu_solver_count_constraints` does the per-body-group constraint counting. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_solver_init_constraints( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(num_workgroups)] num_workgroups: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &[IndexedManifold], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] - constraints: &mut [TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] - constraint_builders: &mut [TwoBodyConstraintBuilder], - #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] body_is_multibody: &[u32], - #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] collider_world_poses: &[Pose], - #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_body_poses: &[Pose], - #[spirv(storage_buffer, descriptor_set = 1, binding = 2)] vels: &[Velocity], - #[spirv(storage_buffer, descriptor_set = 1, binding = 3)] mprops: &[WorldMassProperties], - #[spirv(uniform, descriptor_set = 1, binding = 4)] params: &RbdSimParams, -) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - - let total = contact_plan.bound; - let collider_world_poses = Slice(collider_world_poses, 0); - let solver_body_poses = Slice(solver_body_poses, 0); - let vels = Slice(vels, 0); - let mprops = Slice(mprops, 0); - let body_is_multibody = Slice(body_is_multibody, 0); - - for i in StepRng::new(invocation_id.x..total, num_threads) { - let i = i as usize; - let im = contacts.at(i); - // Manifolds touching a multibody link belong to the multibody contact - // solver (`gpu_mb_init_contact_constraints`); building a rigid-body - // constraint for them too would solve the contact twice, the second - // time against a zero-inverse-mass copy of the link that never moves. - if im.contact.len == 0 || touches_multibody(im, &body_is_multibody) { - // Gap, inert or multibody-owned slot: clear the (stale) - // constraint so every flat consumer skips it. - constraints.at_mut(i).len = 0; - continue; - } - im.contact_to_constraint( - &mprops, - &collider_world_poses, - &solver_body_poses, - &vels, - params, - constraints.at_mut(i), - constraint_builders.at_mut(i), - ); - } -} - -/// Companion pass to `gpu_solver_init_constraints`: counts, per body-group, how -/// many constraints touch each body (used to size the graph-coloring graph). +/// Counts, per body, how many constraints touch it (the adjacency lists of mass splitting, +/// graph coloring and the warmstart gather). #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_solver_count_constraints( @@ -108,7 +46,7 @@ pub fn gpu_solver_count_constraints( continue; } // Multibody-owned manifolds have no rigid-body constraint (see - // `gpu_solver_init_constraints`). + // `gpu_init_contact_links`). if touches_multibody(im, &body_is_multibody) { continue; } @@ -134,75 +72,6 @@ fn touches_multibody(im: &IndexedManifold, body_is_multibody: &Slice<'_, u32>) - body_is_multibody[im.bodies.x as usize] != 0 || body_is_multibody[im.bodies.y as usize] != 0 } -/// Updates constraints for a new substep. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_solver_update_constraints( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(num_workgroups)] num_workgroups: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] - constraints: &mut [TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] - constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(uniform, descriptor_set = 0, binding = 2)] contact_plan: &ContactPlan, - #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] solver_body_poses: &[Pose], - #[spirv(uniform, descriptor_set = 1, binding = 1)] params: &RbdSimParams, -) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - - let total = contact_plan.bound; - let mut constraints = SliceMut(constraints, 0); - let constraint_builders = Slice(constraint_builders, 0); - let solver_body_poses = Slice(solver_body_poses, 0); - - for i in StepRng::new(invocation_id.x..total, num_threads) { - if constraints[i as usize].len == 0 { - // Gap / inert slot. - continue; - } - constraints[i as usize].update_constraint( - &constraint_builders[i as usize], - &solver_body_poses, - params, - ); - } -} - -/// Relax-pass refresh: recomputes the unbiased normal rhs of every contact -/// constraint from the current (post-integration) solver poses. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_solver_refresh_rhs_wo_bias( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(num_workgroups)] num_workgroups: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] - constraints: &mut [TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] - constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(uniform, descriptor_set = 0, binding = 2)] contact_plan: &ContactPlan, - #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] solver_body_poses: &[Pose], - #[spirv(uniform, descriptor_set = 1, binding = 1)] params: &RbdSimParams, -) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - - let total = contact_plan.bound; - let mut constraints = SliceMut(constraints, 0); - let constraint_builders = Slice(constraint_builders, 0); - let solver_body_poses = Slice(solver_body_poses, 0); - - for i in StepRng::new(invocation_id.x..total, num_threads) { - if constraints[i as usize].len == 0 { - // Gap / inert slot. - continue; - } - constraints[i as usize].refresh_rhs_wo_bias( - &constraint_builders[i as usize], - &solver_body_poses, - params, - ); - } -} - #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_solver_sort_constraints( @@ -255,13 +124,21 @@ pub fn gpu_solver_cleanup( #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] vels: &[Velocity], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] mprops: &[WorldMassProperties], #[spirv(uniform, descriptor_set = 1, binding = 2)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 1, binding = 3)] hub_first_slot: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 1, binding = 4)] hub_counts: &mut HubCounts, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let num_slots = batch_ids.colliders_batch_capacity * batch_ids.num_batches; + if invocation_id.x == 0 { + hub_counts.hubs = 0; + hub_counts.slots = 0; + } + for i in StepRng::new(invocation_id.x..num_slots, num_threads) { let idx = i as usize; body_constraint_counts.write(idx, 0); + hub_first_slot.write(idx, NOT_A_HUB); // HACK: to handle static bodies. if mprops.at(idx).inv_mass != Vector::ZERO { @@ -326,263 +203,6 @@ pub fn gpu_apply_solver_vels_inc( } } -/// Applies warmstart impulses without graph coloring (gather-style per body). -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_warmstart_without_colors( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(num_workgroups)] num_workgroups: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] body_constraint_counts: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_constraint_ids: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] solver_vels: &mut [Velocity], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, -) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let num_bodies = batch_ids.bodies_len * batch_ids.num_batches; - - let body_constraint_counts = Slice(body_constraint_counts, 0); - let body_constraint_ids = Slice(body_constraint_ids, 0); - let constraints = Slice(constraints, 0); - let mut solver_vels = SliceMut(solver_vels, 0); - - for body_id in StepRng::new(invocation_id.x..num_bodies, num_threads) { - let mut solver_vel = solver_vels[body_id as usize]; - warmstart_body( - body_id, - &body_constraint_counts, - &body_constraint_ids, - &constraints, - &mut solver_vel, - ); - solver_vels[body_id as usize] = solver_vel; - } -} - -/// Applies warmstart impulses with graph coloring (scatter-style per constraint). -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_warmstart( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(num_workgroups)] num_workgroups: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_starts: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] curr_color: &u32, - #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, -) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let nb = batch_ids.num_batches; - - let constraints = Slice(constraints, 0); - let color_sorted_ids = Slice(color_sorted_ids, 0); - let mut solver_vels = SliceMut(solver_vels, 0); - let color = *curr_color; - - // Buckets are color-major (`color * num_batches + batch`) and the buffer - // holds post-scatter exclusive ENDS, so color `c` (over every batch) spans - // `[ends[c*nb - 1], ends[(c+1)*nb - 1])`. `c >= 1` keeps the index valid. - let start = color_starts.read((color * nb - 1) as usize); - let end = color_starts.read(((color + 1) * nb - 1) as usize); - - for k in StepRng::new(start + invocation_id.x..end, num_threads) { - let i = color_sorted_ids[k as usize]; - let constraint = &constraints[i as usize]; - let solver_id1 = constraint.solver_body_a as usize; - let solver_id2 = constraint.solver_body_b as usize; - - let mut solver_vel1 = solver_vels[solver_id1]; - let mut solver_vel2 = solver_vels[solver_id2]; - - constraint.warmstart_constraint(&mut solver_vel1, &mut solver_vel2); - - solver_vels[solver_id1] = solver_vel1; - solver_vels[solver_id2] = solver_vel2; - } -} - -/// Main constraint solver iteration kernel (Projected Gauss-Seidel). -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_step_gauss_seidel( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(num_workgroups)] num_workgroups: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] - constraints: &mut [TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_starts: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] curr_color: &u32, - #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, - #[spirv(uniform, descriptor_set = 0, binding = 6)] use_bias: &u32, -) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let nb = batch_ids.num_batches; - - let mut constraints = SliceMut(constraints, 0); - let color_sorted_ids = Slice(color_sorted_ids, 0); - let mut solver_vels = SliceMut(solver_vels, 0); - let color = *curr_color; - let (use_bias, solve_friction) = decode_bias_mode(*use_bias); - - // Color-major bucket ends; see `gpu_warmstart`. - let start = color_starts.read((color * nb - 1) as usize); - let end = color_starts.read(((color + 1) * nb - 1) as usize); - - for k in StepRng::new(start + invocation_id.x..end, num_threads) { - let i = color_sorted_ids[k as usize]; - let solver_id1 = constraints[i as usize].solver_body_a as usize; - let solver_id2 = constraints[i as usize].solver_body_b as usize; - - let mut solver_vel1 = solver_vels[solver_id1]; - let mut solver_vel2 = solver_vels[solver_id2]; - - constraints[i as usize].solve_constraint_gauss_seidel( - &mut solver_vel1, - &mut solver_vel2, - use_bias, - solve_friction, - ); - - solver_vels[solver_id1] = solver_vel1; - solver_vels[solver_id2] = solver_vel2; - } -} - -/// Fused colored warmstart: only one 64-lane workgroup per batch walks every color -/// bucket. -/// -/// Used for small scenes where the contact count is small wrt. the environment count. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_warmstart_fused( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_starts: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] num_colors: &u32, - #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, -) { - let lane = invocation_id.x; - let batch_id = invocation_id.y; - let nb = batch_ids.num_batches; - - let constraints = Slice(constraints, 0); - let color_sorted_ids = Slice(color_sorted_ids, 0); - let mut solver_vels = SliceMut(solver_vels, 0); - let num_colors = *num_colors; - - for color in 1..=num_colors { - // This batch's bucket for `color` (color-major layout, post-scatter - // exclusive ends; the index is >= num_batches >= 1 for color >= 1). - let bucket = (color * nb + batch_id) as usize; - let start = color_starts.read(bucket - 1); - let end = color_starts.read(bucket); - #[cfg(not(feature = "web-compat"))] - if start == end { - // Empty color. - continue; - } - - if start != end { - for k in StepRng::new(start + lane..end, WORKGROUP_SIZE) { - let i = color_sorted_ids[k as usize]; - let constraint = &constraints[i as usize]; - let solver_id1 = constraint.solver_body_a as usize; - let solver_id2 = constraint.solver_body_b as usize; - - let mut solver_vel1 = solver_vels[solver_id1]; - let mut solver_vel2 = solver_vels[solver_id2]; - - constraint.warmstart_constraint(&mut solver_vel1, &mut solver_vel2); - - solver_vels[solver_id1] = solver_vel1; - solver_vels[solver_id2] = solver_vel2; - } - } - - control_barrier::< - { khal_std::memory::Scope::Workgroup as u32 }, - { khal_std::memory::Scope::QueueFamily as u32 }, - { - khal_std::memory::Semantics::UNIFORM_MEMORY.bits() - | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() - }, - >(); - } -} - -/// Fused colored Gauss-Seidel sweep: only one 64-lane workgroup per batch walks -/// every color bucket with a storage barrier between colors. -/// -/// Used for small scenes where the contact count is small wrt. the environment count. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_step_gauss_seidel_fused( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] - constraints: &mut [TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_starts: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] num_colors: &u32, - #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, - #[spirv(uniform, descriptor_set = 0, binding = 6)] use_bias: &u32, -) { - let lane = invocation_id.x; - let batch_id = invocation_id.y; - let nb = batch_ids.num_batches; - - let mut constraints = SliceMut(constraints, 0); - let color_sorted_ids = Slice(color_sorted_ids, 0); - let mut solver_vels = SliceMut(solver_vels, 0); - let num_colors = *num_colors; - let (use_bias, solve_friction) = decode_bias_mode(*use_bias); - - for color in 1..=num_colors { - // Empty-color skip: see `gpu_warmstart_fused`. - let bucket = (color * nb + batch_id) as usize; - let start = color_starts.read(bucket - 1); - let end = color_starts.read(bucket); - #[cfg(not(feature = "web-compat"))] - if start == end { - continue; - } - - if start != end { - for k in StepRng::new(start + lane..end, WORKGROUP_SIZE) { - let i = color_sorted_ids[k as usize]; - let solver_id1 = constraints[i as usize].solver_body_a as usize; - let solver_id2 = constraints[i as usize].solver_body_b as usize; - - let mut solver_vel1 = solver_vels[solver_id1]; - let mut solver_vel2 = solver_vels[solver_id2]; - - constraints[i as usize].solve_constraint_gauss_seidel( - &mut solver_vel1, - &mut solver_vel2, - use_bias, - solve_friction, - ); - - solver_vels[solver_id1] = solver_vel1; - solver_vels[solver_id2] = solver_vel2; - } - } - - control_barrier::< - { khal_std::memory::Scope::Workgroup as u32 }, - { khal_std::memory::Scope::QueueFamily as u32 }, - { - khal_std::memory::Semantics::UNIFORM_MEMORY.bits() - | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() - }, - >(); - } -} - /// Integrates velocity to update poses. #[spirv_bindgen] #[spirv(compute(threads(64)))] diff --git a/src_rbd_shaders/dynamics/solver_utils.rs b/src_rbd_shaders/dynamics/solver_utils.rs index eeaff880..bf100dfe 100644 --- a/src_rbd_shaders/dynamics/solver_utils.rs +++ b/src_rbd_shaders/dynamics/solver_utils.rs @@ -5,29 +5,21 @@ use super::body::{Velocity, WorldMassProperties}; use super::constraint::{ - SUB_LEN, TwoBodyConstraint, TwoBodyConstraintBuilder, TwoBodyConstraintNormalPart, - TwoBodyConstraintTangentPart, + ContactLink, ContactPoint, MAX_CONSTRAINTS_PER_MANIFOLD, SUB_LEN, TwoBodyConstraint, }; use super::sim_params::RbdSimParams; -use crate::{Pad, Pose, Vector, gcross, gdot}; +use crate::{Pose, Vector, gcross, gcross_av, gdot}; use khal_std::index::MaybeIndexUnchecked; -#[cfg(feature = "dim2")] -use glamx::Vec2; #[cfg(feature = "dim3")] use glamx::{Vec2, Vec3}; use crate::queries::IndexedManifold; use crate::utils::Slice; -/// Helper function for safe inverse. -fn inv(x: f32) -> f32 { - if x == 0.0 { 0.0 } else { 1.0 / x } -} - /// Helper function for maybe inverse with threshold. #[cfg(feature = "dim3")] -fn maybe_inv(a: f32) -> f32 { +pub(crate) fn maybe_inv(a: f32) -> f32 { const INV_EPSILON: f32 = 1.0e-20; if a < -INV_EPSILON || a > INV_EPSILON { 1.0 / a @@ -38,17 +30,11 @@ fn maybe_inv(a: f32) -> f32 { /// Cap the magnitude of a 2D vector. #[cfg(feature = "dim3")] -fn cap_magnitude(v: Vec2, limit: f32) -> Vec2 { +pub(crate) fn cap_magnitude(v: Vec2, limit: f32) -> Vec2 { let n = v.length(); if n > limit { v * (limit / n) } else { v } } -#[cfg(feature = "dim2")] -/// Computes an orthonormal vector perpendicular to the input (2D). -fn orthonormal_vector(vec: Vector) -> Vector { - Vec2::new(-vec.y, vec.x) -} - #[cfg(feature = "dim3")] /// Computes an orthonormal vector perpendicular to the input (3D). fn orthonormal_vector(vec: Vector) -> Vector { @@ -64,35 +50,88 @@ fn orthonormal_vector(vec: Vector) -> Vector { Vec3::new(b, sign + vec.y * vec.y * a, -vec.y) } -#[cfg(feature = "dim2")] -/// Computes the tangent contact directions for friction (2D version). -fn compute_tangent_contact_directions(force_dir1: Vector) -> [Vector; SUB_LEN] { - [orthonormal_vector(force_dir1)] -} - #[cfg(feature = "dim3")] /// Computes the tangent contact directions for friction (3D version). /// /// A deterministic basis derived from the contact normal. -fn compute_tangent_contact_directions(force_dir1: Vector) -> [Vector; SUB_LEN] { +pub(crate) fn compute_tangent_contact_directions(force_dir1: Vector) -> [Vector; SUB_LEN] { let tangent1 = orthonormal_vector(force_dir1); let bitangent1 = force_dir1.cross(tangent1); [tangent1, bitangent1] } -#[cfg(feature = "dim2")] -/// Cross product for angular velocity (2D). -fn gcross_(a: f32, b: Vector) -> Vector { - a * b +// Constant indices let the shader compiler scalarize each manifold point. +// Preserve Gauss-Seidel point order; only the loop control changes. +macro_rules! for_contact_point { + ($k:ident, $len:expr, $body:block) => {{ + let len = $len as usize; + { + let $k = 0usize; + if $k < len { + $body + } + } + { + let $k = 1usize; + if $k < len { + $body + } + } + #[cfg(feature = "dim3")] + { + let $k = 2usize; + if $k < len { + $body + } + } + #[cfg(feature = "dim3")] + { + let $k = 3usize; + if $k < len { + $body + } + } + }}; } -#[cfg(feature = "dim3")] -/// Cross product for angular velocity (3D). -fn gcross_(a: Vector, b: Vector) -> Vector { - a.cross(b) -} +pub(crate) use for_contact_point; impl IndexedManifold { + /// Initializes the manifold identity and current mass/material data without rebuilding points. + #[inline(always)] + pub fn contact_to_constraint_header( + &self, + mprops: &Slice, + constraint: &mut TwoBodyConstraint, + ) { + let bid1 = self.bodies.x; + let bid2 = self.bodies.y; + let mprops1 = &mprops[bid1 as usize]; + let mprops2 = &mprops[bid2 as usize]; + #[cfg(feature = "dim3")] + { + constraint.ii_a = super::constraint::SymInertia::from_mat4(mprops1.inv_inertia); + constraint.ii_b = super::constraint::SymInertia::from_mat4(mprops2.inv_inertia); + } + #[cfg(feature = "dim2")] + { + constraint.ii_a = mprops1.inv_inertia; + constraint.ii_b = mprops2.inv_inertia; + } + constraint.im_a = mprops1.inv_mass; + constraint.im_b = mprops2.inv_mass; + constraint.limit = self.friction; + constraint.restitution = self.restitution; + constraint.solver_body_a = bid1; + constraint.solver_body_b = bid2; + constraint.vel_slot_a = bid1; + constraint.vel_slot_b = bid2; + constraint.warmstart_collider_a = self.colliders.x; + constraint.warmstart_collider_b = self.colliders.y; + constraint.warmstart_subshape = self.subshape; + constraint.len = self.contact.len; + } + /// Converts a contact manifold to a solver constraint. /// /// `collider_world_poses` are used to recover the world-space contact normal @@ -106,23 +145,32 @@ impl IndexedManifold { collider_world_poses: &Slice, solver_body_poses: &Slice, vels: &Slice, - params: &RbdSimParams, constraint: &mut TwoBodyConstraint, - builder: &mut TwoBodyConstraintBuilder, ) { - let id1 = self.colliders.x; + self.contact_to_constraint_header(mprops, constraint); + self.contact_to_constraint_points( + collider_world_poses[self.colliders.x as usize], + solver_body_poses[self.bodies.x as usize], + solver_body_poses[self.bodies.y as usize], + vels, + constraint, + ); + } + + /// Initializes the contact points of `constraint`, whose header was initialized from this + /// manifold, and their effective masses. + #[inline(always)] + pub fn contact_to_constraint_points( + &self, + cpose1: Pose, + spose1: Pose, + spose2: Pose, + vels: &Slice, + constraint: &mut TwoBodyConstraint, + ) { let contact = &self.contact; let bid1 = self.bodies.x; let bid2 = self.bodies.y; - - let mprops1 = &mprops[bid1 as usize]; - let mprops2 = &mprops[bid2 as usize]; - // Contact features (`points_a`, `normal_a`) are stored in collider A's - // local space, so only `cpose1` is needed to recover their world-space - // forms; collider B's pose isn't read here. - let cpose1 = collider_world_poses[id1 as usize]; - let spose1 = solver_body_poses[bid1 as usize]; - let spose2 = solver_body_poses[bid2 as usize]; let vel1 = &vels[bid1 as usize]; let vel2 = &vels[bid2 as usize]; @@ -131,27 +179,13 @@ impl IndexedManifold { // Combined friction / restitution of the two colliders, resolved at // narrow-phase time (rapier's `CoefficientCombineRule`) and stored in the // manifold. - let friction = self.friction; let restitution = self.restitution; - let tangents1 = compute_tangent_contact_directions(force_dir1); constraint.dir_a = force_dir1; - constraint.im_a = mprops1.inv_mass; - constraint.im_b = mprops2.inv_mass; - // Softness (cfm, erp) and the rhs are (re)computed by - // `update_constraint` before the first solve of every substep. - constraint.cfm_factor = params.contact_cfm_factor(); - constraint.limit = friction; - constraint.solver_body_a = bid1; - constraint.solver_body_b = bid2; - #[cfg(feature = "dim3")] { - constraint.tangent_a = tangents1.read(0); + constraint.tangent_a = compute_tangent_contact_directions(force_dir1).read(0); } - constraint.warmstart_collider_a = self.colliders.x; - constraint.warmstart_collider_b = self.colliders.y; - constraint.warmstart_subshape = self.subshape; for k in 0..(contact.len as usize) { let pt = cpose1 @@ -161,511 +195,441 @@ impl IndexedManifold { // the world COM); use the latter to mirror rapier's solver convention. let dp1 = pt - spose1.translation; let dp2 = pt - spose2.translation; - let contact_vel1 = vel1.linear + gcross_(vel1.angular, dp1); - let contact_vel2 = vel2.linear + gcross_(vel2.angular, dp2); - - // - // Normal part: - // - let torque_dir1 = gcross(dp1, force_dir1); - let torque_dir2 = gcross(dp2, -force_dir1); - - let ii_torque_dir1 = mprops1.inv_inertia_mul(torque_dir1); - let ii_torque_dir2 = mprops2.inv_inertia_mul(torque_dir2); - - let imsum = mprops1.inv_mass + mprops2.inv_mass; - let projected_mass = inv(force_dir1.dot(imsum * force_dir1) - + gdot(ii_torque_dir1, torque_dir1) - + gdot(ii_torque_dir2, torque_dir2)); + let contact_vel1 = vel1.linear + gcross_av(vel1.angular, dp1); + let contact_vel2 = vel2.linear + gcross_av(vel2.angular, dp2); // TODO: handle is_bouncy? - let dist = contact.points_a.at(k).dist; - let normal_rhs_wo_bias = restitution * (contact_vel1 - contact_vel2).dot(force_dir1); - - constraint.elements.at_mut(k).normal_part = TwoBodyConstraintNormalPart { - torque_dir_a: torque_dir1, - ii_torque_dir_a: ii_torque_dir1, - torque_dir_b: torque_dir2, - ii_torque_dir_b: ii_torque_dir2, - rhs: normal_rhs_wo_bias, - rhs_wo_bias: normal_rhs_wo_bias, - impulse: 0.0, - r: projected_mass, - cfm_factor: 1.0, - #[cfg(feature = "dim3")] - _padding0: 0.0, - #[cfg(feature = "dim3")] - _padding1: 0.0, - #[cfg(feature = "dim3")] - _padding2: 0.0, - }; - - // - // Tangent part: - // + let point = constraint.points.at_mut(k); + point.r_a = dp1; + point.r_b = dp2; + // Anchors are stored in the solver-body (COM-centered) frame, matching rapier: + // the iterations recover the world point as `solver_body_pose * local_pt`. + point.local_pt_a = spose1.inverse_transform_point(pt); + point.local_pt_b = spose2.inverse_transform_point(pt); + point.dist = contact.points_a.at(k).dist; + point.normal_vel = restitution * (contact_vel1 - contact_vel2).dot(force_dir1); + point.normal_impulse = 0.0; #[cfg(feature = "dim2")] { - let t_torque_dir1 = gcross(dp1, tangents1.read(0)); - let t_torque_dir2 = gcross(dp2, -tangents1.read(0)); - let t_ii_torque_dir1 = mprops1.inv_inertia * t_torque_dir1; - let t_ii_torque_dir2 = mprops2.inv_inertia * t_torque_dir2; - let r = tangents1.read(0).dot(imsum * tangents1.read(0)) - + gdot(t_ii_torque_dir1, t_torque_dir1) - + gdot(t_ii_torque_dir2, t_torque_dir2); - let rhs_wo_bias = 0.0; - - constraint.elements.at_mut(k).tangent_part = TwoBodyConstraintTangentPart { - torque_dir_a: [Pad::new(t_torque_dir1)], - ii_torque_dir_a: [Pad::new(t_ii_torque_dir1)], - torque_dir_b: [Pad::new(t_torque_dir2)], - ii_torque_dir_b: [Pad::new(t_ii_torque_dir2)], - rhs: [rhs_wo_bias], - rhs_wo_bias: [rhs_wo_bias], - impulse: [0.0], - r: [inv(r)], - }; + point.tangent_impulse = [0.0]; } - #[cfg(feature = "dim3")] { - let mut tangent_part = TwoBodyConstraintTangentPart::default(); - - for j in 0..SUB_LEN { - let t_torque_dir1 = gcross(dp1, tangents1.read(j)); - let t_torque_dir2 = gcross(dp2, -tangents1.read(j)); - let t_ii_torque_dir1 = mprops1.inv_inertia_mul(t_torque_dir1); - let t_ii_torque_dir2 = mprops2.inv_inertia_mul(t_torque_dir2); - let r = tangents1.read(j).dot(imsum * tangents1.read(j)) - + gdot(t_ii_torque_dir1, t_torque_dir1) - + gdot(t_ii_torque_dir2, t_torque_dir2); - let rhs_wo_bias = 0.0; - - tangent_part.torque_dir_a.write(j, Pad::new(t_torque_dir1)); - tangent_part - .ii_torque_dir_a - .write(j, Pad::new(t_ii_torque_dir1)); - tangent_part.torque_dir_b.write(j, Pad::new(t_torque_dir2)); - tangent_part - .ii_torque_dir_b - .write(j, Pad::new(t_ii_torque_dir2)); - tangent_part.rhs.write(j, rhs_wo_bias); - tangent_part.rhs_wo_bias.write(j, rhs_wo_bias); - tangent_part.r.write(j, r); - } - - tangent_part.impulse = Vec2::ZERO; - - // Compute r[2] for 3D (cross term) - tangent_part.r.write( - 2, - 2.0 * ((**tangent_part.torque_dir_a.at(0)) - .dot(**tangent_part.ii_torque_dir_a.at(1)) - + (**tangent_part.torque_dir_b.at(0)) - .dot(**tangent_part.ii_torque_dir_b.at(1))), - ); - - constraint.elements.at_mut(k).tangent_part = tangent_part; + point.tangent_impulse = Vec2::ZERO; } - - // Builder info for warmstarting. Anchors are stored in the solver-body - // (COM-centered) frame, matching rapier — `update_constraint` then - // recovers the world point as `solver_body_pose * local_pt`. - builder.infos.at_mut(k).local_pt_a = spose1.inverse_transform_point(pt); - builder.infos.at_mut(k).local_pt_b = spose2.inverse_transform_point(pt); - builder.infos.at_mut(k).dist = dist; - builder.infos.at_mut(k).normal_vel = normal_rhs_wo_bias; } constraint.len = contact.len; + constraint.compute_effective_masses(); } } impl TwoBodyConstraint { - /// Updates constraint coefficients for a new substep. + /// Scales the accumulated impulses (the warmstart coefficient). #[inline(always)] - pub fn update_constraint( - &mut self, - builder: &TwoBodyConstraintBuilder, - poses: &Slice, - params: &RbdSimParams, - ) { - let body1 = self.solver_body_a as usize; - let body2 = self.solver_body_b as usize; - - let is_static = self.im_a == Vector::ZERO || self.im_b == Vector::ZERO; - let cfm_factor = if is_static { - params.static_contact_cfm_factor() - } else { - params.contact_cfm_factor() - }; - let erp_inv_dt = if is_static { - params.static_contact_erp_inv_dt() - } else { - params.contact_erp_inv_dt() - }; - let inv_dt_val = params.inv_dt(); - let max_corr_velocity = params.max_corrective_velocity(); - let warmstart_coeff = params.warmstart_coefficient; - - let pose1 = poses[body1]; - let pose2 = poses[body2]; - let num_contacts = self.len as usize; - - #[cfg(feature = "dim2")] - let tangents1 = Vec2::new(-self.dir_a.y, self.dir_a.x); - - #[cfg(feature = "dim3")] - let tangents1 = [self.tangent_a, self.dir_a.cross(self.tangent_a)]; - - for j in 0..num_contacts { - // NOTE: the tangent velocity is equivalent to an additional movement of the first body's surface. - let info = builder.infos.at(j); - let p1 = pose1 * info.local_pt_a; // TODO (conveyor belts): + info.tangent_vel * solved_dt; - let p2 = pose2 * info.local_pt_b; - let dist = info.dist + (p1 - p2).dot(self.dir_a); - - // Normal part. + pub fn scale_impulses(&mut self, coeff: f32) { + for k in 0..self.len as usize { + let point = self.points.at_mut(k); + point.normal_impulse *= coeff; + #[cfg(feature = "dim2")] { - let rhs_wo_bias = info.normal_vel + dist.max(0.0) * inv_dt_val; - let rhs_bias = (dist * erp_inv_dt).clamp(-max_corr_velocity, 0.0); - let new_rhs = rhs_wo_bias + rhs_bias; - - self.elements.at_mut(j).normal_part.rhs_wo_bias = rhs_wo_bias; - self.elements.at_mut(j).normal_part.rhs = new_rhs; - // Separated (speculative) points are solved rigidly. - self.elements.at_mut(j).normal_part.cfm_factor = - if dist <= 0.0 { cfm_factor } else { 1.0 }; - self.elements.at_mut(j).normal_part.impulse *= warmstart_coeff; + point.tangent_impulse = [point.tangent_impulse[0] * coeff]; } - - // Tangent parts. + #[cfg(feature = "dim3")] { - #[cfg(feature = "dim2")] - { - let impulse_val = - self.elements.at(j).tangent_part.impulse.read(0) * warmstart_coeff; - let rhs_wo_bias_val = self.elements.at(j).tangent_part.rhs_wo_bias.read(0); - self.elements - .at_mut(j) - .tangent_part - .impulse - .write(0, impulse_val); - let bias = (p1 - p2).dot(tangents1) * inv_dt_val; - self.elements - .at_mut(j) - .tangent_part - .rhs - .write(0, rhs_wo_bias_val + bias); - } - #[cfg(feature = "dim3")] - { - self.elements.at_mut(j).tangent_part.impulse *= warmstart_coeff; - let rhs_wo_bias_0 = self.elements.at(j).tangent_part.rhs_wo_bias.read(0); - let rhs_wo_bias_1 = self.elements.at(j).tangent_part.rhs_wo_bias.read(1); - let bias0 = (p1 - p2).dot(tangents1.read(0)) * inv_dt_val; - self.elements - .at_mut(j) - .tangent_part - .rhs - .write(0, rhs_wo_bias_0 + bias0); - let bias1 = (p1 - p2).dot(tangents1.read(1)) * inv_dt_val; - self.elements - .at_mut(j) - .tangent_part - .rhs - .write(1, rhs_wo_bias_1 + bias1); - } + point.tangent_impulse *= coeff; } } + } - self.cfm_factor = cfm_factor; + /// The total accumulated impulse of the contact point `k`, in world space, as applied to + /// body A (body B receives its opposite). + #[inline(always)] + fn point_impulse(&self, k: usize, tangents: &[Vector; SUB_LEN]) -> Vector { + let point = self.points.at(k); + #[cfg(feature = "dim2")] + return self.dir_a * point.normal_impulse + tangents.read(0) * point.tangent_impulse[0]; + #[cfg(feature = "dim3")] + return self.dir_a * point.normal_impulse + + tangents.read(0) * point.tangent_impulse.x + + tangents.read(1) * point.tangent_impulse.y; } +} - /// Relax-pass refresh: recompute the unbiased normal rhs (speculative term - /// included) from the current solver poses. +impl TwoBodyConstraint { + /// Applies warmstart impulses to a constraint (scatter-style, requires graph coloring). #[inline(always)] - pub fn refresh_rhs_wo_bias( + pub fn warmstart_constraint(&self, solver_vel1: &mut Velocity, solver_vel2: &mut Velocity) { + let tangents = self.tangents(); + let mut force = Vector::ZERO; + #[cfg(feature = "dim2")] + let (mut torque_a, mut torque_b) = (0.0, 0.0); + #[cfg(feature = "dim3")] + let (mut torque_a, mut torque_b) = (Vec3::ZERO, Vec3::ZERO); + + for_contact_point!(k, self.len, { + let f = self.point_impulse(k, &tangents); + force += f; + torque_a += gcross(self.points.at(k).r_a, f); + torque_b += gcross(self.points.at(k).r_b, f); + }); + + solver_vel1.linear += self.im_a * force; + solver_vel1.angular += self.ii_a_mul(torque_a); + solver_vel2.linear -= self.im_b * force; + solver_vel2.angular -= self.ii_b_mul(torque_b); + } + + /// Main constraint solver iteration (Projected Gauss-Seidel). `solve_friction` + /// gates the tangent rows: the stabilization iteration always solves them, the + /// biased pass only when `RbdSimParams::friction_in_bias_pass` is set. + /// + /// The normal right-hand sides are computed here from the current solver poses (rather + /// than stored by a separate pass): with bias for the biased iteration, without for the + /// relaxation iteration, which runs after the positions are integrated. + #[inline(always)] + pub fn solve_constraint_gauss_seidel( &mut self, - builder: &TwoBodyConstraintBuilder, poses: &Slice, params: &RbdSimParams, + solver_vel1: &mut Velocity, + solver_vel2: &mut Velocity, + use_bias: bool, + solve_friction: bool, ) { - let body1 = self.solver_body_a as usize; - let body2 = self.solver_body_b as usize; - let inv_dt_val = params.inv_dt(); - - let pose1 = poses[body1]; - let pose2 = poses[body2]; - let num_contacts = self.len as usize; - - for j in 0..num_contacts { - let info = builder.infos.at(j); - let p1 = pose1 * info.local_pt_a; - let p2 = pose2 * info.local_pt_b; - let dist = info.dist + (p1 - p2).dot(self.dir_a); - self.elements.at_mut(j).normal_part.rhs_wo_bias = - info.normal_vel + dist.max(0.0) * inv_dt_val; - } + let data = ContactSolverData::from_constraint(self); + data.solve_points( + &mut self.points, + poses, + params, + solver_vel1, + solver_vel2, + use_bias, + solve_friction, + ); } } -/// Applies warmstart impulses to a single body (gather-style, no graph coloring required). -#[inline(always)] -pub fn warmstart_body( - body_id: u32, - body_constraint_counts: &Slice, - body_constraint_ids: &Slice, - constraints: &Slice, - solver_vel: &mut Velocity, -) { - let first_constraint_id = if body_id != 0 { - body_constraint_counts[body_id as usize - 1] as usize - } else { - 0 - }; - let last_constraint_id = body_constraint_counts[body_id as usize] as usize; +#[cfg(feature = "dim3")] +type SolverInertia = super::constraint::SymInertia; +#[cfg(feature = "dim2")] +type SolverInertia = f32; + +/// Contact geometry shared by all point rows, cached independently of mutable impulses. +#[derive(Clone, Copy)] +pub(crate) struct ContactSolverData { + pub dir_a: Vector, + pub len: u32, + #[cfg(feature = "dim3")] + pub tangent_a: Vector, + pub limit: f32, + pub im_a: Vector, + pub solver_body_a: u32, + pub im_b: Vector, + pub solver_body_b: u32, + pub ii_a: SolverInertia, + pub ii_b: SolverInertia, + pub vel_slot_a: u32, + pub vel_slot_b: u32, +} - for i in first_constraint_id..last_constraint_id { - let cid = body_constraint_ids[i] as usize; - let constraint = &constraints[cid]; - let solver_body_1 = constraint.solver_body_a; - let dir_a = constraint.dir_a; - let im_a = constraint.im_a; - let im_b = constraint.im_b; - let len = constraint.len as usize; +#[cfg(feature = "dim3")] +type TangentImpulse = Vec2; +#[cfg(feature = "dim2")] +type TangentImpulse = [f32; 1]; +pub(crate) trait ContactPointAccess { + fn read_point(&self, k: usize) -> ContactPoint; + fn write_normal_impulse(&mut self, k: usize, value: f32); + fn write_tangent_impulse(&mut self, k: usize, value: TangentImpulse); +} +impl ContactPointAccess for [ContactPoint; MAX_CONSTRAINTS_PER_MANIFOLD] { + #[inline(always)] + fn read_point(&self, k: usize) -> ContactPoint { + *self.at(k) + } + #[inline(always)] + fn write_normal_impulse(&mut self, k: usize, value: f32) { + self.at_mut(k).normal_impulse = value; + } + #[inline(always)] + fn write_tangent_impulse(&mut self, k: usize, value: TangentImpulse) { + self.at_mut(k).tangent_impulse = value; + } +} +impl ContactSolverData { + #[inline(always)] + pub(crate) fn from_constraint(c: &TwoBodyConstraint) -> Self { + Self { + dir_a: c.dir_a, + len: c.len, + #[cfg(feature = "dim3")] + tangent_a: c.tangent_a, + limit: c.limit, + im_a: c.im_a, + solver_body_a: c.solver_body_a, + im_b: c.im_b, + solver_body_b: c.solver_body_b, + ii_a: c.ii_a, + ii_b: c.ii_b, + vel_slot_a: c.vel_slot_a, + vel_slot_b: c.vel_slot_b, + } + } + #[inline(always)] + pub(crate) fn ii_a_mul(&self, v: crate::AngVector) -> crate::AngVector { #[cfg(feature = "dim2")] - let tangent_a = Vec2::new(-dir_a.y, dir_a.x); + return self.ii_a * v; #[cfg(feature = "dim3")] - let tangent_a = constraint.tangent_a; - - for k in 0..len { - // Warmstart the normal part of the constraint. - { - let c = &constraint.elements.at(k).normal_part; - if solver_body_1 == body_id { - solver_vel.linear += dir_a * im_a * c.impulse; - solver_vel.angular += c.ii_torque_dir_a * c.impulse; - } else { - solver_vel.linear += dir_a * im_b * -c.impulse; - solver_vel.angular += c.ii_torque_dir_b * c.impulse; - } - } - - // Warmstart the tangent parts of the constraint. - { - let c = &constraint.elements.at(k).tangent_part; - - #[cfg(feature = "dim2")] - { - if solver_body_1 == body_id { - solver_vel.linear += tangent_a * im_a * c.impulse.read(0); - solver_vel.angular += c.ii_torque_dir_a.at(0).0 * c.impulse.read(0); - } else { - solver_vel.linear += tangent_a * im_b * -c.impulse.read(0); - solver_vel.angular += c.ii_torque_dir_b.at(0).0 * c.impulse.read(0); - } - } - #[cfg(feature = "dim3")] - { - let tangents_a = [tangent_a, dir_a.cross(tangent_a)]; - if solver_body_1 == body_id { - solver_vel.linear += (tangents_a.read(0) * c.impulse.x - + tangents_a.read(1) * c.impulse.y) - * im_a; - solver_vel.angular += **c.ii_torque_dir_a.at(0) * c.impulse.x - + **c.ii_torque_dir_a.at(1) * c.impulse.y; - } else { - solver_vel.linear += (tangents_a.read(0) * -c.impulse.x - + tangents_a.read(1) * -c.impulse.y) - * im_b; - solver_vel.angular += **c.ii_torque_dir_b.at(0) * c.impulse.x - + **c.ii_torque_dir_b.at(1) * c.impulse.y; - } - } - } - } + return self.ii_a.mul(v); } -} - -impl TwoBodyConstraint { - /// Applies warmstart impulses to a constraint (scatter-style, requires graph coloring). - pub fn warmstart_constraint(&self, solver_vel1: &mut Velocity, solver_vel2: &mut Velocity) { - let dir_a = self.dir_a; - let im_a = self.im_a; - let im_b = self.im_b; - + #[inline(always)] + pub(crate) fn ii_b_mul(&self, v: crate::AngVector) -> crate::AngVector { #[cfg(feature = "dim2")] - let tangent_a = Vec2::new(-dir_a.y, dir_a.x); + return self.ii_b * v; #[cfg(feature = "dim3")] - let tangent_a = self.tangent_a; - - for k in 0..(self.len as usize) { - // Warmstart the normal part of the constraint. - { - let c = &self.elements.at(k).normal_part; - solver_vel1.linear += dir_a * im_a * c.impulse; - solver_vel1.angular += c.ii_torque_dir_a * c.impulse; - solver_vel2.linear += dir_a * im_b * -c.impulse; - solver_vel2.angular += c.ii_torque_dir_b * c.impulse; - } - - // Warmstart the tangent parts of the constraint. - { - let c = &self.elements.at(k).tangent_part; - - #[cfg(feature = "dim2")] - { - solver_vel1.linear += tangent_a * im_a * c.impulse.read(0); - solver_vel1.angular += **c.ii_torque_dir_a.at(0) * c.impulse.read(0); - solver_vel2.linear += tangent_a * im_b * -c.impulse.read(0); - solver_vel2.angular += **c.ii_torque_dir_b.at(0) * c.impulse.read(0); - } - #[cfg(feature = "dim3")] - { - let tangents_a = [tangent_a, dir_a.cross(tangent_a)]; - solver_vel1.linear += (tangents_a.read(0) * c.impulse.x - + tangents_a.read(1) * c.impulse.y) - * im_a; - solver_vel1.angular += **c.ii_torque_dir_a.at(0) * c.impulse.x - + **c.ii_torque_dir_a.at(1) * c.impulse.y; - solver_vel2.linear += (tangents_a.read(0) * -c.impulse.x - + tangents_a.read(1) * -c.impulse.y) - * im_b; - solver_vel2.angular += **c.ii_torque_dir_b.at(0) * c.impulse.x - + **c.ii_torque_dir_b.at(1) * c.impulse.y; - } - } - } + return self.ii_b.mul(v); + } + #[inline(always)] + pub(crate) fn tangents(&self) -> [Vector; SUB_LEN] { + #[cfg(feature = "dim2")] + return [Vector::new(-self.dir_a.y, self.dir_a.x)]; + #[cfg(feature = "dim3")] + return [self.tangent_a, self.dir_a.cross(self.tangent_a)]; } - /// Main constraint solver iteration (Projected Gauss-Seidel). `solve_friction` - /// gates the tangent rows: the stabilization sweep always solves them, the - /// biased pass only when `RbdSimParams::friction_in_bias_pass` is set. #[inline(always)] - pub fn solve_constraint_gauss_seidel( - &mut self, + pub fn solve_points( + &self, + points: &mut P, + poses: &Slice, + params: &RbdSimParams, solver_vel1: &mut Velocity, solver_vel2: &mut Velocity, use_bias: bool, solve_friction: bool, ) { let dir_a = self.dir_a; - let friction_coeff = self.limit; let im_a = self.im_a; let im_b = self.im_b; + let is_static = im_a == Vector::ZERO || im_b == Vector::ZERO; + let (cfm_factor, erp_inv_dt) = if is_static { + ( + params.static_contact_cfm_factor(), + params.static_contact_erp_inv_dt(), + ) + } else { + (params.contact_cfm_factor(), params.contact_erp_inv_dt()) + }; + let inv_dt = params.inv_dt(); + let max_corr_velocity = params.max_corrective_velocity(); + let pose1 = poses[self.solver_body_a as usize]; + let pose2 = poses[self.solver_body_b as usize]; + // Solve the normal parts of the constraint. - for k in 0..(self.len as usize) { - let c = &self.elements.at(k).normal_part; - // Copy values we need after the assignment - let ii_torque_dir_a = c.ii_torque_dir_a; - let ii_torque_dir_b = c.ii_torque_dir_b; + for_contact_point!(k, self.len, { + let point = points.read_point(k); + let p1 = pose1 * point.local_pt_a; + let p2 = pose2 * point.local_pt_b; + let dist = point.dist + (p1 - p2).dot(dir_a); + let rhs_wo_bias = point.normal_vel + dist.max(0.0) * inv_dt; let (rhs, cfm_factor) = if use_bias { - (c.rhs, c.cfm_factor) + // Not `clamp`: its `min <= max` assertion exits the kernel early, which breaks + // the uniform control flow of the fused iterations' barriers on the web. + let rhs_bias = (dist * erp_inv_dt).max(-max_corr_velocity).min(0.0); + // Separated (speculative) points are solved rigidly. + let cfm = if dist <= 0.0 { cfm_factor } else { 1.0 }; + (rhs_wo_bias + rhs_bias, cfm) } else { - (c.rhs_wo_bias, 1.0) + (rhs_wo_bias, 1.0) }; - let dvel = dir_a.dot(solver_vel1.linear) + gdot(c.torque_dir_a, solver_vel1.angular) + + let torque_dir_a = gcross(point.r_a, dir_a); + let torque_dir_b = gcross(point.r_b, -dir_a); + let dvel = dir_a.dot(solver_vel1.linear) + gdot(torque_dir_a, solver_vel1.angular) - dir_a.dot(solver_vel2.linear) - + gdot(c.torque_dir_b, solver_vel2.angular) + + gdot(torque_dir_b, solver_vel2.angular) + rhs; - let new_impulse = cfm_factor * (c.impulse - c.r * dvel).max(0.0); - let delta_impulse = new_impulse - c.impulse; + let impulse = point.normal_impulse; + let new_impulse = cfm_factor * (impulse - point.normal_mass * dvel).max(0.0); + let delta_impulse = new_impulse - impulse; - self.elements.at_mut(k).normal_part.impulse = new_impulse; + points.write_normal_impulse(k, new_impulse); solver_vel1.linear += dir_a * im_a * delta_impulse; - solver_vel1.angular += ii_torque_dir_a * delta_impulse; + solver_vel1.angular += self.ii_a_mul(torque_dir_a) * delta_impulse; solver_vel2.linear += dir_a * im_b * -delta_impulse; - solver_vel2.angular += ii_torque_dir_b * delta_impulse; - } + solver_vel2.angular += self.ii_b_mul(torque_dir_b) * delta_impulse; + }); - // Friction is solved during the stabilization sweep, and during the + // Friction is solved during the stabilization iteration, and during the // biased pass only when `friction_in_bias_pass` is set. if !solve_friction { return; } - #[cfg(feature = "dim2")] - let tangent_a = Vec2::new(-dir_a.y, dir_a.x); - #[cfg(feature = "dim3")] - let tangent_a = self.tangent_a; + let friction_coeff = self.limit; + let tangents = self.tangents(); // Solve the tangent parts of the constraint. - for k in 0..(self.len as usize) { - let limit = friction_coeff * self.elements.at(k).normal_part.impulse; - let c = &self.elements.at(k).tangent_part; - // Copy values we need after the assignment - let ii_torque_dir_a = c.ii_torque_dir_a; - let ii_torque_dir_b = c.ii_torque_dir_b; + for_contact_point!(k, self.len, { + let point = points.read_point(k); + let limit = friction_coeff * point.normal_impulse; #[cfg(feature = "dim2")] { - let dvel = tangent_a.dot(solver_vel1.linear) - + gdot(**c.torque_dir_a.at(0), solver_vel1.angular) - - tangent_a.dot(solver_vel2.linear) - + gdot(**c.torque_dir_b.at(0), solver_vel2.angular) - + c.rhs_wo_bias.read(0); + let t = tangents.read(0); + let torque_dir_a = gcross(point.r_a, t); + let torque_dir_b = gcross(point.r_b, -t); + let dvel = t.dot(solver_vel1.linear) + gdot(torque_dir_a, solver_vel1.angular) + - t.dot(solver_vel2.linear) + + gdot(torque_dir_b, solver_vel2.angular); + let impulse = point.tangent_impulse[0]; // NOTE: don’t use clamp since it can panic. - let new_impulse = (c.impulse.read(0) - c.r.read(0) * dvel) - .max(-limit) - .min(limit); - let delta_impulse = new_impulse - c.impulse.read(0); - - self.elements - .at_mut(k) - .tangent_part - .impulse - .write(0, new_impulse); - - solver_vel1.linear += tangent_a * im_a * delta_impulse; - solver_vel1.angular += **ii_torque_dir_a.at(0) * delta_impulse; - - solver_vel2.linear += tangent_a * im_b * -delta_impulse; - solver_vel2.angular += **ii_torque_dir_b.at(0) * delta_impulse; + let new_impulse = (impulse - point.tangent_mass * dvel).max(-limit).min(limit); + let delta_impulse = new_impulse - impulse; + + points.write_tangent_impulse(k, [new_impulse]); + + solver_vel1.linear += t * im_a * delta_impulse; + solver_vel1.angular += self.ii_a_mul(torque_dir_a) * delta_impulse; + + solver_vel2.linear += t * im_b * -delta_impulse; + solver_vel2.angular += self.ii_b_mul(torque_dir_b) * delta_impulse; } #[cfg(feature = "dim3")] { - let tangents_a = [tangent_a, dir_a.cross(tangent_a)]; - let dvel_0 = tangents_a.read(0).dot(solver_vel1.linear) - + gdot(**c.torque_dir_a.at(0), solver_vel1.angular) - - tangents_a.read(0).dot(solver_vel2.linear) - + gdot(**c.torque_dir_b.at(0), solver_vel2.angular) - + c.rhs_wo_bias.read(0); - let dvel_1 = tangents_a.read(1).dot(solver_vel1.linear) - + gdot(**c.torque_dir_a.at(1), solver_vel1.angular) - - tangents_a.read(1).dot(solver_vel2.linear) - + gdot(**c.torque_dir_b.at(1), solver_vel2.angular) - + c.rhs_wo_bias.read(1); - - let k11 = c.r.read(0); - let k22 = c.r.read(1); - let k12 = c.r.read(2) * 0.5; + let t0 = tangents.read(0); + let t1 = tangents.read(1); + let torque_dir_a0 = gcross(point.r_a, t0); + let torque_dir_b0 = gcross(point.r_b, -t0); + let torque_dir_a1 = gcross(point.r_a, t1); + let torque_dir_b1 = gcross(point.r_b, -t1); + let dvel_0 = t0.dot(solver_vel1.linear) + gdot(torque_dir_a0, solver_vel1.angular) + - t0.dot(solver_vel2.linear) + + gdot(torque_dir_b0, solver_vel2.angular); + let dvel_1 = t1.dot(solver_vel1.linear) + gdot(torque_dir_a1, solver_vel1.angular) + - t1.dot(solver_vel2.linear) + + gdot(torque_dir_b1, solver_vel2.angular); + + let k11 = point.tangent_k[0]; + let k22 = point.tangent_k[1]; + let k12 = point.tangent_k[2] * 0.5; let inv_det = maybe_inv(k11 * k22 - k12 * k12); let delta_impulse = Vec2::new( (k22 * dvel_0 - k12 * dvel_1) * inv_det, (k11 * dvel_1 - k12 * dvel_0) * inv_det, ); - let mut new_impulse = c.impulse - delta_impulse; - new_impulse = cap_magnitude(new_impulse, limit); - - let delta_impulse = new_impulse - c.impulse; - self.elements.at_mut(k).tangent_part.impulse = new_impulse; - - solver_vel1.linear += (tangents_a.read(0) * delta_impulse.x - + tangents_a.read(1) * delta_impulse.y) - * im_a; - solver_vel1.angular += **ii_torque_dir_a.at(0) * delta_impulse.x - + **ii_torque_dir_a.at(1) * delta_impulse.y; - - solver_vel2.linear += (tangents_a.read(0) * -delta_impulse.x - + tangents_a.read(1) * -delta_impulse.y) - * im_b; - solver_vel2.angular += **ii_torque_dir_b.at(0) * delta_impulse.x - + **ii_torque_dir_b.at(1) * delta_impulse.y; + let impulse = point.tangent_impulse; + let new_impulse = cap_magnitude(impulse - delta_impulse, limit); + let delta_impulse = new_impulse - impulse; + points.write_tangent_impulse(k, new_impulse); + + let lin = t0 * delta_impulse.x + t1 * delta_impulse.y; + solver_vel1.linear += lin * im_a; + solver_vel1.angular += self + .ii_a_mul(torque_dir_a0 * delta_impulse.x + torque_dir_a1 * delta_impulse.y); + + solver_vel2.linear -= lin * im_b; + solver_vel2.angular += self + .ii_b_mul(torque_dir_b0 * delta_impulse.x + torque_dir_b1 * delta_impulse.y); } + }); + } +} + +#[cfg(all(test, not(target_arch_is_gpu)))] +#[path = "../tests/contact_solver.rs"] +mod tests; + +impl TwoBodyConstraint { + #[inline(always)] + pub fn init_header(&mut self, manifold: &IndexedManifold, mprops: &Slice) { + manifold.contact_to_constraint_header(mprops, self); + } + /// Initializes the contact points of a constraint whose header was initialized from + /// `manifold`. + #[inline(always)] + pub fn init_points_with_poses( + &mut self, + manifold: &IndexedManifold, + collider_a: Pose, + solver_a: Pose, + solver_b: Pose, + vels: &Slice, + ) { + manifold.contact_to_constraint_points(collider_a, solver_a, solver_b, vels, self); + } + /// Applies the velocity slots and mass scales chosen by mass splitting. + #[inline(always)] + pub fn apply_link(&mut self, link: &ContactLink) { + self.vel_slot_a = link.vel_slot_a; + self.vel_slot_b = link.vel_slot_b; + self.restitution = link.restitution; + self.im_a *= link.mass_scale_a; + self.im_b *= link.mass_scale_b; + #[cfg(feature = "dim2")] + { + self.ii_a *= link.mass_scale_a; + self.ii_b *= link.mass_scale_b; + } + #[cfg(feature = "dim3")] + { + self.ii_a.scale(link.mass_scale_a); + self.ii_b.scale(link.mass_scale_b); + } + } + #[inline(always)] + pub fn recycle_from(&mut self, old: &Self, vel1: &Velocity, vel2: &Velocity) { + self.dir_a = old.dir_a; + #[cfg(feature = "dim3")] + { + self.tangent_a = old.tangent_a; } + self.len = old.len; + for k in 0..old.len as usize { + let mut point = *old.points.at(k); + let v1 = vel1.linear + gcross_av(vel1.angular, point.r_a); + let v2 = vel2.linear + gcross_av(vel2.angular, point.r_b); + point.normal_vel = self.restitution * (v1 - v2).dot(old.dir_a); + self.points.write(k, point); + } + self.compute_effective_masses(); + } + #[inline(always)] + pub fn local_anchors(&self, k: usize) -> (Vector, Vector) { + let p = self.points.at(k); + (p.local_pt_a, p.local_pt_b) + } + // Pass only the friction impulse through the matching loop, rather than + // copying the entire old manifold for every matched point. + #[cfg(feature = "dim3")] + #[inline(always)] + pub fn friction_warmstart(&self, k: usize) -> glamx::Vec4 { + let t = self.points.at(k).tangent_impulse; + (self.tangent_a * t.x + self.dir_a.cross(self.tangent_a) * t.y).extend(0.0) } + #[cfg(feature = "dim2")] + #[inline(always)] + pub fn friction_warmstart(&self, k: usize) -> f32 { + self.points.at(k).tangent_impulse[0] + } + #[cfg(feature = "dim3")] + #[inline(always)] + pub fn transfer_friction_point(&mut self, impulse: glamx::Vec4, k: usize) { + let world = impulse.truncate(); + self.points.at_mut(k).tangent_impulse = Vec2::new( + world.dot(self.tangent_a), + world.dot(self.dir_a.cross(self.tangent_a)), + ); + } + #[cfg(feature = "dim2")] + #[inline(always)] + pub fn transfer_friction_point(&mut self, impulse: f32, k: usize) { + self.points.at_mut(k).tangent_impulse = [impulse]; + } + #[inline(always)] + pub fn finish_friction_transfer(&mut self) {} } diff --git a/src_rbd_shaders/dynamics/tile_kernels.rs b/src_rbd_shaders/dynamics/tile_kernels.rs new file mode 100644 index 00000000..2dbb1374 --- /dev/null +++ b/src_rbd_shaders/dynamics/tile_kernels.rs @@ -0,0 +1,469 @@ +//! Contact-tile kernels shared by the friction models: building the constraints, solving them +//! and warmstarting from them. +//! +//! Included by each model's module, which provides `ContactTile`, `TwoBodyConstraint` (its +//! constraint type) and the tile accessors (`header`, `read_constraint`, `write_constraint`, +//! `solve_constraint`, `tile_warmstart`, `constraint_warmstart`). + +use super::*; +use crate::broad_phase::ContactPlan; +use crate::dynamics::contact_tiles::{TAIL_COLOR, tile_lane}; +use crate::dynamics::mass_splitting::{HubCounts, NOT_A_HUB}; +use crate::dynamics::{ + ContactLink, ContactRecycleOffsets, ContactRecycleState, WorldMassProperties, decode_bias_mode, +}; +use crate::queries::IndexedManifold; +use crate::utils::BatchIndices; +use khal_std::glamx::UVec3; +use khal_std::iter::StepRng; +use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::control_barrier; + +const WORKGROUP_SIZE: u32 = 64; + +/// The number of constraints in the tiles: the colored ones, then the unsolved ones (in the +/// last color bucket). +#[inline(always)] +fn constraint_count(buckets: &[u32], ids: &BatchIndices) -> u32 { + buckets.read((ids.solver_color_buckets_stride * ids.num_batches - 1) as usize) +} + +/// The constraints of `color`, over every batch. Buckets are color-major and hold exclusive +/// ends. +#[inline(always)] +fn color_constraints(buckets: &[u32], ids: &BatchIndices, color: u32) -> (u32, u32) { + let nb = ids.num_batches; + ( + buckets.read((color * nb - 1) as usize), + buckets.read(((color + 1) * nb - 1) as usize), + ) +} + +/// Makes every lane's storage writes visible to the workgroup before the next color. +#[inline(always)] +fn color_barrier() { + control_barrier::< + { khal_std::memory::Scope::Workgroup as u32 }, + { khal_std::memory::Scope::QueueFamily as u32 }, + { + khal_std::memory::Semantics::UNIFORM_MEMORY.bits() + | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() + }, + >(); +} + +/// Caches the warmstart of the constraint at `index`: the velocity changes of its bodies, +/// tagged with their index. +#[inline(always)] +fn write_warmstart(tiles: &mut [ContactTile], index: usize, (a, b): (Velocity, Velocity)) { + let (tile, lane) = tile_lane(index); + tiles.at_mut(tile).warmstart.bodies.write(lane, [a, b]); +} + +/// Builds every constraint in color order, gives it the impulses of its previous-frame +/// constraint (or that constraint's contacts when they are recycled), and caches its warmstart +/// for the first substep. Runs after the links are sorted by color. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_prepare_constraints( + #[spirv(global_invocation_id)] gid: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] sorted_links: &[ContactLink], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts: &[IndexedManifold], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] mprops: &[WorldMassProperties], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] solver_poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] vels: &[Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] states: &[ContactRecycleState], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] old_tiles: &[ContactTile], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] tiles: &mut [ContactTile], + #[spirv(uniform, descriptor_set = 0, binding = 8)] plan: &ContactPlan, + #[spirv(uniform, descriptor_set = 0, binding = 9)] recycle_offsets: &ContactRecycleOffsets, +) { + let index = gid.x as usize; + if gid.x >= plan.bound { + return; + } + // The links past the last constraint are inactive (see `gpu_color_buckets_count`). + let link = sorted_links.read(index); + if link.len == 0 { + return; + } + let i = link.contact as usize; + let im = contacts.at(i); + let mut c = TwoBodyConstraint::default(); + c.init_header(im, &Slice(mprops, 0)); + c.apply_link(&link); + if link.recycled == 0 { + let pose_a = states.read(recycle_offsets.new_base as usize + i).pose_a; + c.init_points_with_poses( + im, + pose_a, + solver_poses.read(link.solver_body_a as usize), + solver_poses.read(link.solver_body_b as usize), + &Slice(vels, 0), + ); + } + if link.previous_constraint != u32::MAX { + let old = read_constraint(old_tiles, link.previous_constraint_index as usize); + if link.recycled != 0 { + c.recycle_from( + &old, + vels.at(link.solver_body_a as usize), + vels.at(link.solver_body_b as usize), + ); + } else { + transfer_contact_impulses(&old, &mut c); + } + } + write_constraint(tiles, index, &c); + write_warmstart(tiles, index, constraint_warmstart(&c)); +} + +/// Gives each new contact point the impulses of the nearest previous one (by local anchors, +/// within 10cm). +#[inline(always)] +fn transfer_contact_impulses(old: &TwoBodyConstraint, new: &mut TwoBodyConstraint) { + let sq_threshold = 1.0e-1 * 1.0e-1; + + // The points of one small manifold are all within 10cm of each other: taking the first + // candidate instead of the nearest would hand every point the same impulse. + for k_new in 0..(new.len as usize) { + let (pt_new_a, pt_new_b) = new.local_anchors(k_new); + let mut best_k_old = usize::MAX; + let mut best_sq = sq_threshold; + for k_old in 0..(old.len as usize) { + let (pt_old_a, pt_old_b) = old.local_anchors(k_old); + let dpt_a = pt_old_a - pt_new_a; + let dpt_b = pt_old_b - pt_new_b; + let sq = dpt_a.dot(dpt_a).max(dpt_b.dot(dpt_b)); + if sq < best_sq { + best_sq = sq; + best_k_old = k_old; + } + } + if best_k_old != usize::MAX { + new.points.at_mut(k_new).normal_impulse = old.points.at(best_k_old).normal_impulse; + new.transfer_friction_point(old.friction_warmstart(best_k_old), k_new); + } + } + new.finish_friction_transfer(); +} + +/// Scales the accumulated impulses of every constraint by the warmstart coefficient and caches +/// their warmstart (only dispatched when that coefficient isn't 1). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_scale_constraint_impulses( + #[spirv(global_invocation_id)] gid: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] buckets: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] tiles: &mut [ContactTile], + #[spirv(uniform, descriptor_set = 0, binding = 2)] ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 3)] params: &RbdSimParams, +) { + let index = gid.x as usize; + if gid.x >= constraint_count(buckets, ids) { + return; + } + let mut c = read_constraint(tiles, index); + c.scale_impulses(params.warmstart_coefficient); + write_constraint(tiles, index, &c); + write_warmstart(tiles, index, constraint_warmstart(&c)); +} + +/// Solves the constraints of one color, or of every color from [`TAIL_COLOR`] for the tail +/// variants (one workgroup, with a barrier between colors). The caching variants also store +/// each constraint's warmstart for the next substep. +macro_rules! tile_solver_kernel { + ($name:ident, $mode:ident, $mode_value:expr, $cache:expr, $tail:expr) => { + #[spirv_bindgen] + #[spirv(compute(threads(64)))] + pub fn $name( + #[spirv(global_invocation_id)] gid: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] tiles: &mut [ContactTile], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] buckets: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] solver_body_poses: &[Pose], + #[spirv(uniform, descriptor_set = 0, binding = 4)] curr_color: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 6)] $mode: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &RbdSimParams, + ) { + let _ = $mode; + let (use_bias, solve_friction) = $mode_value; + let poses = Slice(solver_body_poses, 0); + let first = if $tail { TAIL_COLOR } else { *curr_color }; + for color in first..=*curr_color { + let (start, end) = color_constraints(buckets, ids, color); + for index in StepRng::new(start + gid.x..end, num_workgroups.x * WORKGROUP_SIZE) { + let index = index as usize; + let h = header(tiles, index); + let mut a = solver_vels.read(h.vel_slot_a as usize); + let mut b = solver_vels.read(h.vel_slot_b as usize); + solve_constraint( + &h, + tiles, + index, + &poses, + params, + &mut a, + &mut b, + use_bias, + solve_friction, + ); + solver_vels.write(h.vel_slot_a as usize, a); + solver_vels.write(h.vel_slot_b as usize, b); + if $cache { + let warmstart = tile_warmstart(tiles, index, &h); + write_warmstart(tiles, index, warmstart); + } + } + if $tail { + color_barrier(); + } + } + } + }; +} + +tile_solver_kernel!( + gpu_solve_constraints, + mode, + decode_bias_mode(*mode), + false, + false +); +tile_solver_kernel!( + gpu_solve_constraints_biased, + mode, + (true, false), + false, + false +); +tile_solver_kernel!( + gpu_solve_constraints_unbiased, + mode, + (false, true), + false, + false +); +tile_solver_kernel!( + gpu_solve_constraints_unbiased_cached, + mode, + (false, true), + true, + false +); +tile_solver_kernel!( + gpu_solve_constraints_tail, + mode, + decode_bias_mode(*mode), + false, + true +); +tile_solver_kernel!( + gpu_solve_constraints_tail_biased, + mode, + (true, false), + false, + true +); +tile_solver_kernel!( + gpu_solve_constraints_tail_unbiased, + mode, + (false, true), + false, + true +); +tile_solver_kernel!( + gpu_solve_constraints_tail_unbiased_cached, + mode, + (false, true), + true, + true +); + +/// Solves every color with one 64-lane workgroup per batch, with a barrier between colors. +/// Used for small scenes where the contact count is small wrt. the environment count. +macro_rules! fused_tile_solver_kernel { + ($name:ident, $cache:expr) => { + #[spirv_bindgen] + #[spirv(compute(threads(64)))] + pub fn $name( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] tiles: &mut [ContactTile], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] buckets: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] solver_body_poses: &[Pose], + #[spirv(uniform, descriptor_set = 0, binding = 4)] num_colors: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 6)] mode: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &RbdSimParams, + ) { + let lane = invocation_id.x; + let batch = invocation_id.y; + let nb = ids.num_batches; + let poses = Slice(solver_body_poses, 0); + let (use_bias, solve_friction) = decode_bias_mode(*mode); + for color in 1..=*num_colors { + let bucket = (color * nb + batch) as usize; + let start = buckets.read(bucket - 1); + let end = buckets.read(bucket); + #[cfg(not(feature = "web-compat"))] + if start == end { + continue; + } + for index in StepRng::new(start + lane..end, WORKGROUP_SIZE) { + let index = index as usize; + let h = header(tiles, index); + let mut a = solver_vels.read(h.vel_slot_a as usize); + let mut b = solver_vels.read(h.vel_slot_b as usize); + solve_constraint( + &h, + tiles, + index, + &poses, + params, + &mut a, + &mut b, + use_bias, + solve_friction, + ); + solver_vels.write(h.vel_slot_a as usize, a); + solver_vels.write(h.vel_slot_b as usize, b); + if $cache { + let warmstart = tile_warmstart(tiles, index, &h); + write_warmstart(tiles, index, warmstart); + } + } + color_barrier(); + } + } + }; +} + +fused_tile_solver_kernel!(gpu_solve_constraints_fused, false); +fused_tile_solver_kernel!(gpu_solve_constraints_fused_cached, true); + +/// Adds a constraint's warmstart to the velocities of its bodies. +#[inline(always)] +fn apply_warmstart(tiles: &[ContactTile], solver_vels: &mut [Velocity], index: usize) { + let h = header(tiles, index); + let (da, db) = tile_warmstart(tiles, index, &h); + let mut a = solver_vels.read(h.vel_slot_a as usize); + let mut b = solver_vels.read(h.vel_slot_b as usize); + a.linear += da.linear; + a.angular += da.angular; + b.linear += db.linear; + b.angular += db.angular; + solver_vels.write(h.vel_slot_a as usize, a); + solver_vels.write(h.vel_slot_b as usize, b); +} + +/// Applies the warmstart of one color's constraints to their bodies (scatter-style, for scenes +/// with multibodies, whose warmstart can't be gathered per body). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_warmstart_constraints( + #[spirv(global_invocation_id)] gid: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] tiles: &[ContactTile], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] buckets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] curr_color: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 4)] ids: &BatchIndices, +) { + let (start, end) = color_constraints(buckets, ids, *curr_color); + for index in StepRng::new(start + gid.x..end, num_workgroups.x * WORKGROUP_SIZE) { + apply_warmstart(tiles, solver_vels, index as usize); + } +} + +/// [`gpu_warmstart_constraints`] for every color, with one 64-lane workgroup per batch. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_warmstart_constraints_fused( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] tiles: &[ContactTile], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] buckets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] num_colors: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 4)] ids: &BatchIndices, +) { + let lane = invocation_id.x; + let batch = invocation_id.y; + let nb = ids.num_batches; + for color in 1..=*num_colors { + let bucket = (color * nb + batch) as usize; + let start = buckets.read(bucket - 1); + let end = buckets.read(bucket); + #[cfg(not(feature = "web-compat"))] + if start == end { + continue; + } + for index in StepRng::new(start + lane..end, WORKGROUP_SIZE) { + apply_warmstart(tiles, solver_vels, index as usize); + } + color_barrier(); + } +} + +/// Sums the cached warmstart of each body's constraints, in its adjacency-list order (split +/// bodies are warmstarted through their sub-bodies). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_gather_warmstart_velocities( + #[spirv(global_invocation_id)] gid: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] counts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraint_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraint_indices: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] tiles: &[ContactTile], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] hub_first_slot: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 6)] ids: &BatchIndices, +) { + let body = gid.x as usize; + if gid.x < ids.bodies_len * ids.num_batches && hub_first_slot.read(body) == NOT_A_HUB { + let start = if body == 0 { 0 } else { counts.read(body - 1) }; + let end = counts.read(body); + let mut vel = solver_vels.read(body); + for entry in start..end { + let index = + constraint_indices.read(constraint_ids.read(entry as usize) as usize) as usize; + let (tile, lane) = tile_lane(index); + let [a, b] = tiles.at(tile).warmstart.bodies.read(lane); + let dv = if a.padding1 == gid.x { a } else { b }; + vel.linear += dv.linear; + vel.angular += dv.angular; + } + solver_vels.write(body, vel); + } +} + +/// Applies the warmstart of each split body's constraint to its sub-body (the body's own +/// warmstart gather skips them). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_hub_warmstart_constraints( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] tiles: &[ContactTile], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] hub_slot_constraint: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] hub_counts: &HubCounts, + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] constraint_indices: &[u32], +) { + let slot = invocation_id.x; + if slot < hub_counts.live_slots() { + let vel_slot = hub_counts.base + slot; + let index = + constraint_indices.read(hub_slot_constraint.read(slot as usize) as usize) as usize; + let h = header(tiles, index); + let (da, db) = tile_warmstart(tiles, index, &h); + let d = if h.vel_slot_a == vel_slot { da } else { db }; + let mut vel = solver_vels.read(vel_slot as usize); + vel.linear += d.linear; + vel.angular += d.angular; + solver_vels.write(vel_slot as usize, vel); + } +} diff --git a/src_rbd_shaders/dynamics/twist_friction.rs b/src_rbd_shaders/dynamics/twist_friction.rs new file mode 100644 index 00000000..7b1e7145 --- /dev/null +++ b/src_rbd_shaders/dynamics/twist_friction.rs @@ -0,0 +1,582 @@ +//! Manifold twist friction with dedicated compact storage. +//! +//! Normal rows store one frozen lever arm. The other differs by the same +//! frozen COM offset for every point, including recycled contacts. Sliding and +//! twist impulses/masses are stored once per manifold, never per point. +use super::solver_utils::{ + ContactPointAccess, ContactSolverData, cap_magnitude, compute_tangent_contact_directions, + maybe_inv, +}; +use super::{ContactLink, ContactPoint, RbdSimParams, SymInertia, Velocity, WorldMassProperties}; +use crate::queries::IndexedManifold; +use crate::utils::Slice; +use crate::{Pose, Vector}; +use crunchy::unroll; +use glamx::{Quat, Vec2, Vec3, Vec4}; +use khal_std::index::MaybeIndexUnchecked; + +/// A normal-only contact row (32 bytes), in the frozen world frame. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct TwistContactPoint { + pub r_a: Vec3, + pub normal_impulse: f32, + pub normal_mass: f32, + pub dist: f32, + pub normal_vel: f32, + pub radius: f32, +} + +/// One sliding row pair and one normal-axis angular row for the manifold. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct TwistFriction { + pub center_a: Vec3, + pub twist_mass: f32, + /// k00, k11, 2*k01 of the central sliding rows. + pub tangent_k: Vec3, + pub twist_impulse: f32, + pub tangent_impulse: Vec2, + pub _padding: [u32; 2], +} + +/// Compact twist-friction manifold (384 bytes versus Coulomb's 544). +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct TwistConstraint { + /// Number of active contact points in this manifold. + pub len: u32, + /// Index of body A (poses, mass properties). + pub solver_body_a: u32, + /// Index of body B. + pub solver_body_b: u32, + /// Index of body A's velocity in the solver velocity buffer: `solver_body_a`, or one of + /// the sub-body slots when body A is split (see `mass_splitting`). + pub vel_slot_a: u32, + /// Index of body B's velocity in the solver velocity buffer. + pub vel_slot_b: u32, + /// Collider A of the source manifold, which identifies it across steps with + /// `warmstart_collider_b` and `warmstart_subshape` (for the warmstart and color seeding). + pub warmstart_collider_a: u32, + /// Collider B of the source manifold. + pub warmstart_collider_b: u32, + /// [`IndexedManifold::subshape`] of the source manifold. + /// + /// [`IndexedManifold::subshape`]: crate::queries::IndexedManifold::subshape + pub warmstart_subshape: u32, + /// Contact normal from body A's perspective (points away from A). + pub dir_a: Vector, + /// Friction coefficient. + pub limit: f32, + /// First friction direction (orthogonal to the normal); the second is `dir_a × tangent_a`. + pub tangent_a: Vector, + /// Combined restitution coefficient. + pub restitution: f32, + /// Inverse mass of body A along each axis. + pub im_a: Vector, + pub _padding_a: u32, + /// Inverse mass of body B along each axis. + pub im_b: Vector, + pub _padding_b: u32, + /// World-space inverse inertia of body A. + pub ii_a: SymInertia, + /// World-space inverse inertia of body B. + pub ii_b: SymInertia, + pub points: [TwistContactPoint; 4], + /// Frozen r_b - r_a, shared by all contact points. + pub offset_b: Vec3, + pub _padding_offset: u32, + /// Body rotations when the contact geometry was frozen. Local anchors are + /// recovered from these and the frozen lever arms, including after recycling. + pub frame_a: Quat, + pub frame_b: Quat, + pub friction: TwistFriction, +} + +impl TwistConstraint { + #[inline(always)] + pub fn init_header(&mut self, manifold: &IndexedManifold, mprops: &Slice) { + let a = mprops[manifold.bodies.x as usize]; + let b = mprops[manifold.bodies.y as usize]; + self.im_a = a.inv_mass; + self.im_b = b.inv_mass; + self.ii_a = SymInertia::from_mat4(a.inv_inertia); + self.ii_b = SymInertia::from_mat4(b.inv_inertia); + self.solver_body_a = manifold.bodies.x; + self.solver_body_b = manifold.bodies.y; + self.vel_slot_a = self.solver_body_a; + self.vel_slot_b = self.solver_body_b; + self.limit = manifold.friction; + self.restitution = manifold.restitution; + self.warmstart_collider_a = manifold.colliders.x; + self.warmstart_collider_b = manifold.colliders.y; + self.warmstart_subshape = manifold.subshape; + self.len = manifold.contact.len; + } + #[inline(always)] + pub fn init_contact( + &mut self, + manifold: &IndexedManifold, + mprops: &Slice, + collider_poses: &Slice, + solver_poses: &Slice, + vels: &Slice, + ) { + self.init_contact_with_poses( + manifold, + mprops, + collider_poses[manifold.colliders.x as usize], + solver_poses[manifold.bodies.x as usize], + solver_poses[manifold.bodies.y as usize], + vels, + ); + } + #[inline(always)] + pub fn init_contact_with_poses( + &mut self, + manifold: &IndexedManifold, + mprops: &Slice, + ca: Pose, + a: Pose, + b: Pose, + vels: &Slice, + ) { + self.init_header(manifold, mprops); + self.init_points_with_poses(manifold, ca, a, b, vels); + } + /// Initializes the contact points of a constraint whose header was initialized from + /// `manifold`, and its effective masses. + #[inline(always)] + pub fn init_points_with_poses( + &mut self, + manifold: &IndexedManifold, + ca: Pose, + a: Pose, + b: Pose, + vels: &Slice, + ) { + let va = vels[self.solver_body_a as usize]; + let vb = vels[self.solver_body_b as usize]; + self.dir_a = -(ca.rotation * manifold.contact.normal_a); + self.tangent_a = compute_tangent_contact_directions(self.dir_a).read(0); + self.offset_b = a.translation - b.translation; + self.friction = TwistFriction::default(); + self.frame_a = a.rotation; + self.frame_b = b.rotation; + for k in 0..self.len as usize { + let contact = manifold.contact.points_a.at(k); + let pt = ca * (contact.pt + manifold.contact.normal_a * contact.dist / 2.0); + let ra = pt - a.translation; + let rb = ra + self.offset_b; + self.points.write( + k, + TwistContactPoint { + r_a: ra, + normal_impulse: 0.0, + normal_mass: 0.0, + dist: contact.dist, + normal_vel: self.restitution + * (va.linear + va.angular.cross(ra) - vb.linear - vb.angular.cross(rb)) + .dot(self.dir_a), + radius: 0.0, + }, + ); + } + self.compute_effective_masses(); + } + /// Applies the velocity slots and mass scales chosen by mass splitting. + #[inline(always)] + pub fn apply_link(&mut self, link: &ContactLink) { + self.vel_slot_a = link.vel_slot_a; + self.vel_slot_b = link.vel_slot_b; + self.restitution = link.restitution; + self.im_a *= link.mass_scale_a; + self.im_b *= link.mass_scale_b; + self.ii_a.scale(link.mass_scale_a); + self.ii_b.scale(link.mass_scale_b); + } + #[inline(always)] + pub fn compute_effective_masses(&mut self) { + if self.len == 0 { + return; + } + let mut center = Vec3::ZERO; + let im = self.im_a + self.im_b; + for k in 0..self.len as usize { + let p = self.points.at_mut(k); + let a = p.r_a.cross(self.dir_a); + let b = (p.r_a + self.offset_b).cross(-self.dir_a); + let kn = + self.dir_a.dot(im * self.dir_a) + self.ii_a.mul(a).dot(a) + self.ii_b.mul(b).dot(b); + p.normal_mass = if kn == 0.0 { 0.0 } else { 1.0 / kn }; + center += p.r_a; + } + center /= self.len as f32; + self.friction.center_a = center; + for k in 0..self.len as usize { + let p = self.points.at_mut(k); + p.radius = (p.r_a - center).length(); + } + let t0 = self.tangent_a; + let t1 = self.dir_a.cross(t0); + let ra = center; + let rb = center + self.offset_b; + let a0 = ra.cross(t0); + let a1 = ra.cross(t1); + let b0 = rb.cross(-t0); + let b1 = rb.cross(-t1); + self.friction.tangent_k = Vec3::new( + t0.dot(im * t0) + a0.dot(self.ii_a.mul(a0)) + b0.dot(self.ii_b.mul(b0)), + t1.dot(im * t1) + a1.dot(self.ii_a.mul(a1)) + b1.dot(self.ii_b.mul(b1)), + 2.0 * (self.ii_a.mul(a0).dot(a1) + self.ii_b.mul(b0).dot(b1)), + ); + self.friction.twist_mass = maybe_inv( + self.dir_a.dot(self.ii_a.mul(self.dir_a)) + self.dir_a.dot(self.ii_b.mul(self.dir_a)), + ); + if self.len == 1 { + self.friction.twist_impulse = 0.0; + } + } + #[inline(always)] + pub fn recycle_from(&mut self, old: &Self, a: &Velocity, b: &Velocity) { + self.dir_a = old.dir_a; + self.tangent_a = old.tangent_a; + self.len = old.len; + self.offset_b = old.offset_b; + self.friction = old.friction; + self.frame_a = old.frame_a; + self.frame_b = old.frame_b; + for k in 0..self.len as usize { + let mut p = old.points.read(k); + p.normal_vel = self.restitution + * (a.linear + a.angular.cross(p.r_a) + - b.linear + - b.angular.cross(p.r_a + self.offset_b)) + .dot(self.dir_a); + self.points.write(k, p); + } + self.compute_effective_masses(); + } + #[inline(always)] + pub fn local_anchors(&self, k: usize) -> (Vec3, Vec3) { + let r = self.points.at(k).r_a; + ( + self.frame_a.inverse() * r, + self.frame_b.inverse() * (r + self.offset_b), + ) + } + #[inline(always)] + pub fn friction_warmstart(&self, _k: usize) -> Vec4 { + let t = self.friction.tangent_impulse; + (self.tangent_a * t.x + self.dir_a.cross(self.tangent_a) * t.y) + .extend(self.friction.twist_impulse) + } + #[inline(always)] + pub fn transfer_friction_point(&mut self, impulse: Vec4, _k: usize) { + let world = impulse.truncate(); + self.friction.tangent_impulse += Vec2::new( + world.dot(self.tangent_a), + world.dot(self.dir_a.cross(self.tangent_a)), + ); + self.friction.twist_impulse += impulse.w; + } + #[inline(always)] + pub fn finish_friction_transfer(&mut self) { + if self.len != 0 { + self.friction.tangent_impulse /= self.len as f32; + self.friction.twist_impulse = if self.len > 1 { + self.friction.twist_impulse / self.len as f32 + } else { + 0.0 + }; + } + } + #[inline(always)] + pub fn scale_impulses(&mut self, scale: f32) { + for k in 0..self.len as usize { + self.points.at_mut(k).normal_impulse *= scale; + } + self.friction.tangent_impulse *= scale; + self.friction.twist_impulse *= scale; + } + #[inline(always)] + pub(crate) fn solver_data(&self) -> ContactSolverData { + ContactSolverData { + dir_a: self.dir_a, + len: self.len, + tangent_a: self.tangent_a, + limit: self.limit, + im_a: self.im_a, + im_b: self.im_b, + ii_a: self.ii_a, + ii_b: self.ii_b, + solver_body_a: self.solver_body_a, + solver_body_b: self.solver_body_b, + vel_slot_a: self.vel_slot_a, + vel_slot_b: self.vel_slot_b, + } + } + #[inline(always)] + pub fn warmstart_constraint(&self, a: &mut Velocity, b: &mut Velocity) { + warmstart(&self.solver_data(), self, a, b); + } + #[inline(always)] + pub fn solve_constraint_gauss_seidel( + &mut self, + poses: &Slice, + params: &RbdSimParams, + a: &mut Velocity, + b: &mut Velocity, + use_bias: bool, + friction: bool, + ) { + solve( + &self.solver_data(), + self, + poses, + params, + a, + b, + use_bias, + friction, + ); + } +} + +/// Access to a twist constraint, in registers or in a tile, for the shared twist arithmetic. +/// This is never used by the Coulomb kernels and contains no runtime friction-model branch. +pub(crate) trait TwistAccess: ContactPointAccess { + fn friction(&self) -> TwistFriction; + fn offset_b(&self) -> Vec3; + fn frames(&self) -> (Quat, Quat); + fn radius(&self, k: usize) -> f32; + fn write_friction(&mut self, tangent: Vec2, twist: f32); +} +impl ContactPointAccess for TwistConstraint { + #[inline(always)] + fn read_point(&self, k: usize) -> ContactPoint { + let p = self.points.at(k); + ContactPoint { + r_a: p.r_a, + r_b: p.r_a + self.offset_b, + local_pt_a: self.frame_a.inverse() * p.r_a, + local_pt_b: self.frame_b.inverse() * (p.r_a + self.offset_b), + normal_mass: p.normal_mass, + normal_impulse: p.normal_impulse, + normal_vel: p.normal_vel, + dist: p.dist, + ..ContactPoint::default() + } + } + #[inline(always)] + fn write_normal_impulse(&mut self, k: usize, value: f32) { + self.points.at_mut(k).normal_impulse = value; + } + #[inline(always)] + fn write_tangent_impulse(&mut self, _k: usize, _value: Vec2) {} +} +impl TwistAccess for TwistConstraint { + #[inline(always)] + fn friction(&self) -> TwistFriction { + self.friction + } + #[inline(always)] + fn offset_b(&self) -> Vec3 { + self.offset_b + } + #[inline(always)] + fn frames(&self) -> (Quat, Quat) { + (self.frame_a, self.frame_b) + } + #[inline(always)] + fn radius(&self, k: usize) -> f32 { + self.points.at(k).radius + } + #[inline(always)] + fn write_friction(&mut self, tangent: Vec2, twist: f32) { + self.friction.tangent_impulse = tangent; + self.friction.twist_impulse = twist; + } +} +#[inline(always)] +pub(crate) fn warmstart( + h: &ContactSolverData, + points: &P, + a: &mut Velocity, + b: &mut Velocity, +) { + if h.len == 0 { + return; + } + let f = points.friction(); + let tangent = + h.tangent_a * f.tangent_impulse.x + h.dir_a.cross(h.tangent_a) * f.tangent_impulse.y; + let twist = h.dir_a * f.twist_impulse; + let mut force = Vec3::ZERO; + let (mut ta, mut tb) = (Vec3::ZERO, Vec3::ZERO); + crunchy::unroll! { for k in 0..4 { if k < h.len as usize { + let p = points.read_point(k); + let normal = h.dir_a * p.normal_impulse; + force += normal; + ta += p.r_a.cross(normal); + tb += p.r_b.cross(normal); + } }} + force += tangent; + ta += f.center_a.cross(tangent) + twist; + tb += (f.center_a + points.offset_b()).cross(tangent) + twist; + a.linear += h.im_a * force; + a.angular += h.ii_a.mul(ta); + b.linear -= h.im_b * force; + b.angular -= h.ii_b.mul(tb); +} +/// Transform the normal back into the frozen world frame once per manifold. +/// The two local anchors derive from r_a and r_a + offset_b, so their projected +/// separation needs just one dot product per point. No per-point anchors are stored. +/// The frozen Jacobians and Gauss-Seidel point order stay unchanged. +#[inline(always)] +fn solve_normals( + h: &ContactSolverData, + points: &mut P, + poses: &Slice, + params: &RbdSimParams, + a: &mut Velocity, + b: &mut Velocity, + use_bias: bool, +) { + if h.len == 0 { + return; + } + let pa = poses[h.solver_body_a as usize]; + let pb = poses[h.solver_body_b as usize]; + let (fa, fb) = points.frames(); + let na = fa * pa.inverse_transform_vector(h.dir_a); + let nb = fb * pb.inverse_transform_vector(h.dir_a); + let delta_n = na - nb; + let separation = (pa.translation - pb.translation).dot(h.dir_a) - points.offset_b().dot(nb); + let is_static = h.im_a == Vec3::ZERO || h.im_b == Vec3::ZERO; + let (cfm, erp) = if is_static { + ( + params.static_contact_cfm_factor(), + params.static_contact_erp_inv_dt(), + ) + } else { + (params.contact_cfm_factor(), params.contact_erp_inv_dt()) + }; + let inv_dt = params.inv_dt(); + let max_corr = params.max_corrective_velocity(); + crunchy::unroll! { for k in 0..4 { if k < h.len as usize { + let p = points.read_point(k); + let dist = p.dist + (separation + p.r_a.dot(delta_n)); + let rhs = p.normal_vel + dist.max(0.0) * inv_dt; + let (rhs, softness) = if use_bias { + (rhs + (dist * erp).max(-max_corr).min(0.0), if dist <= 0.0 { cfm } else { 1.0 }) + } else { + (rhs, 1.0) + }; + let ta = p.r_a.cross(h.dir_a); + let tb = p.r_b.cross(-h.dir_a); + let dvel = h.dir_a.dot(a.linear) + ta.dot(a.angular) + - h.dir_a.dot(b.linear) + tb.dot(b.angular) + rhs; + let impulse = softness * (p.normal_impulse - p.normal_mass * dvel).max(0.0); + let delta = impulse - p.normal_impulse; + points.write_normal_impulse(k, impulse); + a.linear += h.dir_a * h.im_a * delta; + a.angular += h.ii_a.mul(ta) * delta; + b.linear += h.dir_a * h.im_b * -delta; + b.angular += h.ii_b.mul(tb) * delta; + } }} +} +#[inline(always)] +pub(crate) fn solve( + h: &ContactSolverData, + points: &mut P, + poses: &Slice, + params: &RbdSimParams, + a: &mut Velocity, + b: &mut Velocity, + use_bias: bool, + solve_friction: bool, +) { + solve_normals(h, points, poses, params, a, b, use_bias); + if !solve_friction || h.len == 0 { + return; + } + let mut rhs = Vec2::ZERO; + if use_bias { + let pa = poses[h.solver_body_a as usize]; + let pb = poses[h.solver_body_b as usize]; + let (fa, fb) = points.frames(); + let center = points.friction().center_a; + let drift = (pa * (fa.inverse() * center) + - pb * (fb.inverse() * (center + points.offset_b()))) + * params.inv_dt(); + rhs = Vec2::new( + drift.dot(h.tangent_a), + drift.dot(h.dir_a.cross(h.tangent_a)), + ); + } + solve_friction_rows(h, points, a, b, rhs); +} +#[inline(always)] +pub(crate) fn solve_friction_rows( + h: &ContactSolverData, + points: &mut P, + a: &mut Velocity, + b: &mut Velocity, + rhs: Vec2, +) { + if h.len == 0 { + return; + } + let f = points.friction(); + let (ra, rb) = (f.center_a, f.center_a + points.offset_b()); + let (mut tangent_limit, mut twist_limit) = (0.0, 0.0); + crunchy::unroll! { for k in 0..4 { if k < h.len as usize { + let p = points.read_point(k); + tangent_limit += p.normal_impulse; + twist_limit += p.normal_impulse * points.radius(k); + } }} + tangent_limit *= h.limit; + twist_limit *= h.limit; + let mut twist = 0.0; + if h.len > 1 { + let dvel = h.dir_a.dot(a.angular - b.angular); + twist = (f.twist_impulse - f.twist_mass * dvel) + .max(-twist_limit) + .min(twist_limit); + let delta = twist - f.twist_impulse; + a.angular += h.ii_a.mul(h.dir_a) * delta; + b.angular -= h.ii_b.mul(h.dir_a) * delta; + } + let t0 = h.tangent_a; + let t1 = h.dir_a.cross(t0); + let a0 = ra.cross(t0); + let a1 = ra.cross(t1); + let b0 = rb.cross(-t0); + let b1 = rb.cross(-t1); + let d0 = t0.dot(a.linear) + a0.dot(a.angular) - t0.dot(b.linear) + b0.dot(b.angular) + rhs.x; + let d1 = t1.dot(a.linear) + a1.dot(a.angular) - t1.dot(b.linear) + b1.dot(b.angular) + rhs.y; + let k00 = f.tangent_k.x; + let k11 = f.tangent_k.y; + let k01 = f.tangent_k.z * 0.5; + let inv_det = maybe_inv(k00 * k11 - k01 * k01); + let delta = Vec2::new( + (k11 * d0 - k01 * d1) * inv_det, + (k00 * d1 - k01 * d0) * inv_det, + ); + let tangent = cap_magnitude(f.tangent_impulse - delta, tangent_limit); + let delta = tangent - f.tangent_impulse; + let force = t0 * delta.x + t1 * delta.y; + a.linear += h.im_a * force; + a.angular += h.ii_a.mul(a0 * delta.x + a1 * delta.y); + b.linear -= h.im_b * force; + b.angular += h.ii_b.mul(b0 * delta.x + b1 * delta.y); + points.write_friction(tangent, twist); +} + +#[cfg(all(test, not(target_arch_is_gpu)))] +#[path = "../tests/twist_friction.rs"] +mod tests; diff --git a/src_rbd_shaders/dynamics/twist_tiles.rs b/src_rbd_shaders/dynamics/twist_tiles.rs new file mode 100644 index 00000000..37ed73a7 --- /dev/null +++ b/src_rbd_shaders/dynamics/twist_tiles.rs @@ -0,0 +1,433 @@ +//! Twist contact constraints (central sliding friction and twist friction per manifold) in +//! typed tiles, and their kernels: 21 16-byte solver fields per constraint instead of +//! Coulomb's 29. +use super::contact_tiles::{HeaderTile, TILE_LEN, VecScalar, WarmstartTile, tile_lane}; +use super::solver_utils::{ContactPointAccess, ContactSolverData}; +use super::twist_friction::{self, TwistAccess, TwistContactPoint}; +use super::{ + ContactPoint, RbdSimParams, TwistConstraint as TwoBodyConstraint, TwistFriction, Velocity, +}; +use crate::Pose; +use crate::utils::Slice; +use crunchy::unroll; +use khal_std::glamx::{Quat, Vec2, Vec3}; +use khal_std::index::MaybeIndexUnchecked; + +/// Accumulated central friction impulse, and the coupling term `2 k01` of its effective mass. +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct TwistFrictionRows { + pub impulse: Vec2, + pub k01: f32, + pub _padding: f32, +} + +/// [`TILE_LEN`] consecutive twist constraints. The header stores the twist mass next to body +/// B's inertia diagonal, and the twist impulse next to its off-diagonal terms. +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct TwistTile { + pub header: HeaderTile, + /// Offset from the frozen body-A to body-B points, and the friction mass term `k00`. + pub offset_b: [VecScalar; TILE_LEN], + /// Friction center relative to body A, and the friction mass term `k11`. + pub center_a: [VecScalar; TILE_LEN], + pub friction: [TwistFrictionRows; TILE_LEN], + /// Body A's frozen frame. + pub frame_a: [Quat; TILE_LEN], + /// Body B's frozen frame. + pub frame_b: [Quat; TILE_LEN], + /// Each point's accumulated normal impulse. + pub normal_impulse: [[f32; 4]; TILE_LEN], + /// Each point's signed distance at the start of the step. + pub dist: [[f32; 4]; TILE_LEN], + /// Each point's target normal velocity. + pub normal_vel: [[f32; 4]; TILE_LEN], + /// Each point's twist friction radius. + pub radius: [[f32; 4]; TILE_LEN], + /// Each point's lever arm from body A's center of mass, and inverse effective normal mass. + pub arm_a: [[VecScalar; TILE_LEN]; 4], + pub warmstart: WarmstartTile, +} + +impl TwistTile { + #[inline(always)] + fn point(&self, lane: usize, k: usize) -> ContactPoint { + let a = self.arm_a.at(k).read(lane); + let rb = a.vec + self.offset_b.read(lane).vec; + ContactPoint { + r_a: a.vec, + r_b: rb, + normal_impulse: self.normal_impulse.at(lane).read(k), + normal_mass: a.scalar, + local_pt_a: self.frame_a.read(lane).inverse() * a.vec, + local_pt_b: self.frame_b.read(lane).inverse() * rb, + dist: self.dist.at(lane).read(k), + normal_vel: self.normal_vel.at(lane).read(k), + ..ContactPoint::default() + } + } + + #[inline(always)] + fn friction(&self, lane: usize) -> TwistFriction { + let offset_b = self.offset_b.read(lane); + let center_a = self.center_a.read(lane); + let friction = self.friction.read(lane); + TwistFriction { + center_a: center_a.vec, + tangent_k: Vec3::new(offset_b.scalar, center_a.scalar, friction.k01), + tangent_impulse: friction.impulse, + twist_mass: self.header.inertia_diag_b.read(lane).scalar, + twist_impulse: self.header.inertia_off_b.read(lane).scalar, + _padding: [0; 2], + } + } + + /// Writes the constraint of `lane`. + #[inline(always)] + pub fn write_constraint(&mut self, lane: usize, c: &TwoBodyConstraint) { + let f = &c.friction; + self.header + .write(lane, &c.solver_data(), [f.twist_mass, f.twist_impulse]); + self.offset_b.write( + lane, + VecScalar { + vec: c.offset_b, + scalar: f.tangent_k.x, + }, + ); + self.center_a.write( + lane, + VecScalar { + vec: f.center_a, + scalar: f.tangent_k.y, + }, + ); + self.friction.write( + lane, + TwistFrictionRows { + impulse: f.tangent_impulse, + k01: f.tangent_k.z, + _padding: 0.0, + }, + ); + self.frame_a.write(lane, c.frame_a); + self.frame_b.write(lane, c.frame_b); + let mut impulses = [0.0; 4]; + let mut distances = [0.0; 4]; + let mut velocities = [0.0; 4]; + let mut radii = [0.0; 4]; + unroll! { for k in 0..4 { if k < c.len as usize { + let p = c.points.at(k); + self.arm_a.at_mut(k).write(lane, VecScalar { vec: p.r_a, scalar: p.normal_mass }); + impulses.write(k, p.normal_impulse); + distances.write(k, p.dist); + velocities.write(k, p.normal_vel); + radii.write(k, p.radius); + } }} + self.normal_impulse.write(lane, impulses); + self.dist.write(lane, distances); + self.normal_vel.write(lane, velocities); + self.radius.write(lane, radii); + } + + /// The constraint of `lane`. Its identity fields are left to their defaults: they live in the + /// contact links. + #[inline(always)] + pub fn constraint(&self, lane: usize) -> TwoBodyConstraint { + let h = self.header.read(lane); + let mut c = TwoBodyConstraint { + dir_a: h.dir_a, + len: h.len, + tangent_a: h.tangent_a, + limit: h.limit, + im_a: h.im_a, + im_b: h.im_b, + ii_a: h.ii_a, + ii_b: h.ii_b, + solver_body_a: h.solver_body_a, + solver_body_b: h.solver_body_b, + vel_slot_a: h.vel_slot_a, + vel_slot_b: h.vel_slot_b, + offset_b: self.offset_b.read(lane).vec, + frame_a: self.frame_a.read(lane), + frame_b: self.frame_b.read(lane), + friction: self.friction(lane), + ..Default::default() + }; + let impulses = self.normal_impulse.read(lane); + let distances = self.dist.read(lane); + let velocities = self.normal_vel.read(lane); + let radii = self.radius.read(lane); + unroll! { for k in 0..4 { if k < h.len as usize { + let a = self.arm_a.at(k).read(lane); + c.points.write(k, TwistContactPoint { + r_a: a.vec, + normal_impulse: impulses.read(k), + normal_mass: a.scalar, + dist: distances.read(k), + normal_vel: velocities.read(k), + radius: radii.read(k), + }); + } }} + c + } +} + +/// The tile type of this model's kernels (see `tile_kernels.rs`). +pub type ContactTile = TwistTile; + +#[inline(always)] +fn header(tiles: &[TwistTile], index: usize) -> ContactSolverData { + let (tile, lane) = tile_lane(index); + tiles.at(tile).header.read(lane) +} + +#[inline(always)] +fn read_constraint(tiles: &[TwistTile], index: usize) -> TwoBodyConstraint { + let (tile, lane) = tile_lane(index); + tiles.at(tile).constraint(lane) +} + +#[inline(always)] +fn write_constraint(tiles: &mut [TwistTile], index: usize, c: &TwoBodyConstraint) { + let (tile, lane) = tile_lane(index); + tiles.at_mut(tile).write_constraint(lane, c); +} + +/// The velocity changes of the bodies of `c` applying its accumulated impulses, tagged with +/// their body index (for the warmstart gather). +#[inline(always)] +fn constraint_warmstart(c: &TwoBodyConstraint) -> (Velocity, Velocity) { + let (mut a, mut b) = (Velocity::default(), Velocity::default()); + c.warmstart_constraint(&mut a, &mut b); + a.padding1 = c.solver_body_a; + b.padding1 = c.solver_body_b; + (a, b) +} + +struct TilePoints<'a> { + tiles: &'a mut [TwistTile], + index: usize, +} +impl TilePoints<'_> { + #[inline(always)] + fn tile(&self) -> (&TwistTile, usize) { + let (tile, lane) = tile_lane(self.index); + (self.tiles.at(tile), lane) + } +} +impl ContactPointAccess for TilePoints<'_> { + #[inline(always)] + fn read_point(&self, k: usize) -> ContactPoint { + let (tile, lane) = self.tile(); + tile.point(lane, k) + } + #[inline(always)] + fn write_normal_impulse(&mut self, k: usize, value: f32) { + let (tile, lane) = tile_lane(self.index); + self.tiles + .at_mut(tile) + .normal_impulse + .at_mut(lane) + .write(k, value); + } + #[inline(always)] + fn write_tangent_impulse(&mut self, _k: usize, _value: Vec2) {} +} +impl TwistAccess for TilePoints<'_> { + #[inline(always)] + fn friction(&self) -> TwistFriction { + let (tile, lane) = self.tile(); + tile.friction(lane) + } + #[inline(always)] + fn offset_b(&self) -> Vec3 { + let (tile, lane) = self.tile(); + tile.offset_b.read(lane).vec + } + #[inline(always)] + fn frames(&self) -> (Quat, Quat) { + let (tile, lane) = self.tile(); + (tile.frame_a.read(lane), tile.frame_b.read(lane)) + } + #[inline(always)] + fn radius(&self, k: usize) -> f32 { + let (tile, lane) = self.tile(); + tile.radius.at(lane).read(k) + } + #[inline(always)] + fn write_friction(&mut self, tangent: Vec2, twist: f32) { + let (tile, lane) = tile_lane(self.index); + let tile = self.tiles.at_mut(tile); + tile.friction.at_mut(lane).impulse = tangent; + tile.header.inertia_off_b.at_mut(lane).scalar = twist; + } +} + +#[inline(always)] +#[allow(clippy::too_many_arguments)] +fn solve_constraint( + h: &ContactSolverData, + tiles: &mut [TwistTile], + index: usize, + poses: &Slice, + params: &RbdSimParams, + a: &mut Velocity, + b: &mut Velocity, + bias: bool, + friction: bool, +) { + twist_friction::solve( + h, + &mut TilePoints { tiles, index }, + poses, + params, + a, + b, + bias, + friction, + ); +} + +/// The velocity changes of the bodies of the constraint at `index` applying its accumulated +/// impulses, tagged with their body index (for the warmstart gather). +#[inline(always)] +fn tile_warmstart( + tiles: &[TwistTile], + index: usize, + h: &ContactSolverData, +) -> (Velocity, Velocity) { + let (mut a, mut b) = (Velocity::default(), Velocity::default()); + let (tile, lane) = tile_lane(index); + twist_friction::warmstart( + h, + &TileConstraint { + tile: tiles.at(tile), + lane, + }, + &mut a, + &mut b, + ); + a.padding1 = h.solver_body_a; + b.padding1 = h.solver_body_b; + (a, b) +} + +/// Read-only access to a constraint in a tile, for the warmstart. +struct TileConstraint<'a> { + tile: &'a TwistTile, + lane: usize, +} +impl ContactPointAccess for TileConstraint<'_> { + #[inline(always)] + fn read_point(&self, k: usize) -> ContactPoint { + self.tile.point(self.lane, k) + } + #[inline(always)] + fn write_normal_impulse(&mut self, _k: usize, _value: f32) {} + #[inline(always)] + fn write_tangent_impulse(&mut self, _k: usize, _value: Vec2) {} +} +impl TwistAccess for TileConstraint<'_> { + #[inline(always)] + fn friction(&self) -> TwistFriction { + self.tile.friction(self.lane) + } + #[inline(always)] + fn offset_b(&self) -> Vec3 { + self.tile.offset_b.read(self.lane).vec + } + #[inline(always)] + fn frames(&self) -> (Quat, Quat) { + ( + self.tile.frame_a.read(self.lane), + self.tile.frame_b.read(self.lane), + ) + } + #[inline(always)] + fn radius(&self, k: usize) -> f32 { + self.tile.radius.at(self.lane).read(k) + } + #[inline(always)] + fn write_friction(&mut self, _tangent: Vec2, _twist: f32) {} +} + +// The kernels are shared by both friction models, each instantiating them with its own tiles. +#[allow(clippy::duplicate_mod)] +#[path = "tile_kernels.rs"] +mod kernels; +pub use kernels::*; + +/// Asserts that solving `initial` from a tile gives `expected` and its body velocities, for +/// constraints at the start and the end of tiles. +#[cfg(all(test, not(target_arch_is_gpu)))] +#[allow(clippy::too_many_arguments)] +pub(crate) fn assert_tile_equivalent( + initial: TwoBodyConstraint, + poses: &Slice, + params: &RbdSimParams, + initial_a: Velocity, + initial_b: Velocity, + use_bias: bool, + solve_friction: bool, + expected: &TwoBodyConstraint, + expected_a: Velocity, + expected_b: Velocity, +) { + for index in [0usize, 31, 32, 63] { + let mut tiles = vec![TwistTile::default(); 2]; + write_constraint(&mut tiles, index, &initial); + let h = header(&tiles, index); + let (mut a, mut b) = (initial_a, initial_b); + solve_constraint( + &h, + &mut tiles, + index, + poses, + params, + &mut a, + &mut b, + use_bias, + solve_friction, + ); + let solved = read_constraint(&tiles, index); + for k in 0..expected.len as usize { + assert_eq!( + solved.points[k].normal_impulse, + expected.points[k].normal_impulse + ); + } + assert_eq!( + solved.friction.tangent_impulse, + expected.friction.tangent_impulse + ); + assert_eq!( + solved.friction.twist_impulse, + expected.friction.twist_impulse + ); + assert_eq!( + (a.linear, a.angular, b.linear, b.angular), + ( + expected_a.linear, + expected_a.angular, + expected_b.linear, + expected_b.angular + ) + ); + let (cached_a, cached_b) = tile_warmstart(&tiles, index, &h); + let (mut warm_a, mut warm_b) = (Velocity::default(), Velocity::default()); + expected.warmstart_constraint(&mut warm_a, &mut warm_b); + assert_eq!( + ( + cached_a.linear, + cached_a.angular, + cached_b.linear, + cached_b.angular + ), + (warm_a.linear, warm_a.angular, warm_b.linear, warm_b.angular) + ); + } +} diff --git a/src_rbd_shaders/dynamics/warmstart.rs b/src_rbd_shaders/dynamics/warmstart.rs deleted file mode 100644 index 5141b0ae..00000000 --- a/src_rbd_shaders/dynamics/warmstart.rs +++ /dev/null @@ -1,304 +0,0 @@ -//! Constraint warmstarting (impulse caching across frames). -//! -//! Transfers impulses from frame `n-1` to frame `n` as initial guesses for the solver. -//! Contacts are matched by proximity in local coordinates (threshold: 10cm). - -use crate::broad_phase::ContactPlan; -use khal_std::glamx::UVec3; -use khal_std::macros::{spirv, spirv_bindgen}; - -use super::constraint::{TwoBodyConstraint, TwoBodyConstraintBuilder}; -use crate::utils::{Slice, SliceMut}; -use khal_std::index::MaybeIndexUnchecked; - -/// Transfers warmstart impulses from previous frame to current frame. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_transfer_warmstart_impulses( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] old_body_constraint_counts: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] old_body_constraint_ids: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] - old_constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] - old_constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] - new_constraints: &mut [TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] - new_constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(uniform, descriptor_set = 0, binding = 6)] contact_plan: &ContactPlan, -) { - let total = contact_plan.bound; - let old_body_constraint_counts = Slice(old_body_constraint_counts, 0); - let old_body_constraint_ids = Slice(old_body_constraint_ids, 0); - let old_constraints = Slice(old_constraints, 0); - let old_constraint_builders = Slice(old_constraint_builders, 0); - let mut new_constraints = SliceMut(new_constraints, 0); - let new_constraint_builders = Slice(new_constraint_builders, 0); - - let cid_new = invocation_id.x; - - if cid_new < total { - // Gap / inert slot: stale body ids must not be dereferenced. - if new_constraints[cid_new as usize].len == 0 { - return; - } - transfer_warmstart_impulses( - cid_new, - &old_body_constraint_counts, - &old_body_constraint_ids, - &old_constraints, - &old_constraint_builders, - &mut new_constraints, - &new_constraint_builders, - ); - } -} - -/// Seeds the topo-gc coloring from the previous frame's colors (same -/// body-pair matching as the warmstart impulse transfer), between the -/// topo-gc reset and its iterations. -/// -/// Contacts persist across frames, so most constraints can reuse last -/// frame's color and the iterations only color the genuinely new ones; the -/// fix-conflicts pass still validates every seed, so a stale seed is simply -/// uncolored and recomputed. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_seed_colors_from_warmstart( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] old_body_constraint_counts: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] old_body_constraint_ids: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] - old_constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] - new_constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] old_constraints_colors: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] constraints_colors: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] colored: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 7)] contact_plan: &ContactPlan, -) { - let total = contact_plan.bound; - let old_body_constraint_counts = Slice(old_body_constraint_counts, 0); - let old_body_constraint_ids = Slice(old_body_constraint_ids, 0); - let old_constraints = Slice(old_constraints, 0); - let old_constraints_colors = Slice(old_constraints_colors, 0); - let mut constraints_colors = SliceMut(constraints_colors, 0); - let mut colored = SliceMut(colored, 0); - let new_constraints = Slice(new_constraints, 0); - - let i = invocation_id.x as usize; - - if (i as u32) < total { - // Gap / inert slot: stale body ids must not be dereferenced. - if new_constraints[i].len == 0 { - return; - } - let body_a = new_constraints[i].solver_body_a; - let body_b = new_constraints[i].solver_body_b; - - let first_a = if body_a != 0 { - old_body_constraint_counts[body_a as usize - 1] as usize - } else { - 0 - }; - let last_a = old_body_constraint_counts[body_a as usize] as usize; - let first_b = if body_b != 0 { - old_body_constraint_counts[body_b as usize - 1] as usize - } else { - 0 - }; - let last_b = old_body_constraint_counts[body_b as usize] as usize; - - let len_a = last_a - first_a; - let len_b = last_b - first_b; - let (first_ref, last_ref) = if len_a != 0 && len_a < len_b { - (first_a, last_a) - } else { - (first_b, last_b) - }; - - for j in first_ref..last_ref { - let cid_old = old_body_constraint_ids[j] as usize; - // Same manifold identity as `transfer_warmstart_impulses`: the manifolds of one body - // pair had distinct colors, and seeding them all with the first one's color makes - // them conflict. - let same_manifold = old_constraints[cid_old].warmstart_collider_a - == new_constraints[i].warmstart_collider_a - && old_constraints[cid_old].warmstart_collider_b - == new_constraints[i].warmstart_collider_b - && old_constraints[cid_old].warmstart_subshape - == new_constraints[i].warmstart_subshape; - if old_constraints[cid_old].solver_body_a == body_a - && old_constraints[cid_old].solver_body_b == body_b - && same_manifold - { - let old_color = old_constraints_colors[cid_old]; - // Colors 1..64 are the valid topo-gc range; anything else - // (e.g. stale data after a buffer resize) stays uncolored. - if old_color > 0 && old_color < 64 { - constraints_colors[i] = old_color; - colored[i] = 1; - } - break; - } - } - } -} - -/// Transfers warmstart impulses from previous frame to current frame. -/// -/// NOTE: this assumes that the solver body ids in the constraints match the index of the body itself. -/// This also assumes that bodies in a given constraint pair are always in the same order (they don't -/// get swapped from one frame to another). -pub fn transfer_warmstart_impulses( - cid_new: u32, - old_body_constraint_counts: &Slice, - old_body_constraint_ids: &Slice, - old_constraints: &Slice, - old_constraint_builders: &Slice, - new_constraints: &mut SliceMut, - new_constraint_builders: &Slice, -) { - let i = cid_new as usize; - - // Get the two bodies involved in this new constraint - let body_a = new_constraints[i].solver_body_a; - let body_b = new_constraints[i].solver_body_b; - - // Find the range of old constraints involving body_a - // old_body_constraint_counts is a prefix sum, so the range is [counts[i-1], counts[i]) - let first_constraint_id_a = if body_a != 0 { - old_body_constraint_counts[body_a as usize - 1] as usize - } else { - 0 - }; - let last_constraint_id_a = old_body_constraint_counts[body_a as usize] as usize; - - // Find the range of old constraints involving body_b - let first_constraint_id_b = if body_b != 0 { - old_body_constraint_counts[body_b as usize - 1] as usize - } else { - 0 - }; - let last_constraint_id_b = old_body_constraint_counts[body_b as usize] as usize; - - let len_a = last_constraint_id_a - first_constraint_id_a; - let len_b = last_constraint_id_b - first_constraint_id_b; - - // Optimization: search the smaller constraint list to minimize iterations - // Also avoid static bodies which may have zero-length lists - // Select the smallest list with a nonzero size (for example static bodies would have - // a zero-length list despite having some constraints). - // TODO: compare this approach with just using a hashmap. - let (first_constraint_id_ref, last_constraint_id_ref) = if len_a != 0 && len_a < len_b { - (first_constraint_id_a, last_constraint_id_a) - } else { - (first_constraint_id_b, last_constraint_id_b) - }; - - // Search through old constraints for matching body pair - for j in first_constraint_id_ref..last_constraint_id_ref { - let cid_old = old_body_constraint_ids[j] as usize; - - // Check if this old constraint comes from the same manifold: same body pair, collider pair - // and sub-shape. A body pair alone is ambiguous (one manifold per trimesh triangle or - // polyline segment, or per collider of a compound body), and matching the first one - // hands the same impulses to every manifold of the pair. - let same_manifold = old_constraints[cid_old].warmstart_collider_a - == new_constraints[i].warmstart_collider_a - && old_constraints[cid_old].warmstart_collider_b - == new_constraints[i].warmstart_collider_b - && old_constraints[cid_old].warmstart_subshape == new_constraints[i].warmstart_subshape; - if old_constraints[cid_old].solver_body_a == body_a - && old_constraints[cid_old].solver_body_b == body_b - && same_manifold - { - // Body pair match found! Now match individual contact points. - // We don't have feature IDs, so matching is done by proximity in local space. - - // Distance threshold for matching contact points (10cm) - let dist_threshold = 1.0e-1; // 10cm - let sq_threshold = dist_threshold * dist_threshold; - - // Match each new contact point with the NEAREST old contact point - // within the threshold: the points of one small manifold are all - // within 10cm of each other, so taking the first candidate would hand - // every point the same (first) impulse. - for k_new in 0..(new_constraints[i].len as usize) { - let pt_new_a = new_constraint_builders[i].infos.at(k_new).local_pt_a; - let pt_new_b = new_constraint_builders[i].infos.at(k_new).local_pt_b; - - let mut best_k_old = usize::MAX; - let mut best_sq = sq_threshold; - for k_old in 0..(old_constraints[cid_old].len as usize) { - let pt_old_a = old_constraint_builders[cid_old].infos.at(k_old).local_pt_a; - let pt_old_b = old_constraint_builders[cid_old].infos.at(k_old).local_pt_b; - let dpt_a = pt_old_a - pt_new_a; - let dpt_b = pt_old_b - pt_new_b; - let sq = dpt_a.dot(dpt_a).max(dpt_b.dot(dpt_b)); - if sq < best_sq { - best_sq = sq; - best_k_old = k_old; - } - } - - { - let k_old = best_k_old; - if k_old != usize::MAX { - // Contact point match found! Transfer the accumulated impulse. - // The impulse field contains the last substep's impulse, which serves - // as the warmstart value for this frame. - new_constraints[i] - .elements - .at_mut(k_new) - .normal_part - .impulse = old_constraints[cid_old] - .elements - .at(k_old) - .normal_part - .impulse; - - // 2D: the tangent is derived deterministically from the - // normal, so the scalar component carries over as-is. - #[cfg(feature = "dim2")] - { - new_constraints[i] - .elements - .at_mut(k_new) - .tangent_part - .impulse = old_constraints[cid_old] - .elements - .at(k_old) - .tangent_part - .impulse; - } - - // 3D: reproject the friction impulse through world space - // onto the new constraint's tangent basis. - #[cfg(feature = "dim3")] - { - let old_c = &old_constraints[cid_old]; - let old_t0 = old_c.tangent_a; - let old_t1 = old_c.dir_a.cross(old_t0); - let old_impulse = old_c.elements.at(k_old).tangent_part.impulse; - let world = old_t0 * old_impulse.x + old_t1 * old_impulse.y; - - let new_t0 = new_constraints[i].tangent_a; - let new_t1 = new_constraints[i].dir_a.cross(new_t0); - new_constraints[i] - .elements - .at_mut(k_new) - .tangent_part - .impulse = - khal_std::glamx::Vec2::new(world.dot(new_t0), world.dot(new_t1)); - } - } - } - } - - // Since we found a matching body pair, no need to search further - break; - } - } -} diff --git a/src_rbd_shaders/queries/contact.rs b/src_rbd_shaders/queries/contact.rs index 4301a68b..a222c31e 100644 --- a/src_rbd_shaders/queries/contact.rs +++ b/src_rbd_shaders/queries/contact.rs @@ -126,8 +126,9 @@ pub struct IndexedManifold { /// this manifold, or 0 for a whole convex shape. With `colliders`, it identifies the /// manifold across frames for warmstarting. pub subshape: u32, - /// Padding so the struct size stays a multiple of 16 bytes (std430 array stride). - pub _padding: f32, + /// Bound on the distance between each collider's local origin and its points (rapier's + /// contact-recycling `max_extent`), or a negative value if the pair can't be recycled. + pub recycle_extent: f32, } /// Computes the contact between two balls. diff --git a/src_rbd_shaders/queries/polygonal_feature.rs b/src_rbd_shaders/queries/polygonal_feature.rs index c1f23dbe..56684828 100644 --- a/src_rbd_shaders/queries/polygonal_feature.rs +++ b/src_rbd_shaders/queries/polygonal_feature.rs @@ -582,10 +582,110 @@ mod dim3 { result } + // The functions below index their small arrays with constants only (loops are unrolled, + // runtime indices become selects): an array indexed dynamically lives in private memory, + // which is slow on the GPU. + + /// Expands `$body` once per listed index, with `$i` bound to it. + macro_rules! unroll { + ($i:ident in [$($n:literal),*] $body:block) => { + $({ + let $i: usize = $n; + $body + })* + }; + } + + /// `arr[i]` for a runtime `i < 4`. + #[inline(always)] + fn pick4(arr: &[T; 4], i: usize) -> T { + let mut r = arr[0]; + if i == 1 { + r = arr[1]; + } + if i == 2 { + r = arr[2]; + } + if i == 3 { + r = arr[3]; + } + r + } + + /// `arr[i]` for a runtime `i < MAX_CANDIDATE_POINTS`. + #[inline(always)] + fn pick_candidate(arr: &[ContactPoint; MAX_CANDIDATE_POINTS], i: usize) -> ContactPoint { + let mut r = arr[0]; + unroll!(k in [1, 2, 3, 4, 5, 6, 7] { + if i == k { + r = arr[k]; + } + }); + r + } + + /// Whether `p` is outside the (convex) polygon `poly` of `len` vertices: the signs of its + /// side tests along the edges differ. + #[inline(always)] + #[allow(unused_assignments)] // The last unrolled iteration's updates aren't read. + fn is_outside(poly: &[Vec2; 4], len: usize, p: Vec2) -> bool { + let last = pick4(poly, len - 1); + let mut sign = perp(poly[0] - last, p - last); + let mut outside = false; + unroll!(j in [0, 1, 2] { + if j + 1 < len && !outside { + let new_sign = perp(poly[j + 1] - poly[j], p - poly[j]); + if sign == 0.0 { + sign = new_sign; + } else if sign * new_sign < 0.0 { + outside = true; + } + } + }); + outside + } + + /// Appends a candidate (dropped when the candidates are full). + #[inline(always)] + fn push_candidate( + candidates: &mut [ContactPoint; MAX_CANDIDATE_POINTS], + num_candidates: &mut u32, + point: ContactPoint, + ) { + let n = *num_candidates as usize; + unroll!(k in [0, 1, 2, 3, 4, 5, 6, 7] { + if n == k { + candidates[k] = point; + } + }); + if n < MAX_CANDIDATE_POINTS { + *num_candidates += 1; + } + } + + /// The candidate for the points `local_p1` (in face 1's frame) and `local_p2_1` (face 2's, + /// expressed in face 1's frame), on the side of the flipped features if `flipped`. + #[inline(always)] + fn candidate( + pose12: glamx::Pose3, + local_p1: Vec3, + local_p2_1: Vec3, + dist: f32, + flipped: bool, + ) -> ContactPoint { + let pt = if !flipped { + local_p1 + } else { + pose12.inverse_transform_point(local_p2_1) + }; + ContactPoint { pt, dist } + } + /// Reduces the candidate set to at most `MAX_MANIFOLD_POINTS` solver /// contacts. Mirrors `reduce_manifold_naive`: pick the deepest point, then /// the one furthest from it, then the two extremes along the tangent of /// that segment, considering only points within `prediction`. + #[allow(unused_assignments)] // The last unrolled iteration's updates aren't read. pub fn manifold_reduction( candidates: &[ContactPoint; MAX_CANDIDATE_POINTS], num_candidates: u32, @@ -596,9 +696,11 @@ mod dim3 { let num = num_candidates as usize; if num <= MAX_MANIFOLD_POINTS { - for i in 0..num { - result.points_a.write(i, candidates.read(i)); - } + unroll!(i in [0, 1, 2, 3] { + if i < num { + result.points_a[i] = candidates[i]; + } + }); result.len = num_candidates; return result; } @@ -606,80 +708,87 @@ mod dim3 { const NONE: usize = MAX_CANDIDATE_POINTS; // 1. Find the deepest contact. - let mut selected = [NONE, NONE, NONE, NONE]; + let mut selected0 = NONE; let mut deepest_dist = MAX_FLT; - for i in 0..num { - if candidates.at(i).dist < deepest_dist { - deepest_dist = candidates.at(i).dist; - selected.write(0, i); + unroll!(i in [0, 1, 2, 3, 4, 5, 6, 7] { + if i < num && candidates[i].dist < deepest_dist { + deepest_dist = candidates[i].dist; + selected0 = i; } - } + }); - if selected.read(0) == NONE { + if selected0 == NONE { return result; } // 2. Find the point that is the furthest from the deepest one. - let selected_a = candidates.at(selected.read(0)).pt; + let point0 = pick_candidate(candidates, selected0); + let selected_a = point0.pt; + let mut selected1 = NONE; let mut furthest_dist = -MAX_FLT; - for i in 0..num { - let d = candidates.at(i).pt - selected_a; - let dist = d.dot(d); - if i != selected.read(0) && candidates.at(i).dist <= prediction && dist > furthest_dist - { - furthest_dist = dist; - selected.write(1, i); + unroll!(i in [0, 1, 2, 3, 4, 5, 6, 7] { + if i < num { + let d = candidates[i].pt - selected_a; + let dist = d.dot(d); + if i != selected0 && candidates[i].dist <= prediction && dist > furthest_dist { + furthest_dist = dist; + selected1 = i; + } } - } + }); - result.points_a.write(0, candidates.read(selected.read(0))); + result.points_a[0] = point0; result.len = 1; - if selected.read(1) == NONE { + if selected1 == NONE { return result; } // 3. Now find the two points furthest from the segment we built so far. // A zero-length segment has no tangent, so it stays a single contact. - let selected_b = candidates.at(selected.read(1)).pt; + let point1 = pick_candidate(candidates, selected1); + let selected_b = point1.pt; if selected_a == selected_b { return result; } let tangent = (selected_b - selected_a).cross(normal); + let mut selected2 = NONE; + let mut selected3 = NONE; let mut min_dot = MAX_FLT; let mut max_dot = -MAX_FLT; - for i in 0..num { - if i == selected.read(0) || i == selected.read(1) || candidates.at(i).dist > prediction + unroll!(i in [0, 1, 2, 3, 4, 5, 6, 7] { + if i < num + && i != selected0 + && i != selected1 + && candidates[i].dist <= prediction { - continue; - } - - let d = (candidates.at(i).pt - selected_a).dot(tangent); - if d < min_dot { - min_dot = d; - selected.write(2, i); - } - if d > max_dot { - max_dot = d; - selected.write(3, i); + let d = (candidates[i].pt - selected_a).dot(tangent); + if d < min_dot { + min_dot = d; + selected2 = i; + } + if d > max_dot { + max_dot = d; + selected3 = i; + } } - } + }); - result.points_a.write(1, candidates.read(selected.read(1))); + result.points_a[1] = point1; result.len = 2; - if selected.read(2) == NONE { + if selected2 == NONE { return result; } - result.points_a.write(2, candidates.read(selected.read(2))); + result.points_a[2] = pick_candidate(candidates, selected2); result.len = 3; // The min and max extremes come from one pass, so a single remaining // candidate is picked for both; keeping it once leaves three contacts. - if selected.read(2) == selected.read(3) { + if selected2 == selected3 { return result; } - result.points_a.write(3, candidates.read(selected.read(3))); + result.points_a[3] = pick_candidate(candidates, selected3); result.len = 4; result } @@ -696,208 +805,132 @@ mod dim3 { let mut num_candidates = 0u32; let basis = orthonormal_basis3(sep_axis1); - let projected_face1 = [ - Vec2::new( - face1.vertices.read(0).dot(basis.read(0)), - face1.vertices.read(0).dot(basis.read(1)), - ), - Vec2::new( - face1.vertices.read(1).dot(basis.read(0)), - face1.vertices.read(1).dot(basis.read(1)), - ), - Vec2::new( - face1.vertices.read(2).dot(basis.read(0)), - face1.vertices.read(2).dot(basis.read(1)), - ), - Vec2::new( - face1.vertices.read(3).dot(basis.read(0)), - face1.vertices.read(3).dot(basis.read(1)), - ), - ]; - + let (basis0, basis1) = (basis[0], basis[1]); + let project = |p: Vec3| Vec2::new(p.dot(basis0), p.dot(basis1)); + let vertices1 = face1.vertices; let vertices2_1 = [ - pose12.transform_point(face2.vertices.read(0)), - pose12.transform_point(face2.vertices.read(1)), - pose12.transform_point(face2.vertices.read(2)), - pose12.transform_point(face2.vertices.read(3)), + pose12.transform_point(face2.vertices[0]), + pose12.transform_point(face2.vertices[1]), + pose12.transform_point(face2.vertices[2]), + pose12.transform_point(face2.vertices[3]), + ]; + let projected_face1 = [ + project(vertices1[0]), + project(vertices1[1]), + project(vertices1[2]), + project(vertices1[3]), ]; let projected_face2 = [ - Vec2::new( - vertices2_1.read(0).dot(basis.read(0)), - vertices2_1.read(0).dot(basis.read(1)), - ), - Vec2::new( - vertices2_1.read(1).dot(basis.read(0)), - vertices2_1.read(1).dot(basis.read(1)), - ), - Vec2::new( - vertices2_1.read(2).dot(basis.read(0)), - vertices2_1.read(2).dot(basis.read(1)), - ), - Vec2::new( - vertices2_1.read(3).dot(basis.read(0)), - vertices2_1.read(3).dot(basis.read(1)), - ), + project(vertices2_1[0]), + project(vertices2_1[1]), + project(vertices2_1[2]), + project(vertices2_1[3]), ]; + let num_vertices1 = face1.num_vertices as usize; + let num_vertices2 = face2.num_vertices as usize; + + // Set once a vertex test found every contact. + let mut done = false; // Check vertices of face1 inside face2 - if face2.num_vertices > 2 { - let normal2_1 = (vertices2_1.read(2) - vertices2_1.read(1)) - .cross(vertices2_1.read(0) - vertices2_1.read(1)); + if num_vertices2 > 2 { + let normal2_1 = + (vertices2_1[2] - vertices2_1[1]).cross(vertices2_1[0] - vertices2_1[1]); let denom = normal2_1.dot(sep_axis1); if !relative_eq_scalar(denom, 0.0) { - let last_index2 = face2.num_vertices as usize - 1; let mut any_point_is_outside = false; - for i in 0..face1.num_vertices as usize { - let p1 = projected_face1.read(i); - - let mut sign = perp( - projected_face2.read(0) - projected_face2.read(last_index2), - p1 - projected_face2.read(last_index2), - ); - - let mut point_is_outside = false; - for j in 0..last_index2 { - let new_sign = perp( - projected_face2.read(j + 1) - projected_face2.read(j), - p1 - projected_face2.read(j), - ); - - if sign == 0.0 { - sign = new_sign; - } else if sign * new_sign < 0.0 { - point_is_outside = true; - break; - } - } - - any_point_is_outside = any_point_is_outside || point_is_outside; - let dist = - (vertices2_1.read(0) - face1.vertices.read(i)).dot(normal2_1) / denom; - - if !point_is_outside && dist <= prediction { - let local_p1 = face1.vertices.read(i); - let local_p2_1 = face1.vertices.read(i) + dist * sep_axis1; - - if !flipped { - candidates.at_mut(num_candidates as usize).pt = local_p1; - candidates.at_mut(num_candidates as usize).dist = dist; - } else { - let local_p2 = pose12.inverse_transform_point(local_p2_1); - candidates.at_mut(num_candidates as usize).pt = local_p2; - candidates.at_mut(num_candidates as usize).dist = dist; + unroll!(i in [0, 1, 2, 3] { + if i < num_vertices1 { + let point_is_outside = + is_outside(&projected_face2, num_vertices2, projected_face1[i]); + any_point_is_outside = any_point_is_outside || point_is_outside; + let dist = (vertices2_1[0] - vertices1[i]).dot(normal2_1) / denom; + + if !point_is_outside && dist <= prediction { + let local_p1 = vertices1[i]; + let local_p2_1 = vertices1[i] + dist * sep_axis1; + push_candidate( + &mut candidates, + &mut num_candidates, + candidate(pose12, local_p1, local_p2_1, dist, flipped), + ); } - num_candidates += 1; } - } + }); - if !any_point_is_outside { - return manifold_reduction(&candidates, num_candidates, sep_axis1, prediction); - } + done = !any_point_is_outside; } } // Check vertices of face2 inside face1 - if face1.num_vertices > 2 { - let normal1 = (face1.vertices.read(2) - face1.vertices.read(1)) - .cross(face1.vertices.read(0) - face1.vertices.read(1)); + if !done && num_vertices1 > 2 { + let normal1 = (vertices1[2] - vertices1[1]).cross(vertices1[0] - vertices1[1]); let denom = -normal1.dot(sep_axis1); if !relative_eq_scalar(denom, 0.0) { - let last_index1 = face1.num_vertices as usize - 1; let mut any_point_is_outside = false; - for i in 0..face2.num_vertices as usize { - let p2 = projected_face2.read(i); - - let mut sign = perp( - projected_face1.read(0) - projected_face1.read(last_index1), - p2 - projected_face1.read(last_index1), - ); - - let mut point_is_outside = false; - for j in 0..last_index1 { - let new_sign = perp( - projected_face1.read(j + 1) - projected_face1.read(j), - p2 - projected_face1.read(j), - ); - - if sign == 0.0 { - sign = new_sign; - } else if sign * new_sign < 0.0 { - point_is_outside = true; - break; - } - } - - any_point_is_outside = any_point_is_outside || point_is_outside; - let dist = (face1.vertices.read(0) - vertices2_1.read(i)).dot(normal1) / denom; - - if !point_is_outside && dist <= prediction { - let local_p2_1 = vertices2_1.read(i); - let local_p1 = vertices2_1.read(i) - dist * sep_axis1; - - if !flipped { - candidates.at_mut(num_candidates as usize).pt = local_p1; - candidates.at_mut(num_candidates as usize).dist = dist; - } else { - let local_p2 = pose12.inverse_transform_point(local_p2_1); - candidates.at_mut(num_candidates as usize).pt = local_p2; - candidates.at_mut(num_candidates as usize).dist = dist; + unroll!(i in [0, 1, 2, 3] { + if i < num_vertices2 { + let point_is_outside = + is_outside(&projected_face1, num_vertices1, projected_face2[i]); + any_point_is_outside = any_point_is_outside || point_is_outside; + let dist = (vertices1[0] - vertices2_1[i]).dot(normal1) / denom; + + if !point_is_outside && dist <= prediction { + let local_p2_1 = vertices2_1[i]; + let local_p1 = vertices2_1[i] - dist * sep_axis1; + push_candidate( + &mut candidates, + &mut num_candidates, + candidate(pose12, local_p1, local_p2_1, dist, flipped), + ); } - num_candidates += 1; } - } + }); - if !any_point_is_outside { - return manifold_reduction(&candidates, num_candidates, sep_axis1, prediction); - } + done = !any_point_is_outside; } } - // Check edge-edge intersections - for j in 0..face2.num_vertices as usize { - for i in 0..face1.num_vertices as usize { - let bcoords = closest_points_line2d( - projected_face1.read(i), - projected_face1.read((i + 1) % face1.num_vertices as usize), - projected_face2.read(j), - projected_face2.read((j + 1) % face2.num_vertices as usize), - ); - if bcoords.x > 0.0 && bcoords.x < 1.0 && bcoords.y > 0.0 && bcoords.y < 1.0 { - let edge1_a = face1.vertices.read(i); - let edge1_b = face1.vertices.read((i + 1) % face1.num_vertices as usize); - let edge2_a = vertices2_1.read(j); - let edge2_b = vertices2_1.read((j + 1) % face2.num_vertices as usize); - let local_p1 = edge1_a * (1.0 - bcoords.x) + edge1_b * bcoords.x; - let local_p2_1 = edge2_a * (1.0 - bcoords.y) + edge2_b * bcoords.y; - let dist = (local_p2_1 - local_p1).dot(sep_axis1); - - if dist <= prediction { - if !flipped { - candidates.at_mut(num_candidates as usize).pt = local_p1; - candidates.at_mut(num_candidates as usize).dist = dist; - } else { - let local_p2 = pose12.inverse_transform_point(local_p2_1); - candidates.at_mut(num_candidates as usize).pt = local_p2; - candidates.at_mut(num_candidates as usize).dist = dist; - } - num_candidates += 1; - } - - if num_candidates as usize == MAX_CANDIDATE_POINTS { - return manifold_reduction( - &candidates, - num_candidates, - sep_axis1, - prediction, + // Check edge-edge intersections (the candidates stop growing once full). + unroll!(j in [0, 1, 2, 3] { + if !done && j < num_vertices2 { + let next_j = if j + 1 < num_vertices2 { j + 1 } else { 0 }; + let edge2_a = vertices2_1[j]; + let edge2_b = pick4(&vertices2_1, next_j); + let proj2_a = projected_face2[j]; + let proj2_b = pick4(&projected_face2, next_j); + unroll!(i in [0, 1, 2, 3] { + if i < num_vertices1 && (num_candidates as usize) < MAX_CANDIDATE_POINTS { + let next_i = if i + 1 < num_vertices1 { i + 1 } else { 0 }; + let bcoords = closest_points_line2d( + projected_face1[i], + pick4(&projected_face1, next_i), + proj2_a, + proj2_b, ); + if bcoords.x > 0.0 && bcoords.x < 1.0 && bcoords.y > 0.0 && bcoords.y < 1.0 { + let edge1_a = vertices1[i]; + let edge1_b = pick4(&vertices1, next_i); + let local_p1 = edge1_a * (1.0 - bcoords.x) + edge1_b * bcoords.x; + let local_p2_1 = edge2_a * (1.0 - bcoords.y) + edge2_b * bcoords.y; + let dist = (local_p2_1 - local_p1).dot(sep_axis1); + + if dist <= prediction { + push_candidate( + &mut candidates, + &mut num_candidates, + candidate(pose12, local_p1, local_p2_1, dist, flipped), + ); + } + } } - } + }); } - } + }); manifold_reduction(&candidates, num_candidates, sep_axis1, prediction) } diff --git a/src_rbd_shaders/queries/sat.rs b/src_rbd_shaders/queries/sat.rs index c018ab79..03313a02 100644 --- a/src_rbd_shaders/queries/sat.rs +++ b/src_rbd_shaders/queries/sat.rs @@ -9,8 +9,6 @@ use parry::shape::SupportMap; #[cfg(feature = "dim3")] use glamx::Vec3; -#[cfg(feature = "dim3")] -use khal_std::index::MaybeIndexUnchecked; /// Result of a separating axis test. #[derive(Clone, Copy, Default)] @@ -78,25 +76,9 @@ pub fn cuboid_cuboid_find_local_separating_edge_twoway( let y2 = pos12.rotation * Vec3::new(0.0, 1.0, 0.0); let z2 = pos12.rotation * Vec3::new(0.0, 0.0, 1.0); - // We have 3 * 3 = 9 axes to test. - let axes = [ - // Vector::{x, y ,z}().cross(x2) - Vec3::new(0.0, -x2.z, x2.y), - Vec3::new(x2.z, 0.0, -x2.x), - Vec3::new(-x2.y, x2.x, 0.0), - // Vector::{x, y ,z}().cross(y2) - Vec3::new(0.0, -y2.z, y2.y), - Vec3::new(y2.z, 0.0, -y2.x), - Vec3::new(-y2.y, y2.x, 0.0), - // Vector::{x, y ,z}().cross(z2) - Vec3::new(0.0, -z2.z, z2.y), - Vec3::new(z2.z, 0.0, -z2.x), - Vec3::new(-z2.y, z2.x, 0.0), - ]; - - // TODO: unroll loop - for i in 0..9 { - let axis1 = axes.read(i); + // We have 3 * 3 = 9 axes to test, unrolled: a loop over an array of axes would index it + // dynamically, which puts it in private memory. + let mut test_axis = |axis1: Vec3| { let norm1 = axis1.length(); if norm1 > EPSILON { let sep = cuboid_cuboid_compute_separation_wrt_local_line( @@ -110,7 +92,19 @@ pub fn cuboid_cuboid_find_local_separating_edge_twoway( best_sep = sep; } } - } + }; + // Vector::{x, y ,z}().cross(x2) + test_axis(Vec3::new(0.0, -x2.z, x2.y)); + test_axis(Vec3::new(x2.z, 0.0, -x2.x)); + test_axis(Vec3::new(-x2.y, x2.x, 0.0)); + // Vector::{x, y ,z}().cross(y2) + test_axis(Vec3::new(0.0, -y2.z, y2.y)); + test_axis(Vec3::new(y2.z, 0.0, -y2.x)); + test_axis(Vec3::new(-y2.y, y2.x, 0.0)); + // Vector::{x, y ,z}().cross(z2) + test_axis(Vec3::new(0.0, -z2.z, z2.y)); + test_axis(Vec3::new(z2.z, 0.0, -z2.x)); + test_axis(Vec3::new(-z2.y, z2.x, 0.0)); best_sep } diff --git a/src_rbd_shaders/shapes/shape.rs b/src_rbd_shaders/shapes/shape.rs index 38d20059..81fcba5e 100644 --- a/src_rbd_shaders/shapes/shape.rs +++ b/src_rbd_shaders/shapes/shape.rs @@ -640,6 +640,46 @@ impl Shape { Aabb::default() } + /// The distance from the local origin to the farthest corner of the local AABB (rapier's + /// contact-recycling extent). + pub fn local_extent(&self, vertices: &[PaddedVector]) -> f32 { + let aabb = self.compute_aabb(Pose::IDENTITY, vertices); + aabb.mins.length().max(aabb.maxs.length()) + } + + /// [`Self::local_extent`] of a ball, cuboid, capsule, cone or cylinder, without vertex + /// buffer (0 for the other shapes). + pub fn primitive_local_extent(&self) -> f32 { + let ty = self.shape_type(); + if ty == SHAPE_TYPE_BALL { + #[cfg(feature = "dim2")] + return self.to_ball().radius * core::f32::consts::SQRT_2; + #[cfg(feature = "dim3")] + return self.to_ball().radius * 1.732_050_8; + } + if ty == SHAPE_TYPE_CUBOID { + return self.to_cuboid().half_extents.length(); + } + if ty == SHAPE_TYPE_CAPSULE { + let capsule = self.to_capsule(); + let (a, b) = (capsule.segment.a, capsule.segment.b); + let r = Vector::splat(capsule.radius); + return (a.min(b) - r).length().max((a.max(b) + r).length()); + } + #[cfg(feature = "dim3")] + { + if ty == SHAPE_TYPE_CONE { + let cone = self.to_cone(); + return Vector::new(cone.radius, cone.half_height, cone.radius).length(); + } + if ty == SHAPE_TYPE_CYLINDER { + let cyl = self.to_cylinder(); + return Vector::new(cyl.radius, cyl.half_height, cyl.radius).length(); + } + } + 0.0 + } + /// Creates a ball (sphere/circle) shape. /// /// # Parameters diff --git a/src_rbd_shaders/tests/coloring.rs b/src_rbd_shaders/tests/coloring.rs new file mode 100644 index 00000000..3e73ba5e --- /dev/null +++ b/src_rbd_shaders/tests/coloring.rs @@ -0,0 +1,126 @@ +use super::*; + +#[test] +fn first_free_color_across_both_words() { + // In particular, an occupied color 32 must not shift a free low color. + assert_eq!(first_free_color((0b1011, 1)), 2); + for free in 0..64u32 { + let mask = !(1u64 << free); + assert_eq!(first_free_color((mask as u32, (mask >> 32) as u32)), free); + } + let mut mask = 0x6a09e667f3bcc909u64; + for _ in 0..10000 { + mask = mask.wrapping_mul(6364136223846793005).wrapping_add(1); + assert_eq!( + first_free_color((mask as u32, (mask >> 32) as u32)), + (!mask).trailing_zeros() + ); + } +} +#[test] +fn conflict_pruning_matches_full_graph_with_seeded_and_pending_colors() { + use super::*; + // Include static nodes, multibody groups, split hubs, inert slots and both color words. + for deterministic in [0, 1] { + for seed in 0..32u32 { + let groups = [0u32, 1, 1, 2, 3, 4, 5, 5]; + let mut constraints = vec![ContactLink::default(); 96]; + let mut colors = vec![0; 96]; + let mut pending = vec![MAX_U32; 96]; + let mut colored = vec![1; 96]; + for (i, c) in constraints.iter_mut().enumerate() { + c.len = u32::from(i % 13 != 0); + c.solver_body_a = (i % 8) as u32; + c.solver_body_b = ((i * 3 + 5) % 8) as u32; + c.vel_slot_a = c.solver_body_a + + if groups[c.solver_body_a as usize] == 4 { + 100 + } else { + 0 + }; + c.vel_slot_b = c.solver_body_b + + if groups[c.solver_body_b as usize] == 4 { + 100 + } else { + 0 + }; + colors[i] = 1 + ((i as u32 * 17 + seed) % 7) * 9; + if deterministic != 0 && (i as u32 + seed).is_multiple_of(3) { + pending[i] = 1 + ((i as u32 * 13 + seed) % 7) * 9; + colored[i] = 0; + } + } + let touches = |i: usize, group: u32| { + let c = &constraints[i]; + c.len != 0 + && group != 0 + && ((groups[c.solver_body_a as usize] == group + && c.vel_slot_a == c.solver_body_a) + || (groups[c.solver_body_b as usize] == group + && c.vel_slot_b == c.solver_body_b)) + }; + let mut counts = Vec::new(); + let mut adjacency = Vec::new(); + for group in 0..6 { + for i in 0..constraints.len() { + if touches(i, group) { + adjacency.push(i as u32); + } + } + counts.push(adjacency.len() as u32); + } + let effective: Vec<_> = colors + .iter() + .zip(&pending) + .map(|(&c, &p)| if p == MAX_U32 { c } else { p }) + .collect(); + let mut expected_colors = colors.clone(); + let mut expected_colored = colored.clone(); + for i in 0..constraints.len() { + if constraints[i].len == 0 { + continue; + } + let conflict = (i + 1..constraints.len()).any(|j| { + effective[i] == effective[j] + && (1..6).any(|group| touches(i, group) && touches(j, group)) + }); + if conflict { + expected_colored[i] = 0; + } else if pending[i] != MAX_U32 { + expected_colored[i] = 1; + expected_colors[i] = pending[i]; + } + } + for lane in 0..64 { + gpu_fix_conflicts_topo_gc( + UVec3::new(lane, 0, 0), + UVec3::ONE, + &counts, + &adjacency, + &constraints, + &mut colors, + &mut colored, + &mut 0, + &ContactPlan { + bound: 96, + ..Default::default() + }, + &groups, + &BatchIndices { + deterministic, + ..Default::default() + }, + &pending, + ); + } + assert_eq!( + colors, expected_colors, + "seed {seed}, deterministic {deterministic}" + ); + assert_eq!( + colored, expected_colored, + "seed {seed}, deterministic {deterministic}" + ); + } + } +} diff --git a/src_rbd_shaders/tests/contact_constraints.rs b/src_rbd_shaders/tests/contact_constraints.rs new file mode 100644 index 00000000..0eddd3f2 --- /dev/null +++ b/src_rbd_shaders/tests/contact_constraints.rs @@ -0,0 +1,228 @@ +//! The contact links' matching with the previous step, and the constraints built from them. +use super::*; +use crate::broad_phase::ContactPlan; +use crate::dynamics::{ + ContactLink, ContactRecycleOffsets, ContactRecycleState, MAX_CONSTRAINTS_PER_MANIFOLD, + WorldMassProperties, gpu_init_contact_links, gpu_match_contact_links, +}; +use crate::queries::{ContactManifold, IndexedManifold}; +use glamx::{UVec2, UVec3}; + +fn plan(bound: u32) -> ContactPlan { + ContactPlan { + bound, + ..Default::default() + } +} + +/// The links of `manifolds`, as initialized at the start of a step. +fn links_of(manifolds: &[IndexedManifold]) -> Vec { + let mut links = vec![ContactLink::default(); manifolds.len()]; + for i in 0..manifolds.len() { + gpu_init_contact_links( + UVec3::new(i as u32, 0, 0), + manifolds, + &mut links, + &[0; 3], + &plan(manifolds.len() as u32), + ); + } + links +} + +#[test] +fn matching_distinguishes_subshapes_and_clears_stale_matches() { + // Two manifolds of the same body/collider pair, in opposite adjacency order. + let mut manifold = IndexedManifold { + bodies: UVec2::new(1, 2), + colliders: UVec2::new(1, 2), + ..Default::default() + }; + manifold.contact.len = 1; + let mut old_manifolds = [manifold; 2]; + old_manifolds[1].subshape = 1; + let old_links = links_of(&old_manifolds); + let counts = [0, 2, 4]; + let ids = [1, 0, 0, 1]; + let offsets = ContactRecycleOffsets { + old_base: 0, + new_base: 2, + ..Default::default() + }; + let poses = [Pose::IDENTITY; 3]; + for subshape in 0..3 { + let mut current = manifold; + current.subshape = subshape; + let mut links = links_of(&[current]); + links[0].previous_constraint = 123; + let mut states = [ContactRecycleState::default(); 4]; + gpu_match_contact_links( + UVec3::ZERO, + &[current], + &mut links, + &counts, + &ids, + &old_links, + &[0, 1], + &mut states, + &poses, + &plan(1), + &RbdSimParams::tgs_soft(), + &offsets, + ); + assert_eq!( + links[0].previous_constraint, + if subshape < 2 { subshape } else { u32::MAX } + ); + assert_eq!( + links[0].previous_constraint_index, + links[0].previous_constraint + ); + assert_eq!(links[0].recycled, 0); + } +} + +#[test] +fn constraints_recycle_transfer_or_build_fresh_contacts() { + let params = RbdSimParams::tgs_soft(); + let mut manifold = IndexedManifold { + bodies: UVec2::new(0, 1), + colliders: UVec2::new(0, 1), + friction: 0.6, + restitution: 0.2, + subshape: 3, + recycle_extent: 1.0, + contact: ContactManifold { + normal_a: Vector::Y, + len: MAX_CONSTRAINTS_PER_MANIFOLD as u32, + ..Default::default() + }, + }; + for k in 0..MAX_CONSTRAINTS_PER_MANIFOLD { + manifold.contact.points_a[k].pt = Vector::splat(0.05 * k as f32); + manifold.contact.points_a[k].dist = -0.01; + } + let old_poses = [Pose::IDENTITY; 2]; + let vels = [Velocity::default(); 2]; + let mut mprops = [WorldMassProperties::default(); 2]; + let mut old = TwoBodyConstraint::default(); + manifold.contact_to_constraint( + &Slice(&mprops, 0), + &Slice(&old_poses, 0), + &Slice(&old_poses, 0), + &Slice(&vels, 0), + &mut old, + ); + for k in 0..MAX_CONSTRAINTS_PER_MANIFOLD { + old.points[k].normal_impulse = 0.2 * k as f32; + } + let mut old_tiles = [CoulombTile::default()]; + old_tiles[0].write_constraint(0, &old); + let old_links = links_of(&[manifold]); + let old_state = ContactRecycleState { + pose_a: old_poses[0], + pose_b: old_poses[1], + colliders: manifold.colliders, + max_extent: 1.0, + max_drift: 0.05, + }; + let offsets = ContactRecycleOffsets { + old_base: 0, + new_base: 1, + ..Default::default() + }; + // Recycled contacts, fresh geometry, and a previously unseen manifold. + for drift in [0.0, 0.001, 0.3] { + for matched in [false, true] { + for len in 1..=MAX_CONSTRAINTS_PER_MANIFOLD { + let mut current = manifold; + current.contact.len = len as u32; + current.subshape += u32::from(!matched); + let poses = [Pose::IDENTITY, Pose::from_translation(Vector::X * drift)]; + mprops[0].inv_mass = Vector::splat(0.3); + mprops[1].inv_mass = Vector::new( + 0.0, + 0.6, + #[cfg(feature = "dim3")] + 0.8, + ); + let mut links = links_of(&[current]); + let mut states = [old_state, ContactRecycleState::default()]; + gpu_match_contact_links( + UVec3::ZERO, + &[current], + &mut links, + &[1, 2], + &[0, 0], + &old_links, + &[0], + &mut states, + &poses, + &plan(1), + ¶ms, + &offsets, + ); + // The link of contact 0 is its own sorted link. + let mut tiles = [CoulombTile::default()]; + gpu_prepare_constraints( + UVec3::ZERO, + &links, + &[current], + &mprops, + &poses, + &vels, + &states, + &old_tiles, + &mut tiles, + &plan(1), + &offsets, + ); + + // Recycling keeps the previous contacts with their impulses. Otherwise, the + // drifted anchors are too far from the previous ones to inherit impulses. + let recycled = matched && drift < 0.05; + let mut expected = TwoBodyConstraint::default(); + current.contact_to_constraint( + &Slice(&mprops, 0), + &Slice(&poses, 0), + &Slice(&poses, 0), + &Slice(&vels, 0), + &mut expected, + ); + if recycled { + expected.recycle_from(&old, &vels[0], &vels[1]); + } + if matched && !recycled && drift == 0.0 { + for k in 0..len { + expected.points[k].normal_impulse = old.points[k].normal_impulse; + } + } + let mut expected_tiles = [CoulombTile::default()]; + expected_tiles[0].write_constraint(0, &expected); + // The cached warmstart is tagged with the body indices. + let (mut a, mut b) = (Velocity::default(), Velocity::default()); + expected.warmstart_constraint(&mut a, &mut b); + a.padding1 = 0; + b.padding1 = 1; + expected_tiles[0].warmstart.bodies[0] = [a, b]; + assert_eq!(tiles, expected_tiles, "drift {drift}, matched {matched}"); + assert_eq!(links[0].recycled, u32::from(recycled)); + let expected_state = if recycled { + old_state + } else { + ContactRecycleState { + pose_a: poses[0], + pose_b: poses[1], + colliders: current.colliders, + max_extent: 1.0, + max_drift: params.contact_recycle_distance(), + } + }; + assert_eq!( + bytemuck::bytes_of(&states[1]), + bytemuck::bytes_of(&expected_state) + ); + } + } + } +} diff --git a/src_rbd_shaders/tests/contact_solver.rs b/src_rbd_shaders/tests/contact_solver.rs new file mode 100644 index 00000000..9a4f0a3b --- /dev/null +++ b/src_rbd_shaders/tests/contact_solver.rs @@ -0,0 +1,255 @@ +//! Regression oracle: the pre-optimization point loop. +use super::*; +use crate::dynamics::MAX_CONSTRAINTS_PER_MANIFOLD; + +fn reference_solve( + c: &mut TwoBodyConstraint, + poses: &Slice, + params: &RbdSimParams, + solver_vel1: &mut Velocity, + solver_vel2: &mut Velocity, + use_bias: bool, + solve_friction: bool, +) { + let dir_a = c.dir_a; + let im_a = c.im_a; + let im_b = c.im_b; + + let is_static = im_a == Vector::ZERO || im_b == Vector::ZERO; + let (cfm_factor, erp_inv_dt) = if is_static { + ( + params.static_contact_cfm_factor(), + params.static_contact_erp_inv_dt(), + ) + } else { + (params.contact_cfm_factor(), params.contact_erp_inv_dt()) + }; + let inv_dt = params.inv_dt(); + let max_corr_velocity = params.max_corrective_velocity(); + let pose1 = poses[c.solver_body_a as usize]; + let pose2 = poses[c.solver_body_b as usize]; + + // Solve the normal parts of the constraint. + for k in 0..(c.len as usize) { + let point = c.points.at(k); + let p1 = pose1 * point.local_pt_a; + let p2 = pose2 * point.local_pt_b; + let dist = point.dist + (p1 - p2).dot(dir_a); + let rhs_wo_bias = point.normal_vel + dist.max(0.0) * inv_dt; + let (rhs, cfm_factor) = if use_bias { + // Not `clamp`: its `min <= max` assertion exits the kernel early, which breaks + // the uniform control flow of the fused iterations' barriers on the web. + let rhs_bias = (dist * erp_inv_dt).max(-max_corr_velocity).min(0.0); + // Separated (speculative) points are solved rigidly. + let cfm = if dist <= 0.0 { cfm_factor } else { 1.0 }; + (rhs_wo_bias + rhs_bias, cfm) + } else { + (rhs_wo_bias, 1.0) + }; + + let torque_dir_a = gcross(point.r_a, dir_a); + let torque_dir_b = gcross(point.r_b, -dir_a); + let dvel = dir_a.dot(solver_vel1.linear) + gdot(torque_dir_a, solver_vel1.angular) + - dir_a.dot(solver_vel2.linear) + + gdot(torque_dir_b, solver_vel2.angular) + + rhs; + let impulse = point.normal_impulse; + let new_impulse = cfm_factor * (impulse - point.normal_mass * dvel).max(0.0); + let delta_impulse = new_impulse - impulse; + + c.points.at_mut(k).normal_impulse = new_impulse; + + solver_vel1.linear += dir_a * im_a * delta_impulse; + solver_vel1.angular += c.ii_a_mul(torque_dir_a) * delta_impulse; + + solver_vel2.linear += dir_a * im_b * -delta_impulse; + solver_vel2.angular += c.ii_b_mul(torque_dir_b) * delta_impulse; + } + + // Friction is solved during the stabilization iteration, and during the + // biased pass only when `friction_in_bias_pass` is set. + if !solve_friction { + return; + } + + let friction_coeff = c.limit; + let tangents = c.tangents(); + + // Solve the tangent parts of the constraint. + for k in 0..(c.len as usize) { + let point = c.points.at(k); + let limit = friction_coeff * point.normal_impulse; + + #[cfg(feature = "dim2")] + { + let t = tangents.read(0); + let torque_dir_a = gcross(point.r_a, t); + let torque_dir_b = gcross(point.r_b, -t); + let dvel = t.dot(solver_vel1.linear) + gdot(torque_dir_a, solver_vel1.angular) + - t.dot(solver_vel2.linear) + + gdot(torque_dir_b, solver_vel2.angular); + let impulse = point.tangent_impulse[0]; + // NOTE: don’t use clamp since it can panic. + let new_impulse = (impulse - point.tangent_mass * dvel).max(-limit).min(limit); + let delta_impulse = new_impulse - impulse; + + c.points.at_mut(k).tangent_impulse = [new_impulse]; + + solver_vel1.linear += t * im_a * delta_impulse; + solver_vel1.angular += c.ii_a_mul(torque_dir_a) * delta_impulse; + + solver_vel2.linear += t * im_b * -delta_impulse; + solver_vel2.angular += c.ii_b_mul(torque_dir_b) * delta_impulse; + } + #[cfg(feature = "dim3")] + { + let t0 = tangents.read(0); + let t1 = tangents.read(1); + let torque_dir_a0 = gcross(point.r_a, t0); + let torque_dir_b0 = gcross(point.r_b, -t0); + let torque_dir_a1 = gcross(point.r_a, t1); + let torque_dir_b1 = gcross(point.r_b, -t1); + let dvel_0 = t0.dot(solver_vel1.linear) + gdot(torque_dir_a0, solver_vel1.angular) + - t0.dot(solver_vel2.linear) + + gdot(torque_dir_b0, solver_vel2.angular); + let dvel_1 = t1.dot(solver_vel1.linear) + gdot(torque_dir_a1, solver_vel1.angular) + - t1.dot(solver_vel2.linear) + + gdot(torque_dir_b1, solver_vel2.angular); + + let k11 = point.tangent_k[0]; + let k22 = point.tangent_k[1]; + let k12 = point.tangent_k[2] * 0.5; + let inv_det = maybe_inv(k11 * k22 - k12 * k12); + let delta_impulse = Vec2::new( + (k22 * dvel_0 - k12 * dvel_1) * inv_det, + (k11 * dvel_1 - k12 * dvel_0) * inv_det, + ); + let impulse = point.tangent_impulse; + let new_impulse = cap_magnitude(impulse - delta_impulse, limit); + let delta_impulse = new_impulse - impulse; + c.points.at_mut(k).tangent_impulse = new_impulse; + + let lin = t0 * delta_impulse.x + t1 * delta_impulse.y; + solver_vel1.linear += lin * im_a; + solver_vel1.angular += + c.ii_a_mul(torque_dir_a0 * delta_impulse.x + torque_dir_a1 * delta_impulse.y); + + solver_vel2.linear -= lin * im_b; + solver_vel2.angular += + c.ii_b_mul(torque_dir_b0 * delta_impulse.x + torque_dir_b1 * delta_impulse.y); + } + } +} + +#[test] +fn optimized_iterations_match_original_point_order() { + let poses = [Pose::IDENTITY, Pose::from_translation(Vector::splat(0.25))]; + let poses = Slice(&poses, 0); + let params = RbdSimParams::tgs_soft(); + for len in 0..=MAX_CONSTRAINTS_PER_MANIFOLD { + for static_side in 0..3 { + for locked in [false, true] { + for use_bias in [false, true] { + for solve_friction in [false, true] { + let mut c = TwoBodyConstraint { + len: len as u32, + dir_a: Vector::splat(1.0).normalize(), + ..Default::default() + }; + #[cfg(feature = "dim3")] + { + c.tangent_a = compute_tangent_contact_directions(c.dir_a)[0]; + c.ii_a = super::super::constraint::SymInertia { + diag: Vec3::new(0.7, 0.8, 0.9), + off: Vec3::splat(0.05), + ..Default::default() + }; + c.ii_b = c.ii_a; + } + #[cfg(feature = "dim2")] + { + c.ii_a = 0.7; + c.ii_b = 0.8; + } + c.im_a = if static_side == 1 { + Vector::ZERO + } else { + Vector::splat(0.7) + }; + c.im_b = if static_side == 2 { + Vector::ZERO + } else { + Vector::splat(0.4) + }; + if locked { + c.im_a.x = 0.0; + c.im_b.y = 0.0; + } + c.solver_body_a = 0; + c.solver_body_b = 1; + c.limit = 0.6; + for k in 0..len { + let p = &mut c.points[k]; + p.r_a = Vector::splat(0.1 * (k + 1) as f32); + p.r_b = -p.r_a; + p.local_pt_a = p.r_a; + p.local_pt_b = p.r_b; + p.dist = -0.1 + k as f32 * 0.05; + p.normal_vel = -0.03 * k as f32; + p.normal_impulse = 0.1 * k as f32; + #[cfg(feature = "dim3")] + { + p.tangent_impulse = Vec2::new(0.03, -0.01); + } + #[cfg(feature = "dim2")] + { + p.tangent_impulse = [0.03]; + } + } + c.compute_effective_masses(); + let mut expected = c; + let mut a = Velocity::default(); + let mut b = Velocity::default(); + a.linear = Vector::splat(0.3); + b.linear = Vector::splat(-0.2); + let (mut expected_a, mut expected_b) = (a, b); + reference_solve( + &mut expected, + &poses, + ¶ms, + &mut expected_a, + &mut expected_b, + use_bias, + solve_friction, + ); + super::super::coulomb_tiles::assert_tile_equivalent( + c, + &poses, + ¶ms, + a, + b, + use_bias, + solve_friction, + &expected, + expected_a, + expected_b, + ); + c.solve_constraint_gauss_seidel( + &poses, + ¶ms, + &mut a, + &mut b, + use_bias, + solve_friction, + ); + assert_eq!(a.linear, expected_a.linear); + assert_eq!(a.angular, expected_a.angular); + assert_eq!(b.linear, expected_b.linear); + assert_eq!(b.angular, expected_b.angular); + assert_eq!(bytemuck::bytes_of(&c), bytemuck::bytes_of(&expected)); + } + } + } + } + } +} diff --git a/src_rbd_shaders/tests/hub_split.rs b/src_rbd_shaders/tests/hub_split.rs new file mode 100644 index 00000000..873aa718 --- /dev/null +++ b/src_rbd_shaders/tests/hub_split.rs @@ -0,0 +1,101 @@ +use super::*; + +#[test] +fn hub_list_split_matches_full_adjacency_scan() { + // Two hubs share one constraint. The remaining two dynamic bodies are unsplit. + // Multiple workgroup counts cover striding, empty groups, and both constraint sides. + let bodies = 6usize; + let mut initial = vec![ContactLink::default(); 167]; + let mut lists = vec![Vec::::new(); bodies]; + for (i, c) in initial.iter_mut().enumerate() { + let (a, b) = if i < 83 { + (1, 3) + } else if i < 166 { + (4, 2) + } else { + (1, 2) + }; + c.len = 1; + c.solver_body_a = a; + c.solver_body_b = b; + c.vel_slot_a = a; + c.vel_slot_b = b; + c.mass_scale_a = 1.0; + c.mass_scale_b = 1.0; + lists[a as usize].push(i as u32); + lists[b as usize].push(i as u32); + } + let mut ids = Vec::new(); + let counts: Vec = lists + .iter() + .map(|list| { + ids.extend_from_slice(list); + ids.len() as u32 + }) + .collect(); + let first = [NOT_A_HUB, 0, 84, NOT_A_HUB, NOT_A_HUB, NOT_A_HUB]; + for hubs in [0, 1, 2] { + let hub_counts = HubCounts { + hubs, + slots: hubs * 84, + base: bodies as u32, + pool: 192, + }; + let hub_list = [1u32, 2]; + let mut expected = initial.clone(); + let mut expected_bodies = vec![u32::MAX; 192]; + let mut expected_constraints = expected_bodies.clone(); + for body in 0..bodies { + if !hub_list[..hubs as usize].contains(&(body as u32)) { + continue; + } + let start = if body == 0 { 0 } else { counts[body - 1] }; + let end = counts[body]; + let n = (end - start) as f32; + for entry in start..end { + let slot = first[body] + entry - start; + let cid = ids[entry as usize]; + let c = &mut expected[cid as usize]; + if c.solver_body_a == body as u32 { + c.vel_slot_a = bodies as u32 + slot; + c.mass_scale_a = n; + } else { + c.vel_slot_b = bodies as u32 + slot; + c.mass_scale_b = n; + } + expected_bodies[slot as usize] = body as u32; + expected_constraints[slot as usize] = cid; + } + } + for groups in [1, 2, 4] { + let mut actual = initial.clone(); + let mut actual_bodies = vec![u32::MAX; 192]; + let mut actual_constraints = actual_bodies.clone(); + for group in 0..groups { + for lane in 0..64 { + gpu_hub_split_constraints( + UVec3::new(lane, 0, 0), + UVec3::new(group, 0, 0), + UVec3::new(groups, 1, 1), + &counts, + &ids, + &first, + &mut actual, + &mut actual_bodies, + &mut actual_constraints, + &hub_counts, + &hub_list, + ); + } + } + assert_eq!(actual, expected); + assert_eq!(actual_bodies, expected_bodies); + assert_eq!(actual_constraints, expected_constraints); + let mut slot_grid = [[99u32; 3]]; + let mut hub_grid = [[99u32; 3]]; + gpu_hub_dispatch(&hub_counts, &mut slot_grid, &mut hub_grid); + assert_eq!(slot_grid, [[(hubs * 84).div_ceil(64), 1, 1]]); + assert_eq!(hub_grid, [[hubs, 1, 1]]); + } + } +} diff --git a/src_rbd_shaders/tests/twist_friction.rs b/src_rbd_shaders/tests/twist_friction.rs new file mode 100644 index 00000000..ce290940 --- /dev/null +++ b/src_rbd_shaders/tests/twist_friction.rs @@ -0,0 +1,518 @@ +use super::*; +use crate::dynamics::{FrictionModel, RbdSimParams}; +use crate::{Pose, utils::Slice}; + +fn patch(len: usize) -> TwistConstraint { + let mut c = TwistConstraint { + len: len as u32, + dir_a: Vec3::Y, + tangent_a: Vec3::X, + im_a: Vec3::ONE, + solver_body_b: 1, + limit: 0.5, + ..Default::default() + }; + c.ii_a.diag = Vec3::ONE; + for (i, x) in [-1.0, 1.0, -1.0, 1.0].into_iter().enumerate().take(len) { + let p = &mut c.points[i]; + p.r_a = Vec3::new(x, -1.0, if i < 2 { -1.0 } else { 1.0 }); + p.normal_impulse = 1.0; + } + c.compute_effective_masses(); + c +} + +#[test] +fn twist_radius_and_load_bound_the_angular_impulse() { + // Four unit loads at radius sqrt(2), mu=0.5: max torque impulse=2sqrt(2). + let mut c = patch(4); + let mut a = Velocity::default(); + let mut b = Velocity::default(); + a.angular = Vec3::Y * 10.0; + solve_friction_rows(&c.solver_data(), &mut c, &mut a, &mut b, Vec2::ZERO); + let limit = 2.0 * 2.0f32.sqrt(); + assert!((c.friction.twist_impulse + limit).abs() < 1e-6); + assert!((a.angular.y - (10.0 - limit)).abs() < 1e-6); + assert_eq!(a.linear, Vec3::ZERO); + assert_eq!(b.angular, Vec3::ZERO); + // With an attainable target, the twist row cancels relative spin exactly. + let mut c = patch(4); + a.angular = Vec3::Y; + solve_friction_rows(&c.solver_data(), &mut c, &mut a, &mut b, Vec2::ZERO); + assert_eq!(a.angular, Vec3::ZERO); + assert_eq!(c.friction.twist_impulse, -1.0); +} + +#[test] +fn one_point_zero_load_and_zero_friction_do_not_add_twist_resistance() { + for (len, mu, load) in [(1, 0.5, 1.0), (4, 0.0, 1.0), (4, 0.5, 0.0)] { + let mut c = patch(len); + c.limit = mu; + for p in &mut c.points[..len] { + p.normal_impulse = load; + } + let mut a = Velocity::default(); + let mut b = Velocity::default(); + a.angular = Vec3::Y * 10.0; + solve_friction_rows(&c.solver_data(), &mut c, &mut a, &mut b, Vec2::ZERO); + assert_eq!(c.friction.twist_impulse, 0.0); + // A single offset point can still resist spin through its sliding row. + if len > 1 { + assert_eq!(a.angular.y, 10.0); + } + } +} + +#[test] +fn central_friction_uses_total_load_and_warmstarts_once() { + let mut c = patch(4); + let mut a = Velocity::default(); + let mut b = Velocity::default(); + a.linear = Vec3::X * 10.0; + solve_friction_rows(&c.solver_data(), &mut c, &mut a, &mut b, Vec2::ZERO); + // Central effective mass=1/2, capped by mu*four unit normal impulses=2. + assert_eq!(c.friction.tangent_impulse, Vec2::new(-2.0, 0.0)); + assert_eq!(a.linear.x, 8.0); + assert_eq!(a.angular.z, -2.0); + c.scale_impulses(0.25); + let mut warm = Velocity::default(); + c.warmstart_constraint(&mut warm, &mut b); + assert_eq!(warm.linear, Vec3::new(-0.5, 1.0, 0.0)); + assert_eq!(warm.angular.z, -0.5); + assert_eq!(c.friction.tangent_impulse.x, -0.5); +} + +#[test] +fn surviving_contact_warmstarts_are_averaged_and_single_point_twist_is_cleared() { + let mut c = patch(4); + let mut old = c; + old.friction.tangent_impulse = Vec2::splat(2.0); + old.friction.twist_impulse = 2.0; + // Three matched points and one fresh point: average over all four. + for k in 0..3 { + c.transfer_friction_point(old.friction_warmstart(k), k); + } + c.finish_friction_transfer(); + assert_eq!(c.friction.tangent_impulse, Vec2::splat(1.5)); + assert_eq!(c.friction.twist_impulse, 1.5); + c.len = 1; + c.compute_effective_masses(); + assert_eq!(c.friction.twist_impulse, 0.0); + assert_eq!(c.points[0].radius, 0.0); +} + +#[test] +fn twist_tiles_and_cached_warmstarts_match_twist_constraints() { + let poses_array = [Pose::IDENTITY; 2]; + let poses = Slice(&poses_array, 0); + let params = RbdSimParams::tgs_soft(); + for len in 0..=4 { + for locked in [false, true] { + for biased in [false, true] { + for friction in [false, true] { + let mut initial = patch(len); + initial.frame_a = Quat::from_rotation_x(0.13); + initial.frame_b = Quat::from_rotation_y(-0.08); + initial.im_b = Vec3::splat(0.4); + initial.ii_b.diag = Vec3::splat(0.3); + initial.ii_a.off = Vec3::splat(0.03); + if locked { + initial.im_a.x = 0.0; + initial.ii_a.diag = Vec3::ZERO; + initial.ii_a.off = Vec3::ZERO; + } + initial.compute_effective_masses(); + let mut a = Velocity::default(); + let mut b = Velocity::default(); + a.linear = Vec3::new(0.2, -0.5, 0.4); + a.angular = Vec3::new(0.3, 0.6, 0.1); + b.angular = Vec3::splat(-0.2); + let (ia, ib) = (a, b); + let mut expected = initial; + expected.solve_constraint_gauss_seidel( + &poses, ¶ms, &mut a, &mut b, biased, friction, + ); + crate::dynamics::twist_tiles::assert_tile_equivalent( + initial, &poses, ¶ms, ia, ib, biased, friction, &expected, a, b, + ); + } + } + } + } +} + +#[test] +fn layouts_and_default_remain_compatible() { + assert_eq!(core::mem::size_of::(), 76); + assert_eq!( + core::mem::size_of::(), + 96 + ); + assert_eq!(core::mem::size_of::(), 384); + assert_eq!( + core::mem::size_of::(), + 544 + ); + assert_eq!(core::mem::size_of::(), 32); + // 21 16-byte solver fields per constraint, and the 4 of its cached warmstart. + assert_eq!( + core::mem::size_of::(), + (21 + 4) * 16 * crate::dynamics::contact_tiles::TILE_LEN + ); + assert_eq!( + RbdSimParams::default().friction_model, + FrictionModel::Coulomb + ); +} + +#[test] +fn twist_is_solved_before_the_off_center_sliding_row() { + let mut c = patch(2); + c.limit = 10.0; + c.points[0].r_a = Vec3::new(0.0, -1.0, 0.0); + c.points[1].r_a = Vec3::new(2.0, -1.0, 0.0); + c.compute_effective_masses(); + let mut a = Velocity::default(); + let mut b = Velocity::default(); + a.angular = Vec3::Y; + solve_friction_rows(&c.solver_data(), &mut c, &mut a, &mut b, Vec2::ZERO); + assert_eq!(a.angular, Vec3::ZERO); + assert_eq!(a.linear, Vec3::ZERO); + assert_eq!(c.friction.tangent_impulse, Vec2::ZERO); + assert_eq!(c.friction.twist_impulse, -1.0); +} + +#[test] +fn twist_effective_mass_uses_both_bodies() { + let mut c = patch(4); + c.limit = 10.0; + c.im_b = Vec3::ONE; + c.ii_b.diag = Vec3::splat(3.0); + c.compute_effective_masses(); + let mut a = Velocity::default(); + let mut b = Velocity::default(); + a.angular = Vec3::Y * 5.0; + b.angular = Vec3::Y; + solve_friction_rows(&c.solver_data(), &mut c, &mut a, &mut b, Vec2::ZERO); + assert_eq!(a.angular, Vec3::Y * 4.0); + assert_eq!(b.angular, a.angular); + assert_eq!(c.friction.twist_impulse, -1.0); +} + +#[test] +fn central_drift_bias_is_used_only_in_the_biased_friction_pass() { + let poses = [Pose::IDENTITY, Pose::from_translation(Vec3::X * 0.1)]; + for biased in [false, true] { + let mut c = patch(4); + let mut a = Velocity::default(); + let mut b = Velocity::default(); + c.solve_constraint_gauss_seidel( + &Slice(&poses, 0), + &RbdSimParams::tgs_soft(), + &mut a, + &mut b, + biased, + true, + ); + if biased { + let limit = c.limit * c.points.iter().map(|p| p.normal_impulse).sum::(); + assert!(c.friction.tangent_impulse.x > 1.5); + assert!((c.friction.tangent_impulse.length() - limit).abs() < 1e-6); + } else { + assert_eq!(c.friction.tangent_impulse, Vec2::ZERO); + } + } +} + +#[test] +fn twist_constraints_recycle_or_partially_transfer_previous_contacts() { + use crate::broad_phase::ContactPlan; + use crate::dynamics::twist_tiles::{self, TwistTile}; + use crate::dynamics::{ + ContactLink, ContactRecycleOffsets, ContactRecycleState, gpu_init_contact_links, + gpu_match_contact_links, + }; + use crate::queries::IndexedManifold; + use glamx::{UVec2, UVec3}; + let params = RbdSimParams::tgs_soft(); + let poses = [Pose::IDENTITY; 2]; + let velocities = [Velocity::default(); 2]; + let mut props = [WorldMassProperties::default(); 2]; + props[0].inv_mass = Vec3::ONE; + props[0].inv_inertia = glamx::Mat4::IDENTITY; + let mut manifold = IndexedManifold { + bodies: UVec2::new(0, 1), + colliders: UVec2::new(0, 1), + friction: 0.5, + recycle_extent: 1.0, + ..Default::default() + }; + manifold.contact.len = 4; + manifold.contact.normal_a = -Vec3::Y; + for k in 0..4 { + manifold.contact.points_a[k].pt = Vec3::new((k % 2) as f32, 0.0, (k / 2) as f32); + } + let mut old = TwistConstraint::default(); + old.init_contact( + &manifold, + &Slice(&props, 0), + &Slice(&poses, 0), + &Slice(&poses, 0), + &Slice(&velocities, 0), + ); + for k in 0..4 { + old.points[k].normal_impulse = (k + 1) as f32; + } + old.friction.tangent_impulse = Vec2::new(0.4, 0.2); + old.friction.twist_impulse = 0.3; + let mut old_tile = TwistTile::default(); + old_tile.write_constraint(0, &old); + let mut old_links = [ContactLink::default()]; + gpu_init_contact_links( + UVec3::ZERO, + &[manifold], + &mut old_links, + &[0, 0], + &ContactPlan { + bound: 1, + ..Default::default() + }, + ); + let mut changed = manifold; + changed.contact.points_a[3].pt += Vec3::X * 0.3; + for recycle in [false, true] { + let state = ContactRecycleState { + pose_a: Pose::IDENTITY, + pose_b: Pose::IDENTITY, + colliders: manifold.colliders, + max_extent: 1.0, + max_drift: if recycle { 0.05 } else { 0.0 }, + }; + let plan = ContactPlan { + bound: 1, + ..Default::default() + }; + let offsets = ContactRecycleOffsets { + old_base: 0, + new_base: 1, + ..Default::default() + }; + let mut states = [state; 2]; + let mut links = [ContactLink::default()]; + gpu_init_contact_links(UVec3::ZERO, &[changed], &mut links, &[0, 0], &plan); + gpu_match_contact_links( + UVec3::ZERO, + &[changed], + &mut links, + &[1, 1], + &[0], + &old_links, + &[0], + &mut states, + &poses, + &plan, + ¶ms, + &offsets, + ); + assert_eq!(links[0].previous_constraint, 0); + assert_eq!(links[0].recycled, u32::from(recycle)); + // The link of contact 0 is its own sorted link. + let mut tiles = [TwistTile::default()]; + twist_tiles::gpu_prepare_constraints( + UVec3::ZERO, + &links, + &[changed], + &props, + &poses, + &velocities, + &states, + &[old_tile], + &mut tiles, + &plan, + &offsets, + ); + let new = tiles[0].constraint(0); + let factor = if recycle { 1.0 } else { 0.75 }; + assert!((new.friction.twist_impulse - old.friction.twist_impulse * factor).abs() < 1e-6); + assert!( + (new.friction.tangent_impulse - old.friction.tangent_impulse * factor).length() < 1e-6 + ); + assert_eq!( + new.points[3].normal_impulse, + if recycle { 4.0 } else { 0.0 } + ); + let (mut a, mut b) = (Velocity::default(), Velocity::default()); + new.warmstart_constraint(&mut a, &mut b); + let [wa, wb] = tiles[0].warmstart.bodies[0]; + assert_eq!((wa.linear, wa.angular), (a.linear, a.angular)); + assert_eq!((wb.linear, wb.angular), (b.linear, b.angular)); + } +} + +#[test] +fn projected_normals_match_world_anchor_reference() { + // Compare against the unchanged world-anchor normal solver, including rotated + // bodies, anisotropic/locked masses, restitution, and speculative contacts. + for len in 0..=4 { + for use_bias in [false, true] { + for origin in [Vec3::ZERO, Vec3::new(42.5, -13.0, 20.0)] { + for static_a in [false, true] { + let mut poses_array = [ + Pose { + translation: origin, + rotation: glamx::Quat::from_scaled_axis(Vec3::new(0.2, -0.3, 0.1)), + ..Pose::IDENTITY + }, + Pose { + translation: origin + Vec3::new(0.3, -0.7, 0.2), + rotation: glamx::Quat::from_scaled_axis(Vec3::new(-0.4, 0.1, 0.3)), + ..Pose::IDENTITY + }, + ]; + let mut c = patch(len); + c.dir_a = Vec3::new(0.2, 0.9, -0.3).normalize(); + c.im_a = if static_a { + Vec3::ZERO + } else { + Vec3::new(0.0, 0.8, 1.2) + }; + c.im_b = Vec3::new(0.3, 0.7, 0.2); + c.ii_a.diag = if static_a { + Vec3::ZERO + } else { + Vec3::new(0.6, 0.4, 0.8) + }; + c.ii_b.diag = Vec3::new(0.4, 0.7, 0.9); + c.ii_b.off = Vec3::new(0.03, -0.04, 0.01); + c.offset_b = poses_array[0].translation - poses_array[1].translation; + c.frame_a = poses_array[0].rotation; + c.frame_b = poses_array[1].rotation; + for (k, p) in c.points.iter_mut().enumerate().take(len) { + p.r_a = poses_array[0].rotation * p.r_a; + p.normal_vel = [-0.1, 0.3, 0.0, -0.2][k]; + p.dist = [-0.1, 0.08, -0.04, 0.2][k]; + } + c.compute_effective_masses(); + poses_array[0].rotation = + glamx::Quat::from_rotation_x(0.03) * poses_array[0].rotation; + poses_array[1].translation += Vec3::Y * 0.002; + let h = c.solver_data(); + let mut reference = c; + let mut a = Velocity::default(); + let mut b = Velocity::default(); + a.linear = Vec3::new(0.2, -2.0, 0.4); + a.angular = Vec3::new(0.3, 0.6, 0.1); + b.angular = Vec3::splat(-0.2); + let (mut ra, mut rb) = (a, b); + let poses = Slice(&poses_array, 0); + let params = RbdSimParams::tgs_soft(); + h.solve_points( + &mut reference, + &poses, + ¶ms, + &mut ra, + &mut rb, + use_bias, + false, + ); + solve_normals(&h, &mut c, &poses, ¶ms, &mut a, &mut b, use_bias); + for (actual, expected) in [ + (a.linear, ra.linear), + (a.angular, ra.angular), + (b.linear, rb.linear), + (b.angular, rb.angular), + ] { + assert!( + (actual - expected).length() < 5e-4, + "{actual:?} != {expected:?}" + ); + } + for k in 0..len { + assert!( + (c.points[k].normal_impulse - reference.points[k].normal_impulse).abs() + < 5e-4 + ); + } + } + } + } + } +} + +#[test] +fn frozen_frames_survive_recycling_and_preserve_local_anchors() { + let mut old = patch(4); + old.frame_a = Quat::from_scaled_axis(Vec3::new(0.2, -0.4, 0.1)); + old.frame_b = Quat::from_scaled_axis(Vec3::new(-0.1, 0.3, 0.5)); + old.offset_b = Vec3::new(0.4, 0.7, -0.2); + old.compute_effective_masses(); + let mut new = old; + new.frame_a = Quat::from_rotation_y(0.7); + new.frame_b = Quat::from_rotation_x(-0.4); + new.offset_b += Vec3::splat(0.01); + new.recycle_from(&old, &Velocity::default(), &Velocity::default()); + assert_eq!(new.frame_a, old.frame_a); + assert_eq!(new.frame_b, old.frame_b); + assert_eq!(new.offset_b, old.offset_b); + for k in 0..4 { + assert_eq!(new.local_anchors(k), old.local_anchors(k)); + let (a, b) = new.local_anchors(k); + assert!((new.frame_a * a - new.points[k].r_a).length() < 1e-6); + assert!((new.frame_b * b - new.points[k].r_a - new.offset_b).length() < 1e-6); + } +} + +#[test] +fn central_drift_matches_averaging_rotated_world_anchors() { + for len in 1..=4 { + let mut c = patch(len); + c.frame_a = Quat::from_scaled_axis(Vec3::new(0.2, -0.4, 0.1)); + c.frame_b = Quat::from_scaled_axis(Vec3::new(-0.1, 0.3, 0.5)); + c.offset_b = Vec3::new(0.4, 0.7, -0.2); + c.compute_effective_masses(); + let poses_array = [ + Pose::from_parts(Vec3::new(0.2, 0.3, 0.4), Quat::from_rotation_x(0.3)), + Pose::from_parts(Vec3::new(0.4, 0.6, 0.2), Quat::from_rotation_y(-0.1)), + ]; + let poses = Slice(&poses_array, 0); + let params = RbdSimParams::tgs_soft(); + let h = c.solver_data(); + let mut reference = c; + let (mut a, mut b) = (Velocity::default(), Velocity::default()); + let (mut ra, mut rb) = (a, b); + // Independent old arithmetic: transform and sum every local anchor. + h.solve_points( + &mut reference, + &poses, + ¶ms, + &mut ra, + &mut rb, + true, + false, + ); + let mut drift = Vec3::ZERO; + for k in 0..len { + let (anchor_a, anchor_b) = reference.local_anchors(k); + drift += poses_array[0] * anchor_a - poses_array[1] * anchor_b; + } + drift *= params.inv_dt() / len as f32; + let rhs = Vec2::new( + drift.dot(h.tangent_a), + drift.dot(h.dir_a.cross(h.tangent_a)), + ); + solve_friction_rows(&h, &mut reference, &mut ra, &mut rb, rhs); + solve(&h, &mut c, &poses, ¶ms, &mut a, &mut b, true, true); + for (actual, expected) in [ + (a.linear, ra.linear), + (a.angular, ra.angular), + (b.linear, rb.linear), + (b.angular, rb.angular), + ] { + assert!( + (actual - expected).length() < 1e-4, + "{actual:?} != {expected:?}" + ); + } + assert!((c.friction.tangent_impulse - reference.friction.tangent_impulse).length() < 1e-4); + } +} diff --git a/src_rbd_shaders/utils/indices.rs b/src_rbd_shaders/utils/indices.rs index 386cff47..8ce58456 100644 --- a/src_rbd_shaders/utils/indices.rs +++ b/src_rbd_shaders/utils/indices.rs @@ -55,7 +55,7 @@ pub struct BatchIndices { /// `M⁻¹·Jᵀ` column). pub mb_contact_constraints_capacity: u32, /// Per-batch multibody-touching impulse-joint slot count (loop bound for - /// the flat sweeps). + /// the flat dispatches). pub mb_imp_joints_batch_capacity: u32, /// Actual max `ndofs` across every multibody in every batch (often smaller /// than the fixed `MAX_MB_DOFS` limit). @@ -69,7 +69,7 @@ pub struct BatchIndices { pub mb_max_joint_constraints: u32, /// Per-batch stride of the contact-solver color-bucket buffers /// (`color_counts` / `color_starts` / `color_cursors`), = `max_colors + 3` - /// so that `starts[c + 1]` is in bounds for every swept color. + /// so that `starts[c + 1]` is in bounds for every solved color. pub solver_color_buckets_stride: u32, /* diff --git a/src_viewer/lib.rs b/src_viewer/lib.rs index f81f7e62..4401f394 100644 --- a/src_viewer/lib.rs +++ b/src_viewer/lib.rs @@ -48,6 +48,8 @@ pub enum UiSection { pub enum DemoKind { Rbd, Mpm, + /// Large rigid-body scenes for performance work. + Stress, } /// A loop transition requested from the UI: stop entirely, or switch to another diff --git a/src_viewer/sensors.rs b/src_viewer/sensors.rs index 11af3c32..0bb06005 100644 --- a/src_viewer/sensors.rs +++ b/src_viewer/sensors.rs @@ -15,6 +15,7 @@ use glamx::{Mat4, Pose3, Vec3}; use kiss3d::camera::Camera3d; use kiss3d::event::WindowEvent; use kiss3d::prelude::Color; +#[cfg(not(target_arch = "wasm32"))] use kiss3d::scene::SceneNode3d; use kiss3d::window::{Canvas, NumSamples, OffscreenSurface}; use nexus::rbd::math::Pose; @@ -248,6 +249,7 @@ impl SensorCamera { /// Renders the shaded scene and returns it as row-major, top-left origin /// RGB bytes (`width * height * 3`). + #[cfg(not(target_arch = "wasm32"))] pub async fn render_rgb(&mut self, scene: &mut SceneNode3d) -> Vec { self.surface.render_3d(scene, &mut self.camera).await; self.surface.snap_image().into_raw() @@ -255,6 +257,7 @@ impl SensorCamera { /// Renders linear eye-space depth in world units, row-major with a top-left /// origin; background pixels are `0.0`. + #[cfg(not(target_arch = "wasm32"))] pub fn render_depth(&mut self, scene: &mut SceneNode3d) -> Vec { self.surface.snap_depth_raw(scene, &mut self.camera) } @@ -262,6 +265,7 @@ impl SensorCamera { /// Renders the per-pixel segmentation id (`0` for background), row-major /// with a top-left origin. Ids are the objects' `segmentation_id`s, which /// the viewer sets per body. + #[cfg(not(target_arch = "wasm32"))] pub fn render_segmentation(&mut self, scene: &mut SceneNode3d) -> Vec { self.surface.snap_segmentation(scene, &mut self.camera) } diff --git a/src_viewer/ui.rs b/src_viewer/ui.rs index 15ea877d..10d0f8f9 100644 --- a/src_viewer/ui.rs +++ b/src_viewer/ui.rs @@ -161,6 +161,7 @@ pub fn main_panel(ctx: &egui::Context, state: &mut UiState, gpu_available: bool) &state.run_stats, state.sync_time, state.steps, + &mut state.wait_for_render, ); } Some(UiSection::Examples) if !state.demos.is_empty() => { @@ -252,6 +253,20 @@ fn simulation_settings(ui: &mut egui::Ui, state: &mut UiState) { } ui.label("Rigid bodies"); ui.add(egui::Slider::new(&mut s.rbd_steps_per_frame, 1..=20).text("steps / frame")); + ui.add(egui::Slider::new(&mut s.rbd_substeps, 1..=20).text("iterations (substeps)")) + .on_hover_text( + "TGS substeps per physics step. More substeps improve accuracy at a higher \ + cost, without changing the amount of simulated time.", + ); + #[cfg(feature = "dim3")] + ui.horizontal(|ui| { + use nexus::rbd::dynamics::FrictionModel; + ui.label("Friction"); + ui.selectable_value(&mut s.rbd_friction_model, FrictionModel::Coulomb, "Coulomb") + .on_hover_text("Friction at each contact point"); + ui.selectable_value(&mut s.rbd_friction_model, FrictionModel::Simplified, "Twist") + .on_hover_text("Central friction and twist resistance per contact manifold. Multibody contacts use Coulomb."); + }); } ui.add_space(4.0); @@ -459,6 +474,7 @@ fn performance_ui( run_stats: &RunStats, sync_time: Duration, steps: u64, + wait_for_render: &mut bool, ) { // Scene entity counts. ui.label(RichText::new("Scene").strong()); @@ -510,6 +526,11 @@ fn performance_ui( )) .strong(), ); + ui.checkbox(wait_for_render, "Wait for render before stepping") + .on_hover_text( + "Wait for each frame's rendering to finish on the GPU before stepping, so the GPU \ + pass timings don't include rendering work running concurrently.", + ); if !run_stats.gpu_pass_times.is_empty() { CollapsingHeader::new(format!("GPU passes: {:.2}ms", run_stats.gpu_total_time_ms)) .id_salt("rbd_gpu_passes") @@ -542,13 +563,13 @@ fn performance_ui( } impl UiState { - /// Order in which demos appear in the picker: grouped by kind (Rbd, Mpm), + /// Order in which demos appear in the picker: grouped by kind (Rbd, Mpm, Stress), /// preserving each group's listing order. Prev/Next walks this sequence so it /// matches the visible list rather than the raw (lexicographically-sorted) /// `demos` index order. fn demo_display_order(&self) -> Vec { let mut order = Vec::with_capacity(self.demos.len()); - for kind in [DemoKind::Rbd, DemoKind::Mpm] { + for kind in [DemoKind::Rbd, DemoKind::Mpm, DemoKind::Stress] { for (i, (_, k)) in self.demos.iter().enumerate() { if *k == kind { order.push(i); @@ -598,6 +619,7 @@ fn demo_navigation(ui: &mut egui::Ui, state: &mut UiState) { fn examples_section(ui: &mut egui::Ui, state: &mut UiState) { demo_group(ui, state, DemoKind::Rbd, "Rigid Bodies"); demo_group(ui, state, DemoKind::Mpm, "MPM"); + demo_group(ui, state, DemoKind::Stress, "Stress tests"); } fn demo_group(ui: &mut egui::Ui, state: &mut UiState, kind: DemoKind, label: &str) { diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 5a46fa2a..7f3ceaae 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -111,6 +111,9 @@ pub struct UiState { pub backend_type: BackendType, /// Replay each frame's GPU work through a compute graph (backends with graph support only). pub compute_graphs: bool, + /// Wait for each frame's rendering to finish on the GPU before stepping, so the GPU pass + /// timings don't include rendering work running concurrently. + pub wait_for_render: bool, pub gpu_init_error: Option, /// Names + kinds of all registered demos, used to populate the demo picker. pub demos: Vec<(String, DemoKind)>, @@ -155,6 +158,11 @@ pub struct SimSettings { pub mpm_gravity: Vector, /// Rigid-body solver steps advanced per rendered frame. pub rbd_steps_per_frame: u32, + /// TGS substeps per rigid-body solver step. + pub rbd_substeps: u32, + /// Simulation-wide rigid-body contact friction model. + #[cfg(feature = "dim3")] + pub rbd_friction_model: nexus::rbd::dynamics::FrictionModel, /// Identical runs (see [`NexusState::set_deterministic`]). A viewer setting, /// kept across demos. pub deterministic: bool, @@ -167,6 +175,9 @@ impl Default for SimSettings { mpm_use_cpic: true, mpm_gravity: Vector::ZERO, rbd_steps_per_frame: 1, + rbd_substeps: nexus::rbd::dynamics::RbdSimParams::default().num_solver_iterations, + #[cfg(feature = "dim3")] + rbd_friction_model: nexus::rbd::dynamics::FrictionModel::Coulomb, deterministic: false, } } @@ -263,6 +274,11 @@ pub struct NexusViewer { /// Whether [`Self::render_frame`] draws the built-in egui panel. Disable /// for clean frame capture ([`Self::snap_rgb`]). draw_ui: bool, + /// Frame to save as an image, and where (see [`Self::with_snapshot`]). + #[cfg(not(target_arch = "wasm32"))] + snapshot: Option<(u32, std::path::PathBuf)>, + /// Number of frames rendered so far. + frames_rendered: u32, /// Lazily-created path tracer for [`Self::raytrace_frame`]. Kept across /// frames so samples keep accumulating while the scene is static. #[cfg(feature = "dim3")] @@ -408,6 +424,9 @@ impl NexusViewer { last_gpu_pass_times: Vec::new(), last_gpu_total_time_ms: 0.0, draw_ui: true, + #[cfg(not(target_arch = "wasm32"))] + snapshot: None, + frames_rendered: 0, #[cfg(feature = "dim3")] raytracer: None, #[cfg(feature = "dim3")] @@ -431,6 +450,7 @@ impl NexusViewer { ui_section: Some(UiSection::Examples), backend_type: BackendType::Gpu, compute_graphs: false, + wait_for_render: false, gpu_init_error: None, demos, selected_demo: 0, @@ -465,6 +485,13 @@ impl NexusViewer { self } + /// Waits for each frame's rendering to finish on the GPU before stepping. Also toggled from + /// the performance panel. + pub fn with_wait_for_render(mut self, enabled: bool) -> Self { + self.ui.wait_for_render = enabled; + self + } + pub fn with_cpu(mut self) -> Self { self.ui.backend_type = BackendType::Cpu; self @@ -476,6 +503,13 @@ impl NexusViewer { self } + /// Saves the `frame`-th rendered frame (counting from 1) as an image at `path`, then quits. + #[cfg(not(target_arch = "wasm32"))] + pub fn with_snapshot(mut self, frame: u32, path: impl Into) -> Self { + self.snapshot = Some((frame, path.into())); + self + } + pub fn with_running(mut self) -> Self { self.ui.run_state = RunState::Running; self @@ -1054,6 +1088,7 @@ impl NexusViewer { /// Renders sensor camera `id`'s shaded RGB image (row-major, top-left /// origin, `width * height * 3` bytes). #[cfg(feature = "dim3")] + #[cfg(not(target_arch = "wasm32"))] pub async fn render_sensor_rgb(&mut self, id: usize) -> Option> { let sensor = self.sensors.get_mut(id)?.as_mut()?; Some(sensor.render_rgb(&mut self.scene3d).await) @@ -1061,6 +1096,7 @@ impl NexusViewer { /// Renders sensor camera `id`'s linear metric depth (`0.0` = background). #[cfg(feature = "dim3")] + #[cfg(not(target_arch = "wasm32"))] pub fn render_sensor_depth(&mut self, id: usize) -> Option> { let sensor = self.sensors.get_mut(id)?.as_mut()?; Some(sensor.render_depth(&mut self.scene3d)) @@ -1068,6 +1104,7 @@ impl NexusViewer { /// Renders sensor camera `id`'s per-pixel segmentation ids (`0` = background). #[cfg(feature = "dim3")] + #[cfg(not(target_arch = "wasm32"))] pub fn render_sensor_segmentation(&mut self, id: usize) -> Option> { let sensor = self.sensors.get_mut(id)?.as_mut()?; Some(sensor.render_segmentation(&mut self.scene3d)) @@ -1384,6 +1421,14 @@ impl NexusViewer { self.ui.sim_settings.mpm_use_cpic = state.mpm_use_cpic(); self.ui.sim_settings.mpm_gravity = state.mpm_gravity(); self.ui.sim_settings.rbd_steps_per_frame = state.rbd_steps_per_frame(); + self.ui.sim_settings.rbd_substeps = state.rbd_substeps(); + #[cfg(feature = "dim3")] + { + self.ui.sim_settings.rbd_friction_model = state + .rbd_sim_params(0) + .map(|params| params.friction_model) + .unwrap_or_default(); + } self.ui.settings_demo = Some(self.ui.selected_demo); } else { let s = self.ui.sim_settings.clone(); @@ -1391,6 +1436,9 @@ impl NexusViewer { state.set_mpm_use_cpic(s.mpm_use_cpic); state.set_mpm_gravity(s.mpm_gravity); state.set_rbd_steps_per_frame(s.rbd_steps_per_frame); + state.set_rbd_substeps(self.backend(), s.rbd_substeps); + #[cfg(feature = "dim3")] + state.set_rbd_friction_model(s.rbd_friction_model); } // The compute-graph toggle is a testbed-wide choice, not a scene // setting: it always flows from the backend panel into the scene. @@ -1590,6 +1638,21 @@ impl NexusViewer { return false; } + self.frames_rendered += 1; + if self.ui.wait_for_render { + Self::wait_for_render().await; + } + #[cfg(not(target_arch = "wasm32"))] + if let Some((frame, path)) = &self.snapshot + && self.frames_rendered == *frame + { + if let Err(err) = self.window.snap_image().save(path) { + eprintln!("Failed to save the snapshot to {}: {err}", path.display()); + } + self.ui.transition = Some(Transition::Quit); + return false; + } + let gpu_available = self.webgpu.is_some(); // Disjoint closure capture (edition 2024): the closure borrows `self.ui` // and `scene_ui` while `self.window` is the receiver. @@ -1603,6 +1666,23 @@ impl NexusViewer { self.ui.transition.is_none() } + /// Waits until the work submitted to kiss3d's queue (this frame's rendering) is done. + async fn wait_for_render() { + if !kiss3d::context::Context::is_initialized() { + return; + } + let ctxt = kiss3d::context::Context::get(); + let (sender, receiver) = futures::channel::oneshot::channel(); + ctxt.queue.on_submitted_work_done(move || { + let _ = sender.send(()); + }); + #[cfg(not(target_arch = "wasm32"))] + let _ = ctxt + .device + .poll(khal::re_exports::wgpu::PollType::wait_indefinitely()); + let _ = receiver.await; + } + /// Captures the last rendered frame as `(width, height, rgb)`, where `rgb` /// is row-major, top-to-bottom, 3 bytes (R, G, B) per pixel. ///