poc_kokkos_rs/functor/
policies.rs

1use std::marker::ConstParamTy;
2
3#[cfg(feature = "rayon")]
4use rayon::prelude::*;
5
6use crate::functor::ForFunctor;
7
8// -- scheduling
9
10/// Scheduling enum. CURRENTLY IGNORED.
11///
12/// Used to set the workload scheduling policy. Defaults to [Schedule::Static].
13#[derive(ConstParamTy, Debug, Default, Clone, Copy, PartialEq, Eq)]
14pub enum Schedule {
15    #[default]
16    /// Default value. Workload is divided once and split equally between
17    /// computational ressources.
18    Static,
19    /// Dynamic scheduling. Workload is divided at start, split between computational
20    /// ressources and work stealing is enabled.
21    Dynamic,
22}
23
24// -- exec space
25
26/// Execution Space enum.
27///
28/// Used to specify the target device of execution for the dispatch.
29/// Defaults to [ExecutionSpace::Serial].
30#[derive(ConstParamTy, Debug, Default, Clone, Copy, PartialEq, Eq)]
31pub enum ExecutionSpace {
32    #[default]
33    /// Default value. Execute the kernel sequentially.
34    Serial,
35    /// Target the CPU. Execute the kernel in parallel by using a feature-determined
36    /// backend.
37    DeviceCPU,
38    /// Target the GPU. UNIMPLEMENTED.
39    DeviceGPU,
40}
41
42// -- exec policy
43
44pub(crate) trait ExecutionPolicy
45where
46    Self: Sized,
47{
48    type KernelArgType;
49
50    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, functor: F);
51
52    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, functor: F);
53
54    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, functor: F);
55}
56
57pub struct TeamHandle {
58    _league_rank: usize,
59    _team_rank: usize,
60    _team_size: usize,
61}
62
63pub struct Range(pub usize);
64pub struct MDRange<const N: usize>(pub [usize; N]);
65
66pub struct TeamPolicy(pub usize, pub usize, pub usize);
67pub struct PerTeam;
68pub struct PerThread;
69
70pub struct TeamThreadRange(pub TeamHandle, pub usize);
71pub struct TeamThreadMDRange<const N: usize>(pub TeamHandle, pub [usize; N]);
72
73pub struct TeamVectorRange(pub TeamHandle, pub usize);
74pub struct TeamVectorMDRange<const N: usize>(pub TeamHandle, pub [usize; N]);
75
76pub struct ThreadVectorRange(pub TeamHandle, pub usize);
77pub struct ThreadVectorMDRange<const N: usize>(pub TeamHandle, pub [usize; N]);
78
79impl ExecutionPolicy for Range {
80    type KernelArgType = usize;
81
82    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, functor: F) {
83        let _ = SCHEDULE;
84        let Range(n) = self;
85
86        (0..n).for_each(|i| functor.execute(i));
87    }
88
89    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, functor: F) {
90        let Range(n) = self;
91
92        let n_threads = 4; // TODO: read it properly
93
94        match SCHEDULE {
95            #[cfg(not(feature = "rayon"))]
96            Schedule::Static => {
97                let chunk_size = n / n_threads + 1;
98
99                std::thread::scope(|s| {
100                    let mut handles = Vec::with_capacity(n_threads);
101                    let f = &functor;
102                    let blocks = (0..n_threads)
103                        .map(|tid| tid * chunk_size..(tid * chunk_size + chunk_size).min(n));
104
105                    for b in blocks {
106                        let h = s.spawn(move || {
107                            b.into_iter().for_each(|i| f.execute(i));
108                        });
109                        handles.push(h);
110                    }
111
112                    for h in handles {
113                        h.join().unwrap();
114                    }
115                });
116            }
117            #[cfg(not(feature = "rayon"))]
118            Schedule::Dynamic => {
119                unimplemented!("E: Dynamic dispatch isn't supported without rayon");
120            }
121            #[cfg(feature = "rayon")]
122            Schedule::Static => {
123                let chunk_size = n / n_threads + 1;
124
125                // mimic static scheduling by generating one item to process per thread
126                (0..n_threads)
127                    .into_par_iter()
128                    .map(|tid| tid * chunk_size..(tid * chunk_size + chunk_size).min(n))
129                    .for_each(|b| {
130                        b.into_iter().for_each(|i| functor.execute(i));
131                    });
132            }
133            #[cfg(feature = "rayon")]
134            Schedule::Dynamic => {
135                (0..n).into_par_iter().for_each(|i| functor.execute(i));
136            }
137        }
138    }
139
140    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
141        todo!()
142    }
143}
144impl<const N: usize> ExecutionPolicy for MDRange<N> {
145    type KernelArgType = [usize; N];
146
147    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
148        todo!()
149    }
150
151    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
152        todo!()
153    }
154
155    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
156        todo!()
157    }
158}
159
160impl ExecutionPolicy for TeamPolicy {
161    type KernelArgType = TeamHandle;
162
163    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
164        todo!()
165    }
166
167    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
168        todo!()
169    }
170
171    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
172        todo!()
173    }
174}
175impl ExecutionPolicy for PerTeam {
176    type KernelArgType = usize; // ?
177
178    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
179        todo!()
180    }
181
182    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
183        todo!()
184    }
185
186    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
187        todo!()
188    }
189}
190impl ExecutionPolicy for PerThread {
191    type KernelArgType = usize; // ?
192
193    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
194        todo!()
195    }
196
197    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
198        todo!()
199    }
200
201    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
202        todo!()
203    }
204}
205
206impl ExecutionPolicy for TeamThreadRange {
207    type KernelArgType = usize;
208
209    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
210        todo!()
211    }
212
213    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
214        todo!()
215    }
216
217    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
218        todo!()
219    }
220}
221impl<const N: usize> ExecutionPolicy for TeamThreadMDRange<N> {
222    type KernelArgType = usize; // ?
223
224    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
225        todo!()
226    }
227
228    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
229        todo!()
230    }
231
232    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
233        todo!()
234    }
235}
236
237impl ExecutionPolicy for ThreadVectorRange {
238    type KernelArgType = usize;
239
240    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
241        todo!()
242    }
243
244    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
245        todo!()
246    }
247
248    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
249        todo!()
250    }
251}
252impl<const N: usize> ExecutionPolicy for ThreadVectorMDRange<N> {
253    type KernelArgType = usize; // ?
254
255    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
256        todo!()
257    }
258
259    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
260        todo!()
261    }
262
263    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
264        todo!()
265    }
266}
267
268impl ExecutionPolicy for TeamVectorRange {
269    type KernelArgType = usize;
270
271    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
272        todo!()
273    }
274
275    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
276        todo!()
277    }
278
279    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
280        todo!()
281    }
282}
283impl<const N: usize> ExecutionPolicy for TeamVectorMDRange<N> {
284    type KernelArgType = usize; // ?
285
286    fn dispatch_seq<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
287        todo!()
288    }
289
290    fn dispatch_cpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
291        todo!()
292    }
293
294    fn dispatch_gpu<const SCHEDULE: Schedule, F: ForFunctor<Self>>(self, _functor: F) {
295        todo!()
296    }
297}
298
299impl From<usize> for Range {
300    fn from(value: usize) -> Self {
301        Self(value)
302    }
303}
304
305impl<const N: usize> From<[usize; N]> for MDRange<N> {
306    fn from(value: [usize; N]) -> Self {
307        Self(value)
308    }
309}
poc_kokkos_rs/functor/policies.rs

poc_kokkos_rs/functor/
policies.rs
