Skip to content

Commit

Permalink
std_algos: for_each: try condense the impl
Browse files Browse the repository at this point in the history
  • Loading branch information
@fnrizzi
fnrizzi committed Jul 24, 2023
1 parent 96922d9 commit a17ca7b
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 42 deletions.
23 changes: 11 additions & 12 deletions algorithms/src/std_algorithms/Kokkos_ForEach.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,17 +32,17 @@ template <
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));
return Impl::for_each_impl(label.c_str(), 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));
return Impl::for_each_impl("Kokkos::for_each_iterator_api_default", ex, first,
last, std::move(functor));
}

template <
Expand All @@ -55,8 +55,8 @@ UnaryFunctorType for_each(const std::string& label, const ExecutionSpace& ex,
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));
return Impl::for_each_impl(label.c_str(), ex, KE::begin(v), KE::end(v),
std::move(functor));
}

template <
Expand All @@ -69,9 +69,8 @@ UnaryFunctorType for_each(const ExecutionSpace& ex,
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));
return Impl::for_each_impl("Kokkos::for_each_view_api_default", ex,
KE::begin(v), KE::end(v), std::move(functor));
}

//
Expand All @@ -85,7 +84,7 @@ template <class TeamHandleType, class IteratorType, class UnaryFunctorType,
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));
return Impl::for_each_impl("", teamHandle, first, last, std::move(functor));
}

template <class TeamHandleType, class DataType, class... Properties,
Expand All @@ -98,8 +97,8 @@ for_each(const TeamHandleType& teamHandle,
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));
return Impl::for_each_impl("", teamHandle, KE::begin(v), KE::end(v),
std::move(functor));
}

} // namespace Experimental
Expand Down
63 changes: 33 additions & 30 deletions algorithms/src/std_algorithms/impl/Kokkos_ForEachForEachN.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,24 +41,27 @@ struct StdForEachFunctor {
: m_first(std::move(_first)), m_functor(std::move(_functor)) {}
};

//
// exespace impl
//
template <class ExecutionSpace, class IteratorType, class UnaryFunctorType>
UnaryFunctorType for_each_exespace_impl(const std::string& label,
const ExecutionSpace& ex,
IteratorType first, IteratorType last,
UnaryFunctorType functor) {
template <class HandleType, class IteratorType, class UnaryFunctorType>
KOKKOS_FUNCTION UnaryFunctorType
for_each_impl([[maybe_unused]] const char* label, const HandleType& handle,
IteratorType first, IteratorType last, UnaryFunctorType functor) {
// checks
Impl::static_assert_random_access_and_accessible(ex, first);
Impl::static_assert_random_access_and_accessible(handle, first);
Impl::expect_valid_range(first, last);

// run
const auto num_elements = Kokkos::Experimental::distance(first, last);
::Kokkos::parallel_for(
label, RangePolicy<ExecutionSpace>(ex, 0, num_elements),
StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));
ex.fence("Kokkos::for_each: fence after operation");
if constexpr (is_execution_space_v<HandleType>) {
::Kokkos::parallel_for(
label, RangePolicy<HandleType>(handle, 0, num_elements),
StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));
handle.fence("Kokkos::for_each: fence after operation");
} else {
::Kokkos::parallel_for(
TeamThreadRange(handle, 0, num_elements),
StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));
handle.team_barrier();
}

return functor;
}
Expand All @@ -77,7 +80,7 @@ IteratorType for_each_n_exespace_impl(const std::string& label,
return first;
}

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

return last;
Expand All @@ -86,24 +89,24 @@ 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) {
// checks
Impl::static_assert_random_access_and_accessible(teamHandle, first);
Impl::expect_valid_range(first, last);
// template <class TeamHandleType, class IteratorType, class UnaryFunctorType>
// KOKKOS_FUNCTION UnaryFunctorType
// 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);

// run
const auto num_elements = Kokkos::Experimental::distance(first, last);
::Kokkos::parallel_for(
TeamThreadRange(teamHandle, 0, num_elements),
StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));
// // run
// const auto num_elements = Kokkos::Experimental::distance(first, last);
// ::Kokkos::parallel_for(
// TeamThreadRange(teamHandle, 0, num_elements),
// StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));

teamHandle.team_barrier();
// teamHandle.team_barrier();

return functor;
}
// return functor;
// }

template <class TeamHandleType, class IteratorType, class SizeType,
class UnaryFunctorType>
Expand All @@ -118,7 +121,7 @@ for_each_n_team_impl(const TeamHandleType& teamHandle, IteratorType first,
return first;
}

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

return last;
Expand Down

0 comments on commit a17ca7b

Please sign in to comment.