Refactor the constraints solver code.

This commit is contained in:
Crozet Sébastien
2020-12-31 10:32:34 +01:00
parent 348a339fe3
commit 64507a68e1
6 changed files with 435 additions and 651 deletions

View File

@@ -5,8 +5,8 @@ use super::{
use super::{WVelocityConstraint, WVelocityGroundConstraint};
use crate::dynamics::solver::categorization::{categorize_contacts, categorize_joints};
use crate::dynamics::solver::{
AnyPositionConstraint, PositionConstraint, PositionGroundConstraint, WPositionConstraint,
WPositionGroundConstraint,
AnyJointPositionConstraint, AnyPositionConstraint, PositionConstraint,
PositionGroundConstraint, WPositionConstraint, WPositionGroundConstraint,
};
use crate::dynamics::{
solver::{AnyVelocityConstraint, DeltaVel},
@@ -19,54 +19,24 @@ use crate::utils::WAngularInertia;
pub(crate) struct VelocitySolver {
pub mj_lambdas: Vec<DeltaVel<f32>>,
pub contact_part: VelocitySolverPart<AnyVelocityConstraint>,
pub joint_part: VelocitySolverPart<AnyJointVelocityConstraint>,
}
impl VelocitySolver {
pub fn new() -> Self {
Self {
mj_lambdas: Vec::new(),
contact_part: VelocitySolverPart::new(),
joint_part: VelocitySolverPart::new(),
}
}
pub fn init_constraints(
&mut self,
island_id: usize,
params: &IntegrationParameters,
bodies: &RigidBodySet,
manifolds: &[&mut ContactManifold],
manifold_indices: &[ContactManifoldIndex],
joints: &[JointGraphEdge],
joint_constraint_indices: &[JointIndex],
position_constraints: &mut Vec<AnyPositionConstraint>,
) {
self.contact_part.init_constraints(
island_id,
params,
bodies,
manifolds,
manifold_indices,
position_constraints,
);
self.joint_part.init_constraints(
island_id,
params,
bodies,
joints,
joint_constraint_indices,
)
}
pub fn solve_constraints(
pub fn solve(
&mut self,
island_id: usize,
params: &IntegrationParameters,
bodies: &mut RigidBodySet,
manifolds_all: &mut [&mut ContactManifold],
joints_all: &mut [JointGraphEdge],
contact_constraints: &mut [AnyVelocityConstraint],
joint_constraints: &mut [AnyJointVelocityConstraint],
) {
self.mj_lambdas.clear();
self.mj_lambdas
@@ -75,11 +45,11 @@ impl VelocitySolver {
/*
* Warmstart constraints.
*/
for constraint in &self.joint_part.constraints {
for constraint in &*joint_constraints {
constraint.warmstart(&mut self.mj_lambdas[..]);
}
for constraint in &self.contact_part.constraints {
for constraint in &*contact_constraints {
constraint.warmstart(&mut self.mj_lambdas[..]);
}
@@ -87,11 +57,11 @@ impl VelocitySolver {
* Solve constraints.
*/
for _ in 0..params.max_velocity_iterations {
for constraint in &mut self.joint_part.constraints {
for constraint in &mut *joint_constraints {
constraint.solve(&mut self.mj_lambdas[..]);
}
for constraint in &mut self.contact_part.constraints {
for constraint in &mut *contact_constraints {
constraint.solve(&mut self.mj_lambdas[..]);
}
}
@@ -104,327 +74,12 @@ impl VelocitySolver {
});
// Write impulses back into the manifold structures.
for constraint in &self.joint_part.constraints {
for constraint in &*joint_constraints {
constraint.writeback_impulses(joints_all);
}
for constraint in &self.contact_part.constraints {
for constraint in &*contact_constraints {
constraint.writeback_impulses(manifolds_all);
}
}
}
pub(crate) struct VelocitySolverPart<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>,
}
impl<Constraint> VelocitySolverPart<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(),
}
}
}
impl VelocitySolverPart<AnyVelocityConstraint> {
pub fn init_constraint_groups(
&mut self,
island_id: usize,
bodies: &RigidBodySet,
manifolds: &[&mut ContactManifold],
manifold_indices: &[ContactManifoldIndex],
) {
self.not_ground_interactions.clear();
self.ground_interactions.clear();
categorize_contacts(
bodies,
manifolds,
manifold_indices,
&mut self.ground_interactions,
&mut self.not_ground_interactions,
);
self.interaction_groups.clear_groups();
self.interaction_groups.group_manifolds(
island_id,
bodies,
manifolds,
&self.not_ground_interactions,
);
self.ground_interaction_groups.clear_groups();
self.ground_interaction_groups.group_manifolds(
island_id,
bodies,
manifolds,
&self.ground_interactions,
);
// NOTE: uncomment this do disable SIMD contact resolution.
// self.interaction_groups
// .nongrouped_interactions
// .append(&mut self.interaction_groups.grouped_interactions);
// self.ground_interaction_groups
// .nongrouped_interactions
// .append(&mut self.ground_interaction_groups.grouped_interactions);
}
pub fn init_constraints(
&mut self,
island_id: usize,
params: &IntegrationParameters,
bodies: &RigidBodySet,
manifolds: &[&mut ContactManifold],
manifold_indices: &[ContactManifoldIndex],
position_constraints: &mut Vec<AnyPositionConstraint>,
) {
self.init_constraint_groups(island_id, bodies, manifolds, manifold_indices);
self.constraints.clear();
#[cfg(feature = "simd-is-enabled")]
{
self.compute_grouped_constraints(params, bodies, manifolds, position_constraints);
}
self.compute_nongrouped_constraints(params, bodies, manifolds, position_constraints);
#[cfg(feature = "simd-is-enabled")]
{
self.compute_grouped_ground_constraints(
params,
bodies,
manifolds,
position_constraints,
);
}
self.compute_nongrouped_ground_constraints(params, bodies, manifolds, position_constraints);
}
#[cfg(feature = "simd-is-enabled")]
fn compute_grouped_constraints(
&mut self,
params: &IntegrationParameters,
bodies: &RigidBodySet,
manifolds_all: &[&mut ContactManifold],
constraints_all: &mut Vec<AnyPositionConstraint>,
) {
for manifolds_i in self
.interaction_groups
.grouped_interactions
.chunks_exact(SIMD_WIDTH)
{
let manifold_id = array![|ii| manifolds_i[ii]; SIMD_WIDTH];
let manifolds = array![|ii| &*manifolds_all[manifolds_i[ii]]; SIMD_WIDTH];
WVelocityConstraint::generate(
params,
manifold_id,
manifolds,
bodies,
&mut self.constraints,
true,
);
WPositionConstraint::generate(params, manifolds, bodies, constraints_all, true);
}
}
fn compute_nongrouped_constraints(
&mut self,
params: &IntegrationParameters,
bodies: &RigidBodySet,
manifolds_all: &[&mut ContactManifold],
constraints_all: &mut Vec<AnyPositionConstraint>,
) {
for manifold_i in &self.interaction_groups.nongrouped_interactions {
let manifold = &manifolds_all[*manifold_i];
VelocityConstraint::generate(
params,
*manifold_i,
manifold,
bodies,
&mut self.constraints,
true,
);
PositionConstraint::generate(params, manifold, bodies, constraints_all, true);
}
}
#[cfg(feature = "simd-is-enabled")]
fn compute_grouped_ground_constraints(
&mut self,
params: &IntegrationParameters,
bodies: &RigidBodySet,
manifolds_all: &[&mut ContactManifold],
constraints_all: &mut Vec<AnyPositionConstraint>,
) {
for manifolds_i in self
.ground_interaction_groups
.grouped_interactions
.chunks_exact(SIMD_WIDTH)
{
let manifold_id = array![|ii| manifolds_i[ii]; SIMD_WIDTH];
let manifolds = array![|ii| &*manifolds_all[manifolds_i[ii]]; SIMD_WIDTH];
WVelocityGroundConstraint::generate(
params,
manifold_id,
manifolds,
bodies,
&mut self.constraints,
true,
);
WPositionGroundConstraint::generate(params, manifolds, bodies, constraints_all, true);
}
}
fn compute_nongrouped_ground_constraints(
&mut self,
params: &IntegrationParameters,
bodies: &RigidBodySet,
manifolds_all: &[&mut ContactManifold],
constraints_all: &mut Vec<AnyPositionConstraint>,
) {
for manifold_i in &self.ground_interaction_groups.nongrouped_interactions {
let manifold = &manifolds_all[*manifold_i];
VelocityGroundConstraint::generate(
params,
*manifold_i,
manifold,
bodies,
&mut self.constraints,
true,
);
PositionGroundConstraint::generate(params, manifold, bodies, constraints_all, true)
}
}
}
impl VelocitySolverPart<AnyJointVelocityConstraint> {
pub fn init_constraints(
&mut self,
island_id: usize,
params: &IntegrationParameters,
bodies: &RigidBodySet,
joints: &[JointGraphEdge],
joint_constraint_indices: &[JointIndex],
) {
// Generate constraints for joints.
self.not_ground_interactions.clear();
self.ground_interactions.clear();
categorize_joints(
bodies,
joints,
joint_constraint_indices,
&mut self.ground_interactions,
&mut self.not_ground_interactions,
);
self.constraints.clear();
self.interaction_groups.clear_groups();
self.interaction_groups.group_joints(
island_id,
bodies,
joints,
&self.not_ground_interactions,
);
self.ground_interaction_groups.clear_groups();
self.ground_interaction_groups.group_joints(
island_id,
bodies,
joints,
&self.ground_interactions,
);
// NOTE: uncomment this do disable SIMD joint resolution.
// self.interaction_groups
// .nongrouped_interactions
// .append(&mut self.interaction_groups.grouped_interactions);
// self.ground_interaction_groups
// .nongrouped_interactions
// .append(&mut self.ground_interaction_groups.grouped_interactions);
self.compute_nongrouped_joint_ground_constraints(params, bodies, joints);
#[cfg(feature = "simd-is-enabled")]
{
self.compute_grouped_joint_ground_constraints(params, bodies, joints);
}
self.compute_nongrouped_joint_constraints(params, bodies, joints);
#[cfg(feature = "simd-is-enabled")]
{
self.compute_grouped_joint_constraints(params, bodies, joints);
}
}
fn compute_nongrouped_joint_ground_constraints(
&mut self,
params: &IntegrationParameters,
bodies: &RigidBodySet,
joints_all: &[JointGraphEdge],
) {
for joint_i in &self.ground_interaction_groups.nongrouped_interactions {
let joint = &joints_all[*joint_i].weight;
let vel_constraint =
AnyJointVelocityConstraint::from_joint_ground(params, *joint_i, joint, bodies);
self.constraints.push(vel_constraint);
}
}
#[cfg(feature = "simd-is-enabled")]
fn compute_grouped_joint_ground_constraints(
&mut self,
params: &IntegrationParameters,
bodies: &RigidBodySet,
joints_all: &[JointGraphEdge],
) {
for joints_i in self
.ground_interaction_groups
.grouped_interactions
.chunks_exact(SIMD_WIDTH)
{
let joints_id = array![|ii| joints_i[ii]; SIMD_WIDTH];
let joints = array![|ii| &joints_all[joints_i[ii]].weight; SIMD_WIDTH];
let vel_constraint = AnyJointVelocityConstraint::from_wide_joint_ground(
params, joints_id, joints, bodies,
);
self.constraints.push(vel_constraint);
}
}
fn compute_nongrouped_joint_constraints(
&mut self,
params: &IntegrationParameters,
bodies: &RigidBodySet,
joints_all: &[JointGraphEdge],
) {
for joint_i in &self.interaction_groups.nongrouped_interactions {
let joint = &joints_all[*joint_i].weight;
let vel_constraint =
AnyJointVelocityConstraint::from_joint(params, *joint_i, joint, bodies);
self.constraints.push(vel_constraint);
}
}
#[cfg(feature = "simd-is-enabled")]
fn compute_grouped_joint_constraints(
&mut self,
params: &IntegrationParameters,
bodies: &RigidBodySet,
joints_all: &[JointGraphEdge],
) {
for joints_i in self
.interaction_groups
.grouped_interactions
.chunks_exact(SIMD_WIDTH)
{
let joints_id = array![|ii| joints_i[ii]; SIMD_WIDTH];
let joints = array![|ii| &joints_all[joints_i[ii]].weight; SIMD_WIDTH];
let vel_constraint =
AnyJointVelocityConstraint::from_wide_joint(params, joints_id, joints, bodies);
self.constraints.push(vel_constraint);
}
}
}