Skip to content

Commit

Permalink
bring back previous code as discussed in meeting
Browse files Browse the repository at this point in the history
  • Loading branch information
@fnrizzi
fnrizzi committed Jul 27, 2023
1 parent 5a7039f commit aaf00c0
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 67 deletions.
23 changes: 12 additions & 11 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_impl(label.c_str(), ex, first, last,
std::move(functor));
return 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_impl("Kokkos::for_each_iterator_api_default", ex, first,
last, std::move(functor));
return Impl::for_each_exespace_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_impl(label.c_str(), ex, KE::begin(v), KE::end(v),
std::move(functor));
return Impl::for_each_exespace_impl(label, ex, KE::begin(v), KE::end(v),
std::move(functor));
}

template <
Expand All @@ -69,8 +69,9 @@ UnaryFunctorType for_each(const ExecutionSpace& ex,
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

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

//
Expand All @@ -84,7 +85,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_impl("", teamHandle, first, last, std::move(functor));
return Impl::for_each_team_impl(teamHandle, first, last, std::move(functor));
}

template <class TeamHandleType, class DataType, class... Properties,
Expand All @@ -97,8 +98,8 @@ for_each(const TeamHandleType& teamHandle,
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

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

} // namespace Experimental
Expand Down
81 changes: 25 additions & 56 deletions algorithms/src/std_algorithms/impl/Kokkos_ForEachForEachN.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,23 +23,6 @@
#include <std_algorithms/Kokkos_Distance.hpp>
#include <string>

#define KOKKOS_DO_PRAGMA(x) _Pragma(#x)
#if defined __NVCC__ && defined __NVCC_DIAG_PRAGMA_SUPPORT__
#define KOKKOS_NVCC_PRAGMA_PUSH KOKKOS_DO_PRAGMA(nv_diagnostic push)
#define KOKKOS_NVCC_PRAGMA_SUPPRESS(WARNING) \
KOKKOS_DO_PRAGMA(nv_diag_suppress = WARNING)
#define KOKKOS_NVCC_PRAGMA_POP KOKKOS_DO_PRAGMA(nv_diagnostic pop)
#elif defined __NVCC__
#define KOKKOS_NVCC_PRAGMA_PUSH KOKKOS_DO_PRAGMA(diagnostic push)
#define KOKKOS_NVCC_PRAGMA_SUPPRESS(WARNING) \
KOKKOS_DO_PRAGMA(diag_suppress WARNING)
#define KOKKOS_NVCC_PRAGMA_POP KOKKOS_DO_PRAGMA(diagnostic pop)
#else
#define KOKKOS_NVCC_PRAGMA_PUSH
#define KOKKOS_NVCC_PRAGMA_SUPPRESS(WARNING)
#define KOKKOS_NVCC_PRAGMA_POP
#endif

namespace Kokkos {
namespace Experimental {
namespace Impl {
Expand All @@ -58,39 +41,25 @@ struct StdForEachFunctor {
: m_first(std::move(_first)), m_functor(std::move(_functor)) {}
};

KOKKOS_NVCC_PRAGMA_PUSH
KOKKOS_NVCC_PRAGMA_SUPPRESS(20011)
KOKKOS_NVCC_PRAGMA_SUPPRESS(20014)

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) {
UnaryFunctorType 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);

// run
const auto num_elements = Kokkos::Experimental::distance(first, last);
if constexpr (is_execution_space_v<HandleType>) {
KOKKOS_IF_ON_DEVICE(
(Kokkos::abort("Can't use an execution space on the device!");))
::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();
}
::Kokkos::parallel_for(
label, RangePolicy<HandleType>(handle, 0, num_elements),
StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));
handle.fence("Kokkos::for_each: fence after operation");

return functor;
}

KOKKOS_NVCC_PRAGMA_POP

template <class ExecutionSpace, class IteratorType, class SizeType,
class UnaryFunctorType>
IteratorType for_each_n_exespace_impl(const std::string& label,
Expand All @@ -105,7 +74,7 @@ IteratorType for_each_n_exespace_impl(const std::string& label,
return first;
}

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

return last;
Expand All @@ -114,21 +83,21 @@ 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);
// // 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();
// return functor;
// }
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));
teamHandle.team_barrier();
return functor;
}

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

for_each_impl("", teamHandle, first, last, std::move(functor));
for_each_team_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 aaf00c0

Please sign in to comment.