diff --git a/crates/nexus_mpm_shaders2d/Cargo.toml b/crates/nexus_mpm_shaders2d/Cargo.toml index 34a45ded..bde234f4 100644 --- a/crates/nexus_mpm_shaders2d/Cargo.toml +++ b/crates/nexus_mpm_shaders2d/Cargo.toml @@ -26,7 +26,7 @@ dim2 = [] unsafe_remove_boundchecks = ["vortx-shaders/unsafe_remove_boundchecks"] push_constants = [] # Enables some changes in the shaders for compatibility with web platforms. -web-compat = [] +web-compat = ["nexus_rbd_shaders2d/web-compat"] cpu = [] cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"] cuda = [] diff --git a/crates/nexus_mpm_shaders3d/Cargo.toml b/crates/nexus_mpm_shaders3d/Cargo.toml index f3ef990a..ba2e2fd1 100644 --- a/crates/nexus_mpm_shaders3d/Cargo.toml +++ b/crates/nexus_mpm_shaders3d/Cargo.toml @@ -25,7 +25,7 @@ dim3 = [] unsafe_remove_boundchecks = ["vortx-shaders/unsafe_remove_boundchecks"] push_constants = [] # Enables some changes in the shaders for compatibility with web platforms. -web-compat = [] +web-compat = ["nexus_rbd_shaders3d/web-compat"] cpu = [] cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"] cuda = [] diff --git a/crates/nexus_rbd2d/build.rs b/crates/nexus_rbd2d/build.rs index 77b2a85d..fc1104cf 100644 --- a/crates/nexus_rbd2d/build.rs +++ b/crates/nexus_rbd2d/build.rs @@ -5,9 +5,16 @@ fn main() { let output_dir = PathBuf::from(std::env::var_os("OUT_DIR").expect("OUT_DIR not set by cargo")) .join("shaders-spirv"); - KhalBuilder::from_dependency("nexus_rbd_shaders2d", true) + let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH").unwrap(); + + let mut builder = KhalBuilder::from_dependency("nexus_rbd_shaders2d", true) .feature("dim2") // Feature enabled unconditionally for the radix-sort device lost issue (see comment in the radix sort shader code). - .feature("unsafe_remove_boundchecks") - .build(output_dir); + .feature("unsafe_remove_boundchecks"); + + if target_arch == "wasm32" { + builder = builder.feature("web-compat"); + } + + builder.build(output_dir); } diff --git a/crates/nexus_rbd3d/build.rs b/crates/nexus_rbd3d/build.rs index 1b0bd163..064f44be 100644 --- a/crates/nexus_rbd3d/build.rs +++ b/crates/nexus_rbd3d/build.rs @@ -5,9 +5,16 @@ fn main() { let output_dir = PathBuf::from(std::env::var_os("OUT_DIR").expect("OUT_DIR not set by cargo")) .join("shaders-spirv"); - KhalBuilder::from_dependency("nexus_rbd_shaders3d", true) + let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH").unwrap(); + + let mut builder = KhalBuilder::from_dependency("nexus_rbd_shaders3d", true) .feature("dim3") // Feature enabled unconditionally for the radix-sort device lost issue (see comment in the radix sort shader code). - .feature("unsafe_remove_boundchecks") - .build(output_dir); + .feature("unsafe_remove_boundchecks"); + + if target_arch == "wasm32" { + builder = builder.feature("web-compat"); + } + + builder.build(output_dir); } diff --git a/crates/nexus_rbd_shaders2d/Cargo.toml b/crates/nexus_rbd_shaders2d/Cargo.toml index a2ae70c0..d8672afd 100644 --- a/crates/nexus_rbd_shaders2d/Cargo.toml +++ b/crates/nexus_rbd_shaders2d/Cargo.toml @@ -24,8 +24,13 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = [ [features] default = ["dim2"] dim2 = [] -unsafe_remove_boundchecks = ["vortx-shaders/unsafe_remove_boundchecks"] +unsafe_remove_boundchecks = [ + "vortx-shaders/unsafe_remove_boundchecks", + "khal-std/unsafe_remove_boundchecks", +] push_constants = [] +# Enables some changes in the shaders for compatibility with web platforms. +web-compat = [] cpu = [] cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"] cuda = [] diff --git a/crates/nexus_rbd_shaders3d/Cargo.toml b/crates/nexus_rbd_shaders3d/Cargo.toml index 413688ea..1b306203 100644 --- a/crates/nexus_rbd_shaders3d/Cargo.toml +++ b/crates/nexus_rbd_shaders3d/Cargo.toml @@ -24,8 +24,13 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = [ [features] default = ["dim3"] dim3 = [] -unsafe_remove_boundchecks = ["vortx-shaders/unsafe_remove_boundchecks"] +unsafe_remove_boundchecks = [ + "vortx-shaders/unsafe_remove_boundchecks", + "khal-std/unsafe_remove_boundchecks", +] push_constants = [] +# Enables some changes in the shaders for compatibility with web platforms. +web-compat = [] cpu = [] cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"] cuda = [] diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 37842371..ce742aa9 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -54,6 +54,7 @@ impl GpuMultibodySet { // Per-multibody maxima (not per-env sums) for the uniform loop bounds. let mut max_mb_ndofs = 0u32; let mut max_mb_links = 0u32; + let mut max_mb_joint_constraints = 0u32; let mut global_max_dofs = 0u32; let mut global_max_jac = 0u32; let mut global_max_mm = 0u32; @@ -151,6 +152,7 @@ impl GpuMultibodySet { // among the joint constraints). let num_couplings = mb.couplings().len() as u32; let max_constraints = max_constraints + num_couplings; + max_mb_joint_constraints = max_mb_joint_constraints.max(max_constraints); infos.push(MultibodyInfo { first_link, @@ -499,6 +501,12 @@ impl GpuMultibodySet { has_joint_constraints: all_infos.iter().any(|info| info.max_constraints > 0), multibody_info: Tensor::vector(backend, &all_infos, storage).unwrap(), + max_contact_constraints: Tensor::scalar( + backend, + 0u32, + BufferUsages::STORAGE | BufferUsages::UNIFORM, + ) + .unwrap(), links_static: Tensor::vector(backend, &all_statics, storage | BufferUsages::COPY_DST) .unwrap(), links_static_mirror: all_statics.clone(), @@ -678,6 +686,7 @@ impl GpuMultibodySet { mb_imp_joint_max_color_group_len: 0, max_ndofs: max_mb_ndofs, max_links: max_mb_links, + max_joint_constraints: max_mb_joint_constraints, joint_constraints_per_batch: cons_cap, joint_constraint_columns_per_batch: cons_col_cap, contact_constraints_per_batch: contact_cons_cap, diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 5f7c61db..5a017157 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -51,6 +51,9 @@ pub struct GpuMultibodySet { /// Per-batch multibody descriptors. pub(super) multibody_info: Tensor, + /// Max `contact_constraint_count` across every multibody, written each + /// step by `gpu_mb_compute_solve_bounds`. + pub(super) max_contact_constraints: Tensor, /// Per-batch static link data. pub(super) links_static: Tensor, /// CPU-side mirror of [`Self::links_static`] used to support runtime @@ -139,6 +142,9 @@ pub struct GpuMultibodySet { /// Max link count across every multibody in every batch (CPU mirror of /// `BatchIndices::mb_max_links`). pub(super) max_links: u32, + /// Max joint-constraint slot count across every multibody in every batch + /// (CPU mirror of `BatchIndices::mb_max_joint_constraints`). + pub(super) max_joint_constraints: u32, /// Largest color group across batches — the per-color dispatch width. pub(super) mb_imp_joint_max_color_group_len: u32, /// Per-batch capacities of the joint / contact constraint slabs (CPU-side @@ -455,6 +461,7 @@ impl GpuMultibodySet { dst.mb_imp_joint_color_groups_batch_capacity = self.mb_imp_joint_num_colors.max(1); dst.mb_max_ndofs = self.max_ndofs; dst.mb_max_links = self.max_links; + dst.mb_max_joint_constraints = self.max_joint_constraints; dst.mb_pack_lanes = self.pack_lanes(); dst.coriolis_w_section_offset = self.coriolis_entries_per_batch * self.num_batches; dst.i_coriolis_dt_section_offset = 2 * self.coriolis_entries_per_batch * self.num_batches; diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 11f2d1a1..f52a414a 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -5,9 +5,9 @@ use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ GpuMbApplyContactRestitution, GpuMbBuildContactDelassus, GpuMbComputeDynamicsPre, - GpuMbFinalizeContactConstraints, GpuMbFinalizeImpulseJointConstraints, GpuMbGravityAndLu, - GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, - GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, + GpuMbComputeSolveBounds, GpuMbFinalizeContactConstraints, GpuMbFinalizeImpulseJointConstraints, + GpuMbGravityAndLu, GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, + GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, GpuMbRemoveImpulseJointConstraintBias, GpuMbSeedContactRestitution, GpuMbSnapshotContactWarmstart, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, GpuMbSolveJoints, GpuMbStashContactsLen, @@ -44,6 +44,9 @@ pub struct GpuMultibodySolver { /// Fills the per-multibody Delassus blocks (`D = J M⁻¹ Jᵀ` + free-body /// 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. + compute_solve_bounds: GpuMbComputeSolveBounds, /// Constraint-space contact sweep: `a = J·u` tracked incrementally in /// shared memory via the Delassus rows, breaking the per-iteration /// dof-space latency chain. @@ -360,6 +363,17 @@ impl GpuMultibodySolver { )?; } + { + let mut pass = encoder.begin_pass("[RBD] mbb/solve-bounds", timestamps.as_deref_mut()); + self.compute_solve_bounds.call( + &mut pass, + MB_LU_LANES, + &mb.multibody_info, + &mut mb.max_contact_constraints, + args.batch_indices, + )?; + } + // Delassus blocks for the constraint-space contact sweep (consumes // the columns finalized just above). if let Some(delassus) = &mut mb.contact_delassus { @@ -417,6 +431,7 @@ impl GpuMultibodySolver { delassus, use_bias, args.batch_indices, + &mb.max_contact_constraints, &mut mb.dof_state, args.solver_vels, )?; @@ -432,6 +447,7 @@ impl GpuMultibodySolver { &mb.contact_constraint_columns, use_bias, args.batch_indices, + &mb.max_contact_constraints, &mut mb.dof_state, args.solver_vels, )?; @@ -450,6 +466,7 @@ impl GpuMultibodySolver { &mb.contact_constraint_columns, use_bias, args.batch_indices, + &mb.max_contact_constraints, &mut mb.dof_state, args.solver_vels, )?; @@ -649,6 +666,7 @@ impl GpuMultibodySolver { &mb.contact_constraint_jacs, &mb.contact_constraint_columns, args.batch_indices, + &mb.max_contact_constraints, &mut mb.dof_state, args.solver_vels, ) diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 075e07f6..46211483 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -1042,6 +1042,7 @@ pub fn gpu_mb_apply_contact_restitution( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_constraint_jacs: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraint_columns: &[f32], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 5)] max_contact_constraints: &u32, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] dof_state: &mut [f32], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_vels: &mut [Velocity], #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS], @@ -1051,16 +1052,21 @@ pub fn gpu_mb_apply_contact_restitution( let batch_id = workgroup_id.y; let mb_idx = workgroup_id.x; let lane = local_id.x; - if mb_idx >= batch_ids.multibodies_len { + let in_range = mb_idx < batch_ids.multibodies_len; + #[cfg(not(feature = "web-compat"))] + if !in_range { return; } + let slot = if in_range { mb_idx } else { 0 }; - let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); + let mb = multibody_info.read(batch_ids.mbi(batch_id, slot as usize)); let ndofs = mb.ndofs; let count = mb.contact_constraint_count; + #[cfg(not(feature = "web-compat"))] if ndofs == 0 || count == 0 { return; } + let active = in_range && ndofs != 0 && count != 0; let colliders_start = batch_ids.coll_start(batch_id); let v_base = mb.first_dof as usize; @@ -1070,34 +1076,58 @@ pub fn gpu_mb_apply_contact_restitution( let col_base = batch_ids.mb_contact_constraint_columns_start(batch_id) + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; - if lane < ndofs { - dof_v[lane as usize] = dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)); + if active && lane < ndofs { + dof_v.write( + lane as usize, + dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)), + ); } workgroup_memory_barrier_with_group_sync(); - for s in 0..count { - let cons = contact_constraints.read(cons_base + s as usize); + #[cfg(feature = "web-compat")] + let contact_sweep_len = *max_contact_constraints; + #[cfg(not(feature = "web-compat"))] + let contact_sweep_len = count; + #[cfg(not(feature = "web-compat"))] + let _ = max_contact_constraints; + + for s in 0..contact_sweep_len { + let slot_active = active && s < count; + let cons_idx = if slot_active { + cons_base + s as usize + } else { + 0 + }; + let cons = contact_constraints.read(cons_idx); // Only approaching, load-bearing points bounce. - if cons.kind != MB_CONTACT_KIND_NORMAL - || cons.restitution_seed >= 0.0 - || cons.impulse <= 0.0 - { + let solve = slot_active + && cons.kind == MB_CONTACT_KIND_NORMAL + && cons.restitution_seed < 0.0 + && cons.impulse > 0.0; + #[cfg(not(feature = "web-compat"))] + if !solve { continue; } let col_offset = col_base + (s as usize) * dofs_stride; let is_self = cons.free_body_id == u32::MAX; - scratch[lane as usize] = if lane < ndofs { - contact_constraint_jacs.read(col_offset + lane as usize) * dof_v[lane as usize] - } else { - 0.0 - }; + if solve { + scratch.write( + lane as usize, + if lane < ndofs { + contact_constraint_jacs.read(col_offset + lane as usize) + * dof_v.read(lane as usize) + } else { + 0.0 + }, + ); + } workgroup_memory_barrier_with_group_sync(); - if lane == 0 { + if solve && lane == 0 { let mut j_dot_v = 0.0f32; for i in 0..ndofs { - j_dot_v += scratch[i as usize]; + j_dot_v += scratch.read(i as usize); } let free = if is_self { Velocity::default() @@ -1127,17 +1157,17 @@ pub fn gpu_mb_apply_contact_restitution( workgroup_memory_barrier_with_group_sync(); let delta = *delta_shared; - if delta != 0.0 && lane < ndofs { + if solve && delta != 0.0 && lane < ndofs { let col = contact_constraint_columns.read(col_offset + lane as usize); - dof_v[lane as usize] += delta * col; + dof_v.write(lane as usize, dof_v.read(lane as usize) + delta * col); } workgroup_memory_barrier_with_group_sync(); } - if lane < ndofs { + if active && lane < ndofs { dof_state.write( batch_ids.mbi(batch_id, v_base + lane as usize), - dof_v[lane as usize], + dof_v.read(lane as usize), ); } } diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index 19f669d2..b0dd32a5 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -121,8 +121,8 @@ fn emit_joint_constraints( if (motor_axes & (1 << axis)) != 0 { let has_limits = (limit_axes & (1 << axis)) != 0; - let limit_min = stat.data.limits[axis as usize].min; - let limit_max = stat.data.limits[axis as usize].max; + let limit_min = stat.data.limits.read(axis as usize).min; + let limit_max = stat.data.limits.read(axis as usize).max; let cons = build_motor_constraint( abs_dof, k, @@ -130,7 +130,7 @@ fn emit_joint_constraints( curr_pos, inv_dt, dt, - &stat.data.motors[axis as usize], + stat.data.motors.at(axis as usize), has_limits, limit_min, limit_max, @@ -145,8 +145,8 @@ fn emit_joint_constraints( axis, curr_pos, [ - stat.data.limits[axis as usize].min, - stat.data.limits[axis as usize].max, + stat.data.limits.read(axis as usize).min, + stat.data.limits.read(axis as usize).max, ], joint_erp_inv_dt, joint_cfm_coeff, @@ -172,8 +172,8 @@ fn emit_joint_constraints( axis, curr_pos, [ - stat.data.limits[axis as usize].min, - stat.data.limits[axis as usize].max, + stat.data.limits.read(axis as usize).min, + stat.data.limits.read(axis as usize).max, ], joint_erp_inv_dt, joint_cfm_coeff, @@ -183,8 +183,8 @@ fn emit_joint_constraints( } if (motor_axes & (1 << axis)) != 0 { let has_limits = (limit_axes & (1 << axis)) != 0; - let limit_min = stat.data.limits[axis as usize].min; - let limit_max = stat.data.limits[axis as usize].max; + let limit_min = stat.data.limits.read(axis as usize).min; + let limit_max = stat.data.limits.read(axis as usize).max; let cons = build_motor_constraint( abs_dof, k, @@ -192,7 +192,7 @@ fn emit_joint_constraints( curr_pos, inv_dt, dt, - &stat.data.motors[axis as usize], + stat.data.motors.at(axis as usize), has_limits, limit_min, limit_max, @@ -452,16 +452,21 @@ pub fn gpu_mb_init_joint_constraints( let mb_idx = workgroup_id.x; let lane = local_id.x; let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + let in_range = mb_idx < num_mb; + #[cfg(not(feature = "web-compat"))] + if !in_range { return; } + let slot = if in_range { mb_idx } else { 0 }; - let mb = batch_ids.ib(batch_id, multibody_info).read(mb_idx as usize); + let mb = batch_ids.ib(batch_id, multibody_info).read(slot as usize); let ndofs = mb.ndofs; // Uniform per workgroup: every lane of this group returns together. + #[cfg(not(feature = "web-compat"))] if ndofs == 0 { return; } + let active = in_range && ndofs != 0; let mb_mm_base = mb.mass_matrix_offset as usize; let piv = batch_ids.ivec(batch_id, mb.first_dof as usize); @@ -473,11 +478,13 @@ pub fn gpu_mb_init_joint_constraints( let m = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); // Stage 1: lane-parallel slot reset. - for s in StepRng::new(lane..mb.max_constraints, LANES) { - let mut cz: MultibodyJointConstraint = joint_constraints.read(cons_base + s as usize); - cz.kind = 0; - cz.impulse = 0.0; - joint_constraints.write(cons_base + s as usize, cz); + if active { + for s in StepRng::new(lane..mb.max_constraints, LANES) { + let mut cz: MultibodyJointConstraint = joint_constraints.read(cons_base + s as usize); + cz.kind = 0; + cz.impulse = 0.0; + joint_constraints.write(cons_base + s as usize, cz); + } } control_barrier::< @@ -490,7 +497,7 @@ pub fn gpu_mb_init_joint_constraints( >(); // Stage 2: serial metadata emission on lane 0. - if lane == 0 { + if active && lane == 0 { emit_joint_constraints( links_static, links_workspace, @@ -516,28 +523,30 @@ pub fn gpu_mb_init_joint_constraints( >(); // Stage 3: lane-parallel finalize. - for s in StepRng::new(lane..mb.max_constraints, LANES) { - let mut cons = joint_constraints.read(cons_base + s as usize); - if cons.kind == 0 { - continue; + if active { + for s in StepRng::new(lane..mb.max_constraints, LANES) { + let mut cons = joint_constraints.read(cons_base + s as usize); + if cons.kind == 0 { + continue; + } + let lhs = compute_constraint_column( + joint_constraint_columns, + col_base, + s, + dofs_stride, + ndofs, + cons.dof_id, + cons.dof2_id, + cons.coupling_coeff, + mass_matrices, + m, + lu_pivots, + piv, + ); + let cfm_gain = lhs * cons.cfm_coeff + cons.cfm_gain; + cons.cfm_gain = cfm_gain; + cons.inv_lhs = inv(lhs + cfm_gain); + joint_constraints.write(cons_base + s as usize, cons); } - let lhs = compute_constraint_column( - joint_constraint_columns, - col_base, - s, - dofs_stride, - ndofs, - cons.dof_id, - cons.dof2_id, - cons.coupling_coeff, - mass_matrices, - m, - lu_pivots, - piv, - ); - let cfm_gain = lhs * cons.cfm_coeff + cons.cfm_gain; - cons.cfm_gain = cfm_gain; - cons.inv_lhs = inv(lhs + cfm_gain); - joint_constraints.write(cons_base + s as usize, cons); } } diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index 0c898790..3110926f 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -34,6 +34,45 @@ fn cap_friction(t0: f32, t1: f32, limit: f32) -> (f32, f32) { } } +/// Calculate the maximum `contact_constraint_count` over every (multibody, batch). +/// +/// The output value is written into a uniform that will be passed to the other kernels +/// and used when `web-compat` is enabled. (Since it’s a uniform it can be used in conditions +/// without breaking uniform control flow.) +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_compute_solve_bounds( + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] max_contact_constraints: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, + #[spirv(workgroup)] scratch: &mut [u32; LANES as usize], +) { + let lane = local_id.x; + let total = batch_ids.multibodies_len * batch_ids.num_batches; + + let mut lane_max = 0u32; + for i in StepRng::new(lane..total, LANES) { + let count = multibody_info.read(i as usize).contact_constraint_count; + if count > lane_max { + lane_max = count; + } + } + scratch.write(lane as usize, lane_max); + workgroup_memory_barrier_with_group_sync(); + + if lane == 0 { + let mut max_count = 0u32; + for i in 0..LANES { + let v = scratch.read(i as usize); + if v > max_count { + max_count = v; + } + } + max_contact_constraints.write(0, max_count); + } +} + /// One PGS iteration over a multibody's joint (limit/motor) constraints followed /// by its contact constraints. /// @@ -53,6 +92,7 @@ pub fn gpu_mb_solve_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] contact_constraint_columns: &[f32], #[spirv(uniform, descriptor_set = 0, binding = 6)] use_bias: &u32, #[spirv(uniform, descriptor_set = 0, binding = 7)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 8)] max_contact_constraints: &u32, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] dof_state: &mut [f32], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_vels: &mut [Velocity], #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS], @@ -65,13 +105,17 @@ pub fn gpu_mb_solve_constraints( let mb_idx = workgroup_id.x; let lane = local_id.x; let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + let in_range = mb_idx < num_mb; + #[cfg(not(feature = "web-compat"))] + if !in_range { return; } + let slot = if in_range { mb_idx } else { 0 }; - let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); + let mb = multibody_info.read(batch_ids.mbi(batch_id, slot as usize)); let ndofs = mb.ndofs; // Uniform per workgroup: every lane of this group returns together. + #[cfg(not(feature = "web-compat"))] if ndofs == 0 { return; } @@ -91,60 +135,87 @@ pub fn gpu_mb_solve_constraints( + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; let contact_count = mb.contact_constraint_count; + #[cfg(not(feature = "web-compat"))] if mb.max_constraints == 0 && contact_count == 0 { // Nothing to solve. return; } + let active = in_range && ndofs != 0 && (mb.max_constraints != 0 || contact_count != 0); // Load the generalized velocities and accumulated contact impulses into // workgroup memory. - if lane < ndofs { - dof_v[lane as usize] = dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)); - } - for s in StepRng::new(lane..contact_count, LANES) { - imp_shared[s as usize] = contact_constraints.read(ccons_base + s as usize).impulse; + if active { + if lane < ndofs { + dof_v.write( + lane as usize, + dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)), + ); + } + for s in StepRng::new(lane..contact_count, LANES) { + imp_shared.write( + s as usize, + contact_constraints.read(ccons_base + s as usize).impulse, + ); + } } workgroup_memory_barrier_with_group_sync(); + #[cfg(feature = "web-compat")] + let joint_sweep_len = batch_ids.mb_max_joint_constraints; + #[cfg(not(feature = "web-compat"))] + let joint_sweep_len = mb.max_constraints; + // Joint limits/motors - for s in 0..mb.max_constraints { - let cons = joint_constraints.read(jcons_base + s as usize); - if cons.kind != MB_JOINT_KIND_LIMIT - && cons.kind != MB_JOINT_KIND_MOTOR - && cons.kind != MB_JOINT_KIND_COUPLING - { + for s in 0..joint_sweep_len { + let slot_active = active && s < mb.max_constraints; + let cons_idx = if slot_active { + jcons_base + s as usize + } else { + 0 + }; + let cons = joint_constraints.read(cons_idx); + let solve = slot_active + && (cons.kind == MB_JOINT_KIND_LIMIT + || cons.kind == MB_JOINT_KIND_MOTOR + || cons.kind == MB_JOINT_KIND_COUPLING); + #[cfg(not(feature = "web-compat"))] + if !solve { // Unused slot or inactive limit. continue; } - let rhs = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; - // Generalized `J·v` for `J = e_{dof_id} - coupling_coeff*e_{dof2_id}` - // (coupling rows); collapses to `v[dof_id]` for limit / motor rows - // (their `coupling_coeff` is 0). - let v_d = dof_v[cons.dof_id as usize] - cons.coupling_coeff * dof_v[cons.dof2_id as usize]; - let rhs_total = v_d + rhs; - let raw_imp = cons.impulse + cons.inv_lhs * (rhs_total - cons.cfm_gain * cons.impulse); - let mut new_imp = raw_imp; - if new_imp < cons.impulse_lo { - new_imp = cons.impulse_lo; - } - if new_imp > cons.impulse_hi { - new_imp = cons.impulse_hi; - } - let delta = new_imp - cons.impulse; + let mut delta = 0.0f32; + if solve { + let rhs = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; + // Generalized `J·v` for `J = e_{dof_id} - coupling_coeff*e_{dof2_id}` + // (coupling rows); collapses to `v[dof_id]` for limit / motor rows + // (their `coupling_coeff` is 0). + let v_d = dof_v.read(cons.dof_id as usize) + - cons.coupling_coeff * dof_v.read(cons.dof2_id as usize); + let rhs_total = v_d + rhs; + let raw_imp = cons.impulse + cons.inv_lhs * (rhs_total - cons.cfm_gain * cons.impulse); + let mut new_imp = raw_imp; + if new_imp < cons.impulse_lo { + new_imp = cons.impulse_lo; + } + if new_imp > cons.impulse_hi { + new_imp = cons.impulse_hi; + } + delta = new_imp - cons.impulse; - if lane == 0 { - let mut cons = cons; - cons.impulse = new_imp; - joint_constraints.write(jcons_base + s as usize, cons); + if lane == 0 { + let mut cons = cons; + cons.impulse = new_imp; + joint_constraints.write(jcons_base + s as usize, cons); + } } - // All lanes read `dof_v[dof_id]` above; sync before overwriting it. + // All lanes read `dof_v.read(dof_id)` above; sync before overwriting it. workgroup_memory_barrier_with_group_sync(); - if lane < ndofs { + if solve && lane < ndofs { let col = joint_constraint_columns .read(jcol_base + (s as usize) * dofs_stride + lane as usize); - dof_v[lane as usize] -= delta * col; + dof_v.write(lane as usize, dof_v.read(lane as usize) - delta * col); } workgroup_memory_barrier_with_group_sync(); } @@ -152,17 +223,35 @@ pub fn gpu_mb_solve_constraints( // Contacts. In 3D the two friction rows of a contact point are solved // together so their impulse can be capped to the friction cone; the second // row is handled by its sibling and skipped here. - for s in 0..contact_count { - let cons = contact_constraints.read(ccons_base + s as usize); + #[cfg(feature = "web-compat")] + let contact_sweep_len = *max_contact_constraints; + #[cfg(not(feature = "web-compat"))] + let contact_sweep_len = contact_count; + #[cfg(not(feature = "web-compat"))] + let _ = max_contact_constraints; + + for s in 0..contact_sweep_len { + let slot_active = active && s < contact_count; + let cons_idx = if slot_active { + ccons_base + s as usize + } else { + 0 + }; + let cons = contact_constraints.read(cons_idx); let is_tangent = cons.kind == MB_CONTACT_KIND_TANGENT; - // Friction is only solved during the relaxation phase. - if use_bias && is_tangent { - continue; - } + // Friction is only solved during the relaxation phase, and in 3D a + // tangent pair is solved by its first row only. #[cfg(feature = "dim3")] - if is_tangent && s != cons.normal_constraint_slot + 1 { + let solve = slot_active + && !(use_bias && is_tangent) + && !(is_tangent && s != cons.normal_constraint_slot + 1); + #[cfg(feature = "dim2")] + let solve = slot_active && !(use_bias && is_tangent); + #[cfg(not(feature = "web-compat"))] + if !solve { continue; } + #[cfg(feature = "dim3")] let has_pair = is_tangent; #[cfg(feature = "dim2")] @@ -174,36 +263,46 @@ pub fn gpu_mb_solve_constraints( // Multibody side of J · u, one product per lane; lane 0 sums them in // DOF order. - scratch[lane as usize] = if lane < ndofs { - contact_constraint_jacs.read(col_offset + lane as usize) * dof_v[lane as usize] - } else { - 0.0 - }; + if solve { + scratch.write( + lane as usize, + if lane < ndofs { + contact_constraint_jacs.read(col_offset + lane as usize) + * dof_v.read(lane as usize) + } else { + 0.0 + }, + ); + } workgroup_memory_barrier_with_group_sync(); let mut j_dot_v0 = 0.0f32; - if lane == 0 { + if solve && lane == 0 { for i in 0..ndofs { - j_dot_v0 += scratch[i as usize]; + j_dot_v0 += scratch.read(i as usize); } } workgroup_memory_barrier_with_group_sync(); - if has_pair { - scratch[lane as usize] = if lane < ndofs { - contact_constraint_jacs.read(col_offset2 + lane as usize) * dof_v[lane as usize] - } else { - 0.0 - }; + if solve && has_pair { + scratch.write( + lane as usize, + if lane < ndofs { + contact_constraint_jacs.read(col_offset2 + lane as usize) + * dof_v.read(lane as usize) + } else { + 0.0 + }, + ); } workgroup_memory_barrier_with_group_sync(); - if lane == 0 { + if solve && lane == 0 { let cons2 = contact_constraints.read(ccons_base + (s + 1) as usize); let mut j_dot_v1 = 0.0f32; if has_pair { for i in 0..ndofs { - j_dot_v1 += scratch[i as usize]; + j_dot_v1 += scratch.read(i as usize); } } // Free-body side stays lane-0-local. @@ -220,12 +319,12 @@ pub fn gpu_mb_solve_constraints( } let cfm_factor = if use_bias { cons.cfm_factor } else { 1.0 }; - let impulse0 = imp_shared[s as usize]; + let impulse0 = imp_shared.read(s as usize); let rhs0 = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; let raw0 = cfm_factor * (impulse0 - cons.inv_lhs * (j_dot_v0 + rhs0)); let impulse1 = if has_pair { - imp_shared[(s + 1) as usize] + imp_shared.read((s + 1) as usize) } else { 0.0 }; @@ -243,7 +342,8 @@ 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 { - let limit = cons.friction_coeff * imp_shared[cons.normal_constraint_slot as usize]; + let limit = + cons.friction_coeff * imp_shared.read(cons.normal_constraint_slot as usize); cap_friction(raw0, raw1, limit) } else if raw0 < 0.0 { (0.0, 0.0) @@ -253,9 +353,9 @@ pub fn gpu_mb_solve_constraints( let delta0 = new0 - impulse0; let delta1 = if has_pair { new1 - impulse1 } else { 0.0 }; - imp_shared[s as usize] = new0; + imp_shared.write(s as usize, new0); if has_pair { - imp_shared[(s + 1) as usize] = new1; + imp_shared.write((s + 1) as usize, new1); } *delta_shared = delta0; *delta2_shared = delta1; @@ -273,33 +373,35 @@ pub fn gpu_mb_solve_constraints( } workgroup_memory_barrier_with_group_sync(); - // Per-lane `dof_v[lane]` update. + // Per-lane `dof_v.read(lane)` update. let delta0 = *delta_shared; let delta1 = *delta2_shared; - if lane < ndofs { + if solve && lane < ndofs { if delta0 != 0.0 { let col = contact_constraint_columns.read(col_offset + lane as usize); - dof_v[lane as usize] += delta0 * col; + dof_v.write(lane as usize, dof_v.read(lane as usize) + delta0 * col); } if has_pair && delta1 != 0.0 { let col = contact_constraint_columns.read(col_offset2 + lane as usize); - dof_v[lane as usize] += delta1 * col; + dof_v.write(lane as usize, dof_v.read(lane as usize) + delta1 * col); } } workgroup_memory_barrier_with_group_sync(); } // Writeback - if lane < ndofs { - dof_state.write( - batch_ids.mbi(batch_id, v_base + lane as usize), - dof_v[lane as usize], - ); - } - for s in StepRng::new(lane..contact_count, LANES) { - let mut cons = contact_constraints.read(ccons_base + s as usize); - cons.impulse = imp_shared[s as usize]; - contact_constraints.write(ccons_base + s as usize, cons); + if active { + if lane < ndofs { + dof_state.write( + batch_ids.mbi(batch_id, v_base + lane as usize), + dof_v.read(lane as usize), + ); + } + for s in StepRng::new(lane..contact_count, LANES) { + let mut cons = contact_constraints.read(ccons_base + s as usize); + cons.impulse = imp_shared.read(s as usize); + contact_constraints.write(ccons_base + s as usize, cons); + } } } @@ -324,16 +426,21 @@ pub fn gpu_mb_solve_joints( let mb_idx = workgroup_id.x; let lane = local_id.x; let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + let in_range = mb_idx < num_mb; + #[cfg(not(feature = "web-compat"))] + if !in_range { return; } + let slot = if in_range { mb_idx } else { 0 }; - let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); + let mb = multibody_info.read(batch_ids.mbi(batch_id, slot as usize)); let ndofs = mb.ndofs; // Uniform per workgroup: every lane of this group returns together. + #[cfg(not(feature = "web-compat"))] if ndofs == 0 || mb.max_constraints == 0 { return; } + let active = in_range && ndofs != 0 && mb.max_constraints != 0; let use_bias = *use_bias != 0; let v_base = mb.first_dof as usize; @@ -342,56 +449,76 @@ pub fn gpu_mb_solve_joints( let jcol_base = batch_ids.mb_joint_constraint_columns_start(batch_id) + (mb.first_constraint as usize) * dofs_stride; - if lane < ndofs { - dof_v[lane as usize] = dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)); + if active && lane < ndofs { + dof_v.write( + lane as usize, + dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)), + ); } workgroup_memory_barrier_with_group_sync(); - for s in 0..mb.max_constraints { - let cons = joint_constraints.read(jcons_base + s as usize); - if cons.kind != MB_JOINT_KIND_LIMIT - && cons.kind != MB_JOINT_KIND_MOTOR - && cons.kind != MB_JOINT_KIND_COUPLING - { + #[cfg(feature = "web-compat")] + let joint_sweep_len = batch_ids.mb_max_joint_constraints; + #[cfg(not(feature = "web-compat"))] + let joint_sweep_len = mb.max_constraints; + + for s in 0..joint_sweep_len { + let slot_active = active && s < mb.max_constraints; + let cons_idx = if slot_active { + jcons_base + s as usize + } else { + 0 + }; + let cons = joint_constraints.read(cons_idx); + let solve = slot_active + && (cons.kind == MB_JOINT_KIND_LIMIT + || cons.kind == MB_JOINT_KIND_MOTOR + || cons.kind == MB_JOINT_KIND_COUPLING); + #[cfg(not(feature = "web-compat"))] + if !solve { // Unused slot or inactive limit. continue; } - let rhs = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; - // Generalized `J·v` for `J = e_{dof_id} - coupling_coeff*e_{dof2_id}` - // (coupling rows); collapses to `v[dof_id]` for limit / motor rows - // (their `coupling_coeff` is 0). - let v_d = dof_v[cons.dof_id as usize] - cons.coupling_coeff * dof_v[cons.dof2_id as usize]; - let rhs_total = v_d + rhs; - let raw_imp = cons.impulse + cons.inv_lhs * (rhs_total - cons.cfm_gain * cons.impulse); - let mut new_imp = raw_imp; - if new_imp < cons.impulse_lo { - new_imp = cons.impulse_lo; - } - if new_imp > cons.impulse_hi { - new_imp = cons.impulse_hi; - } - let delta = new_imp - cons.impulse; + let mut delta = 0.0f32; + if solve { + let rhs = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; + // Generalized `J·v` for `J = e_{dof_id} - coupling_coeff*e_{dof2_id}` + // (coupling rows); collapses to `v[dof_id]` for limit / motor rows + // (their `coupling_coeff` is 0). + let v_d = dof_v.read(cons.dof_id as usize) + - cons.coupling_coeff * dof_v.read(cons.dof2_id as usize); + let rhs_total = v_d + rhs; + let raw_imp = cons.impulse + cons.inv_lhs * (rhs_total - cons.cfm_gain * cons.impulse); + let mut new_imp = raw_imp; + if new_imp < cons.impulse_lo { + new_imp = cons.impulse_lo; + } + if new_imp > cons.impulse_hi { + new_imp = cons.impulse_hi; + } + delta = new_imp - cons.impulse; - if lane == 0 { - let mut cons = cons; - cons.impulse = new_imp; - joint_constraints.write(jcons_base + s as usize, cons); + if lane == 0 { + let mut cons = cons; + cons.impulse = new_imp; + joint_constraints.write(jcons_base + s as usize, cons); + } } workgroup_memory_barrier_with_group_sync(); - if lane < ndofs { + if solve && lane < ndofs { let col = joint_constraint_columns .read(jcol_base + (s as usize) * dofs_stride + lane as usize); - dof_v[lane as usize] -= delta * col; + dof_v.write(lane as usize, dof_v.read(lane as usize) - delta * col); } workgroup_memory_barrier_with_group_sync(); } - if lane < ndofs { + if active && lane < ndofs { dof_state.write( batch_ids.mbi(batch_id, v_base + lane as usize), - dof_v[lane as usize], + dof_v.read(lane as usize), ); } } @@ -489,6 +616,7 @@ pub fn gpu_mb_solve_contacts_delassus( #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] delassus: &[f32], #[spirv(uniform, descriptor_set = 0, binding = 5)] use_bias: &u32, #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 7)] max_contact_constraints: &u32, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] dof_state: &mut [f32], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_vels: &mut [Velocity], #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS], @@ -505,17 +633,22 @@ pub fn gpu_mb_solve_contacts_delassus( let mb_idx = workgroup_id.x; let lane = local_id.x; let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + let in_range = mb_idx < num_mb; + #[cfg(not(feature = "web-compat"))] + if !in_range { return; } + let slot = if in_range { mb_idx } else { 0 }; - let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); + let mb = multibody_info.read(batch_ids.mbi(batch_id, slot as usize)); let ndofs = mb.ndofs; let count = mb.contact_constraint_count; // Uniform per workgroup: every lane of this group returns together. + #[cfg(not(feature = "web-compat"))] if ndofs == 0 || count == 0 { return; } + let active = in_range && ndofs != 0 && count != 0; let use_bias = *use_bias != 0; let v_base = mb.first_dof as usize; @@ -529,154 +662,193 @@ pub fn gpu_mb_solve_contacts_delassus( * (MAXC as usize) * (MAXC as usize); - if lane < ndofs { - dof_v[lane as usize] = dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)); + if active && lane < ndofs { + dof_v.write( + lane as usize, + dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)), + ); } // Preload the per-constraint solve scalars into shared SoA arrays so the // serial recurrence below never touches storage on its critical path. // `meta` packs the kind, the paired normal slot, and whether the // free-body side needs the fire-and-forget storage velocity update. - for s in StepRng::new(lane..count, LANES) { - let cons = contact_constraints.read(cons_base + s as usize); - imp_shared[s as usize] = cons.impulse; - rhs_shared[s as usize] = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; - inv_lhs_shared[s as usize] = cons.inv_lhs; - cfm_shared[s as usize] = if use_bias { cons.cfm_factor } else { 1.0 }; - friction_shared[s as usize] = cons.friction_coeff; - let is_self = cons.free_body_id == u32::MAX; - let free_active = - !is_self && (cons.free_body_im != 0.0 || gdot(cons.ii_ang_jac, cons.ii_ang_jac) != 0.0); - meta_shared[s as usize] = (cons.kind & 0xff) - | ((cons.normal_constraint_slot & 0xffff) << 8) - | (if free_active { 1 << 24 } else { 0 }); + if active { + for s in StepRng::new(lane..count, LANES) { + let cons = contact_constraints.read(cons_base + s as usize); + imp_shared.write(s as usize, cons.impulse); + rhs_shared.write( + s as usize, + if use_bias { cons.rhs } else { cons.rhs_wo_bias }, + ); + inv_lhs_shared.write(s as usize, cons.inv_lhs); + cfm_shared.write(s as usize, if use_bias { cons.cfm_factor } else { 1.0 }); + friction_shared.write(s as usize, cons.friction_coeff); + let is_self = cons.free_body_id == u32::MAX; + let free_active = !is_self + && (cons.free_body_im != 0.0 || gdot(cons.ii_ang_jac, cons.ii_ang_jac) != 0.0); + meta_shared.write( + s as usize, + (cons.kind & 0xff) + | ((cons.normal_constraint_slot & 0xffff) << 8) + | (if free_active { 1 << 24 } else { 0 }), + ); + } } workgroup_memory_barrier_with_group_sync(); // Fresh `a[s] = J_s · u` under the current (post-joint-sweep, post- // warmstart) velocities. - for s in StepRng::new(lane..count, LANES) { - let jac_off = col_base + (s as usize) * dofs_stride; - let mut dot = 0.0f32; - for i in 0..ndofs { - dot += contact_constraint_jacs.read(jac_off + i as usize) * dof_v[i as usize]; - } - let cons = contact_constraints.read(cons_base + s as usize); - if cons.free_body_id != u32::MAX { - let free = solver_vels.read(colliders_start + cons.free_body_id as usize); - dot += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); + if active { + for s in StepRng::new(lane..count, LANES) { + let jac_off = col_base + (s as usize) * dofs_stride; + let mut dot = 0.0f32; + for i in 0..ndofs { + dot += contact_constraint_jacs.read(jac_off + i as usize) * dof_v.read(i as usize); + } + let cons = contact_constraints.read(cons_base + s as usize); + if cons.free_body_id != u32::MAX { + let free = solver_vels.read(colliders_start + cons.free_body_id as usize); + dot += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); + } + a_shared.write(s as usize, dot); } - a_shared[s as usize] = dot; } workgroup_memory_barrier_with_group_sync(); // In 3D the two friction rows of a contact point are solved together so // their impulse can be capped to the friction cone; the second row is // handled by its sibling and skipped here. - for s in 0..count { - let meta = meta_shared[s as usize]; + #[cfg(feature = "web-compat")] + let contact_sweep_len = *max_contact_constraints; + #[cfg(not(feature = "web-compat"))] + let contact_sweep_len = count; + #[cfg(not(feature = "web-compat"))] + let _ = max_contact_constraints; + + for s in 0..contact_sweep_len { + let slot_active = active && s < count; + let meta = if slot_active { + meta_shared.read(s as usize) + } else { + 0 + }; let kind = meta & 0xff; let normal_slot = (meta >> 8) & 0xffff; let free_active = (meta >> 24) != 0; let is_tangent = kind == MB_CONTACT_KIND_TANGENT; - if use_bias && is_tangent { - // Friction is only solved during the stabilization sweep. - continue; - } + // Friction is only solved during the stabilization sweep, and in 3D a + // tangent pair is solved by its first row only. #[cfg(feature = "dim3")] - if is_tangent && s != normal_slot + 1 { + let solve = + slot_active && !(use_bias && is_tangent) && !(is_tangent && s != normal_slot + 1); + #[cfg(feature = "dim2")] + let solve = slot_active && !(use_bias && is_tangent); + #[cfg(not(feature = "web-compat"))] + if !solve { continue; } + #[cfg(feature = "dim3")] let has_pair = is_tangent; #[cfg(feature = "dim2")] let has_pair = false; - let impulse0 = imp_shared[s as usize]; - let raw0 = cfm_shared[s as usize] - * (impulse0 - - inv_lhs_shared[s as usize] * (a_shared[s as usize] + rhs_shared[s as usize])); - let impulse1 = if has_pair { - imp_shared[(s + 1) as usize] - } else { - 0.0 - }; - let raw1 = if has_pair { - cfm_shared[(s + 1) as usize] - * (impulse1 - - inv_lhs_shared[(s + 1) as usize] - * (a_shared[(s + 1) as usize] + rhs_shared[(s + 1) as usize])) - } else { - 0.0 - }; - - let (new0, new1) = if is_tangent { - let limit = friction_shared[s as usize] * imp_shared[normal_slot as usize]; - cap_friction(raw0, raw1, limit) - } else if raw0 < 0.0 { - (0.0, 0.0) - } else { - (raw0, 0.0) - }; - let delta0 = new0 - impulse0; - let delta1 = if has_pair { new1 - impulse1 } else { 0.0 }; + if solve { + let impulse0 = imp_shared.read(s as usize); + let raw0 = cfm_shared.read(s as usize) + * (impulse0 + - inv_lhs_shared.read(s as usize) + * (a_shared.read(s as usize) + rhs_shared.read(s as usize))); + let impulse1 = if has_pair { + imp_shared.read((s + 1) as usize) + } else { + 0.0 + }; + let raw1 = if has_pair { + cfm_shared.read((s + 1) as usize) + * (impulse1 + - inv_lhs_shared.read((s + 1) as usize) + * (a_shared.read((s + 1) as usize) + rhs_shared.read((s + 1) as usize))) + } else { + 0.0 + }; - if delta0 != 0.0 || delta1 != 0.0 { - if lane == 0 { - imp_shared[s as usize] = new0; - if has_pair { - imp_shared[(s + 1) as usize] = new1; - } + let (new0, new1) = if is_tangent { + let limit = + friction_shared.read(s as usize) * imp_shared.read(normal_slot as usize); + cap_friction(raw0, raw1, limit) + } else if raw0 < 0.0 { + (0.0, 0.0) + } else { + (raw0, 0.0) + }; + let delta0 = new0 - impulse0; + let delta1 = if has_pair { new1 - impulse1 } else { 0.0 }; - if free_active { - let cons = contact_constraints.read(cons_base + s as usize); - let mut free = solver_vels.read(colliders_start + cons.free_body_id as usize); - free.linear += cons.lin_jac * (cons.free_body_im * delta0); - free.angular += cons.ii_ang_jac * delta0; + if delta0 != 0.0 || delta1 != 0.0 { + if lane == 0 { + imp_shared.write(s as usize, new0); if has_pair { - let cons2 = contact_constraints.read(cons_base + (s + 1) as usize); - free.linear += cons2.lin_jac * (cons2.free_body_im * delta1); - free.angular += cons2.ii_ang_jac * delta1; + imp_shared.write((s + 1) as usize, new1); + } + + if free_active { + let cons = contact_constraints.read(cons_base + s as usize); + let mut free = + solver_vels.read(colliders_start + cons.free_body_id as usize); + free.linear += cons.lin_jac * (cons.free_body_im * delta0); + free.angular += cons.ii_ang_jac * delta0; + if has_pair { + let cons2 = contact_constraints.read(cons_base + (s + 1) as usize); + free.linear += cons2.lin_jac * (cons2.free_body_im * delta1); + free.angular += cons2.ii_ang_jac * delta1; + } + solver_vels.write(colliders_start + cons.free_body_id as usize, free); } - solver_vels.write(colliders_start + cons.free_body_id as usize, free); } - } - // Lane-parallel Delassus row update (row `s` is contiguous), plus - // the off-path dof update (each lane owns its DOF). - let d_row = d_base + (s * MAXC) as usize; - let d_row2 = d_base + ((s + 1) * MAXC) as usize; - for j in StepRng::new(lane..count, LANES) { - let mut acc = delta0 * delassus.read(d_row + j as usize); - if has_pair { - acc += delta1 * delassus.read(d_row2 + j as usize); + // Lane-parallel Delassus row update (row `s` is contiguous), plus + // the off-path dof update (each lane owns its DOF). + let d_row = d_base + (s * MAXC) as usize; + let d_row2 = d_base + ((s + 1) * MAXC) as usize; + for j in StepRng::new(lane..count, LANES) { + let mut acc = delta0 * delassus.read(d_row + j as usize); + if has_pair { + acc += delta1 * delassus.read(d_row2 + j as usize); + } + a_shared.write(j as usize, a_shared.read(j as usize) + acc); } - a_shared[j as usize] += acc; - } - if lane < ndofs { - let col = contact_constraint_columns - .read(col_base + (s as usize) * dofs_stride + lane as usize); - dof_v[lane as usize] += delta0 * col; - if has_pair { - let col2 = contact_constraint_columns - .read(col_base + ((s + 1) as usize) * dofs_stride + lane as usize); - dof_v[lane as usize] += delta1 * col2; + if lane < ndofs { + let col = contact_constraint_columns + .read(col_base + (s as usize) * dofs_stride + lane as usize); + dof_v.write(lane as usize, dof_v.read(lane as usize) + delta0 * col); + if has_pair { + let col2 = contact_constraint_columns + .read(col_base + ((s + 1) as usize) * dofs_stride + lane as usize); + dof_v.write(lane as usize, dof_v.read(lane as usize) + delta1 * col2); + } } + #[cfg(not(feature = "web-compat"))] + workgroup_memory_barrier_with_group_sync(); } - workgroup_memory_barrier_with_group_sync(); } + #[cfg(feature = "web-compat")] + workgroup_memory_barrier_with_group_sync(); } // Writeback. - if lane < ndofs { - dof_state.write( - batch_ids.mbi(batch_id, v_base + lane as usize), - dof_v[lane as usize], - ); - } - for s in StepRng::new(lane..count, LANES) { - let mut cons = contact_constraints.read(cons_base + s as usize); - cons.impulse = imp_shared[s as usize]; - contact_constraints.write(cons_base + s as usize, cons); + if active { + if lane < ndofs { + dof_state.write( + batch_ids.mbi(batch_id, v_base + lane as usize), + dof_v.read(lane as usize), + ); + } + for s in StepRng::new(lane..count, LANES) { + let mut cons = contact_constraints.read(cons_base + s as usize); + cons.impulse = imp_shared.read(s as usize); + contact_constraints.write(cons_base + s as usize, cons); + } } } diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 327ddaf9..596a8fa5 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -476,7 +476,9 @@ pub fn gpu_warmstart_fused( let num_colors = *num_colors; let base = (batch_id * stride) as usize; - if color_starts.read(base + 1) == color_starts.read(base + num_colors as usize + 1) { + let any_work = color_starts.read(base + 1) != color_starts.read(base + num_colors as usize + 1); + #[cfg(not(feature = "web-compat"))] + if !any_work { // Every color bucket is empty. return; } @@ -485,24 +487,27 @@ pub fn gpu_warmstart_fused( let bucket = base + color as usize; let start = color_starts.read(bucket); let end = color_starts.read(bucket + 1); + #[cfg(not(feature = "web-compat"))] if start == end { // Empty color. continue; } - 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; + if any_work && 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]; + 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); + constraint.warmstart_constraint(&mut solver_vel1, &mut solver_vel2); - solver_vels[solver_id1] = solver_vel1; - solver_vels[solver_id2] = solver_vel2; + solver_vels[solver_id1] = solver_vel1; + solver_vels[solver_id2] = solver_vel2; + } } control_barrier::< @@ -545,7 +550,9 @@ pub fn gpu_step_gauss_seidel_fused( // Early-out / empty-color skip: see `gpu_warmstart_fused`. let base = (batch_id * stride) as usize; - if color_starts.read(base + 1) == color_starts.read(base + num_colors as usize + 1) { + let any_work = color_starts.read(base + 1) != color_starts.read(base + num_colors as usize + 1); + #[cfg(not(feature = "web-compat"))] + if !any_work { return; } @@ -553,26 +560,29 @@ pub fn gpu_step_gauss_seidel_fused( let bucket = base + color as usize; let start = color_starts.read(bucket); let end = color_starts.read(bucket + 1); + #[cfg(not(feature = "web-compat"))] if start == end { continue; } - 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; + if any_work && 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]; + 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, - ); + constraints[i as usize].solve_constraint_gauss_seidel( + &mut solver_vel1, + &mut solver_vel2, + use_bias, + ); - solver_vels[solver_id1] = solver_vel1; - solver_vels[solver_id2] = solver_vel2; + solver_vels[solver_id1] = solver_vel1; + solver_vels[solver_id2] = solver_vel2; + } } control_barrier::< diff --git a/src_rbd_shaders/dynamics/solver_utils.rs b/src_rbd_shaders/dynamics/solver_utils.rs index 869a6b77..1d199979 100644 --- a/src_rbd_shaders/dynamics/solver_utils.rs +++ b/src_rbd_shaders/dynamics/solver_utils.rs @@ -601,7 +601,10 @@ impl TwoBodyConstraint { - tangent_a.dot(solver_vel2.linear) + gdot(**c.torque_dir_b.at(0), solver_vel2.angular) + c.rhs_wo_bias.read(0); - let new_impulse = (c.impulse.read(0) - c.r.read(0) * dvel).clamp(-limit, limit); + // 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 diff --git a/src_rbd_shaders/utils/indices.rs b/src_rbd_shaders/utils/indices.rs index 806a48de..8c1627c4 100644 --- a/src_rbd_shaders/utils/indices.rs +++ b/src_rbd_shaders/utils/indices.rs @@ -65,6 +65,8 @@ pub struct BatchIndices { /// Lanes per multibody for the packed per-multibody workgroup kernels: /// `next_power_of_two(mb_max_ndofs).clamp(8, 64)`. pub mb_pack_lanes: u32, + /// Max `max_constraints` across every multibody in every batch. + 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.