First public release of Rapier.
This commit is contained in:
485
src/dynamics/solver/parallel_velocity_solver.rs
Normal file
485
src/dynamics/solver/parallel_velocity_solver.rs
Normal file
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user