poc_kokkos_rs/view/mod.rs
1//! data structure related code
2//!
3//! This module contains code used for the implementations of `Views`, a data structure
4//! defined and used by the Kokkos library. There are different types of views, all
5//! implemented using the same backend, [ViewBase].
6//!
7//! Eventually, the different types of Views should be removed and replaced by a single
8//! type. The distinction between original and mirrors doesn't seem necessary in a Rust
9//! implementation where the ownership system handles all memory transaction.
10//!
11//! In order to have thread-safe structures to use in parallel statement, the inner data
12//! type of views is adjusted implicitly when compiling using parallelization features.
13//! To match the adjusted data type, view access is done through `get` and `set` methods,
14//! allowing for feature-specific mutability in signatures while keeping a consistent user
15//! API.
16//!
17//! Parameters of aforementionned views are defined in the [`parameters`] sub-module.
18//!
19//! ### Example
20//!
21//! Initialize and fill a 2D matrix:
22//! ```rust
23//! use poc_kokkos_rs::view::{
24//! parameters::Layout,
25//! ViewOwned,
26//! };
27//!
28//! let mut viewA: ViewOwned<'_, 2, f64> = ViewOwned::new(
29//! Layout::Right, // see parameters & Kokkos doc
30//! [3, 5], // 3 rows, 5 columns
31//! );
32//!
33//! for row in 0..3 {
34//! for col in 0..5 {
35//! viewA.set([row, col], row as f64);
36//! }
37//! }
38//!
39//! // viewA:
40//! // (0.0 0.0 0.0 0.0 0.0)
41//! // (1.0 1.0 1.0 1.0 1.0)
42//! // (2.0 2.0 2.0 2.0 2.0)
43//! ```
44
45pub mod parameters;
46
47#[cfg(any(feature = "rayon", feature = "threads", feature = "gpu"))]
48use atomic::{Atomic, Ordering};
49
50#[cfg(any(doc, not(any(feature = "rayon", feature = "threads", feature = "gpu"))))]
51use std::ops::IndexMut;
52
53use self::parameters::{compute_stride, DataTraits, DataType, InnerDataType, Layout};
54use std::{fmt::Debug, ops::Index};
55
56#[derive(Debug)]
57/// Enum used to classify view-related errors.
58///
59/// In all variants, the internal value is a description of the error.
60pub enum ViewError<'a> {
61 ValueError(&'a str),
62 DoubleMirroring(&'a str),
63}
64
65#[derive(Debug, PartialEq)]
66/// Common structure used as the backend of all `View` types. The main differences between
67/// usable types is the type of the `data` field.
68pub struct ViewBase<'a, const N: usize, T>
69where
70 T: DataTraits,
71{
72 /// Data container. Depending on the type, it can be a vector (`Owned`), a reference
73 /// (`ReadOnly`) or a mutable reference (`ReadWrite`).
74 pub data: DataType<'a, T>,
75 /// Memory layout of the view. Refer to Kokkos documentation for more information.
76 pub layout: Layout<N>,
77 /// Dimensions of the data represented by the view. The view can:
78 ///
79 /// - be a vector (1 dimension)
80 /// - be a multi-dimensionnal array (up to 8 dimensions)
81 ///
82 /// The number of dimensions is referred to as the _depth_. Dimension 0, i.e. scalar,
83 /// is not directly supported.
84 pub dim: [usize; N],
85 /// Stride between each element of a given dimension. Computed automatically for
86 /// [Layout::Left] and [Layout::Right].
87 pub stride: [usize; N],
88}
89
90#[cfg(not(any(feature = "rayon", feature = "threads", feature = "gpu")))]
91// ~~~~~~~~ Constructors
92impl<const N: usize, T> ViewBase<'_, N, T>
93where
94 T: DataTraits, // fair assumption imo
95{
96 /// Constructor used to create owned views. See dedicated methods for others.
97 pub fn new(layout: Layout<N>, dim: [usize; N]) -> Self {
98 // compute stride & capacity
99 let stride = compute_stride(&dim, &layout);
100 let capacity: usize = dim.iter().product();
101
102 // build & return
103 Self {
104 data: DataType::Owned(vec![T::default(); capacity]), // should this be allocated though?
105 layout,
106 dim,
107 stride,
108 }
109 }
110
111 /// Constructor used to create owned views. See dedicated methods for others.
112 pub fn new_from_data(data: Vec<T>, layout: Layout<N>, dim: [usize; N]) -> Self {
113 // compute stride if necessary
114 let stride = compute_stride(&dim, &layout);
115
116 // checks
117 let capacity: usize = dim.iter().product();
118 assert_eq!(capacity, data.len());
119
120 // build & return
121 Self {
122 data: DataType::Owned(data),
123 layout,
124 dim,
125 stride,
126 }
127 }
128}
129
130#[cfg(any(feature = "rayon", feature = "threads", feature = "gpu"))]
131// ~~~~~~~~ Constructors
132impl<const N: usize, T> ViewBase<'_, N, T>
133where
134 T: DataTraits, // fair assumption imo
135{
136 /// Constructor used to create owned views. See dedicated methods for others.
137 pub fn new(layout: Layout<N>, dim: [usize; N]) -> Self {
138 // compute stride & capacity
139 let stride = compute_stride(&dim, &layout);
140 let capacity: usize = dim.iter().product();
141
142 // build & return
143 Self {
144 data: DataType::Owned((0..capacity).map(|_| Atomic::new(T::default())).collect()), // should this be allocated though?
145 layout,
146 dim,
147 stride,
148 }
149 }
150
151 /// Constructor used to create owned views. See dedicated methods for others.
152 pub fn new_from_data(data: Vec<T>, layout: Layout<N>, dim: [usize; N]) -> Self {
153 // compute stride if necessary
154 let stride = compute_stride(&dim, &layout);
155
156 // checks
157 let capacity: usize = dim.iter().product();
158 assert_eq!(capacity, data.len());
159
160 // build & return
161 Self {
162 data: DataType::Owned(data.into_iter().map(|elem| Atomic::new(elem)).collect()),
163 layout,
164 dim,
165 stride,
166 }
167 }
168}
169
170// ~~~~~~~~ Uniform writing interface across all features
171impl<'a, const N: usize, T> ViewBase<'a, N, T>
172where
173 T: DataTraits,
174{
175 #[inline(always)]
176 #[cfg(not(any(feature = "rayon", feature = "threads", feature = "gpu")))]
177 /// Writing interface.
178 ///
179 /// Two different implementations of this method are defined in order to satisfy
180 /// the (im)mutability requirements when using parallelization features & keep a
181 /// consistent user API:
182 ///
183 /// - any feature enabled: implictly use an atomic store operation on top of the
184 /// regular [Index] trait implementation to prevent a mutable borrow. The store
185 /// currently uses relaxed ordering, this may change.
186 /// - no feature enabled: uses a regular [IndexMut] trait implementation.
187 ///
188 /// Note that [Index] is always implemented while [IndexMut] only is when no
189 /// features are enabled.
190 ///
191 /// **Current version**: no feature
192 pub fn set(&mut self, index: [usize; N], val: T) {
193 self[index] = val;
194 }
195
196 #[inline(always)]
197 #[cfg(any(feature = "rayon", feature = "threads", feature = "gpu"))]
198 /// Writing interface.
199 ///
200 /// Two different implementations of this method are defined in order to satisfy
201 /// the (im)mutability requirements when using parallelization features & keep a
202 /// consistent user API:
203 ///
204 /// - any feature enabled: implictly use an atomic store operation on top of the
205 /// regular [Index] trait implementation to prevent a mutable borrow. The store
206 /// currently uses relaxed ordering, this may change.
207 /// - no feature enabled: uses a regular [IndexMut] trait implementation.
208 ///
209 /// Note that [Index] is always implemented while [IndexMut] only is when no
210 /// features are enabled.
211 ///
212 /// **Current version**: thread-safe
213 pub fn set(&self, index: [usize; N], val: T) {
214 self[index].store(val, Ordering::Relaxed);
215 }
216
217 #[inline(always)]
218 #[cfg(not(any(feature = "rayon", feature = "threads", feature = "gpu")))]
219 /// Reading interface.
220 ///
221 /// Two different implementations of this method are defined in order to keep a
222 /// consistent user API across features:
223 ///
224 /// - any feature enabled: implictly use an atomic load operation on top of the
225 /// regular [Index] trait implementation. The load currently uses relaxed ordering,
226 /// this may change.
227 /// - no feature enabled: uses the regular [Index] trait implementation.
228 ///
229 /// Note that [Index] is always implemented while [IndexMut] only is when no
230 /// features are enabled.
231 ///
232 /// **Current version**: no feature
233 pub fn get(&self, index: [usize; N]) -> T {
234 self[index]
235 }
236
237 #[inline(always)]
238 #[cfg(any(feature = "rayon", feature = "threads", feature = "gpu"))]
239 /// Reading interface.
240 ///
241 /// Two different implementations of this method are defined in order to keep a
242 /// consistent user API across features:
243 ///
244 /// - any feature enabled: implictly use an atomic load operation on top of the
245 /// regular [Index] trait implementation. The load currently uses relaxed ordering,
246 /// this may change.
247 /// - no feature enabled: uses the regular [Index] trait implementation.
248 ///
249 /// Note that [Index] is always implemented while [IndexMut] only is when no
250 /// features are enabled.
251 ///
252 /// **Current version**: thread-safe
253 pub fn get(&self, index: [usize; N]) -> T {
254 self[index].load(atomic::Ordering::Relaxed)
255 }
256
257 // ~~~~~~~~ Mirrors
258
259 /// Create a new View mirroring `self`, i.e. referencing the same data. This mirror
260 /// is always immutable, but it inner values might still be writable if they are
261 /// atomic types.
262 ///
263 /// Note that mirrors currently can only be created from the "original" view,
264 /// i.e. the view owning the data.
265 pub fn create_mirror<'b>(&'a self) -> Result<ViewRO<'b, N, T>, ViewError<'a>>
266 where
267 'a: 'b, // 'a outlives 'b
268 {
269 let inner = if let DataType::Owned(v) = &self.data {
270 v
271 } else {
272 return Err(ViewError::DoubleMirroring(
273 "Cannot create a mirror from a non-data-owning View",
274 ));
275 };
276
277 Ok(Self {
278 data: DataType::Borrowed(inner),
279 layout: self.layout,
280 dim: self.dim,
281 stride: self.stride,
282 })
283 }
284
285 #[cfg(not(any(feature = "rayon", feature = "threads", feature = "gpu")))]
286 /// Create a new View mirroring `self`, i.e. referencing the same data. This mirror
287 /// uses a mutable reference, hence the serial-only definition
288 ///
289 /// Note that mirrors currently can only be created from the "original" view,
290 /// i.e. the view owning the data.
291 ///
292 /// Only defined when no feature are enabled since all interfaces should be immutable
293 /// otherwise.
294 pub fn create_mutable_mirror<'b>(&'a mut self) -> Result<ViewRW<'b, N, T>, ViewError<'a>>
295 where
296 'a: 'b, // 'a outlives 'b
297 {
298 let inner = if let DataType::Owned(v) = &mut self.data {
299 v
300 } else {
301 return Err(ViewError::DoubleMirroring(
302 "Cannot create a mirror from a non-data-owning View",
303 ));
304 };
305
306 Ok(Self {
307 data: DataType::MutBorrowed(inner),
308 layout: self.layout,
309 dim: self.dim,
310 stride: self.stride,
311 })
312 }
313
314 // ~~~~~~~~ Convenience
315
316 #[cfg(all(
317 test,
318 not(any(feature = "rayon", feature = "threads", feature = "gpu"))
319 ))]
320 /// Consumes the view to return a `Vec` containing its raw data content.
321 ///
322 /// This method is meant to be used in tests
323 pub fn raw_val<'b>(self) -> Result<Vec<T>, ViewError<'b>> {
324 if let DataType::Owned(v) = self.data {
325 Ok(v)
326 } else {
327 Err(ViewError::ValueError(
328 "Cannot fetch raw values of a non-data-owning views",
329 ))
330 }
331 }
332
333 #[cfg(all(test, any(feature = "rayon", feature = "threads", feature = "gpu")))]
334 /// Consumes the view to return a `Vec` containing its raw data content.
335 ///
336 /// This method is meant to be used in tests
337 pub fn raw_val<'b>(self) -> Result<Vec<T>, ViewError<'b>> {
338 if let DataType::Owned(v) = self.data {
339 Ok(v.iter()
340 .map(|elem| elem.load(atomic::Ordering::Relaxed))
341 .collect::<Vec<T>>())
342 } else {
343 Err(ViewError::ValueError(
344 "Cannot fetch raw values of a non-data-owning views",
345 ))
346 }
347 }
348
349 #[inline(always)]
350 /// Mapping function between N-indices and the flat offset.
351 pub fn flat_idx(&self, index: [usize; N]) -> usize {
352 index
353 .iter()
354 .zip(self.stride.iter())
355 .map(|(i, s_i)| *i * *s_i)
356 .sum()
357 }
358}
359
360/// **Read-only access is always implemented.**
361impl<const N: usize, T> Index<[usize; N]> for ViewBase<'_, N, T>
362where
363 T: DataTraits,
364{
365 type Output = InnerDataType<T>;
366
367 fn index(&self, index: [usize; N]) -> &Self::Output {
368 let flat_idx: usize = self.flat_idx(index);
369 match &self.data {
370 DataType::Owned(v) => {
371 assert!(flat_idx < v.len()); // remove bounds check
372 &v[flat_idx]
373 }
374 DataType::Borrowed(slice) => {
375 assert!(flat_idx < slice.len()); // remove bounds check
376 &slice[flat_idx]
377 }
378 DataType::MutBorrowed(mut_slice) => {
379 assert!(flat_idx < mut_slice.len()); // remove bounds check
380 &mut_slice[flat_idx]
381 }
382 }
383 }
384}
385
386#[cfg(not(any(feature = "rayon", feature = "threads", feature = "gpu")))]
387/// **Read-write access is only implemented when no parallel features are enabled.**
388impl<const N: usize, T> IndexMut<[usize; N]> for ViewBase<'_, N, T>
389where
390 T: DataTraits,
391{
392 fn index_mut(&mut self, index: [usize; N]) -> &mut Self::Output {
393 let flat_idx: usize = self.flat_idx(index);
394 match &mut self.data {
395 DataType::Owned(v) => {
396 assert!(flat_idx < v.len()); // remove bounds check
397 &mut v[flat_idx]
398 }
399 DataType::Borrowed(_) => unimplemented!("Cannot mutably access a read-only view!"),
400 DataType::MutBorrowed(mut_slice) => {
401 assert!(flat_idx < mut_slice.len()); // remove bounds check
402 &mut mut_slice[flat_idx]
403 }
404 }
405 }
406}
407
408/// View type owning the data it yields access to, i.e. "original" view.
409pub type ViewOwned<'a, const N: usize, T> = ViewBase<'a, N, T>;
410
411/// View type owning a read-only borrow to the data it yields access to, i.e. a
412/// read-only mirror.
413pub type ViewRO<'a, const N: usize, T> = ViewBase<'a, N, T>;
414
415/// View type owning a mutable borrow to the data it yields access to, i.e. a
416/// read-write mirror.
417pub type ViewRW<'a, const N: usize, T> = ViewBase<'a, N, T>;