slosh3d 0.1.0

Cross-platform GPU 3D Material Point Method implementation.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use crate::grid::grid::{GpuGrid, GpuGridMetadata};
use crate::solver::params::GpuSimulationParams;
use crate::solver::SimulationParams;
use encase::ShaderType;
use nexus::dynamics::{GpuBodySet, GpuMassProperties, GpuVelocity};
use nexus::math::GpuSim;
use rapier::math::{AngVector, Point, Vector};
use slang_hal::backend::Backend;
use slang_hal::function::GpuFunction;
use slang_hal::{Shader, ShaderArgs};
use stensor::tensor::{GpuScalar, GpuVector};
use wgpu::BufferUsages;

#[derive(Shader)]
#[shader(module = "slosh::solver::rigid_impulses")]
pub struct WgRigidImpulses<B: Backend> {
    pub update: GpuFunction<B>,
    pub update_world_mass_properties: GpuFunction<B>,
}

#[derive(Copy, Clone, PartialEq, Debug, Default, ShaderType)]
#[repr(C)]
pub struct RigidImpulse {
    pub com: Point<f32>, // For convenience, to reduce the number of bindings
    pub linear: Vector<f32>,
    pub angular: AngVector<f32>,
}

pub struct GpuImpulses<B: Backend> {
    pub incremental_impulses: GpuVector<RigidImpulse, B>,
    pub total_impulses: GpuVector<RigidImpulse, B>,
    pub total_impulses_staging: GpuVector<RigidImpulse, B>,
}

impl<B: Backend> GpuImpulses<B> {
    pub fn new(backend: &B) -> Result<Self, B::Error> {
        const MAX_BODY_COUNT: usize = 16; // CPIC doesnt support more.
        let impulses = [RigidImpulse::default(); MAX_BODY_COUNT];
        Ok(Self {
            incremental_impulses: GpuVector::vector_encased(
                backend,
                impulses,
                BufferUsages::STORAGE,
            )?,
            total_impulses: GpuVector::vector_encased(
                backend,
                impulses,
                BufferUsages::STORAGE | BufferUsages::COPY_SRC,
            )?,
            total_impulses_staging: GpuVector::vector_encased(
                backend,
                impulses,
                BufferUsages::MAP_READ | BufferUsages::COPY_DST,
            )?,
        })
    }
}

#[derive(ShaderArgs)]
struct RigidImpulsesArgs<'a, B: Backend> {
    params: Option<&'a GpuScalar<SimulationParams, B>>,
    grid: Option<&'a GpuVector<GpuGridMetadata, B>>,
    local_mprops: &'a GpuVector<GpuMassProperties, B>,
    poses: &'a GpuVector<GpuSim, B>,
    vels: &'a GpuVector<GpuVelocity, B>,
    mprops: &'a GpuVector<GpuMassProperties, B>,
    incremental_impulses: &'a GpuVector<RigidImpulse, B>,
}

impl<B: Backend> WgRigidImpulses<B> {
    pub fn launch(
        &self,
        backend: &B,
        pass: &mut B::Pass,
        grid: &GpuGrid<B>,
        sim_params: &GpuSimulationParams<B>,
        impulses: &GpuImpulses<B>,
        bodies: &GpuBodySet<B>,
    ) -> Result<(), B::Error> {
        let args = RigidImpulsesArgs {
            params: Some(&sim_params.params),
            grid: Some(&grid.meta),
            local_mprops: bodies.local_mprops(),
            poses: bodies.poses(),
            vels: bodies.vels(),
            mprops: bodies.mprops(),
            incremental_impulses: &impulses.incremental_impulses,
        };
        self.update.launch_grid(backend, pass, &args, 1)
    }

    pub fn launch_update_world_mass_properties(
        &self,
        backend: &B,
        pass: &mut B::Pass,
        impulses: &GpuImpulses<B>,
        bodies: &GpuBodySet<B>,
    ) -> Result<(), B::Error> {
        let args = RigidImpulsesArgs {
            params: None,
            grid: None,
            local_mprops: bodies.local_mprops(),
            poses: bodies.poses(),
            vels: bodies.vels(),
            mprops: bodies.mprops(),
            incremental_impulses: &impulses.incremental_impulses,
        };

        self.update_world_mass_properties
            .launch_grid(backend, pass, &args, 1)
    }
}