DDC 0.7.0
Loading...
Searching...
No Matches
parallel_for_each.hpp
1// Copyright (C) The DDC development team, see COPYRIGHT.md file
2//
3// SPDX-License-Identifier: MIT
4
5#pragma once
6
7#include <cstddef>
8#include <string>
9#include <type_traits>
10#include <utility>
11
12#include <Kokkos_Core.hpp>
13
14#include "detail/kokkos.hpp"
15
16#include "ddc_to_kokkos_execution_policy.hpp"
17#include "discrete_domain.hpp"
18#include "discrete_element.hpp"
19
20namespace ddc {
21
22namespace detail {
23
24template <class F, class Support, class IndexSequence>
25class ForEachKokkosLambdaAdapter;
26
27template <class F, class Support, std::size_t... Idx>
28class ForEachKokkosLambdaAdapter<F, Support, std::index_sequence<Idx...>>
29{
30 template <std::size_t I>
31 using index_type = DiscreteElementType;
32
33 F m_f;
34
35 Support m_support;
36
37public:
38 explicit ForEachKokkosLambdaAdapter(F const& f, Support const& support)
39 : m_f(f)
40 , m_support(support)
41 {
42 }
43
44 template <std::size_t N = sizeof...(Idx), std::enable_if_t<(N == 0), bool> = true>
45 KOKKOS_FUNCTION void operator()([[maybe_unused]] index_type<0> unused_id) const
46 {
47 m_f(DiscreteElement<>());
48 }
49
50 template <std::size_t N = sizeof...(Idx), std::enable_if_t<(N > 0), bool> = true>
51 KOKKOS_FUNCTION void operator()(index_type<Idx>... ids) const
52 {
53 m_f(m_support(typename Support::discrete_vector_type(ids...)));
54 }
55};
56
57template <class ExecSpace, class Support, class Functor>
58void for_each_kokkos(
59 std::string const& label,
60 ExecSpace const& execution_space,
61 Support const& domain,
62 Functor const& f) noexcept
63{
64 Kokkos::parallel_for(
65 label,
66 ddc_to_kokkos_execution_policy(execution_space, domain),
67 ForEachKokkosLambdaAdapter<
68 Functor,
69 Support,
70 std::make_index_sequence<Support::rank()>>(f, domain));
71}
72
73} // namespace detail
74
75/** iterates over a nD domain using a given `Kokkos` execution space
76 * @param[in] label name for easy identification of the parallel_for_each algorithm
77 * @param[in] execution_space a Kokkos execution space where the loop will be executed on
78 * @param[in] domain the domain over which to iterate
79 * @param[in] f a functor taking an index as parameter
80 */
81template <class ExecSpace, class Support, class Functor>
83 std::string const& label,
84 ExecSpace const& execution_space,
85 Support const& domain,
86 Functor&& f) noexcept
87{
88 detail::for_each_kokkos(label, execution_space, domain, std::forward<Functor>(f));
89}
90
91/** iterates over a nD domain using a given `Kokkos` execution space
92 * @param[in] execution_space a Kokkos execution space where the loop will be executed on
93 * @param[in] domain the domain over which to iterate
94 * @param[in] f a functor taking an index as parameter
95 */
96template <class ExecSpace, class Support, class Functor>
97std::enable_if_t<Kokkos::is_execution_space_v<ExecSpace>> parallel_for_each(
98 ExecSpace const& execution_space,
99 Support const& domain,
100 Functor&& f) noexcept
101{
102 detail::for_each_kokkos(
103 "ddc_for_each_default",
104 execution_space,
105 domain,
106 std::forward<Functor>(f));
107}
108
109/** iterates over a nD domain using the `Kokkos` default execution space
110 * @param[in] label name for easy identification of the parallel_for_each algorithm
111 * @param[in] domain the domain over which to iterate
112 * @param[in] f a functor taking an index as parameter
113 */
114template <class Support, class Functor>
115void parallel_for_each(std::string const& label, Support const& domain, Functor&& f) noexcept
116{
117 parallel_for_each(label, Kokkos::DefaultExecutionSpace(), domain, std::forward<Functor>(f));
118}
119
120/** iterates over a nD domain using the `Kokkos` default execution space
121 * @param[in] domain the domain over which to iterate
122 * @param[in] f a functor taking an index as parameter
123 */
124template <class Support, class Functor>
125void parallel_for_each(Support const& domain, Functor&& f) noexcept
126{
127 parallel_for_each(
128 "ddc_for_each_default",
129 Kokkos::DefaultExecutionSpace(),
130 domain,
131 std::forward<Functor>(f));
132}
133
134} // namespace ddc
KOKKOS_DEFAULTED_FUNCTION constexpr DiscreteElement()=default
The top-level namespace of DDC.
void parallel_for_each(std::string const &label, Support const &domain, Functor &&f) noexcept
iterates over a nD domain using the Kokkos default execution space
void parallel_for_each(Support const &domain, Functor &&f) noexcept
iterates over a nD domain using the Kokkos default execution space
std::enable_if_t< Kokkos::is_execution_space_v< ExecSpace > > parallel_for_each(ExecSpace const &execution_space, Support const &domain, Functor &&f) noexcept
iterates over a nD domain using a given Kokkos execution space
void parallel_for_each(std::string const &label, ExecSpace const &execution_space, Support const &domain, Functor &&f) noexcept
iterates over a nD domain using a given Kokkos execution space