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}