fips_md/runtime/
particle_data.rs

1//! Structs for managing runtime data for particle types
2
3use anyhow::{Result, anyhow};
4use aligned_box::AlignedBox;
5
6use core::panic;
7use std::{collections::HashMap, sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}, unimplemented, usize};
8
9use crate::{runtime::{ParticleID, MemberID, MemberIndexEntry, ParticleIndexEntry}, utils::FipsValue};
10
11// For now we just enforce 32 byte alignment (enough for AVX2 and older)
12// TODO: Better allocation management once custom allocators become stable
13const ALIGNMENT: usize = 32;
14
15/// Options for allocating data for constant members
16#[derive(Clone)]
17pub struct UniformMembers {
18    single_allocations: HashMap<String, FipsValue>
19    // /// Only allocate a single instance for all particles of the same kind
20    // Single(HashMap<String, parser::SubstitutionValue>),
21    // /// Allocate a separate instance for every particle
22    // PerParticle
23}
24
25impl UniformMembers {
26    pub fn new() -> Self {
27        Self {
28            single_allocations: HashMap::new()
29        }
30    }
31
32    pub fn set_uniform_member<S: Into<String>, T: Into<FipsValue>>(&mut self, name: S, value: T) {
33        self.single_allocations.insert(name.into(), value.into());
34    }
35
36    fn get_member(&self, name: &str) -> Option<&FipsValue> {
37        self.single_allocations.get(name)
38    }
39}
40
41/// Store (and thus root owner) for all particle data used in a simulation
42pub struct ParticleStore {
43    /// All stored particle data
44    particles: HashMap<ParticleID, RwLock<ParticleData>>
45}
46
47impl ParticleStore {
48    /// Create new particle store
49    pub(crate) fn new() -> Self {
50        Self {
51            particles: HashMap::new()
52        }
53    }
54
55    /// Allocate data for a given number of particles
56    ///
57    /// Note that you have to define a default method for allocating constant members.
58    /// If mixed allocations of constants are desired, manually reallocate the
59    /// corresponding members after this call.
60    pub(crate) fn create_particles<'a>(&'a mut self, particle: (ParticleID, &ParticleIndexEntry),
61        count: usize, uniform_members: &UniformMembers) -> Result<RwLockWriteGuard<'a, ParticleData>>
62    {
63        let (particle_id, particle_definition) = particle;
64        // Check if particle is already allocated
65        if self.particles.contains_key(&particle_id) {
66            return Err(anyhow!("Particle with name {} has already been allocated",
67                particle_definition.get_name()));
68        };
69        // Create new particle allocation
70        let particle_data = ParticleData::new(particle_definition, count, uniform_members)?;
71        self.particles.insert(particle_id, RwLock::new(particle_data));
72        Ok(self.particles.get(&particle_id).unwrap().write().unwrap())
73    }
74
75    /// Get the (approximate) memory used for particle data in bytes
76    pub(crate) fn get_memory_usage(&self) -> usize {
77        let mut memory = 0;
78        for (_, particle_data) in &self.particles {
79            memory += particle_data.read().unwrap().get_memory_usage();
80        }
81        memory
82    }
83
84    /// Return particles for which data is stored as key-value pairs
85    pub(crate) fn get_particles(&self) -> impl Iterator<Item = (ParticleID, &RwLock<ParticleData>)> {
86        self.particles.iter().map(|(k,v)| (*k,v))
87    }
88
89    /// Get particle data for a given particle ID (automatically borrows)
90    pub(crate) fn get_particle<'a>(&self, particle_id: ParticleID) -> Option<RwLockReadGuard<ParticleData>> {
91        self.particles.get(&particle_id).map(|x| x.read().unwrap())
92    }
93
94    /// Get mutable particle data for a given particle ID (automatically borrows)
95    pub(crate) fn get_particle_mut<'a>(&self, particle_id: ParticleID) -> Option<RwLockWriteGuard<ParticleData>> {
96        self.particles.get(&particle_id).map(|x| x.write().unwrap())
97    }
98}
99
100/// Data for all instances of a single particle type
101pub struct ParticleData {
102    /// Number of particles
103    count: usize,
104    // /// Array of particle positions
105    // positions: Pin<AlignedBox<[f64]>>,
106    /// Array of byte arrays containing other per-particle member data
107    members: HashMap<MemberID, RwLock<MemberData>>
108}
109
110impl ParticleData {
111    /// Create new allocations for a single particle type
112    pub(crate) fn new(particle: &ParticleIndexEntry, count: usize, uniform_members: &UniformMembers) -> Result<Self>
113    {
114        let mut members = HashMap::new();
115        for (member_id, member) in particle.get_members() {
116            members.insert(member_id, RwLock::new(MemberData::new(member, count, uniform_members)?));
117        }
118        Ok(Self {
119            count,
120            // positions: Pin::new(AlignedBox::slice_from_value(ALIGNMENT, count*domain.get_dim(), 0.0)
121            //     .map_err(|e| anyhow!("Cannot allocate particle positions: {}", e))?),
122            members 
123        })
124    }
125
126    /// Get number of particles
127    pub(crate) fn get_particle_count(&self) -> usize {
128        self.count
129    }
130
131    /// Get memory used for the data of this particle type
132    pub(crate) fn get_memory_usage(&self) -> usize {
133        let mut memory = 0;
134        //memory += self.positions.len() * std::mem::size_of::<f64>();
135        for (_, member_data) in &self.members {
136            memory += member_data.read().unwrap().get_memory_usage();
137        }
138        memory
139    }
140
141    pub(crate) fn borrow_member(&self, member_id: &MemberID) -> Option<RwLockReadGuard<MemberData>> {
142        match self.members.get(member_id) {
143            None => None,
144            Some(member) => {
145                Some(member.read().unwrap())
146            }
147        }
148    }
149
150    pub(crate) fn borrow_member_mut(&self, member_id: &MemberID) -> Option<RwLockWriteGuard<MemberData>> {
151        match self.members.get(member_id) {
152            None => None,
153            Some(member) => {
154                Some(member.write().unwrap())
155            }
156        }
157    }
158}
159
160pub enum MemberData {
161    /// One value for every particle
162    PerParticle {
163        /// Data as byte array
164        data: AlignedBox<[u8]>,
165        /// Flag for constant members
166        mutable: bool
167    },
168    /// The same value for all particles (this must be constant)
169    Uniform(FipsValue)
170
171    // /// Mutable particle data (definitely per particle)
172    // Mutable {
173    //     data: AlignedBox<[u8]>
174    // },
175    // /// Constant particle data (maybe per particle, but maybe constant for all particles)
176    // Constant(ConstMemberData)
177}
178
179impl MemberData {
180    /// Create a new member data storage
181    pub(crate) fn new(member: &MemberIndexEntry, count: usize, uniform_members: &UniformMembers) -> Result<Self> {
182        // Is member in 
183        match uniform_members.get_member(member.get_name()) {
184            Some(value) => {
185                Ok(Self::Uniform(value.clone()))
186            }
187            None => {
188                let data = AlignedBox::slice_from_value(ALIGNMENT, count*member.get_member_size()?, 0u8)
189                    .map_err(|e| anyhow!("Cannot allocate particle member data: {}", e))?;
190                let mutable = member.is_mutable();
191                Ok(Self::PerParticle {
192                    data, mutable
193                })
194            }
195        }
196    }
197
198    /// Get allocated memory size in bytes
199    pub(crate) fn get_memory_usage(&self) -> usize {
200        match self {
201            MemberData::PerParticle { data, .. } => { data.len() }
202            MemberData::Uniform(value) => { value.get_size() }
203        }
204    }
205
206    // /// Returns true if allocations are made per-particle
207    // pub(crate) fn is_per_particle(&self) -> bool {
208    //     match self {
209    //         MemberData::PerParticle{..} => true,
210    //         MemberData::Uniform(_) => false
211    //     }
212    // }
213
214    /// Returns true if this member is allocated uniformly
215    /// (i.e. only one instance for all particles of this type)
216    pub(crate) fn is_uniform(&self) -> bool {
217        match self {
218            MemberData::PerParticle{..} => false,
219            MemberData::Uniform(_) => true
220        }
221    }
222
223    pub(crate) fn as_i64(&self) -> i64 {
224        match self {
225            Self::PerParticle {data, ..} => *bytemuck::from_bytes::<i64>(data),
226            Self::Uniform(value) => match value {
227                FipsValue::Int64(value) => *value,
228                _ => panic!("Cannot access member data as i64 (type is {})", value.get_type()) 
229            }
230        }
231    }
232
233    pub(crate) fn as_f64(&self) -> f64 {
234        match self {
235            Self::PerParticle {data, ..} => *bytemuck::from_bytes::<f64>(data),
236            Self::Uniform(value) => match value {
237                FipsValue::Double(value) => *value,
238                _ => panic!("Cannot access member data as f64 (type is {})", value.get_type()) 
239            }
240        }
241    }
242
243    pub(crate) fn as_i64_slice(&self) -> &[i64] {
244        match self {
245            Self::PerParticle {data, ..} => bytemuck::cast_slice::<u8,i64>(data),
246            Self::Uniform(_) => unimplemented!()
247        }
248    }
249    
250    pub(crate) fn as_f64_slice(&self) -> &[f64] {
251        match self {
252            Self::PerParticle {data, ..} => bytemuck::cast_slice::<u8,f64>(data),
253            Self::Uniform(_) => unimplemented!()
254        }
255    }
256
257    pub(crate) fn as_i64_slice_mut(&mut self) -> &mut [i64] {
258        match self {
259            Self::PerParticle {data, ..} => bytemuck::cast_slice_mut::<u8,i64>(data),
260            Self::Uniform(_) => unimplemented!()
261        }
262    }
263
264    pub(crate) fn as_f64_slice_mut(&mut self) -> &mut [f64] {
265        match self {
266            Self::PerParticle {data, ..} => bytemuck::cast_slice_mut::<u8,f64>(data),
267            Self::Uniform(_) => unimplemented!()
268        }
269    }
270}
fips_md/runtime/particle_data.rs

fips_md/runtime/
particle_data.rs
