dynamics/
simd.rs

1//! Packing, unpacking, etc.
2
3use std::array::from_fn;
4
5use lin_alg::f32::{Vec3x8, Vec3x16, f32x8, f32x16};
6
7use crate::{AtomDynamicsx8, AtomDynamicsx16, MdState};
8
9impl MdState {
10    pub fn pack_atoms(&mut self) {
11        if is_x86_feature_detected!("avx512f") {
12            self.atoms_x16 = Vec::new();
13
14            for chunk in self.atoms.chunks(16) {
15                let serial_number = from_fn(|i| chunk.get(i).map(|a| a.serial_number).unwrap_or(0));
16                let static_ = from_fn(|i| chunk.get(i).map(|a| a.static_).unwrap_or(false));
17                let bonded_only = from_fn(|i| chunk.get(i).map(|a| a.bonded_only).unwrap_or(false));
18                let force_field_type = from_fn(|i| {
19                    chunk
20                        .get(i)
21                        .map(|a| a.force_field_type.clone())
22                        .unwrap_or_default()
23                });
24                let element = from_fn(|i| chunk.get(i).map(|a| a.element).unwrap_or_default());
25
26                let px = from_fn(|i| chunk.get(i).map(|a| a.posit.x).unwrap_or(0.0));
27                let py = from_fn(|i| chunk.get(i).map(|a| a.posit.y).unwrap_or(0.0));
28                let pz = from_fn(|i| chunk.get(i).map(|a| a.posit.z).unwrap_or(0.0));
29
30                let vx = from_fn(|i| chunk.get(i).map(|a| a.vel.x).unwrap_or(0.0));
31                let vy = from_fn(|i| chunk.get(i).map(|a| a.vel.y).unwrap_or(0.0));
32                let vz = from_fn(|i| chunk.get(i).map(|a| a.vel.z).unwrap_or(0.0));
33
34                let ax = from_fn(|i| chunk.get(i).map(|a| a.accel.x).unwrap_or(0.0));
35                let ay = from_fn(|i| chunk.get(i).map(|a| a.accel.y).unwrap_or(0.0));
36                let az = from_fn(|i| chunk.get(i).map(|a| a.accel.z).unwrap_or(0.0));
37
38                let mass = from_fn(|i| chunk.get(i).map(|a| a.mass).unwrap_or(0.0));
39                let partial_charge =
40                    from_fn(|i| chunk.get(i).map(|a| a.partial_charge).unwrap_or(0.0));
41                let lj_sigma = from_fn(|i| chunk.get(i).map(|a| a.lj_sigma).unwrap_or(0.0));
42                let lj_eps = from_fn(|i| chunk.get(i).map(|a| a.lj_eps).unwrap_or(0.0));
43
44                self.atoms_x16.push(AtomDynamicsx16 {
45                    serial_number,
46                    static_,
47                    bonded_only,
48                    force_field_type,
49                    element,
50                    posit: Vec3x16 {
51                        x: f32x16::from_array(px),
52                        y: f32x16::from_array(py),
53                        z: f32x16::from_array(pz),
54                    },
55                    vel: Vec3x16 {
56                        x: f32x16::from_array(vx),
57                        y: f32x16::from_array(vy),
58                        z: f32x16::from_array(vz),
59                    },
60                    accel: Vec3x16 {
61                        x: f32x16::from_array(ax),
62                        y: f32x16::from_array(ay),
63                        z: f32x16::from_array(az),
64                    },
65                    mass: f32x16::from_array(mass),
66                    partial_charge: f32x16::from_array(partial_charge),
67                    lj_sigma: f32x16::from_array(lj_sigma),
68                    lj_eps: f32x16::from_array(lj_eps),
69                });
70            }
71        } else {
72            self.atoms_x8 = Vec::new();
73
74            for chunk in self.atoms.chunks(8) {
75                let serial_number = from_fn(|i| chunk.get(i).map(|a| a.serial_number).unwrap_or(0));
76                let static_ = from_fn(|i| chunk.get(i).map(|a| a.static_).unwrap_or(false));
77                let bonded_only = from_fn(|i| chunk.get(i).map(|a| a.bonded_only).unwrap_or(false));
78                let force_field_type = from_fn(|i| {
79                    chunk
80                        .get(i)
81                        .map(|a| a.force_field_type.clone())
82                        .unwrap_or_default()
83                });
84                let element = from_fn(|i| chunk.get(i).map(|a| a.element).unwrap_or_default());
85
86                let px = from_fn(|i| chunk.get(i).map(|a| a.posit.x).unwrap_or(0.0));
87                let py = from_fn(|i| chunk.get(i).map(|a| a.posit.y).unwrap_or(0.0));
88                let pz = from_fn(|i| chunk.get(i).map(|a| a.posit.z).unwrap_or(0.0));
89
90                let vx = from_fn(|i| chunk.get(i).map(|a| a.vel.x).unwrap_or(0.0));
91                let vy = from_fn(|i| chunk.get(i).map(|a| a.vel.y).unwrap_or(0.0));
92                let vz = from_fn(|i| chunk.get(i).map(|a| a.vel.z).unwrap_or(0.0));
93
94                let ax = from_fn(|i| chunk.get(i).map(|a| a.accel.x).unwrap_or(0.0));
95                let ay = from_fn(|i| chunk.get(i).map(|a| a.accel.y).unwrap_or(0.0));
96                let az = from_fn(|i| chunk.get(i).map(|a| a.accel.z).unwrap_or(0.0));
97
98                let mass = from_fn(|i| chunk.get(i).map(|a| a.mass).unwrap_or(0.0));
99                let partial_charge =
100                    from_fn(|i| chunk.get(i).map(|a| a.partial_charge).unwrap_or(0.0));
101                let lj_sigma = from_fn(|i| chunk.get(i).map(|a| a.lj_sigma).unwrap_or(0.0));
102                let lj_eps = from_fn(|i| chunk.get(i).map(|a| a.lj_eps).unwrap_or(0.0));
103
104                self.atoms_x8.push(AtomDynamicsx8 {
105                    serial_number,
106                    static_,
107                    bonded_only,
108                    force_field_type,
109                    element,
110                    posit: Vec3x8 {
111                        x: f32x8::from_array(px),
112                        y: f32x8::from_array(py),
113                        z: f32x8::from_array(pz),
114                    },
115                    vel: Vec3x8 {
116                        x: f32x8::from_array(vx),
117                        y: f32x8::from_array(vy),
118                        z: f32x8::from_array(vz),
119                    },
120                    accel: Vec3x8 {
121                        x: f32x8::from_array(ax),
122                        y: f32x8::from_array(ay),
123                        z: f32x8::from_array(az),
124                    },
125                    mass: f32x8::from_array(mass),
126                    partial_charge: f32x8::from_array(partial_charge),
127                    lj_sigma: f32x8::from_array(lj_sigma),
128                    lj_eps: f32x8::from_array(lj_eps),
129                });
130            }
131        }
132    }
133}
dynamics/simd.rs

dynamics/
simd.rs
