nexus_rbd2d 0.4.0

//! GPU-parallel constraint solver using graph coloring.
//!
//! Constraint-based physics solver running entirely on the GPU, using graph
//! coloring to solve constraints in parallel without data races. Uses the
//! `Soft-TGS` algorithm (as in Rapier).

use crate::dynamics::joint::{GpuJointSolver, JointSolverArgs};
#[cfg(feature = "dim3")]
use crate::dynamics::multibody::{GpuMultibodySet, GpuMultibodySolver, MultibodySolverArgs};
use crate::math::Pose;
use crate::queries::GpuIndexedContact;
use crate::shaders::dynamics::{
    GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized,
    GpuRemoveCfmAndBiasKernel, GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize,
    GpuSolverIncColor, GpuSolverInitConstraints, GpuSolverResetColor, GpuSolverSortConstraints,
    GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuWarmstart, LocalMassProperties,
    RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity, WorldMassProperties,
};
use crate::utils::{GpuPrefixSum, PrefixSumWorkspace};
use khal::Shader;
use khal::backend::{GpuBackend, GpuBackendError, GpuPass};
use vortx::tensor::Tensor;

/// GPU shader bundle for the constraint solver.
#[derive(Shader)]
pub struct GpuSolver {
    sort_constraints: GpuSolverSortConstraints,
    reset_color: GpuSolverResetColor,
    inc_color: GpuSolverIncColor,
    /// Initializes constraints from contact manifolds.
    init_constraints: GpuSolverInitConstraints,
    /// Companion counting pass to `init_constraints` (split out to keep each
    /// pass within the 8-storage-buffer WebGPU limit).
    count_constraints: GpuSolverCountConstraints,
    /// Updates nonlinear constraint terms during substeps.
    update_constraints: GpuSolverUpdateConstraints,
    /// Clears solver velocities and constraint counts.
    cleanup: GpuSolverCleanup,
    /// Applies warmstart impulses from previous frame.
    warmstart: GpuWarmstart,
    /// Gauss-Seidel iteration step (sequential per color).
    step_gauss_seidel: GpuStepGaussSeidel,
    /// Initializes solver velocity increments.
    init_solver_vels_inc: GpuInitSolverVelsInc,
    /// Seeds the COM-centered solver poses from the body world poses
    /// (rapier's `SolverBodies::copy_from`). Run once per step.
    init_solver_bodies: GpuInitSolverBodies,
    /// Applies accumulated solver velocity increments.
    apply_solver_vels_inc: GpuApplySolverVelsInc,
    /// Integrates positions from velocities.
    integrate_linearized: GpuIntegrateLinearized,
    /// Writes solver velocities and converts the COM-centered solver poses
    /// back to body-origin poses.
    finalize: GpuSolverFinalize,
    /// Removes CFM and bias terms for velocity-only solving.
    remove_cfm_and_bias_kernel: GpuRemoveCfmAndBiasKernel,
}

/// Arguments for constraint solver dispatch, used by [`GpuSolver::prepare`] and
/// [`GpuSolver::solve_tgs`].
pub struct SolverArgs<'a> {
    /// Total number of colors from graph coloring.
    pub num_colors: u32,
    /// Number of simulation batches.
    pub num_batches: u32,
    /// Number of colliders.
    pub num_colliders: u32,
    /// Contact manifolds generated by narrow-phase.
    pub contacts: &'a Tensor<GpuIndexedContact>,
    /// Number of contacts (per batch).
    pub contacts_len: &'a Tensor<u32>,
    /// Indirect dispatch arguments based on contact count.
    pub contacts_len_indirect: &'a Tensor<[u32; 3]>,
    /// Solver constraints (output from constraint initialization).
    pub constraints: &'a mut Tensor<TwoBodyConstraint>,
    /// Builder data for initializing constraints.
    pub constraint_builders: &'a mut Tensor<TwoBodyConstraintBuilder>,
    /// Global simulation parameters.
    pub sim_params: &'a Tensor<RbdSimParams>,
    /// Rigid body world-origin poses. Mirrors rapier's `RigidBody::position`.
    /// Read at the start of each step to seed [`Self::solver_body_poses`] and
    /// written back at the end of the substep loop by `finalize`.
    pub body_poses: &'a mut Tensor<Pose>,
    /// Rigid-body poses centered at the rigid-body center-of-mass — rapier's
    /// `SolverPose`. This is the only pose buffer the solver substep loop
    /// touches. Seeded from `body_poses` at step start and written back to
    /// `body_poses` at step end.
    pub solver_body_poses: &'a mut Tensor<Pose>,
    /// Per-collider local pose relative to the rigid-body it is attached to.
    pub collider_local_poses: &'a Tensor<Pose>,
    /// Per-collider world poses (= `body_poses[i] * collider_local_poses[i]`),
    /// kept up-to-date once per step before broad/narrow-phase.
    pub collider_world_poses: &'a Tensor<Pose>,
    /// Rigid body velocities.
    pub vels: &'a mut Tensor<Velocity>,
    /// Solver working velocities.
    pub solver_vels: &'a mut Tensor<Velocity>,
    /// Solver output velocities (currently unused).
    pub solver_vels_out: &'a Tensor<Velocity>,
    /// Accumulated velocity increments during substeps.
    pub solver_vels_inc: &'a mut Tensor<Velocity>,
    /// World-space mass properties.
    pub mprops: &'a Tensor<WorldMassProperties>,
    /// Local-space mass properties.
    pub local_mprops: &'a Tensor<LocalMassProperties>,
    /// Number of constraints per body.
    ///
    /// All constraints of all the bodies part of the same multibody are counted in a single
    /// entry at its root’s body index.
    pub body_constraint_counts: &'a mut Tensor<u32>,
    /// Constraint IDs associated with each body.
    ///
    /// All constraints of all the bodies part of the same multibody are in the same list associated
    /// to the multibody’s root.
    pub body_constraint_ids: &'a mut Tensor<u32>,
    /// Color assigned to each constraint by graph coloring.
    pub constraints_colors: &'a Tensor<u32>,
    /// Current color being processed.
    pub curr_color: &'a mut Tensor<u32>,
    /// Prefix sum shader for building constraint ranges.
    pub prefix_sum: &'a GpuPrefixSum,
    /// Number of solver iterations (max across all environments).
    pub num_solver_iterations: u32,
    /// Per-body graph-coloring group id (multibody-aware).
    pub body_group: &'a Tensor<u32>,
    /// Shared per-batch capacity / section-offset uniform — see
    /// [`crate::shaders::utils::BatchIndices`]. Consumed by the (refactored)
    /// multibody kernels via `MultibodySolverArgs::batch_indices`; the RBD
    /// constraint-solver kernels will migrate to it next.
    pub batch_indices: &'a Tensor<crate::shaders::utils::BatchIndices>,
}

impl GpuSolver {
    /// Prepares constraints for solving (init, prefix sum, sort).
    pub fn prepare<'a>(
        &self,
        backend: &GpuBackend,
        pass: &mut GpuPass,
        args: SolverArgs<'a>,
        prefix_sum_workspace: &'a mut PrefixSumWorkspace,
    ) -> Result<(), GpuBackendError> {
        // Cleanup zeroes body_constraint_counts, solver_vels, vels, mprops.
        self.cleanup.call(
            pass,
            [args.num_colliders, args.num_batches, 1],
            args.body_constraint_counts,
            args.solver_vels,
            args.vels,
            args.mprops,
            args.batch_indices,
        )?;

        // Seed `solver_body_poses` from `body_poses`: rapier's
        // `SolverBodies::copy_from`. After this, only the COM-centered solver
        // poses are touched until the final `finalize` writeback.
        // Independent of `cleanup` (different buffers), but `init_constraints`
        // below reads `solver_body_poses` so we barrier before it.
        self.init_solver_bodies.call(
            pass,
            [args.num_colliders, args.num_batches, 1],
            args.body_poses,
            args.local_mprops,
            args.solver_body_poses,
            args.batch_indices,
        )?;

        self.init_constraints.call(
            pass,
            args.contacts_len_indirect,
            args.contacts,
            args.constraints,
            args.constraint_builders,
            args.collider_world_poses,
            args.solver_body_poses,
            args.vels,
            args.mprops,
            args.sim_params,
            args.batch_indices,
        )?;

        // Counting runs as a separate dispatch (same indirect grid) so the
        // build pass above stays within 8 storage buffers.
        self.count_constraints.call(
            pass,
            args.contacts_len_indirect,
            args.contacts,
            args.body_constraint_counts,
            args.body_group,
            args.mprops,
            args.batch_indices,
        )?;

        args.prefix_sum.launch(
            backend,
            pass,
            prefix_sum_workspace,
            args.body_constraint_counts,
            args.num_batches,
        )?;

        self.sort_constraints.call(
            pass,
            args.contacts_len_indirect,
            args.body_constraint_counts,
            args.mprops,
            args.contacts,
            args.contacts_len,
            args.body_constraint_ids,
            args.body_group,
            args.batch_indices,
        )?;

        Ok(())
    }

    /// Solves constraints using the TGS (Total Gauss-Seidel) algorithm.
    ///
    /// When `multibody` is `Some`, multibody substep work is interleaved with
    /// the rigid-body substep work — one `multibody.apply_substep` call inside
    /// each iteration of the substep loop, just like rapier's
    /// `velocity_solver::solve_constraints`.
    pub fn solve_tgs<'a>(
        &self,
        pass: &mut GpuPass,
        joint_solver: &GpuJointSolver,
        args: SolverArgs<'a>,
        mut joint_args: JointSolverArgs<'a>,
        #[cfg(feature = "dim3")] multibody: Option<(&GpuMultibodySolver, &mut GpuMultibodySet)>,
    ) -> Result<(), GpuBackendError> {
        let num_substeps = args.num_solver_iterations;
        #[cfg(feature = "dim3")]
        let (mb_solver, mut mb_state) = match multibody {
            Some((s, st)) => (Some(s), Some(st)),
            None => (None, None),
        };

        /*
         * Init solver vel increments.
         */
        self.init_solver_vels_inc.call(
            pass,
            [args.num_colliders, args.num_batches, 1],
            args.solver_vels_inc,
            args.mprops,
            args.sim_params,
            args.batch_indices,
        )?;

        joint_solver.init(pass, &mut joint_args)?;

        // Per substep, the multibody work is split into five phases that are
        // INTERLEAVED with the matching rigid-body phases, mirroring rapier's
        // `velocity_solver::solve_constraints` order:
        //   P1/F1  integrate all velocities
        //   P2/F2  build + warmstart all constraints
        //   P3/F3  one PGS sweep over ALL joints + contacts WITH bias
        //   P4/F4  integrate ALL positions ONCE
        //   P5/F5  one PGS sweep over ALL joints + contacts WITHOUT bias
        // (Previously the entire multibody solve — including its own position
        // integration — ran as a monolithic block BEFORE the rigid-body solve,
        // which de-synchronised mb↔free contact coupling.)
        //
        // Each `mb_phase!($method $(, $extra)*)` invocation runs one multibody
        // phase. It is `#[cfg(feature = "dim3")]`-gated (the multibody solver
        // only exists in 3D) and reconstructs `mb_args` each time — the borrow
        // of `args.solver_vels` / `args.solver_body_poses` must be released
        // before the interleaved rigid-body call that touches the same buffers.
        macro_rules! mb_phase {
            ($method:ident $(, $extra:expr)*) => {{
                #[cfg(feature = "dim3")]
                if let (Some(solver), Some(state)) = (mb_solver, mb_state.as_deref_mut()) {
                    let mut mb_args = MultibodySolverArgs {
                        poses: &mut *args.solver_body_poses,
                        collider_world_poses: args.collider_world_poses,
                        mprops: args.mprops,
                        contacts: args.contacts,
                        contacts_len: args.contacts_len,
                        solver_vels: &mut *args.solver_vels,
                        batch_indices: args.batch_indices,
                    };
                    solver.$method(pass, state, &mut mb_args $(, $extra)*)?;
                }
            }};
        }

        for substep_id in 0..num_substeps {
            let is_last_substep = substep_id == num_substeps - 1;
            // Only consumed by the dim3-only multibody phases.
            #[cfg(not(feature = "dim3"))]
            let _ = is_last_substep;

            /*
             * P1/F1 — integrate velocities (apply `a · dt'` / gravity increment).
             */
            mb_phase!(substep_integrate_velocities);
            self.apply_solver_vels_inc.call(
                pass,
                [args.num_colliders, args.num_batches, 1],
                args.solver_vels,
                args.solver_vels_inc,
                args.batch_indices,
            )?;

            /*
             * P2/F2 — build + warmstart constraints.
             */
            mb_phase!(substep_build_constraints);
            self.update_constraints.call(
                pass,
                args.contacts_len_indirect,
                args.constraints,
                args.constraint_builders,
                args.contacts_len,
                args.solver_body_poses,
                args.sim_params,
                args.batch_indices,
            )?;
            joint_solver.update(pass, &mut joint_args, args.solver_body_poses)?;
            self.reset_color.call(pass, 1u32, args.curr_color)?;
            for _ in 0..args.num_colors {
                self.warmstart.call(
                    pass,
                    args.contacts_len_indirect,
                    args.constraints,
                    args.solver_vels,
                    args.constraints_colors,
                    args.contacts_len,
                    args.curr_color,
                    args.batch_indices,
                )?;
                self.inc_color.call(pass, 1u32, args.curr_color)?
            }

            /*
             * P3/F3 — solve ALL joints + contacts WITH bias.
             */
            mb_phase!(substep_solve_with_bias);
            joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?;
            self.reset_color.call(pass, 1u32, args.curr_color)?;
            for _ in 0..args.num_colors {
                self.step_gauss_seidel.call(
                    pass,
                    args.contacts_len_indirect,
                    args.constraints,
                    args.solver_vels,
                    args.constraints_colors,
                    args.contacts_len,
                    args.curr_color,
                    args.batch_indices,
                )?;
                self.inc_color.call(pass, 1u32, args.curr_color)?
            }

            /*
             * P4/F4 — integrate ALL positions once.
             */
            mb_phase!(substep_integrate_positions, is_last_substep);
            self.integrate_linearized.call(
                pass,
                [args.num_colliders, args.num_batches, 1],
                args.solver_body_poses,
                args.solver_vels,
                args.sim_params,
                args.batch_indices,
            )?;

            /*
             * P5/F5 — solve ALL joints + contacts WITHOUT bias (stabilization).
             */
            mb_phase!(substep_solve_no_bias);
            joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?;
            self.remove_cfm_and_bias_kernel.call(
                pass,
                args.contacts_len_indirect,
                args.constraints,
                args.contacts_len,
                args.batch_indices,
            )?;
            self.reset_color.call(pass, 1u32, args.curr_color)?;
            for _ in 0..args.num_colors {
                self.step_gauss_seidel.call(
                    pass,
                    args.contacts_len_indirect,
                    args.constraints,
                    args.solver_vels,
                    args.constraints_colors,
                    args.contacts_len,
                    args.curr_color,
                    args.batch_indices,
                )?;
                self.inc_color.call(pass, 1u32, args.curr_color)?
            }
        }

        /*
         * Writeback body velocities and convert COM-centered solver poses
         * back to body-origin poses.
         */
        self.finalize.call(
            pass,
            [args.num_colliders, args.num_batches, 1],
            args.vels,
            args.solver_vels,
            args.body_poses,
            args.solver_body_poses,
            args.local_mprops,
            args.batch_indices,
        )?;

        Ok(())
    }
}