Refactor the parallel solver code the same way we did with the non-parallel solver.

This commit is contained in:
Crozet Sébastien
2021-01-04 17:59:51 +01:00
parent aa61fe65e3
commit c28b14d31c
14 changed files with 691 additions and 873 deletions

View File

@@ -1,325 +1,41 @@
use super::ParallelInteractionGroups;
use super::{AnyJointVelocityConstraint, AnyVelocityConstraint, DeltaVel, ThreadContext};
use crate::dynamics::solver::categorization::{categorize_contacts, categorize_joints};
use crate::dynamics::solver::{InteractionGroups, VelocityConstraint, VelocityGroundConstraint};
use crate::dynamics::solver::parallel_solver_constraints::ConstraintDesc;
use crate::dynamics::solver::{
AnyJointPositionConstraint, AnyPositionConstraint, InteractionGroups,
ParallelSolverConstraints, 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,
math::{Real, 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 ParallelSolverConstraints<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> ParallelSolverConstraints<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 ParallelSolverConstraints<$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_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 ParallelSolverConstraints<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 ParallelSolverConstraints<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].data.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].data.constraint_index] = constraint;
}
}
}
}
}
}
pub(crate) struct ParallelVelocitySolver {
part: ParallelSolverConstraints<AnyVelocityConstraint>,
joint_part: ParallelSolverConstraints<AnyJointVelocityConstraint>,
}
pub(crate) struct ParallelVelocitySolver {}
impl ParallelVelocitySolver {
pub fn new() -> Self {
Self {
part: ParallelSolverConstraints::new(),
joint_part: ParallelSolverConstraints::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,
pub fn solve(
thread: &ThreadContext,
params: &IntegrationParameters,
manifolds_all: &mut [&mut ContactManifold],
joints_all: &mut [JointGraphEdge],
mj_lambdas: &mut [DeltaVel<Real>],
contact_constraints: &mut ParallelSolverConstraints<
AnyVelocityConstraint,
AnyPositionConstraint,
>,
joint_constraints: &mut ParallelSolverConstraints<
AnyJointVelocityConstraint,
AnyJointPositionConstraint,
>,
) {
if self.part.constraint_descs.len() == 0 && self.joint_part.constraint_descs.len() == 0 {
if contact_constraints.constraint_descs.is_empty()
&& joint_constraints.constraint_descs.is_empty()
{
return;
}
@@ -348,9 +64,9 @@ impl ParallelVelocitySolver {
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..]
&mut $part.velocity_constraints[$part.constraint_descs[start_index].0..]
} else {
&mut $part.constraints[$part.constraint_descs[start_index].0..$part.constraint_descs[end_index].0]
&mut $part.velocity_constraints[$part.constraint_descs[start_index].0..$part.constraint_descs[end_index].0]
};
for constraint in constraints {
@@ -380,10 +96,10 @@ impl ParallelVelocitySolver {
}
);
warmstart!(self.joint_part);
shift = self.joint_part.constraint_descs.len();
warmstart!(joint_constraints);
shift = joint_constraints.constraint_descs.len();
start_index -= shift;
warmstart!(self.part);
warmstart!(contact_constraints);
}
/*
@@ -398,8 +114,8 @@ impl ParallelVelocitySolver {
.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 contact_descs = &contact_constraints.constraint_descs[..];
let joint_descs = &joint_constraints.constraint_descs[..];
let mut target_num_desc = 0;
let mut shift = 0;
@@ -416,9 +132,12 @@ impl ParallelVelocitySolver {
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..]
&mut $part.velocity_constraints
[$part.constraint_descs[start_index].0..]
} else {
&mut $part.constraints[$part.constraint_descs[start_index].0
&mut $part.velocity_constraints[$part.constraint_descs
[start_index]
.0
..$part.constraint_descs[end_index].0]
};
@@ -455,10 +174,10 @@ impl ParallelVelocitySolver {
};
}
solve!(self.joint_part);
solve!(joint_constraints);
shift += joint_descs.len();
start_index -= joint_descs.len();
solve!(self.part);
solve!(contact_constraints);
shift += contact_descs.len();
start_index -= contact_descs.len();
}
@@ -467,8 +186,9 @@ impl ParallelVelocitySolver {
/*
* Writeback impulses.
*/
let joint_constraints = &self.joint_part.constraints;
let contact_constraints = &self.part.constraints;
let joint_constraints = &joint_constraints.velocity_constraints;
let contact_constraints = &contact_constraints.velocity_constraints;
crate::concurrent_loop! {
let batch_size = thread.batch_size;
for constraint in joint_constraints[thread.joint_writeback_index] {