Skip to content

Commit

Permalink
team-level std algos: part 12 (kokkos#6350)
Browse files Browse the repository at this point in the history 
Add team-level exclusive scan and transform exclusive scan algorithms
  • Loading branch information
@fnrizzi
fnrizzi committed Oct 4, 2023
1 parent 41cf2e5 commit 60e4d13
Show file tree
Hide file tree
Showing 10 changed files with 1,158 additions and 297 deletions.
243 changes: 175 additions & 68 deletions algorithms/src/std_algorithms/Kokkos_ExclusiveScan.hpp
View file

Large diffs are not rendered by default.

130 changes: 92 additions & 38 deletions algorithms/src/std_algorithms/Kokkos_TransformExclusiveScan.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,44 +23,52 @@
namespace Kokkos {
namespace Experimental {

template <class ExecutionSpace, class InputIteratorType,
class OutputIteratorType, class ValueType, class BinaryOpType,
class UnaryOpType>
std::enable_if_t<::Kokkos::Experimental::Impl::are_iterators<
InputIteratorType, OutputIteratorType>::value,
OutputIteratorType>
transform_exclusive_scan(const ExecutionSpace& ex, InputIteratorType first,
InputIteratorType last, OutputIteratorType first_dest,
ValueType init_value, BinaryOpType binary_op,
UnaryOpType unary_op) {
//
// overload set accepting execution space
//
template <typename ExecutionSpace, typename InputIteratorType,
typename OutputIteratorType, typename ValueType,
typename BinaryOpType, typename UnaryOpType,
std::enable_if_t<
Impl::are_iterators_v<InputIteratorType, OutputIteratorType>&& ::
Kokkos::is_execution_space_v<ExecutionSpace>,
int> = 0>
OutputIteratorType transform_exclusive_scan(
const ExecutionSpace& ex, InputIteratorType first, InputIteratorType last,
OutputIteratorType first_dest, ValueType init_value, BinaryOpType binary_op,
UnaryOpType unary_op) {
Impl::static_assert_is_not_openmptarget(ex);
static_assert(std::is_move_constructible<ValueType>::value,
static_assert(std::is_move_constructible_v<ValueType>,
"ValueType must be move constructible.");
return Impl::transform_exclusive_scan_impl(
return Impl::transform_exclusive_scan_exespace_impl(
"Kokkos::transform_exclusive_scan_custom_functors_iterator_api", ex,
first, last, first_dest, init_value, binary_op, unary_op);
first, last, first_dest, std::move(init_value), binary_op, unary_op);
}

template <class ExecutionSpace, class InputIteratorType,
class OutputIteratorType, class ValueType, class BinaryOpType,
class UnaryOpType>
std::enable_if_t<::Kokkos::Experimental::Impl::are_iterators<
InputIteratorType, OutputIteratorType>::value,
OutputIteratorType>
transform_exclusive_scan(const std::string& label, const ExecutionSpace& ex,
InputIteratorType first, InputIteratorType last,
OutputIteratorType first_dest, ValueType init_value,
BinaryOpType binary_op, UnaryOpType unary_op) {
template <typename ExecutionSpace, typename InputIteratorType,
typename OutputIteratorType, typename ValueType,
typename BinaryOpType, typename UnaryOpType,
std::enable_if_t<
Impl::are_iterators_v<InputIteratorType, OutputIteratorType>&& ::
Kokkos::is_execution_space_v<ExecutionSpace>,
int> = 0>
OutputIteratorType transform_exclusive_scan(
const std::string& label, const ExecutionSpace& ex, InputIteratorType first,
InputIteratorType last, OutputIteratorType first_dest, ValueType init_value,
BinaryOpType binary_op, UnaryOpType unary_op) {
Impl::static_assert_is_not_openmptarget(ex);
static_assert(std::is_move_constructible<ValueType>::value,
static_assert(std::is_move_constructible_v<ValueType>,
"ValueType must be move constructible.");
return Impl::transform_exclusive_scan_impl(label, ex, first, last, first_dest,
init_value, binary_op, unary_op);
return Impl::transform_exclusive_scan_exespace_impl(
label, ex, first, last, first_dest, std::move(init_value), binary_op,
unary_op);
}

template <class ExecutionSpace, class DataType1, class... Properties1,
class DataType2, class... Properties2, class ValueType,
class BinaryOpType, class UnaryOpType>
template <
typename ExecutionSpace, typename DataType1, typename... Properties1,
typename DataType2, typename... Properties2, typename ValueType,
typename BinaryOpType, typename UnaryOpType,
std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
auto transform_exclusive_scan(
const ExecutionSpace& ex,
const ::Kokkos::View<DataType1, Properties1...>& view_from,
Expand All @@ -69,18 +77,20 @@ auto transform_exclusive_scan(
Impl::static_assert_is_not_openmptarget(ex);
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_from);
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_dest);
static_assert(std::is_move_constructible<ValueType>::value,
static_assert(std::is_move_constructible_v<ValueType>,
"ValueType must be move constructible.");
namespace KE = ::Kokkos::Experimental;
return Impl::transform_exclusive_scan_impl(
return Impl::transform_exclusive_scan_exespace_impl(
"Kokkos::transform_exclusive_scan_custom_functors_view_api", ex,
KE::cbegin(view_from), KE::cend(view_from), KE::begin(view_dest),
init_value, binary_op, unary_op);
std::move(init_value), binary_op, unary_op);
}

template <class ExecutionSpace, class DataType1, class... Properties1,
class DataType2, class... Properties2, class ValueType,
class BinaryOpType, class UnaryOpType>
template <
typename ExecutionSpace, typename DataType1, typename... Properties1,
typename DataType2, typename... Properties2, typename ValueType,
typename BinaryOpType, typename UnaryOpType,
std::enable_if_t<::Kokkos::is_execution_space_v<ExecutionSpace>, int> = 0>
auto transform_exclusive_scan(
const std::string& label, const ExecutionSpace& ex,
const ::Kokkos::View<DataType1, Properties1...>& view_from,
Expand All @@ -89,12 +99,56 @@ auto transform_exclusive_scan(
Impl::static_assert_is_not_openmptarget(ex);
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_from);
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_dest);
static_assert(std::is_move_constructible<ValueType>::value,
static_assert(std::is_move_constructible_v<ValueType>,
"ValueType must be move constructible.");
namespace KE = ::Kokkos::Experimental;
return Impl::transform_exclusive_scan_impl(
return Impl::transform_exclusive_scan_exespace_impl(
label, ex, KE::cbegin(view_from), KE::cend(view_from),
KE::begin(view_dest), init_value, binary_op, unary_op);
KE::begin(view_dest), std::move(init_value), binary_op, unary_op);
}

//
// 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 <typename TeamHandleType, typename InputIteratorType,
typename OutputIteratorType, typename ValueType,
typename BinaryOpType, typename UnaryOpType,
std::enable_if_t<
Impl::are_iterators_v<InputIteratorType, OutputIteratorType>&& ::
Kokkos::is_team_handle_v<TeamHandleType>,
int> = 0>
KOKKOS_FUNCTION OutputIteratorType transform_exclusive_scan(
const TeamHandleType& teamHandle, InputIteratorType first,
InputIteratorType last, OutputIteratorType first_dest, ValueType init_value,
BinaryOpType binary_op, UnaryOpType unary_op) {
Impl::static_assert_is_not_openmptarget(teamHandle);
static_assert(std::is_move_constructible_v<ValueType>,
"ValueType must be move constructible.");
return Impl::transform_exclusive_scan_team_impl(
teamHandle, first, last, first_dest, std::move(init_value), binary_op,
unary_op);
}

template <typename TeamHandleType, typename DataType1, typename... Properties1,
typename DataType2, typename... Properties2, typename ValueType,
typename BinaryOpType, typename UnaryOpType,
std::enable_if_t<::Kokkos::is_team_handle_v<TeamHandleType>, int> = 0>
KOKKOS_FUNCTION auto transform_exclusive_scan(
const TeamHandleType& teamHandle,
const ::Kokkos::View<DataType1, Properties1...>& view_from,
const ::Kokkos::View<DataType2, Properties2...>& view_dest,
ValueType init_value, BinaryOpType binary_op, UnaryOpType unary_op) {
Impl::static_assert_is_not_openmptarget(teamHandle);
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_from);
Impl::static_assert_is_admissible_to_kokkos_std_algorithms(view_dest);
static_assert(std::is_move_constructible_v<ValueType>,
"ValueType must be move constructible.");
namespace KE = ::Kokkos::Experimental;
return Impl::transform_exclusive_scan_team_impl(
teamHandle, KE::cbegin(view_from), KE::cend(view_from),
KE::begin(view_dest), std::move(init_value), binary_op, unary_op);
}

} // namespace Experimental
Expand Down

0 comments on commit 60e4d13

Please sign in to comment.