Skip to content

Commit

Permalink
Do not return a copy of the input functor for Kokkos::Experimental::f…
Browse files Browse the repository at this point in the history
…or_each
  • Loading branch information
tpadioleau committed Apr 2, 2024
1 parent 6355510 commit 391e040
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 42 deletions.
56 changes: 26 additions & 30 deletions algorithms/src/std_algorithms/Kokkos_ForEach.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -29,49 +29,46 @@ namespace Experimental {
template <
class ExecutionSpace, class IteratorType, class UnaryFunctorType,
std::enable_if_t<Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
UnaryFunctorType for_each(const std::string& label, const ExecutionSpace& ex,
IteratorType first, IteratorType last,
UnaryFunctorType functor) {
return Impl::for_each_exespace_impl(label, ex, first, last,
std::move(functor));
void for_each(const std::string& label, const ExecutionSpace& ex,
IteratorType first, IteratorType last, UnaryFunctorType functor) {
Impl::for_each_exespace_impl(label, ex, first, last, std::move(functor));
}

template <
class ExecutionSpace, class IteratorType, class UnaryFunctorType,
std::enable_if_t<Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
UnaryFunctorType for_each(const ExecutionSpace& ex, IteratorType first,
IteratorType last, UnaryFunctorType functor) {
return Impl::for_each_exespace_impl("Kokkos::for_each_iterator_api_default",
ex, first, last, std::move(functor));
void for_each(const ExecutionSpace& ex, IteratorType first, IteratorType last,
UnaryFunctorType functor) {
Impl::for_each_exespace_impl("Kokkos::for_each_iterator_api_default", ex,
first, last, std::move(functor));
}

template <
class ExecutionSpace, class DataType, class... Properties,
class UnaryFunctorType,
std::enable_if_t<Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
UnaryFunctorType for_each(const std::string& label, const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v,
UnaryFunctorType functor) {
void for_each(const std::string& label, const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v,
UnaryFunctorType functor) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

namespace KE = ::Kokkos::Experimental;
return Impl::for_each_exespace_impl(label, ex, KE::begin(v), KE::end(v),
std::move(functor));
Impl::for_each_exespace_impl(label, ex, KE::begin(v), KE::end(v),
std::move(functor));
}

template <
class ExecutionSpace, class DataType, class... Properties,
class UnaryFunctorType,
std::enable_if_t<Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
UnaryFunctorType for_each(const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v,
UnaryFunctorType functor) {
void for_each(const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v,
UnaryFunctorType functor) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

namespace KE = ::Kokkos::Experimental;
return Impl::for_each_exespace_impl("Kokkos::for_each_view_api_default", ex,
KE::begin(v), KE::end(v),
std::move(functor));
Impl::for_each_exespace_impl("Kokkos::for_each_view_api_default", ex,
KE::begin(v), KE::end(v), std::move(functor));
}

//
Expand All @@ -82,24 +79,23 @@ UnaryFunctorType for_each(const ExecutionSpace& ex,

template <class TeamHandleType, class IteratorType, class UnaryFunctorType,
std::enable_if_t<Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
KOKKOS_FUNCTION UnaryFunctorType for_each(const TeamHandleType& teamHandle,
IteratorType first, IteratorType last,
UnaryFunctorType functor) {
return Impl::for_each_team_impl(teamHandle, first, last, std::move(functor));
KOKKOS_FUNCTION void for_each(const TeamHandleType& teamHandle,
IteratorType first, IteratorType last,
UnaryFunctorType functor) {
Impl::for_each_team_impl(teamHandle, first, last, std::move(functor));
}

template <class TeamHandleType, class DataType, class... Properties,
class UnaryFunctorType,
std::enable_if_t<Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
KOKKOS_FUNCTION UnaryFunctorType
for_each(const TeamHandleType& teamHandle,
const ::Kokkos::View<DataType, Properties...>& v,
UnaryFunctorType functor) {
KOKKOS_FUNCTION void for_each(const TeamHandleType& teamHandle,
const ::Kokkos::View<DataType, Properties...>& v,
UnaryFunctorType functor) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

namespace KE = ::Kokkos::Experimental;
return Impl::for_each_team_impl(teamHandle, KE::begin(v), KE::end(v),
std::move(functor));
Impl::for_each_team_impl(teamHandle, KE::begin(v), KE::end(v),
std::move(functor));
}

} // namespace Experimental
Expand Down
20 changes: 8 additions & 12 deletions algorithms/src/std_algorithms/impl/Kokkos_ForEachForEachN.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -42,10 +42,9 @@ struct StdForEachFunctor {
};

template <class HandleType, class IteratorType, class UnaryFunctorType>
UnaryFunctorType for_each_exespace_impl(const std::string& label,
const HandleType& handle,
IteratorType first, IteratorType last,
UnaryFunctorType functor) {
void for_each_exespace_impl(const std::string& label, const HandleType& handle,
IteratorType first, IteratorType last,
UnaryFunctorType functor) {
// checks
Impl::static_assert_random_access_and_accessible(handle, first);
Impl::expect_valid_range(first, last);
Expand All @@ -56,8 +55,6 @@ UnaryFunctorType for_each_exespace_impl(const std::string& label,
label, RangePolicy<HandleType>(handle, 0, num_elements),
StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));
handle.fence("Kokkos::for_each: fence after operation");

return functor;
}

template <class ExecutionSpace, class IteratorType, class SizeType,
Expand All @@ -75,7 +72,7 @@ IteratorType for_each_n_exespace_impl(const std::string& label,
}

for_each_exespace_impl(label, ex, first, last, std::move(functor));
// no neeed to fence since for_each_exespace_impl fences already
// no need to fence since for_each_exespace_impl fences already

return last;
}
Expand All @@ -84,9 +81,9 @@ IteratorType for_each_n_exespace_impl(const std::string& label,
// team impl
//
template <class TeamHandleType, class IteratorType, class UnaryFunctorType>
KOKKOS_FUNCTION UnaryFunctorType
for_each_team_impl(const TeamHandleType& teamHandle, IteratorType first,
IteratorType last, UnaryFunctorType functor) {
KOKKOS_FUNCTION void for_each_team_impl(const TeamHandleType& teamHandle,
IteratorType first, IteratorType last,
UnaryFunctorType functor) {
// checks
Impl::static_assert_random_access_and_accessible(teamHandle, first);
Impl::expect_valid_range(first, last);
Expand All @@ -96,7 +93,6 @@ for_each_team_impl(const TeamHandleType& teamHandle, IteratorType first,
TeamThreadRange(teamHandle, 0, num_elements),
StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));
teamHandle.team_barrier();
return functor;
}

template <class TeamHandleType, class IteratorType, class SizeType,
Expand All @@ -113,7 +109,7 @@ for_each_n_team_impl(const TeamHandleType& teamHandle, IteratorType first,
}

for_each_team_impl(teamHandle, first, last, std::move(functor));
// no neeed to fence since for_each_team_impl fences already
// no need to fence since for_each_team_impl fences already

return last;
}
Expand Down

0 comments on commit 391e040

Please sign in to comment.