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>;