Skip to content

Commit

Permalink
subset of team level impl of std algorithms
Browse files Browse the repository at this point in the history
Co-authored-by: Cezary Skrzyński
Co-authored-by: Jakub Strzebonski
  • Loading branch information
fnrizzi committed Jul 17, 2023
1 parent 1c20b36 commit 5721206
Show file tree
Hide file tree
Showing 20 changed files with 2,452 additions and 232 deletions.
119 changes: 94 additions & 25 deletions algorithms/src/std_algorithms/Kokkos_AdjacentFind.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -23,71 +23,140 @@
namespace Kokkos {
namespace Experimental {

//
// overload set accepting execution space
//

// overload set1
template <class ExecutionSpace, class IteratorType>
IteratorType adjacent_find(const ExecutionSpace& ex, IteratorType first,
IteratorType last) {
return Impl::adjacent_find_impl("Kokkos::adjacent_find_iterator_api_default",
ex, first, last);
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
IteratorType>
adjacent_find(const ExecutionSpace& ex, IteratorType first, IteratorType last) {
return Impl::adjacent_find_exespace_impl(
"Kokkos::adjacent_find_iterator_api_default", ex, first, last);
}

template <class ExecutionSpace, class IteratorType>
IteratorType adjacent_find(const std::string& label, const ExecutionSpace& ex,
IteratorType first, IteratorType last) {
return Impl::adjacent_find_impl(label, ex, first, last);
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
IteratorType>
adjacent_find(const std::string& label, const ExecutionSpace& ex,
IteratorType first, IteratorType last) {
return Impl::adjacent_find_exespace_impl(label, ex, first, last);
}

template <class ExecutionSpace, class DataType, class... Properties>
template <class ExecutionSpace, class DataType, class... Properties,
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
int> = 0>
auto adjacent_find(const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
namespace KE = ::Kokkos::Experimental;
return Impl::adjacent_find_impl("Kokkos::adjacent_find_view_api_default", ex,
KE::begin(v), KE::end(v));
return Impl::adjacent_find_exespace_impl(
"Kokkos::adjacent_find_view_api_default", ex, KE::begin(v), KE::end(v));
}

template <class ExecutionSpace, class DataType, class... Properties>
template <class ExecutionSpace, class DataType, class... Properties,
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
int> = 0>
auto adjacent_find(const std::string& label, const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
namespace KE = ::Kokkos::Experimental;
return Impl::adjacent_find_impl(label, ex, KE::begin(v), KE::end(v));
return Impl::adjacent_find_exespace_impl(label, ex, KE::begin(v), KE::end(v));
}

// overload set2
template <class ExecutionSpace, class IteratorType, class BinaryPredicateType>
IteratorType adjacent_find(const ExecutionSpace& ex, IteratorType first,
IteratorType last, BinaryPredicateType pred) {
return Impl::adjacent_find_impl("Kokkos::adjacent_find_iterator_api_default",
ex, first, last, pred);
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
IteratorType>
adjacent_find(const ExecutionSpace& ex, IteratorType first, IteratorType last,
BinaryPredicateType pred) {
return Impl::adjacent_find_exespace_impl(
"Kokkos::adjacent_find_iterator_api_default", ex, first, last, pred);
}

template <class ExecutionSpace, class IteratorType, class BinaryPredicateType>
IteratorType adjacent_find(const std::string& label, const ExecutionSpace& ex,
IteratorType first, IteratorType last,
BinaryPredicateType pred) {
return Impl::adjacent_find_impl(label, ex, first, last, pred);
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
IteratorType>
adjacent_find(const std::string& label, const ExecutionSpace& ex,
IteratorType first, IteratorType last, BinaryPredicateType pred) {
return Impl::adjacent_find_exespace_impl(label, ex, first, last, pred);
}

template <class ExecutionSpace, class DataType, class... Properties,
class BinaryPredicateType>
class BinaryPredicateType,
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
int> = 0>
auto adjacent_find(const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v,
BinaryPredicateType pred) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
namespace KE = ::Kokkos::Experimental;
return Impl::adjacent_find_impl("Kokkos::adjacent_find_view_api_default", ex,
KE::begin(v), KE::end(v), pred);
return Impl::adjacent_find_exespace_impl(
"Kokkos::adjacent_find_view_api_default", ex, KE::begin(v), KE::end(v),
pred);
}

template <class ExecutionSpace, class DataType, class... Properties,
class BinaryPredicateType>
class BinaryPredicateType,
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
int> = 0>
auto adjacent_find(const std::string& label, const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v,
BinaryPredicateType pred) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
namespace KE = ::Kokkos::Experimental;
return Impl::adjacent_find_impl(label, ex, KE::begin(v), KE::end(v), pred);
return Impl::adjacent_find_exespace_impl(label, ex, KE::begin(v), KE::end(v),
pred);
}

//
// overload set accepting a team handle
// Note: for now omit the overloads accepting a label
// since they cause issues on device because of the string allocation.
//

// overload set1
template <class TeamHandleType, class IteratorType>
KOKKOS_FUNCTION std::enable_if_t<
::Kokkos::is_team_handle<TeamHandleType>::value, IteratorType>
adjacent_find(const TeamHandleType& teamHandle, IteratorType first,
IteratorType last) {
return Impl::adjacent_find_team_impl(teamHandle, first, last);
}

template <
class TeamHandleType, class DataType, class... Properties,
std::enable_if_t<::Kokkos::is_team_handle<TeamHandleType>::value, int> = 0>
KOKKOS_FUNCTION auto adjacent_find(
const TeamHandleType& teamHandle,
const ::Kokkos::View<DataType, Properties...>& v) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
namespace KE = ::Kokkos::Experimental;
return Impl::adjacent_find_team_impl(teamHandle, KE::begin(v), KE::end(v));
}

// overload set2
template <class TeamHandleType, class IteratorType, class BinaryPredicateType>
KOKKOS_FUNCTION std::enable_if_t<
::Kokkos::is_team_handle<TeamHandleType>::value, IteratorType>
adjacent_find(const TeamHandleType& teamHandle, IteratorType first,
IteratorType last, BinaryPredicateType pred) {
return Impl::adjacent_find_team_impl(teamHandle, first, last, pred);
}

template <
class TeamHandleType, class DataType, class... Properties,
class BinaryPredicateType,
std::enable_if_t<::Kokkos::is_team_handle<TeamHandleType>::value, int> = 0>
KOKKOS_FUNCTION auto adjacent_find(
const TeamHandleType& teamHandle,
const ::Kokkos::View<DataType, Properties...>& v,
BinaryPredicateType pred) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);
namespace KE = ::Kokkos::Experimental;
return Impl::adjacent_find_team_impl(teamHandle, KE::begin(v), KE::end(v),
pred);
}

} // namespace Experimental
Expand Down
68 changes: 51 additions & 17 deletions algorithms/src/std_algorithms/Kokkos_Count.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -23,41 +23,75 @@
namespace Kokkos {
namespace Experimental {

//
// overload set accepting execution space
//
template <class ExecutionSpace, class IteratorType, class T>
typename IteratorType::difference_type count(const ExecutionSpace& ex,
IteratorType first,
IteratorType last,
const T& value) {
return Impl::count_impl("Kokkos::count_iterator_api_default", ex, first, last,
value);
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
typename IteratorType::difference_type>
count(const ExecutionSpace& ex, IteratorType first, IteratorType last,
const T& value) {
return Impl::count_exespace_impl("Kokkos::count_iterator_api_default", ex,
first, last, value);
}

template <class ExecutionSpace, class IteratorType, class T>
typename IteratorType::difference_type count(const std::string& label,
const ExecutionSpace& ex,
IteratorType first,
IteratorType last,
const T& value) {
return Impl::count_impl(label, ex, first, last, value);
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
typename IteratorType::difference_type>
count(const std::string& label, const ExecutionSpace& ex, IteratorType first,
IteratorType last, const T& value) {
return Impl::count_exespace_impl(label, ex, first, last, value);
}

template <class ExecutionSpace, class DataType, class... Properties, class T>
template <class ExecutionSpace, class DataType, class... Properties, class T,
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
int> = 0>
auto count(const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v, const T& value) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

namespace KE = ::Kokkos::Experimental;
return Impl::count_impl("Kokkos::count_view_api_default", ex, KE::cbegin(v),
KE::cend(v), value);
return Impl::count_exespace_impl("Kokkos::count_view_api_default", ex,
KE::cbegin(v), KE::cend(v), value);
}

template <class ExecutionSpace, class DataType, class... Properties, class T>
template <class ExecutionSpace, class DataType, class... Properties, class T,
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
int> = 0>
auto count(const std::string& label, const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v, const T& value) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

namespace KE = ::Kokkos::Experimental;
return Impl::count_impl(label, ex, KE::cbegin(v), KE::cend(v), value);
return Impl::count_exespace_impl(label, ex, KE::cbegin(v), KE::cend(v),
value);
}

//
// overload set accepting a team handle
// Note: for now omit the overloads accepting a label
// since they cause issues on device because of the string allocation.
//

template <class TeamHandleType, class IteratorType, class T>
KOKKOS_FUNCTION
std::enable_if_t<::Kokkos::is_team_handle<TeamHandleType>::value,
typename IteratorType::difference_type>
count(const TeamHandleType& teamHandle, IteratorType first,
IteratorType last, const T& value) {
return Impl::count_team_impl(teamHandle, first, last, value);
}

template <
class TeamHandleType, class DataType, class... Properties, class T,
std::enable_if_t<::Kokkos::is_team_handle<TeamHandleType>::value, int> = 0>
KOKKOS_FUNCTION auto count(const TeamHandleType& teamHandle,
const ::Kokkos::View<DataType, Properties...>& v,
const T& value) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

namespace KE = ::Kokkos::Experimental;
return Impl::count_team_impl(teamHandle, KE::cbegin(v), KE::cend(v), value);
}

} // namespace Experimental
Expand Down
73 changes: 53 additions & 20 deletions algorithms/src/std_algorithms/Kokkos_CountIf.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -23,46 +23,79 @@
namespace Kokkos {
namespace Experimental {

//
// overload set accepting execution space
//
template <class ExecutionSpace, class IteratorType, class Predicate>
typename IteratorType::difference_type count_if(const ExecutionSpace& ex,
IteratorType first,
IteratorType last,
Predicate predicate) {
return Impl::count_if_impl("Kokkos::count_if_iterator_api_default", ex, first,
last, std::move(predicate));
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
typename IteratorType::difference_type>
count_if(const ExecutionSpace& ex, IteratorType first, IteratorType last,
Predicate predicate) {
return Impl::count_if_exespace_impl("Kokkos::count_if_iterator_api_default",
ex, first, last, std::move(predicate));
}

template <class ExecutionSpace, class IteratorType, class Predicate>
typename IteratorType::difference_type count_if(const std::string& label,
const ExecutionSpace& ex,
IteratorType first,
IteratorType last,
Predicate predicate) {
return Impl::count_if_impl(label, ex, first, last, std::move(predicate));
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value,
typename IteratorType::difference_type>
count_if(const std::string& label, const ExecutionSpace& ex, IteratorType first,
IteratorType last, Predicate predicate) {
return Impl::count_if_exespace_impl(label, ex, first, last,
std::move(predicate));
}

template <class ExecutionSpace, class DataType, class... Properties,
class Predicate>
template <
class ExecutionSpace, class DataType, class... Properties, class Predicate,
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value, int> =
0>
auto count_if(const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v,
Predicate predicate) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

namespace KE = ::Kokkos::Experimental;
return Impl::count_if_impl("Kokkos::count_if_view_api_default", ex,
KE::cbegin(v), KE::cend(v), std::move(predicate));
return Impl::count_if_exespace_impl("Kokkos::count_if_view_api_default", ex,
KE::cbegin(v), KE::cend(v),
std::move(predicate));
}

template <class ExecutionSpace, class DataType, class... Properties,
class Predicate>
template <
class ExecutionSpace, class DataType, class... Properties, class Predicate,
std::enable_if_t<::Kokkos::is_execution_space<ExecutionSpace>::value, int> =
0>
auto count_if(const std::string& label, const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v,
Predicate predicate) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

namespace KE = ::Kokkos::Experimental;
return Impl::count_if_impl(label, ex, KE::cbegin(v), KE::cend(v),
std::move(predicate));
return Impl::count_if_exespace_impl(label, ex, KE::cbegin(v), KE::cend(v),
std::move(predicate));
}

//
// overload set accepting team handle
//
template <class ExecutionSpace, class IteratorType, class Predicate>
KOKKOS_FUNCTION
std::enable_if_t<::Kokkos::is_team_handle<ExecutionSpace>::value,
typename IteratorType::difference_type>
count_if(const ExecutionSpace& ex, IteratorType first, IteratorType last,
Predicate predicate) {
return Impl::count_if_team_impl(ex, first, last, std::move(predicate));
}

template <
class ExecutionSpace, class DataType, class... Properties, class Predicate,
std::enable_if_t<::Kokkos::is_team_handle<ExecutionSpace>::value, int> = 0>
KOKKOS_FUNCTION auto count_if(const ExecutionSpace& ex,
const ::Kokkos::View<DataType, Properties...>& v,
Predicate predicate) {
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(v);

namespace KE = ::Kokkos::Experimental;
return Impl::count_if_team_impl(ex, KE::cbegin(v), KE::cend(v),
std::move(predicate));
}

} // namespace Experimental
Expand Down

0 comments on commit 5721206

Please sign in to comment.