poc_kokkos_rs/functor/
for_kernel.rs

1use crate::functor::{ExecutionPolicy, ExecutionSpace, Schedule};
2
3#[allow(private_bounds)]
4pub trait ForFunctor<P: ExecutionPolicy>: Send + Sync {
5    /// Kernel to execute.
6    #[allow(private_interfaces)]
7    fn execute(&self, args: P::KernelArgType);
8}
9
10#[allow(private_interfaces)]
11impl<P: ExecutionPolicy, Closure: Fn(P::KernelArgType) + Send + Sync> ForFunctor<P> for Closure {
12    fn execute(&self, args: <P as ExecutionPolicy>::KernelArgType) {
13        self(args)
14    }
15}
16
17#[allow(private_bounds)]
18pub fn parallel_for<
19    const EXECUTION_SPACE: ExecutionSpace,
20    const SCHEDULE: Schedule,
21    P: ExecutionPolicy,
22    F: ForFunctor<P>,
23>(
24    _label: Option<&str>, // TODO: debug-level log?
25    policy: P,
26    functor: F,
27) {
28    match EXECUTION_SPACE {
29        ExecutionSpace::Serial => {
30            policy.dispatch_seq::<SCHEDULE, F>(functor);
31        }
32        ExecutionSpace::DeviceCPU => {
33            policy.dispatch_cpu::<SCHEDULE, F>(functor);
34        }
35        ExecutionSpace::DeviceGPU => {
36            policy.dispatch_gpu::<SCHEDULE, F>(functor);
37        }
38    }
39}
poc_kokkos_rs/functor/for_kernel.rs

poc_kokkos_rs/functor/
for_kernel.rs
