First public release of Rapier.v0.1.0

1 files changed, 485 insertions, 0 deletions

diff --git a/src/dynamics/solver/parallel_velocity_solver.rs b/src/dynamics/solver/parallel_velocity_solver.rs
new file mode 100644
index 0000000..303a18f
--- /dev/null
+++ b/src/dynamics/solver/parallel_velocity_solver.rs
@@ -0,0 +1,485 @@
+use super::ParallelInteractionGroups;
+use super::{AnyJointVelocityConstraint, AnyVelocityConstraint, DeltaVel, ThreadContext};
+use crate::dynamics::solver::categorization::{categorize_joints, categorize_velocity_contacts};
+use crate::dynamics::solver::{InteractionGroups, VelocityConstraint, VelocityGroundConstraint};
+use crate::dynamics::{IntegrationParameters, JointGraphEdge, RigidBodySet};
+use crate::geometry::ContactManifold;
+#[cfg(feature = "simd-is-enabled")]
+use crate::{
+    dynamics::solver::{WVelocityConstraint, WVelocityGroundConstraint},
+    simd::SIMD_WIDTH,
+};
+use std::sync::atomic::Ordering;
+
+pub(crate) enum VelocityConstraintDesc {
+    NongroundNongrouped(usize),
+    GroundNongrouped(usize),
+    #[cfg(feature = "simd-is-enabled")]
+    NongroundGrouped([usize; SIMD_WIDTH]),
+    #[cfg(feature = "simd-is-enabled")]
+    GroundGrouped([usize; SIMD_WIDTH]),
+}
+
+pub(crate) struct ParallelVelocitySolverPart<Constraint> {
+    pub not_ground_interactions: Vec<usize>,
+    pub ground_interactions: Vec<usize>,
+    pub interaction_groups: InteractionGroups,
+    pub ground_interaction_groups: InteractionGroups,
+    pub constraints: Vec<Constraint>,
+    pub constraint_descs: Vec<(usize, VelocityConstraintDesc)>,
+    pub parallel_desc_groups: Vec<usize>,
+}
+
+impl<Constraint> ParallelVelocitySolverPart<Constraint> {
+    pub fn new() -> Self {
+        Self {
+            not_ground_interactions: Vec::new(),
+            ground_interactions: Vec::new(),
+            interaction_groups: InteractionGroups::new(),
+            ground_interaction_groups: InteractionGroups::new(),
+            constraints: Vec::new(),
+            constraint_descs: Vec::new(),
+            parallel_desc_groups: Vec::new(),
+        }
+    }
+}
+
+macro_rules! impl_init_constraints_group {
+    ($Constraint: ty, $Interaction: ty, $categorize: ident, $group: ident, $num_active_constraints: path, $empty_constraint: expr $(, $weight: ident)*) => {
+        impl ParallelVelocitySolverPart<$Constraint> {
+            pub fn init_constraints_groups(
+                &mut self,
+                island_id: usize,
+                bodies: &RigidBodySet,
+                interactions: &mut [$Interaction],
+                interaction_groups: &ParallelInteractionGroups,
+            ) {
+                let mut total_num_constraints = 0;
+                let num_groups = interaction_groups.num_groups();
+
+                self.interaction_groups.clear_groups();
+                self.ground_interaction_groups.clear_groups();
+                self.parallel_desc_groups.clear();
+                self.constraint_descs.clear();
+                self.parallel_desc_groups.push(0);
+
+                for i in 0..num_groups {
+                    let group = interaction_groups.group(i);
+
+                    self.not_ground_interactions.clear();
+                    self.ground_interactions.clear();
+                    $categorize(
+                        bodies,
+                        interactions,
+                        group,
+                        &mut self.ground_interactions,
+                        &mut self.not_ground_interactions,
+                    );
+
+                    #[cfg(feature = "simd-is-enabled")]
+                    let start_grouped = self.interaction_groups.grouped_interactions.len();
+                    let start_nongrouped = self.interaction_groups.nongrouped_interactions.len();
+                    #[cfg(feature = "simd-is-enabled")]
+                    let start_grouped_ground = self.ground_interaction_groups.grouped_interactions.len();
+                    let start_nongrouped_ground = self.ground_interaction_groups.nongrouped_interactions.len();
+
+                    self.interaction_groups.$group(
+                        island_id,
+                        bodies,
+                        interactions,
+                        &self.not_ground_interactions,
+                    );
+                    self.ground_interaction_groups.$group(
+                        island_id,
+                        bodies,
+                        interactions,
+                        &self.ground_interactions,
+                    );
+
+                    // Compute constraint indices.
+                    for interaction_i in &self.interaction_groups.nongrouped_interactions[start_nongrouped..] {
+                        let interaction = &mut interactions[*interaction_i]$(.$weight)*;
+                        interaction.constraint_index = total_num_constraints;
+                        self.constraint_descs.push((
+                            total_num_constraints,
+                            VelocityConstraintDesc::NongroundNongrouped(*interaction_i),
+                        ));
+                        total_num_constraints += $num_active_constraints(interaction);
+                    }
+
+                    #[cfg(feature = "simd-is-enabled")]
+                    for interaction_i in
+                        self.interaction_groups.grouped_interactions[start_grouped..].chunks(SIMD_WIDTH)
+                    {
+                        let interaction = &mut interactions[interaction_i[0]]$(.$weight)*;
+                        interaction.constraint_index = total_num_constraints;
+                        self.constraint_descs.push((
+                            total_num_constraints,
+                            VelocityConstraintDesc::NongroundGrouped(
+                                array![|ii| interaction_i[ii]; SIMD_WIDTH],
+                            ),
+                        ));
+                        total_num_constraints += $num_active_constraints(interaction);
+                    }
+
+                    for interaction_i in
+                        &self.ground_interaction_groups.nongrouped_interactions[start_nongrouped_ground..]
+                    {
+                        let interaction = &mut interactions[*interaction_i]$(.$weight)*;
+                        interaction.constraint_index = total_num_constraints;
+                        self.constraint_descs.push((
+                            total_num_constraints,
+                            VelocityConstraintDesc::GroundNongrouped(*interaction_i),
+                        ));
+                        total_num_constraints += $num_active_constraints(interaction);
+                    }
+
+                    #[cfg(feature = "simd-is-enabled")]
+                    for interaction_i in self.ground_interaction_groups.grouped_interactions
+                        [start_grouped_ground..]
+                        .chunks(SIMD_WIDTH)
+                    {
+                        let interaction = &mut interactions[interaction_i[0]]$(.$weight)*;
+                        interaction.constraint_index = total_num_constraints;
+                        self.constraint_descs.push((
+                            total_num_constraints,
+                            VelocityConstraintDesc::GroundGrouped(
+                                array![|ii| interaction_i[ii]; SIMD_WIDTH],
+                            ),
+                        ));
+                        total_num_constraints += $num_active_constraints(interaction);
+                    }
+
+                    self.parallel_desc_groups.push(self.constraint_descs.len());
+                }
+
+                // Resize the constraints set.
+                self.constraints.clear();
+                self.constraints
+                    .resize_with(total_num_constraints, || $empty_constraint)
+            }
+        }
+    }
+}
+
+impl_init_constraints_group!(
+    AnyVelocityConstraint,
+    &mut ContactManifold,
+    categorize_velocity_contacts,
+    group_manifolds,
+    VelocityConstraint::num_active_constraints,
+    AnyVelocityConstraint::Empty
+);
+
+impl_init_constraints_group!(
+    AnyJointVelocityConstraint,
+    JointGraphEdge,
+    categorize_joints,
+    group_joints,
+    AnyJointVelocityConstraint::num_active_constraints,
+    AnyJointVelocityConstraint::Empty,
+    weight
+);
+
+impl ParallelVelocitySolverPart<AnyVelocityConstraint> {
+    fn fill_constraints(
+        &mut self,
+        thread: &ThreadContext,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        manifolds_all: &[&mut ContactManifold],
+    ) {
+        let descs = &self.constraint_descs;
+
+        crate::concurrent_loop! {
+            let batch_size = thread.batch_size;
+            for desc in descs[thread.constraint_initialization_index, thread.num_initialized_constraints] {
+                match &desc.1 {
+                    VelocityConstraintDesc::NongroundNongrouped(manifold_id) => {
+                        let manifold = &*manifolds_all[*manifold_id];
+                        VelocityConstraint::generate(params, *manifold_id, manifold, bodies, &mut self.constraints, false);
+                    }
+                    VelocityConstraintDesc::GroundNongrouped(manifold_id) => {
+                        let manifold = &*manifolds_all[*manifold_id];
+                        VelocityGroundConstraint::generate(params, *manifold_id, manifold, bodies, &mut self.constraints, false);
+                    }
+                    #[cfg(feature = "simd-is-enabled")]
+                    VelocityConstraintDesc::NongroundGrouped(manifold_id) => {
+                        let manifolds = array![|ii| &*manifolds_all[manifold_id[ii]]; SIMD_WIDTH];
+                        WVelocityConstraint::generate(params, *manifold_id, manifolds, bodies, &mut self.constraints, false);
+                    }
+                    #[cfg(feature = "simd-is-enabled")]
+                    VelocityConstraintDesc::GroundGrouped(manifold_id) => {
+                        let manifolds = array![|ii| &*manifolds_all[manifold_id[ii]]; SIMD_WIDTH];
+                        WVelocityGroundConstraint::generate(params, *manifold_id, manifolds, bodies, &mut self.constraints, false);
+                    }
+                }
+            }
+        }
+    }
+}
+
+impl ParallelVelocitySolverPart<AnyJointVelocityConstraint> {
+    fn fill_constraints(
+        &mut self,
+        thread: &ThreadContext,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        joints_all: &[JointGraphEdge],
+    ) {
+        let descs = &self.constraint_descs;
+
+        crate::concurrent_loop! {
+            let batch_size = thread.batch_size;
+            for desc in descs[thread.joint_constraint_initialization_index, thread.num_initialized_joint_constraints] {
+                match &desc.1 {
+                    VelocityConstraintDesc::NongroundNongrouped(joint_id) => {
+                        let joint = &joints_all[*joint_id].weight;
+                        let constraint = AnyJointVelocityConstraint::from_joint(params, *joint_id, joint, bodies);
+                        self.constraints[joint.constraint_index] = constraint;
+                    }
+                    VelocityConstraintDesc::GroundNongrouped(joint_id) => {
+                        let joint = &joints_all[*joint_id].weight;
+                        let constraint = AnyJointVelocityConstraint::from_joint_ground(params, *joint_id, joint, bodies);
+                        self.constraints[joint.constraint_index] = constraint;
+                    }
+                    #[cfg(feature = "simd-is-enabled")]
+                    VelocityConstraintDesc::NongroundGrouped(joint_id) => {
+                        let joints = array![|ii| &joints_all[joint_id[ii]].weight; SIMD_WIDTH];
+                        let constraint = AnyJointVelocityConstraint::from_wide_joint(params, *joint_id, joints, bodies);
+                        self.constraints[joints[0].constraint_index] = constraint;
+                    }
+                    #[cfg(feature = "simd-is-enabled")]
+                    VelocityConstraintDesc::GroundGrouped(joint_id) => {
+                        let joints = array![|ii| &joints_all[joint_id[ii]].weight; SIMD_WIDTH];
+                        let constraint = AnyJointVelocityConstraint::from_wide_joint_ground(params, *joint_id, joints, bodies);
+                        self.constraints[joints[0].constraint_index] = constraint;
+                    }
+                }
+            }
+        }
+    }
+}
+
+pub(crate) struct ParallelVelocitySolver {
+    part: ParallelVelocitySolverPart<AnyVelocityConstraint>,
+    joint_part: ParallelVelocitySolverPart<AnyJointVelocityConstraint>,
+}
+
+impl ParallelVelocitySolver {
+    pub fn new() -> Self {
+        Self {
+            part: ParallelVelocitySolverPart::new(),
+            joint_part: ParallelVelocitySolverPart::new(),
+        }
+    }
+
+    pub fn init_constraint_groups(
+        &mut self,
+        island_id: usize,
+        bodies: &RigidBodySet,
+        manifolds: &mut [&mut ContactManifold],
+        manifold_groups: &ParallelInteractionGroups,
+        joints: &mut [JointGraphEdge],
+        joint_groups: &ParallelInteractionGroups,
+    ) {
+        self.part
+            .init_constraints_groups(island_id, bodies, manifolds, manifold_groups);
+        self.joint_part
+            .init_constraints_groups(island_id, bodies, joints, joint_groups);
+    }
+
+    pub fn fill_constraints(
+        &mut self,
+        thread: &ThreadContext,
+        params: &IntegrationParameters,
+        bodies: &RigidBodySet,
+        manifolds: &[&mut ContactManifold],
+        joints: &[JointGraphEdge],
+    ) {
+        self.part
+            .fill_constraints(thread, params, bodies, manifolds);
+        self.joint_part
+            .fill_constraints(thread, params, bodies, joints);
+        ThreadContext::lock_until_ge(
+            &thread.num_initialized_constraints,
+            self.part.constraint_descs.len(),
+        );
+        ThreadContext::lock_until_ge(
+            &thread.num_initialized_joint_constraints,
+            self.joint_part.constraint_descs.len(),
+        );
+    }
+
+    pub fn solve_constraints(
+        &mut self,
+        thread: &ThreadContext,
+        params: &IntegrationParameters,
+        manifolds_all: &mut [&mut ContactManifold],
+        joints_all: &mut [JointGraphEdge],
+        mj_lambdas: &mut [DeltaVel<f32>],
+    ) {
+        if self.part.constraint_descs.len() == 0 && self.joint_part.constraint_descs.len() == 0 {
+            return;
+        }
+
+        /*
+         * Warmstart constraints.
+         */
+        {
+            // Each thread will concurrently grab thread.batch_size constraint desc to
+            // solve. If the batch size is large enough for to cross the boundary of
+            // a parallel_desc_group, we have to wait util the current group is finished
+            // before starting the next one.
+            let mut target_num_desc = 0;
+            let mut start_index = thread
+                .warmstart_contact_index
+                .fetch_add(thread.batch_size, Ordering::SeqCst);
+            let mut batch_size = thread.batch_size;
+            let mut shift = 0;
+
+            macro_rules! warmstart(
+                ($part: expr) => {
+                    for group in $part.parallel_desc_groups.windows(2) {
+                        let num_descs_in_group = group[1] - group[0];
+                        target_num_desc += num_descs_in_group;
+
+                        while start_index < group[1] {
+                            let end_index = (start_index + batch_size).min(group[1]);
+
+                            let constraints = if end_index == $part.constraint_descs.len() {
+                                &mut $part.constraints[$part.constraint_descs[start_index].0..]
+                            } else {
+                                &mut $part.constraints[$part.constraint_descs[start_index].0..$part.constraint_descs[end_index].0]
+                            };
+
+                            for constraint in constraints {
+                                constraint.warmstart(mj_lambdas);
+                            }
+
+                            let num_solved = end_index - start_index;
+                            batch_size -= num_solved;
+
+                            thread
+                                .num_warmstarted_contacts
+                                .fetch_add(num_solved, Ordering::SeqCst);
+
+                            if batch_size == 0 {
+                                start_index = thread
+                                    .warmstart_contact_index
+                                    .fetch_add(thread.batch_size, Ordering::SeqCst);
+                                start_index -= shift;
+                                batch_size = thread.batch_size;
+                            } else {
+                                start_index += num_solved;
+                            }
+                        }
+
+                        ThreadContext::lock_until_ge(&thread.num_warmstarted_contacts, target_num_desc);
+                    }
+                }
+            );
+
+            warmstart!(self.joint_part);
+            shift = self.joint_part.constraint_descs.len();
+            start_index -= shift;
+            warmstart!(self.part);
+        }
+
+        /*
+         * Solve constraints.
+         */
+        {
+            // Each thread will concurrently grab thread.batch_size constraint desc to
+            // solve. If the batch size is large enough for to cross the boundary of
+            // a parallel_desc_group, we have to wait util the current group is finished
+            // before starting the next one.
+            let mut start_index = thread
+                .solve_interaction_index
+                .fetch_add(thread.batch_size, Ordering::SeqCst);
+            let mut batch_size = thread.batch_size;
+            let contact_descs = &self.part.constraint_descs[..];
+            let joint_descs = &self.joint_part.constraint_descs[..];
+            let mut target_num_desc = 0;
+            let mut shift = 0;
+
+            for _ in 0..params.max_velocity_iterations {
+                macro_rules! solve {
+                    ($part: expr) => {
+                        // Joint groups.
+                        for group in $part.parallel_desc_groups.windows(2) {
+                            let num_descs_in_group = group[1] - group[0];
+
+                            target_num_desc += num_descs_in_group;
+
+                            while start_index < group[1] {
+                                let end_index = (start_index + batch_size).min(group[1]);
+
+                                let constraints = if end_index == $part.constraint_descs.len() {
+                                    &mut $part.constraints[$part.constraint_descs[start_index].0..]
+                                } else {
+                                    &mut $part.constraints[$part.constraint_descs[start_index].0
+                                        ..$part.constraint_descs[end_index].0]
+                                };
+
+                                //                                println!(
+                                //                                    "Solving a constraint {:?}.",
+                                //                                    rayon::current_thread_index()
+                                //                                );
+                                for constraint in constraints {
+                                    constraint.solve(mj_lambdas);
+                                }
+
+                                let num_solved = end_index - start_index;
+                                batch_size -= num_solved;
+
+                                thread
+                                    .num_solved_interactions
+                                    .fetch_add(num_solved, Ordering::SeqCst);
+
+                                if batch_size == 0 {
+                                    start_index = thread
+                                        .solve_interaction_index
+                                        .fetch_add(thread.batch_size, Ordering::SeqCst);
+                                    start_index -= shift;
+                                    batch_size = thread.batch_size;
+                                } else {
+                                    start_index += num_solved;
+                                }
+                            }
+                            ThreadContext::lock_until_ge(
+                                &thread.num_solved_interactions,
+                                target_num_desc,
+                            );
+                        }
+                    };
+                }
+
+                solve!(self.joint_part);
+                shift += joint_descs.len();
+                start_index -= joint_descs.len();
+                solve!(self.part);
+                shift += contact_descs.len();
+                start_index -= contact_descs.len();
+            }
+        }
+
+        /*
+         * Writeback impulses.
+         */
+        let joint_constraints = &self.joint_part.constraints;
+        let contact_constraints = &self.part.constraints;
+        crate::concurrent_loop! {
+             let batch_size = thread.batch_size;
+             for constraint in joint_constraints[thread.joint_writeback_index] {
+                 constraint.writeback_impulses(joints_all);
+             }
+        }
+        crate::concurrent_loop! {
+             let batch_size = thread.batch_size;
+             for constraint in contact_constraints[thread.impulse_writeback_index] {
+                 constraint.writeback_impulses(manifolds_all);
+             }
+        }
+    }
+}