poc_kokkos_rs/routines/
parameters.rs

1//! parallel statement parameterization code
2//!
3//! This module contains all code used to parameterize parallel statements. Some
4//! parameters are direct Kokkos replicates while others are Rust-specific.
5//!
6//! Notable differences include:
7//! - [ExecutionPolicy] struct: Instead of having multiple types of execution policies
8//!   for each range, with re-occuring parameters, range specification is now a
9//!   subparameter of execution policies.
10//!
11
12use std::ops::Range;
13
14/// Execution Space enum.
15///
16/// Used to specify the target device of execution for the dispatch.
17/// Defaults to [ExecutionSpace::Serial].
18#[derive(Debug, Default, Clone)]
19pub enum ExecutionSpace {
20    #[default]
21    /// Default value. Execute the kernel sequentially.
22    Serial,
23    /// Target the CPU. Execute the kernel in parallel by using a feature-determined
24    /// backend.
25    DeviceCPU,
26    /// Target the GPU. UNIMPLEMENTED.
27    DeviceGPU,
28}
29
30#[derive(Debug, Clone)]
31/// Range Policy enum.
32///
33/// This holds information related to the looping structure adopted by the routine.
34pub enum RangePolicy<const N: usize> {
35    // Outer range
36    /// 1D iteration range.
37    RangePolicy(Range<usize>),
38    /// N-dimensional iteration range.
39    MDRangePolicy([Range<usize>; N]),
40    /// Team-based iteration policy.
41    TeamPolicy {
42        /// Number of team.
43        league_size: usize,
44        /// Number of threads per team.
45        team_size: usize,
46        /// Number of vector
47        vector_size: usize,
48    },
49
50    // Specifics
51    /// Policy used to ensure each team execute the body once and only once.
52    PerTeam,
53    /// Policy used to ensure each thread execute the body once and only once.
54    PerThread,
55
56    // Medium range
57    /// Medium-level depth. Can host further nests using vectors.
58    TeamThreadRange,
59    /// Medium-level depth. Can host further nests using vectors.
60    TeamThreadMDRange,
61
62    /// Medium-level depth. Cannot host further nests.
63    TeamVectorRange,
64    /// Medium-level depth. Cannot host further nests.
65    TeamVectorMDRange,
66
67    // Inner Range
68    /// Inner-level depth. Cannot host further nests.
69    ThreadVectorRange,
70    /// Inner-level depth. Cannot host further nests.
71    ThreadVectorMDRange,
72}
73
74/// Scheduling enum. CURRENTLY IGNORED.
75///
76/// Used to set the workload scheduling policy. Defaults to [Schedule::Static].
77#[derive(Debug, Default, Clone)]
78pub enum Schedule {
79    #[default]
80    /// Default value. Workload is divided once and split equally between
81    /// computational ressources.
82    Static,
83    /// Dynamic scheduling. Workload is divided at start, split between computational
84    /// ressources and work stealing is enabled.
85    Dynamic,
86}
87
88#[derive(Debug, Clone)]
89/// Execution Policy enum. See Kokkos documentation for explanation on their model.
90///
91/// ### Example
92///
93/// ```rust
94/// use poc_kokkos_rs::routines::parameters::{
95///     ExecutionPolicy,
96///     ExecutionSpace,
97///     RangePolicy,
98///     Schedule
99/// };
100///
101/// let length: usize = 8;
102///
103/// let execp =  ExecutionPolicy::<1> {
104///         space: ExecutionSpace::DeviceCPU, // will try to parallelize code on CPU
105///         range: RangePolicy::RangePolicy(0..length), // equivalent to "for i in 0..length"
106///         schedule: Schedule::Static, // static division of workload
107///     };
108/// ```
109pub struct ExecutionPolicy<const N: usize> {
110    /// Execution space targetted by the dispatch.
111    pub space: ExecutionSpace,
112    /// Iteration pattern used to handle the workload.
113    pub range: RangePolicy<N>,
114    /// Scheduling policy for the dispatch. CURRENTLY IGNORED.
115    pub schedule: Schedule,
116}