slosh2d 0.4.1

Cross-platform GPU 2D 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

use crate::grid::grid::{GpuGrid, GpuGridMetadata};
use crate::math::Matrix;
use crate::solver::{GpuParticleModelData, GpuParticles, Kinematics, ParticleProperties};
use slang_hal::backend::{Backend, Encoder};
use slang_hal::function::GpuFunction;
use slang_hal::{Shader, ShaderArgs};
use stensor::tensor::{GpuScalar, GpuTensor};

#[derive(Copy, Clone, PartialEq, Debug, Default, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
/// The GPU representation of a maximum timestep estimate.
pub struct GpuTimestepBounds {
    compute_max_dt_as_uint: u32,
}

impl GpuTimestepBounds {
    // NOTE: this **MUST** match the constant in the GPU slang shader.
    const FLOAT_TO_INT: f32 = 1.0e12;

    /// Initializes the timestep bound (defaults to 0).
    pub fn new() -> GpuTimestepBounds {
        Self::default()
    }

    /// The time estimate, in seconds.
    pub fn computed_dt(&self) -> f32 {
        self.compute_max_dt_as_uint as f32 / Self::FLOAT_TO_INT
    }
}

#[derive(Shader)]
#[shader(
    module = "slosh::solver::timestep_bound",
    specialize = ["slosh::models::specializations"]
)]
/// GPU kernel responsible for computing an estimate on the maximum timestep that can
/// be taken by the simulation without blowing up.
///
/// Note that this a best-effort estimate. It does not completely eliminate risks of divergence.
pub struct WgTimestepBounds<B: Backend> {
    reset_timestep_bound: GpuFunction<B>,
    estimate_timestep_bound: GpuFunction<B>,
}

#[derive(ShaderArgs)]
struct TimestepBoundsArgs<'a, B: Backend, GpuModel: GpuParticleModelData> {
    grid: &'a GpuTensor<GpuGridMetadata, B>,
    particles_model: &'a GpuTensor<GpuModel, B>,
    particles_kin: &'a GpuTensor<Kinematics, B>,
    particles_def_grad: &'a GpuTensor<Matrix<f32>, B>,
    particles_props: &'a GpuTensor<ParticleProperties, B>,
    particles_len: &'a GpuScalar<u32, B>,
    result: &'a GpuScalar<GpuTimestepBounds, B>,
}

impl<B: Backend> WgTimestepBounds<B> {
    /// Launches the timestep bounds estimation, and return the estimated maximum timestep length.
    pub async fn compute_bounds<GpuModel: GpuParticleModelData>(
        &self,
        backend: &B,
        grid: &GpuGrid<B>,
        particles: &GpuParticles<B, GpuModel>,
        bounds: &GpuScalar<GpuTimestepBounds, B>,
        bounds_staging: &mut GpuScalar<GpuTimestepBounds, B>,
    ) -> Result<f32, B::Error> {
        let mut encoder = backend.begin_encoding();
        let mut pass = encoder.begin_pass("timestep_bounds", None);
        self.launch(backend, &mut pass, grid, particles, bounds)?;
        drop(pass);
        bounds_staging.copy_from_view(&mut encoder, bounds)?;
        backend.submit(encoder)?;

        let mut result = [GpuTimestepBounds::new()];
        backend
            .read_buffer(bounds_staging.buffer(), &mut result)
            .await?;
        Ok(result[0].computed_dt())
    }

    fn launch<GpuModel: GpuParticleModelData>(
        &self,
        backend: &B,
        pass: &mut B::Pass,
        grid: &GpuGrid<B>,
        particles: &GpuParticles<B, GpuModel>,
        bounds: &GpuScalar<GpuTimestepBounds, B>,
    ) -> Result<(), B::Error> {
        let args = TimestepBoundsArgs {
            grid: &grid.meta,
            particles_model: particles.models(),
            particles_kin: &particles.kinematics,
            particles_def_grad: &particles.def_grad,
            particles_props: &particles.properties,
            particles_len: particles.gpu_len(),
            result: bounds,
        };
        self.reset_timestep_bound
            .launch(backend, pass, &args, [1; 3])?;
        self.estimate_timestep_bound
            .launch(backend, pass, &args, [particles.len() as u32, 1, 1])
    }
}