path: root/src/dynamics/solver/parallel_island_solver.rs

use super::{DeltaVel, ParallelInteractionGroups, ParallelVelocitySolver};
use crate::dynamics::solver::{
    AnyJointPositionConstraint, AnyJointVelocityConstraint, AnyPositionConstraint,
    AnyVelocityConstraint, ParallelPositionSolver, ParallelSolverConstraints,
};
use crate::dynamics::{IntegrationParameters, JointGraphEdge, JointIndex, RigidBodySet};
use crate::geometry::{ContactManifold, ContactManifoldIndex};
use crate::math::{Isometry, Real};
use crate::utils::WAngularInertia;
use rayon::Scope;
use std::sync::atomic::{AtomicUsize, Ordering};

#[macro_export]
#[doc(hidden)]
macro_rules! concurrent_loop {
    (let batch_size = $batch_size: expr;
     for $elt: ident in $array: ident[$index_stream:expr,$index_count:expr] $f: expr) => {
        let max_index = $array.len();

        if max_index > 0 {
            loop {
                let start_index = $index_stream.fetch_add($batch_size, Ordering::SeqCst);
                if start_index > max_index {
                    break;
                }

                let end_index = (start_index + $batch_size).min(max_index);
                for $elt in &$array[start_index..end_index] {
                    $f
                }

                $index_count.fetch_add(end_index - start_index, Ordering::SeqCst);
            }
        }
    };

    (let batch_size = $batch_size: expr;
     for $elt: ident in $array: ident[$index_stream:expr] $f: expr) => {
        let max_index = $array.len();

        if max_index > 0 {
            loop {
                let start_index = $index_stream.fetch_add($batch_size, Ordering::SeqCst);
                if start_index > max_index {
                    break;
                }

                let end_index = (start_index + $batch_size).min(max_index);
                for $elt in &$array[start_index..end_index] {
                    $f
                }
            }
        }
    };
}

pub(crate) struct ThreadContext {
    pub batch_size: usize,
    // Velocity solver.
    pub constraint_initialization_index: AtomicUsize,
    pub num_initialized_constraints: AtomicUsize,
    pub joint_constraint_initialization_index: AtomicUsize,
    pub num_initialized_joint_constraints: AtomicUsize,
    pub warmstart_contact_index: AtomicUsize,
    pub num_warmstarted_contacts: AtomicUsize,
    pub warmstart_joint_index: AtomicUsize,
    pub num_warmstarted_joints: AtomicUsize,
    pub solve_interaction_index: AtomicUsize,
    pub num_solved_interactions: AtomicUsize,
    pub impulse_writeback_index: AtomicUsize,
    pub joint_writeback_index: AtomicUsize,
    pub body_integration_index: AtomicUsize,
    pub num_integrated_bodies: AtomicUsize,
    // Position solver.
    pub position_constraint_initialization_index: AtomicUsize,
    pub num_initialized_position_constraints: AtomicUsize,
    pub position_joint_constraint_initialization_index: AtomicUsize,
    pub num_initialized_position_joint_constraints: AtomicUsize,
    pub solve_position_interaction_index: AtomicUsize,
    pub num_solved_position_interactions: AtomicUsize,
    pub position_writeback_index: AtomicUsize,
}

impl ThreadContext {
    pub fn new(batch_size: usize) -> Self {
        ThreadContext {
            batch_size, // TODO perhaps there is some optimal value we can compute depending on the island size?
            constraint_initialization_index: AtomicUsize::new(0),
            num_initialized_constraints: AtomicUsize::new(0),
            joint_constraint_initialization_index: AtomicUsize::new(0),
            num_initialized_joint_constraints: AtomicUsize::new(0),
            num_warmstarted_contacts: AtomicUsize::new(0),
            warmstart_contact_index: AtomicUsize::new(0),
            num_warmstarted_joints: AtomicUsize::new(0),
            warmstart_joint_index: AtomicUsize::new(0),
            solve_interaction_index: AtomicUsize::new(0),
            num_solved_interactions: AtomicUsize::new(0),
            impulse_writeback_index: AtomicUsize::new(0),
            joint_writeback_index: AtomicUsize::new(0),
            body_integration_index: AtomicUsize::new(0),
            num_integrated_bodies: AtomicUsize::new(0),
            position_constraint_initialization_index: AtomicUsize::new(0),
            num_initialized_position_constraints: AtomicUsize::new(0),
            position_joint_constraint_initialization_index: AtomicUsize::new(0),
            num_initialized_position_joint_constraints: AtomicUsize::new(0),
            solve_position_interaction_index: AtomicUsize::new(0),
            num_solved_position_interactions: AtomicUsize::new(0),
            position_writeback_index: AtomicUsize::new(0),
        }
    }

    pub fn lock_until_ge(val: &AtomicUsize, target: usize) {
        if target > 0 {
            //        let backoff = crossbeam::utils::Backoff::new();
            std::sync::atomic::fence(Ordering::SeqCst);
            while val.load(Ordering::Relaxed) < target {
                //  backoff.spin();
                // std::thread::yield_now();
            }
        }
    }
}

pub struct ParallelIslandSolver {
    mj_lambdas: Vec<DeltaVel<Real>>,
    positions: Vec<Isometry<Real>>,
    parallel_groups: ParallelInteractionGroups,
    parallel_joint_groups: ParallelInteractionGroups,
    parallel_contact_constraints:
        ParallelSolverConstraints<AnyVelocityConstraint, AnyPositionConstraint>,
    parallel_joint_constraints:
        ParallelSolverConstraints<AnyJointVelocityConstraint, AnyJointPositionConstraint>,
    thread: ThreadContext,
}

impl ParallelIslandSolver {
    pub fn new() -> Self {
        Self {
            mj_lambdas: Vec::new(),
            positions: Vec::new(),
            parallel_groups: ParallelInteractionGroups::new(),
            parallel_joint_groups: ParallelInteractionGroups::new(),
            parallel_contact_constraints: ParallelSolverConstraints::new(),
            parallel_joint_constraints: ParallelSolverConstraints::new(),
            thread: ThreadContext::new(8),
        }
    }

    pub fn solve_island<'s>(
        &'s mut self,
        scope: &Scope<'s>,
        island_id: usize,
        params: &'s IntegrationParameters,
        bodies: &'s mut RigidBodySet,
        manifolds: &'s mut Vec<&'s mut ContactManifold>,
        manifold_indices: &'s [ContactManifoldIndex],
        joints: &'s mut Vec<JointGraphEdge>,
        joint_indices: &[JointIndex],
    ) {
        let num_threads = rayon::current_num_threads();
        let num_task_per_island = num_threads; // (num_threads / num_islands).max(1); // TODO: not sure this is the best value. Also, perhaps it is better to interleave tasks of each island?
        self.thread = ThreadContext::new(8); // TODO: could we compute some kind of optimal value here?
        self.parallel_groups
            .group_interactions(island_id, bodies, manifolds, manifold_indices);
        self.parallel_joint_groups
            .group_interactions(island_id, bodies, joints, joint_indices);
        self.parallel_contact_constraints.init_constraint_groups(
            island_id,
            bodies,
            manifolds,
            &self.parallel_groups,
        );
        self.parallel_joint_constraints.init_constraint_groups(
            island_id,
            bodies,
            joints,
            &self.parallel_joint_groups,
        );

        self.mj_lambdas.clear();
        self.mj_lambdas
            .resize(bodies.active_island(island_id).len(), DeltaVel::zero());
        self.positions.clear();
        self.positions
            .resize(bodies.active_island(island_id).len(), Isometry::identity());

        {
            // Initialize `mj_lambdas` (per-body velocity deltas) with external accelerations (gravity etc):

            let island_range = bodies.active_island_range(island_id);
            let active_bodies = &bodies.active_dynamic_set[island_range];
            let bodies = &mut bodies.bodies;

            let thread = &self.thread;

            concurrent_loop! {
                let batch_size = thread.batch_size;
                for handle in active_bodies[thread.body_integration_index, thread.num_integrated_bodies] {
                    let rb = &mut bodies[handle.0];
                    let dvel = &mut self.mj_lambdas[rb.active_set_offset];

                    dvel.linear += rb.force * (rb.effective_inv_mass * params.dt);
                    rb.force = na::zero();

                    // dvel.angular is actually storing angular velocity delta multiplied by the square root of the inertia tensor:
                    dvel.angular += rb.effective_world_inv_inertia_sqrt * rb.torque * params.dt;
                    rb.torque = na::zero();
                }
            }
        }

        for _ in 0..num_task_per_island {
            // We use AtomicPtr because it is Send+Sync while *mut is not.
            // See https://internals.rust-lang.org/t/shouldnt-pointers-be-send-sync-or/8818
            let thread = &self.thread;
            let