fix(rbd) web compatibility

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 = []

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 = []

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);
 }

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);
 }

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 = []

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 = []

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,

src_rbd/dynamics/multibody/multibody_set.rs

@@ -51,6 +51,9 @@ pub struct GpuMultibodySet {
 
     /// Per-batch multibody descriptors.
     pub(super) multibody_info: Tensor<MultibodyInfo>,
+    /// Max `contact_constraint_count` across every multibody, written each
+    /// step by `gpu_mb_compute_solve_bounds`.
+    pub(super) max_contact_constraints: Tensor<u32>,
     /// Per-batch static link data.
     pub(super) links_static: Tensor<MultibodyLinkStatic>,
     /// 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;

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,
         )

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),
         );
     }
 }