Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow functors as reducers for nested team parallel reduce #6921

Open
wants to merge 9 commits into
base: develop
Choose a base branch
from

Conversation

ldh4
Copy link
Contributor

@ldh4 ldh4 commented Apr 5, 2024

Resolves #6317.

This PR allows functors to be used as reducers in nested parallel reduction, specifically for parallel_reduce that takes in TeamThreadRange, ThreadVectorRange or TeamVectorRange.
This capability is already available for reductions with Range/MDRange/Team policies as noted in #6774.

To stay consistent with the existing 'functors as reducers' behavior , changes in this PR expects:

  • ReducerArgument is not a reducer type
  • Functor has init() and join() defined
  • Functor does not necessarily have to satisfy the ReducerConcept (i.e. having typedef reducer = [ReducerStruct] is not required)

@ldh4 ldh4 force-pushed the functors_as_reducers_nested_policies branch from 56716f9 to cd00bd8 Compare April 5, 2024 20:14
@ldh4 ldh4 force-pushed the functors_as_reducers_nested_policies branch from 05b6edf to 684bb69 Compare April 9, 2024 23:53
@masterleinad
Copy link
Contributor

Functor has init() and join() defined

Only join is required. We fall back to the default constructor for init if the reducer/functor doesn't implement it.

Copy link
Contributor

@masterleinad masterleinad left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I only commented on the Cuda backend but the same comments apply to all other backends.

Comment on lines +533 to +535
constexpr bool is_reducer_closure =
functor_analysis_type::has_join_member_function &&
functor_analysis_type::has_init_member_function;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
constexpr bool is_reducer_closure =
functor_analysis_type::has_join_member_function &&
functor_analysis_type::has_init_member_function;
constexpr bool is_reducer_closure =
functor_analysis_type::has_join_member_function;

Comment on lines +537 to +539
using ReducerSelector =
typename Kokkos::Impl::if_c<is_reducer_closure, Closure,
Sum<ValueType>>::type;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are you not using functor_analysis_type::Reducer? Because that type is not compatible with team_reduce?

Either way, ReducerSelector is only used below in the !is_reducer_closure branch. So there is no need for this alias.

Comment on lines +617 to +619
using ReducerSelector =
typename Kokkos::Impl::if_c<is_reducer_closure, Closure,
Sum<ValueType>>::type;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See above.

Comment on lines +303 to +305
if constexpr (is_reducer_v<ReducerType>) {
reducer.reference() = tmp2;
}))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why aren't we using final here if defined?

Comment on lines +737 to +739
using ReducerSelector =
typename Kokkos::Impl::if_c<is_reducer_closure, Closure,
Sum<ValueType>>::type;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See above.

Comment on lines +1797 to +1798
KOKKOS_INLINE_FUNCTION
void final(value_type &) const {}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should test with a non-trivial final as well.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants